tm_back/pkg/elasticsearch/client.go

package elasticsearch

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"github.com/elastic/go-elasticsearch/v8"
	"tm/pkg/logger"
)

type client struct {
	es     *elasticsearch.Client
	config Config
	logger logger.Logger
	queue  chan AuditDocument
	done   chan struct{}
	wg     sync.WaitGroup
}

// NewClient creates an Elasticsearch client for audit log indexing and search.
func NewClient(config Config, log logger.Logger) (Client, error) {
	if !config.Enabled || strings.TrimSpace(config.URL) == "" {
		return NewNoopClient(), nil
	}

	if config.RequestTimeout == 0 {
		config.RequestTimeout = 10 * time.Second
	}
	if config.BulkFlushPeriod == 0 {
		config.BulkFlushPeriod = 5 * time.Second
	}
	if config.QueueSize == 0 {
		config.QueueSize = 1000
	}
	if config.IndexPrefix == "" {
		config.IndexPrefix = "tm-audit-logs"
	}

	esCfg := elasticsearch.Config{
		Addresses: []string{config.URL},
		Username:  config.Username,
		Password:  config.Password,
	}

	es, err := elasticsearch.NewClient(esCfg)
	if err != nil {
		return nil, fmt.Errorf("create elasticsearch client: %w", err)
	}

	c := &client{
		es:     es,
		config: config,
		logger: log,
		queue:  make(chan AuditDocument, config.QueueSize),
		done:   make(chan struct{}),
	}

	ctx, cancel := context.WithTimeout(context.Background(), config.RequestTimeout)
	defer cancel()
	if err := c.ping(ctx); err != nil {
		return nil, err
	}

	c.wg.Add(1)
	go c.runBulkWorker()

	log.Info("Successfully connected to Elasticsearch", map[string]interface{}{
		"url":          config.URL,
		"index_prefix": config.IndexPrefix,
	})

	return c, nil
}

