Initial commit

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-15 13:57:10 +08:00
commit 8679200f41
1897 changed files with 257900 additions and 0 deletions
@@ -0,0 +1,267 @@
+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
+}