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 <haijin.chn@gmail.com>
This commit is contained in:
wangyunlai
2026-08-10 15:06:32 +08:00
committed by GitHub
parent c0bc146fcb
commit 73d006fa0e
39 changed files with 5830 additions and 47 deletions

View File

@@ -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 {

View File

@@ -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:

View File

@@ -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, &current); 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
}

View File

@@ -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)
}
}
}

View File

@@ -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)
}

View File

@@ -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)
}

View File

@@ -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
}

View File

@@ -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)
}
}

View File

@@ -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
}

View File

@@ -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)
}
}

View File

@@ -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
}

View File

@@ -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_<dim>_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
}
}

View File

@@ -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": "<em>Apple</em> <em>cat</em> concatenate.",
"chinese": "这是<em>数据库</em>系统",
}
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": "<em>OceanBase</em> search"}
if got := engine.GetHighlight(chunks, []string{"oceanbase"}, "content"); !reflect.DeepEqual(got, want) {
t.Fatalf("GetHighlight() = %#v, want %#v", got, want)
}
}

View File

@@ -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("<em>")
result.WriteString(text[match.start:match.end])
result.WriteString("</em>")
position = match.end
}
result.WriteString(text[position:])
return result.String()
}

View File

@@ -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 := "<em>Apple</em>, <em>cat</em> and <em>dog</em>; 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 != "这是<em>数据库</em>系统" {
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)
}
}

View File

@@ -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)
}
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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")
}

View File

@@ -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)
}
}

View File

@@ -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
}

View File

@@ -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)
}
}

View File

@@ -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)
}

View File

@@ -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)
}
}