migrate from github
This commit is contained in:
@@ -0,0 +1,295 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/XingfenD/yoresee_doc/pkg/errs"
|
||||
)
|
||||
|
||||
type ElasticsearchClient struct {
|
||||
addresses []string
|
||||
username string
|
||||
password string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
var ES *ElasticsearchClient
|
||||
|
||||
func NewElasticsearchClient(cfg *ElasticsearchOptions) (*ElasticsearchClient, error) {
|
||||
if cfg == nil || !cfg.Enabled {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
addresses := make([]string, 0, len(cfg.Addresses))
|
||||
for _, addr := range cfg.Addresses {
|
||||
trimmed := strings.TrimSpace(addr)
|
||||
if trimmed == "" {
|
||||
continue
|
||||
}
|
||||
addresses = append(addresses, strings.TrimRight(trimmed, "/"))
|
||||
}
|
||||
if len(addresses) == 0 {
|
||||
return nil, errs.ErrElasticAddressesEmpty
|
||||
}
|
||||
|
||||
timeout := time.Duration(cfg.Timeout) * time.Second
|
||||
if timeout <= 0 {
|
||||
timeout = 5 * time.Second
|
||||
}
|
||||
|
||||
client := &ElasticsearchClient{
|
||||
addresses: addresses,
|
||||
username: cfg.Username,
|
||||
password: cfg.Password,
|
||||
client: &http.Client{
|
||||
Timeout: timeout,
|
||||
},
|
||||
}
|
||||
if err := client.Ping(context.Background()); err != nil {
|
||||
return nil, errs.Wrap(errs.ErrElasticPingFailed, err)
|
||||
}
|
||||
return client, nil
|
||||
}
|
||||
|
||||
func (c *ElasticsearchClient) Ping(ctx context.Context) error {
|
||||
if c == nil {
|
||||
return errs.ErrElasticClientNil
|
||||
}
|
||||
var lastErr error
|
||||
for _, address := range c.addresses {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, address, nil)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
if c.username != "" {
|
||||
req.SetBasicAuth(c.username, c.password)
|
||||
}
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
body := struct {
|
||||
Tagline string `json:"tagline"`
|
||||
}{}
|
||||
decodeErr := json.NewDecoder(resp.Body).Decode(&body)
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 && decodeErr == nil {
|
||||
return nil
|
||||
}
|
||||
lastErr = errs.Detailf(errs.ErrElasticStatusDecode, "status=%d decodeErr=%v", resp.StatusCode, decodeErr)
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
func (c *ElasticsearchClient) UpsertDocument(ctx context.Context, index string, docID string, body map[string]interface{}) error {
|
||||
if c == nil {
|
||||
return errs.ErrElasticClientNil
|
||||
}
|
||||
if strings.TrimSpace(index) == "" || strings.TrimSpace(docID) == "" {
|
||||
return errs.ErrElasticIndexOrDocID
|
||||
}
|
||||
payload, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for _, address := range c.addresses {
|
||||
url := fmt.Sprintf("%s/%s/_doc/%s", address, index, docID)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if c.username != "" {
|
||||
req.SetBasicAuth(c.username, c.password)
|
||||
}
|
||||
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
return nil
|
||||
}
|
||||
lastErr = errs.Detailf(errs.ErrElasticStatus, "%d", resp.StatusCode)
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
func (c *ElasticsearchClient) SearchIDs(ctx context.Context, index string, searchBody map[string]interface{}) ([]int64, error) {
|
||||
if c == nil {
|
||||
return nil, errs.ErrElasticClientNil
|
||||
}
|
||||
payload, err := json.Marshal(searchBody)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for _, address := range c.addresses {
|
||||
url := fmt.Sprintf("%s/%s/_search", address, index)
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
if c.username != "" {
|
||||
httpReq.SetBasicAuth(c.username, c.password)
|
||||
}
|
||||
|
||||
resp, err := c.client.Do(httpReq)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
bodyBytes, readErr := io.ReadAll(resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if readErr != nil {
|
||||
lastErr = readErr
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
lastErr = errs.Detailf(errs.ErrElasticStatusWithBody, "status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(bodyBytes)))
|
||||
continue
|
||||
}
|
||||
|
||||
var parsed struct {
|
||||
Hits struct {
|
||||
Hits []struct {
|
||||
ID string `json:"_id"`
|
||||
Source struct {
|
||||
ID int64 `json:"id"`
|
||||
} `json:"_source"`
|
||||
} `json:"hits"`
|
||||
} `json:"hits"`
|
||||
}
|
||||
if err = json.Unmarshal(bodyBytes, &parsed); err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
|
||||
ids := make([]int64, 0, len(parsed.Hits.Hits))
|
||||
for _, hit := range parsed.Hits.Hits {
|
||||
if hit.Source.ID > 0 {
|
||||
ids = append(ids, hit.Source.ID)
|
||||
continue
|
||||
}
|
||||
id, parseErr := strconv.ParseInt(hit.ID, 10, 64)
|
||||
if parseErr != nil || id <= 0 {
|
||||
continue
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids, nil
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
func (c *ElasticsearchClient) IndexExists(ctx context.Context, index string) (bool, error) {
|
||||
if c == nil {
|
||||
return false, errs.ErrElasticClientNil
|
||||
}
|
||||
index = strings.TrimSpace(index)
|
||||
if index == "" {
|
||||
return false, errs.ErrElasticIndexEmpty
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for _, address := range c.addresses {
|
||||
url := fmt.Sprintf("%s/%s", address, index)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
if c.username != "" {
|
||||
req.SetBasicAuth(c.username, c.password)
|
||||
}
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusNotFound {
|
||||
return false, nil
|
||||
}
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
return true, nil
|
||||
}
|
||||
lastErr = errs.Detailf(errs.ErrElasticStatus, "%d", resp.StatusCode)
|
||||
}
|
||||
return false, lastErr
|
||||
}
|
||||
|
||||
func (c *ElasticsearchClient) CreateIndex(ctx context.Context, index string, body map[string]interface{}) error {
|
||||
if c == nil {
|
||||
return errs.ErrElasticClientNil
|
||||
}
|
||||
index = strings.TrimSpace(index)
|
||||
if index == "" {
|
||||
return errs.ErrElasticIndexEmpty
|
||||
}
|
||||
|
||||
payload := []byte("{}")
|
||||
if body != nil {
|
||||
encoded, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload = encoded
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for _, address := range c.addresses {
|
||||
url := fmt.Sprintf("%s/%s", address, index)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if c.username != "" {
|
||||
req.SetBasicAuth(c.username, c.password)
|
||||
}
|
||||
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
bodyBytes, readErr := io.ReadAll(resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if readErr != nil {
|
||||
lastErr = readErr
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
return nil
|
||||
}
|
||||
lastErr = errs.Detailf(errs.ErrElasticStatusWithBody, "status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(bodyBytes)))
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
func CloseElasticsearch() error {
|
||||
ES = nil
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user