sorbet/pkg/zinc/zinc.go

package zinc

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/go-resty/resty/v2"
	"net/http"
	"time"
)

type Client struct {
	Host     string
	User     string
	Password string
}

type Index struct {
	Name        string         `json:"name"`
	StorageType string         `json:"storage_type"`
	Mappings    *IndexMappings `json:"mappings"`
}

type IndexMappings struct {
	Properties *IndexProperty `json:"properties"`
}

type IndexProperty map[string]*IndexPropertyT

type IndexPropertyT struct {
	Type           string `json:"type"`
	Index          bool   `json:"index"`
	Store          bool   `json:"store"`
	Sortable       bool   `json:"sortable"`
	Aggregatable   bool   `json:"aggregatable"`
	Highlightable  bool   `json:"highlightable"`
	Analyzer       string `json:"analyzer"`
	SearchAnalyzer string `json:"search_analyzer"`
	Format         string `json:"format"`
}

type QueryResultT struct {
	Took     int          `json:"took"`
	TimedOut bool         `json:"timed_out"`
	Hits     *HitsResultT `json:"hits"`
}

type HitsResultT struct {
	Total    *HitsResultTotalT `json:"total"`
	MaxScore float64           `json:"max_score"`
	Hits     []*HitItem        `json:"hits"`
}

type HitsResultTotalT struct {
	Value int64 `json:"value"`
}

type HitItem struct {
	Index     string    `json:"_index"`
	Type      string    `json:"_type"`
	ID        string    `json:"_id"`
	Score     float64   `json:"_score"`
	Timestamp time.Time `json:"@timestamp"`
	Source    any       `json:"_source"`
}

// NewClient 获取ZincClient新实例
func NewClient(host, user, passwd string) *Client {
	return &Client{
		Host:     host,
		User:     user,
		Password: passwd,
	}
}

// CreateIndex 创建索引
func (c *Client) CreateIndex(name string, p *IndexProperty) bool {
	data := &Index{
		Name:        name,
		StorageType: "disk",
		Mappings: &IndexMappings{
			Properties: p,
		},
	}
	resp, err := c.request().SetBody(data).Put("/api/index")

	if err != nil || resp.StatusCode() != http.StatusOK {
		return false
	}

	return true
}

// ExistIndex 检查索引是否存在
func (c *Client) ExistIndex(name string) bool {
	resp, err := c.request().Get("/api/index")

	if err != nil || resp.StatusCode() != http.StatusOK {
		return false
	}

	retData := &map[string]any{}
	err = json.Unmarshal([]byte(resp.String()), retData)
	if err != nil {
		return false
	}

	if _, ok := (*retData)[name]; ok {
		return true
	}

	return false
}

// PutDoc 新增/更新文档
func (c *Client) PutDoc(name string, id int64, doc any) (bool, error) {
	resp, err := c.request().SetBody(doc).Put(fmt.Sprintf("/api/%s/_doc/%d", name, id))

	if err != nil {
		return false, err
	}

	if resp.StatusCode() != http.StatusOK {
		return false, errors.New(resp.Status())
	}

	return true, nil
}

// BulkPushDoc 批量新增文档
func (c *Client) BulkPushDoc(docs []map[string]any) (bool, error) {
	dataStr := ""
	for _, doc := range docs {
		str, err := json.Marshal(doc)
		if err == nil {
			dataStr = dataStr + string(str) + "\n"
		}
	}

	resp, err := c.request().SetBody(dataStr).Post("/api/_bulk")
	if err != nil {
		return false, err
	}

	if resp.StatusCode() != http.StatusOK {
		return false, errors.New(resp.Status())
	}

	return true, nil
}

func (c *Client) EsQuery(indexName string, q any) (*QueryResultT, error) {
	resp, err := c.request().SetBody(q).Post(fmt.Sprintf("/es/%s/_search", indexName))
	if err != nil {
		return nil, err
	}

	if resp.StatusCode() != http.StatusOK {
		return nil, errors.New(resp.Status())
	}

	result := &QueryResultT{}
	err = json.Unmarshal(resp.Body(), result)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (c *Client) ApiQuery(indexName string, q any) (*QueryResultT, error) {
	resp, err := c.request().SetBody(q).Post(fmt.Sprintf("/api/%s/_search", indexName))
	if err != nil {
		return nil, err
	}

	if resp.StatusCode() != http.StatusOK {
		return nil, errors.New(resp.Status())
	}

	result := &QueryResultT{}
	err = json.Unmarshal(resp.Body(), result)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (c *Client) DelDoc(indexName, id string) error {
	resp, err := c.request().Delete(fmt.Sprintf("/api/%s/_doc/%s", indexName, id))
	if err != nil {
		return err
	}

	if resp.StatusCode() != http.StatusOK {
		return errors.New(resp.Status())
	}

	return nil
}

func (c *Client) request() *resty.Request {
	client := resty.New()
	client.DisableWarn = true
	client.SetBaseURL(c.Host)
	client.SetBasicAuth(c.User, c.Password)
	return client.R()
}