ragflow/internal/engine/elasticsearch/client.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package elasticsearch

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"ragflow/internal/server"
	"time"

	"github.com/elastic/go-elasticsearch/v8"
	"github.com/elastic/go-elasticsearch/v8/esapi"
)

// Engine Elasticsearch engine implementation
type elasticsearchEngine struct {
	client *elasticsearch.Client
	config *server.ElasticsearchConfig
}

// NewEngine creates an Elasticsearch engine
func NewEngine(cfg interface{}) (*elasticsearchEngine, error) {
	esConfig, ok := cfg.(*server.ElasticsearchConfig)
	if !ok {
		return nil, fmt.Errorf("invalid Elasticsearch config type, expected *config.ElasticsearchConfig")
	}

	// Create ES client
	client, err := elasticsearch.NewClient(elasticsearch.Config{
		Addresses: []string{esConfig.Hosts},
		Username:  esConfig.Username,
		Password:  esConfig.Password,
		Transport: &http.Transport{
			MaxIdleConnsPerHost:   10,
			ResponseHeaderTimeout: 30 * time.Second,
		},
	})
	if err != nil {
		return nil, fmt.Errorf("failed to create Elasticsearch client: %w", err)
	}

	// Check connection
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	req := esapi.InfoRequest{}
	res, err := req.Do(ctx, client)
	if err != nil {
		return nil, fmt.Errorf("failed to ping Elasticsearch: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		return nil, fmt.Errorf("Elasticsearch returned error: %s", res.Status())
	}

	engine := &elasticsearchEngine{
		client: client,
		config: esConfig,
	}

	return engine, nil
}

// Type returns the engine type
func (e *elasticsearchEngine) Type() string {
	return "elasticsearch"
}

// Ping health check
func (e *elasticsearchEngine) Ping(ctx context.Context) error {
	req := esapi.InfoRequest{}
	res, err := req.Do(ctx, e.client)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.IsError() {
		return fmt.Errorf("elasticsearch ping failed: %s", res.Status())
	}
	return nil
}

// Close closes the connection
func (e *elasticsearchEngine) Close() error {
	// Go-elasticsearch client doesn't have a Close method, connection is managed by the transport
	return nil
}

// GetClusterStats gets Elasticsearch cluster statistics
// Reference: curl -XGET "http://{es_host}/_cluster/stats" -H "kbn-xsrf: reporting"
func (e *elasticsearchEngine) GetClusterStats() (map[string]interface{}, error) {
	req := esapi.ClusterStatsRequest{}
	res, err := req.Do(context.Background(), e.client)
	if err != nil {
		return nil, fmt.Errorf("failed to get cluster stats: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		return nil, fmt.Errorf("elasticsearch cluster stats returned error: %s", res.Status())
	}

	var rawStats map[string]interface{}
	if err := json.NewDecoder(res.Body).Decode(&rawStats); err != nil {
		return nil, fmt.Errorf("failed to decode cluster stats: %w", err)
	}

	result := make(map[string]interface{})

	// Basic cluster info
	if clusterName, ok := rawStats["cluster_name"].(string); ok {
		result["cluster_name"] = clusterName
	}
	if status, ok := rawStats["status"].(string); ok {
		result["status"] = status
	}

	// Indices info
	if indices, ok := rawStats["indices"].(map[string]interface{}); ok {
		if count, ok := indices["count"].(float64); ok {
			result["indices"] = int(count)
		}
		if shards, ok := indices["shards"].(map[string]interface{}); ok {
			if total, ok := shards["total"].(float64); ok {
				result["indices_shards"] = int(total)
			}
		}
		if docs, ok := indices["docs"].(map[string]interface{}); ok {
			if docCount, ok := docs["count"].(float64); ok {
				result["docs"] = int64(docCount)
			}
			if deleted, ok := docs["deleted"].(float64); ok {
				result["docs_deleted"] = int64(deleted)
			}
		}
		if store, ok := indices["store"].(map[string]interface{}); ok {
			if sizeInBytes, ok := store["size_in_bytes"].(float64); ok {
				result["store_size"] = convertBytes(int64(sizeInBytes))
			}
			if totalDataSetSize, ok := store["total_data_set_size_in_bytes"].(float64); ok {
				result["total_dataset_size"] = convertBytes(int64(totalDataSetSize))
			}
		}
		if mappings, ok := indices["mappings"].(map[string]interface{}); ok {
			if fieldCount, ok := mappings["total_field_count"].(float64); ok {
				result["mappings_fields"] = int(fieldCount)
			}
			if dedupFieldCount, ok := mappings["total_deduplicated_field_count"].(float64); ok {
				result["mappings_deduplicated_fields"] = int(dedupFieldCount)
			}
			if dedupSize, ok := mappings["total_deduplicated_mapping_size_in_bytes"].(float64); ok {
				result["mappings_deduplicated_size"] = convertBytes(int64(dedupSize))
			}
		}
	}

	// Nodes info
	if nodes, ok := rawStats["nodes"].(map[string]interface{}); ok {
		if count, ok := nodes["count"].(map[string]interface{}); ok {
			if total, ok := count["total"].(float64); ok {
				result["nodes"] = int(total)
			}
		}
		if versions, ok := nodes["versions"].([]interface{}); ok {
			result["nodes_version"] = versions
		}
		if os, ok := nodes["os"].(map[string]interface{}); ok {
			if mem, ok := os["mem"].(map[string]interface{}); ok {
				if totalInBytes, ok := mem["total_in_bytes"].(float64); ok {
					result["os_mem"] = convertBytes(int64(totalInBytes))
				}
				if usedInBytes, ok := mem["used_in_bytes"].(float64); ok {
					result["os_mem_used"] = convertBytes(int64(usedInBytes))
				}
				if usedPercent, ok := mem["used_percent"].(float64); ok {
					result["os_mem_used_percent"] = usedPercent
				}
			}
		}
		if jvm, ok := nodes["jvm"].(map[string]interface{}); ok {
			if versions, ok := jvm["versions"].([]interface{}); ok && len(versions) > 0 {
				if version0, ok := versions[0].(map[string]interface{}); ok {
					if vmVersion, ok := version0["vm_version"].(string); ok {
						result["jvm_versions"] = vmVersion
					}
				}
			}
			if mem, ok := jvm["mem"].(map[string]interface{}); ok {
				if heapUsed, ok := mem["heap_used_in_bytes"].(float64); ok {
					result["jvm_heap_used"] = convertBytes(int64(heapUsed))
				}
				if heapMax, ok := mem["heap_max_in_bytes"].(float64); ok {
					result["jvm_heap_max"] = convertBytes(int64(heapMax))
				}
			}
		}
	}

	return result, nil
}

// convertBytes converts bytes to human readable format
func convertBytes(bytes int64) string {
	const (
		KB = 1024
		MB = 1024 * KB
		GB = 1024 * MB
		TB = 1024 * GB
		PB = 1024 * TB
	)

	if bytes >= PB {
		return fmt.Sprintf("%.2f pb", float64(bytes)/float64(PB))
	}
	if bytes >= TB {
		return fmt.Sprintf("%.2f tb", float64(bytes)/float64(TB))
	}
	if bytes >= GB {
		return fmt.Sprintf("%.2f gb", float64(bytes)/float64(GB))
	}
	if bytes >= MB {
		return fmt.Sprintf("%.2f mb", float64(bytes)/float64(MB))
	}
	if bytes >= KB {
		return fmt.Sprintf("%.2f kb", float64(bytes)/float64(KB))
	}
	return fmt.Sprintf("%d b", bytes)
}