Create & Init Project...

library/database/elastic/query.go

@@ -0,0 +1,578 @@
+package elastic
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/url"
+	"strings"
+	"time"
+
+	"go-common/library/ecode"
+	httpx "go-common/library/net/http/blademaster"
+	timex "go-common/library/time"
+
+	"github.com/pkg/errors"
+)
+
+const (
+	// OrderAsc order ascend
+	OrderAsc = "asc"
+	// OrderDesc order descend
+	OrderDesc = "desc"
+
+	// RangeScopeLoRo left open & right open
+	RangeScopeLoRo rangeScope = "( )"
+	// RangeScopeLoRc left open & right close
+	RangeScopeLoRc rangeScope = "( ]"
+	// RangeScopeLcRo left close & right open
+	RangeScopeLcRo rangeScope = "[ )"
+	// RangeScopeLcRc lect close & right close
+	RangeScopeLcRc rangeScope = "[ ]"
+
+	// NotTypeEq not type eq
+	NotTypeEq notType = "eq"
+	// NotTypeIn not type in
+	NotTypeIn notType = "in"
+	// NotTypeRange not type range
+	NotTypeRange notType = "range"
+
+	// LikeLevelHigh wildcard keyword
+	LikeLevelHigh likeLevel = "high"
+	// LikeLevelMiddle ngram(1,2)
+	LikeLevelMiddle likeLevel = "middle"
+	// LikeLevelLow match split word
+	LikeLevelLow likeLevel = "low"
+
+	// IndexTypeYear index by year
+	IndexTypeYear indexType = "year"
+	// IndexTypeMonth index by month
+	IndexTypeMonth indexType = "month"
+	// IndexTypeWeek index by week
+	IndexTypeWeek indexType = "week"
+	// IndexTypeDay index by day
+	IndexTypeDay indexType = "day"
+
+	// EnhancedModeGroupBy group by mode
+	EnhancedModeGroupBy enhancedType = "group_by"
+	// EnhancedModeDistinct distinct mode
+	EnhancedModeDistinct enhancedType = "distinct"
+	// EnhancedModeSum sum mode
+	EnhancedModeSum enhancedType = "sum"
+)
+
+type (
+	notType      string
+	rangeScope   string
+	likeLevel    string
+	indexType    string
+	enhancedType string
+)
+
+var (
+	_defaultHost = "http://manager.bilibili.co"
+	_pathQuery   = "/x/admin/search/query"
+	_pathUpsert  = "/x/admin/search/upsert"
+
+	_defaultHTTPClient = &httpx.ClientConfig{
+		App: &httpx.App{
+			Key:    "3c4e41f926e51656",
+			Secret: "26a2095b60c24154521d24ae62b885bb",
+		},
+		Dial:    timex.Duration(time.Second),
+		Timeout: timex.Duration(time.Second),
+	}
+
+	// for index by week
+	weeks = map[int]string{0: "0107", 1: "0815", 2: "1623", 3: "2431"}
+)
+
+// Config Elastic config
+type Config struct {
+	Host       string
+	HTTPClient *httpx.ClientConfig
+}
+
+// response query elastic response
+type response struct {
+	Code    int             `json:"code"`
+	Message string          `json:"message"`
+	Data    json.RawMessage `json:"data"`
+}
+
+type query struct {
+	Fields          []string            `json:"fields"`
+	From            string              `json:"from"`
+	OrderScoreFirst bool                `json:"order_score_first"`
+	OrderRandomSeed string              `json:"order_random_seed"`
+	Highlight       bool                `json:"highlight"`
+	Pn              int                 `json:"pn"`
+	Ps              int                 `json:"ps"`
+	Order           []map[string]string `json:"order,omitempty"`
+	Where           where               `json:"where,omitempty"`
+}
+
+func (q *query) string() (string, error) {
+	var (
+		sli []string
+		m   = make(map[string]bool)
+	)
+	for _, i := range strings.Split(q.From, ",") {
+		if m[i] {
+			continue
+		}
+		m[i] = true
+		sli = append(sli, i)
+	}
+	q.From = strings.Join(sli, ",")
+	bs, err := json.Marshal(q)
+	if err != nil {
+		return "", err
+	}
+	return string(bs), nil
+}
+
+type where struct {
+	GroupBy  string                      `json:"group_by,omitempty"`
+	Like     []whereLike                 `json:"like,omitempty"`
+	Eq       map[string]interface{}      `json:"eq,omitempty"`
+	Or       map[string]interface{}      `json:"or,omitempty"`
+	In       map[string][]interface{}    `json:"in,omitempty"`
+	Range    map[string]string           `json:"range,omitempty"`
+	Combo    []*Combo                    `json:"combo,omitempty"`
+	Not      map[notType]map[string]bool `json:"not,omitempty"`
+	Enhanced []interface{}               `json:"enhanced,omitempty"`
+}
+
+type whereLike struct {
+	Fields []string  `json:"kw_fields"`
+	Words  []string  `json:"kw"`
+	Or     bool      `json:"or"`
+	Level  likeLevel `json:"level"`
+}
+
+// Combo mix eq & in & range
+type Combo struct {
+	EQ       []map[string]interface{}   `json:"eq,omitempty"`
+	In       []map[string][]interface{} `json:"in,omitempty"`
+	Range    []map[string]string        `json:"range,omitempty"`
+	NotEQ    []map[string]interface{}   `json:"not_eq,omitempty"`
+	NotIn    []map[string][]interface{} `json:"not_in,omitempty"`
+	NotRange []map[string]string        `json:"not_range,omitempty"`
+	Min      struct {
+		EQ       int `json:"eq,omitempty"`
+		In       int `json:"in,omitempty"`
+		Range    int `json:"range,omitempty"`
+		NotEQ    int `json:"not_eq,omitempty"`
+		NotIn    int `json:"not_in,omitempty"`
+		NotRange int `json:"not_range,omitempty"`
+		Min      int `json:"min"`
+	} `json:"min"`
+}
+
+// ComboEQ .
+func (cmb *Combo) ComboEQ(eq []map[string]interface{}) *Combo {
+	cmb.EQ = append(cmb.EQ, eq...)
+	return cmb
+}
+
+// ComboRange .
+func (cmb *Combo) ComboRange(r []map[string]string) *Combo {
+	cmb.Range = append(cmb.Range, r...)
+	return cmb
+}
+
+// ComboIn .
+func (cmb *Combo) ComboIn(in []map[string][]interface{}) *Combo {
+	cmb.In = append(cmb.In, in...)
+	return cmb
+}
+
+// ComboNotEQ .
+func (cmb *Combo) ComboNotEQ(eq []map[string]interface{}) *Combo {
+	cmb.NotEQ = append(cmb.NotEQ, eq...)
+	return cmb
+}
+
+// ComboNotRange .
+func (cmb *Combo) ComboNotRange(r []map[string]string) *Combo {
+	cmb.NotRange = append(cmb.NotRange, r...)
+	return cmb
+}
+
+// ComboNotIn .
+func (cmb *Combo) ComboNotIn(in []map[string][]interface{}) *Combo {
+	cmb.NotIn = append(cmb.NotIn, in...)
+	return cmb
+}
+
+// MinEQ .
+func (cmb *Combo) MinEQ(min int) *Combo {
+	cmb.Min.EQ = min
+	return cmb
+}
+
+// MinIn .
+func (cmb *Combo) MinIn(min int) *Combo {
+	cmb.Min.In = min
+	return cmb
+}
+
+// MinRange .
+func (cmb *Combo) MinRange(min int) *Combo {
+	cmb.Min.Range = min
+	return cmb
+}
+
+// MinNotEQ .
+func (cmb *Combo) MinNotEQ(min int) *Combo {
+	cmb.Min.NotEQ = min
+	return cmb
+}
+
+// MinNotIn .
+func (cmb *Combo) MinNotIn(min int) *Combo {
+	cmb.Min.NotIn = min
+	return cmb
+}
+
+// MinNotRange .
+func (cmb *Combo) MinNotRange(min int) *Combo {
+	cmb.Min.NotRange = min
+	return cmb
+}
+
+// MinAll .
+func (cmb *Combo) MinAll(min int) *Combo {
+	cmb.Min.Min = min
+	return cmb
+}
+
+type groupBy struct {
+	Mode  enhancedType        `json:"mode"`
+	Field string              `json:"field"`
+	Order []map[string]string `json:"order"`
+}
+
+type enhance struct {
+	Mode  enhancedType        `json:"mode"`
+	Field string              `json:"field"`
+	Order []map[string]string `json:"order,omitempty"`
+	Size  int                 `json:"size,omitempty"`
+}
+
+// Elastic clastic instance
+type Elastic struct {
+	c      *Config
+	client *httpx.Client
+}
+
+// NewElastic .
+func NewElastic(c *Config) *Elastic {
+	if c == nil {
+		c = &Config{
+			Host:       _defaultHost,
+			HTTPClient: _defaultHTTPClient,
+		}
+	}
+	return &Elastic{
+		c:      c,
+		client: httpx.NewClient(c.HTTPClient),
+	}
+}
+
+// Request request to elastic
+type Request struct {
+	*Elastic
+	q        query
+	business string
+}
+
+// NewRequest new a request every search query
+func (e *Elastic) NewRequest(business string) *Request {
+	return &Request{
+		Elastic:  e,
+		business: business,
+		q: query{
+			Fields:          []string{},
+			Highlight:       false,
+			OrderScoreFirst: true,
+			OrderRandomSeed: "",
+			Pn:              1,
+			Ps:              10,
+		},
+	}
+}
+
+// Fields add query fields
+func (r *Request) Fields(fields ...string) *Request {
+	r.q.Fields = append(r.q.Fields, fields...)
+	return r
+}
+
+// Index add query index
+func (r *Request) Index(indexes ...string) *Request {
+	r.q.From = strings.Join(indexes, ",")
+	return r
+}
+
+// IndexByMod index by mod
+func (r *Request) IndexByMod(prefix string, val, mod int64) *Request {
+	tmp := mod - 1
+	var digit int
+	for tmp > 0 {
+		tmp /= 10
+		digit++
+	}
+	format := fmt.Sprintf("%s_%%0%dd", prefix, digit)
+	r.q.From = fmt.Sprintf(format, val%mod)
+	return r
+}
+
+// IndexByTime index by time
+func (r *Request) IndexByTime(prefix string, typ indexType, begin, end time.Time) *Request {
+	var (
+		buf     bytes.Buffer
+		index   string
+		indexes = make(map[string]struct{})
+	)
+	for {
+		year := begin.Format("2006")
+		month := begin.Format("01")
+		switch typ {
+		case IndexTypeYear:
+			index = strings.Join([]string{prefix, year}, "_")
+		case IndexTypeMonth:
+			index = strings.Join([]string{prefix, year, month}, "_")
+		case IndexTypeDay:
+			day := begin.Format("02")
+			index = strings.Join([]string{prefix, year, month, day}, "_")
+		case IndexTypeWeek:
+			index = strings.Join([]string{prefix, year, month, weeks[begin.Day()/8]}, "_")
+		}
+		if begin.After(end) && begin.Day() != end.Day() {
+			break
+		}
+		indexes[index] = struct{}{}
+		begin = begin.AddDate(0, 0, 1)
+	}
+	for i := range indexes {
+		buf.WriteString(i)
+		buf.WriteString(",")
+	}
+	r.q.From = strings.TrimSuffix(buf.String(), ",")
+	return r
+}
+
+// OrderScoreFirst switch for order score first
+func (r *Request) OrderScoreFirst(v bool) *Request {
+	r.q.OrderScoreFirst = v
+	return r
+}
+
+// OrderRandomSeed switch for order random
+func (r *Request) OrderRandomSeed(v string) *Request {
+	r.q.OrderRandomSeed = v
+	return r
+}
+
+// Highlight switch from highlight
+func (r *Request) Highlight(v bool) *Request {
+	r.q.Highlight = v
+	return r
+}
+
+// Pn page number
+func (r *Request) Pn(v int) *Request {
+	r.q.Pn = v
+	return r
+}
+
+// Ps page size
+func (r *Request) Ps(v int) *Request {
+	r.q.Ps = v
+	return r
+}
+
+// Order filed sort
+func (r *Request) Order(field, sort string) *Request {
+	if sort != OrderAsc {
+		sort = OrderDesc
+	}
+	r.q.Order = append(r.q.Order, map[string]string{field: sort})
+	return r
+}
+
+// WhereEq where qual
+func (r *Request) WhereEq(field string, eq interface{}) *Request {
+	if r.q.Where.Eq == nil {
+		r.q.Where.Eq = make(map[string]interface{})
+	}
+	r.q.Where.Eq[field] = eq
+	return r
+}
+
+// WhereOr where or
+func (r *Request) WhereOr(field string, or interface{}) *Request {
+	if r.q.Where.Or == nil {
+		r.q.Where.Or = make(map[string]interface{})
+	}
+	r.q.Where.Or[field] = or
+	return r
+}
+
+// WhereIn where in
+func (r *Request) WhereIn(field string, in interface{}) *Request {
+	if r.q.Where.In == nil {
+		r.q.Where.In = make(map[string][]interface{})
+	}
+	switch v := in.(type) {
+	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, string:
+		r.q.Where.In[field] = append(r.q.Where.In[field], v)
+	case []int:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []int64:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []string:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []int8:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []int16:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []int32:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []uint:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []uint8:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []uint16:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []uint32:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	case []uint64:
+		for _, i := range v {
+			r.q.Where.In[field] = append(r.q.Where.In[field], i)
+		}
+	}
+	return r
+}
+
+// WhereRange where range
+func (r *Request) WhereRange(field string, start, end interface{}, scope rangeScope) *Request {
+	if r.q.Where.Range == nil {
+		r.q.Where.Range = make(map[string]string)
+	}
+	if start == nil {
+		start = ""
+	}
+	if end == nil {
+		end = ""
+	}
+	switch scope {
+	case RangeScopeLoRo:
+		r.q.Where.Range[field] = fmt.Sprintf("(%v,%v)", start, end)
+	case RangeScopeLoRc:
+		r.q.Where.Range[field] = fmt.Sprintf("(%v,%v]", start, end)
+	case RangeScopeLcRo:
+		r.q.Where.Range[field] = fmt.Sprintf("[%v,%v)", start, end)
+	case RangeScopeLcRc:
+		r.q.Where.Range[field] = fmt.Sprintf("[%v,%v]", start, end)
+	}
+	return r
+}
+
+// WhereNot where not
+func (r *Request) WhereNot(typ notType, fields ...string) *Request {
+	if r.q.Where.Not == nil {
+		r.q.Where.Not = make(map[notType]map[string]bool)
+	}
+	if r.q.Where.Not[typ] == nil {
+		r.q.Where.Not[typ] = make(map[string]bool)
+	}
+	for _, v := range fields {
+		r.q.Where.Not[typ][v] = true
+	}
+	return r
+}
+
+// WhereLike where like
+func (r *Request) WhereLike(fields, words []string, or bool, level likeLevel) *Request {
+	if len(fields) == 0 || len(words) == 0 {
+		return r
+	}
+	l := whereLike{Fields: fields, Words: words, Or: or, Level: level}
+	r.q.Where.Like = append(r.q.Where.Like, l)
+	return r
+}
+
+// WhereCombo where combo
+func (r *Request) WhereCombo(cmb ...*Combo) *Request {
+	r.q.Where.Combo = append(r.q.Where.Combo, cmb...)
+	return r
+}
+
+// GroupBy where group by
+func (r *Request) GroupBy(mode enhancedType, field string, order []map[string]string) *Request {
+	for _, i := range order {
+		for k, v := range i {
+			if v != OrderAsc {
+				i[k] = OrderDesc
+			}
+		}
+	}
+	r.q.Where.Enhanced = append(r.q.Where.Enhanced, groupBy{Mode: mode, Field: field, Order: order})
+	return r
+}
+
+// Sum where enhance sum
+func (r *Request) Sum(field string) *Request {
+	r.q.Where.Enhanced = append(r.q.Where.Enhanced, enhance{Mode: EnhancedModeSum, Field: field})
+	return r
+}
+
+// Scan parse the query response data
+func (r *Request) Scan(ctx context.Context, result interface{}) (err error) {
+	q, err := r.q.string()
+	if err != nil {
+		return
+	}
+	params := url.Values{}
+	params.Add("business", r.business)
+	params.Add("query", q)
+	response := new(response)
+	if err = r.client.Get(ctx, r.c.Host+_pathQuery, "", params, &response); err != nil {
+		return
+	}
+	if !ecode.Int(response.Code).Equal(ecode.OK) {
+		return ecode.Int(response.Code)
+	}
+	err = errors.Wrapf(json.Unmarshal(response.Data, &result), "scan(%s)", response.Data)
+	return
+}
+
+// Params get query parameters
+func (r *Request) Params() string {
+	q, _ := r.q.string()
+	return fmt.Sprintf("business=%s&query=%s", r.business, q)
+}