golang操作elastic的一些小工具

个人封装的，框架是goframe。

package es

import (
    "encoding/json"
    "errors"
    "fmt"

    "github.com/gogf/gf/encoding/gbase64"
    "github.com/gogf/gf/frame/g"
    "github.com/gogf/gf/os/gcfg"
    "github.com/gogf/gf/util/gconv"
)

type Error struct {
    RootCause interface{} `json:"root_cause"`
    Reason    string      `json:"reason"`
}

type Shards struct {
    Total      int `json:"total"`
    Successful int `json:"successful"`
    Skipped    int `json:"skipped"`
    Failed     int `json:"failed"`
}

// Elastic响应结果数量
type HitsTotal struct {
    Value int `json:"value"`
}

// Elastic相应结果文档集合
type Hits struct {
    Total HitsTotal     `json:"total"`
    Hits  []interface{} `json:"hits"`
}

type Mappings struct {
    Properties map[string]interface{} `json:"properties"`
}

// MappingsResponse
type MappingsResponse struct {
    Mappings Mappings `json:"mappings"`
}

// Elastic的响应结果结构体
type ElasticResponse struct {
    Took         interface{} `json:"took"`
    TimedOut     interface{} `json:"timed_out"`
    Shards       Shards      `json:"_shards"`
    Status       int         `json:"status"`
    Error        Error       `json:"error"`
    Hits         Hits        `json:"hits"`
    Aggregations interface{} `json:"aggregations"`
}

// AnalyseElasticResponse 解析Elastic的响应结果
func AnalyseElasticResponse(response string) (ElasticResponse, error) {
    var r ElasticResponse
    if len(response) == 0 {
        return r, nil
    }
    err := json.Unmarshal([]byte(response), &r)
    if err != nil {
        return r, err
    }
    return r, nil
}

// AnalyseElasticResponse 解析Elastic的响应结果
func AnalyseMappingsResponse(data map[string]interface{}) (*MappingsResponse, error) {
    var res *MappingsResponse
    if err := gconv.Struct(data, &res); err != nil {
        panic(err)
    }
    return res, nil
}

// ParseQuery 格式化前端查询参数
// params["must"] = map[string]interface{}{"range":map[string]interface{}{"name":map[string]interface{}{"lt":"2021-01-01",}}}
// params["should"] = map[string]interface{}{"term":map[string]interface{}{"age":"1",}}}
// sort []map[string]interface{}{{"age":map[string]interface{"order": "asc",}}}
func ParseQuery(params interface{}, page int, limit int, sort interface{}, aggs interface{}, source ...interface{}) string {
    requests := map[string]interface{}{
        "query":            map[string]interface{}{},
        "from":             (page - 1) * limit,
        "size":             limit,
        "track_total_hits": true,
        "sort":             []interface{}{},
        "aggs":             map[string]interface{}{},
    }
    // query查询结构 - 使用Bool查询:must|must_not|should
    query := map[string]interface{}{}
    queryBool := map[string]interface{}{}
    queryBoolList := []string{"must", "must_not", "should"}
    for _, tag := range queryBoolList {
        formatTagList := []interface{}{}
        if tagItems, ok := gconv.Map(params)[tag]; ok {
            for _, tagItem := range gconv.SliceMap(tagItems) {
                formatTagList = append(formatTagList, gconv.MapDeep(tagItem))
            }
        }

        queryBool[tag] = formatTagList
    }
    if tagItems, ok := gconv.Map(params)["minimum_should_match"]; ok {
        queryBool["minimum_should_match"] = tagItems
    }
    query["bool"] = queryBool
    // 拼接query参数
    requests["query"] = query
    // 排序
    sortList := []map[string]interface{}{}
    if s := gconv.SliceMap(sort); len(s) != 0 {
        for _, sortItem := range s {
            sortList = append(sortList, gconv.MapDeep(sortItem))
        }
    }
    requests["sort"] = sortList
    // 聚合查找
    requests["aggs"] = gconv.MapDeep(aggs)
    // 字段选择
    if len(source) > 0 {
        requests["_source"] = source[0]
    }
    requestJSON, err := json.Marshal(requests)
    if err != nil {
        panic(err)
    }
    fmt.Println("query --> ", string(requestJSON))
    return string(requestJSON)
}

