feat: manually mirror opencoze's code from bytedance

Change-Id: I09a73aadda978ad9511264a756b2ce51f5761adf

backend/infra/impl/es/es7.go

@@ -0,0 +1,371 @@
+/*
+ * Copyright 2025 coze-dev Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package es
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/elastic/go-elasticsearch/v7"
+	"github.com/elastic/go-elasticsearch/v7/esapi"
+	"github.com/elastic/go-elasticsearch/v7/esutil"
+
+	"github.com/coze-dev/coze-studio/backend/infra/contract/es"
+	"github.com/coze-dev/coze-studio/backend/pkg/lang/conv"
+	"github.com/coze-dev/coze-studio/backend/pkg/lang/ptr"
+	"github.com/coze-dev/coze-studio/backend/pkg/logs"
+)
+
+type es7Client struct {
+	esClient *elasticsearch.Client
+}
+
+func newES7() (Client, error) {
+	esAddr := os.Getenv("ES_ADDR")
+	esUsername := os.Getenv("ES_USERNAME")
+	esPassword := os.Getenv("ES_PASSWORD")
+
+	esClient, err := elasticsearch.NewClient(elasticsearch.Config{
+		Addresses: []string{esAddr},
+		Username:  esUsername,
+		Password:  esPassword,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &es7Client{esClient: esClient}, nil
+}
+
+func (c *es7Client) Create(ctx context.Context, index, id string, document any) error {
+	body, err := json.Marshal(document)
+	if err != nil {
+		return err
+	}
+
+	req := esapi.IndexRequest{
+		Index:      index,
+		DocumentID: id,
+		Body:       bytes.NewReader(body),
+		Refresh:    "true",
+	}
+
+	logs.CtxDebugf(ctx, "[Create] req : %s", conv.DebugJsonToStr(req))
+	_, err = req.Do(ctx, c.esClient)
+	return err
+}
+
+func (c *es7Client) Update(ctx context.Context, index, id string, document any) error {
+	bodyMap := map[string]any{"doc": document}
+	body, err := json.Marshal(bodyMap)
+	if err != nil {
+		return err
+	}
+	req := esapi.UpdateRequest{
+		Index:      index,
+		DocumentID: id,
+		Body:       bytes.NewReader(body),
+	}
+
+	logs.CtxDebugf(ctx, "[Update] req : %s", conv.DebugJsonToStr(req))
+
+	_, err = req.Do(ctx, c.esClient)
+	return err
+}
+
+func (c *es7Client) Delete(ctx context.Context, index, id string) error {
+	req := esapi.DeleteRequest{
+		Index:      index,
+		DocumentID: id,
+	}
+
+	logs.CtxDebugf(ctx, "[Delete] req : %s", conv.DebugJsonToStr(req))
+
+	_, err := req.Do(ctx, c.esClient)
+	return err
+}
+
+func (c *es7Client) Exists(ctx context.Context, index string) (bool, error) {
+	req := esapi.IndicesExistsRequest{Index: []string{index}}
+	logs.CtxDebugf(ctx, "[Exists] req : %s", conv.DebugJsonToStr(req))
+
+	res, err := req.Do(ctx, c.esClient)
+	if err != nil {
+		return false, err
+	}
+	defer res.Body.Close()
+	return res.StatusCode == 200, nil
+}
+
+func (c *es7Client) CreateIndex(ctx context.Context, index string, properties map[string]any) error {
+	mapping := map[string]any{
+		"mappings": map[string]any{
+			"properties": properties,
+		},
+	}
+
+	body, err := json.Marshal(mapping)
+	if err != nil {
+		return err
+	}
+
+	req := esapi.IndicesCreateRequest{
+		Index: index,
+		Body:  bytes.NewReader(body),
+	}
+
+	logs.CtxDebugf(ctx, "[CreateIndex] req : %s", conv.DebugJsonToStr(req))
+	_, err = req.Do(ctx, c.esClient)
+	return err
+}
+
+func (c *es7Client) DeleteIndex(ctx context.Context, index string) error {
+	req := esapi.IndicesDeleteRequest{
+		Index:             []string{index},
+		IgnoreUnavailable: ptr.Of(true),
+	}
+
+	logs.CtxDebugf(ctx, "[DeleteIndex] req : %s", conv.DebugJsonToStr(req))
+	_, err := req.Do(ctx, c.esClient)
+	return err
+}
+
+func (c *es7Client) Search(ctx context.Context, index string, req *Request) (*Response, error) {
+	queryBody := map[string]any{}
+	if q := c.query2ESQuery(req.Query); q != nil {
+		queryBody["query"] = q
+	}
+	if req.Size != nil {
+		queryBody["size"] = *req.Size
+	}
+	if req.MinScore != nil {
+		queryBody["min_score"] = *req.MinScore
+	}
+	if len(req.Sort) > 0 {
+		var sorts []map[string]any
+		for _, s := range req.Sort {
+			order := "asc"
+			if !s.Asc {
+				order = "desc"
+			}
+			sorts = append(sorts, map[string]any{
+				s.Field: map[string]string{"order": order},
+			})
+		}
+		queryBody["sort"] = sorts
+	}
+
+	if req.From != nil {
+		queryBody["from"] = *req.From
+	} else {
+		if len(req.SearchAfter) > 0 {
+			queryBody["search_after"] = req.SearchAfter
+		}
+	}
+
+	body, err := json.Marshal(queryBody)
+	if err != nil {
+		return nil, err
+	}
+
+	res, err := c.esClient.Search(
+		c.esClient.Search.WithContext(ctx),
+		c.esClient.Search.WithIndex(index),
+		c.esClient.Search.WithBody(bytes.NewReader(body)),
+	)
+
+	logs.CtxDebugf(ctx, "[Search] req : %s", string(body))
+
+	if err != nil {
+		return nil, err
+	}
+	defer res.Body.Close()
+
+	respBytes, err := io.ReadAll(res.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	var esResp Response
+	if err := json.Unmarshal(respBytes, &esResp); err != nil {
+		return nil, err
+	}
+	return &esResp, nil
+}
+
+func (c *es7Client) query2ESQuery(q *Query) map[string]any {
+	if q == nil {
+		return nil
+	}
+
+	var base map[string]any
+
+	switch q.Type {
+	case es.QueryTypeEqual:
+		base = map[string]any{
+			"term": map[string]any{
+				q.KV.Key: q.KV.Value,
+			},
+		}
+	case es.QueryTypeMatch:
+		base = map[string]any{
+			"match": map[string]any{
+				q.KV.Key: fmt.Sprint(q.KV.Value),
+			},
+		}
+	case es.QueryTypeMultiMatch:
+		base = map[string]any{
+			"multi_match": map[string]any{
+				"fields":   q.MultiMatchQuery.Fields,
+				"operator": q.MultiMatchQuery.Operator,
+				"query":    q.MultiMatchQuery.Query,
+				"type":     q.MultiMatchQuery.Type,
+			},
+		}
+	case es.QueryTypeNotExists:
+		base = map[string]any{
+			"bool": map[string]any{
+				"must_not": []map[string]any{
+					{"exists": map[string]any{"field": q.KV.Key}},
+				},
+			},
+		}
+	case es.QueryTypeContains:
+		base = map[string]any{
+			"wildcard": map[string]any{
+				q.KV.Key: map[string]any{
+					"value":            fmt.Sprintf("*%s*", q.KV.Value),
+					"case_insensitive": true,
+				},
+			},
+		}
+	case es.QueryTypeIn:
+		base = map[string]any{
+			"terms": map[string]any{
+				q.KV.Key: q.KV.Value,
+			},
+		}
+	default:
+		base = map[string]any{}
+	}
+
+	// 若没有 BoolQuery，直接返回 base query
+	if q.Bool == nil {
+		return base
+	}
+
+	// 如果有 BoolQuery，把 base 作为 BoolQuery 的一部分（或为空）
+	boolQuery := map[string]any{}
+
+	appendBool := func(key string, queries []Query) {
+		if len(queries) == 0 {
+			return
+		}
+		var arr []map[string]any
+		for i := range queries {
+			sub := c.query2ESQuery(&queries[i])
+			if sub != nil {
+				arr = append(arr, sub)
+			}
+		}
+		if len(arr) > 0 {
+			boolQuery[key] = arr
+		}
+	}
+
+	appendBool("filter", q.Bool.Filter)
+	appendBool("must", q.Bool.Must)
+	appendBool("must_not", q.Bool.MustNot)
+	appendBool("should", q.Bool.Should)
+
+	// 如果 base 不是空，作为一个 filter 附加进去
+	if len(base) > 0 {
+		if _, ok := boolQuery["filter"]; !ok {
+			boolQuery["filter"] = []map[string]any{}
+		}
+		boolQuery["filter"] = append(boolQuery["filter"].([]map[string]any), base)
+	}
+
+	if q.Bool.MinimumShouldMatch != nil {
+		boolQuery["minimum_should_match"] = *q.Bool.MinimumShouldMatch
+	}
+
+	return map[string]any{"bool": boolQuery}
+}
+
+func (c *es7Client) NewBulkIndexer(index string) (BulkIndexer, error) {
+	bi, err := esutil.NewBulkIndexer(esutil.BulkIndexerConfig{
+		Client: c.esClient,
+		Index:  index,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &es7BulkIndexer{bi: bi}, nil
+}
+
+type es7BulkIndexer struct {
+	bi esutil.BulkIndexer
+}
+
+func (b *es7BulkIndexer) Add(ctx context.Context, item BulkIndexerItem) error {
+	var buf bytes.Buffer
+	if item.Body != nil {
+		data, err := json.Marshal(item.Body)
+		if err != nil {
+			return err
+		}
+		buf.Write(data)
+	}
+
+	return b.bi.Add(ctx, esutil.BulkIndexerItem{
+		Action:          item.Action,
+		DocumentID:      item.DocumentID,
+		Body:            &buf,
+		Routing:         item.Routing,
+		Version:         item.Version,
+		VersionType:     item.VersionType,
+		RetryOnConflict: item.RetryOnConflict,
+	},
+	)
+}
+
+func (b *es7BulkIndexer) Close(ctx context.Context) error {
+	return b.bi.Close(ctx)
+}
+
+func (c *es7Client) Types() Types {
+	return &es7Types{}
+}
+
+type es7Types struct{}
+
+func (t *es7Types) NewLongNumberProperty() any {
+	return map[string]string{"type": "long"}
+}
+
+func (t *es7Types) NewTextProperty() any {
+	return map[string]string{"type": "text"}
+}
+
+func (t *es7Types) NewUnsignedLongNumberProperty() any {
+	return map[string]string{"type": "unsigned_long"}
+}

backend/infra/impl/es/es8.go

@@ -0,0 +1,301 @@
+/*
+ * Copyright 2025 coze-dev Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package es
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/elastic/go-elasticsearch/v8"
+	"github.com/elastic/go-elasticsearch/v8/esutil"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/core/search"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/indices/create"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/indices/delete"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/indices/exists"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/types"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/types/enums/operator"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/types/enums/sortorder"
+	"github.com/elastic/go-elasticsearch/v8/typedapi/types/enums/textquerytype"
+
+	"github.com/coze-dev/coze-studio/backend/infra/contract/es"
+	"github.com/coze-dev/coze-studio/backend/pkg/lang/conv"
+	"github.com/coze-dev/coze-studio/backend/pkg/lang/ptr"
+	"github.com/coze-dev/coze-studio/backend/pkg/logs"
+	"github.com/coze-dev/coze-studio/backend/pkg/sonic"
+)
+
+type es8Client struct {
+	esClient *elasticsearch.TypedClient
+	types    *es8Types
+}
+
+type es8BulkIndexer struct {
+	bi esutil.BulkIndexer
+}
+
+type es8Types struct{}
+
+func newES8() (Client, error) {
+	esAddr := os.Getenv("ES_ADDR")
+	esUsername := os.Getenv("ES_USERNAME")
+	esPassword := os.Getenv("ES_PASSWORD")
+	esClient, err := elasticsearch.NewTypedClient(elasticsearch.Config{
+		Addresses: []string{esAddr},
+		Username:  esUsername,
+		Password:  esPassword,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &es8Client{
+		esClient: esClient,
+		types:    &es8Types{},
+	}, nil
+}
+
+func (c *es8Client) Create(ctx context.Context, index, id string, document any) error {
+	_, err := c.esClient.Index(index).Id(id).Document(document).Do(ctx)
+	return err
+}
+
+func (c *es8Client) Update(ctx context.Context, index, id string, document any) error {
+	_, err := c.esClient.Update(index, id).Doc(document).Do(ctx)
+	return err
+}
+
+func (c *es8Client) Delete(ctx context.Context, index, id string) error {
+	_, err := c.esClient.Delete(index, id).Do(ctx)
+	return err
+}
+
+func (c *es8Client) Exists(ctx context.Context, index string) (bool, error) {
+	exist, err := exists.NewExistsFunc(c.esClient)(index).Do(ctx)
+	if err != nil {
+		return false, err
+	}
+
+	return exist, nil
+}
+
+func (c *es8Client) query2ESQuery(q *Query) *types.Query {
+	if q == nil {
+		return nil
+	}
+
+	var typesQ *types.Query
+	switch q.Type {
+	case es.QueryTypeEqual:
+		typesQ = &types.Query{
+			Term: map[string]types.TermQuery{
+				q.KV.Key: {Value: q.KV.Value},
+			},
+		}
+	case es.QueryTypeMatch:
+		typesQ = &types.Query{
+			Match: map[string]types.MatchQuery{
+				q.KV.Key: {Query: fmt.Sprint(q.KV.Value)},
+			},
+		}
+	case es.QueryTypeMultiMatch:
+		typesQ = &types.Query{
+			MultiMatch: &types.MultiMatchQuery{
+				Fields:   q.MultiMatchQuery.Fields,
+				Operator: &operator.Operator{Name: q.MultiMatchQuery.Operator},
+				Query:    q.MultiMatchQuery.Query,
+				Type:     &textquerytype.TextQueryType{Name: q.MultiMatchQuery.Type},
+			},
+		}
+	case es.QueryTypeNotExists:
+		typesQ = &types.Query{
+			Bool: &types.BoolQuery{
+				MustNot: []types.Query{{Exists: &types.ExistsQuery{Field: q.KV.Key}}},
+			},
+		}
+	case es.QueryTypeContains:
+		typesQ = &types.Query{
+			Wildcard: map[string]types.WildcardQuery{
+				q.KV.Key: {
+					Value:           ptr.Of(fmt.Sprintf("*%s*", q.KV.Value)),
+					CaseInsensitive: ptr.Of(true), // 忽略大小写
+				},
+			},
+		}
+	case es.QueryTypeIn:
+		typesQ = &types.Query{
+			Terms: &types.TermsQuery{
+				TermsQuery: map[string]types.TermsQueryField{
+					q.KV.Key: q.KV.Value,
+				},
+			},
+		}
+	default:
+		typesQ = &types.Query{}
+	}
+
+	if q.Bool == nil {
+		return typesQ
+	}
+
+	typesQ.Bool = &types.BoolQuery{}
+	for idx := range q.Bool.Filter {
+		v := q.Bool.Filter[idx]
+		typesQ.Bool.Filter = append(typesQ.Bool.Filter, *c.query2ESQuery(&v))
+	}
+
+	for idx := range q.Bool.Must {
+		v := q.Bool.Must[idx]
+		typesQ.Bool.Must = append(typesQ.Bool.Must, *c.query2ESQuery(&v))
+	}
+
+	for idx := range q.Bool.MustNot {
+		v := q.Bool.MustNot[idx]
+		typesQ.Bool.MustNot = append(typesQ.Bool.MustNot, *c.query2ESQuery(&v))
+	}
+
+	for idx := range q.Bool.Should {
+		v := q.Bool.Should[idx]
+		typesQ.Bool.Should = append(typesQ.Bool.Should, *c.query2ESQuery(&v))
+	}
+
+	if q.Bool.MinimumShouldMatch != nil {
+		typesQ.Bool.MinimumShouldMatch = q.Bool.MinimumShouldMatch
+	}
+
+	return typesQ
+}
+
+func (c *es8Client) Search(ctx context.Context, index string, req *Request) (*Response, error) {
+	esReq := &search.Request{
+		Query:    c.query2ESQuery(req.Query),
+		Size:     req.Size,
+		MinScore: (*types.Float64)(req.MinScore),
+	}
+
+	for _, sort := range req.Sort {
+		order := sortorder.Asc
+		if !sort.Asc {
+			order = sortorder.Desc
+		}
+		esReq.Sort = append(esReq.Sort, types.SortCombinations(types.SortOptions{
+			SortOptions: map[string]types.FieldSort{
+				sort.Field: {
+					Order: ptr.Of(order),
+				},
+			},
+		}))
+	}
+
+	if req.From != nil {
+		esReq.From = req.From
+	} else {
+		for _, v := range req.SearchAfter {
+			esReq.SearchAfter = append(esReq.SearchAfter, types.FieldValue(v))
+		}
+	}
+
+	logs.CtxDebugf(ctx, "Elasticsearch Request: %s\n", conv.DebugJsonToStr(esReq))
+
+	resp, err := c.esClient.Search().Request(esReq).Index(index).Do(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	respJson, err := sonic.MarshalString(resp)
+	if err != nil {
+		return nil, err
+	}
+
+	var esResp Response
+	if err := sonic.UnmarshalString(respJson, &esResp); err != nil {
+		return nil, err
+	}
+
+	return &esResp, nil
+}
+
+func (c *es8Client) CreateIndex(ctx context.Context, index string, properties map[string]any) error {
+	propertiesMap := make(map[string]types.Property)
+	for k, v := range properties {
+		propertiesMap[k] = v
+	}
+
+	if _, err := create.NewCreateFunc(c.esClient)(index).Request(&create.Request{
+		Mappings: &types.TypeMapping{
+			Properties: propertiesMap,
+		},
+	}).Do(ctx); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (c *es8Client) DeleteIndex(ctx context.Context, index string) error {
+	_, err := delete.NewDeleteFunc(c.esClient)(index).
+		IgnoreUnavailable(true).Do(ctx)
+	return err
+}
+
+func (c *es8Client) NewBulkIndexer(index string) (BulkIndexer, error) {
+	bi, err := esutil.NewBulkIndexer(esutil.BulkIndexerConfig{
+		Client: c.esClient,
+		Index:  index,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &es8BulkIndexer{bi}, nil
+}
+
+func (c *es8Client) Types() Types {
+	return c.types
+}
+
+func (t *es8Types) NewLongNumberProperty() any {
+	return types.NewLongNumberProperty()
+}
+
+func (t *es8Types) NewTextProperty() any {
+	return types.NewTextProperty()
+}
+
+func (t *es8Types) NewUnsignedLongNumberProperty() any {
+	return types.NewUnsignedLongNumberProperty()
+}
+
+func (b *es8BulkIndexer) Add(ctx context.Context, item BulkIndexerItem) error {
+	return b.bi.Add(ctx, esutil.BulkIndexerItem{
+		Index:           item.Index,
+		Action:          item.Action,
+		DocumentID:      item.DocumentID,
+		Routing:         item.Routing,
+		Version:         item.Version,
+		VersionType:     item.VersionType,
+		Body:            item.Body,
+		RetryOnConflict: item.RetryOnConflict,
+		// not support in es7
+		// RequireAlias:    item.RequireAlias,
+		// IfSeqNo:         item.IfSeqNo,
+		// IfPrimaryTerm:   item.IfPrimaryTerm,
+	})
+}
+
+func (b *es8BulkIndexer) Close(ctx context.Context) error {
+	return b.bi.Close(ctx)
+}

backend/infra/impl/es/es_impl.go

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2025 coze-dev Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package es
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/coze-dev/coze-studio/backend/infra/contract/es"
+)
+
+type (
+	Client          = es.Client
+	Types           = es.Types
+	BulkIndexer     = es.BulkIndexer
+	BulkIndexerItem = es.BulkIndexerItem
+	BoolQuery       = es.BoolQuery
+	Query           = es.Query
+	Response        = es.Response
+	Request         = es.Request
+)
+
+func New() (Client, error) {
+	v := os.Getenv("ES_VERSION")
+	if v == "v8" {
+		return newES8()
+	} else if v == "v7" {
+		return newES7()
+	}
+
+	return nil, fmt.Errorf("unsupported es version %s", v)
+}