openbilibili/library/database/elastic/query.go

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)
}