// QueryElastic 查询ES并获取解析后的结果
// ElasticIndex 索引ID
// QerryJSON 请求参数
func QueryElastic(ElasticIndex string, QerryJSON string) (*ElasticResponse, error) {
    result := &ElasticResponse{}
    Host := gcfg.Instance().GetString("elastic.Host")
    Port := gcfg.Instance().GetString("elastic.Port")
    UserName := gcfg.Instance().GetString("elastic.UserName")
    Password := gcfg.Instance().GetString("elastic.Password")
    URL := Host + ":" + gconv.String(Port) + "/" + ElasticIndex + "/_doc/_search"
    authEncode := gbase64.EncodeString(UserName + ":" + Password)
    h := map[string]string{
        "Authorization": "Basic " + authEncode,
    }
    if r, err := g.Client().Header(h).Post(URL, QerryJSON); err != nil {
        panic(err)
    } else {
        defer r.Close()
        responseString := r.ReadAllString()
        r, err := AnalyseElasticResponse(responseString)
        if err != nil {
            return result, err
        }
        return &r, nil
    }
}

// GetMapping 获取index对应的mapping
func GetMapping(ElasticIndex string) (g.Map, error) {
    str := `/_mapping/`
    Host := gcfg.Instance().GetString("elastic.Host")
    Port := gcfg.Instance().GetString("elastic.Port")
    UserName := gcfg.Instance().GetString("elastic.UserName")
    Password := gcfg.Instance().GetString("elastic.Password")
    URL := Host + ":" + gconv.String(Port) + "/" + ElasticIndex + str
    authEncode := gbase64.EncodeString(UserName + ":" + Password)
    h := map[string]string{
        "Authorization": "Basic " + authEncode,
    }
    if r, err := g.Client().Header(h).Get(URL); err != nil {
        panic(err)
    } else {
        defer r.Close()
        responseString := r.ReadAllString()
        return gconv.Map(responseString), nil
    }
}

// GetAggregations 查询ES并获取解析后的结果
// response *ElasticResponse ES查询后的结果
// 返回结果集,error对象,如果error对象为空表示读取正常
func GetAggregations(response *ElasticResponse) ([]map[string]interface{}, error) {
    if response.Error.Reason != "" {
        return nil, errors.New(response.Error.Reason)
    }
    Aggregations := gconv.Map(response.Aggregations)
    if len(Aggregations) == 0 {
        return nil, nil
    }
    // 递归格式化
    var formatAgg []map[string]interface{}
    FormatAgg(&formatAgg, Aggregations)
    return formatAgg, nil
}

// FormatAgg  无限递归拉取直至最后一级别只有key和doc_count
// direct 解析后的结果集
// Aggregations 用于解析的结果集
// key|doc_count
// wystanxu@tencent.com
// 2021-08-11
func FormatAgg(direct *[]map[string]interface{}, Aggregations interface{}) bool {
    tmp := gconv.Map(Aggregations)
    for idx, v := range tmp {
        if buckets, exist := gconv.Map(v)["buckets"]; exist {
            b := gconv.SliceMap(buckets)
            for _, vv := range b {
                children := []map[string]interface{}{}
                for kkk, vvv := range vv {
                    kkk_str := gconv.String(kkk)
                    if kkk_str != "key" && kkk_str != "doc_count" && kkk_str != "key_as_string" {
                        FormatAgg(&children, map[string]interface{}{
                            kkk_str: vvv,
                        })
                    }
                }
                item := map[string]interface{}{
                    "key_name":  idx,
                    "key_value": gconv.String(vv["key"]),
                    "count":     gconv.String(vv["doc_count"]),
                }
                if tmp, exist := vv["key_as_string"]; exist {
                    item["key_value"] = tmp
                }
                if len(children) != 0 {
                    item["children"] = children
                }
                *direct = append(*direct, item)
            }
        }
    }
    return true
}

 

posted @ 2022-02-19 00:28  许伟强  阅读(197)  评论(0)    收藏  举报