zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-25 01:20:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: fix context (#18061)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-08-10 20:49:32 +08:00
committed by GitHub
parent a0438517bc
commit 38c5d7c338
9 changed files with 39 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
internal/engine/elasticsearch/client.go
View File

@@ -40,7 +40,7 @@ type Engine struct {
}
// NewEngine creates an Elasticsearch engine
func NewEngine(esConfig config.ElasticsearchConfig) (*Engine, error) {
func NewEngine(ctx context.Context, esConfig config.ElasticsearchConfig) (*Engine, error) {
// Create ES client
client, err := elasticsearch.NewClient(elasticsearch.Config{
Addresses: []string{esConfig.Hosts},
@@ -57,11 +57,11 @@ func NewEngine(esConfig config.ElasticsearchConfig) (*Engine, error) {
}
// Check connection
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
newCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
req := esapi.InfoRequest{}
res, err := req.Do(ctx, client)
res, err := req.Do(newCtx, client)
if err != nil {
return nil, fmt.Errorf("failed to ping Elasticsearch: %w", err)
}
@@ -78,11 +78,11 @@ func NewEngine(esConfig config.ElasticsearchConfig) (*Engine, error) {
// Create two index templates for different index types
// Template for chunk indices (ragflow_*) - priority 1
if err = engine.CreateIndexTemplate(context.Background(), "ragflow_mapping", "ragflow_*", "mapping.json", 1); err != nil {
if err = engine.CreateIndexTemplate(newCtx, "ragflow_mapping", "ragflow_*", "mapping.json", 1); err != nil {
return nil, fmt.Errorf("failed to create chunk index template: %w", err)
}
// Template for doc_meta indices (ragflow_doc_meta_*) - priority 2 (higher than ragflow_*)
if err = engine.CreateIndexTemplate(context.Background(), "ragflow_doc_meta_mapping", "ragflow_doc_meta_*", "doc_meta_es_mapping.json", 2); err != nil {
if err = engine.CreateIndexTemplate(newCtx, "ragflow_doc_meta_mapping", "ragflow_doc_meta_*", "doc_meta_es_mapping.json", 2); err != nil {
return nil, fmt.Errorf("failed to create doc_meta index template: %w", err)
}
@@ -201,9 +201,9 @@ func (e *Engine) CreateIndexTemplate(ctx context.Context, templateName, indexPat
// GetClusterStats gets Elasticsearch cluster statistics
// Reference: curl -XGET "http://{es_host}/_cluster/stats" -H "kbn-xsrf: reporting"
func (e *Engine) GetClusterStats() (map[string]interface{}, error) {
func (e *Engine) GetClusterStats(ctx context.Context) (map[string]interface{}, error) {
req := esapi.ClusterStatsRequest{}
res, err := req.Do(context.Background(), e.client)
res, err := req.Do(ctx, e.client)
if err != nil {
return nil, fmt.Errorf("failed to get cluster stats: %w", err)
}
@@ -378,7 +378,7 @@ func extractErrorReason(bodyBytes []byte) string {
// GetIndexStats gets statistics for specified indices using the _cat/indices API
// Returns index, health, status, docs.count, store.size, dataset.size for each index
func (e *Engine) GetIndexStats(indices []string) ([]map[string]interface{}, error) {
func (e *Engine) GetIndexStats(ctx context.Context, indices []string) ([]map[string]interface{}, error) {
if len(indices) == 0 {
return []map[string]interface{}{}, nil
}
@@ -389,7 +389,7 @@ func (e *Engine) GetIndexStats(indices []string) ([]map[string]interface{}, erro
H: []string{"index", "health", "status", "docs.count", "store.size", "dataset.size"},
}
res, err := req.Do(context.Background(), e.client)
res, err := req.Do(ctx, e.client)
if err != nil {
return nil, fmt.Errorf("failed to get index stats: %w", err)
}

7
internal/engine/global.go
View File

@@ -17,6 +17,7 @@
package engine
import (
"context"
"fmt"
"sync"
@@ -40,7 +41,7 @@ var (
)
// InitDocEngine initializes document engine
func InitDocEngine() error {
func InitDocEngine(ctx context.Context) error {
var initErr error
once.Do(func() {
@@ -50,9 +51,9 @@ func InitDocEngine() error {
var err error
switch engineType {
case "elasticsearch":
globalEngine, err = elasticsearch.NewEngine(globalConfig.GetElasticsearchConfig())
globalEngine, err = elasticsearch.NewEngine(ctx, globalConfig.GetElasticsearchConfig())
case "infinity":
globalEngine, err = infinity.NewEngine(globalConfig.GetInfinityConfig())
globalEngine, err = infinity.NewEngine(ctx, globalConfig.GetInfinityConfig())
case "oceanbase", "seekdb":
connectionConfig, resolveErr := globalConfig.ResolveOceanBaseConnection(engineType)
if resolveErr != nil {

3
internal/engine/infinity/chunk.go
View File

@@ -553,9 +553,6 @@ func (e *Engine) AdjustChunkPagerank(ctx context.Context, baseName, chunkID, dat
if chunkID == "" {
return fmt.Errorf("chunk id cannot be empty")
}
if ctx == nil {
ctx = context.Background()
}
if e.client == nil || e.client.pool == nil {
return fmt.Errorf("infinity client not initialized")
}

6
internal/engine/infinity/client.go
View File

@@ -340,7 +340,7 @@ type Engine struct {
}
// NewEngine creates an Infinity engine
func NewEngine(infinityConfig config.InfinityConfig) (*Engine, error) {
func NewEngine(ctx context.Context, infinityConfig config.InfinityConfig) (*Engine, error) {
client, err := NewInfinityClient(infinityConfig)
if err != nil {
@@ -364,12 +364,12 @@ func NewEngine(infinityConfig config.InfinityConfig) (*Engine, error) {
}
// Wait for Infinity to be healthy
if err = client.WaitForHealthy(context.Background(), 120*time.Second); err != nil {
if err = client.WaitForHealthy(ctx, 120*time.Second); err != nil {
return nil, fmt.Errorf("infinity not healthy: %w", err)
}
// MigrateDB creates the database if it doesn't exist
if err = engine.MigrateDB(context.Background()); err != nil {
if err = engine.MigrateDB(ctx); err != nil {
return nil, fmt.Errorf("failed to migrate database: %w", err)
}