backend/pkg/storage/elasticsearch.go

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
}