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,110 @@
+package elastic
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+
+	"github.com/elastic/go-elasticsearch/v8"
+)
+
+var (
+	AnalyzerIkSmart   = "ik_smart"
+	AnalyzerIkMaxWord = "ik_max_word"
+	NumberOfShards    = 1
+	NumberOfReplicas  = 1
+)
+
+// Client
+type Client struct {
+	Client *elasticsearch.Client
+}
+
+func unmarshalAggregateBody(bind interface{}, resp *Response) error {
+	var v map[string]interface{}
+	if bind == nil { //传入空表示只执行 defer
+		return nil
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&v); err != nil {
+		return err
+	}
+	if len(v) > 0 {
+		if v, ok := v["aggregations"].(map[string]interface{}); ok {
+			b, _ := json.Marshal(v)
+			_ = json.Unmarshal(b, &bind)
+		}
+	}
+	return nil
+}
+
+func unmarshalBodyWithTotal(bind interface{}, resp *Response) error {
+	var v map[string]interface{}
+	if bind == nil { //传入空表示只执行 defer
+		return nil
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&v); err != nil {
+		return err
+	}
+	if len(v) > 0 {
+		if h, ok := v["hits"].(map[string]interface{}); ok {
+			b, err := json.Marshal(h)
+			if err != nil {
+				return err
+			}
+			if err = json.Unmarshal(b, &bind); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func unmarshalBody(bind interface{}, resp *Response) error {
+	var v map[string]interface{}
+	if bind == nil { //传入空表示只执行 defer
+		return nil
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&v); err != nil {
+		return err
+	}
+	if len(v) > 0 {
+		if h, ok := v["hits"].(map[string]interface{}); ok {
+			if hh, ok := h["hits"].([]interface{}); ok {
+				b, _ := json.Marshal(hh)
+				_ = json.Unmarshal(b, &bind)
+			}
+		} else if _, ok := v["docs"].([]interface{}); ok {
+			b, _ := json.Marshal(v["docs"])
+			_ = json.Unmarshal(b, &bind)
+		} else if _, ok := v["_source"].(map[string]interface{}); ok {
+			b, _ := json.Marshal(v)
+			_ = json.Unmarshal(b, &bind)
+		}
+	}
+	return nil
+}
+
+func marshalSearchM(query M) io.Reader {
+	var buf bytes.Buffer
+	if len(query) == 0 {
+		return nil
+	}
+	if err := json.NewEncoder(&buf).Encode(query); err != nil {
+		return nil
+	}
+	return bytes.NewReader(buf.Bytes())
+}
+
+func marshalBulkM(source M) io.Reader {
+	var buf bytes.Buffer
+	for k, v := range source {
+		var meta = []byte(fmt.Sprintf(`{ "index" : { "_id" : "%s" } }%s`, k, "\n"))
+		var data, _ = json.Marshal(v)
+		data = append(data, "\n"...)
+		buf.Grow(len(meta) + len(data))
+		buf.Write(meta)
+		buf.Write(data)
+	}
+	return bytes.NewReader(buf.Bytes())
+}
@@ -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
+}
@@ -0,0 +1,39 @@
+package elastic
+
+import (
+	"errors"
+
+	"91porn-server/common/log"
+
+	"github.com/elastic/go-elasticsearch/v8"
+)
+
+var esClient *Client
+
+// Client 获取elastic客户端
+func InitElastic(opt Options) (*Client, error) {
+	if len(opt.Address) == 0 {
+		return nil, errors.New("elasticSearch  address is empty")
+	}
+	cfg := elasticsearch.Config{
+		Addresses: opt.Address,
+		Username:  opt.UserName,
+		Password:  opt.PassWord,
+	}
+	es, err := elasticsearch.NewClient(cfg)
+	if err != nil {
+		log.Error("create elasticSearch client occour error", log.E(err))
+		return nil, err
+	}
+	esClient = &Client{Client: es}
+	if err = esClient.Ping(); err != nil {
+		log.Warn("elasticSarch ping failed", log.E(err))
+		return nil, err
+	}
+	log.Info("init elasticSearch successful")
+	return esClient, nil
+}
+
+func Init() *Client {
+	return esClient
+}