web_hgdjb_server/pub/esmanager/client.go

package esmanager

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"

	"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 {
	existsRes, err := m.client.Indices.Exists([]string{indexName})
	if err != nil {
		return err
	}
	defer existsRes.Body.Close()
	if existsRes.StatusCode == http.StatusOK {
		return nil
	}
	if existsRes.StatusCode != http.StatusNotFound {
		body, _ := io.ReadAll(existsRes.Body)
		return errors.New(fmt.Sprintf("%+v body:%s", existsRes, string(body)))
	}
	res, err := m.client.Indices.Create(indexName, func(request *esapi.IndicesCreateRequest) {
		request.Body = marshalSearchM(setting)
	})
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode == http.StatusOK || res.StatusCode == http.StatusCreated {
		return nil
	}
	body, _ := io.ReadAll(res.Body)
	if res.StatusCode == http.StatusBadRequest && strings.Contains(string(body), "resource_already_exists_exception") {
		return nil
	}
	return errors.New(fmt.Sprintf("%+v body:%s", res, string(body)))
}

// 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()
	if res.StatusCode != http.StatusOK {
		return errors.New(fmt.Sprintf("%+v", res))
	}
	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()),
		m.client.Bulk.WithRefresh("true"),
	)
	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))
	}
	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
}