huangguo_server/common/elastic/elasticClient.go

package elastic

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"

	"91porn-server/common/log"

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

type M map[string]interface{}

type A []map[string]interface{}

type MGetBody struct {
	ID string `json:"_id"`
}

// ElasticResp 响应
type Response = esapi.Response

type Options struct {
	Address             []string
	MaxIdleConnsPerHost int           //每个client z最多允许的空闲连接数
	IdleConnTimeout     time.Duration //空闲连接超时时间
	UserName            string
	PassWord            string
}

// 查看集群信息
func (c *Client) Info() (*Response, error) {
	resp, err := c.Client.Info(c.Client.Info.WithContext(context.Background()))
	if err != nil {
		return nil, err
	}
	resp.Body.Close()
	return resp, err
}

// 初始化索引
func (c *Client) CreateIndices(index string, setting M) error {
	resp, err := c.Client.Indices.Create(index, func(request *esapi.IndicesCreateRequest) {
		request.Body = marshalSearchM(setting)
	})
	if err != nil {
		return err
	}
	resp.Body.Close()
	return nil
}

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

func (c *Client) SearchWithTotal(index string, bind interface{}, query M, o ...func(*esapi.SearchRequest)) error {
	o = append(o,
		c.Client.Search.WithContext(context.Background()),
		c.Client.Search.WithIndex(index),
		c.Client.Search.WithBody(marshalSearchM(query)),
		c.Client.Search.WithTrackTotalHits(true),
		c.Client.Search.WithPretty(),
	)
	resp, err := c.Client.Search(o...)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	return unmarshalBodyWithTotal(bind, resp)
}

// Aggregate Aggregate 聚合
func (c *Client) Aggregate(index string, bind interface{}, query M, o ...func(*esapi.SearchRequest)) (int, error) {
	o = append(o,
		c.Client.Search.WithContext(context.Background()),
		c.Client.Search.WithIndex(index),
		c.Client.Search.WithBody(marshalSearchM(query)),
		c.Client.Search.WithTrackTotalHits(true),
		c.Client.Search.WithPretty(),
	)
	resp, err := c.Client.Search(o...)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	return resp.StatusCode, unmarshalAggregateBody(bind, resp)
}

// Get 根据ID搜索
func (c *Client) Get(index string, bind interface{}, id string, o ...func(*esapi.GetRequest)) (int, error) {
	o = append(o,
		c.Client.Get.WithContext(context.Background()),
		c.Client.Get.WithPretty(),
	)
	resp, err := c.Client.Get(index, id, o...)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	return resp.StatusCode, unmarshalBody(bind, resp)
}

// MGet MGet 批量查询
func (c *Client) MGet(index string, bind interface{}, ids []string, o ...func(*esapi.MgetRequest)) error {
	m := make([]MGetBody, len(ids))
	for i, v := range ids {
		m[i] = MGetBody{ID: v}
	}
	o = append(o,
		c.Client.Mget.WithIndex(index),
		c.Client.Mget.WithContext(context.Background()),
	)
	resp, err := c.Client.Mget(marshalSearchM(M{"docs": m}), o...)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("%+v", resp)
	}
	return unmarshalBody(bind, resp)
}

// BulkDelete 函数
func (c *Client) BulkDelete(index string, source M, o ...func(*esapi.BulkRequest)) error {
	o = append(o,
		c.Client.Bulk.WithIndex(index),
		c.Client.Bulk.WithContext(context.Background()),
		c.Client.Bulk.WithRefresh("true"),
	)
	resp, err := c.Client.Bulk(marshalDeleteBulkM(source), o...)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.IsError() {
		return fmt.Errorf("error response: %s", resp.String())
	}
	return nil
}

func marshalDeleteBulkM(source M) io.Reader {
	var buf bytes.Buffer
	for k, _ := range source {
		var meta = []byte(fmt.Sprintf(`{ "delete" : { "_id" : "%s" } }%s`, k, "\n"))
		buf.Grow(len(meta))
		buf.Write(meta)
	}
	return bytes.NewReader(buf.Bytes())
}

// 批量插入
func (c *Client) Bulk(index string, source M, o ...func(*esapi.BulkRequest)) error {
	o = append(o,
		c.Client.Bulk.WithIndex(index),
		c.Client.Bulk.WithContext(context.Background()),
	)
	resp, err := c.Client.Bulk(marshalBulkM(source), o...)
	if err != nil {
		return err
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("%+v", resp)
	}
	return nil
}

// BulkChecked 在 HTTP 成功后继续检查每条写入结果，供需要可靠推进同步进度的任务使用。
func (c *Client) BulkChecked(index string, source M) error {
	if len(source) == 0 {
		return nil
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	resp, err := c.Client.Bulk(marshalBulkM(source),
		c.Client.Bulk.WithIndex(index), c.Client.Bulk.WithContext(ctx))
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("bulk %s: HTTP %d", index, resp.StatusCode)
	}
	var result struct {
		Errors bool `json:"errors"`
		Items  []map[string]struct {
			ID     string          `json:"_id"`
			Status int             `json:"status"`
			Error  json.RawMessage `json:"error"`
		} `json:"items"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return fmt.Errorf("decode bulk %s response: %w", index, err)
	}
	if len(result.Items) != len(source) {
		return fmt.Errorf("bulk %s: expected %d results, got %d", index, len(source), len(result.Items))
	}
	for _, item := range result.Items {
		entry, ok := item["index"]
		if !ok || len(item) != 1 {
			return fmt.Errorf("bulk %s: missing index result", index)
		}
		if entry.Status < 200 || entry.Status >= 300 || (len(entry.Error) > 0 && string(entry.Error) != "null") {
			return fmt.Errorf("bulk %s: document %s failed, status %d", index, entry.ID, entry.Status)
		}
	}
	if result.Errors {
		return fmt.Errorf("bulk %s: response contains errors", index)
	}
	return nil
}

// Count  获取数量
func (c *Client) Count(index string, query M, o ...func(*esapi.CountRequest)) (cnt int, err error) {
	o = append(o,
		c.Client.Count.WithIndex(index),
		c.Client.Count.WithBody(marshalSearchM(query)),
		c.Client.Count.WithContext(context.Background()),
	)
	resp, err := c.Client.Count(o...)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return 0, fmt.Errorf("%+v", resp)
	}
	var bind struct {
		Count int `json:"count"`
	}
	return bind.Count, unmarshalBody(&bind, resp)
}

func (c *Client) Ping() error {
	resp, err := c.Client.Ping(c.Client.Ping.WithContext(context.Background()))
	if err != nil {
		log.Warn("elasticSarch ping error ", log.E(err))
		return err
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		log.Warn("elasticSarch ping failed ", log.Any("status code", resp.StatusCode))
		return errors.New("elasticSarch ping failed")
	}
	return nil
}