webp/models/zincsearch.go

package models

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
)

var (
	user     string
	password string
	index    string
	zincHost string
)

func ZincInit() {
	user = config.GetString("zincsearch.user")
	password = config.GetString("zincsearch.password")
	index = config.GetString("zincsearch.index")
	zincHost = config.GetString("zincsearch.host")
}

// 返回结果
type Response struct {
	Took     int  `json:"took"`
	TimedOut bool `json:"timed_out"`
	Shards   struct {
		Total      int `json:"total"`
		Successful int `json:"successful"`
		Skipped    int `json:"skipped"`
		Failed     int `json:"failed"`
	} `json:"_shards"`
	Hits struct {
		Total struct {
			Value int `json:"value"`
		} `json:"total"`
		MaxScore float64 `json:"max_score"`
		Hits     []struct {
			Index     string  `json:"_index"`
			Type      string  `json:"_type"`
			ID        string  `json:"_id"`
			Score     float64 `json:"_score"`
			Timestamp string  `json:"@timestamp"`
			Source    struct {
				Timestamp  string `json:"@timestamp"`
				Athlete    string `json:"Athlete"`
				City       string `json:"City"`
				Country    string `json:"Country"`
				Discipline string `json:"Discipline"`
				Event      string `json:"Event"`
				Gender     string `json:"Gender"`
				Medal      string `json:"Medal"`
				Season     string `json:"Season"`
				Sport      string `json:"Sport"`
				Year       int    `json:"Year"`
			} `json:"_source"`
		} `json:"hits"`
	} `json:"hits"`
}

func (res Response) ToIDList(first, last int, after, before string) (id_list []string) {
	for _, hit := range res.Hits.Hits {
		id_list = append(id_list, hit.ID)
	}

	// 如果 after 不为 0, 从这个ID开始向后取切片
	if after != "" {
		for i, id := range id_list {
			if id == after {
				id_list = id_list[i+1:]
				break
			}
		}
	}

	// 如果 before 不为 0, 从这个ID开始向前取切片
	if before != "" {
		for i, id := range id_list {
			if id == before {
				id_list = id_list[:i]
				break
			}
		}
	}

	// 如果 first 不为 0, 取切片的前 first 个元素
	if first != 0 {
		if first > len(id_list) {
			first = len(id_list)
		}
		id_list = id_list[:first]
	}

	// 如果 last 不为 0, 取切片的后 last 个元素
	if last != 0 {
		if last > len(id_list) {
			last = len(id_list)
		}
		id_list = id_list[len(id_list)-last:]
	}

	return
}

func ZincSearch(query map[string]interface{}) (rest Response, err error) {
	bas64encodedCreds := base64.StdEncoding.EncodeToString([]byte(user + ":" + password))
	zincURL := zincHost + "/es/" + index + "/_search"
	queryBytes, _ := json.Marshal(query)
	req, err := http.NewRequest("POST", zincURL, bytes.NewBuffer(queryBytes))
	if err != nil {
		log.Println("创建请求失败", err)
		return
	}
	req.Header.Set("Content-type", "application/json")
	req.Header.Set("Authorization", "Basic "+bas64encodedCreds)

	// 创建一个新的 HTTP 客户端
	customClient := &http.Client{
		Transport: &http.Transport{
			Proxy: func(*http.Request) (*url.URL, error) {
				return nil, nil
			},
		},
	}

	resp, err := customClient.Do(req)
	if err != nil {
		log.Println("发起请求失败", err)
		return
	}
	defer resp.Body.Close()

	// 检查状态码
	if resp.StatusCode != http.StatusOK {
		log.Println("请求未成功，状态码：", resp.StatusCode)
		return rest, fmt.Errorf("请求未成功，状态码：%d", resp.StatusCode)
	}

	err = json.NewDecoder(resp.Body).Decode(&rest)
	if err != nil {
		log.Println("解析响应失败", err)
		return
	}
	return
}

func ZincPut(data map[string]interface{}) (err error) {
	bas64encodedCreds := base64.StdEncoding.EncodeToString([]byte(user + ":" + password))
	zincURL := zincHost + "/api/" + index + "/_doc"
	queryBytes, _ := json.Marshal(data)
	req, err := http.NewRequest("POST", zincURL, bytes.NewBuffer(queryBytes))
	if err != nil {
		return
	}
	req.Header.Set("Content-type", "application/json")
	req.Header.Set("Authorization", "Basic "+bas64encodedCreds)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	// 直接作为文本打印 body
	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return
	}
	fmt.Println(string(bodyBytes))
	return
}