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 }