func (c *client) ping(ctx context.Context) error {
	res, err := c.es.Ping(c.es.Ping.WithContext(ctx))
	if err != nil {
		return fmt.Errorf("ping elasticsearch: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		body, _ := io.ReadAll(res.Body)
		return fmt.Errorf("elasticsearch ping failed: %s", string(body))
	}

	return nil
}

func (c *client) Index(_ context.Context, doc AuditDocument) error {
	select {
	case c.queue <- doc:
		return nil
	default:
		c.logger.Warn("Audit log queue full, dropping entry", map[string]interface{}{
			"action": doc.Action,
		})
		return nil
	}
}

func (c *client) Search(ctx context.Context, filter SearchFilter) (*SearchResult, error) {
	indexPattern := c.config.IndexPrefix + "-*"
	query := buildSearchQuery(filter)

	payload, err := json.Marshal(query)
	if err != nil {
		return nil, fmt.Errorf("marshal search query: %w", err)
	}

	res, err := c.es.Search(
		c.es.Search.WithContext(ctx),
		c.es.Search.WithIndex(indexPattern),
		c.es.Search.WithBody(bytes.NewReader(payload)),
		c.es.Search.WithTrackTotalHits(true),
	)
	if err != nil {
		return nil, fmt.Errorf("search audit logs: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		body, _ := io.ReadAll(res.Body)
		return nil, fmt.Errorf("elasticsearch search error: %s", string(body))
	}

	return decodeSearchResponse(res.Body)
}

func (c *client) Close() error {
	close(c.done)
	c.wg.Wait()
	return nil
}

func (c *client) runBulkWorker() {
	defer c.wg.Done()

	ticker := time.NewTicker(c.config.BulkFlushPeriod)
	defer ticker.Stop()

	batch := make([]AuditDocument, 0, 64)

	flush := func() {
		if len(batch) == 0 {
			return
		}
		if err := c.flushBatch(batch); err != nil {
			c.logger.Error("Failed to flush audit logs to Elasticsearch", map[string]interface{}{
				"error": err.Error(),
				"count": len(batch),
			})
		}
		batch = batch[:0]
	}

	for {
		select {
		case <-c.done:
			for {
				select {
				case doc := <-c.queue:
					batch = append(batch, doc)
				default:
					flush()
					return
				}
			}
		case doc := <-c.queue:
			batch = append(batch, doc)
			if len(batch) >= 64 {
				flush()
			}
		case <-ticker.C:
			flush()
		}
	}
}

func (c *client) flushBatch(batch []AuditDocument) error {
	if len(batch) == 0 {
		return nil
	}

	var buf bytes.Buffer
	for _, doc := range batch {
		indexName := indexNameFor(c.config.IndexPrefix, doc.At)
		meta := fmt.Sprintf(`{"index":{"_index":"%s"}}`, indexName)
		buf.WriteString(meta)
		buf.WriteByte('\n')

		body, err := json.Marshal(doc)
		if err != nil {
			return fmt.Errorf("marshal audit document: %w", err)
		}
		buf.Write(body)
		buf.WriteByte('\n')
	}

	ctx, cancel := context.WithTimeout(context.Background(), c.config.RequestTimeout)
	defer cancel()

	res, err := c.es.Bulk(bytes.NewReader(buf.Bytes()), c.es.Bulk.WithContext(ctx))
	if err != nil {
		return fmt.Errorf("bulk index audit logs: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		responseBody, _ := io.ReadAll(res.Body)
		return fmt.Errorf("elasticsearch bulk error: %s", string(responseBody))
	}

	return nil
}

func indexNameFor(prefix string, at int64) string {
	t := time.Unix(at, 0).UTC()
	return fmt.Sprintf("%s-%04d.%02d", prefix, t.Year(), int(t.Month()))
}

func buildSearchQuery(filter SearchFilter) map[string]interface{} {
	must := make([]map[string]interface{}, 0)
	filterClauses := make([]map[string]interface{}, 0)

	addTerm := func(field, value string) {
		if value == "" {
			return
		}
		filterClauses = append(filterClauses, map[string]interface{}{
			"term": map[string]interface{}{field: value},
		})
	}

	addTerm("actor_id", filter.ActorID)
	addTerm("actor_type", filter.ActorType)
	addTerm("action", filter.Action)
	addTerm("target_type", filter.TargetType)
	addTerm("target_id", filter.TargetID)

	if filter.From > 0 || filter.To > 0 {
		rangeQuery := map[string]interface{}{}
		if filter.From > 0 {
			rangeQuery["gte"] = filter.From
		}
		if filter.To > 0 {
			rangeQuery["lte"] = filter.To
		}
		filterClauses = append(filterClauses, map[string]interface{}{
			"range": map[string]interface{}{"at": rangeQuery},
		})
	}

	boolQuery := map[string]interface{}{}
	if len(must) > 0 {
		boolQuery["must"] = must
	}
	if len(filterClauses) > 0 {
		boolQuery["filter"] = filterClauses
	}

	limit := filter.Limit
	if limit <= 0 {
		limit = 20
	}
	offset := filter.Offset
	if offset < 0 {
		offset = 0
	}

	return map[string]interface{}{
		"query": map[string]interface{}{
			"bool": boolQuery,
		},
		"sort": []map[string]interface{}{
			{"at": map[string]interface{}{"order": "desc"}},
		},
		"from": offset,
		"size": limit,
	}
}

func decodeSearchResponse(body io.Reader) (*SearchResult, error) {
	var parsed struct {
		Hits struct {
			Total struct {
				Value int64 `json:"value"`
			} `json:"total"`
			Hits []struct {
				Source AuditDocument `json:"_source"`
			} `json:"hits"`
		} `json:"hits"`
	}

	if err := json.NewDecoder(body).Decode(&parsed); err != nil {
		return nil, fmt.Errorf("decode search response: %w", err)
	}

	items := make([]AuditDocument, 0, len(parsed.Hits.Hits))
	for _, hit := range parsed.Hits.Hits {
		items = append(items, hit.Source)
	}

	return &SearchResult{
		Items:      items,
		TotalCount: parsed.Hits.Total.Value,
	}, nil
}