mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-14 20:54:30 +08:00
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:
@@ -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 {
|
||||
|
||||
@@ -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:
|
||||
|
||||
306
internal/engine/oceanbase/client.go
Normal file
306
internal/engine/oceanbase/client.go
Normal 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, ¤t); err == nil && current >= target {
|
||||
return
|
||||
}
|
||||
if _, err := e.db.ExecContext(ctx, fmt.Sprintf("SET GLOBAL ob_query_timeout=%d", target)); err != nil {
|
||||
common.Warn("Failed to set OceanBase query timeout", zap.Error(err))
|
||||
return
|
||||
}
|
||||
// Existing sessions retain the old global value. Closing idle sessions
|
||||
// mirrors Python's engine.dispose() while keeping in-flight work alive.
|
||||
maxIdle := e.maxIdleConns
|
||||
if maxIdle <= 0 {
|
||||
maxIdle = e.db.Stats().MaxOpenConnections
|
||||
}
|
||||
e.db.SetMaxIdleConns(0)
|
||||
e.db.SetMaxIdleConns(maxIdle)
|
||||
}
|
||||
|
||||
// Ping checks that the database connection is alive.
|
||||
func (e *Engine) Ping(ctx context.Context) error {
|
||||
if e == nil || e.db == nil {
|
||||
return fmt.Errorf("OceanBase client is not initialized")
|
||||
}
|
||||
return e.db.PingContext(ctx)
|
||||
}
|
||||
|
||||
// Close closes the SQL connection pool.
|
||||
func (e *Engine) Close() error {
|
||||
if e == nil || e.db == nil {
|
||||
return nil
|
||||
}
|
||||
return e.db.Close()
|
||||
}
|
||||
|
||||
// GetType returns oceanbase or seekdb, preserving the configured engine name.
|
||||
func (e *Engine) GetType() string { return e.engineType }
|
||||
|
||||
// SupportsPageRank reports that OceanBase applies pagerank during ranking.
|
||||
func (e *Engine) SupportsPageRank() bool { return true }
|
||||
|
||||
// AdjustChunkPagerank atomically updates the legacy pagerank_fea column.
|
||||
func (e *Engine) AdjustChunkPagerank(ctx context.Context, tableName, chunkID, datasetID string, delta, minWeight, maxWeight float64) error {
|
||||
if err := validateIdentifier(tableName); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := e.db.ExecContext(ctx, "UPDATE "+quoteIdentifier(tableName)+
|
||||
" SET pagerank_fea = GREATEST(?, LEAST(?, COALESCE(pagerank_fea, 0) + ?)) WHERE id = ? AND kb_id = ?",
|
||||
minWeight, maxWeight, delta, chunkID, datasetID)
|
||||
return err
|
||||
}
|
||||
|
||||
func envBool(name string, defaultValue bool) bool {
|
||||
raw := strings.ToLower(strings.TrimSpace(os.Getenv(name)))
|
||||
if raw == "" {
|
||||
return defaultValue
|
||||
}
|
||||
return raw == "true" || raw == "1" || raw == "yes" || raw == "y"
|
||||
}
|
||||
|
||||
func envInt64(name string, defaultValue int64) int64 {
|
||||
raw := strings.TrimSpace(os.Getenv(name))
|
||||
if raw == "" {
|
||||
return defaultValue
|
||||
}
|
||||
value, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil || value <= 0 {
|
||||
return defaultValue
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func envInt(name string, defaultValue int) int {
|
||||
raw := strings.TrimSpace(os.Getenv(name))
|
||||
if raw == "" {
|
||||
return defaultValue
|
||||
}
|
||||
value, err := strconv.Atoi(raw)
|
||||
if err != nil || value < 0 {
|
||||
return defaultValue
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func compareVersions(left, right string) int {
|
||||
parse := func(raw string) []int {
|
||||
parts := strings.FieldsFunc(raw, func(r rune) bool { return r < '0' || r > '9' })
|
||||
values := make([]int, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
value, err := strconv.Atoi(part)
|
||||
if err == nil {
|
||||
values = append(values, value)
|
||||
}
|
||||
}
|
||||
return values
|
||||
}
|
||||
a, b := parse(left), parse(right)
|
||||
for i := 0; i < max(len(a), len(b)); i++ {
|
||||
var av, bv int
|
||||
if i < len(a) {
|
||||
av = a[i]
|
||||
}
|
||||
if i < len(b) {
|
||||
bv = b[i]
|
||||
}
|
||||
if av < bv {
|
||||
return -1
|
||||
}
|
||||
if av > bv {
|
||||
return 1
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
128
internal/engine/oceanbase/client_test.go
Normal file
128
internal/engine/oceanbase/client_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
393
internal/engine/oceanbase/codec.go
Normal file
393
internal/engine/oceanbase/codec.go
Normal 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)
|
||||
}
|
||||
344
internal/engine/oceanbase/compatibility_test.go
Normal file
344
internal/engine/oceanbase/compatibility_test.go
Normal 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)
|
||||
}
|
||||
350
internal/engine/oceanbase/crud.go
Normal file
350
internal/engine/oceanbase/crud.go
Normal 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
|
||||
}
|
||||
153
internal/engine/oceanbase/crud_test.go
Normal file
153
internal/engine/oceanbase/crud_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
144
internal/engine/oceanbase/document.go
Normal file
144
internal/engine/oceanbase/document.go
Normal 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
|
||||
}
|
||||
59
internal/engine/oceanbase/document_test.go
Normal file
59
internal/engine/oceanbase/document_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
234
internal/engine/oceanbase/filter.go
Normal file
234
internal/engine/oceanbase/filter.go
Normal 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
|
||||
}
|
||||
200
internal/engine/oceanbase/helpers.go
Normal file
200
internal/engine/oceanbase/helpers.go
Normal 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
|
||||
}
|
||||
}
|
||||
83
internal/engine/oceanbase/helpers_test.go
Normal file
83
internal/engine/oceanbase/helpers_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
199
internal/engine/oceanbase/highlight.go
Normal file
199
internal/engine/oceanbase/highlight.go
Normal 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()
|
||||
}
|
||||
44
internal/engine/oceanbase/highlight_test.go
Normal file
44
internal/engine/oceanbase/highlight_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
229
internal/engine/oceanbase/integration_test.go
Normal file
229
internal/engine/oceanbase/integration_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
406
internal/engine/oceanbase/metadata.go
Normal file
406
internal/engine/oceanbase/metadata.go
Normal 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
|
||||
}
|
||||
232
internal/engine/oceanbase/rows.go
Normal file
232
internal/engine/oceanbase/rows.go
Normal 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
|
||||
}
|
||||
531
internal/engine/oceanbase/schema.go
Normal file
531
internal/engine/oceanbase/schema.go
Normal 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")
|
||||
}
|
||||
162
internal/engine/oceanbase/schema_test.go
Normal file
162
internal/engine/oceanbase/schema_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
840
internal/engine/oceanbase/search.go
Normal file
840
internal/engine/oceanbase/search.go
Normal 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
|
||||
}
|
||||
105
internal/engine/oceanbase/search_test.go
Normal file
105
internal/engine/oceanbase/search_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
51
internal/engine/oceanbase/sql.go
Normal file
51
internal/engine/oceanbase/sql.go
Normal 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)
|
||||
}
|
||||
144
internal/engine/oceanbase/sql_test.go
Normal file
144
internal/engine/oceanbase/sql_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user