gitee.com/h79/goutils@v1.22.10/dao/db/elastic.go (about)

     1  /**
     2   * @author huqiuyun
     3   * 数据Elasticsearch
     4   */
     5  package db
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	commonconfig "gitee.com/h79/goutils/common/config"
    11  	"gitee.com/h79/goutils/common/logger"
    12  	"gitee.com/h79/goutils/common/result"
    13  	"gitee.com/h79/goutils/dao/config"
    14  	"gitee.com/h79/goutils/dao/util"
    15  	"github.com/olivere/elastic/v7"
    16  	"go.uber.org/zap"
    17  	"time"
    18  )
    19  
    20  type esClient struct {
    21  	Client *elastic.Client
    22  }
    23  
    24  var _ EsDatabase = (*elasticGroup)(nil)
    25  
    26  type elasticGroup struct {
    27  	selectName string
    28  	esMap      map[string]*esClient
    29  }
    30  
    31  func NewES(conf []config.Elastic) (EsDatabase, error) {
    32  	logger.L().Info("Elastic", zap.Any("config", conf))
    33  	ess := &elasticGroup{esMap: map[string]*esClient{}}
    34  	for _, cfg := range conf {
    35  		if len(cfg.Host) <= 0 {
    36  			continue
    37  		}
    38  		interval := cfg.Interval
    39  		if interval <= 0 {
    40  			interval = 10
    41  		} else if interval > 3600 {
    42  			interval = 3600
    43  		}
    44  		interval = interval * time.Second
    45  
    46  		url := cfg.Host[0]
    47  		errLog := &esLogger{
    48  			LogLevel: logger.ErrorLevel,
    49  			Level:    cfg.Logger.LogLevel,
    50  		}
    51  		infoLog := &esLogger{
    52  			LogLevel: logger.InfoLevel,
    53  			Level:    cfg.Logger.LogLevel,
    54  		}
    55  		traceLog := &esLogger{
    56  			LogLevel: logger.DebugLevel,
    57  			Level:    cfg.Logger.LogLevel,
    58  		}
    59  		var client, err = elastic.NewClient(
    60  			elastic.SetURL(cfg.Host...),
    61  			elastic.SetBasicAuth(cfg.User, cfg.Pwd),
    62  			elastic.SetHealthcheckInterval(interval),
    63  			elastic.SetSniff(cfg.Sniff),
    64  			elastic.SetErrorLog(errLog),
    65  			elastic.SetInfoLog(infoLog),
    66  			elastic.SetTraceLog(traceLog))
    67  
    68  		if err != nil {
    69  			util.Alarm(result.ErrEsClientInternal, "", fmt.Sprintf("New Elastic(%s)", cfg.Name), err)
    70  			return nil, err
    71  		}
    72  		_, _, err = client.Ping(url).Do(context.Background())
    73  		if err != nil {
    74  			util.Alarm(result.ErrEsPingInternal, "", fmt.Sprintf("Ping Elastic(%s)", cfg.Name), err)
    75  			return nil, err
    76  		}
    77  		ver, err := client.ElasticsearchVersion(url)
    78  		logger.L().Info("Elastic", zap.String("Name", cfg.Name), zap.String("Version", ver))
    79  		ess.esMap[cfg.Name] = &esClient{client}
    80  		if commonconfig.RegisterConfig != nil {
    81  			commonconfig.RegisterConfig("ES:Err|"+cfg.Name, errLog.handlerConfig)
    82  			commonconfig.RegisterConfig("ES:Info|"+cfg.Name, infoLog.handlerConfig)
    83  			commonconfig.RegisterConfig("ES:Trace|"+cfg.Name, traceLog.handlerConfig)
    84  		}
    85  	}
    86  	return ess, nil
    87  }
    88  func (conns *elasticGroup) Get(name string) (*elastic.Client, string, error) {
    89  	if len(name) == 0 {
    90  		name = conns.selectName
    91  	}
    92  	es, exists := conns.esMap[name]
    93  	if exists == true {
    94  		return es.Client, name, nil
    95  	}
    96  	return nil, "", result.Error(result.ErrNotFound, "Not found")
    97  }
    98  
    99  func (conns *elasticGroup) Do(name string, call func(es *elastic.Client) error) error {
   100  	if len(name) == 0 {
   101  		name = conns.selectName
   102  	}
   103  	if es, exists := conns.esMap[name]; exists == true {
   104  		return call(es.Client)
   105  	}
   106  	return result.Error(result.ErrNotFound, "Not found")
   107  }
   108  
   109  func (conns *elasticGroup) Close(name string) error {
   110  	if len(name) == 0 {
   111  		name = conns.selectName
   112  	}
   113  	es, exists := conns.esMap[name]
   114  	if exists == true {
   115  		es.Client.Stop()
   116  		delete(conns.esMap, name)
   117  	}
   118  	return nil
   119  }
   120  
   121  func (conns *elasticGroup) CloseAll() {
   122  	for _, es := range conns.esMap {
   123  		es.Client.Stop()
   124  	}
   125  	conns.esMap = nil
   126  }
   127  
   128  func (conns *elasticGroup) Select(name string) {
   129  	conns.selectName = name
   130  }
   131  
   132  func (conns *elasticGroup) GetSelector() string {
   133  	return conns.selectName
   134  }