From 73d006fa0e3cd2870c4cce507c02a616aa7150ee Mon Sep 17 00:00:00 2001 From: wangyunlai Date: Mon, 10 Aug 2026 15:06:32 +0800 Subject: [PATCH] Add OceanBase and SeekDB Go document engine (#17780) ## What changed - add an OceanBase/SeekDB Go document engine using `database/sql` and the existing MySQL driver - preserve the Python connector's configuration, physical table names, schema, index names, and ARRAY/JSON/VECTOR encodings - implement chunk, memory, document metadata, skill, SQL, full-text, vector, and fusion search paths - support `DBMS_HYBRID_SEARCH.SEARCH` behind the existing feature flag, with SQL fallback only when the package is unavailable - wire the engine into retrieval, memory, metadata, vector hydration, and SQL chat flows - add Python/Go compatibility contracts, SQL mock tests, and an integration-tagged round-trip test --------- Co-authored-by: Jin Hai --- internal/engine/engine.go | 16 +- internal/engine/global.go | 15 +- internal/engine/oceanbase/client.go | 306 +++++++ internal/engine/oceanbase/client_test.go | 128 +++ internal/engine/oceanbase/codec.go | 393 ++++++++ .../engine/oceanbase/compatibility_test.go | 344 +++++++ internal/engine/oceanbase/crud.go | 350 ++++++++ internal/engine/oceanbase/crud_test.go | 153 ++++ internal/engine/oceanbase/document.go | 144 +++ internal/engine/oceanbase/document_test.go | 59 ++ internal/engine/oceanbase/filter.go | 234 +++++ internal/engine/oceanbase/helpers.go | 200 +++++ internal/engine/oceanbase/helpers_test.go | 83 ++ internal/engine/oceanbase/highlight.go | 199 +++++ internal/engine/oceanbase/highlight_test.go | 44 + internal/engine/oceanbase/integration_test.go | 229 +++++ internal/engine/oceanbase/metadata.go | 406 +++++++++ internal/engine/oceanbase/rows.go | 232 +++++ internal/engine/oceanbase/schema.go | 531 +++++++++++ internal/engine/oceanbase/schema_test.go | 162 ++++ internal/engine/oceanbase/search.go | 840 ++++++++++++++++++ internal/engine/oceanbase/search_test.go | 105 +++ internal/engine/oceanbase/sql.go | 51 ++ internal/engine/oceanbase/sql_test.go | 144 +++ internal/server/config.go | 8 +- internal/server/config/database_config.go | 44 + internal/server/config/doc_engine_config.go | 130 ++- internal/server/config/environments.go | 8 +- internal/server/config/general_config.go | 3 + .../server/config/oceanbase_config_test.go | 147 +++ internal/service/chat_pipeline.go | 6 +- internal/service/chat_pipeline_test.go | 3 + internal/service/chunk/vector.go | 7 +- internal/service/chunk/vector_test.go | 20 +- .../service/document/document_metadata.go | 4 +- internal/service/memory.go | 21 +- internal/service/memory_message_test.go | 70 +- internal/service/metadata_filter_test.go | 14 + internal/service/nlp/retrieval.go | 24 +- 39 files changed, 5830 insertions(+), 47 deletions(-) create mode 100644 internal/engine/oceanbase/client.go create mode 100644 internal/engine/oceanbase/client_test.go create mode 100644 internal/engine/oceanbase/codec.go create mode 100644 internal/engine/oceanbase/compatibility_test.go create mode 100644 internal/engine/oceanbase/crud.go create mode 100644 internal/engine/oceanbase/crud_test.go create mode 100644 internal/engine/oceanbase/document.go create mode 100644 internal/engine/oceanbase/document_test.go create mode 100644 internal/engine/oceanbase/filter.go create mode 100644 internal/engine/oceanbase/helpers.go create mode 100644 internal/engine/oceanbase/helpers_test.go create mode 100644 internal/engine/oceanbase/highlight.go create mode 100644 internal/engine/oceanbase/highlight_test.go create mode 100644 internal/engine/oceanbase/integration_test.go create mode 100644 internal/engine/oceanbase/metadata.go create mode 100644 internal/engine/oceanbase/rows.go create mode 100644 internal/engine/oceanbase/schema.go create mode 100644 internal/engine/oceanbase/schema_test.go create mode 100644 internal/engine/oceanbase/search.go create mode 100644 internal/engine/oceanbase/search_test.go create mode 100644 internal/engine/oceanbase/sql.go create mode 100644 internal/engine/oceanbase/sql_test.go create mode 100644 internal/server/config/oceanbase_config_test.go diff --git a/internal/engine/engine.go b/internal/engine/engine.go index c9a903d535..a23fa1dfb5 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -32,6 +32,8 @@ type EngineType string const ( EngineElasticsearch EngineType = "elasticsearch" EngineInfinity EngineType = "infinity" + EngineOceanBase EngineType = "oceanbase" + EngineSeekDB EngineType = "seekdb" EngineSereneDB EngineType = "serenedb" ) @@ -94,10 +96,16 @@ type DocEngine interface { // Type returns the engine type (helper method for runtime type checking) // This is a workaround since we can't import elasticsearch or infinity packages directly func Type(docEngine DocEngine) EngineType { - // Type checking through interface methods is not straightforward - // This is a placeholder that should be implemented differently - // or rely on configuration to know the type - return EngineType("unknown") + if docEngine == nil { + return EngineType("unknown") + } + return EngineType(docEngine.GetType()) +} + +// IsOceanBaseFamily reports whether a configured engine uses the shared +// OceanBase/SeekDB SQL implementation. +func IsOceanBaseFamily(engineName string) bool { + return engineName == string(EngineOceanBase) || engineName == string(EngineSeekDB) } type MessageQueue interface { diff --git a/internal/engine/global.go b/internal/engine/global.go index af6ab046d3..f7a9f5ae8c 100644 --- a/internal/engine/global.go +++ b/internal/engine/global.go @@ -18,15 +18,15 @@ package engine import ( "fmt" - "ragflow/internal/common" - "ragflow/internal/engine/nats" - "ragflow/internal/server" "sync" + "ragflow/internal/common" "ragflow/internal/engine/elasticsearch" "ragflow/internal/engine/infinity" + "ragflow/internal/engine/nats" + "ragflow/internal/engine/oceanbase" "ragflow/internal/engine/serenedb" - + "ragflow/internal/server" "ragflow/internal/tokenizer" "go.uber.org/zap" @@ -53,6 +53,13 @@ func InitDocEngine() error { globalEngine, err = elasticsearch.NewEngine(globalConfig.GetElasticsearchConfig()) case "infinity": globalEngine, err = infinity.NewEngine(globalConfig.GetInfinityConfig()) + case "oceanbase", "seekdb": + connectionConfig, resolveErr := globalConfig.ResolveOceanBaseConnection(engineType) + if resolveErr != nil { + err = resolveErr + } else { + globalEngine, err = oceanbase.NewEngine(engineType, connectionConfig) + } case "serenedb": globalEngine, err = serenedb.NewEngine(globalConfig.GetSereneDBConfig()) default: diff --git a/internal/engine/oceanbase/client.go b/internal/engine/oceanbase/client.go new file mode 100644 index 0000000000..c41e6314e1 --- /dev/null +++ b/internal/engine/oceanbase/client.go @@ -0,0 +1,306 @@ +// +// 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 oceanbase implements the OceanBase and SeekDB document engines over +// their MySQL-compatible SQL protocol. +package oceanbase + +import ( + "context" + "database/sql" + "fmt" + "os" + "regexp" + "strconv" + "strings" + "sync/atomic" + "time" + + "ragflow/internal/common" + "ragflow/internal/server/config" + + mysql "github.com/go-sql-driver/mysql" + "go.uber.org/zap" +) + +const ( + connectionAttempts = 2 + connectionRetryInterval = 5 * time.Second + defaultOperationTimeout = 100 * time.Second + minimumOceanBaseVersion = "4.3.5.1" + minimumHybridVersion = "4.4.1.0" + minimumSeekDBIndexRefreshVersion = "1.3.0.0" +) + +var seekDBVersionPattern = regexp.MustCompile(`(?i)\bseekdb[-\s]v?(\d+\.\d+\.\d+(?:\.\d+)?)`) + +type featureFlags struct { + enableFullTextSearch bool + useFullTextHint bool + searchOriginalContent bool + enableHybridSearch bool + useFullTextFirstFusionSearch bool +} + +// Engine is an OceanBase-family document engine backed by database/sql. +type Engine struct { + db *sql.DB + dbName string + engineType string + flags featureFlags + maxIdleConns int + hybridAvailable atomic.Bool + indexRefreshEnabled bool +} + +// NewEngine creates an OceanBase or SeekDB document engine. +func NewEngine(engineType string, cfg config.OceanBaseConnectionConfig) (*Engine, error) { + if engineType != "oceanbase" && engineType != "seekdb" { + return nil, fmt.Errorf("invalid OceanBase-family engine type: %s", engineType) + } + if cfg.MaxConnections <= 0 { + cfg.MaxConnections = 300 + } + + driverConfig := mysql.NewConfig() + driverConfig.User = cfg.User + driverConfig.Passwd = cfg.Password + driverConfig.Net = "tcp" + driverConfig.Addr = fmt.Sprintf("%s:%d", cfg.Host, cfg.Port) + driverConfig.DBName = cfg.DBName + driverConfig.ParseTime = true + driverConfig.Collation = "utf8mb4_unicode_ci" + driverConfig.Params = map[string]string{"charset": "utf8mb4"} + driverConfig.Timeout = 30 * time.Second + driverConfig.ReadTimeout = defaultOperationTimeout + driverConfig.WriteTimeout = defaultOperationTimeout + + var db *sql.DB + var lastErr error + for attempt := 0; attempt < connectionAttempts; attempt++ { + db, lastErr = sql.Open("mysql", driverConfig.FormatDSN()) + if lastErr == nil { + maxOverflow := envInt("OB_MAX_OVERFLOW", max(cfg.MaxConnections/2, 10)) + db.SetMaxOpenConns(cfg.MaxConnections + maxOverflow) + db.SetMaxIdleConns(cfg.MaxConnections) + db.SetConnMaxLifetime(3600 * time.Second) + pingCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + lastErr = db.PingContext(pingCtx) + cancel() + if lastErr == nil { + break + } + _ = db.Close() + } + if attempt+1 < connectionAttempts { + time.Sleep(connectionRetryInterval) + } + } + if lastErr != nil { + return nil, fmt.Errorf("connect to %s %s:%d: %w", engineType, cfg.Host, cfg.Port, lastErr) + } + + engine := newEngineWithDB(engineType, cfg.DBName, db) + engine.maxIdleConns = cfg.MaxConnections + if err := engine.initialize(context.Background()); err != nil { + _ = db.Close() + return nil, err + } + return engine, nil +} + +func newEngineWithDB(engineType, dbName string, db *sql.DB) *Engine { + flags := featureFlags{ + enableFullTextSearch: envBool("ENABLE_FULLTEXT_SEARCH", true), + useFullTextHint: envBool("USE_FULLTEXT_HINT", true), + searchOriginalContent: envBool("SEARCH_ORIGINAL_CONTENT", true), + enableHybridSearch: envBool("ENABLE_HYBRID_SEARCH", false), + useFullTextFirstFusionSearch: envBool("USE_FULLTEXT_FIRST_FUSION_SEARCH", true), + } + return &Engine{db: db, dbName: dbName, engineType: engineType, flags: flags} +} + +func (e *Engine) initialize(ctx context.Context) error { + var version string + if err := e.db.QueryRowContext(ctx, "SELECT OB_VERSION()").Scan(&version); err != nil { + return fmt.Errorf("get OceanBase version: %w", err) + } + if compareVersions(version, minimumOceanBaseVersion) < 0 { + return fmt.Errorf("OceanBase version must be at least %s, current version is %s", minimumOceanBaseVersion, version) + } + if err := e.initializeIndexRefresh(ctx); err != nil { + return err + } + + e.ensureQueryTimeout(ctx) + if e.flags.enableHybridSearch { + available := e.engineType == "seekdb" || compareVersions(version, minimumHybridVersion) >= 0 + e.hybridAvailable.Store(available) + if available { + // The DBMS hybrid-search DSL uses the tokenized FTS fields. This + // is the same switch made by the Python HybridSearch client path. + e.flags.searchOriginalContent = false + } else { + common.Warn("OceanBase DBMS hybrid search is unavailable for this version", + zap.String("version", version), zap.String("minimumVersion", minimumHybridVersion)) + } + } + return nil +} + +func (e *Engine) initializeIndexRefresh(ctx context.Context) error { + if e.engineType != "seekdb" { + return nil + } + + var serverVersion string + if err := e.db.QueryRowContext(ctx, "SELECT VERSION()").Scan(&serverVersion); err != nil { + return fmt.Errorf("get SeekDB version: %w", err) + } + seekDBVersion, ok := extractSeekDBVersion(serverVersion) + if !ok { + common.Warn("Could not parse SeekDB version; index refresh is disabled", + zap.String("version", serverVersion)) + return nil + } + e.indexRefreshEnabled = compareVersions(seekDBVersion, minimumSeekDBIndexRefreshVersion) >= 0 + return nil +} + +func extractSeekDBVersion(serverVersion string) (string, bool) { + match := seekDBVersionPattern.FindStringSubmatch(serverVersion) + if len(match) != 2 { + return "", false + } + return match[1], true +} + +func (e *Engine) ensureQueryTimeout(ctx context.Context) { + target := envInt64("OB_QUERY_TIMEOUT", 100_000_000) + var name string + var current int64 + if err := e.db.QueryRowContext(ctx, "SHOW VARIABLES LIKE 'ob_query_timeout'").Scan(&name, ¤t); err == nil && current >= target { + return + } + if _, err := e.db.ExecContext(ctx, fmt.Sprintf("SET GLOBAL ob_query_timeout=%d", target)); err != nil { + common.Warn("Failed to set OceanBase query timeout", zap.Error(err)) + return + } + // Existing sessions retain the old global value. Closing idle sessions + // mirrors Python's engine.dispose() while keeping in-flight work alive. + maxIdle := e.maxIdleConns + if maxIdle <= 0 { + maxIdle = e.db.Stats().MaxOpenConnections + } + e.db.SetMaxIdleConns(0) + e.db.SetMaxIdleConns(maxIdle) +} + +// Ping checks that the database connection is alive. +func (e *Engine) Ping(ctx context.Context) error { + if e == nil || e.db == nil { + return fmt.Errorf("OceanBase client is not initialized") + } + return e.db.PingContext(ctx) +} + +// Close closes the SQL connection pool. +func (e *Engine) Close() error { + if e == nil || e.db == nil { + return nil + } + return e.db.Close() +} + +// GetType returns oceanbase or seekdb, preserving the configured engine name. +func (e *Engine) GetType() string { return e.engineType } + +// SupportsPageRank reports that OceanBase applies pagerank during ranking. +func (e *Engine) SupportsPageRank() bool { return true } + +// AdjustChunkPagerank atomically updates the legacy pagerank_fea column. +func (e *Engine) AdjustChunkPagerank(ctx context.Context, tableName, chunkID, datasetID string, delta, minWeight, maxWeight float64) error { + if err := validateIdentifier(tableName); err != nil { + return err + } + _, err := e.db.ExecContext(ctx, "UPDATE "+quoteIdentifier(tableName)+ + " SET pagerank_fea = GREATEST(?, LEAST(?, COALESCE(pagerank_fea, 0) + ?)) WHERE id = ? AND kb_id = ?", + minWeight, maxWeight, delta, chunkID, datasetID) + return err +} + +func envBool(name string, defaultValue bool) bool { + raw := strings.ToLower(strings.TrimSpace(os.Getenv(name))) + if raw == "" { + return defaultValue + } + return raw == "true" || raw == "1" || raw == "yes" || raw == "y" +} + +func envInt64(name string, defaultValue int64) int64 { + raw := strings.TrimSpace(os.Getenv(name)) + if raw == "" { + return defaultValue + } + value, err := strconv.ParseInt(raw, 10, 64) + if err != nil || value <= 0 { + return defaultValue + } + return value +} + +func envInt(name string, defaultValue int) int { + raw := strings.TrimSpace(os.Getenv(name)) + if raw == "" { + return defaultValue + } + value, err := strconv.Atoi(raw) + if err != nil || value < 0 { + return defaultValue + } + return value +} + +func compareVersions(left, right string) int { + parse := func(raw string) []int { + parts := strings.FieldsFunc(raw, func(r rune) bool { return r < '0' || r > '9' }) + values := make([]int, 0, len(parts)) + for _, part := range parts { + value, err := strconv.Atoi(part) + if err == nil { + values = append(values, value) + } + } + return values + } + a, b := parse(left), parse(right) + for i := 0; i < max(len(a), len(b)); i++ { + var av, bv int + if i < len(a) { + av = a[i] + } + if i < len(b) { + bv = b[i] + } + if av < bv { + return -1 + } + if av > bv { + return 1 + } + } + return 0 +} diff --git a/internal/engine/oceanbase/client_test.go b/internal/engine/oceanbase/client_test.go new file mode 100644 index 0000000000..6106466e54 --- /dev/null +++ b/internal/engine/oceanbase/client_test.go @@ -0,0 +1,128 @@ +// +// 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 oceanbase + +import ( + "context" + "regexp" + "testing" + + "github.com/DATA-DOG/go-sqlmock" +) + +func TestInitializeConfiguresIndexRefresh(t *testing.T) { + tests := []struct { + name string + engineType string + serverVersion string + wantEnabled bool + }{ + { + name: "SeekDB 1.3 enables refresh", + engineType: "seekdb", + serverVersion: "5.7.25-OceanBase seekdb-v1.3.0.0", + wantEnabled: true, + }, + { + name: "SeekDB 1.2 keeps refresh disabled", + engineType: "seekdb", + serverVersion: "5.7.25-OceanBase seekdb-v1.2.0.0", + wantEnabled: false, + }, + { + name: "OceanBase keeps refresh disabled", + engineType: "oceanbase", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Setenv("ENABLE_HYBRID_SEARCH", "false") + t.Setenv("OB_QUERY_TIMEOUT", "100000000") + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + + engine := newEngineWithDB(test.engineType, "legacy_doc", db) + mock.ExpectQuery(regexp.QuoteMeta("SELECT OB_VERSION()")). + WillReturnRows(sqlmock.NewRows([]string{"OB_VERSION()"}).AddRow("4.3.5.3")) + if test.engineType == "seekdb" { + mock.ExpectQuery(regexp.QuoteMeta("SELECT VERSION()")). + WillReturnRows(sqlmock.NewRows([]string{"VERSION()"}).AddRow(test.serverVersion)) + } + mock.ExpectQuery(regexp.QuoteMeta("SHOW VARIABLES LIKE 'ob_query_timeout'")). + WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).AddRow("ob_query_timeout", 100000000)) + + if err := engine.initialize(context.Background()); err != nil { + t.Fatal(err) + } + if engine.indexRefreshEnabled != test.wantEnabled { + t.Fatalf("indexRefreshEnabled = %t, want %t", engine.indexRefreshEnabled, test.wantEnabled) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } + }) + } +} + +func TestInitializeDisablesRefreshForUnrecognizedSeekDBVersion(t *testing.T) { + t.Setenv("ENABLE_HYBRID_SEARCH", "false") + t.Setenv("OB_QUERY_TIMEOUT", "100000000") + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + + engine := newEngineWithDB("seekdb", "legacy_doc", db) + mock.ExpectQuery(regexp.QuoteMeta("SELECT OB_VERSION()")). + WillReturnRows(sqlmock.NewRows([]string{"OB_VERSION()"}).AddRow("4.3.5.3")) + mock.ExpectQuery(regexp.QuoteMeta("SELECT VERSION()")). + WillReturnRows(sqlmock.NewRows([]string{"VERSION()"}).AddRow("5.7.25-OceanBase")) + mock.ExpectQuery(regexp.QuoteMeta("SHOW VARIABLES LIKE 'ob_query_timeout'")). + WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).AddRow("ob_query_timeout", 100000000)) + + if err := engine.initialize(context.Background()); err != nil { + t.Fatalf("initialize() error = %v", err) + } + if engine.indexRefreshEnabled { + t.Fatal("indexRefreshEnabled = true, want false") + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestExtractSeekDBVersion(t *testing.T) { + for _, test := range []struct { + raw string + want string + ok bool + }{ + {raw: "5.7.25-OceanBase seekdb-v1.3.0.0", want: "1.3.0.0", ok: true}, + {raw: "OceanBase SEEKDB 1.3.0", want: "1.3.0", ok: true}, + {raw: "OceanBase_CE 4.3.5.3", ok: false}, + } { + got, ok := extractSeekDBVersion(test.raw) + if got != test.want || ok != test.ok { + t.Errorf("extractSeekDBVersion(%q) = (%q, %t), want (%q, %t)", test.raw, got, ok, test.want, test.ok) + } + } +} diff --git a/internal/engine/oceanbase/codec.go b/internal/engine/oceanbase/codec.go new file mode 100644 index 0000000000..f25c2d5284 --- /dev/null +++ b/internal/engine/oceanbase/codec.go @@ -0,0 +1,393 @@ +// +// 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 oceanbase + +import ( + "encoding/json" + "fmt" + "math" + "reflect" + "regexp" + "sort" + "strconv" + "strings" + + "ragflow/internal/tokenizer" +) + +var vectorColumnPattern = regexp.MustCompile(`^q_(\d+)_vec$`) + +var arrayColumns = map[string]bool{ + "important_kwd": true, "question_kwd": true, "tag_kwd": true, + "position_int": true, "page_num_int": true, "top_int": true, + "source_id": true, "entities_kwd": true, +} + +var jsonColumns = map[string]bool{ + "tag_feas": true, "chunk_data": true, "metadata": true, "extra": true, "meta_fields": true, +} + +var knownChunkColumns = func() map[string]bool { + known := make(map[string]bool, len(chunkColumns)) + for _, column := range chunkColumns { + known[column.name] = true + } + return known +}() + +var memoryFieldToColumn = map[string]string{ + "message_type": "message_type_kwd", + "status": "status_int", + "content": "content_ltks", +} + +var memoryColumnToField = map[string]string{ + "message_type_kwd": "message_type", + "status_int": "status", + "content_ltks": "content", +} + +func normalizeChunk(document map[string]interface{}) (map[string]interface{}, error) { + result := make(map[string]interface{}, len(chunkColumns)+1) + extra := make(map[string]interface{}) + if existing, ok := document["extra"].(map[string]interface{}); ok { + for key, value := range existing { + extra[key] = value + } + } + for key, value := range document { + key = mapChunkField(key) + if vectorColumnPattern.MatchString(key) { + encoded, err := encodeVector(value) + if err != nil { + return nil, fmt.Errorf("encode %s: %w", key, err) + } + result[key] = encoded + continue + } + if !knownChunkColumns[key] { + extra[key] = value + continue + } + if value == nil { + switch key { + case "available_int": + result[key] = 1 + case "removed_kwd": + result[key] = "N" + case "_order_id": + result[key] = 0 + default: + result[key] = nil + } + continue + } + encoded, err := encodeColumnValue(key, value) + if err != nil { + return nil, fmt.Errorf("encode %s: %w", key, err) + } + result[key] = encoded + } + for _, column := range chunkColumns { + if _, ok := result[column.name]; ok { + continue + } + switch column.name { + case "available_int": + result[column.name] = 1 + case "removed_kwd": + result[column.name] = "N" + case "_order_id": + result[column.name] = 0 + default: + result[column.name] = nil + } + } + if len(extra) > 0 { + encoded, err := json.Marshal(extra) + if err != nil { + return nil, err + } + result["extra"] = string(encoded) + } + metadata := asStringMap(document["metadata"]) + if docID := stringValue(document["doc_id"]); docID != "" { + result["group_id"] = docID + if groupID := stringValue(metadata["_group_id"]); groupID != "" { + result["group_id"] = groupID + } + if title := stringValue(metadata["_title"]); title != "" { + result["docnm_kwd"] = title + } + } + return result, nil +} + +func mapChunkField(field string) string { + if field == "chunk_order_int" { + return "_order_id" + } + return field +} + +func normalizeMemory(document map[string]interface{}) (map[string]interface{}, error) { + result := make(map[string]interface{}, len(memoryColumns)+1) + for _, column := range memoryColumns { + result[column.name] = nil + } + for key, value := range document { + if mapped, ok := memoryFieldToColumn[key]; ok { + key = mapped + } + if key == "content_embed" { + vector, ok := floatSlice(value) + if !ok || len(vector) == 0 { + continue + } + key = fmt.Sprintf("q_%d_vec", len(vector)) + value = vector + } + if vectorColumnPattern.MatchString(key) { + encoded, err := encodeVector(value) + if err != nil { + return nil, fmt.Errorf("encode %s: %w", key, err) + } + result[key] = encoded + continue + } + if _, ok := result[key]; !ok { + continue + } + result[key] = value + } + if status, ok := result["status_int"].(bool); ok { + if status { + result["status_int"] = 1 + } else { + result["status_int"] = 0 + } + } + if result["status_int"] == nil { + result["status_int"] = 1 + } + if result["zone_id"] == nil { + result["zone_id"] = 0 + } + if content := stringValue(result["content_ltks"]); content != "" { + result["tokenized_content_ltks"] = tokenizeMemoryContent(content) + } + return result, nil +} + +func tokenizeMemoryContent(content string) string { + tokens, err := tokenizer.Tokenize(content) + if err != nil { + return content + } + fineTokens, err := tokenizer.FineGrainedTokenize(tokens) + if err != nil { + return tokens + } + return fineTokens +} + +func normalizeSkill(document map[string]interface{}, documentID string) (map[string]interface{}, error) { + result := make(map[string]interface{}, len(skillColumns)+1) + for _, column := range skillColumns { + result[column.name] = nil + } + for key, value := range document { + if vectorColumnPattern.MatchString(key) { + encoded, err := encodeVector(value) + if err != nil { + return nil, err + } + result[key] = encoded + continue + } + if _, ok := result[key]; ok { + result[key] = value + } + } + if stringValue(result["skill_id"]) == "" { + result["skill_id"] = documentID + } + for _, pair := range [][2]string{{"name", "name_tks"}, {"tags", "tags_tks"}, {"description", "description_tks"}, {"content", "content_tks"}} { + if result[pair[1]] != nil { + continue + } + original := stringValue(result[pair[0]]) + tokens, err := tokenizer.Tokenize(original) + if err != nil { + tokens = original + } + result[pair[1]] = tokens + } + return result, nil +} + +func encodeColumnValue(columnName string, value interface{}) (interface{}, error) { + if value == nil { + return nil, nil + } + if columnName == "kb_id" { + if values, ok := interfaceSlice(value); ok { + if len(values) == 0 { + return nil, nil + } + return values[0], nil + } + } + if columnName == "content_with_weight" { + if _, ok := value.(map[string]interface{}); ok { + encoded, err := json.Marshal(value) + return string(encoded), err + } + } + if arrayColumns[columnName] { + return encodeArray(value) + } + if jsonColumns[columnName] { + if raw, ok := value.(string); ok { + return raw, nil + } + encoded, err := json.Marshal(value) + return string(encoded), err + } + return value, nil +} + +func encodeUpdateValue(kind, columnName string, value interface{}) (interface{}, error) { + if vectorColumnPattern.MatchString(columnName) { + return encodeVector(value) + } + if kind == "chunk" || kind == "metadata" { + return encodeColumnValue(columnName, value) + } + return value, nil +} + +func encodeArray(value interface{}) (string, error) { + values, ok := interfaceSlice(value) + if !ok { + encoded, err := json.Marshal(value) + return string(encoded), err + } + cleaned := make([]interface{}, len(values)) + for i, item := range values { + if text, ok := item.(string); ok { + text = strings.TrimSpace(text) + text = strings.ReplaceAll(text, `\`, `\\`) + text = strings.ReplaceAll(text, "\n", `\n`) + text = strings.ReplaceAll(text, "\r", `\r`) + text = strings.ReplaceAll(text, "\t", `\t`) + cleaned[i] = text + } else { + cleaned[i] = item + } + } + encoded, err := json.Marshal(cleaned) + return string(encoded), err +} + +func encodeVector(value interface{}) (string, error) { + values, ok := floatSlice(value) + if !ok { + return "", fmt.Errorf("expected numeric vector, got %T", value) + } + parts := make([]string, len(values)) + for i, number := range values { + // pyobvector converts every component to float32 before serializing. + parts[i] = strconv.FormatFloat(float64(float32(number)), 'g', -1, 32) + } + return "[" + strings.Join(parts, ",") + "]", nil +} + +func vectorDimension(document map[string]interface{}) int { + for key, value := range document { + if matches := vectorColumnPattern.FindStringSubmatch(key); len(matches) == 2 { + dimension, _ := strconv.Atoi(matches[1]) + return dimension + } + if key == "content_embed" { + if vector, ok := floatSlice(value); ok { + return len(vector) + } + } + } + return 0 +} + +func sortedColumns(document map[string]interface{}) []string { + columns := make([]string, 0, len(document)) + for column := range document { + columns = append(columns, column) + } + sort.Strings(columns) + return columns +} + +func interfaceSlice(value interface{}) ([]interface{}, bool) { + if value == nil { + return nil, false + } + rv := reflect.ValueOf(value) + if rv.Kind() != reflect.Slice && rv.Kind() != reflect.Array { + return nil, false + } + result := make([]interface{}, rv.Len()) + for i := 0; i < rv.Len(); i++ { + result[i] = rv.Index(i).Interface() + } + return result, true +} + +func floatSlice(value interface{}) ([]float64, bool) { + values, ok := interfaceSlice(value) + if !ok { + return nil, false + } + result := make([]float64, len(values)) + for i, item := range values { + number, ok := numberToFloat(item) + if !ok { + return nil, false + } + result[i] = number + if math.IsNaN(result[i]) || math.IsInf(result[i], 0) { + return nil, false + } + } + return result, true +} + +func asStringMap(value interface{}) map[string]interface{} { + if result, ok := value.(map[string]interface{}); ok { + return result + } + return map[string]interface{}{} +} + +func stringValue(value interface{}) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} diff --git a/internal/engine/oceanbase/compatibility_test.go b/internal/engine/oceanbase/compatibility_test.go new file mode 100644 index 0000000000..a93101e775 --- /dev/null +++ b/internal/engine/oceanbase/compatibility_test.go @@ -0,0 +1,344 @@ +// +// 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 oceanbase + +import ( + "encoding/json" + "os" + "path/filepath" + "reflect" + "runtime" + "strings" + "testing" + + "ragflow/internal/engine/types" +) + +func TestPythonChunkSchemaContract(t *testing.T) { + expected := []string{ + "id", "kb_id", "doc_id", "docnm_kwd", "doc_type_kwd", "title_tks", "title_sm_tks", + "content_with_weight", "content_ltks", "content_sm_ltks", "pagerank_fea", "important_kwd", + "important_tks", "question_kwd", "question_tks", "tag_kwd", "tag_feas", "available_int", + "create_time", "create_timestamp_flt", "img_id", "position_int", "page_num_int", "top_int", + "knowledge_graph_kwd", "source_id", "entity_kwd", "entity_type_kwd", "from_entity_kwd", + "to_entity_kwd", "weight_int", "weight_flt", "entities_kwd", "rank_flt", "n_hop_with_weight", + "removed_kwd", "raptor_kwd", "raptor_layer_int", "chunk_data", "metadata", "extra", "_order_id", + "group_id", "mom_id", + } + if got := columnNames(chunkColumns); !reflect.DeepEqual(got, expected) { + t.Fatalf("Go chunk schema columns changed:\n got: %v\nwant: %v", got, expected) + } + + python := readRepoFile(t, "rag", "utils", "ob_conn.py") + for _, column := range expected { + if !strings.Contains(python, "Column(\""+column+"\"") { + t.Errorf("Python chunk schema no longer declares %q", column) + } + } + for _, snippet := range []string{ + `Column("important_kwd", ARRAY(String(256))`, + `Column("question_kwd", ARRAY(String(1024))`, + `Column("position_int", ARRAY(ARRAY(Integer))`, + `Column("metadata", JSON`, + `Column("extra", JSON`, + `server_default="1"`, + `server_default="'N'"`, + } { + if !strings.Contains(python, snippet) { + t.Errorf("Python chunk schema contract is missing %q", snippet) + } + } + + wantTypes := map[string]string{ + "important_kwd": "ARRAY(VARCHAR(256)) NULL", + "question_kwd": "ARRAY(VARCHAR(1024)) NULL", + "position_int": "ARRAY(ARRAY(INTEGER)) NULL", + "metadata": "JSON NULL", + "extra": "JSON NULL", + "available_int": "INTEGER NOT NULL DEFAULT 1", + "removed_kwd": "VARCHAR(256) NULL DEFAULT 'N'", + } + for name, want := range wantTypes { + if got := columnType(chunkColumns, name); got != want { + t.Errorf("column %s type = %q, want %q", name, got, want) + } + } +} + +func TestPythonMemoryAndMetadataSchemaContract(t *testing.T) { + memoryExpected := []string{ + "id", "message_id", "message_type_kwd", "source_id", "memory_id", "user_id", "agent_id", + "session_id", "zone_id", "valid_at", "invalid_at", "forget_at", "status_int", "content_ltks", + "tokenized_content_ltks", + } + if got := columnNames(memoryColumns); !reflect.DeepEqual(got, memoryExpected) { + t.Fatalf("Go memory schema columns changed:\n got: %v\nwant: %v", got, memoryExpected) + } + memoryPython := readRepoFile(t, "memory", "utils", "ob_conn.py") + for _, column := range memoryExpected { + if !strings.Contains(memoryPython, "Column(\""+column+"\"") { + t.Errorf("Python memory schema no longer declares %q", column) + } + } + + metadataExpected := []string{"id", "kb_id", "meta_fields"} + if got := columnNames(metadataColumns); !reflect.DeepEqual(got, metadataExpected) { + t.Fatalf("Go metadata schema columns = %v, want %v", got, metadataExpected) + } + basePython := readRepoFile(t, "common", "doc_store", "ob_conn_base.py") + for _, column := range metadataExpected { + if !strings.Contains(basePython, "Column(\""+column+"\"") { + t.Errorf("Python metadata schema no longer declares %q", column) + } + } +} + +func TestPythonPhysicalTableNameContract(t *testing.T) { + if got := metadataTableName("tenant-1"); got != "ragflow_doc_meta_tenant-1" { + t.Fatalf("metadata table = %q", got) + } + if tableKind("ragflow_tenant-1") != "chunk" || tableKind("memory_tenant-1") != "memory" || tableKind("custom_table", "skill") != "skill" { + t.Fatal("physical table kinds are not recognized") + } + + contracts := map[string][]string{ + filepath.Join("rag", "nlp", "search.py"): {`return f"ragflow_{uid}"`}, + filepath.Join("memory", "services", "messages.py"): {`f"memory_{uid}"`}, + filepath.Join("api", "db", "services", "doc_metadata_service.py"): {`f"ragflow_doc_meta_{tenant_id}"`}, + } + for path, snippets := range contracts { + content := readRepoFile(t, strings.Split(path, string(filepath.Separator))...) + for _, snippet := range snippets { + if !strings.Contains(content, snippet) { + t.Errorf("%s no longer contains table-name contract %q", path, snippet) + } + } + } +} + +func TestLegacyChunkEncodingContract(t *testing.T) { + document := map[string]interface{}{ + "id": "chunk-1", "kb_id": []string{"kb-1"}, "doc_id": "doc-1", + "available_int": nil, "removed_kwd": nil, "chunk_order_int": 7, + "metadata": map[string]interface{}{"_group_id": "group-1", "_title": "renamed"}, + "important_kwd": []string{" alpha\t", "beta\n"}, + "q_3_vec": []float64{0.1, 0.2, 0.3}, + "custom_field": "preserved", + } + got, err := normalizeChunk(document) + if err != nil { + t.Fatal(err) + } + if got["kb_id"] != "kb-1" || got["available_int"] != 1 || got["removed_kwd"] != "N" || got["_order_id"] != 7 { + t.Fatalf("legacy scalar/default encoding changed: %#v", got) + } + if got["group_id"] != "group-1" || got["docnm_kwd"] != "renamed" { + t.Fatalf("metadata denormalization changed: %#v", got) + } + if got["q_3_vec"] != "[0.1,0.2,0.3]" { + t.Fatalf("vector encoding = %q", got["q_3_vec"]) + } + var extra map[string]interface{} + if err := json.Unmarshal([]byte(got["extra"].(string)), &extra); err != nil { + t.Fatal(err) + } + if extra["custom_field"] != "preserved" { + t.Fatalf("unknown field was not preserved in extra: %#v", extra) + } +} + +func TestLegacyMemoryAliasesAndVectorEncoding(t *testing.T) { + got, err := normalizeMemory(map[string]interface{}{ + "id": "memory-1_1", "message_id": "1", "memory_id": "memory-1", + "message_type": "raw", "status": false, "content": "hello world", + "content_embed": []float64{0.25, 0.5}, + }) + if err != nil { + t.Fatal(err) + } + if got["message_type_kwd"] != "raw" || got["status_int"] != 0 || got["content_ltks"] != "hello world" { + t.Fatalf("memory aliases changed: %#v", got) + } + if got["q_2_vec"] != "[0.25,0.5]" { + t.Fatalf("memory vector encoding = %q", got["q_2_vec"]) + } + decoded := decodeLogicalRow(map[string]interface{}{ + "message_type_kwd": "raw", "status_int": int64(0), "content_ltks": "hello", + "q_2_vec": "[0.25,0.5]", + }, "memory") + if decoded["message_type"] != "raw" || decoded["status"] != false || !reflect.DeepEqual(decoded["content_embed"], []interface{}{0.25, 0.5}) { + t.Fatalf("memory read aliases changed: %#v", decoded) + } +} + +func TestSkillRowsKeepStringStatusAndSkillRowID(t *testing.T) { + decoded := decodeLogicalRow(map[string]interface{}{ + "skill_id": "skill-1", + "status": "draft", + }, "skill") + if decoded["status"] != "draft" { + t.Fatalf("skill status = %#v, want draft", decoded["status"]) + } + + expression, alias, err := selectExpression("row_id()", "skill") + if err != nil { + t.Fatal(err) + } + if expression != "`skill_id` AS `row_id`" || alias != "row_id" { + t.Fatalf("skill row ID projection = (%q, %q)", expression, alias) + } +} + +func TestDBMSHybridBodyMatchesPythonSemantics(t *testing.T) { + plan := searchPlan{ + text: &types.MatchTextExpr{MatchingText: "hello", TopN: 10, ExtraOptions: map[string]interface{}{"minimum_should_match": 0.3}}, + dense: &types.MatchDenseExpr{VectorColumnName: "q_2_vec", EmbeddingData: []float64{0.1, 0.2}, EmbeddingDataType: "float", TopN: 8, ExtraOptions: map[string]interface{}{"similarity": 0.42}}, + fusion: &types.FusionExpr{Method: "weighted_sum", FusionParams: map[string]interface{}{"weights": "0.25,0.75"}}, + } + body, ok := buildDBMSBody("chunk", map[string]interface{}{"kb_id": []string{"kb-1"}, "available_int": 0}, &types.SearchRequest{ + Offset: 2, Limit: 5, RankFeature: map[string]float64{"pagerank_fea": 0.1}, + }, plan) + if !ok { + t.Fatal("hybrid body unexpectedly required SQL fallback") + } + root, ok := body["query"].(map[string]interface{}) + if !ok { + t.Fatalf("hybrid body query leg = %#v", body["query"]) + } + query, ok := root["bool"].(map[string]interface{}) + if !ok { + t.Fatalf("hybrid body bool leg = %#v", root) + } + mustClauses, ok := query["must"].([]interface{}) + if !ok || len(mustClauses) == 0 { + t.Fatalf("hybrid body must leg = %#v", query["must"]) + } + firstClause, ok := mustClauses[0].(map[string]interface{}) + if !ok { + t.Fatalf("hybrid body must clause = %#v", mustClauses[0]) + } + must, ok := firstClause["query_string"].(map[string]interface{}) + if !ok { + t.Fatalf("hybrid body query_string leg = %#v", firstClause) + } + if must["minimum_should_match"] != "30%" || query["boost"] != 0.25 { + t.Fatalf("hybrid text leg = %#v", query) + } + knn, ok := body["knn"].(map[string]interface{}) + if !ok { + t.Fatalf("hybrid knn leg = %#v", body["knn"]) + } + if knn["k"] != 8 || knn["num_candidates"] != 16 || knn["similarity"] != 0.42 { + t.Fatalf("hybrid vector leg = %#v", knn) + } +} + +func TestMetadataJSONPushdownOperators(t *testing.T) { + predicate, args, err := buildMetaPushdownPredicate([]map[string]interface{}{ + {"key": "author", "op": "contains", "value": "Alice"}, + {"key": "year", "op": "≥", "value": "2024"}, + {"key": "tags", "op": "in", "value": "rag, database"}, + }, "and") + if err != nil { + t.Fatal(err) + } + for _, fragment := range []string{"JSON_UNQUOTE", "DECIMAL(65,20)", "JSON_CONTAINS", " AND "} { + if !strings.Contains(predicate, fragment) { + t.Errorf("metadata predicate %q is missing %q", predicate, fragment) + } + } + wantArgs := []interface{}{"$.author", "Alice", "$.year", int64(2024), "$.tags", `"rag"`, "$.tags", `"database"`} + if !reflect.DeepEqual(args, wantArgs) { + t.Fatalf("metadata args = %#v, want %#v", args, wantArgs) + } + if _, _, err := buildMetaPushdownPredicate([]map[string]interface{}{{"key": "bad-key", "op": "=", "value": "x"}}, "and"); err == nil { + t.Fatal("invalid JSON metadata key must reject push-down") + } +} + +func TestMetadataJSONPushdownRejectsUnsafeNegativeOperators(t *testing.T) { + for _, operator := range []string{"≠", "not in"} { + _, _, err := buildMetaPushdownPredicate([]map[string]interface{}{ + {"key": "tags", "op": operator, "value": []string{"a"}}, + }, "and") + if err == nil { + t.Errorf("operator %q must reject metadata push-down", operator) + } + } + + for _, operator := range []string{"=", "in"} { + if _, _, err := buildMetaPushdownPredicate([]map[string]interface{}{ + {"key": "tags", "op": operator, "value": []string{"a"}}, + }, "and"); err != nil { + t.Errorf("operator %q unexpectedly rejected metadata push-down: %v", operator, err) + } + } +} + +func TestHybridFallbackOnlyForUnavailablePackage(t *testing.T) { + if !isHybridUnavailableError(assertError("ERROR 1305: FUNCTION DBMS_HYBRID_SEARCH.SEARCH does not exist")) { + t.Fatal("missing DBMS package must trigger SQL fallback") + } + if isHybridUnavailableError(assertError("DBMS_HYBRID_SEARCH.SEARCH syntax error in query_string")) { + t.Fatal("query errors must be returned instead of silently falling back") + } + if isHybridUnavailableError(assertError("feature not supported")) { + t.Fatal("unrelated unsupported errors must not trigger SQL fallback") + } +} + +func TestCompareVersions(t *testing.T) { + if compareVersions("OceanBase_CE 4.3.5.1", "4.3.5.1") != 0 || compareVersions("4.4.1.0", "4.3.5.1") <= 0 || compareVersions("4.3.4.0", "4.3.5.1") >= 0 { + t.Fatal("OceanBase version comparison changed") + } +} + +type assertError string + +func (e assertError) Error() string { return string(e) } + +func columnNames(columns []columnDefinition) []string { + result := make([]string, len(columns)) + for i, column := range columns { + result[i] = column.name + } + return result +} + +func columnType(columns []columnDefinition, name string) string { + for _, column := range columns { + if column.name == name { + return column.typeSQL + } + } + return "" +} + +func readRepoFile(t *testing.T, parts ...string) string { + t.Helper() + _, currentFile, _, ok := runtime.Caller(0) + if !ok { + t.Fatal("cannot locate compatibility test source") + } + repoRoot := filepath.Clean(filepath.Join(filepath.Dir(currentFile), "..", "..", "..")) + path := filepath.Join(append([]string{repoRoot}, parts...)...) + content, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read %s: %v", path, err) + } + return string(content) +} diff --git a/internal/engine/oceanbase/crud.go b/internal/engine/oceanbase/crud.go new file mode 100644 index 0000000000..c273b1a849 --- /dev/null +++ b/internal/engine/oceanbase/crud.go @@ -0,0 +1,350 @@ +// +// 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 oceanbase + +import ( + "context" + "database/sql" + "fmt" + "strings" + + "ragflow/internal/engine/types" +) + +// InsertChunks writes chunks or memory messages with legacy REPLACE semantics. +func (e *Engine) InsertChunks(ctx context.Context, chunks []map[string]interface{}, baseName, datasetID string) ([]string, error) { + if len(chunks) == 0 { + return []string{}, nil + } + if err := validateIdentifier(baseName); err != nil { + return nil, err + } + vectorSize := 0 + for _, chunk := range chunks { + vectorSize = vectorDimension(chunk) + if vectorSize > 0 { + break + } + } + exists, err := e.ChunkStoreExists(ctx, baseName, datasetID) + if err != nil { + return nil, err + } + if !exists { + if vectorSize == 0 { + return nil, fmt.Errorf("cannot infer vector size from documents") + } + if err := e.CreateChunkStore(ctx, baseName, datasetID, vectorSize, ""); err != nil { + return nil, err + } + } else if vectorSize > 0 { + if err := e.ensureVectorColumnAndIndex(ctx, baseName, vectorSize, lockPrefix(baseName)); err != nil { + return nil, err + } + } + + tx, err := e.db.BeginTx(ctx, nil) + if err != nil { + return nil, err + } + defer tx.Rollback() + for _, chunk := range chunks { + var normalized map[string]interface{} + switch { + case strings.HasPrefix(baseName, "memory_"): + normalized, err = normalizeMemory(chunk) + if err != nil { + return nil, err + } + if normalized["memory_id"] == nil { + normalized["memory_id"] = datasetID + } + case strings.HasPrefix(baseName, "skill_") || datasetID == "skill": + documentID := stringValue(chunk["skill_id"]) + if documentID == "" { + documentID = stringValue(chunk["id"]) + } + normalized, err = normalizeSkill(chunk, documentID) + if err != nil { + return nil, err + } + default: + normalized, err = normalizeChunk(chunk) + if err != nil { + return nil, err + } + if normalized["kb_id"] == nil || normalized["kb_id"] == "" { + normalized["kb_id"] = datasetID + } + } + if err := replaceRow(ctx, tx, baseName, normalized); err != nil { + return nil, err + } + } + if err := tx.Commit(); err != nil { + return nil, err + } + if vectorSize > 0 { + if err := e.waitForIndexRefresh(ctx); err != nil { + return nil, err + } + } + return []string{}, nil +} + +func (e *Engine) waitForIndexRefresh(ctx context.Context) error { + if !e.indexRefreshEnabled { + return nil + } + if _, err := e.db.ExecContext(ctx, "CALL DBMS_INDEX_MANAGER.REFRESH()"); err != nil { + return fmt.Errorf("wait for SeekDB index refresh: %w", err) + } + return nil +} + +func replaceRow(ctx context.Context, tx *sql.Tx, tableName string, document map[string]interface{}) error { + columns := sortedColumns(document) + quoted := make([]string, len(columns)) + placeholders := make([]string, len(columns)) + values := make([]interface{}, len(columns)) + for i, column := range columns { + if err := validateIdentifier(column); err != nil { + return err + } + quoted[i] = quoteIdentifier(column) + placeholders[i] = "?" + values[i] = document[column] + } + query := fmt.Sprintf("REPLACE INTO %s (%s) VALUES (%s)", quoteIdentifier(tableName), + strings.Join(quoted, ", "), strings.Join(placeholders, ", ")) + if _, err := tx.ExecContext(ctx, query, values...); err != nil { + return fmt.Errorf("replace row in %s: %w", tableName, err) + } + return nil +} + +// GetChunk loads one row, scoped by the requested datasets. +func (e *Engine) GetChunk(ctx context.Context, baseName, chunkID string, datasetIDs []string) (interface{}, error) { + if err := validateIdentifier(baseName); err != nil { + return nil, err + } + kind := tableKind(baseName, datasetIDs...) + exists, err := e.tableExists(ctx, baseName) + if err != nil { + return nil, err + } + if !exists { + if kind == "memory" { + return nil, fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID) + } + return nil, nil + } + identifier := "id" + if kind == "skill" { + identifier = "skill_id" + } + condition := map[string]interface{}{identifier: chunkID} + if kind == "memory" { + condition["memory_id"] = datasetIDs + } else if kind == "chunk" { + condition["kb_id"] = datasetIDs + } + whereSQL, args, err := buildFilter(condition, kind) + if err != nil { + return nil, err + } + rows, err := e.queryRows(ctx, "SELECT * FROM "+quoteIdentifier(baseName)+" WHERE "+whereSQL+" LIMIT 1", args...) + if err != nil { + return nil, err + } + if len(rows) == 0 { + if kind == "memory" { + return nil, fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID) + } + return nil, nil + } + return decodeLogicalRow(rows[0], kind), nil +} + +// UpdateChunks updates rows while retaining the dataset-level scope used by +// the Python connector. +func (e *Engine) UpdateChunks(ctx context.Context, condition, newValue map[string]interface{}, baseName, datasetID string) error { + if err := validateIdentifier(baseName); err != nil { + return err + } + kind := tableKind(baseName, datasetID) + condition = copyMap(condition) + _, scopedByDocument := condition["doc_id"] + if kind == "memory" { + condition["memory_id"] = datasetID + } else if kind == "chunk" { + condition["kb_id"] = datasetID + } + exists, err := e.tableExists(ctx, baseName) + if err != nil { + return err + } + if !exists { + return fmt.Errorf("%w: '%s'", types.ErrIndexNotFound, baseName) + } + whereSQL, args, err := buildFilter(condition, kind) + if err != nil { + return err + } + setParts := make([]string, 0, len(newValue)) + setArgs := make([]interface{}, 0, len(newValue)) + for key, value := range newValue { + if kind == "memory" { + key = mapMemoryField(key) + } else if kind == "chunk" { + key = mapChunkField(key) + } else if kind == "skill" && key == "id" { + key = "skill_id" + } + switch key { + case "id": + continue + case "remove": + switch remove := value.(type) { + case string: + if kind == "chunk" { + remove = mapChunkField(remove) + } + if err := validateIdentifier(remove); err != nil { + return err + } + setParts = append(setParts, quoteIdentifier(remove)+" = NULL") + case map[string]interface{}: + for column, item := range remove { + if kind != "chunk" || !arrayColumns[column] { + return fmt.Errorf("column %s is not an array column", column) + } + setParts = append(setParts, fmt.Sprintf("%s = ARRAY_REMOVE(%s, ?)", quoteIdentifier(column), quoteIdentifier(column))) + setArgs = append(setArgs, item) + } + default: + return fmt.Errorf("remove must be a field name or object") + } + case "add": + items, ok := value.(map[string]interface{}) + if !ok { + return fmt.Errorf("add must be an object") + } + for column, item := range items { + if kind != "chunk" || !arrayColumns[column] { + return fmt.Errorf("column %s is not an array column", column) + } + setParts = append(setParts, fmt.Sprintf("%s = ARRAY_APPEND(%s, ?)", quoteIdentifier(column), quoteIdentifier(column))) + setArgs = append(setArgs, item) + } + default: + if err := validateIdentifier(key); err != nil { + return err + } + encoded, encodeErr := encodeUpdateValue(kind, key, value) + if encodeErr != nil { + return encodeErr + } + setParts = append(setParts, quoteIdentifier(key)+" = ?") + setArgs = append(setArgs, encoded) + if kind == "memory" && key == "content_ltks" { + setParts = append(setParts, quoteIdentifier("tokenized_content_ltks")+" = ?") + setArgs = append(setArgs, tokenizeMemoryContent(stringValue(value))) + } + if key == "metadata" && scopedByDocument { + metadata := asStringMap(value) + if groupID := stringValue(metadata["_group_id"]); groupID != "" { + setParts = append(setParts, quoteIdentifier("group_id")+" = ?") + setArgs = append(setArgs, groupID) + } + if title := stringValue(metadata["_title"]); title != "" { + setParts = append(setParts, quoteIdentifier("docnm_kwd")+" = ?") + setArgs = append(setArgs, title) + } + } + } + } + if len(setParts) == 0 { + return nil + } + setArgs = append(setArgs, args...) + _, err = e.db.ExecContext(ctx, fmt.Sprintf("UPDATE %s SET %s WHERE %s", quoteIdentifier(baseName), strings.Join(setParts, ", "), whereSQL), setArgs...) + return err +} + +// DeleteChunks deletes rows under the requested dataset scope. +func (e *Engine) DeleteChunks(ctx context.Context, condition map[string]interface{}, baseName, datasetID string) (int64, error) { + if err := validateIdentifier(baseName); err != nil { + return 0, err + } + kind := tableKind(baseName, datasetID) + condition = copyMap(condition) + if kind == "memory" { + condition["memory_id"] = datasetID + } else if kind == "chunk" { + condition["kb_id"] = datasetID + } + exists, err := e.tableExists(ctx, baseName) + if err != nil { + return 0, err + } + if !exists { + return 0, nil + } + whereSQL, args, err := buildFilter(condition, kind) + if err != nil { + return 0, err + } + result, err := e.db.ExecContext(ctx, "DELETE FROM "+quoteIdentifier(baseName)+" WHERE "+whereSQL, args...) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + +func tableKind(tableName string, datasetIDs ...string) string { + for _, datasetID := range datasetIDs { + if datasetID == "skill" { + return "skill" + } + } + switch { + case strings.HasPrefix(tableName, "memory_"): + return "memory" + case strings.HasPrefix(tableName, "skill_"): + return "skill" + case strings.HasPrefix(tableName, "ragflow_doc_meta_"): + return "metadata" + default: + return "chunk" + } +} + +func mapMemoryField(field string) string { + if mapped, ok := memoryFieldToColumn[field]; ok { + return mapped + } + return field +} + +func copyMap(source map[string]interface{}) map[string]interface{} { + result := make(map[string]interface{}, len(source)+1) + for key, value := range source { + result[key] = value + } + return result +} diff --git a/internal/engine/oceanbase/crud_test.go b/internal/engine/oceanbase/crud_test.go new file mode 100644 index 0000000000..b21eeb5cf7 --- /dev/null +++ b/internal/engine/oceanbase/crud_test.go @@ -0,0 +1,153 @@ +// +// 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 oceanbase + +import ( + "context" + "regexp" + "testing" + + "github.com/DATA-DOG/go-sqlmock" +) + +func TestInsertChunksFindsLaterVectorAndWaitsForSeekDBIndexRefresh(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + + engine := newEngineWithDB("seekdb", "legacy_doc", db) + engine.flags.enableFullTextSearch = false + engine.indexRefreshEnabled = true + + tableName := "memory_tenant_1" + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?")). + WithArgs("legacy_doc", tableName). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + for _, column := range memoryIndexColumns { + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND INDEX_NAME = ?")). + WithArgs("legacy_doc", tableName, regularIndexName(tableName, column)). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + } + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME = ?")). + WithArgs("legacy_doc", tableName, "q_2_vec"). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND INDEX_NAME = ?")). + WithArgs("legacy_doc", tableName, "q_2_vec_idx"). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + mock.ExpectBegin() + mock.ExpectExec("REPLACE INTO `memory_tenant_1`").WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectExec("REPLACE INTO `memory_tenant_1`").WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectCommit() + mock.ExpectExec(regexp.QuoteMeta("CALL DBMS_INDEX_MANAGER.REFRESH()")). + WillReturnResult(sqlmock.NewResult(0, 0)) + + _, err = engine.InsertChunks(context.Background(), []map[string]interface{}{ + { + "id": "memory-1_1", "message_id": "1", "memory_id": "memory-1", + "content": "without a vector", + }, + { + "id": "memory-1_2", "message_id": "2", "memory_id": "memory-1", + "content": "hello", "content_embed": []float64{0.1, 0.2}, + }, + }, tableName, "memory-1") + if err != nil { + t.Fatal(err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestInsertChunksReturnsNormalizationErrorWithoutPanic(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + + engine := newEngineWithDB("seekdb", "legacy_doc", db) + engine.flags.enableFullTextSearch = false + tableName := "memory_tenant_1" + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?")). + WithArgs("legacy_doc", tableName). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + for _, column := range memoryIndexColumns { + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND INDEX_NAME = ?")). + WithArgs("legacy_doc", tableName, regularIndexName(tableName, column)). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + } + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME = ?")). + WithArgs("legacy_doc", tableName, "q_2_vec"). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND INDEX_NAME = ?")). + WithArgs("legacy_doc", tableName, "q_2_vec_idx"). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + mock.ExpectBegin() + mock.ExpectRollback() + + _, err = engine.InsertChunks(context.Background(), []map[string]interface{}{{ + "id": "memory-1_1", "message_id": "1", "q_2_vec": "invalid", + }}, tableName, "memory-1") + if err == nil { + t.Fatal("InsertChunks() error = nil, want vector normalization error") + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestUpdateChunksRejectsUnsupportedRemoveType(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + + engine := newEngineWithDB("oceanbase", "legacy_doc", db) + tableName := "memory_tenant_1" + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?")). + WithArgs("legacy_doc", tableName). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + + err = engine.UpdateChunks(context.Background(), map[string]interface{}{"id": "memory-1_1"}, + map[string]interface{}{"remove": []string{"forget_at"}}, tableName, "memory-1") + if err == nil { + t.Fatal("UpdateChunks() error = nil, want unsupported remove type error") + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestWaitForIndexRefreshIsNoOpWhenDisabled(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + + engine := newEngineWithDB("seekdb", "legacy_doc", db) + if err := engine.waitForIndexRefresh(context.Background()); err != nil { + t.Fatal(err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} diff --git a/internal/engine/oceanbase/document.go b/internal/engine/oceanbase/document.go new file mode 100644 index 0000000000..f46c1c7bca --- /dev/null +++ b/internal/engine/oceanbase/document.go @@ -0,0 +1,144 @@ +// +// 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 oceanbase + +import ( + "context" + "fmt" + "strings" +) + +// IndexDocument indexes a skill document. Regular chunks use InsertChunks. +func (e *Engine) IndexDocument(ctx context.Context, indexName, docID string, doc interface{}) error { + if err := validateIdentifier(indexName); err != nil { + return err + } + if !strings.HasPrefix(indexName, "skill_") { + return fmt.Errorf("IndexDocument is supported only for skill tables") + } + document, ok := doc.(map[string]interface{}) + if !ok { + return fmt.Errorf("invalid document type %T", doc) + } + ready, err := e.ChunkStoreExists(ctx, indexName, "skill") + if err != nil { + return err + } + vectorSize := vectorDimension(document) + if !ready { + if err := e.CreateChunkStore(ctx, indexName, "skill", vectorSize, ""); err != nil { + return err + } + } else if vectorSize > 0 { + if err := e.ensureVectorColumnAndIndex(ctx, indexName, vectorSize, "ob_"); err != nil { + return err + } + } + normalized, err := normalizeSkill(document, docID) + if err != nil { + return err + } + tx, err := e.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() + if err := replaceRow(ctx, tx, indexName, normalized); err != nil { + return err + } + return tx.Commit() +} + +// BulkIndex indexes skill documents in a single transaction. +func (e *Engine) BulkIndex(ctx context.Context, indexName string, docs []interface{}) (interface{}, error) { + if err := validateIdentifier(indexName); err != nil { + return nil, err + } + if !strings.HasPrefix(indexName, "skill_") { + return nil, fmt.Errorf("BulkIndex is supported only for skill tables") + } + vectorSize := 0 + for _, raw := range docs { + if document, ok := raw.(map[string]interface{}); ok { + docID := stringValue(document["skill_id"]) + if docID == "" { + docID = stringValue(document["id"]) + } + if docID == "" { + return nil, fmt.Errorf("document identifier cannot be empty") + } + if vectorSize == 0 { + vectorSize = vectorDimension(document) + } + } + } + ready, err := e.ChunkStoreExists(ctx, indexName, "skill") + if err != nil { + return nil, err + } + if !ready { + if err := e.CreateChunkStore(ctx, indexName, "skill", vectorSize, ""); err != nil { + return nil, err + } + } else if vectorSize > 0 { + if err := e.ensureVectorColumnAndIndex(ctx, indexName, vectorSize, "ob_"); err != nil { + return nil, err + } + } + tx, err := e.db.BeginTx(ctx, nil) + if err != nil { + return nil, err + } + defer tx.Rollback() + inserted := 0 + for _, raw := range docs { + document, ok := raw.(map[string]interface{}) + if !ok { + continue + } + docID := stringValue(document["skill_id"]) + if docID == "" { + docID = stringValue(document["id"]) + } + normalized, normalizeErr := normalizeSkill(document, docID) + if normalizeErr != nil { + return nil, normalizeErr + } + if err := replaceRow(ctx, tx, indexName, normalized); err != nil { + return nil, err + } + inserted++ + } + if err := tx.Commit(); err != nil { + return nil, err + } + return map[string]interface{}{"inserted": inserted}, nil +} + +// DeleteDocument deletes a skill by primary key. +func (e *Engine) DeleteDocument(ctx context.Context, indexName, docID string) error { + if err := validateIdentifier(indexName); err != nil { + return err + } + if strings.HasPrefix(indexName, "skill_") { + _, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE skill_id = ? OR REPLACE(skill_id, '/', '_') = ?", quoteIdentifier(indexName)), docID, docID) + return err + } + primaryKey := "id" + _, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE %s = ?", quoteIdentifier(indexName), quoteIdentifier(primaryKey)), docID) + return err +} diff --git a/internal/engine/oceanbase/document_test.go b/internal/engine/oceanbase/document_test.go new file mode 100644 index 0000000000..c6299672d5 --- /dev/null +++ b/internal/engine/oceanbase/document_test.go @@ -0,0 +1,59 @@ +// +// 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 oceanbase + +import ( + "context" + "testing" + + "github.com/DATA-DOG/go-sqlmock" +) + +func TestDocumentWritesValidateIdentifiers(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("oceanbase", "legacy_doc", db) + + if err := engine.IndexDocument(context.Background(), "skill_bad`name", "skill-1", map[string]interface{}{}); err == nil { + t.Fatal("IndexDocument() error = nil, want invalid identifier error") + } + if _, err := engine.BulkIndex(context.Background(), "skill_bad`name", nil); err == nil { + t.Fatal("BulkIndex() error = nil, want invalid identifier error") + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestBulkIndexRejectsEmptyIdentifiers(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("oceanbase", "legacy_doc", db) + + if _, err := engine.BulkIndex(context.Background(), "skill_tenant", []interface{}{map[string]interface{}{"name": "missing ID"}}); err == nil { + t.Fatal("BulkIndex() error = nil, want empty identifier error") + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} diff --git a/internal/engine/oceanbase/filter.go b/internal/engine/oceanbase/filter.go new file mode 100644 index 0000000000..ecd129a55f --- /dev/null +++ b/internal/engine/oceanbase/filter.go @@ -0,0 +1,234 @@ +// +// 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 oceanbase + +import ( + "fmt" + "reflect" + "strings" +) + +func buildFilter(condition map[string]interface{}, kind string) (string, []interface{}, error) { + valid := validColumns(kind) + parts := make([]string, 0, len(condition)) + args := make([]interface{}, 0, len(condition)) + for rawKey, value := range condition { + key := rawKey + if kind == "memory" { + key = mapMemoryField(key) + } else if kind == "chunk" { + key = mapChunkField(key) + } else if kind == "skill" && key == "id" { + key = "skill_id" + } else if kind == "metadata" && key == "doc_id" { + key = "id" + } + if isEmptyFilterValue(value) { + continue + } + switch key { + case "exists": + column, ok := value.(string) + if !ok { + continue + } + if kind == "memory" { + column = mapMemoryField(column) + } + if valid[column] { + parts = append(parts, quoteIdentifier(column)+" IS NOT NULL") + } + case "must_not": + object, ok := value.(map[string]interface{}) + if !ok { + continue + } + column, ok := object["exists"].(string) + if !ok { + continue + } + if kind == "memory" { + column = mapMemoryField(column) + } + if valid[column] { + parts = append(parts, quoteIdentifier(column)+" IS NULL") + } + case "metadata_filtering_conditions": + part, partArgs, err := buildMetadataFilter(value) + if err != nil { + return "", nil, err + } + if part != "" { + parts = append(parts, part) + args = append(args, partArgs...) + } + default: + if !valid[key] { + continue + } + if kind == "chunk" && arrayColumns[key] { + values, isSlice := interfaceSlice(value) + if !isSlice { + parts = append(parts, fmt.Sprintf("ARRAY_CONTAINS(%s, ?)", quoteIdentifier(key))) + args = append(args, value) + continue + } + arrayParts := make([]string, 0, len(values)) + for _, item := range values { + arrayParts = append(arrayParts, fmt.Sprintf("ARRAY_CONTAINS(%s, ?)", quoteIdentifier(key))) + args = append(args, item) + } + if len(arrayParts) > 0 { + parts = append(parts, "("+strings.Join(arrayParts, " OR ")+")") + } + continue + } + values, isSlice := interfaceSlice(value) + if isSlice { + if len(values) == 0 { + continue + } + placeholders := make([]string, len(values)) + for i, item := range values { + placeholders[i] = "?" + args = append(args, item) + } + parts = append(parts, fmt.Sprintf("%s IN (%s)", quoteIdentifier(key), strings.Join(placeholders, ", "))) + } else { + parts = append(parts, quoteIdentifier(key)+" = ?") + args = append(args, value) + } + } + } + if len(parts) == 0 { + return "1=1", args, nil + } + return strings.Join(parts, " AND "), args, nil +} + +func buildMetadataFilter(raw interface{}) (string, []interface{}, error) { + filter, ok := raw.(map[string]interface{}) + if !ok { + return "", nil, nil + } + conditions, ok := interfaceSlice(filter["conditions"]) + if !ok || len(conditions) == 0 { + return "", nil, nil + } + logic := strings.ToUpper(stringValue(filter["logical_operator"])) + if logic == "" { + logic = "AND" + } + if logic != "AND" && logic != "OR" { + return "", nil, fmt.Errorf("unsupported metadata logical operator: %s", logic) + } + + parts := make([]string, 0, len(conditions)) + args := make([]interface{}, 0, len(conditions)*2) + for _, item := range conditions { + condition, ok := item.(map[string]interface{}) + if !ok { + continue + } + name := stringValue(condition["name"]) + operator := stringValue(condition["comparison_operator"]) + if name == "" || operator == "" || strings.ContainsAny(name, `"\\`) { + continue + } + path := "$." + name + expression := "JSON_EXTRACT(`metadata`, ?)" + value := condition["value"] + switch operator { + case "is", "is not", "=", "≠", ">", "<", "≥", "≤", "before", "after": + left := expression + comparison := operator + switch operator { + case "is": + comparison = "=" + case "is not": + comparison = "!=" + case "=", "≠", ">", "<", "≥", "≤": + left = "CAST(" + expression + " AS DECIMAL(20,10))" + comparison = map[string]string{"=": "=", "≠": "!=", ">": ">", "<": "<", "≥": ">=", "≤": "<="}[operator] + case "before": + left = "CAST(" + expression + " AS DATETIME)" + comparison = "<" + case "after": + left = "CAST(" + expression + " AS DATETIME)" + comparison = ">" + } + parts = append(parts, left+" "+comparison+" ?") + args = append(args, path, value) + case "contains", "not contains": + prefix := "" + if operator == "not contains" { + prefix = "NOT " + } + parts = append(parts, prefix+"JSON_CONTAINS("+expression+", ?)") + args = append(args, path, value) + case "start with": + parts = append(parts, expression+" LIKE CONCAT(?, '%')") + args = append(args, path, value) + case "end with": + parts = append(parts, expression+" LIKE CONCAT('%', ?)") + args = append(args, path, value) + case "empty": + parts = append(parts, "("+expression+" IS NULL OR "+expression+" = '' OR "+expression+" = '[]' OR "+expression+" = '{}')") + args = append(args, path, path, path, path) + case "not empty": + parts = append(parts, "("+expression+" IS NOT NULL AND "+expression+" != '' AND "+expression+" != '[]' AND "+expression+" != '{}')") + args = append(args, path, path, path, path) + } + } + if len(parts) == 0 { + return "", nil, nil + } + return "(" + strings.Join(parts, " "+logic+" ") + ")", args, nil +} + +func validColumns(kind string) map[string]bool { + result := map[string]bool{"_score": true} + var columns []columnDefinition + switch kind { + case "memory": + columns = memoryColumns + case "metadata": + columns = metadataColumns + case "skill": + columns = skillColumns + default: + columns = chunkColumns + } + for _, column := range columns { + result[column.name] = true + } + return result +} + +func isEmptyFilterValue(value interface{}) bool { + if value == nil { + return true + } + rv := reflect.ValueOf(value) + switch rv.Kind() { + case reflect.String, reflect.Array, reflect.Slice, reflect.Map: + return rv.Len() == 0 + case reflect.Bool: + return !rv.Bool() + } + return false +} diff --git a/internal/engine/oceanbase/helpers.go b/internal/engine/oceanbase/helpers.go new file mode 100644 index 0000000000..729ba3859c --- /dev/null +++ b/internal/engine/oceanbase/helpers.go @@ -0,0 +1,200 @@ +// +// 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 oceanbase + +import ( + "context" + "encoding/json" + "fmt" + "math" + "sort" + "strings" +) + +func (e *Engine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} { + result := make(map[string]map[string]interface{}, len(chunks)) + if len(fields) == 0 { + return result + } + for _, chunk := range chunks { + id := stringValue(chunk["id"]) + if id == "" { + id = stringValue(chunk["skill_id"]) + } + if id == "" { + continue + } + selected := make(map[string]interface{}, len(fields)) + for _, field := range fields { + if value, ok := chunk[field]; ok { + selected[field] = value + } else { + selected[field] = nil + } + } + result[id] = selected + } + return result +} + +func (e *Engine) GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} { + counts := make(map[string]int) + values := make([]string, 0) + addValue := func(value string) { + if counts[value] == 0 { + values = append(values, value) + } + counts[value]++ + } + for _, chunk := range chunks { + value := chunk[fieldName] + if items, ok := interfaceSlice(value); ok { + for _, item := range items { + text, ok := item.(string) + if ok && strings.TrimSpace(text) != "" { + addValue(text) + } + } + continue + } + text, ok := value.(string) + if !ok || strings.TrimSpace(text) == "" { + continue + } + addValue(text) + } + result := make([]map[string]interface{}, 0, len(values)) + for _, value := range values { + result = append(result, map[string]interface{}{"key": value, "count": counts[value]}) + } + sort.SliceStable(result, func(i, j int) bool { + return result[i]["count"].(int) > result[j]["count"].(int) + }) + return result +} + +func (e *Engine) GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string { + result := make(map[string]string) + marker := newHighlightMarker(keywords) + for _, chunk := range chunks { + id := stringValue(chunk["id"]) + if id == "" { + id = stringValue(chunk["skill_id"]) + } + text := stringValue(chunk[fieldName]) + if id == "" || text == "" { + continue + } + tokenizedText := "" + if fieldName == "content_with_weight" { + tokenizedText = stringValue(chunk["content_ltks"]) + } + if highlighted := marker.markText(text, tokenizedText); highlighted != "" { + result[id] = highlighted + } + } + return result +} + +func (e *Engine) GetChunkIDs(chunks []map[string]interface{}) []string { + result := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + if id := stringValue(chunk["id"]); id != "" { + result = append(result, id) + } else if id := stringValue(chunk["skill_id"]); id != "" { + result = append(result, id) + } + } + return result +} + +// KNNScores computes clean cosine scores from vectors selected with the first +// query. Retrieval includes q__vec in OceanBase-family source fields. +func (e *Engine) KNNScores(ctx context.Context, chunks []map[string]interface{}, queryVector []float64, topK int) (map[string]interface{}, error) { + if len(chunks) == 0 || len(queryVector) == 0 { + return nil, nil + } + vectorField := fmt.Sprintf("q_%d_vec", len(queryVector)) + hits := make([]interface{}, 0, len(chunks)) + for _, chunk := range chunks { + vector, ok := floatSlice(chunk[vectorField]) + if !ok { + continue + } + hits = append(hits, map[string]interface{}{"_id": stringValue(chunk["id"]), "_score": cosineSimilarity(queryVector, vector)}) + } + sort.Slice(hits, func(i, j int) bool { + return hits[i].(map[string]interface{})["_score"].(float64) > hits[j].(map[string]interface{})["_score"].(float64) + }) + if topK > 0 && len(hits) > topK { + hits = hits[:topK] + } + return map[string]interface{}{"hits": map[string]interface{}{"hits": hits}}, nil +} + +func (e *Engine) GetScores(knnResult map[string]interface{}) map[string]float64 { + result := make(map[string]float64) + if knnResult == nil { + return result + } + hitsObject, _ := knnResult["hits"].(map[string]interface{}) + hits, _ := hitsObject["hits"].([]interface{}) + for _, raw := range hits { + hit, _ := raw.(map[string]interface{}) + id := stringValue(hit["_id"]) + if score, ok := numberToFloat(hit["_score"]); ok && id != "" { + result[id] = score + } + } + return result +} + +func cosineSimilarity(left, right []float64) float64 { + if len(left) != len(right) || len(left) == 0 { + return 0 + } + var dot, leftNorm, rightNorm float64 + for i := range left { + dot += left[i] * right[i] + leftNorm += left[i] * left[i] + rightNorm += right[i] * right[i] + } + if leftNorm == 0 || rightNorm == 0 { + return 0 + } + return dot / (math.Sqrt(leftNorm) * math.Sqrt(rightNorm)) +} + +func numberToFloat(value interface{}) (float64, bool) { + switch number := value.(type) { + case float64: + return number, true + case float32: + return float64(number), true + case int: + return float64(number), true + case int32: + return float64(number), true + case int64: + return float64(number), true + case json.Number: + parsed, err := number.Float64() + return parsed, err == nil + default: + return 0, false + } +} diff --git a/internal/engine/oceanbase/helpers_test.go b/internal/engine/oceanbase/helpers_test.go new file mode 100644 index 0000000000..78343b8c81 --- /dev/null +++ b/internal/engine/oceanbase/helpers_test.go @@ -0,0 +1,83 @@ +// +// 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 oceanbase + +import ( + "reflect" + "testing" +) + +func TestGetAggregationPreservesScalarStrings(t *testing.T) { + engine := &Engine{} + chunks := []map[string]interface{}{ + {"docnm_kwd": "report,final.pdf"}, + {"docnm_kwd": "report,final.pdf"}, + {"docnm_kwd": "guide.pdf"}, + {"docnm_kwd": ""}, + {}, + } + want := []map[string]interface{}{ + {"key": "report,final.pdf", "count": 2}, + {"key": "guide.pdf", "count": 1}, + } + if got := engine.GetAggregation(chunks, "docnm_kwd"); !reflect.DeepEqual(got, want) { + t.Fatalf("GetAggregation() = %#v, want %#v", got, want) + } +} + +func TestGetAggregationCountsArrayElements(t *testing.T) { + engine := &Engine{} + chunks := []map[string]interface{}{ + {"tag_kwd": []string{"rag", "database"}}, + {"tag_kwd": []interface{}{"rag", "", 7}}, + } + want := []map[string]interface{}{ + {"key": "rag", "count": 2}, + {"key": "database", "count": 1}, + } + if got := engine.GetAggregation(chunks, "tag_kwd"); !reflect.DeepEqual(got, want) { + t.Fatalf("GetAggregation() = %#v, want %#v", got, want) + } +} + +func TestGetHighlightUsesBoundariesAndTokens(t *testing.T) { + engine := &Engine{} + chunks := []map[string]interface{}{ + {"id": "english", "content_with_weight": "Apple cat concatenate.", "content_ltks": "apple cat concatenate"}, + {"id": "chinese", "content_with_weight": "这是数据库系统", "content_ltks": "这是 数据库 系统"}, + {"id": "missing", "content_with_weight": "nothing relevant", "content_ltks": "nothing relevant"}, + } + want := map[string]string{ + "english": "Apple cat concatenate.", + "chinese": "这是数据库系统", + } + if got := engine.GetHighlight(chunks, []string{"apple", "cat", "数据库"}, "content_with_weight"); !reflect.DeepEqual(got, want) { + t.Fatalf("GetHighlight() = %#v, want %#v", got, want) + } +} + +func TestGetHighlightUsesSkillID(t *testing.T) { + engine := &Engine{} + chunks := []map[string]interface{}{{ + "skill_id": "skill-1", + "content": "OceanBase search", + }} + want := map[string]string{"skill-1": "OceanBase search"} + if got := engine.GetHighlight(chunks, []string{"oceanbase"}, "content"); !reflect.DeepEqual(got, want) { + t.Fatalf("GetHighlight() = %#v, want %#v", got, want) + } +} diff --git a/internal/engine/oceanbase/highlight.go b/internal/engine/oceanbase/highlight.go new file mode 100644 index 0000000000..28952ac853 --- /dev/null +++ b/internal/engine/oceanbase/highlight.go @@ -0,0 +1,199 @@ +// +// 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 oceanbase + +import ( + "regexp" + "sort" + "strings" + "unicode" + "unicode/utf8" +) + +type textMatch struct { + start int + end int +} + +// highlightMarker holds the normalized keywords and compiled patterns for one query. +type highlightMarker struct { + keywords []string + keywordSet map[string]struct{} + englishPatterns []*regexp.Regexp + nonEnglishPatterns []*regexp.Regexp +} + +// newHighlightMarker prepares a reusable text marker for a query's keywords. +func newHighlightMarker(keywords []string) *highlightMarker { + keywords = normalizeKeywords(keywords) + marker := &highlightMarker{ + keywords: keywords, + keywordSet: make(map[string]struct{}, len(keywords)), + englishPatterns: make([]*regexp.Regexp, 0, len(keywords)), + nonEnglishPatterns: make([]*regexp.Regexp, 0, len(keywords)), + } + for _, keyword := range keywords { + marker.keywordSet[keyword] = struct{}{} + quoted := regexp.QuoteMeta(keyword) + marker.englishPatterns = append(marker.englishPatterns, regexp.MustCompile("(?i)"+quoted)) + marker.nonEnglishPatterns = append(marker.nonEnglishPatterns, regexp.MustCompile(quoted)) + } + return marker +} + +// markText wraps matching terms in em tags. English terms are matched without +// case sensitivity at word boundaries. Non-English text uses tokenizedText +// when available so highlighting follows the indexed token boundaries. +func (m *highlightMarker) markText(text, tokenizedText string) string { + if m == nil || text == "" || len(m.keywords) == 0 { + return "" + } + + var matches []textMatch + if isMostlyEnglish(text) { + matches = findPatternMatches(text, m.englishPatterns, true) + } else if tokenizedText != "" { + matches = findTokenMatches(text, tokenizedText, m.keywordSet) + } else { + matches = findPatternMatches(text, m.nonEnglishPatterns, false) + } + if len(matches) == 0 { + return "" + } + return applyMatches(text, matches) +} + +func normalizeKeywords(keywords []string) []string { + seen := make(map[string]struct{}, len(keywords)) + result := make([]string, 0, len(keywords)) + for _, keyword := range keywords { + keyword = strings.TrimSpace(keyword) + if keyword == "" { + continue + } + key := strings.ToLower(keyword) + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + result = append(result, keyword) + } + sort.SliceStable(result, func(i, j int) bool { + return len(result[i]) > len(result[j]) + }) + return result +} + +func isMostlyEnglish(text string) bool { + letters := 0 + latinLetters := 0 + for _, r := range text { + if !unicode.IsLetter(r) { + continue + } + letters++ + if unicode.In(r, unicode.Latin) { + latinLetters++ + } + } + return letters > 0 && latinLetters*2 > letters +} + +func findPatternMatches(text string, patterns []*regexp.Regexp, requireBoundary bool) []textMatch { + candidates := make([]textMatch, 0) + for _, pattern := range patterns { + for _, indexes := range pattern.FindAllStringIndex(text, -1) { + candidate := textMatch{start: indexes[0], end: indexes[1]} + if requireBoundary && !hasWordBoundaries(text, candidate) { + continue + } + candidates = append(candidates, candidate) + } + } + return selectMatches(candidates) +} + +func findTokenMatches(text, tokenizedText string, keywordSet map[string]struct{}) []textMatch { + tokens := strings.Fields(tokenizedText) + lastPosition := len(text) + candidates := make([]textMatch, 0) + for i := len(tokens) - 1; i >= 0; i-- { + token := tokens[i] + position := strings.LastIndex(text[:lastPosition], token) + if position < 0 { + continue + } + if _, ok := keywordSet[token]; ok { + candidates = append(candidates, textMatch{start: position, end: position + len(token)}) + } + lastPosition = position + } + return selectMatches(candidates) +} + +func hasWordBoundaries(text string, match textMatch) bool { + if match.start > 0 { + previous, _ := utf8.DecodeLastRuneInString(text[:match.start]) + if isWordRune(previous) { + return false + } + } + if match.end < len(text) { + next, _ := utf8.DecodeRuneInString(text[match.end:]) + if isWordRune(next) { + return false + } + } + return true +} + +func isWordRune(r rune) bool { + return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' +} + +func selectMatches(candidates []textMatch) []textMatch { + sort.Slice(candidates, func(i, j int) bool { + if candidates[i].start != candidates[j].start { + return candidates[i].start < candidates[j].start + } + return candidates[i].end > candidates[j].end + }) + selected := make([]textMatch, 0, len(candidates)) + lastEnd := -1 + for _, candidate := range candidates { + if candidate.start < lastEnd { + continue + } + selected = append(selected, candidate) + lastEnd = candidate.end + } + return selected +} + +func applyMatches(text string, matches []textMatch) string { + var result strings.Builder + position := 0 + for _, match := range matches { + result.WriteString(text[position:match.start]) + result.WriteString("") + result.WriteString(text[match.start:match.end]) + result.WriteString("") + position = match.end + } + result.WriteString(text[position:]) + return result.String() +} diff --git a/internal/engine/oceanbase/highlight_test.go b/internal/engine/oceanbase/highlight_test.go new file mode 100644 index 0000000000..95f4543dff --- /dev/null +++ b/internal/engine/oceanbase/highlight_test.go @@ -0,0 +1,44 @@ +// +// 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 oceanbase + +import "testing" + +func TestMarkTextEnglishBoundariesAndCase(t *testing.T) { + text := "Apple, cat and dog; concatenate is different." + got := newHighlightMarker([]string{"apple", "cat", "dog"}).markText(text, "") + want := "Apple, cat and dog; concatenate is different." + if got != want { + t.Fatalf("markText() = %q, want %q", got, want) + } +} + +func TestMarkTextUsesTokenizedNonEnglishText(t *testing.T) { + text := "这是数据库系统" + if got := newHighlightMarker([]string{"数据库"}).markText(text, "这是 数据库 系统"); got != "这是数据库系统" { + t.Fatalf("markText() = %q", got) + } + if got := newHighlightMarker([]string{"数据"}).markText(text, "这是 数据库 系统"); got != "" { + t.Fatalf("partial token match = %q, want no highlight", got) + } +} + +func TestMarkTextReturnsEmptyWithoutMatch(t *testing.T) { + if got := newHighlightMarker([]string{"missing"}).markText("no relevant text", ""); got != "" { + t.Fatalf("markText() = %q, want empty", got) + } +} diff --git a/internal/engine/oceanbase/integration_test.go b/internal/engine/oceanbase/integration_test.go new file mode 100644 index 0000000000..6b3326d8cb --- /dev/null +++ b/internal/engine/oceanbase/integration_test.go @@ -0,0 +1,229 @@ +//go:build integration + +// +// 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 oceanbase + +import ( + "context" + "encoding/json" + "fmt" + "os" + "reflect" + "strconv" + "testing" + "time" + + "ragflow/internal/engine/types" + "ragflow/internal/server/config" +) + +func TestLegacyStorageRoundTrip(t *testing.T) { + host := os.Getenv("RAGFLOW_TEST_OCEANBASE_HOST") + if host == "" { + t.Skip("RAGFLOW_TEST_OCEANBASE_HOST is not set") + } + port, err := strconv.Atoi(envOr("RAGFLOW_TEST_OCEANBASE_PORT", "2881")) + if err != nil { + t.Fatal(err) + } + engine, err := NewEngine(envOr("RAGFLOW_TEST_OCEANBASE_ENGINE_TYPE", "oceanbase"), config.OceanBaseConnectionConfig{ + DBName: envOr("RAGFLOW_TEST_OCEANBASE_DBNAME", "test"), + User: envOr("RAGFLOW_TEST_OCEANBASE_USER", "root@test"), Password: os.Getenv("RAGFLOW_TEST_OCEANBASE_PASSWORD"), + Host: host, Port: port, MaxConnections: 4, + }) + if err != nil { + t.Fatal(err) + } + defer engine.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + suffix := strconv.FormatInt(time.Now().UnixNano(), 10) + tableName := "ragflow_go_compat_" + suffix + datasetID := "kb-" + suffix + defer cleanupChunkStore(t, engine, tableName, "") + + if err := engine.CreateChunkStore(ctx, tableName, datasetID, 2, "naive"); err != nil { + t.Fatal(err) + } + if _, err := engine.InsertChunks(ctx, []map[string]interface{}{{ + "id": "chunk-1", "kb_id": datasetID, "doc_id": "doc-1", "content_with_weight": "hello oceanbase", + "content_ltks": "hello oceanbase", "important_kwd": []string{"hello"}, + "metadata": map[string]interface{}{"_group_id": "group-1", "custom": "json-value"}, + "custom_field": "kept-in-extra", "q_2_vec": []float64{0.25, 0.5}, + }}, tableName, datasetID); err != nil { + t.Fatal(err) + } + row, err := engine.GetChunk(ctx, tableName, "chunk-1", []string{datasetID}) + if err != nil { + t.Fatal(err) + } + chunk := row.(map[string]interface{}) + if chunk["group_id"] != "group-1" { + t.Fatalf("legacy metadata denormalization failed: %#v", chunk) + } + + result, err := engine.Search(ctx, &types.SearchRequest{ + IndexNames: []string{tableName}, KbIDs: []string{datasetID}, Limit: 10, + SelectFields: []string{"id", "metadata", "extra", "q_2_vec"}, + MatchExprs: []interface{}{&types.MatchDenseExpr{ + VectorColumnName: "q_2_vec", EmbeddingData: []float64{0.25, 0.5}, + EmbeddingDataType: "float", TopN: 10, ExtraOptions: map[string]interface{}{"similarity": 0.1}, + }}, + }) + if err != nil { + t.Fatal(err) + } + if len(result.Chunks) != 1 || result.Chunks[0]["id"] != "chunk-1" { + t.Fatalf("vector round trip returned %#v", result) + } + + pythonChunkID := "python-chunk-1" + _, err = engine.db.ExecContext(ctx, fmt.Sprintf( + "REPLACE INTO %s (id, kb_id, doc_id, content_with_weight, content_ltks, important_kwd, metadata, extra, q_2_vec) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + quoteIdentifier(tableName), + ), pythonChunkID, datasetID, "python-doc-1", "python formatted chunk", "python formatted chunk", + `["python","legacy"]`, `{"_group_id":"python-group","custom":"python-json"}`, + `{"python_extra":"preserved"}`, `[0.75,0.25]`) + if err != nil { + t.Fatal(err) + } + pythonRow, err := engine.GetChunk(ctx, tableName, pythonChunkID, []string{datasetID}) + if err != nil { + t.Fatal(err) + } + pythonChunk := pythonRow.(map[string]interface{}) + if !reflect.DeepEqual(pythonChunk["important_kwd"], []interface{}{"python", "legacy"}) { + t.Fatalf("Python ARRAY decoding failed: %#v", pythonChunk["important_kwd"]) + } + if metadata, ok := pythonChunk["metadata"].(map[string]interface{}); !ok || metadata["custom"] != "python-json" { + t.Fatalf("Python JSON decoding failed: %#v", pythonChunk["metadata"]) + } + if vector, ok := floatSlice(pythonChunk["q_2_vec"]); !ok || !reflect.DeepEqual(vector, []float64{0.75, 0.25}) { + t.Fatalf("Python VECTOR decoding failed: %#v", pythonChunk["q_2_vec"]) + } + + assertStoredJSON(ctx, t, engine, tableName, "important_kwd", "chunk-1", []interface{}{"hello"}) + assertStoredJSON(ctx, t, engine, tableName, "metadata", "chunk-1", map[string]interface{}{"_group_id": "group-1", "custom": "json-value"}) + assertStoredJSON(ctx, t, engine, tableName, "extra", "chunk-1", map[string]interface{}{"custom_field": "kept-in-extra"}) + assertStoredJSON(ctx, t, engine, tableName, "q_2_vec", "chunk-1", []interface{}{0.25, 0.5}) + + memoryTable := "memory_go_compat_" + suffix + memoryA := "memory-a-" + suffix + memoryB := "memory-b-" + suffix + defer cleanupChunkStore(t, engine, memoryTable, "") + if err := engine.CreateChunkStore(ctx, memoryTable, memoryA, 2, ""); err != nil { + t.Fatal(err) + } + if _, err := engine.InsertChunks(ctx, []map[string]interface{}{ + {"id": memoryA + "_1", "message_id": "1", "memory_id": memoryA, "content": "first", "content_embed": []float64{0.1, 0.2}}, + {"id": memoryB + "_1", "message_id": "1", "memory_id": memoryB, "content": "second", "content_embed": []float64{0.3, 0.4}}, + }, memoryTable, memoryA); err != nil { + t.Fatal(err) + } + if err := engine.DropChunkStore(ctx, memoryTable, memoryA); err != nil { + t.Fatal(err) + } + if _, err := engine.GetChunk(ctx, memoryTable, memoryA+"_1", []string{memoryA}); err == nil { + t.Fatal("deleted memory rows are still readable") + } + if _, err := engine.GetChunk(ctx, memoryTable, memoryB+"_1", []string{memoryB}); err != nil { + t.Fatalf("another memory's rows were removed: %v", err) + } + + tenantID := "go_compat_" + suffix + metadataTable := metadataTableName(tenantID) + defer cleanupMetadataStore(t, engine, tenantID) + if err := engine.CreateMetadataStore(ctx, tenantID); err != nil { + t.Fatal(err) + } + if _, err := engine.db.ExecContext(ctx, + "REPLACE INTO "+quoteIdentifier(metadataTable)+" (id, kb_id, meta_fields) VALUES (?, ?, ?)", + "python-meta-1", datasetID, `{"tags":["a","b"],"source":"python"}`, + ); err != nil { + t.Fatal(err) + } + metadataResult, err := engine.SearchMetadata(ctx, &types.SearchMetadataRequest{ + TenantID: tenantID, Limit: 10, SelectFields: []string{"id", "kb_id", "meta_fields"}, + Filter: map[string]interface{}{"id": "python-meta-1"}, + }) + if err != nil { + t.Fatal(err) + } + if len(metadataResult.MetadataRecords) != 1 { + t.Fatalf("Python metadata row returned %#v", metadataResult.MetadataRecords) + } + metaFields, ok := metadataResult.MetadataRecords[0]["meta_fields"].(map[string]interface{}) + if !ok || metaFields["source"] != "python" { + t.Fatalf("Python metadata JSON decoding failed: %#v", metadataResult.MetadataRecords[0]) + } + if err := engine.UpdateMetadata(ctx, "go-meta-1", datasetID, map[string]interface{}{"source": "go", "tags": []string{"c"}}, tenantID); err != nil { + t.Fatal(err) + } + assertStoredJSON(ctx, t, engine, metadataTable, "meta_fields", "go-meta-1", map[string]interface{}{"source": "go", "tags": []interface{}{"c"}}) +} + +func assertStoredJSON(ctx context.Context, t *testing.T, engine *Engine, tableName, columnName, rowID string, want interface{}) { + t.Helper() + var raw interface{} + query := fmt.Sprintf("SELECT %s FROM %s WHERE id = ?", quoteIdentifier(columnName), quoteIdentifier(tableName)) + if err := engine.db.QueryRowContext(ctx, query, rowID).Scan(&raw); err != nil { + t.Fatal(err) + } + var text string + switch value := raw.(type) { + case []byte: + text = string(value) + case string: + text = value + default: + t.Fatalf("stored %s value has type %T", columnName, raw) + } + var got interface{} + if err := json.Unmarshal([]byte(text), &got); err != nil { + t.Fatalf("stored %s value %q is not JSON compatible: %v", columnName, text, err) + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("stored %s value = %#v, want %#v", columnName, got, want) + } +} + +func envOr(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + return fallback +} + +func cleanupChunkStore(t *testing.T, engine *Engine, tableName, datasetID string) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := engine.DropChunkStore(ctx, tableName, datasetID); err != nil { + t.Errorf("clean up chunk store %s: %v", tableName, err) + } +} + +func cleanupMetadataStore(t *testing.T, engine *Engine, tenantID string) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := engine.DropMetadataStore(ctx, tenantID); err != nil { + t.Errorf("clean up metadata store for tenant %s: %v", tenantID, err) + } +} diff --git a/internal/engine/oceanbase/metadata.go b/internal/engine/oceanbase/metadata.go new file mode 100644 index 0000000000..17fd3faa57 --- /dev/null +++ b/internal/engine/oceanbase/metadata.go @@ -0,0 +1,406 @@ +// +// 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 oceanbase + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/engine/types" + + "go.uber.org/zap" + "gorm.io/gorm" +) + +const metadataPushdownMaxSize = 10000 + +func metadataTableName(tenantID string) string { return "ragflow_doc_meta_" + tenantID } + +func validatedMetadataTableName(tenantID string) (string, error) { + tableName := metadataTableName(tenantID) + if err := validateIdentifier(tableName); err != nil { + return "", err + } + return tableName, nil +} + +// CreateMetadataStore creates the per-tenant metadata table. +func (e *Engine) CreateMetadataStore(ctx context.Context, tenantID string) error { + tableName, err := validatedMetadataTableName(tenantID) + if err != nil { + return err + } + if err := e.ensureTableWithLock(ctx, tableName, metadataColumns, "ob_create_doc_meta_table_"+tableName); err != nil { + return err + } + return e.ensureRegularIndex(ctx, tableName, "kb_id", "ob_") +} + +// InsertMetadata stores metadata using the same REPLACE operation as Python. +func (e *Engine) InsertMetadata(ctx context.Context, metadata []map[string]interface{}, tenantID string) ([]string, error) { + if len(metadata) == 0 { + return []string{}, nil + } + tableName, err := validatedMetadataTableName(tenantID) + if err != nil { + return nil, err + } + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return nil, err + } + if !exists { + if err := e.CreateMetadataStore(ctx, tenantID); err != nil { + return nil, err + } + } + tx, err := e.db.BeginTx(ctx, nil) + if err != nil { + return nil, err + } + defer tx.Rollback() + for _, document := range metadata { + metaFields := document["meta_fields"] + var encoded string + switch value := metaFields.(type) { + case string: + encoded = value + case map[string]interface{}: + data, marshalErr := json.Marshal(value) + if marshalErr != nil { + return nil, marshalErr + } + encoded = string(data) + default: + encoded = "{}" + } + row := map[string]interface{}{"id": document["id"], "kb_id": document["kb_id"], "meta_fields": encoded} + if err := replaceRow(ctx, tx, tableName, row); err != nil { + return nil, err + } + } + if err := tx.Commit(); err != nil { + return nil, err + } + return []string{}, nil +} + +// UpdateMetadata replaces the complete JSON object, inserting the row if it +// does not yet exist. This matches the service's replace_meta_fields contract. +func (e *Engine) UpdateMetadata(ctx context.Context, docID, datasetID string, metaFields map[string]interface{}, tenantID string) error { + tableName, err := validatedMetadataTableName(tenantID) + if err != nil { + return err + } + encoded, err := json.Marshal(metaFields) + if err != nil { + return err + } + _, err = e.db.ExecContext(ctx, fmt.Sprintf("REPLACE INTO %s (id, kb_id, meta_fields) VALUES (?, ?, ?)", quoteIdentifier(tableName)), docID, datasetID, string(encoded)) + return err +} + +// DeleteMetadata deletes matching metadata rows. +func (e *Engine) DeleteMetadata(ctx context.Context, condition map[string]interface{}, tenantID string) (int64, error) { + tableName, err := validatedMetadataTableName(tenantID) + if err != nil { + return 0, err + } + exists, err := e.tableExists(ctx, tableName) + if err != nil || !exists { + return 0, err + } + whereSQL, args, err := buildFilter(condition, "metadata") + if err != nil { + return 0, err + } + result, err := e.db.ExecContext(ctx, "DELETE FROM "+quoteIdentifier(tableName)+" WHERE "+whereSQL, args...) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + +// DeleteMetadataKeys removes selected JSON keys and deletes the row if no +// metadata remains. +func (e *Engine) DeleteMetadataKeys(ctx context.Context, docID, datasetID string, keys []string, tenantID string) error { + tableName, err := validatedMetadataTableName(tenantID) + if err != nil { + return err + } + var raw string + if err := e.db.QueryRowContext(ctx, "SELECT meta_fields FROM "+quoteIdentifier(tableName)+" WHERE id = ? AND kb_id = ? LIMIT 1", docID, datasetID).Scan(&raw); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("%w: %s", types.ErrDocumentNotFound, docID) + } + return err + } + fields := make(map[string]interface{}) + if err := json.Unmarshal([]byte(raw), &fields); err != nil { + return fmt.Errorf("decode metadata for document %s: %w", docID, err) + } + for _, key := range keys { + delete(fields, key) + } + if len(fields) == 0 { + _, err := e.db.ExecContext(ctx, "DELETE FROM "+quoteIdentifier(tableName)+" WHERE id = ? AND kb_id = ?", docID, datasetID) + return err + } + encoded, err := json.Marshal(fields) + if err != nil { + return err + } + _, err = e.db.ExecContext(ctx, "UPDATE "+quoteIdentifier(tableName)+" SET meta_fields = ? WHERE id = ? AND kb_id = ?", string(encoded), docID, datasetID) + return err +} + +func (e *Engine) DropMetadataStore(ctx context.Context, tenantID string) error { + tableName, err := validatedMetadataTableName(tenantID) + if err != nil { + return err + } + _, err = e.db.ExecContext(ctx, "DROP TABLE IF EXISTS "+quoteIdentifier(tableName)) + return err +} + +func (e *Engine) MetadataStoreExists(ctx context.Context, tenantID string) (bool, error) { + tableName, err := validatedMetadataTableName(tenantID) + if err != nil { + return false, err + } + return e.tableExists(ctx, tableName) +} + +// SearchMetadata searches a tenant metadata table with exact total count. +func (e *Engine) SearchMetadata(ctx context.Context, req *types.SearchMetadataRequest) (*types.SearchMetadataResult, error) { + if req == nil || req.TenantID == "" { + return nil, fmt.Errorf("tenantID cannot be empty") + } + tableName, err := validatedMetadataTableName(req.TenantID) + if err != nil { + return nil, err + } + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return nil, err + } + if !exists { + return &types.SearchMetadataResult{MetadataRecords: []map[string]interface{}{}}, nil + } + fieldsSQL, _, err := buildSelectFields(req.SelectFields, "metadata") + if err != nil { + return nil, err + } + whereSQL, args, err := buildFilter(req.Filter, "metadata") + if err != nil { + return nil, err + } + total, err := scanCount(e.db.QueryRowContext(ctx, "SELECT COUNT(id) FROM "+quoteIdentifier(tableName)+" WHERE "+whereSQL, args...)) + if err != nil || total == 0 { + return &types.SearchMetadataResult{MetadataRecords: []map[string]interface{}{}, Total: total}, err + } + orderSQL, err := buildOrderBy(req.OrderBy, "metadata") + if err != nil { + return nil, err + } + limit := positiveOr(req.Limit, 30) + query := fmt.Sprintf("SELECT %s FROM %s WHERE %s%s LIMIT %d, %d", fieldsSQL, quoteIdentifier(tableName), whereSQL, orderSQL, max(req.Offset, 0), limit) + rows, err := e.queryRows(ctx, query, args...) + if err != nil { + return nil, err + } + return &types.SearchMetadataResult{MetadataRecords: decodeRows(rows, "metadata"), Total: total}, nil +} + +// FilterDocIdsByMetaPushdown evaluates supported metadata filters in the +// legacy meta_fields JSON column. nil means the caller should fall back. +func (e *Engine) FilterDocIdsByMetaPushdown(ctx context.Context, sqlDB *gorm.DB, kbIDs []string, conditions []map[string]interface{}, logic string) []string { + if len(kbIDs) == 0 || len(conditions) == 0 || (logic != "and" && logic != "or") { + return nil + } + predicate, predicateArgs, err := buildMetaPushdownPredicate(conditions, logic) + if err != nil { + common.Debug("OceanBase metadata push-down is unsupported", zap.Error(err)) + return nil + } + tenantID, err := dao.GetTenantIDByKBID(ctx, sqlDB, kbIDs[0]) + if err != nil { + return nil + } + tableName, err := validatedMetadataTableName(tenantID) + if err != nil { + common.Debug("OceanBase metadata table name is invalid", zap.Error(err)) + return nil + } + exists, err := e.tableExists(ctx, tableName) + if err != nil || !exists { + return nil + } + kbPlaceholders := make([]string, len(kbIDs)) + args := make([]interface{}, 0, len(kbIDs)+len(predicateArgs)) + for i, kbID := range kbIDs { + kbPlaceholders[i] = "?" + args = append(args, kbID) + } + whereSQL := "kb_id IN (" + strings.Join(kbPlaceholders, ", ") + ") AND (" + predicate + ")" + args = append(args, predicateArgs...) + total, err := scanCount(e.db.QueryRowContext(ctx, "SELECT COUNT(id) FROM "+quoteIdentifier(tableName)+" WHERE "+whereSQL, args...)) + if err != nil || total > metadataPushdownMaxSize { + return nil + } + if total == 0 { + return []string{} + } + rows, err := e.queryRows(ctx, fmt.Sprintf("SELECT id FROM %s WHERE %s LIMIT %d", quoteIdentifier(tableName), whereSQL, metadataPushdownMaxSize), args...) + if err != nil { + return nil + } + ids := make([]string, 0, len(rows)) + for _, row := range rows { + if id := stringValue(row["id"]); id != "" { + ids = append(ids, id) + } + } + return ids +} + +func buildMetaPushdownPredicate(conditions []map[string]interface{}, logic string) (string, []interface{}, error) { + logic = strings.ToLower(strings.TrimSpace(logic)) + if logic != "and" && logic != "or" { + return "", nil, fmt.Errorf("unsupported metadata logic: %s", logic) + } + parts := make([]string, 0, len(conditions)) + args := make([]interface{}, 0, len(conditions)*4) + for _, condition := range conditions { + key := stringValue(condition["key"]) + op := stringValue(condition["op"]) + if key == "" || !metadataKeyPattern.MatchString(key) { + return "", nil, fmt.Errorf("invalid metadata key") + } + path := "$." + key + value := condition["value"] + expression := "JSON_EXTRACT(meta_fields, ?)" + if op == "≠" || op == "not in" { + return "", nil, fmt.Errorf("metadata operator %s is unsafe for multi-valued fields", op) + } + switch op { + case "=": + candidate, err := encodeJSONCandidate(value) + if err != nil { + return "", nil, err + } + contains := "JSON_CONTAINS(" + expression + ", ?)" + parts = append(parts, contains) + args = append(args, path, candidate) + case ">", "<", "≥", "≤": + operator := map[string]string{">": ">", "<": "<", "≥": ">=", "≤": "<="}[op] + coerced := coerceMetadataScalar(value) + if _, numeric := numberToFloat(coerced); numeric { + parts = append(parts, "CAST(JSON_UNQUOTE("+expression+") AS DECIMAL(65,20)) "+operator+" ?") + } else { + parts = append(parts, "LOWER(JSON_UNQUOTE("+expression+")) "+operator+" LOWER(?)") + } + args = append(args, path, coerced) + case "in": + values := metadataMembers(value) + if len(values) == 0 { + return "", nil, fmt.Errorf("metadata %s requires at least one value", op) + } + memberParts := make([]string, 0, len(values)) + for _, member := range values { + candidate, err := encodeJSONCandidate(member) + if err != nil { + return "", nil, err + } + memberParts = append(memberParts, "JSON_CONTAINS("+expression+", ?)") + args = append(args, path, candidate) + } + parts = append(parts, "("+strings.Join(memberParts, " OR ")+")") + case "contains", "not contains", "start with", "end with": + text := stringValue(value) + if text == "" { + return "", nil, fmt.Errorf("metadata %s requires a value", op) + } + like := "LOWER(JSON_UNQUOTE(" + expression + ")) LIKE " + switch op { + case "contains", "not contains": + like += "CONCAT('%', LOWER(?), '%')" + case "start with": + like += "CONCAT(LOWER(?), '%')" + case "end with": + like += "CONCAT('%', LOWER(?))" + } + if op == "not contains" { + like = "NOT (" + like + ")" + } + parts = append(parts, like) + args = append(args, path, text) + case "empty": + parts = append(parts, "("+expression+" IS NULL OR JSON_TYPE("+expression+") = 'NULL' OR JSON_UNQUOTE("+expression+") = '' OR JSON_LENGTH("+expression+") = 0)") + args = append(args, path, path, path, path) + case "not empty": + parts = append(parts, "NOT ("+expression+" IS NULL OR JSON_TYPE("+expression+") = 'NULL' OR JSON_UNQUOTE("+expression+") = '' OR JSON_LENGTH("+expression+") = 0)") + args = append(args, path, path, path, path) + default: + return "", nil, fmt.Errorf("unsupported metadata operator: %s", op) + } + } + if len(parts) == 0 { + return "", nil, fmt.Errorf("empty metadata predicate") + } + return strings.Join(parts, " "+strings.ToUpper(logic)+" "), args, nil +} + +func encodeJSONCandidate(value interface{}) (string, error) { + encoded, err := json.Marshal(coerceMetadataScalar(value)) + return string(encoded), err +} + +func coerceMetadataScalar(value interface{}) interface{} { + if value == nil { + return nil + } + text := strings.TrimSpace(stringValue(value)) + if integer, err := strconv.ParseInt(text, 10, 64); err == nil { + return integer + } + if number, err := strconv.ParseFloat(text, 64); err == nil { + return number + } + return text +} + +func metadataMembers(value interface{}) []interface{} { + if values, ok := interfaceSlice(value); ok { + return values + } + parts := strings.Split(stringValue(value), ",") + result := make([]interface{}, 0, len(parts)) + for _, part := range parts { + if part = strings.TrimSpace(part); part != "" { + result = append(result, part) + } + } + return result +} diff --git a/internal/engine/oceanbase/rows.go b/internal/engine/oceanbase/rows.go new file mode 100644 index 0000000000..1908e8b909 --- /dev/null +++ b/internal/engine/oceanbase/rows.go @@ -0,0 +1,232 @@ +// +// 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 oceanbase + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strings" +) + +func (e *Engine) queryRows(ctx context.Context, query string, args ...interface{}) ([]map[string]interface{}, error) { + rows, err := e.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + columns, err := rows.Columns() + if err != nil { + return nil, err + } + result := make([]map[string]interface{}, 0) + for rows.Next() { + values := make([]interface{}, len(columns)) + destinations := make([]interface{}, len(columns)) + for i := range values { + destinations[i] = &values[i] + } + if err := rows.Scan(destinations...); err != nil { + return nil, err + } + row := make(map[string]interface{}, len(columns)) + for i, column := range columns { + if values[i] == nil { + continue + } + if raw, ok := values[i].([]byte); ok { + row[column] = string(raw) + } else { + row[column] = values[i] + } + } + result = append(result, row) + } + return result, rows.Err() +} + +func decodeLogicalRow(row map[string]interface{}, kind string) map[string]interface{} { + result := make(map[string]interface{}, len(row)) + for rawColumn, value := range row { + column := rawColumn + if kind == "memory" { + if vectorColumnPattern.MatchString(column) { + column = "content_embed" + } else if mapped, ok := memoryColumnToField[column]; ok { + column = mapped + } + } + if kind == "chunk" && column == "_order_id" { + column = "chunk_order_int" + } + if kind == "memory" && column == "status" { + switch status := value.(type) { + case int64: + value = status != 0 + case int: + value = status != 0 + case string: + value = status != "" && status != "0" + } + } + storedColumn := rawColumn + if (kind == "chunk" && arrayColumns[storedColumn]) || jsonColumns[storedColumn] || vectorColumnPattern.MatchString(storedColumn) { + if text, ok := value.(string); ok { + var decoded interface{} + if err := json.Unmarshal([]byte(text), &decoded); err == nil { + value = decoded + } + } + } + result[column] = value + } + return result +} + +func scanCount(row *sql.Row) (int64, error) { + var count int64 + if err := row.Scan(&count); err != nil { + return 0, err + } + return count, nil +} + +type selectField struct { + column string + alias string +} + +func (field selectField) expression(tableAlias string) string { + prefix := "" + if tableAlias != "" { + prefix = quoteIdentifier(tableAlias) + "." + } + expression := prefix + quoteIdentifier(field.column) + if field.alias != field.column { + expression += " AS " + quoteIdentifier(field.alias) + } + return expression +} + +func parseSelectField(field, kind string) (*selectField, error) { + if field == "_score" { + return nil, nil + } + if field == "row_id()" || field == "row_id" { + return &selectField{column: identifierField(kind), alias: "row_id"}, nil + } + column := field + if kind == "memory" { + column = mapMemoryField(field) + } + if kind == "chunk" && field == "chunk_order_int" { + column = "_order_id" + } + if field == "content_embed" && kind == "memory" { + return nil, nil + } + if !vectorColumnPattern.MatchString(column) && !validColumns(kind)[column] { + return nil, fmt.Errorf("unknown %s field: %s", kind, field) + } + return &selectField{column: column, alias: field}, nil +} + +func selectExpression(field, kind string) (string, string, error) { + parsed, err := parseSelectField(field, kind) + if err != nil || parsed == nil { + return "", "", err + } + return parsed.expression(""), parsed.alias, nil +} + +func buildSelectFields(fields []string, kind string) (string, []string, error) { + return buildSelectFieldsWithAlias(fields, kind, "") +} + +func buildQualifiedSelectFields(fields []string, kind, tableAlias string) (string, []string, error) { + return buildSelectFieldsWithAlias(fields, kind, tableAlias) +} + +func buildSelectFieldsWithAlias(fields []string, kind, tableAlias string) (string, []string, error) { + if len(fields) == 0 || containsString(fields, "*") { + fields = defaultFields(kind) + } + if !containsString(fields, identifierField(kind)) { + fields = append([]string{identifierField(kind)}, fields...) + } + expressions := make([]string, 0, len(fields)) + aliases := make([]string, 0, len(fields)) + seen := make(map[string]bool) + for _, field := range fields { + parsed, err := parseSelectField(field, kind) + if err != nil { + return "", nil, err + } + if parsed == nil || seen[parsed.alias] { + continue + } + seen[parsed.alias] = true + expressions = append(expressions, parsed.expression(tableAlias)) + aliases = append(aliases, parsed.alias) + } + return strings.Join(expressions, ", "), aliases, nil +} + +func defaultFields(kind string) []string { + var columns []columnDefinition + switch kind { + case "memory": + columns = memoryColumns + case "metadata": + columns = metadataColumns + case "skill": + columns = skillColumns + default: + columns = chunkColumns + } + fields := make([]string, 0, len(columns)) + for _, column := range columns { + field := column.name + if kind == "memory" { + if mapped, ok := memoryColumnToField[field]; ok { + field = mapped + } + } + if kind == "chunk" && field == "_order_id" { + field = "chunk_order_int" + } + fields = append(fields, field) + } + return fields +} + +func identifierField(kind string) string { + if kind == "skill" { + return "skill_id" + } + return "id" +} + +func containsString(values []string, target string) bool { + for _, value := range values { + if value == target { + return true + } + } + return false +} diff --git a/internal/engine/oceanbase/schema.go b/internal/engine/oceanbase/schema.go new file mode 100644 index 0000000000..b4c94d06eb --- /dev/null +++ b/internal/engine/oceanbase/schema.go @@ -0,0 +1,531 @@ +// +// 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 oceanbase + +import ( + "context" + "crypto/md5" // #nosec G501 -- MD5 provides a deterministic identifier checksum, not security. + "database/sql" + "fmt" + "os" + "regexp" + "strconv" + "strings" + "sync" + "time" + + "ragflow/internal/engine/redis" +) + +const ( + maxIndexNameLength = 64 + indexNameHashLength = 4 + indexNameTruncationSpace = 8 +) + +type columnDefinition struct { + name string + typeSQL string +} + +var ( + identifierPattern = regexp.MustCompile(`^[A-Za-z0-9_-]+$`) + metadataKeyPattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`) + ddlLocks sync.Map +) + +var chunkColumns = []columnDefinition{ + {"id", "VARCHAR(256) NOT NULL PRIMARY KEY"}, + {"kb_id", "VARCHAR(256) NOT NULL"}, + {"doc_id", "VARCHAR(256) NULL"}, + {"docnm_kwd", "VARCHAR(256) NULL"}, + {"doc_type_kwd", "VARCHAR(256) NULL"}, + {"title_tks", "VARCHAR(256) NULL"}, + {"title_sm_tks", "VARCHAR(256) NULL"}, + {"content_with_weight", "LONGTEXT NULL"}, + {"content_ltks", "LONGTEXT NULL"}, + {"content_sm_ltks", "LONGTEXT NULL"}, + {"pagerank_fea", "INTEGER NULL"}, + {"important_kwd", "ARRAY(VARCHAR(256)) NULL"}, + {"important_tks", "TEXT NULL"}, + {"question_kwd", "ARRAY(VARCHAR(1024)) NULL"}, + {"question_tks", "TEXT NULL"}, + {"tag_kwd", "ARRAY(VARCHAR(256)) NULL"}, + {"tag_feas", "JSON NULL"}, + {"available_int", "INTEGER NOT NULL DEFAULT 1"}, + {"create_time", "VARCHAR(19) NULL"}, + {"create_timestamp_flt", "DOUBLE NULL"}, + {"img_id", "VARCHAR(128) NULL"}, + {"position_int", "ARRAY(ARRAY(INTEGER)) NULL"}, + {"page_num_int", "ARRAY(INTEGER) NULL"}, + {"top_int", "ARRAY(INTEGER) NULL"}, + {"knowledge_graph_kwd", "VARCHAR(256) NULL"}, + {"source_id", "ARRAY(VARCHAR(256)) NULL"}, + {"entity_kwd", "VARCHAR(256) NULL"}, + {"entity_type_kwd", "VARCHAR(256) NULL"}, + {"from_entity_kwd", "VARCHAR(256) NULL"}, + {"to_entity_kwd", "VARCHAR(256) NULL"}, + {"weight_int", "INTEGER NULL"}, + {"weight_flt", "DOUBLE NULL"}, + {"entities_kwd", "ARRAY(VARCHAR(256)) NULL"}, + {"rank_flt", "DOUBLE NULL"}, + {"n_hop_with_weight", "LONGTEXT NULL"}, + {"removed_kwd", "VARCHAR(256) NULL DEFAULT 'N'"}, + {"raptor_kwd", "VARCHAR(256) NULL"}, + {"raptor_layer_int", "INTEGER NULL"}, + {"chunk_data", "JSON NULL"}, + {"metadata", "JSON NULL"}, + {"extra", "JSON NULL"}, + {"_order_id", "INTEGER NULL"}, + {"group_id", "VARCHAR(256) NULL"}, + {"mom_id", "VARCHAR(256) NULL"}, +} + +var chunkExtraColumns = selectColumnDefinitions(chunkColumns, + "_order_id", + "group_id", + "mom_id", + "chunk_data", + "raptor_kwd", + "raptor_layer_int", + "n_hop_with_weight", +) + +var memoryColumns = []columnDefinition{ + {"id", "VARCHAR(256) NOT NULL PRIMARY KEY"}, + {"message_id", "VARCHAR(256) NOT NULL"}, + {"message_type_kwd", "VARCHAR(64) NULL"}, + {"source_id", "VARCHAR(256) NULL"}, + {"memory_id", "VARCHAR(256) NOT NULL"}, + {"user_id", "VARCHAR(256) NULL"}, + {"agent_id", "VARCHAR(256) NULL"}, + {"session_id", "VARCHAR(256) NULL"}, + {"zone_id", "INTEGER NULL DEFAULT 0"}, + {"valid_at", "VARCHAR(64) NULL"}, + {"invalid_at", "VARCHAR(64) NULL"}, + {"forget_at", "VARCHAR(64) NULL"}, + {"status_int", "INTEGER NOT NULL DEFAULT 1"}, + {"content_ltks", "LONGTEXT NULL"}, + {"tokenized_content_ltks", "LONGTEXT NULL"}, +} + +var metadataColumns = []columnDefinition{ + {"id", "VARCHAR(256) NOT NULL PRIMARY KEY"}, + {"kb_id", "VARCHAR(256) NOT NULL"}, + {"meta_fields", "JSON NULL"}, +} + +var skillColumns = []columnDefinition{ + {"skill_id", "VARCHAR(256) NOT NULL PRIMARY KEY"}, + {"space_id", "VARCHAR(256) NULL"}, + {"folder_id", "VARCHAR(256) NULL"}, + {"name", "LONGTEXT NULL"}, + {"name_tks", "LONGTEXT NULL"}, + {"tags", "LONGTEXT NULL"}, + {"tags_tks", "LONGTEXT NULL"}, + {"description", "LONGTEXT NULL"}, + {"description_tks", "LONGTEXT NULL"}, + {"content", "LONGTEXT NULL"}, + {"content_tks", "LONGTEXT NULL"}, + {"version", "VARCHAR(64) NULL"}, + {"status", "VARCHAR(64) NULL"}, + {"create_time", "BIGINT NULL DEFAULT 0"}, + {"update_time", "BIGINT NULL DEFAULT 0"}, +} + +var chunkIndexColumns = []string{ + "kb_id", "doc_id", "available_int", "knowledge_graph_kwd", "entity_type_kwd", "removed_kwd", +} + +var memoryIndexColumns = []string{"message_id", "memory_id", "status_int"} + +var originalFullTextFields = []string{"docnm_kwd", "content_with_weight", "important_tks", "question_tks"} +var tokenizedFullTextFields = []string{"title_tks", "title_sm_tks", "important_tks", "question_tks", "content_ltks", "content_sm_ltks"} + +// CreateChunkStore creates or upgrades the legacy shared tenant table. The +// dataset ID is a row-level discriminator for chunk and memory tables. +func (e *Engine) CreateChunkStore(ctx context.Context, baseName, datasetID string, vectorSize int, parserID string) error { + if err := validateIdentifier(baseName); err != nil { + return err + } + + switch { + case strings.HasPrefix(baseName, "skill_") || datasetID == "skill": + if err := e.ensureTable(ctx, baseName, skillColumns, "ob_"); err != nil { + return err + } + for _, field := range []string{"name_tks", "tags_tks", "description_tks", "content_tks"} { + if err := e.ensureFullTextIndex(ctx, baseName, field, "ob_"); err != nil { + return err + } + } + case strings.HasPrefix(baseName, "memory_"): + if err := e.ensureTable(ctx, baseName, memoryColumns, "ob_memory_"); err != nil { + return err + } + for _, field := range memoryIndexColumns { + if err := e.ensureRegularIndex(ctx, baseName, field, "ob_memory_"); err != nil { + return err + } + } + if e.flags.enableFullTextSearch { + for _, field := range []string{"content_ltks", "tokenized_content_ltks"} { + if err := e.ensureFullTextIndex(ctx, baseName, field, "ob_memory_"); err != nil { + return err + } + } + } + default: + if err := e.ensureTable(ctx, baseName, chunkColumns, "ob_"); err != nil { + return err + } + for _, column := range chunkExtraColumns { + if err := e.ensureColumn(ctx, baseName, column, "ob_"); err != nil { + return err + } + } + for _, field := range chunkIndexColumns { + if err := e.ensureRegularIndex(ctx, baseName, field, "ob_"); err != nil { + return err + } + } + if e.flags.enableFullTextSearch { + fields := tokenizedFullTextFields + if e.flags.searchOriginalContent { + fields = originalFullTextFields + } + for _, field := range fields { + if err := e.ensureFullTextIndex(ctx, baseName, field, "ob_"); err != nil { + return err + } + } + } + } + return e.ensureVectorColumnAndIndex(ctx, baseName, vectorSize, lockPrefix(baseName)) +} + +func (e *Engine) ensureTable(ctx context.Context, tableName string, columns []columnDefinition, prefix string) error { + return e.ensureTableWithLock(ctx, tableName, columns, prefix+"create_table_"+tableName) +} + +func (e *Engine) ensureTableWithLock(ctx context.Context, tableName string, columns []columnDefinition, lockName string) error { + return e.withDDLLock(ctx, lockName, func() (bool, error) { + return e.tableExists(ctx, tableName) + }, func() error { + definitions := make([]string, 0, len(columns)) + for _, column := range columns { + definitions = append(definitions, quoteIdentifier(column.name)+" "+column.typeSQL) + } + query := fmt.Sprintf("CREATE TABLE %s (%s) DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci ORGANIZATION=heap", + quoteIdentifier(tableName), strings.Join(definitions, ", ")) + _, err := e.db.ExecContext(ctx, query) + return err + }) +} + +func (e *Engine) ensureColumn(ctx context.Context, tableName string, column columnDefinition, prefix string) error { + return e.withDDLLock(ctx, prefix+"add_"+column.name+"_"+tableName, func() (bool, error) { + return e.columnExists(ctx, tableName, column.name) + }, func() error { + _, err := e.db.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", + quoteIdentifier(tableName), quoteIdentifier(column.name), column.typeSQL)) + return err + }) +} + +func (e *Engine) ensureRegularIndex(ctx context.Context, tableName, columnName, prefix string) error { + indexName := regularIndexName(tableName, columnName) + return e.withDDLLock(ctx, prefix+"add_idx_"+tableName+"_"+columnName, func() (bool, error) { + return e.indexExists(ctx, tableName, indexName) + }, func() error { + _, err := e.db.ExecContext(ctx, fmt.Sprintf("CREATE INDEX %s ON %s (%s)", + quoteIdentifier(indexName), quoteIdentifier(tableName), quoteIdentifier(columnName))) + return err + }) +} + +func regularIndexName(tableName, columnName string) string { + indexName := fmt.Sprintf("ix_%s_%s", tableName, columnName) + if len(indexName) <= maxIndexNameLength { + return indexName + } + digest := fmt.Sprintf("%x", md5.Sum([]byte(indexName))) // #nosec G401 -- This is a non-security identifier checksum. + suffix := "_" + digest[len(digest)-indexNameHashLength:] + return indexName[:maxIndexNameLength-indexNameTruncationSpace] + suffix +} + +func (e *Engine) ensureFullTextIndex(ctx context.Context, tableName, columnName, prefix string) error { + indexName := "fts_idx_" + columnName + return e.withDDLLock(ctx, prefix+"add_fulltext_idx_"+tableName+"_"+columnName, func() (bool, error) { + return e.indexExists(ctx, tableName, indexName) + }, func() error { + _, err := e.db.ExecContext(ctx, fmt.Sprintf("CREATE FULLTEXT INDEX %s ON %s (%s) WITH PARSER IK", + quoteIdentifier(indexName), quoteIdentifier(tableName), quoteIdentifier(columnName))) + return err + }) +} + +func (e *Engine) ensureVectorColumnAndIndex(ctx context.Context, tableName string, vectorSize int, prefix string) error { + if vectorSize <= 0 { + return nil + } + columnName := fmt.Sprintf("q_%d_vec", vectorSize) + if err := e.withDDLLock(ctx, prefix+"add_vector_column_"+tableName+"_"+columnName, func() (bool, error) { + return e.columnExists(ctx, tableName, columnName) + }, func() error { + _, err := e.db.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s VECTOR(%d) NULL", + quoteIdentifier(tableName), quoteIdentifier(columnName), vectorSize)) + return err + }); err != nil { + return err + } + indexName := columnName + "_idx" + return e.withDDLLock(ctx, prefix+"add_vector_idx_"+tableName+"_"+columnName, func() (bool, error) { + return e.indexExists(ctx, tableName, indexName) + }, func() error { + _, err := e.db.ExecContext(ctx, fmt.Sprintf("CREATE VECTOR INDEX %s ON %s (%s) WITH (distance=cosine, type=hnsw, lib=vsag)", + quoteIdentifier(indexName), quoteIdentifier(tableName), quoteIdentifier(columnName))) + return err + }) +} + +func (e *Engine) withDDLLock(ctx context.Context, lockName string, check func() (bool, error), action func() error) error { + value, _ := ddlLocks.LoadOrStore(lockName, &sync.Mutex{}) + lock := value.(*sync.Mutex) + lock.Lock() + defer lock.Unlock() + + exists, err := check() + if err != nil { + return err + } + if exists { + return nil + } + + timeoutSeconds := int64(60) + if raw := strings.TrimSpace(os.Getenv("OB_DDL_TIMEOUT")); raw != "" { + if parsed, parseErr := strconv.ParseInt(raw, 10, 64); parseErr == nil && parsed > 0 { + timeoutSeconds = parsed + } + } + timeout := time.Duration(timeoutSeconds) * time.Second + distributed := redis.NewDistributedLock(lockName, "", timeout, timeout) + if distributed != nil && !distributed.Acquire(ctx) { + deadline := time.NewTimer(timeout) + defer deadline.Stop() + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + waitForLock: + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-deadline.C: + return fmt.Errorf("timeout waiting for DDL %s", lockName) + case <-ticker.C: + exists, err = check() + if err != nil { + return err + } + if exists { + return nil + } + if distributed.Acquire(ctx) { + break waitForLock + } + } + } + } + if distributed != nil { + defer distributed.Release(ctx) + exists, err = check() + if err != nil { + return err + } + if exists { + return nil + } + } + if err := action(); err != nil && !isDuplicateDDLError(err) { + return fmt.Errorf("DDL %s: %w", lockName, err) + } + exists, err = check() + if err != nil { + return err + } + if !exists { + return fmt.Errorf("DDL %s completed without creating the requested object", lockName) + } + return nil +} + +func (e *Engine) tableExists(ctx context.Context, tableName string) (bool, error) { + var count int + err := e.db.QueryRowContext(ctx, + "SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?", + e.dbName, tableName).Scan(&count) + return count > 0, err +} + +func (e *Engine) columnExists(ctx context.Context, tableName, columnName string) (bool, error) { + var count int + err := e.db.QueryRowContext(ctx, + "SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME = ?", + e.dbName, tableName, columnName).Scan(&count) + return count > 0, err +} + +func (e *Engine) indexExists(ctx context.Context, tableName, indexName string) (bool, error) { + var count int + err := e.db.QueryRowContext(ctx, + "SELECT COUNT(*) FROM INFORMATION_SCHEMA.STATISTICS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND INDEX_NAME = ?", + e.dbName, tableName, indexName).Scan(&count) + return count > 0, err +} + +func (e *Engine) findVectorColumn(ctx context.Context, tableName, expectedColumn string) (string, error) { + query := "SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME REGEXP '^q_[0-9]+_vec$'" + args := []interface{}{e.dbName, tableName} + if expectedColumn != "" { + if !vectorColumnPattern.MatchString(expectedColumn) { + return "", fmt.Errorf("invalid vector column: %s", expectedColumn) + } + query += " AND COLUMN_NAME = ?" + args = append(args, expectedColumn) + } + query += " ORDER BY COLUMN_NAME LIMIT 1" + var columnName string + err := e.db.QueryRowContext(ctx, query, args...).Scan(&columnName) + if err == sql.ErrNoRows { + return "", nil + } + return columnName, err +} + +// ChunkStoreExists checks the shared physical table. CreateChunkStore performs +// the additive compatibility upgrade when a caller needs a new vector size. +func (e *Engine) ChunkStoreExists(ctx context.Context, baseName, datasetID string) (bool, error) { + if err := validateIdentifier(baseName); err != nil { + return false, err + } + exists, err := e.tableExists(ctx, baseName) + if err != nil || !exists { + return exists, err + } + kind := tableKind(baseName, datasetID) + var indexColumns, fullTextColumns []string + switch kind { + case "memory": + indexColumns = memoryIndexColumns + if e.flags.enableFullTextSearch { + fullTextColumns = []string{"content_ltks", "tokenized_content_ltks"} + } + case "skill": + fullTextColumns = []string{"name_tks", "tags_tks", "description_tks", "content_tks"} + default: + indexColumns = chunkIndexColumns + if e.flags.enableFullTextSearch { + fullTextColumns = tokenizedFullTextFields + if e.flags.searchOriginalContent { + fullTextColumns = originalFullTextFields + } + } + for _, column := range chunkExtraColumns { + exists, err = e.columnExists(ctx, baseName, column.name) + if err != nil || !exists { + return exists, err + } + } + } + for _, column := range indexColumns { + exists, err = e.indexExists(ctx, baseName, regularIndexName(baseName, column)) + if err != nil || !exists { + return exists, err + } + } + for _, column := range fullTextColumns { + exists, err = e.indexExists(ctx, baseName, "fts_idx_"+column) + if err != nil || !exists { + return exists, err + } + } + return true, nil +} + +// DropChunkStore keeps shared chunk and memory tables alive when only one +// dataset is removed. Skill tables and explicitly unscoped calls are dropped. +func (e *Engine) DropChunkStore(ctx context.Context, baseName, datasetID string) error { + if err := validateIdentifier(baseName); err != nil { + return err + } + if datasetID != "" && datasetID != "skill" { + exists, err := e.tableExists(ctx, baseName) + if err != nil { + return err + } + if !exists { + return nil + } + field := "kb_id" + if strings.HasPrefix(baseName, "memory_") { + field = "memory_id" + } + _, err = e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE %s = ?", + quoteIdentifier(baseName), quoteIdentifier(field)), datasetID) + return err + } + _, err := e.db.ExecContext(ctx, "DROP TABLE IF EXISTS "+quoteIdentifier(baseName)) + return err +} + +func lockPrefix(tableName string) string { + if strings.HasPrefix(tableName, "memory_") { + return "ob_memory_" + } + return "ob_" +} + +func selectColumnDefinitions(columns []columnDefinition, names ...string) []columnDefinition { + byName := make(map[string]columnDefinition, len(columns)) + for _, column := range columns { + byName[column.name] = column + } + selected := make([]columnDefinition, 0, len(names)) + for _, name := range names { + column, ok := byName[name] + if !ok { + panic("missing column definition: " + name) + } + selected = append(selected, column) + } + return selected +} + +func quoteIdentifier(identifier string) string { return "`" + identifier + "`" } + +func validateIdentifier(identifier string) error { + if identifier == "" || !identifierPattern.MatchString(identifier) { + return fmt.Errorf("invalid SQL identifier: %q", identifier) + } + return nil +} + +func isDuplicateDDLError(err error) bool { + message := strings.ToLower(err.Error()) + return strings.Contains(message, "duplicate") || strings.Contains(message, "already exists") +} diff --git a/internal/engine/oceanbase/schema_test.go b/internal/engine/oceanbase/schema_test.go new file mode 100644 index 0000000000..437c94baec --- /dev/null +++ b/internal/engine/oceanbase/schema_test.go @@ -0,0 +1,162 @@ +// +// 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 oceanbase + +import ( + "context" + "regexp" + "strings" + "testing" + + "github.com/DATA-DOG/go-sqlmock" +) + +func TestDropChunkStoreIgnoresMissingSharedTable(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("seekdb", "legacy_doc", db) + + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?")). + WithArgs("legacy_doc", "memory_tenant_1"). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(0)) + + if err := engine.DropChunkStore(context.Background(), "memory_tenant_1", "memory_1"); err != nil { + t.Fatalf("DropChunkStore() error = %v", err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestDropChunkStoreDeletesOnlyScopedRows(t *testing.T) { + tests := []struct { + name string + tableName string + datasetID string + fieldName string + }{ + {name: "memory", tableName: "memory_tenant_1", datasetID: "memory_1", fieldName: "memory_id"}, + {name: "chunk", tableName: "ragflow_tenant_1", datasetID: "kb_1", fieldName: "kb_id"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("oceanbase", "legacy_doc", db) + + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?")). + WithArgs("legacy_doc", test.tableName). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + deleteSQL := "DELETE FROM `" + test.tableName + "` WHERE `" + test.fieldName + "` = ?" + mock.ExpectExec(regexp.QuoteMeta(deleteSQL)). + WithArgs(test.datasetID). + WillReturnResult(sqlmock.NewResult(0, 0)) + + if err := engine.DropChunkStore(context.Background(), test.tableName, test.datasetID); err != nil { + t.Fatalf("DropChunkStore() error = %v", err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } + }) + } +} + +func TestDropChunkStoreDropsTableForUnscopedDataset(t *testing.T) { + tests := []struct { + name string + datasetID string + }{ + {name: "empty dataset", datasetID: ""}, + {name: "skill dataset", datasetID: "skill"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("oceanbase", "legacy_doc", db) + + mock.ExpectExec(regexp.QuoteMeta("DROP TABLE IF EXISTS `ragflow_tenant_1`")). + WillReturnResult(sqlmock.NewResult(0, 0)) + + if err := engine.DropChunkStore(context.Background(), "ragflow_tenant_1", test.datasetID); err != nil { + t.Fatalf("DropChunkStore() error = %v", err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } + }) + } +} + +func TestRegularIndexNamePreservesShortNamesAndCapsLongNames(t *testing.T) { + if got, want := regularIndexName("memory_tenant_1", "memory_id"), "ix_memory_tenant_1_memory_id"; got != want { + t.Fatalf("regularIndexName() = %q, want %q", got, want) + } + + tableName := strings.Repeat("tenant", 10) + first := regularIndexName(tableName, "kb_id") + second := regularIndexName(tableName, "doc_id") + if len(first) > maxIndexNameLength { + t.Fatalf("index name length = %d, want <= %d: %q", len(first), maxIndexNameLength, first) + } + if first == second { + t.Fatalf("different index inputs produced the same name: %q", first) + } + if !identifierPattern.MatchString(first) || first != regularIndexName(tableName, "kb_id") { + t.Fatalf("index name is invalid or nondeterministic: %q", first) + } + + pythonTableName := "ragflow_12345678-1234-1234-1234-123456789012" + if got, want := regularIndexName(pythonTableName, "create_timestamp_flt"), "ix_ragflow_12345678-1234-1234-1234-123456789012_create_t_69b6"; got != want { + t.Fatalf("SQLAlchemy-compatible index name = %q, want %q", got, want) + } +} + +func TestFindVectorColumnUsesExpectedColumn(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("seekdb", "legacy_doc", db) + + query := "SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND COLUMN_NAME REGEXP '^q_[0-9]+_vec$' AND COLUMN_NAME = ? ORDER BY COLUMN_NAME LIMIT 1" + mock.ExpectQuery(regexp.QuoteMeta(query)). + WithArgs("legacy_doc", "memory_tenant_1", "q_1024_vec"). + WillReturnRows(sqlmock.NewRows([]string{"COLUMN_NAME"}).AddRow("q_1024_vec")) + + column, err := engine.findVectorColumn(context.Background(), "memory_tenant_1", "q_1024_vec") + if err != nil { + t.Fatal(err) + } + if column != "q_1024_vec" { + t.Fatalf("findVectorColumn() = %q, want q_1024_vec", column) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} diff --git a/internal/engine/oceanbase/search.go b/internal/engine/oceanbase/search.go new file mode 100644 index 0000000000..939bcdcc5e --- /dev/null +++ b/internal/engine/oceanbase/search.go @@ -0,0 +1,840 @@ +// +// 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 oceanbase + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strconv" + "strings" + + "ragflow/internal/common" + "ragflow/internal/engine/types" + + "go.uber.org/zap" +) + +type searchPlan struct { + text *types.MatchTextExpr + dense *types.MatchDenseExpr + fusion *types.FusionExpr +} + +// Search executes filter, full-text, vector, or fusion search. When explicitly +// enabled and supported, the exact text+dense+fusion form first uses +// DBMS_HYBRID_SEARCH.SEARCH and falls back once to SQL only for recognized +// feature/package availability errors. +func (e *Engine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + if req == nil || len(req.IndexNames) == 0 { + return nil, fmt.Errorf("index names cannot be empty") + } + types.LogSearchRequest("OceanBase", req) + plan := parseSearchPlan(req.MatchExprs) + if !e.flags.enableFullTextSearch && plan.text != nil && plan.dense != nil { + plan.text = nil + plan.fusion = nil + } + if plan.fusion != nil { + weight := fusionVectorWeight(plan.fusion) + if weight <= 0 { + plan.dense, plan.fusion = nil, nil + } else if weight >= 1 { + plan.text, plan.fusion = nil, nil + } + } + + tableNames := uniqueStrings(req.IndexNames) + mergeTables := len(tableNames) > 1 + candidateLimit := globalSearchCandidateLimit(req) + effectivePlan := plan + if mergeTables { + effectivePlan = expandSearchPlan(plan, candidateLimit) + } + hiddenSortFields := searchHiddenSortFields(req, mergeTables) + + result := &types.SearchResult{Chunks: []map[string]interface{}{}} + for _, tableName := range tableNames { + if err := validateIdentifier(tableName); err != nil { + return nil, err + } + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return nil, err + } + if !exists { + continue + } + kind := tableKind(tableName, req.KbIDs...) + effectiveReq := *req + effectiveReq.SelectFields = append([]string(nil), req.SelectFields...) + if mergeTables { + effectiveReq.Offset = 0 + effectiveReq.Limit = candidateLimit + for _, field := range hiddenSortFields { + effectiveReq.SelectFields = append(effectiveReq.SelectFields, field) + } + } + if kind == "memory" && containsString(effectiveReq.SelectFields, "content_embed") { + expectedVectorColumn := "" + if plan.dense != nil { + expectedVectorColumn = plan.dense.VectorColumnName + } + vectorColumn, vectorErr := e.findVectorColumn(ctx, tableName, expectedVectorColumn) + if vectorErr != nil { + return nil, vectorErr + } + fields := make([]string, 0, len(effectiveReq.SelectFields)) + for _, field := range effectiveReq.SelectFields { + if field == "content_embed" { + if vectorColumn != "" { + fields = append(fields, vectorColumn) + } + continue + } + fields = append(fields, field) + } + effectiveReq.SelectFields = fields + } + condition := copyMap(req.Filter) + if kind == "memory" { + if len(req.KbIDs) > 0 { + condition["memory_id"] = req.KbIDs + } + if _, present := condition["must_not"]; !present { + condition["must_not"] = map[string]interface{}{"exists": "forget_at"} + } + } else if kind == "chunk" && len(req.KbIDs) > 0 { + condition["kb_id"] = req.KbIDs + } + + if e.hybridAvailable.Load() && isDBMSHybridPlan(plan) { + chunks, used, hybridErr := e.searchWithDBMS(ctx, tableName, kind, condition, &effectiveReq, effectivePlan) + if hybridErr != nil { + if !isHybridUnavailableError(hybridErr) { + return nil, hybridErr + } + e.hybridAvailable.Store(false) + common.Warn("DBMS hybrid search unavailable; using SQL search", zap.Error(hybridErr)) + } else if used { + result.Chunks = append(result.Chunks, chunks...) + result.Total += int64(len(chunks)) + continue + } + } + + chunks, total, err := e.searchTableWithSQL(ctx, tableName, kind, condition, &effectiveReq, effectivePlan) + if err != nil { + return nil, err + } + result.Chunks = append(result.Chunks, chunks...) + result.Total += total + } + if result.Total == 0 { + result.Total = int64(len(result.Chunks)) + } + if mergeTables { + result.Chunks = mergeSearchChunks(result.Chunks, req, plan) + for _, chunk := range result.Chunks { + for _, field := range hiddenSortFields { + delete(chunk, field) + } + } + } + return result, nil +} + +func (e *Engine) searchTableWithSQL(ctx context.Context, tableName, kind string, condition map[string]interface{}, req *types.SearchRequest, plan searchPlan) ([]map[string]interface{}, int64, error) { + fieldsSQL, _, err := buildSelectFields(req.SelectFields, kind) + if err != nil { + return nil, 0, err + } + filterSQL, filterArgs, err := buildFilter(condition, kind) + if err != nil { + return nil, 0, err + } + offset := max(req.Offset, 0) + limit := req.Limit + if limit <= 0 { + limit = 30 + } + + switch { + case plan.text != nil && plan.dense != nil: + qualifiedFieldsSQL, _, err := buildQualifiedSelectFields(req.SelectFields, kind, "t") + if err != nil { + return nil, 0, err + } + return e.searchFusionSQL(ctx, tableName, kind, fieldsSQL, qualifiedFieldsSQL, filterSQL, filterArgs, offset, limit, plan) + case plan.dense != nil: + return e.searchVectorSQL(ctx, tableName, kind, fieldsSQL, filterSQL, filterArgs, offset, limit, plan.dense) + case plan.text != nil: + return e.searchFullTextSQL(ctx, tableName, kind, fieldsSQL, filterSQL, filterArgs, offset, limit, plan.text) + default: + count, err := scanCount(e.db.QueryRowContext(ctx, "SELECT COUNT("+quoteIdentifier(identifierField(kind))+") FROM "+quoteIdentifier(tableName)+" WHERE "+filterSQL, filterArgs...)) + if err != nil || count == 0 { + return []map[string]interface{}{}, count, err + } + orderSQL, err := buildOrderBy(req.OrderBy, kind) + if err != nil { + return nil, 0, err + } + query := fmt.Sprintf("SELECT %s FROM %s WHERE %s%s LIMIT %d, %d", fieldsSQL, quoteIdentifier(tableName), filterSQL, orderSQL, offset, limit) + rows, err := e.queryRows(ctx, query, filterArgs...) + return decodeRows(rows, kind), count, err + } +} + +func (e *Engine) searchFullTextSQL(ctx context.Context, tableName, kind, fieldsSQL, filterSQL string, filterArgs []interface{}, offset, limit int, text *types.MatchTextExpr) ([]map[string]interface{}, int64, error) { + filterExpr, filterTextArgs, scoreExpr, scoreArgs := e.fullTextExpressions(kind, text) + hint := e.fullTextHint(tableName, kind) + countQuery := fmt.Sprintf("SELECT %sCOUNT(%s) FROM %s WHERE %s AND %s", hint, quoteIdentifier(identifierField(kind)), quoteIdentifier(tableName), filterSQL, filterExpr) + countArgs := appendCopy(filterArgs, filterTextArgs...) + count, err := scanCount(e.db.QueryRowContext(ctx, countQuery, countArgs...)) + if err != nil || count == 0 { + return []map[string]interface{}{}, count, err + } + query := fmt.Sprintf("SELECT %s%s, %s AS _score FROM %s WHERE %s AND %s ORDER BY _score DESC LIMIT %d, %d", + hint, fieldsSQL, scoreExpr, quoteIdentifier(tableName), filterSQL, filterExpr, offset, minPositive(limit, text.TopN)) + args := appendCopy(scoreArgs, filterArgs...) + args = append(args, filterTextArgs...) + rows, err := e.queryRows(ctx, query, args...) + return decodeRows(rows, kind), count, err +} + +func (e *Engine) searchVectorSQL(ctx context.Context, tableName, kind, fieldsSQL, filterSQL string, filterArgs []interface{}, offset, limit int, dense *types.MatchDenseExpr) ([]map[string]interface{}, int64, error) { + if err := validateVectorExpr(dense); err != nil { + return nil, 0, err + } + vector, err := encodeVector(dense.EmbeddingData) + if err != nil { + return nil, 0, err + } + threshold := denseSimilarity(dense) + column := quoteIdentifier(dense.VectorColumnName) + scoreExpr := "(1 - COSINE_DISTANCE(" + column + ", ?))" + countQuery := fmt.Sprintf("SELECT COUNT(%s) FROM %s WHERE %s AND %s >= ?", quoteIdentifier(identifierField(kind)), quoteIdentifier(tableName), filterSQL, scoreExpr) + countArgs := appendCopy(filterArgs, vector, threshold) + count, err := scanCount(e.db.QueryRowContext(ctx, countQuery, countArgs...)) + if err != nil || count == 0 { + return []map[string]interface{}{}, count, err + } + query := fmt.Sprintf("SELECT %s, %s AS _score FROM %s WHERE %s AND %s >= ? ORDER BY COSINE_DISTANCE(%s, ?) APPROXIMATE LIMIT %d OFFSET %d", + fieldsSQL, scoreExpr, quoteIdentifier(tableName), filterSQL, scoreExpr, column, minPositive(limit, dense.TopN), offset) + args := []interface{}{vector} + args = append(args, filterArgs...) + args = append(args, vector, threshold, vector) + rows, err := e.queryRows(ctx, query, args...) + return decodeRows(rows, kind), count, err +} + +func (e *Engine) searchFusionSQL(ctx context.Context, tableName, kind, fieldsSQL, qualifiedFieldsSQL, filterSQL string, filterArgs []interface{}, offset, limit int, plan searchPlan) ([]map[string]interface{}, int64, error) { + if err := validateVectorExpr(plan.dense); err != nil { + return nil, 0, err + } + vector, err := encodeVector(plan.dense.EmbeddingData) + if err != nil { + return nil, 0, err + } + textFilter, textFilterArgs, textScore, textScoreArgs := e.fullTextExpressions(kind, plan.text) + threshold := denseSimilarity(plan.dense) + vectorWeight := fusionVectorWeight(plan.fusion) + textWeight := 1 - vectorWeight + vectorColumn := quoteIdentifier(plan.dense.VectorColumnName) + vectorScore := "(1 - COSINE_DISTANCE(" + vectorColumn + ", ?))" + candidates := positiveOr(plan.text.TopN, limit) + positiveOr(plan.dense.TopN, limit) + hint := e.fullTextHint(tableName, kind) + + if !e.flags.useFullTextFirstFusionSearch { + return e.searchSymmetricFusionSQL(ctx, tableName, kind, qualifiedFieldsSQL, filterSQL, filterArgs, offset, limit, candidates, hint, + textFilter, textFilterArgs, textScore, textScoreArgs, vector, vectorColumn, vectorScore, threshold, textWeight, vectorWeight, plan) + } + + cte := fmt.Sprintf("WITH fulltext_results AS (SELECT %s*, %s AS relevance FROM %s WHERE %s AND %s ORDER BY relevance DESC LIMIT %d)", + hint, textScore, quoteIdentifier(tableName), filterSQL, textFilter, candidates) + countQuery := cte + " SELECT COUNT(*) FROM fulltext_results WHERE " + vectorScore + " >= ?" + countArgs := appendCopy(textScoreArgs, filterArgs...) + countArgs = append(countArgs, textFilterArgs...) + countArgs = append(countArgs, vector, threshold) + count, err := scanCount(e.db.QueryRowContext(ctx, countQuery, countArgs...)) + if err != nil || count == 0 { + return []map[string]interface{}{}, count, err + } + score := fmt.Sprintf("(relevance * %s + %s * %s)", + formatFloat(textWeight), vectorScore, formatFloat(vectorWeight)) + if kind == "chunk" { + score = strings.TrimSuffix(score, ")") + " + (CAST(IFNULL(pagerank_fea, 0) AS DECIMAL(10, 2)) / 100))" + } + query := fmt.Sprintf("%s SELECT %s, %s AS _score FROM fulltext_results WHERE %s >= ? ORDER BY _score DESC LIMIT %d, %d", + cte, fieldsSQL, score, vectorScore, offset, limit) + args := appendCopy(textScoreArgs, filterArgs...) + args = append(args, textFilterArgs...) + args = append(args, vector, vector, threshold) + rows, err := e.queryRows(ctx, query, args...) + return decodeRows(rows, kind), count, err +} + +func (e *Engine) searchSymmetricFusionSQL(ctx context.Context, tableName, kind, fieldsSQL, filterSQL string, filterArgs []interface{}, offset, limit, candidates int, hint, textFilter string, textFilterArgs []interface{}, textScore string, textScoreArgs []interface{}, vector, vectorColumn, vectorScore string, threshold, textWeight, vectorWeight float64, plan searchPlan) ([]map[string]interface{}, int64, error) { + fullTextLimit := positiveOr(plan.text.TopN, candidates) + vectorLimit := positiveOr(plan.dense.TopN, candidates) + pagerankColumn := "" + if kind == "chunk" { + pagerankColumn = ", pagerank_fea" + } + identifier := quoteIdentifier(identifierField(kind)) + cte := fmt.Sprintf("WITH fulltext_results AS (SELECT %s%s AS id%s, %s AS relevance FROM %s WHERE %s AND %s ORDER BY relevance DESC LIMIT %d), "+ + "vector_results AS (SELECT %s AS id%s, %s AS similarity FROM %s WHERE %s AND %s >= ? ORDER BY COSINE_DISTANCE(%s, ?) APPROXIMATE LIMIT %d)", + hint, identifier, pagerankColumn, textScore, quoteIdentifier(tableName), filterSQL, textFilter, fullTextLimit, + identifier, pagerankColumn, vectorScore, quoteIdentifier(tableName), filterSQL, vectorScore, vectorColumn, vectorLimit) + join := " FROM fulltext_results f FULL OUTER JOIN vector_results v ON f.id = v.id" + countArgs := appendCopy(textScoreArgs, filterArgs...) + countArgs = append(countArgs, textFilterArgs...) + countArgs = append(countArgs, vector) + countArgs = append(countArgs, filterArgs...) + countArgs = append(countArgs, vector, threshold, vector) + count, err := scanCount(e.db.QueryRowContext(ctx, cte+" SELECT COUNT(*)"+join, countArgs...)) + if err != nil || count == 0 { + return []map[string]interface{}{}, count, err + } + score := fmt.Sprintf("(IFNULL(f.relevance, 0) * %s + IFNULL(v.similarity, 0) * %s)", + formatFloat(textWeight), formatFloat(vectorWeight)) + if kind == "chunk" { + score = strings.TrimSuffix(score, ")") + " + (CAST(IFNULL(f.pagerank_fea, 0) AS DECIMAL(10, 2)) / 100))" + } + query := cte + fmt.Sprintf(" SELECT %s, %s AS _score FROM (SELECT COALESCE(f.id, v.id) AS id, %s AS score%s) c JOIN %s t ON c.id = t.%s ORDER BY c.score DESC LIMIT %d, %d", + fieldsSQL, "c.score", score, join, quoteIdentifier(tableName), identifier, offset, limit) + rows, err := e.queryRows(ctx, query, countArgs...) + return decodeRows(rows, kind), count, err +} + +func (e *Engine) fullTextExpressions(kind string, text *types.MatchTextExpr) (string, []interface{}, string, []interface{}) { + query := text.MatchingText + if text.ExtraOptions != nil { + if original := stringValue(text.ExtraOptions["original_query"]); original != "" { + query = strings.TrimSpace(original) + } + } + fields, weights := e.fullTextFields(kind, text) + filterParts := make([]string, len(fields)) + scoreParts := make([]string, len(fields)) + filterArgs := make([]interface{}, len(fields)) + scoreArgs := make([]interface{}, len(fields)) + for i, field := range fields { + expression := fmt.Sprintf("MATCH (%s) AGAINST (? IN NATURAL LANGUAGE MODE)", quoteIdentifier(field)) + filterParts[i] = expression + scoreParts[i] = expression + " * " + formatFloat(weights[i]) + filterArgs[i] = query + scoreArgs[i] = query + } + return "(" + strings.Join(filterParts, " OR ") + ")", filterArgs, + "(" + strings.Join(scoreParts, " + ") + ")", scoreArgs +} + +func (e *Engine) fullTextFields(kind string, text *types.MatchTextExpr) ([]string, []float64) { + var specifications []string + switch kind { + case "memory": + specifications = []string{"content_ltks", "tokenized_content_ltks"} + case "skill": + specifications = text.Fields + if len(specifications) == 0 { + specifications = []string{"name_tks^10", "tags_tks^5", "description_tks^3", "content_tks"} + } + default: + if e.flags.searchOriginalContent { + specifications = []string{"docnm_kwd^10", "content_with_weight", "important_tks^20", "question_tks^20"} + } else { + specifications = []string{"title_tks^10", "title_sm_tks^5", "important_tks^20", "question_tks^20", "content_ltks^2", "content_sm_ltks"} + } + } + fields := make([]string, 0, len(specifications)) + weights := make([]float64, 0, len(specifications)) + for _, specification := range specifications { + parts := strings.SplitN(specification, "^", 2) + field := parts[0] + if kind == "skill" && !strings.HasSuffix(field, "_tks") { + field += "_tks" + } + weight := 1.0 + if len(parts) == 2 { + if parsed, err := strconv.ParseFloat(parts[1], 64); err == nil { + weight = parsed + } + } + fields = append(fields, field) + weights = append(weights, weight) + } + var total float64 + for _, weight := range weights { + total += weight + } + if total <= 0 && len(weights) > 0 { + total = float64(len(weights)) + for i := range weights { + weights[i] = 1 + } + } + for i := range weights { + weights[i] /= total + } + return fields, weights +} + +func (e *Engine) fullTextHint(tableName, kind string) string { + if !e.flags.useFullTextHint || kind == "skill" { + return "" + } + fields, _ := e.fullTextFields(kind, &types.MatchTextExpr{}) + indexes := make([]string, len(fields)) + for i, field := range fields { + indexes[i] = "fts_idx_" + field + } + return fmt.Sprintf("/*+ UNION_MERGE(%s %s) */ ", tableName, strings.Join(indexes, " ")) +} + +func buildOrderBy(orderBy *types.OrderByExpr, kind string) (string, error) { + if orderBy == nil || len(orderBy.Fields) == 0 { + return "", nil + } + parts := make([]string, 0, len(orderBy.Fields)) + for _, order := range orderBy.Fields { + column := order.Field + if kind == "memory" { + column = mapMemoryField(column) + } + if kind == "chunk" && column == "chunk_order_int" { + column = "_order_id" + } + if !validColumns(kind)[column] { + return "", fmt.Errorf("unknown order field: %s", order.Field) + } + expression := quoteIdentifier(column) + if kind == "chunk" && arrayColumns[column] { + expression = "ARRAY_AVG(" + expression + ")" + } + direction := "ASC" + if order.Type == types.SortDesc { + direction = "DESC" + } + parts = append(parts, expression+" "+direction) + } + return " ORDER BY " + strings.Join(parts, ", "), nil +} + +func parseSearchPlan(expressions []interface{}) searchPlan { + var plan searchPlan + for _, expression := range expressions { + switch value := expression.(type) { + case string: + if value != "" { + plan.text = &types.MatchTextExpr{MatchingText: value} + } + case *types.MatchTextExpr: + if value != nil && value.MatchingText != "" { + plan.text = value + } + case *types.MatchDenseExpr: + if value != nil && len(value.EmbeddingData) > 0 { + plan.dense = value + } + case *types.FusionExpr: + plan.fusion = value + } + } + return plan +} + +func validateVectorExpr(dense *types.MatchDenseExpr) error { + if dense == nil || len(dense.EmbeddingData) == 0 { + return fmt.Errorf("vector expression is empty") + } + if dense.EmbeddingDataType != "" && dense.EmbeddingDataType != "float" { + return fmt.Errorf("embedding data type %q is not float", dense.EmbeddingDataType) + } + if !vectorColumnPattern.MatchString(dense.VectorColumnName) { + return fmt.Errorf("invalid vector column: %s", dense.VectorColumnName) + } + return nil +} + +func denseSimilarity(dense *types.MatchDenseExpr) float64 { + if dense.ExtraOptions != nil { + switch value := dense.ExtraOptions["similarity"].(type) { + case float64: + return value + case float32: + return float64(value) + case int: + return float64(value) + } + } + return 0 +} + +func fusionVectorWeight(fusion *types.FusionExpr) float64 { + if fusion == nil || fusion.FusionParams == nil { + return 0.5 + } + weights := strings.Split(stringValue(fusion.FusionParams["weights"]), ",") + if len(weights) != 2 { + return 0.5 + } + weight, err := strconv.ParseFloat(strings.TrimSpace(weights[1]), 64) + if err != nil { + return 0.5 + } + return weight +} + +func decodeRows(rows []map[string]interface{}, kind string) []map[string]interface{} { + result := make([]map[string]interface{}, len(rows)) + for i, row := range rows { + result[i] = decodeLogicalRow(row, kind) + } + return result +} + +func formatFloat(value float64) string { return strconv.FormatFloat(value, 'g', -1, 64) } + +func minPositive(first, second int) int { + if first <= 0 { + return second + } + if second <= 0 || first < second { + return first + } + return second +} + +func positiveOr(value, fallback int) int { + if value > 0 { + return value + } + return fallback +} + +func appendCopy(values []interface{}, extra ...interface{}) []interface{} { + result := make([]interface{}, 0, len(values)+len(extra)) + result = append(result, values...) + result = append(result, extra...) + return result +} + +func uniqueStrings(values []string) []string { + seen := make(map[string]bool, len(values)) + result := make([]string, 0, len(values)) + for _, value := range values { + if !seen[value] { + seen[value] = true + result = append(result, value) + } + } + return result +} + +func globalSearchCandidateLimit(req *types.SearchRequest) int { + return max(req.Offset, 0) + positiveOr(req.Limit, 30) +} + +func expandSearchPlan(plan searchPlan, candidateLimit int) searchPlan { + expanded := plan + if plan.text != nil { + text := *plan.text + text.TopN = max(text.TopN, candidateLimit) + expanded.text = &text + } + if plan.dense != nil { + dense := *plan.dense + dense.TopN = max(dense.TopN, candidateLimit) + expanded.dense = &dense + } + return expanded +} + +func searchHiddenSortFields(req *types.SearchRequest, mergeTables bool) []string { + if !mergeTables || req.OrderBy == nil || len(req.OrderBy.Fields) == 0 || len(req.SelectFields) == 0 || containsString(req.SelectFields, "*") { + return nil + } + fields := make([]string, 0, len(req.OrderBy.Fields)) + for _, order := range req.OrderBy.Fields { + if order.Field == "_score" || containsString(req.SelectFields, order.Field) || containsString(fields, order.Field) { + continue + } + fields = append(fields, order.Field) + } + return fields +} + +func mergeSearchChunks(chunks []map[string]interface{}, req *types.SearchRequest, plan searchPlan) []map[string]interface{} { + if req.OrderBy != nil && len(req.OrderBy.Fields) > 0 { + sort.SliceStable(chunks, func(i, j int) bool { + for _, order := range req.OrderBy.Fields { + comparison := compareSearchValues(chunks[i][order.Field], chunks[j][order.Field], order.Field) + if comparison == 0 { + continue + } + if order.Type == types.SortDesc { + return comparison > 0 + } + return comparison < 0 + } + return false + }) + } else if plan.text != nil || plan.dense != nil { + sort.SliceStable(chunks, func(i, j int) bool { + return compareSearchValues(chunks[i]["_score"], chunks[j]["_score"], "_score") > 0 + }) + } + offset := min(max(req.Offset, 0), len(chunks)) + limit := positiveOr(req.Limit, 30) + end := min(offset+limit, len(chunks)) + return chunks[offset:end] +} + +func compareSearchValues(left, right interface{}, field string) int { + if left == nil { + if right == nil { + return 0 + } + return -1 + } + if right == nil { + return 1 + } + leftNumber, leftNumeric := searchSortNumber(left, field) + rightNumber, rightNumeric := searchSortNumber(right, field) + if leftNumeric && rightNumeric { + switch { + case leftNumber < rightNumber: + return -1 + case leftNumber > rightNumber: + return 1 + default: + return 0 + } + } + return strings.Compare(fmt.Sprint(left), fmt.Sprint(right)) +} + +func searchSortNumber(value interface{}, field string) (float64, bool) { + if values, ok := interfaceSlice(value); ok { + if len(values) == 0 { + return 0, false + } + var total float64 + for _, item := range values { + number, ok := searchSortNumber(item, field) + if !ok { + return 0, false + } + total += number + } + return total / float64(len(values)), true + } + if number, ok := numberToFloat(value); ok { + return number, true + } + if number, ok := value.(json.Number); ok { + parsed, err := number.Float64() + return parsed, err == nil + } + if strings.HasSuffix(field, "_int") || strings.HasSuffix(field, "_flt") || field == "_score" { + parsed, err := strconv.ParseFloat(fmt.Sprint(value), 64) + return parsed, err == nil + } + return 0, false +} + +func isDBMSHybridPlan(plan searchPlan) bool { + return plan.text != nil && plan.dense != nil && plan.fusion != nil +} + +func (e *Engine) searchWithDBMS(ctx context.Context, tableName, kind string, condition map[string]interface{}, req *types.SearchRequest, plan searchPlan) ([]map[string]interface{}, bool, error) { + body, ok := buildDBMSBody(kind, condition, req, plan) + if !ok { + return nil, false, nil + } + encoded, err := json.Marshal(body) + if err != nil { + return nil, false, err + } + var raw []byte + if err := e.db.QueryRowContext(ctx, "SELECT DBMS_HYBRID_SEARCH.SEARCH(?, ?)", tableName, string(encoded)).Scan(&raw); err != nil { + return nil, false, fmt.Errorf("DBMS hybrid search: %w", err) + } + if len(raw) == 0 { + return []map[string]interface{}{}, true, nil + } + var documents []map[string]interface{} + if err := json.Unmarshal(raw, &documents); err == nil { + return decodeRows(documents, kind), true, nil + } + var response map[string]interface{} + if err := json.Unmarshal(raw, &response); err != nil { + return nil, false, fmt.Errorf("decode DBMS hybrid search response: %w", err) + } + documents = extractHybridHits(response) + return decodeRows(documents, kind), true, nil +} + +func buildDBMSBody(kind string, condition map[string]interface{}, req *types.SearchRequest, plan searchPlan) (map[string]interface{}, bool) { + filters := make([]interface{}, 0, len(condition)) + valid := validColumns(kind) + for rawField, value := range condition { + field := rawField + if kind == "memory" { + field = mapMemoryField(field) + } else if kind == "chunk" { + field = mapChunkField(field) + } else if kind == "skill" && field == "id" { + field = "skill_id" + } + if !valid[field] { + return nil, false + } + if field == "available_int" { + if fmt.Sprint(value) == "0" { + filters = append(filters, map[string]interface{}{"range": map[string]interface{}{field: map[string]interface{}{"lt": 1}}}) + } else { + filters = append(filters, map[string]interface{}{"bool": map[string]interface{}{"must_not": map[string]interface{}{"range": map[string]interface{}{field: map[string]interface{}{"lt": 1}}}}}) + } + } else if isEmptyFilterValue(value) { + continue + } else if values, ok := interfaceSlice(value); ok { + filters = append(filters, map[string]interface{}{"terms": map[string]interface{}{field: values}}) + } else { + filters = append(filters, map[string]interface{}{"term": map[string]interface{}{field: value}}) + } + } + queryText := plan.text.MatchingText + minimumShouldMatch := interface{}(0.0) + if plan.text.ExtraOptions != nil { + if value := plan.text.ExtraOptions["minimum_should_match"]; value != nil { + minimumShouldMatch = value + } + } + if value, ok := minimumShouldMatch.(float64); ok { + minimumShouldMatch = strconv.Itoa(int(value*100)) + "%" + } else if value, ok := minimumShouldMatch.(float32); ok { + minimumShouldMatch = strconv.Itoa(int(value*100)) + "%" + } + boolQuery := map[string]interface{}{ + "must": []interface{}{map[string]interface{}{"query_string": map[string]interface{}{ + "fields": tokenizedDBMSFields(kind), "type": "best_fields", "query": queryText, + "minimum_should_match": minimumShouldMatch, "boost": 1, + }}}, + "filter": filters, + "boost": 1 - fusionVectorWeight(plan.fusion), + } + if len(req.RankFeature) > 0 { + should := make([]interface{}, 0, len(req.RankFeature)) + for field, boost := range req.RankFeature { + if field != "pagerank_fea" { + field = "tag_feas." + field + } + should = append(should, map[string]interface{}{"rank_feature": map[string]interface{}{ + "field": field, "linear": map[string]interface{}{}, "boost": boost, + }}) + } + boolQuery["should"] = should + } + body := map[string]interface{}{ + "query": map[string]interface{}{"bool": boolQuery}, + "knn": map[string]interface{}{ + "field": plan.dense.VectorColumnName, "k": positiveOr(plan.dense.TopN, req.Limit), + "num_candidates": positiveOr(plan.dense.TopN, req.Limit) * 2, + "query_vector": plan.dense.EmbeddingData, "filter": map[string]interface{}{"bool": boolQuery}, + "similarity": denseSimilarity(plan.dense), + }, + "from": max(req.Offset, 0), "size": positiveOr(req.Limit, 30), + } + if req.OrderBy != nil && len(req.OrderBy.Fields) > 0 { + sorts := make([]interface{}, 0, len(req.OrderBy.Fields)) + for _, order := range req.OrderBy.Fields { + direction := "asc" + if order.Type == types.SortDesc { + direction = "desc" + } + orderInfo := map[string]interface{}{"order": direction} + if order.Field == "page_num_int" || order.Field == "top_int" { + orderInfo["unmapped_type"] = "float" + orderInfo["mode"] = "avg" + orderInfo["numeric_type"] = "double" + } else if strings.HasSuffix(order.Field, "_int") || strings.HasSuffix(order.Field, "_flt") { + orderInfo["unmapped_type"] = "float" + } else { + orderInfo["unmapped_type"] = "text" + } + sorts = append(sorts, map[string]interface{}{order.Field: orderInfo}) + } + body["sort"] = sorts + } + return body, true +} + +func tokenizedDBMSFields(kind string) []string { + switch kind { + case "memory": + return []string{"content_ltks", "tokenized_content_ltks"} + case "skill": + return []string{"name_tks^10", "tags_tks^5", "description_tks^3", "content_tks"} + default: + return []string{"title_tks^10", "title_sm_tks^5", "important_tks^20", "question_tks^20", "content_ltks^2", "content_sm_ltks"} + } +} + +func extractHybridHits(response map[string]interface{}) []map[string]interface{} { + hitsObject, ok := response["hits"].(map[string]interface{}) + if !ok { + return nil + } + hits, ok := hitsObject["hits"].([]interface{}) + if !ok { + return nil + } + result := make([]map[string]interface{}, 0, len(hits)) + for _, item := range hits { + hit, ok := item.(map[string]interface{}) + if !ok { + continue + } + document, _ := hit["_source"].(map[string]interface{}) + if document == nil { + document = hit + } + if score, ok := hit["_score"]; ok { + document["_score"] = score + } + result = append(result, document) + } + return result +} + +func isHybridUnavailableError(err error) bool { + message := strings.ToLower(err.Error()) + if !strings.Contains(message, "dbms_hybrid_search") { + return false + } + markers := []string{"does not exist", "not exist", "unknown", "not supported", "ora-00904", "1305"} + for _, marker := range markers { + if strings.Contains(message, marker) { + return true + } + } + return false +} diff --git a/internal/engine/oceanbase/search_test.go b/internal/engine/oceanbase/search_test.go new file mode 100644 index 0000000000..fd59c2bbbf --- /dev/null +++ b/internal/engine/oceanbase/search_test.go @@ -0,0 +1,105 @@ +// +// 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 oceanbase + +import ( + "context" + "reflect" + "regexp" + "testing" + + "ragflow/internal/engine/types" + + "github.com/DATA-DOG/go-sqlmock" +) + +func TestSearchAppliesPaginationAfterMergingTables(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("seekdb", "legacy_doc", db) + + for _, table := range []struct { + name string + rows *sqlmock.Rows + }{ + { + name: "ragflow_tenant_1", + rows: sqlmock.NewRows([]string{"id", "create_timestamp_flt"}). + AddRow("chunk-1", 10.0). + AddRow("chunk-4", 7.0), + }, + { + name: "ragflow_tenant_2", + rows: sqlmock.NewRows([]string{"id", "create_timestamp_flt"}). + AddRow("chunk-2", 9.0). + AddRow("chunk-3", 8.0), + }, + } { + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?")). + WithArgs("legacy_doc", table.name). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(*)"}).AddRow(1)) + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(`id`) FROM `" + table.name + "` WHERE 1=1")). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(id)"}).AddRow(2)) + mock.ExpectQuery(regexp.QuoteMeta("SELECT `id`, `create_timestamp_flt` FROM `" + table.name + "` WHERE 1=1 ORDER BY `create_timestamp_flt` DESC LIMIT 0, 3")). + WillReturnRows(table.rows) + } + + orderBy := (&types.OrderByExpr{}).Desc("create_timestamp_flt") + result, err := engine.Search(context.Background(), &types.SearchRequest{ + IndexNames: []string{"ragflow_tenant_1", "ragflow_tenant_2"}, + Offset: 1, + Limit: 2, + SelectFields: []string{"id"}, + OrderBy: orderBy, + }) + if err != nil { + t.Fatal(err) + } + if result.Total != 4 { + t.Fatalf("Search() total = %d, want 4", result.Total) + } + if len(result.Chunks) != 2 { + t.Fatalf("Search() returned %d chunks, want 2: %#v", len(result.Chunks), result.Chunks) + } + gotIDs := []interface{}{result.Chunks[0]["id"], result.Chunks[1]["id"]} + if want := []interface{}{"chunk-2", "chunk-3"}; !reflect.DeepEqual(gotIDs, want) { + t.Fatalf("Search() page IDs = %v, want %v", gotIDs, want) + } + if _, exists := result.Chunks[0]["create_timestamp_flt"]; exists { + t.Fatalf("Search() exposed an internal sort field: %#v", result.Chunks[0]) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestBuildQualifiedSelectFieldsUsesStructuredColumns(t *testing.T) { + fields, aliases, err := buildQualifiedSelectFields([]string{"content", "row_id()"}, "memory", "t") + if err != nil { + t.Fatal(err) + } + wantFields := "`t`.`id`, `t`.`content_ltks` AS `content`, `t`.`id` AS `row_id`" + if fields != wantFields { + t.Fatalf("buildQualifiedSelectFields() = %q, want %q", fields, wantFields) + } + if wantAliases := []string{"id", "content", "row_id"}; !reflect.DeepEqual(aliases, wantAliases) { + t.Fatalf("buildQualifiedSelectFields() aliases = %v, want %v", aliases, wantAliases) + } +} diff --git a/internal/engine/oceanbase/sql.go b/internal/engine/oceanbase/sql.go new file mode 100644 index 0000000000..8935d49eea --- /dev/null +++ b/internal/engine/oceanbase/sql.go @@ -0,0 +1,51 @@ +// +// 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 oceanbase + +import ( + "context" + "fmt" + "regexp" + "strings" +) + +var ( + jsonExtractStringPattern = regexp.MustCompile(`(?i)json_extract_string\s*\(\s*([^,]+?)\s*,\s*([^)]+?)\s*\)`) + jsonExtractNullPattern = regexp.MustCompile(`(?i)json_extract_isnull\s*\(\s*([^,]+?)\s*,\s*([^)]+?)\s*\)`) + limitPattern = regexp.MustCompile(`(?i)\blimit\b`) +) + +// RunSQL executes the read-only SQL produced by the chat SQL-retrieval flow. +func (e *Engine) RunSQL(ctx context.Context, tableName, sqlText string, kbIDs []string, format string) ([]map[string]interface{}, error) { + if tableName != "" { + if err := validateIdentifier(tableName); err != nil { + return nil, err + } + } + normalized := strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(sqlText), ";")) + normalized = strings.ReplaceAll(normalized, "`", "") + normalized = jsonExtractStringPattern.ReplaceAllString(normalized, "JSON_UNQUOTE(JSON_EXTRACT($1, $2))") + normalized = jsonExtractNullPattern.ReplaceAllString(normalized, "(JSON_EXTRACT($1, $2) IS NULL)") + lower := strings.ToLower(strings.TrimSpace(normalized)) + if !strings.HasPrefix(lower, "select ") && !strings.HasPrefix(lower, "with ") { + return nil, fmt.Errorf("only SELECT and WITH statements are allowed") + } + if !limitPattern.MatchString(normalized) { + normalized += " LIMIT 1024" + } + return e.queryRows(ctx, normalized) +} diff --git a/internal/engine/oceanbase/sql_test.go b/internal/engine/oceanbase/sql_test.go new file mode 100644 index 0000000000..7ff841f865 --- /dev/null +++ b/internal/engine/oceanbase/sql_test.go @@ -0,0 +1,144 @@ +// +// 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 oceanbase + +import ( + "context" + "regexp" + "strings" + "testing" + + "ragflow/internal/common" + "ragflow/internal/engine/types" + + "github.com/DATA-DOG/go-sqlmock" +) + +func init() { + _ = common.InitLogger("info", common.FileOutput{}, "oceanbase_test") +} + +func TestUpdateMetadataReplacesCompleteLegacyJSON(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("oceanbase", "legacy_doc", db) + + query := "REPLACE INTO `ragflow_doc_meta_tenant-1` (id, kb_id, meta_fields) VALUES (?, ?, ?)" + mock.ExpectExec(regexp.QuoteMeta(query)). + WithArgs("doc-1", "kb-1", `{"author":"Alice"}`). + WillReturnResult(sqlmock.NewResult(0, 1)) + + if err := engine.UpdateMetadata(context.Background(), "doc-1", "kb-1", map[string]interface{}{"author": "Alice"}, "tenant-1"); err != nil { + t.Fatal(err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestMetadataWritesRejectInvalidTenantIdentifier(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("oceanbase", "legacy_doc", db) + + err = engine.UpdateMetadata(context.Background(), "doc-1", "kb-1", + map[string]interface{}{"author": "Alice"}, "tenant`injected") + if err == nil || !strings.Contains(err.Error(), "invalid SQL identifier") { + t.Fatalf("UpdateMetadata() error = %v, want invalid identifier error", err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestDeleteMetadataKeysPreservesInvalidJSONRow(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("oceanbase", "legacy_doc", db) + + query := "SELECT meta_fields FROM `ragflow_doc_meta_tenant-1` WHERE id = ? AND kb_id = ? LIMIT 1" + mock.ExpectQuery(regexp.QuoteMeta(query)). + WithArgs("doc-1", "kb-1"). + WillReturnRows(sqlmock.NewRows([]string{"meta_fields"}).AddRow("{")) + + err = engine.DeleteMetadataKeys(context.Background(), "doc-1", "kb-1", []string{"author"}, "tenant-1") + if err == nil || !strings.Contains(err.Error(), "decode metadata") { + t.Fatalf("DeleteMetadataKeys() error = %v, want JSON decode error", err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestSkillFilterSearchCountsPhysicalPrimaryKey(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + engine := newEngineWithDB("seekdb", "legacy_doc", db) + engine.flags = featureFlags{} + + mock.ExpectQuery(regexp.QuoteMeta("SELECT COUNT(`skill_id`) FROM `skill_tenant-1_default` WHERE 1=1")). + WillReturnRows(sqlmock.NewRows([]string{"COUNT(skill_id)"}).AddRow(0)) + + chunks, total, err := engine.searchTableWithSQL(context.Background(), "skill_tenant-1_default", "skill", map[string]interface{}{}, &types.SearchRequest{Limit: 10}, searchPlan{}) + if err != nil { + t.Fatal(err) + } + if total != 0 || len(chunks) != 0 { + t.Fatalf("empty skill search total=%d chunks=%#v", total, chunks) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +func TestMemorySourceIDUsesScalarSQL(t *testing.T) { + where, args, err := buildFilter(map[string]interface{}{"source_id": "source-1"}, "memory") + if err != nil { + t.Fatal(err) + } + if where != "`source_id` = ?" { + t.Fatalf("memory source_id filter = %q", where) + } + if len(args) != 1 || args[0] != "source-1" { + t.Fatalf("memory source_id args = %#v", args) + } +} + +func TestChunkSourceIDUsesArraySQL(t *testing.T) { + where, args, err := buildFilter(map[string]interface{}{"source_id": "source-1"}, "chunk") + if err != nil { + t.Fatal(err) + } + if where != "ARRAY_CONTAINS(`source_id`, ?)" { + t.Fatalf("chunk source_id filter = %q", where) + } + if len(args) != 1 || args[0] != "source-1" { + t.Fatalf("chunk source_id args = %#v", args) + } +} diff --git a/internal/server/config.go b/internal/server/config.go index ea25b22c42..37b31187c9 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -175,7 +175,7 @@ func GetAllConfigs() ([]map[string]interface{}, error) { // Database databaseType := globalConfig.DatabaseType() switch databaseType { - case "mysql": + case "mysql", "oceanbase": mysqlConfig := globalConfig.GetMySQLConfig() exportedMySQLConfigs := mysqlConfig.ExportConfigs() allConfigs = append(allConfigs, exportedMySQLConfigs) @@ -194,6 +194,12 @@ func GetAllConfigs() ([]map[string]interface{}, error) { infinityConfig := globalConfig.GetInfinityConfig() exportedInfinityConfigs := infinityConfig.ExportConfigs() allConfigs = append(allConfigs, exportedInfinityConfigs) + case "oceanbase": + oceanBaseConfig := globalConfig.GetOceanBaseConfig() + allConfigs = append(allConfigs, oceanBaseConfig.ExportConfigs()) + case "seekdb": + seekDBConfig := globalConfig.GetSeekDBConfig() + allConfigs = append(allConfigs, seekDBConfig.ExportConfigs()) default: return nil, fmt.Errorf("not supported doc engine: %s", docEngineType) } diff --git a/internal/server/config/database_config.go b/internal/server/config/database_config.go index acc39cfe28..eb57232212 100644 --- a/internal/server/config/database_config.go +++ b/internal/server/config/database_config.go @@ -18,6 +18,7 @@ package config import ( "fmt" + "ragflow/internal/common" "github.com/spf13/viper" ) @@ -41,15 +42,58 @@ type MySQLConfig struct { func (c *Config) ParseDatabaseConfig(v *viper.Viper) error { databaseType := c.general.Database + if envType := common.GetEnvSmall(common.EnvDBType); envType != "" { + databaseType = envType + } switch databaseType { case "mysql": c.parseMySQLConfig(v) + case "oceanbase": + c.parseOceanBaseDatabaseConfig(v) default: return fmt.Errorf("database type %s is not supported", databaseType) } return nil } +func (c *Config) parseOceanBaseDatabaseConfig(v *viper.Viper) { + // OceanBase uses the MySQL wire protocol, so the existing DAO can keep + // consuming MySQLConfig. Support both the flat DB_TYPE=oceanbase section and + // the document-engine nested config. + c.parseMySQLConfig(v) + sub := v.Sub("oceanbase") + if sub == nil { + return + } + if nested := sub.Sub("config"); nested != nil { + if nested.IsSet("db_name") { + c.database.MySQL.DatabaseName = nested.GetString("db_name") + } + sub = nested + } + if sub.IsSet("name") { + c.database.MySQL.DatabaseName = sub.GetString("name") + } + if sub.IsSet("db_name") { + c.database.MySQL.DatabaseName = sub.GetString("db_name") + } + if sub.IsSet("user") { + c.database.MySQL.User = sub.GetString("user") + } + if sub.IsSet("password") { + c.database.MySQL.Password = sub.GetString("password") + } + if sub.IsSet("host") { + c.database.MySQL.Host = sub.GetString("host") + } + if sub.IsSet("port") { + c.database.MySQL.Port = sub.GetInt("port") + } + if sub.IsSet("max_connections") { + c.database.MySQL.MaxConnections = sub.GetInt("max_connections") + } +} + func (c *Config) parseMySQLConfig(v *viper.Viper) { // Default MySQL config diff --git a/internal/server/config/doc_engine_config.go b/internal/server/config/doc_engine_config.go index 30c1232e48..41dfc04877 100644 --- a/internal/server/config/doc_engine_config.go +++ b/internal/server/config/doc_engine_config.go @@ -17,13 +17,36 @@ package config import ( + "fmt" + "strings" + "github.com/spf13/viper" ) type DocEngineConfig struct { - ES ElasticsearchConfig `mapstructure:"es"` - Infinity InfinityConfig `mapstructure:"infinity"` - SereneDB SereneDBConfig `mapstructure:"serenedb"` + ES ElasticsearchConfig `mapstructure:"es"` + Infinity InfinityConfig `mapstructure:"infinity"` + OceanBase OceanBaseConfig `mapstructure:"oceanbase"` + SeekDB OceanBaseConfig `mapstructure:"seekdb"` + SereneDB SereneDBConfig `mapstructure:"serenedb"` +} + +// OceanBaseConfig mirrors the existing oceanbase/seekdb service_conf.yaml +// structure used by the Python connector. +type OceanBaseConfig struct { + Scheme string `mapstructure:"scheme"` + Config OceanBaseConnectionConfig `mapstructure:"config"` +} + +// OceanBaseConnectionConfig contains the MySQL-protocol connection settings +// used by both OceanBase and SeekDB document engines. +type OceanBaseConnectionConfig struct { + DBName string `mapstructure:"db_name"` + User string `mapstructure:"user"` + Password string `mapstructure:"password"` + Host string `mapstructure:"host"` + Port int `mapstructure:"port"` + MaxConnections int `mapstructure:"max_connections"` } // ElasticsearchConfig Elasticsearch configuration @@ -58,6 +81,8 @@ type SereneDBConfig struct { func (c *Config) ParseDocEngineConfig(v *viper.Viper) error { c.parseInfinityConfig(v) c.parseElasticsearchConfig(v) + c.parseOceanBaseConfig(v, "oceanbase", &c.docEngine.OceanBase) + c.parseOceanBaseConfig(v, "seekdb", &c.docEngine.SeekDB) c.parseSereneDBConfig(v) return nil } @@ -102,6 +127,105 @@ func (c *Config) parseSereneDBConfig(v *viper.Viper) { } } +func (c *Config) parseOceanBaseConfig(v *viper.Viper, key string, target *OceanBaseConfig) { + defaultPassword := c.database.MySQL.Password + target.Scheme = "oceanbase" + target.Config = OceanBaseConnectionConfig{ + DBName: "test", + User: "root@test", + Password: defaultPassword, + Host: "localhost", + Port: 2881, + MaxConnections: 300, + } + if key == "seekdb" { + target.Config.DBName = "ragflow_doc" + target.Config.User = "root" + } + + if !v.IsSet(key) { + return + } + sub := v.Sub(key) + if sub == nil { + return + } + if sub.IsSet("scheme") { + target.Scheme = sub.GetString("scheme") + } + connection := sub.Sub("config") + if connection == nil { + return + } + if connection.IsSet("db_name") { + target.Config.DBName = connection.GetString("db_name") + } + if connection.IsSet("user") { + target.Config.User = connection.GetString("user") + } + if connection.IsSet("password") { + target.Config.Password = connection.GetString("password") + } + if connection.IsSet("host") { + target.Config.Host = connection.GetString("host") + } + if connection.IsSet("port") { + target.Config.Port = connection.GetInt("port") + } + if connection.IsSet("max_connections") { + target.Config.MaxConnections = connection.GetInt("max_connections") + } +} + +// ResolveOceanBaseConnection returns the effective existing configuration for +// an OceanBase-family document engine. With scheme=mysql, Python takes the +// endpoint and credentials from the mysql section while retaining db_name from +// the nested oceanbase/seekdb config; Go intentionally follows that contract. +func (c *Config) ResolveOceanBaseConnection(engineType string) (OceanBaseConnectionConfig, error) { + var configured OceanBaseConfig + switch strings.ToLower(engineType) { + case "oceanbase": + configured = c.docEngine.OceanBase + case "seekdb": + configured = c.docEngine.SeekDB + default: + return OceanBaseConnectionConfig{}, fmt.Errorf("not an OceanBase-family engine: %s", engineType) + } + + resolved := configured.Config + if strings.EqualFold(configured.Scheme, "mysql") { + mysqlConfig := c.database.MySQL + resolved.User = mysqlConfig.User + resolved.Password = mysqlConfig.Password + resolved.Host = mysqlConfig.Host + resolved.Port = mysqlConfig.Port + resolved.MaxConnections = mysqlConfig.MaxConnections + } + return resolved, nil +} + +func (c *Config) GetOceanBaseConfig() OceanBaseConfig { + return c.docEngine.OceanBase +} + +func (c *Config) GetSeekDBConfig() OceanBaseConfig { + return c.docEngine.SeekDB +} + +func (o OceanBaseConfig) ExportConfigs() map[string]interface{} { + return map[string]interface{}{ + "scheme": o.Scheme, + "config": map[string]interface{}{ + "db_name": o.Config.DBName, + "user": o.Config.User, + "password": o.Config.Password, + "host": o.Config.Host, + "port": o.Config.Port, + "max_connections": o.Config.MaxConnections, + }, + } +} + func (c *Config) parseInfinityConfig(v *viper.Viper) { // Default Infinity config c.docEngine.Infinity.URI = "localhost:23817" diff --git a/internal/server/config/environments.go b/internal/server/config/environments.go index d05e2295cb..aec564bffe 100644 --- a/internal/server/config/environments.go +++ b/internal/server/config/environments.go @@ -155,9 +155,9 @@ func (c *Config) GetEnvironments() error { docEngine := common.GetEnvSmall(common.EnvDocEngine) if docEngine != "" { switch docEngine { - case "infinity", "elasticsearch": + case "infinity", "elasticsearch", "oceanbase", "seekdb": c.environments.DocumentEngineType = docEngine - case "opensearch", "oceanbase": + case "opensearch": return fmt.Errorf("not implemented: %s", docEngine) default: return fmt.Errorf("invalid doc engine: %s", docEngine) @@ -168,8 +168,8 @@ func (c *Config) GetEnvironments() error { databaseType := common.GetEnvSmall(common.EnvDBType) if databaseType != "" { switch databaseType { - case "mysql": - c.environments.DatabaseType = "mysql" + case "mysql", "oceanbase": + c.environments.DatabaseType = databaseType default: return fmt.Errorf("invalid database type: %s", databaseType) } diff --git a/internal/server/config/general_config.go b/internal/server/config/general_config.go index 638268ab4a..d8fe8d3c10 100644 --- a/internal/server/config/general_config.go +++ b/internal/server/config/general_config.go @@ -95,6 +95,9 @@ func (c *Config) GetMode() string { } func (c *Config) DatabaseType() string { + if c.environments.DatabaseType != "" { + return c.environments.DatabaseType + } return c.general.Database } diff --git a/internal/server/config/oceanbase_config_test.go b/internal/server/config/oceanbase_config_test.go new file mode 100644 index 0000000000..dcbf966dec --- /dev/null +++ b/internal/server/config/oceanbase_config_test.go @@ -0,0 +1,147 @@ +// +// 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 config + +import ( + "reflect" + "testing" + + "github.com/spf13/viper" +) + +func TestParseExistingOceanBaseAndSeekDBConfig(t *testing.T) { + t.Setenv("DB_TYPE", "") + mysqlCredential := t.Name() + "-mysql" + oceanBaseCredential := t.Name() + "-oceanbase" + ignoredCredential := t.Name() + "-ignored" + v := viper.New() + v.Set("mysql", map[string]interface{}{ + "name": "rag_flow", "user": "mysql-user", "password": mysqlCredential, + "host": "mysql-host", "port": 3307, "max_connections": 456, + }) + v.Set("oceanbase", map[string]interface{}{ + "scheme": "oceanbase", + "config": map[string]interface{}{ + "db_name": "legacy_doc", "user": "root@ragflow", "password": oceanBaseCredential, + "host": "ob-host", "port": 2881, "max_connections": 123, + }, + }) + v.Set("seekdb", map[string]interface{}{ + "scheme": "mysql", + "config": map[string]interface{}{ + "db_name": "legacy_seekdb", "user": "ignored-user", "password": ignoredCredential, + "host": "ignored-host", "port": 2881, "max_connections": 12, + }, + }) + + config := &Config{} + if err := config.ParseGeneralConfig(v); err != nil { + t.Fatal(err) + } + if err := config.ParseDatabaseConfig(v); err != nil { + t.Fatal(err) + } + if err := config.ParseDocEngineConfig(v); err != nil { + t.Fatal(err) + } + + oceanBase, err := config.ResolveOceanBaseConnection("oceanbase") + if err != nil { + t.Fatal(err) + } + wantOceanBase := OceanBaseConnectionConfig{ + DBName: "legacy_doc", User: "root@ragflow", Password: oceanBaseCredential, + Host: "ob-host", Port: 2881, MaxConnections: 123, + } + if !reflect.DeepEqual(oceanBase, wantOceanBase) { + t.Fatalf("oceanbase config = %#v, want %#v", oceanBase, wantOceanBase) + } + + seekDB, err := config.ResolveOceanBaseConnection("seekdb") + if err != nil { + t.Fatal(err) + } + wantSeekDB := OceanBaseConnectionConfig{ + DBName: "legacy_seekdb", User: "mysql-user", Password: mysqlCredential, + Host: "mysql-host", Port: 3307, MaxConnections: 456, + } + if !reflect.DeepEqual(seekDB, wantSeekDB) { + t.Fatalf("seekdb config = %#v, want %#v", seekDB, wantSeekDB) + } +} + +func TestOceanBaseDefaultCredentialMatchesMySQLWireConfig(t *testing.T) { + t.Setenv("DB_TYPE", "") + v := viper.New() + config := &Config{} + if err := config.ParseGeneralConfig(v); err != nil { + t.Fatal(err) + } + if err := config.ParseDatabaseConfig(v); err != nil { + t.Fatal(err) + } + if err := config.ParseDocEngineConfig(v); err != nil { + t.Fatal(err) + } + + oceanBase, err := config.ResolveOceanBaseConnection("oceanbase") + if err != nil { + t.Fatal(err) + } + if oceanBase.Password != config.GetMySQLConfig().Password { + t.Fatal("OceanBase default credential diverged from the MySQL-wire default") + } +} + +func TestOceanBaseEnvironmentTypesAreAccepted(t *testing.T) { + t.Setenv("DOC_ENGINE", "seekdb") + t.Setenv("DB_TYPE", "oceanbase") + config := &Config{} + if err := config.GetEnvironments(); err != nil { + t.Fatal(err) + } + if got := config.DocEngineType(); got != "seekdb" { + t.Fatalf("doc engine = %q, want seekdb", got) + } + if got := config.DatabaseType(); got != "oceanbase" { + t.Fatalf("database type = %q, want oceanbase", got) + } +} + +func TestParseOceanBaseAsMainDatabaseFromNestedConfig(t *testing.T) { + t.Setenv("DB_TYPE", "oceanbase") + oceanBaseCredential := t.Name() + "-oceanbase" + v := viper.New() + v.Set("oceanbase", map[string]interface{}{ + "scheme": "oceanbase", + "config": map[string]interface{}{ + "db_name": "rag_flow_ob", "user": "root@tenant", "password": oceanBaseCredential, + "host": "ob-main", "port": 2881, "max_connections": 300, + }, + }) + config := &Config{} + if err := config.ParseGeneralConfig(v); err != nil { + t.Fatal(err) + } + if err := config.ParseDatabaseConfig(v); err != nil { + t.Fatal(err) + } + got := config.GetMySQLConfig() + if got.DatabaseName != "rag_flow_ob" || got.User != "root@tenant" || got.Host != "ob-main" || got.Port != 2881 { + t.Fatalf("main OceanBase config was not mapped to MySQL wire config: %#v", got) + } +} diff --git a/internal/service/chat_pipeline.go b/internal/service/chat_pipeline.go index 88fb8223b6..b96ade8ffc 100644 --- a/internal/service/chat_pipeline.go +++ b/internal/service/chat_pipeline.go @@ -3425,7 +3425,7 @@ func buildSQLPrompts(engineName, tableName, question string, fieldMap map[string if isRowCountQuestion(question) { overrideSQL = fmt.Sprintf("SELECT COUNT(*) AS rows FROM %s", tableName) } - case "oceanbase": + case "oceanbase", "seekdb": sysPrompt = oceanbaseSQLSysPrompt bullets := strings.Builder{} for _, n := range names { @@ -3533,7 +3533,7 @@ func sortedFieldNames(fieldMap map[string]interface{}) []string { // OpenSearch share the direct-column template. expectedCol is // "docnm" for Infinity or "docnm_kwd" for everything else. func buildMissingColumnsRepairPrompt(engineName, tableName, question, prevSQL, expectedCol string, fieldMap map[string]interface{}) string { - isJSONEngine := engineName == "infinity" || engineName == "oceanbase" + isJSONEngine := engineName == "infinity" || engine.IsOceanBaseFamily(engineName) names := sortedFieldNames(fieldMap) bullets := strings.Builder{} if isJSONEngine { @@ -3562,7 +3562,7 @@ func buildMissingColumnsRepairPrompt(engineName, tableName, question, prevSQL, e // buildExecutionErrorRepairPrompt returns the engine-specific user // prompt for the execution-error repair flow. func buildExecutionErrorRepairPrompt(engineName, tableName, question, errMsg string, fieldMap map[string]interface{}) string { - isJSONEngine := engineName == "infinity" || engineName == "oceanbase" + isJSONEngine := engineName == "infinity" || engine.IsOceanBaseFamily(engineName) names := sortedFieldNames(fieldMap) bullets := strings.Builder{} if isJSONEngine { diff --git a/internal/service/chat_pipeline_test.go b/internal/service/chat_pipeline_test.go index 0b23341e9c..b47f974e26 100644 --- a/internal/service/chat_pipeline_test.go +++ b/internal/service/chat_pipeline_test.go @@ -1112,6 +1112,9 @@ func TestExpectedDocNameColumn(t *testing.T) { if got := expectedDocNameColumn("oceanbase"); got != "docnm_kwd" { t.Errorf("oceanbase = %q, want docnm_kwd", got) } + if got := expectedDocNameColumn("seekdb"); got != "docnm_kwd" { + t.Errorf("seekdb = %q, want docnm_kwd", got) + } if got := expectedDocNameColumn("elasticsearch"); got != "docnm_kwd" { t.Errorf("elasticsearch = %q, want docnm_kwd", got) } diff --git a/internal/service/chunk/vector.go b/internal/service/chunk/vector.go index 0c9fc2e349..7f88386b48 100644 --- a/internal/service/chunk/vector.go +++ b/internal/service/chunk/vector.go @@ -39,8 +39,8 @@ type vectorFetcher interface { // This is used by citation insertion (insert_citations) to hydrate chunk // vectors on demand, since the main retrieval path skips vector transport. // -// On Infinity / OceanBase the chunks already carry vectors, so we skip -// the round-trip. On ES we query by chunk ID list. +// Infinity search results already carry vectors. Other engines hydrate them +// with a filtered query. // // Degrades gracefully: if the engine returns an error, zero vectors are // returned for all chunk IDs rather than failing the caller. @@ -55,8 +55,7 @@ func FetchChunkVectors(ctx context.Context, engine vectorFetcher, chunkIDs, tena } // Infinity already ships vectors with chunks; no need to fetch. - // TODO: OceanBase engine is not yet implemented — add "oceanbase" here when it lands. - if engine.GetType() == "infinity" || engine.GetType() == "oceanbase" { + if engine.GetType() == "infinity" { for _, cid := range chunkIDs { out[cid] = zeroVector(dim) } diff --git a/internal/service/chunk/vector_test.go b/internal/service/chunk/vector_test.go index bf196ce98e..5809523e7c 100644 --- a/internal/service/chunk/vector_test.go +++ b/internal/service/chunk/vector_test.go @@ -113,14 +113,26 @@ func TestFetchChunkVectors_InfinitySkipsSearch(t *testing.T) { } } -func TestFetchChunkVectors_OceanbaseSkipsSearch(t *testing.T) { - mock := &mockVectorFetcher{engineType: "oceanbase"} +func TestFetchChunkVectors_OceanBaseSearchesVector(t *testing.T) { + mock := &mockVectorFetcher{ + engineType: "oceanbase", + searchResults: map[string]*types.SearchResult{ + "ragflow_t1": { + Chunks: []map[string]interface{}{ + {"id": "c1", "q_3_vec": []interface{}{0.1, 0.2, 0.3}}, + }, + }, + }, + } result := FetchChunkVectors(bg, mock, []string{"c1"}, []string{"t1"}, []string{"kb1"}, 3) if len(result) != 1 { t.Fatalf("expected 1 entry, got %d", len(result)) } - if len(mock.searchCalled) > 0 { - t.Error("Search should not be called for OceanBase engine") + if !reflect.DeepEqual(mock.searchCalled, []string{"ragflow_t1"}) { + t.Errorf("search calls = %v, want [ragflow_t1]", mock.searchCalled) + } + if !reflect.DeepEqual(result["c1"], []float64{0.1, 0.2, 0.3}) { + t.Errorf("c1 = %v, want [0.1 0.2 0.3]", result["c1"]) } } diff --git a/internal/service/document/document_metadata.go b/internal/service/document/document_metadata.go index 4d112513fb..4359c49555 100644 --- a/internal/service/document/document_metadata.go +++ b/internal/service/document/document_metadata.go @@ -77,7 +77,7 @@ func (s *DocumentService) DeleteDocumentMetadata(ctx context.Context, docID stri } // Delete metadata using the document engine - err = s.docEngine.DeleteMetadataKeys(nil, docID, doc.KbID, keys, tenantID) + err = s.docEngine.DeleteMetadataKeys(ctx, docID, doc.KbID, keys, tenantID) if err != nil { return fmt.Errorf("failed to delete metadata: %w", err) } @@ -106,7 +106,7 @@ func (s *DocumentService) DeleteDocumentAllMetadata(ctx context.Context, docID s } // Delete entire document metadata - _, err = s.docEngine.DeleteMetadata(nil, condition, tenantID) + _, err = s.docEngine.DeleteMetadata(ctx, condition, tenantID) if err != nil { return fmt.Errorf("failed to delete document metadata: %w", err) } diff --git a/internal/service/memory.go b/internal/service/memory.go index 6a7b6a8103..8e3faed147 100644 --- a/internal/service/memory.go +++ b/internal/service/memory.go @@ -814,16 +814,17 @@ func sameStringSet(a, b []string) bool { // err := service.DeleteMemory(ctx, "user123", "memory456") func (s *MemoryService) DeleteMemory(ctx context.Context, userID, memoryID string) error { // Verify the caller has access to this memory - if _, err := s.requireMemoryAccess(ctx, userID, memoryID); err != nil { + memory, err := s.requireMemoryAccess(ctx, userID, memoryID) + if err != nil { return err } // TODO: Delete associated message index - Implementation pending MessageService - // messageService := NewMessageService() - // hasIndex, _ := messageService.HasIndex(memory.TenantID, memoryID) - // if hasIndex { - // messageService.DeleteMessage(nil, memory.TenantID, memoryID) - // } + if s.docEngine != nil && engine.IsOceanBaseFamily(s.docEngine.GetType()) { + if err := s.docEngine.DropChunkStore(ctx, memoryIndexName(memory.TenantID), memoryID); err != nil { + return fmt.Errorf("delete memory messages: %w", err) + } + } // Delete memory record if err := s.memoryDAO.DeleteByID(ctx, dao.DB, memoryID); err != nil { @@ -850,8 +851,12 @@ func (s *MemoryService) ForgetMessage(ctx context.Context, userID string, memory forgetTime := now.Format("2006-01-02 15:04:05") messageDocID := fmt.Sprintf("%s_%d", memoryID, messageID) updates := map[string]interface{}{ - "forget_at": forgetTime, - "forget_at_flt": now.UnixMilli(), + "forget_at": forgetTime, + } + // OceanBase/SeekDB memory tables contain forget_at but no forget_at_flt. + // Keep the existing companion-field update for other engines. + if !engine.IsOceanBaseFamily(s.docEngine.GetType()) { + updates["forget_at_flt"] = now.UnixMilli() } condition := map[string]interface{}{ "id": messageDocID, diff --git a/internal/service/memory_message_test.go b/internal/service/memory_message_test.go index e0a21937a9..8c7265240e 100644 --- a/internal/service/memory_message_test.go +++ b/internal/service/memory_message_test.go @@ -27,6 +27,17 @@ func TestIsMessageDocumentNotFound(t *testing.T) { } } +func TestMemoryIndexNameMatchesPythonPrefix(t *testing.T) { + t.Setenv(common.EnvESIndexPrefix, "") + if got := memoryIndexName("tenant-1"); got != "memory_tenant-1" { + t.Fatalf("memoryIndexName() = %q", got) + } + t.Setenv(common.EnvESIndexPrefix, "legacy") + if got := memoryIndexName("tenant-1"); got != "memory_legacy_tenant-1" { + t.Fatalf("memoryIndexName() with prefix = %q", got) + } +} + func TestValidateMemorySearchModels(t *testing.T) { firstTenantEmbeddingID := "tenant-embedding-1" secondTenantEmbeddingID := "tenant-embedding-2" @@ -57,6 +68,7 @@ func TestRequireMemoryAccessReturnsCanceledContext(t *testing.T) { type memoryMessageDocEngine struct { fakeChatDocEngine + engineType string searchReq *enginetypes.SearchRequest searchResp *enginetypes.SearchResult updateCond map[string]interface{} @@ -85,7 +97,12 @@ func (e *memoryMessageDocEngine) FilterDocIdsByMetaPushdown(_ context.Context, _ return nil } -func (e *memoryMessageDocEngine) GetType() string { return "memory" } +func (e *memoryMessageDocEngine) GetType() string { + if e.engineType != "" { + return e.engineType + } + return "memory" +} func setupMemoryMessageTestDB(t *testing.T) { t.Helper() @@ -112,6 +129,57 @@ func setupMemoryMessageTestDB(t *testing.T) { }) } +func TestForgetMessageKeepsCompanionFieldForNonOceanBaseEngines(t *testing.T) { + setupMemoryMessageTestDB(t) + + if err := dao.DB.Create(&entity.Memory{ + ID: "memory-1", + Name: "Test memory", + TenantID: "user-1", + MemoryType: dao.MemoryTypeRaw, + StorageType: "table", + EmbdID: "embd", + LLMID: "llm", + Permissions: string(entity.TenantPermissionMe), + MemorySize: MemorySizeLimit, + }).Error; err != nil { + t.Fatalf("seed memory: %v", err) + } + + for _, test := range []struct { + engineType string + wantForgetAtCompanion bool + }{ + {engineType: "elasticsearch", wantForgetAtCompanion: true}, + {engineType: "infinity", wantForgetAtCompanion: true}, + {engineType: "oceanbase", wantForgetAtCompanion: false}, + {engineType: "seekdb", wantForgetAtCompanion: false}, + } { + t.Run(test.engineType, func(t *testing.T) { + docEngine := &memoryMessageDocEngine{engineType: test.engineType} + service := NewMemoryService() + service.docEngine = docEngine + + if err := service.ForgetMessage(context.Background(), "user-1", "memory-1", 42); err != nil { + t.Fatalf("ForgetMessage() error = %v", err) + } + if docEngine.updateCond["id"] != "memory-1_42" { + t.Fatalf("update condition = %#v", docEngine.updateCond) + } + if docEngine.updateBase != "memory_user-1" || docEngine.updateID != "memory-1" { + t.Fatalf("update target = (%q, %q)", docEngine.updateBase, docEngine.updateID) + } + if forgetAt, ok := docEngine.updateValue["forget_at"].(string); !ok || forgetAt == "" { + t.Fatalf("forget_at = %#v, want non-empty string", docEngine.updateValue["forget_at"]) + } + _, hasCompanion := docEngine.updateValue["forget_at_flt"] + if hasCompanion != test.wantForgetAtCompanion { + t.Fatalf("forget_at_flt present = %t, want %t", hasCompanion, test.wantForgetAtCompanion) + } + }) + } +} + func TestUpdateMemoryTeamMemberCannotChangePermissions(t *testing.T) { setupMemoryMessageTestDB(t) diff --git a/internal/service/metadata_filter_test.go b/internal/service/metadata_filter_test.go index 7e5d46c83f..86a40c32d2 100644 --- a/internal/service/metadata_filter_test.go +++ b/internal/service/metadata_filter_test.go @@ -174,6 +174,20 @@ func TestApplyMetaFilter_NotEqualsAlias(t *testing.T) { } } +func TestApplyMetaFilter_MultiValueNegativeOperators(t *testing.T) { + metas := common.MetaData{"tags": {"a": {"doc1"}, "b": {"doc1"}}} + tests := []MetaFilterCondition{ + {Key: "tags", Value: "a", Op: "≠"}, + {Key: "tags", Value: "a", Op: "not in"}, + } + for _, condition := range tests { + result := ApplyMetaFilter(metas, []MetaFilterCondition{condition}, "and") + if len(result) != 1 || result[0] != "doc1" { + t.Errorf("operator %q fallback result = %v, want [doc1]", condition.Op, result) + } + } +} + func TestApplyMetaFilter_GreaterThan(t *testing.T) { metas := common.MetaData{"score": {"85": {"doc1"}, "70": {"doc2"}}} filters := []MetaFilterCondition{{Key: "score", Value: "80", Op: ">"}} diff --git a/internal/service/nlp/retrieval.go b/internal/service/nlp/retrieval.go index 29b16861b0..2a9f085d66 100644 --- a/internal/service/nlp/retrieval.go +++ b/internal/service/nlp/retrieval.go @@ -167,7 +167,7 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) vtWeight := vtWeightOrig qb := GetQueryBuilder() useInfinity := engine.GetEngineType() == "infinity" - useOceanBase := false // TODO: add OceanBase detection when supported + useOceanBase := engine.IsOceanBaseFamily(s.docEngine.GetType()) // For ES path: call GetScores() for second-pass KNN to get clean cosine similarity // For Infinity path: use _score directly (scores already normalized during fusion) @@ -212,13 +212,19 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) termSimilarity = sim vectorSimilarity = sim } else if useOceanBase { - // OceanBase: extract vectors and compute locally (not implemented) - sim = make([]float64, len(searchResult.IDs)) - for i := range searchResult.IDs { - sim[i] = 0.0 - } - termSimilarity = sim - vectorSimilarity = sim + // OceanBase returns the selected vector column, so reranking can + // reproduce the Python connector's local cosine calculation. + sim, termSimilarity, vectorSimilarity = RerankStandard( + searchResult.Chunks, + nil, + searchResult.QueryVector, + req.Question, + tkWeight, + vtWeight, + "content_ltks", + qb, + *req.RankFeature, + ) } else { // ES PATH: Two-pass KNN approach for clean cosine similarity scores // @@ -665,7 +671,7 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque // Build source with vector column for ES searchSrc := make([]string, len(searchRequest.SelectFields)) copy(searchSrc, searchRequest.SelectFields) - if engine.GetEngineType() == "elasticsearch" { + if engine.GetEngineType() == "elasticsearch" || engine.IsOceanBaseFamily(engine.GetEngineType()) { searchSrc = append(searchSrc, matchDense.VectorColumnName) }