web_hghl_server/pub/esmanager/client.go

package esmanager

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"

	"github.com/elastic/go-elasticsearch/v7/esapi"
)

const (
	AnalyzerIkSmart   = "ik_smart"
	AnalyzerIkMaxWord = "ik_max_word"
	NumberOfShards    = 1
	NumberOfReplicas  = 1
)

func (m *Manager) InitIndices(indexName string, setting M) error {
	res, err := m.client.Indices.Create(indexName, func(request *esapi.IndicesCreateRequest) {
		request.Body = marshalSearchM(setting)
	})
	defer res.Body.Close()
	return err
}

// GetTokens 获取分词
func (m *Manager) Tokens(index string, bind interface{}, query M, o ...func(*esapi.IndicesAnalyzeRequest)) error {
	o = append(o,
		m.client.Indices.Analyze.WithContext(context.Background()),
		m.client.Indices.Analyze.WithIndex(index),
		m.client.Indices.Analyze.WithBody(marshalSearchM(query)),
		m.client.Indices.Analyze.WithPretty(),
	)

	res, err := m.client.Indices.Analyze(o...)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	err = unmarshalAnalyzeBody(bind, res)
	return err
}

// search 搜索
func (m *Manager) Search(index string, bind interface{}, query M, o ...func(*esapi.SearchRequest)) error {
	marshalSearchM(query)
	o = append(o,
		m.client.Search.WithContext(context.Background()),
		m.client.Search.WithIndex(index),
		m.client.Search.WithBody(marshalSearchM(query)),
		m.client.Search.WithTrackTotalHits(true),
		m.client.Search.WithPretty(),
	)
	res, err := m.client.Search(o...)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	err = unmarshalBody(bind, res)
	return err
}

// 批量插入
func (m *Manager) Bulk(index string, source M, o ...func(*esapi.BulkRequest)) error {
	o = append(o,
		m.client.Bulk.WithIndex(index),
		m.client.Bulk.WithContext(context.Background()),
	)
	res, err := m.client.Bulk(marshalBulkM(source), o...)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return errors.New(fmt.Sprintf("%+v", res))
	}
	var bulkResp struct {
		Errors bool `json:"errors"`
		Items  []map[string]struct {
			Status int         `json:"status"`
			Error  interface{} `json:"error"`
		} `json:"items"`
	}
	if err := json.NewDecoder(res.Body).Decode(&bulkResp); err != nil {
		return err
	}
	if bulkResp.Errors {
		errs := make([]string, 0, 3)
		for _, item := range bulkResp.Items {
			for op, result := range item {
				if result.Status >= http.StatusBadRequest {
					errs = append(errs, fmt.Sprintf("%s status=%d error=%v", op, result.Status, result.Error))
					if len(errs) >= 3 {
						return errors.New(fmt.Sprintf("elastic bulk item errors: %v", errs))
					}
				}
			}
		}
		return errors.New(fmt.Sprintf("elastic bulk item errors: %v", errs))
	}
	return nil
}

func (m *Manager) Refresh(index string, o ...func(*esapi.IndicesRefreshRequest)) error {
	o = append(o,
		m.client.Indices.Refresh.WithContext(context.Background()),
		m.client.Indices.Refresh.WithIndex(index),
	)
	res, err := m.client.Indices.Refresh(o...)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode >= http.StatusBadRequest {
		return errors.New(fmt.Sprintf("%+v", res))
	}
	return nil
}

// Count  获取数量
func (m *Manager) Count(index string, query M, o ...func(*esapi.CountRequest)) (int64, error) {
	o = append(o,
		m.client.Count.WithIndex(index),
		m.client.Count.WithBody(marshalSearchM(query)),
		m.client.Count.WithContext(context.Background()),
	)
	res, err := m.client.Count(o...)
	if err != nil {
		return 0, err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return 0, errors.New(fmt.Sprintf("%+v", res))
	}
	type Bind struct {
		Count int64 `json:"count"`
	}
	var bind = Bind{}
	err = unmarshalStruct(&bind, res)
	if err != nil {
		return 0, err
	}
	return bind.Count, err
}

func (m *Manager) Get(index string, id string, bind interface{}, o ...func(*esapi.GetRequest)) error {
	o = append(o,
		m.client.Get.WithContext(context.Background()),
		m.client.Get.WithPretty(),
	)
	res, err := m.client.Get(index, id, o...)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return errors.New(fmt.Sprintf("%+v", res))
	}
	return unmarshalBody(bind, res)
}

func (m *Manager) Add(index string, id string, doc interface{}, o ...func(*esapi.IndexRequest)) error {
	o = append(o,
		m.client.Index.WithDocumentID(id),
		m.client.Index.WithContext(context.Background()),
		m.client.Index.WithPretty(),
		m.client.Index.WithRefresh("true"),
	)
	res, err := m.client.Index(index, marshalStruct(doc), o...)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	return nil
}

func (m *Manager) Del(index string, id string, o ...func(*esapi.DeleteRequest)) error {
	o = append(o,
		m.client.Delete.WithContext(context.Background()),
		m.client.Delete.WithPretty(),
		m.client.Delete.WithRefresh("true"),
	)
	res, err := m.client.Delete(index, id, o...)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	return nil
}

func (m *Manager) DelMany(indexs []string, query M, o ...func(*esapi.DeleteByQueryRequest)) error {
	o = append(o,
		m.client.DeleteByQuery.WithContext(context.Background()),
		m.client.DeleteByQuery.WithPretty(),
		m.client.DeleteByQuery.WithRefresh(true),
	)
	res, err := m.client.DeleteByQuery(indexs, marshalSearchM(query), o...)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	return nil
}