feat(vastbase): add Vastbase G100 document engine support with go (#20419)

This commit is contained in:
kk710
2026-10-08 20:02:37 +08:00
committed by GitHub
parent e73aa1daca
commit ef150bc2cd
42 changed files with 8627 additions and 7 deletions
+8
View File
@@ -66,6 +66,14 @@ oceanbase:
password: 'infini_rag_flow'
host: 'localhost'
port: 2881
vastbase:
host: 'localhost'
port: 5432
user: 'ragflow'
password: 'Infini_Rag@123'
db_name: 'ragflow'
dbcompatibility: 'B'
ssl_mode: '' # empty: engine picks by host (disable on loopback, verify-full otherwise)
kvrocks:
db: 1
username: ''
+39
View File
@@ -32,6 +32,7 @@
# - `opensearch` (https://github.com/opensearch-project/OpenSearch)
# - `seekdb` (https://github.com/oceanbase/seekdb)
# - `gaussdb` GaussDB Centralized / Distributed DocEngine
# - `vastbase` Vastbase G100 (PG-compatible; B or PG dbcompatibility)
DOC_ENGINE=${DOC_ENGINE:-elasticsearch}
# The business metadata database type used by DB_TYPE.
@@ -153,6 +154,44 @@ SEEKDB_DOC_DBNAME=ragflow_doc
# SeekDB memory limit
SEEKDB_MEMORY_LIMIT=2G
# The hostname where the Vastbase service is exposed
VB_HOST=vastbase
# The port used to connect to Vastbase from the RAGFlow container.
# Change this if you use external Vastbase.
VB_PORT=5432
# The port used to expose the Vastbase service to the host machine,
# allowing EXTERNAL access to Vastbase running inside the Docker container.
EXPOSE_VB_PORT=5432
# The host address the exposed port binds to. Defaults to loopback only:
# the database ships default credentials and RAGFlow itself reaches Vastbase
# over the internal docker network. Exposing it to other machines is a
# deliberate choice — set 0.0.0.0 or a specific interface IP here.
# VB_BIND_ADDR=127.0.0.1
# The username for Vastbase
VB_USERNAME=ragflow
# The password for Vastbase.
# Must contain uppercase and lowercase letters, digits, and special
# characters, with a minimum length of 8 (Vastbase password policy).
# WARNING: Change this for production!
VB_PASSWORD=Infini_Rag@123
# The doc database of the Vastbase service to use
VB_VEC_DBNAME=ragflow
# The compatibility mode of Vastbase.
# Available options:
# - `B` (default): MySQL-compatible mode
# - `PG`: PostgreSQL-compatible mode
# Passed to the image's vb_initdb at first initialization and fixed for
# the instance afterwards; changing it on an existing data volume has no
# effect.
VB_DBCOMPATIBILITY=B
# TLS mode for the Vastbase connection (lib/pq sslmode).
# `disable` is the explicit plaintext opt-in for this private container
# network; the bundled container serves no TLS, so keep it set there.
# Pointing RAGFlow at an external Vastbase over an untrusted network
# should use require/verify-ca/verify-full instead — when ssl_mode is
# left unset, non-loopback hosts default to verify-full.
VB_SSL_MODE=disable
# The password for MySQL.
# WARNING: Change this for production!
MYSQL_PASSWORD=infini_rag_flow
+32
View File
@@ -176,6 +176,36 @@ services:
- ragflow
restart: unless-stopped
vastbase:
profiles:
- vastbase
image: ${VB_IMAGE:-vastbase-g100-single-v1.4-3.0.9.31338-linux-x86_64:20260513112809}
env_file: .env
environment:
- VB_PASSWORD=${VB_PASSWORD}
- VB_USERNAME=${VB_USERNAME}
- VB_VEC_DBNAME=${VB_VEC_DBNAME:-ragflow}
- VB_DBCOMPATIBILITY=${VB_DBCOMPATIBILITY:-B}
# Bind the published port to loopback by default: the database ships
# default credentials, and RAGFlow reaches it over the shared `ragflow`
# network. Remote access is an explicit opt-in through VB_BIND_ADDR
# (e.g. 0.0.0.0 or a specific interface IP) in docker/.env.
ports:
- ${VB_BIND_ADDR:-127.0.0.1}:${EXPOSE_VB_PORT:-5432}:5432
volumes:
- vastbase_data:/home/vastbase/data
- ./vastbase/init:/docker-entrypoint-initdb.d
- ./vastbase/license:/home/vastbase/vastbase/lic
mem_limit: ${MEM_LIMIT}
networks:
- ragflow
restart: unless-stopped
healthcheck:
test: ["CMD-SHELL", "vsql -c 'SELECT 1'"]
interval: 10s
timeout: 10s
retries: 120
sandbox-executor-manager:
profiles:
- sandbox
@@ -475,6 +505,8 @@ volumes:
driver: local
seekdb_data:
driver: local
vastbase_data:
driver: local
mysql_data:
driver: local
minio_data:
+11
View File
@@ -91,6 +91,17 @@ seekdb:
password: '${SEEKDB_PASSWORD:-infini_rag_flow}'
host: '${SEEKDB_HOST:-seekdb}'
port: ${SEEKDB_PORT:-2881}
vastbase:
host: '${VB_HOST:-vastbase}'
port: ${VB_PORT:-5432}
user: '${VB_USERNAME:-ragflow}'
password: '${VB_PASSWORD:-Infini_Rag@123}'
db_name: '${VB_VEC_DBNAME:-ragflow}'
dbcompatibility: '${VB_DBCOMPATIBILITY:-B}'
# No disable fallback: an empty ssl_mode reaches the engine's host policy
# (disable on loopback, verify-full otherwise). docker/.env sets disable
# explicitly for the bundled private-network deployment.
ssl_mode: '${VB_SSL_MODE:-}'
kvrocks:
db: 1
username: '${KVROCKS_USERNAME:-}'
+67
View File
@@ -0,0 +1,67 @@
#!/bin/bash
set -e
# 01_init.sh — Vastbase first-time initialization
# Reads VB_USERNAME and VB_PASSWORD from environment variables
# (injected by docker-compose.yml -> .env)
# Creates user and databases, grants all privileges.
VB_USER="${VB_USERNAME:-ragflow}"
# Default must satisfy Vastbase password policy
# (upper + lower + digit + special chars, >= 8 chars)
VB_PASS="${VB_PASSWORD:-Infini_Rag@123}"
VB_VEC_DB="${VB_VEC_DBNAME:-ragflow}"
# SQL encoding — shell quoting and the temp file below perform no SQL
# escaping, so encode here: string literals double embedded single quotes,
# identifiers double embedded double quotes. The SET below pins
# standard_conforming_strings=on so backslashes stay literal, which is what
# the quote-doubling above assumes (PostgreSQL docs, String Constants).
VB_USER_LIT=${VB_USER//\'/\'\'}
VB_USER_IDENT=${VB_USER//\"/\"\"}
VB_PASS_LIT=${VB_PASS//\'/\'\'}
VB_DB_IDENT=${VB_VEC_DB//\"/\"\"}
# ── Step 1: Create user ──────────────────────────────────────────────────
# Role existence is probed with a plain query first so the CREATE can be a
# bare statement in the temp file: no DO $$ block is involved, so a password
# containing "$$" cannot terminate a dollar-quoted body early (doubling
# single quotes is complete escaping for a plain SQL string literal, but it
# does not protect a dollar-quoted one).
SQL_FILE="$(mktemp /tmp/init_vastbase.XXXXXX.sql)"
trap 'rm -f "${SQL_FILE}"' EXIT
if vsql -tA -c "SELECT 1 FROM pg_roles WHERE rolname = '${VB_USER_LIT}'" | grep -qx 1; then
echo "User ${VB_USER} already exists"
else
# Contains the password in cleartext: mktemp gives 0600 perms, and the
# EXIT trap removes it even when `set -e` aborts on a failed vsql call.
cat > "${SQL_FILE}" << SQL
SET standard_conforming_strings = on;
CREATE ROLE "${VB_USER_IDENT}" LOGIN PASSWORD '${VB_PASS_LIT}';
SQL
echo "Creating user ${VB_USER}..."
vsql -f "${SQL_FILE}"
fi
# ── Step 2: Create databases (outside DO block — CREATE DATABASE requires
# a transaction commit and cannot run inside a PL/pgSQL block) ──
echo "Creating database ${VB_VEC_DB}..."
vsql -c "CREATE DATABASE \"${VB_DB_IDENT}\" OWNER \"${VB_USER_IDENT}\"" 2>/dev/null \
|| echo " Database '${VB_VEC_DB}' already exists"
# ── Step 3: Grant database-level privileges ──────────────────────────────
echo "Granting privileges on ${VB_VEC_DB}..."
vsql -c "GRANT ALL PRIVILEGES ON DATABASE \"${VB_DB_IDENT}\" TO \"${VB_USER_IDENT}\";"
# ── Step 4: Grant schema permissions (must connect to each database) ─────
echo "Granting schema permissions on ${VB_VEC_DB}..."
vsql -d "${VB_VEC_DB}" -c "GRANT ALL ON SCHEMA public TO \"${VB_USER_IDENT}\";"
vsql -d "${VB_VEC_DB}" -c "ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO \"${VB_USER_IDENT}\";"
vsql -d "${VB_VEC_DB}" -c "ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO \"${VB_USER_IDENT}\";"
echo "ragflow initialization complete."
+19
View File
@@ -0,0 +1,19 @@
#!/bin/bash
set -e
# 02_bcompat.sh — Vastbase B-mode compatibility options
# Must be run as the vastbase/vb user (who owns $PGDATA)
# The image's entrypoint passes VB_DBCOMPATIBILITY to vb_initdb on first
# initialization, which fixes the instance-wide compatibility mode. The
# GUCs below are B-mode behavior tweaks and must not run on other modes.
VB_MODE="${VB_DBCOMPATIBILITY:-B}"
if [ "${VB_MODE}" != "B" ]; then
echo "VB_DBCOMPATIBILITY=${VB_MODE}: skipping B-mode compatibility options."
exit 0
fi
echo "Setting B-mode compatibility options..."
gs_guc reload -D "$PGDATA" -c "b_format_behavior_compat_options = 'set_keyword_as_colname, show_attalias_as_colname, pg_todate_format'"
echo "Done. Compatibility options applied."
+5
View File
@@ -0,0 +1,5 @@
# License files must never be committed — provide them manually (see README.md).
*
!README.md
!.gitignore
!.gitkeep
View File
+8
View File
@@ -0,0 +1,8 @@
Place the Vastbase G100 license file(s) issued for your deployment in this
directory. `docker/docker-compose-base.yml` bind-mounts it into the container
at `/home/vastbase/vastbase/lic`:
./vastbase/license:/home/vastbase/vastbase/lic
The directory is committed with only this README (and a .gitkeep) so the
mount source always exists; license files themselves must never be committed.
+11
View File
@@ -35,6 +35,7 @@ const (
EngineOceanBase EngineType = "oceanbase"
EngineSeekDB EngineType = "seekdb"
EngineSereneDB EngineType = "serenedb"
EngineVastbase EngineType = "vastbase"
)
// DocEngine document storage engine interface
@@ -108,6 +109,16 @@ func IsOceanBaseFamily(engineName string) bool {
return engineName == string(EngineOceanBase) || engineName == string(EngineSeekDB)
}
// IsVastbase reports whether a configured engine is the Vastbase G100
// implementation. Vastbase is PostgreSQL-wire based and shares several
// service-layer behaviors with the OceanBase family (local rerank, chunk
// vectors in select fields, shared-table memory lifecycle) but is a separate
// SQL implementation, so it keeps its own predicate instead of joining
// IsOceanBaseFamily.
func IsVastbase(engineName string) bool {
return engineName == string(EngineVastbase)
}
type MessageQueue interface {
Init() error
Type() string
+3
View File
@@ -27,6 +27,7 @@ import (
"ragflow/internal/engine/nats"
"ragflow/internal/engine/oceanbase"
"ragflow/internal/engine/serenedb"
"ragflow/internal/engine/vastbase"
"ragflow/internal/server"
"ragflow/internal/tokenizer"
@@ -63,6 +64,8 @@ func InitDocEngine(ctx context.Context) error {
}
case "serenedb":
globalEngine, err = serenedb.NewEngine(globalConfig.GetSereneDBConfig())
case "vastbase":
globalEngine, err = vastbase.NewEngine(globalConfig.GetVastbaseConfig())
default:
err = fmt.Errorf("unsupported doc engine type: %s", engineType)
}
+332
View File
@@ -0,0 +1,332 @@
//
// 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 vastbase implements the Vastbase G100 document engine over its
// PostgreSQL-compatible wire protocol (database/sql + lib/pq). Chunks live in
// shared per-baseName tables with kb_id/memory_id as row-level discriminators;
// vector columns use the floatvector type with graph_index (HNSW) indexes.
package vastbase
import (
"context"
"database/sql"
"fmt"
"os"
"regexp"
"strconv"
"strings"
"time"
"ragflow/internal/common"
"ragflow/internal/server/config"
_ "github.com/lib/pq"
"go.uber.org/zap"
)
const (
connectionAttempts = 24
connectionRetryInterval = 5 * time.Second
defaultPoolSize = 50
defaultMaxOverflow = 100
connMaxLifetime = 30 * time.Minute
)
// connectionTimeout caps the whole retry loop: 24 attempts at the 5 s cadence
// is the Python connector's two-minute contract, and the cap keeps a server
// that accepts but stalls pings from stretching it to 24 x 10 s.
const connectionTimeout = 2 * time.Minute
// initializeTimeout bounds the post-ping initialization queries, so a server
// that answers PingContext but stalls on SELECT vb_version() cannot block
// startup forever.
const initializeTimeout = 30 * time.Second
// Compatibility modes accepted for VastbaseConfig.DBCompatibility.
const (
compatPG = "PG"
compatB = "B"
)
var (
// lib/pq's parseOpts tolerates whitespace around "=" and treats a
// backslash-escaped space as part of an unquoted value, so the value
// alternation must cover both forms or the redaction leaks.
dsnKeywordPwRe = regexp.MustCompile(`(password\s*=\s*)('(?:[^'\\]|\\.)*'|(?:\\.|[^\s\\])+)`)
dsnURLPwRe = regexp.MustCompile(`(://[^:/@\s]+:)[^@\s]+(@)`)
)
// redactDSN hides the password in both the lib/pq keyword form (password=...)
// and the URL form (scheme://user:password@host) used by VASTBASE_DSN.
func redactDSN(dsn string) string {
dsn = dsnKeywordPwRe.ReplaceAllString(dsn, "${1}***")
dsn = dsnURLPwRe.ReplaceAllString(dsn, "${1}***${2}")
return dsn
}
// quoteDSNValue wraps a lib/pq keyword-DSN value in single quotes, escaping
// backslashes and quotes, so hosts/users/passwords with spaces or special
// characters do not break the DSN.
func quoteDSNValue(v string) string {
r := strings.ReplaceAll(v, `\`, `\\`)
r = strings.ReplaceAll(r, `'`, `\'`)
return "'" + r + "'"
}
// isLocalDBHost reports whether the host is loopback or a unix-socket path —
// the only addresses allowed to default to plaintext transport.
func isLocalDBHost(host string) bool {
switch strings.ToLower(strings.TrimSpace(host)) {
case "localhost", "127.0.0.1", "::1", "[::1]":
return true
}
return strings.HasPrefix(host, "/")
}
// Engine is the Vastbase G100 document engine backed by database/sql.
type Engine struct {
db *sql.DB
dbName string
engineType string
dbCompatibility string
dsnSafe string
}
// NewEngine constructs the engine from the Vastbase config. It retries the
// initial connection for up to two minutes (the Python connector contract for
// waiting on a freshly started container).
func NewEngine(cfg config.VastbaseConfig) (*Engine, error) {
dsn := buildDSN(cfg)
db, err := sql.Open("postgres", dsn)
if err != nil {
return nil, fmt.Errorf("vastbase: open: %w", err)
}
poolSize := envInt("VB_POOL_SIZE", defaultPoolSize)
maxOverflow := envInt("VB_MAX_OVERFLOW", defaultMaxOverflow)
db.SetMaxOpenConns(poolSize + maxOverflow)
db.SetMaxIdleConns(poolSize)
db.SetConnMaxLifetime(connMaxLifetime)
e := &Engine{
db: db,
dbName: cfg.DBName,
engineType: "vastbase",
dbCompatibility: cfg.DBCompatibility,
dsnSafe: redactDSN(dsn),
}
var lastErr error
connectCtx, cancelConnect := context.WithTimeout(context.Background(), connectionTimeout)
defer cancelConnect()
for attempt := 0; attempt < connectionAttempts; attempt++ {
if connectCtx.Err() != nil {
break
}
pingCtx, cancel := context.WithTimeout(connectCtx, 10*time.Second)
lastErr = db.PingContext(pingCtx)
cancel()
if lastErr == nil {
break
}
common.Warn("Vastbase is not healthy yet, retrying",
zap.Int("attempt", attempt+1),
zap.Int("max_attempts", connectionAttempts),
zap.String("dsn", e.dsnSafe),
zap.Error(lastErr))
if attempt+1 < connectionAttempts {
select {
case <-time.After(connectionRetryInterval):
case <-connectCtx.Done():
}
}
}
if lastErr != nil {
_ = db.Close()
return nil, fmt.Errorf("vastbase: connect %s: %w", e.dsnSafe, lastErr)
}
initCtx, cancelInit := context.WithTimeout(context.Background(), initializeTimeout)
defer cancelInit()
if err := e.initialize(initCtx); err != nil {
_ = db.Close()
return nil, err
}
common.Info("Vastbase engine initialized",
zap.String("dsn", e.dsnSafe),
zap.String("dbcompatibility", e.dbCompatibility))
return e, nil
}
// newEngineWithDB builds an engine around an already-open database handle,
// for tests that drive the SQL layer with a mock driver.
func newEngineWithDB(cfg config.VastbaseConfig, db *sql.DB) *Engine {
return &Engine{
db: db,
dbName: cfg.DBName,
engineType: "vastbase",
dbCompatibility: cfg.DBCompatibility,
dsnSafe: "redacted",
}
}
// buildDSN assembles a lib/pq keyword DSN. VASTBASE_DSN overrides everything;
// otherwise config values fill in. An unset ssl_mode follows the transport
// policy: plaintext (disable) only for loopback or unix-socket hosts, and
// certificate-verified TLS (verify-full) for every other TCP host — remote
// deployments must not silently fall back to cleartext. Unrecognized keywords
// such as statement_timeout are forwarded to the server as startup runtime
// parameters.
func buildDSN(cfg config.VastbaseConfig) string {
if env := os.Getenv("VASTBASE_DSN"); env != "" {
return env
}
host, port, user, password, dbName := "vastbase", 5432, "ragflow", "", "ragflow"
sslMode := ""
if cfg.Host != "" {
host = cfg.Host
}
if cfg.Port != 0 {
port = cfg.Port
}
if cfg.User != "" {
user = cfg.User
}
password = cfg.Password
if cfg.DBName != "" {
dbName = cfg.DBName
}
if cfg.SSLMode != "" {
sslMode = cfg.SSLMode
}
if sslMode == "" {
if isLocalDBHost(host) {
sslMode = "disable"
} else {
sslMode = "verify-full"
}
}
parts := []string{
fmt.Sprintf("host=%s", quoteDSNValue(host)),
fmt.Sprintf("port=%d", port),
fmt.Sprintf("user=%s", quoteDSNValue(user)),
fmt.Sprintf("dbname=%s", quoteDSNValue(dbName)),
fmt.Sprintf("sslmode=%s", quoteDSNValue(sslMode)),
}
if password != "" {
parts = append(parts, fmt.Sprintf("password=%s", quoteDSNValue(password)))
}
// No keepalive keywords here: lib/pq forwards keywords it does not
// recognize to the server as startup parameters, and Vastbase (like
// vanilla PostgreSQL) rejects `keepalives*` as unknown GUCs. TCP keepalive
// stays on regardless — lib/pq dials through a zero-value net.Dialer,
// which enables keepalives at the OS default interval.
parts = append(parts,
"connect_timeout=30",
"statement_timeout=300000",
)
return strings.Join(parts, " ")
}
// initialize verifies the server answers and fails when the configured
// compatibility mode does not match the instance's datcompatibility. A B-mode
// distance operator or bm25 predicate issued against a PG instance (and vice
// versa) fails at query time, so a confirmed mismatch must stop startup
// before any document operation runs.
func (e *Engine) initialize(ctx context.Context) error {
var version string
if err := e.db.QueryRowContext(ctx, "SELECT vb_version()").Scan(&version); err != nil {
return fmt.Errorf("vastbase: get version: %w", err)
}
common.Info("Vastbase server version", zap.String("version", version))
var datCompat string
if err := e.db.QueryRowContext(ctx,
"SELECT datcompatibility FROM pg_database WHERE datname = current_database()").Scan(&datCompat); err != nil {
return fmt.Errorf("vastbase: read datcompatibility: %w", err)
}
expected := compatPG
if e.dbCompatibility == compatB {
expected = compatB
}
datCompat = strings.TrimSpace(datCompat)
if !strings.EqualFold(datCompat, expected) {
return fmt.Errorf("vastbase: dbcompatibility mismatch: instance reports %q but the engine is configured for %q; "+
"set dbcompatibility to the database's mode (VB_DBCOMPATIBILITY in docker deployments) or recreate the database, "+
"because distance operators and full-text indexes differ between modes", datCompat, e.dbCompatibility)
}
return nil
}
// distanceOp returns the cosine-distance operator for the compatibility mode.
func (e *Engine) distanceOp() string {
if e.dbCompatibility == compatB {
return "<+>"
}
return "<=>"
}
// GetType returns the engine type string.
func (e *Engine) GetType() string {
return e.engineType
}
// SupportsPageRank reports dataset-level pagerank support: pagerank_fea is a
// real column folded into every scored query and adjusted atomically by
// AdjustChunkPagerank.
func (e *Engine) SupportsPageRank() bool {
return true
}
// Ping verifies connectivity.
func (e *Engine) Ping(ctx context.Context) error {
return e.db.PingContext(ctx)
}
// Close releases the connection pool.
func (e *Engine) Close() error {
if e.db == nil {
return nil
}
return e.db.Close()
}
// AdjustChunkPagerank atomically clamps pagerank_fea into [minWeight,
// maxWeight] while applying delta, scoped to the dataset's rows in the shared
// chunk table.
func (e *Engine) AdjustChunkPagerank(ctx context.Context, indexName, chunkID, kbID string, delta, minWeight, maxWeight float64) error {
if err := validateIdentifier(indexName); err != nil {
return err
}
_, err := e.db.ExecContext(ctx, fmt.Sprintf(
"UPDATE %s SET pagerank_fea = GREATEST($1, LEAST($2, COALESCE(pagerank_fea, 0) + $3)) WHERE id = $4 AND kb_id = $5",
quoteIdent(indexName)), minWeight, maxWeight, delta, chunkID, kbID)
if err != nil {
return fmt.Errorf("vastbase: adjust pagerank: %w", err)
}
return nil
}
// envInt reads a positive integer environment variable, falling back to the
// default when unset, empty, or invalid.
func envInt(name string, fallback int) int {
if raw := strings.TrimSpace(os.Getenv(name)); raw != "" {
if parsed, err := strconv.Atoi(raw); err == nil && parsed > 0 {
return parsed
}
}
return fallback
}
+235
View File
@@ -0,0 +1,235 @@
//
// 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 vastbase
import (
"context"
"fmt"
"regexp"
"strings"
"testing"
"ragflow/internal/server/config"
"github.com/DATA-DOG/go-sqlmock"
)
// TestBuildDSNAssemblesKeywordForm pins the keyword DSN, including quoting
// and escaping of values with spaces and special characters.
func TestBuildDSNAssemblesKeywordForm(t *testing.T) {
t.Setenv("VASTBASE_DSN", "")
got := buildDSN(config.VastbaseConfig{
Host: "vb-host", Port: 15432, User: "rag user",
Password: "p@ss word'", DBName: "rag flow", SSLMode: "verify-full",
})
want := `host='vb-host' port=15432 user='rag user' dbname='rag flow' sslmode='verify-full' ` +
`password='p@ss word\'' connect_timeout=30 statement_timeout=300000`
if got != want {
t.Fatalf("DSN = %q", got)
}
}
// TestBuildDSNDefaultsAndOverride checks the host-based sslmode defaults,
// explicit overrides, and the VASTBASE_DSN escape hatch.
func TestBuildDSNDefaultsAndOverride(t *testing.T) {
t.Setenv("VASTBASE_DSN", "")
// Unset ssl_mode: the default host is a TCP service name, so the
// transport policy requires certificate-verified TLS.
got := buildDSN(config.VastbaseConfig{})
if !regexp.MustCompile(`^host='vastbase' port=5432 user='ragflow' dbname='ragflow' sslmode='verify-full' connect_timeout=30`).MatchString(got) {
t.Fatalf("default DSN = %q", got)
}
// No password configured: the keyword must be absent, not empty.
if regexp.MustCompile(`password=''`).MatchString(got) {
t.Fatalf("empty password must be omitted: %q", got)
}
// Loopback and unix sockets may default to plaintext.
for _, host := range []string{"localhost", "127.0.0.1"} {
if got := buildDSN(config.VastbaseConfig{Host: host}); !regexp.MustCompile(`sslmode='disable'`).MatchString(got) {
t.Fatalf("local host %s must default to disable: %q", host, got)
}
}
// An explicit ssl_mode is honored for any host.
if got := buildDSN(config.VastbaseConfig{SSLMode: "require"}); !regexp.MustCompile(`sslmode='require'`).MatchString(got) {
t.Fatalf("explicit ssl_mode must be honored: %q", got)
}
t.Setenv("VASTBASE_DSN", "postgres://u:secret@h:5432/db")
if got := buildDSN(config.VastbaseConfig{Host: "ignored"}); got != "postgres://u:secret@h:5432/db" {
t.Fatalf("VASTBASE_DSN override = %q", got)
}
}
// TestRedactDSN checks password redaction across the keyword and URL DSN
// forms lib/pq accepts, including whitespace and escaped values.
func TestRedactDSN(t *testing.T) {
keyword := redactDSN(`host='h' password='s3cret' user='u'`)
if keyword != `host='h' password=*** user='u'` {
t.Fatalf("keyword redaction = %q", keyword)
}
// lib/pq accepts whitespace around "=" — those DSNs must redact too.
spaced := redactDSN(`host=h user=u password = 'secret' dbname=d`)
if spaced != `host=h user=u password = *** dbname=d` {
t.Fatalf("spaced keyword redaction = %q", spaced)
}
// Escaped characters survive in unquoted values, including escaped
// spaces: the whole token is the password, not just its first word.
escaped := redactDSN(`password=se\ cret dbname=d`)
if escaped != `password=*** dbname=d` {
t.Fatalf("escaped unquoted redaction = %q", escaped)
}
escapedQuote := redactDSN(`host=h password='a\'b' port=5432`)
if escapedQuote != `host=h password=*** port=5432` {
t.Fatalf("escaped quoted redaction = %q", escapedQuote)
}
url := redactDSN("postgres://ragflow:Infini_Rag%40123@vastbase:5432/ragflow")
if url != "postgres://ragflow:***@vastbase:5432/ragflow" {
t.Fatalf("url redaction = %q", url)
}
}
// TestDistanceOpPerCompatibility checks the cosine-distance operator per
// compatibility mode and the PG fallback.
func TestDistanceOpPerCompatibility(t *testing.T) {
if op := (&Engine{dbCompatibility: compatB}).distanceOp(); op != "<+>" {
t.Fatalf("B-mode operator = %q", op)
}
if op := (&Engine{dbCompatibility: compatPG}).distanceOp(); op != "<=>" {
t.Fatalf("PG-mode operator = %q", op)
}
// Anything unrecognized falls back to the PG operator.
if op := (&Engine{dbCompatibility: "weird"}).distanceOp(); op != "<=>" {
t.Fatalf("fallback operator = %q", op)
}
}
// TestInitializeCompatibilityGate checks the datcompatibility gate: a
// matching instance initializes, a mismatched one fails startup before any
// document operation can run, and a failing probe surfaces instead of being
// silently swallowed.
func TestInitializeCompatibilityGate(t *testing.T) {
newCase := func(t *testing.T, configured string) (*Engine, sqlmock.Sqlmock, func()) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: configured}, db)
mock.ExpectQuery("SELECT vb_version()").
WillReturnRows(sqlmock.NewRows([]string{"vb_version"}).AddRow("G100 v1.4"))
return engine, mock, func() { db.Close() }
}
t.Run("matching mode initializes", func(t *testing.T) {
engine, mock, done := newCase(t, compatB)
defer done()
mock.ExpectQuery("SELECT datcompatibility FROM pg_database").
WillReturnRows(sqlmock.NewRows([]string{"datcompatibility"}).AddRow("B"))
if err := engine.initialize(context.Background()); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
})
t.Run("mismatched mode fails startup", func(t *testing.T) {
engine, mock, done := newCase(t, compatPG)
defer done()
mock.ExpectQuery("SELECT datcompatibility FROM pg_database").
WillReturnRows(sqlmock.NewRows([]string{"datcompatibility"}).AddRow("B"))
err := engine.initialize(context.Background())
if err == nil || !strings.Contains(err.Error(), "dbcompatibility mismatch") {
t.Fatalf("err = %v", err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
})
t.Run("probe failure surfaces", func(t *testing.T) {
engine, mock, done := newCase(t, compatB)
defer done()
mock.ExpectQuery("SELECT datcompatibility FROM pg_database").
WillReturnError(fmt.Errorf("pg_database unreadable"))
err := engine.initialize(context.Background())
if err == nil || !strings.Contains(err.Error(), "read datcompatibility") {
t.Fatalf("err = %v", err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
})
}
// TestGetTypeAndSupportsPageRank checks the engine's type string and pagerank
// support flag.
func TestGetTypeAndSupportsPageRank(t *testing.T) {
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, nil)
if engine.GetType() != "vastbase" {
t.Fatalf("engine type = %q", engine.GetType())
}
if !engine.SupportsPageRank() {
t.Fatal("pagerank_fea is a real column; pagerank must be supported")
}
}
// TestAdjustChunkPagerankSQL pins the clamped pagerank UPDATE and the
// rejection of invalid table names.
func TestAdjustChunkPagerankSQL(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
mock.ExpectExec(regexp.QuoteMeta(
`UPDATE "t1" SET pagerank_fea = GREATEST($1, LEAST($2, COALESCE(pagerank_fea, 0) + $3)) WHERE id = $4 AND kb_id = $5`)).
WithArgs(0.0, 10.0, 0.5, "c1", "kb1").
WillReturnResult(sqlmock.NewResult(0, 1))
if err := engine.AdjustChunkPagerank(context.Background(), "t1", "c1", "kb1", 0.5, 0, 10); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
if err := engine.AdjustChunkPagerank(context.Background(), "bad;table", "c1", "kb1", 0.5, 0, 10); err == nil {
t.Fatal("invalid table names must be rejected")
}
}
// TestValidateIdentifier checks accepted and rejected identifier names.
func TestValidateIdentifier(t *testing.T) {
for _, name := range []string{"t1", "memory_x", "skill_abc", "ragflow_doc_meta_tenant-1"} {
if err := validateIdentifier(name); err != nil {
t.Fatalf("validateIdentifier(%q) = %v", name, err)
}
}
for _, name := range []string{"", `bad"quote`, "space name", "semi;colon"} {
if err := validateIdentifier(name); err == nil {
t.Fatalf("validateIdentifier(%q) must fail", name)
}
}
}
// TestQuoteIdentEscapesDoubleQuotes checks identifier quoting.
func TestQuoteIdentEscapesDoubleQuotes(t *testing.T) {
if got := quoteIdent(`we"ird`); got != `"we""ird"` {
t.Fatalf("quoteIdent = %q", got)
}
}
+441
View File
@@ -0,0 +1,441 @@
//
// 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 vastbase
import (
"encoding/json"
"fmt"
"math"
"reflect"
"regexp"
"sort"
"strconv"
"strings"
"ragflow/internal/tokenizer"
"github.com/lib/pq"
)
const keywordSeparator = "###"
var vectorColumnPattern = regexp.MustCompile(`^q_(\d+)_vec$`)
// integerArrayColumns are bound as PG integer[] via pq.Array.
var integerArrayColumns = map[string]bool{
"position_int": true, "page_num_int": true, "top_int": true,
}
var jsonTextColumns = map[string]bool{
"metadata": true, "extra": true, "meta_fields": true,
}
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",
}
var knownMemoryColumns = func() map[string]bool {
known := make(map[string]bool, len(memoryColumns))
for _, column := range memoryColumns {
known[column.name] = true
}
return known
}()
var knownSkillColumns = func() map[string]bool {
known := make(map[string]bool, len(skillColumns))
for _, column := range skillColumns {
known[column.name] = true
}
return known
}()
// fieldKeyword reports whether a field stores its list values ###-joined in a
// single text column: source_id and *_kwd fields, except docnm_kwd and
// knowledge_graph_kwd which store plain text.
func fieldKeyword(field string) bool {
if field == "source_id" {
return true
}
return strings.HasSuffix(field, "_kwd") && field != "docnm_kwd" && field != "knowledge_graph_kwd"
}
// normalizeChunk encodes a chunk document for insert/update. Unknown fields
// are kept as-is: ensureDynamicColumns has (or will) materialize them as real
// text columns — the ES dynamic-mapping equivalent. Unlike the OceanBase
// codec there is no extra folding and no whole-column padding; the table
// DEFAULTs fill the gaps so unwritten fields keep their ES-parity ”.
func normalizeChunk(document map[string]interface{}) (map[string]interface{}, error) {
result := make(map[string]interface{}, len(document))
for key, value := range document {
if vectorColumnPattern.MatchString(key) {
encoded, err := encodeVector(value)
if err != nil {
return nil, fmt.Errorf("encode %s: %w", key, err)
}
result[key] = encoded
continue
}
encoded, err := encodeColumnValue(key, value)
if err != nil {
return nil, fmt.Errorf("encode %s: %w", key, err)
}
result[key] = encoded
}
return result, nil
}
// encodeColumnValue converts one document value into its physical column
// representation: kb_id keeps only the first element of a list, integer
// arrays flatten, maps serialize as JSON, and keyword lists join.
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 integerArrayColumns[columnName] {
return encodeIntegerArray(columnName, value)
}
switch value.(type) {
case map[string]interface{}:
encoded, err := json.Marshal(value)
return string(encoded), err
}
if fieldKeyword(columnName) {
if values, ok := interfaceSlice(value); ok {
parts := make([]string, 0, len(values))
for _, item := range values {
parts = append(parts, stringValue(item))
}
return strings.Join(parts, keywordSeparator), nil
}
}
return value, nil
}
// encodeIntegerArray flattens position_int's two-dimensional [[page, x1, y1,
// x2, y2], ...] into the flat integer[] storage (regrouped on read) and binds
// page_num_int/top_int as plain integer arrays.
func encodeIntegerArray(columnName string, value interface{}) (interface{}, error) {
values, ok := interfaceSlice(value)
if !ok {
return nil, fmt.Errorf("expected integer array for %s, got %T", columnName, value)
}
flat := make([]int64, 0, len(values))
for _, item := range values {
if nested, ok := interfaceSlice(item); ok {
for _, inner := range nested {
number, ok := numberToInt64(inner)
if !ok {
return nil, fmt.Errorf("non-integer element in %s", columnName)
}
flat = append(flat, number)
}
continue
}
number, ok := numberToInt64(item)
if !ok {
return nil, fmt.Errorf("non-integer element in %s", columnName)
}
flat = append(flat, number)
}
return pq.Int64Array(flat), nil
}
// normalizeMemory maps a memory document onto the memory columns, encoding
// content_embed into its q_N_vec column and deriving the tokenized content
// lookup column.
func normalizeMemory(document map[string]interface{}, datasetID string) (map[string]interface{}, error) {
result := make(map[string]interface{}, len(memoryColumns)+1)
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 !knownMemoryColumns[key] {
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 stringValue(result["memory_id"]) == "" && datasetID != "" {
result["memory_id"] = datasetID
}
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
}
// tokenizeMemoryContent produces the fine-grained token string the memory
// full-text lookup queries, falling back to coarser inputs on tokenizer
// failure.
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
}
// normalizeSkill maps a skill document onto the skill columns, passing
// vector fields through encoded and dropping unknown keys.
func normalizeSkill(document map[string]interface{}, documentID string) (map[string]interface{}, error) {
result := make(map[string]interface{}, len(skillColumns)+1)
for key, value := range document {
if vectorColumnPattern.MatchString(key) {
encoded, err := encodeVector(value)
if err != nil {
return nil, err
}
result[key] = encoded
continue
}
if knownSkillColumns[key] {
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
}
// encodeUpdateValue converts one field for an UPDATE, applying the same
// physical representations as encodeColumnValue (vectors, memory status
// bools, JSON maps).
func encodeUpdateValue(kind, columnName string, value interface{}) (interface{}, error) {
if vectorColumnPattern.MatchString(columnName) {
return encodeVector(value)
}
if kind == "memory" {
if columnName == "status_int" {
if status, ok := value.(bool); ok {
if status {
return 1, nil
}
return 0, nil
}
}
return value, nil
}
return encodeColumnValue(columnName, value)
}
// encodeVector renders a vector as the floatvector text form "[0.1,0.2,...]",
// converting every component to float32 first — the format the Python
// psycopg2 connector wrote, kept byte-compatible for shared tables.
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 {
parts[i] = strconv.FormatFloat(float64(float32(number)), 'g', -1, 32)
}
return "[" + strings.Join(parts, ",") + "]", nil
}
// zeroVector is the placeholder written when a document lacks a vector column
// the table already has, mirroring the Python connector's zero backfill.
func zeroVector(dimension int) string {
return "[" + strings.Repeat("0,", dimension-1) + "0]"
}
// vectorDimension reports the embedding size a document carries, from its
// q_N_vec column or content_embed length.
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
}
// sortedColumns lists a document's columns in sorted order, keeping INSERT
// and UPDATE column ordering deterministic.
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
}
// interfaceSlice converts any slice or array value to []interface{}.
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
}
// floatSlice converts a value to []float64 when every element is numeric.
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
}
// numberToFloat coerces the numeric kinds database/sql and JSON decode to
// into a float64.
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
case string:
parsed, err := strconv.ParseFloat(number, 64)
return parsed, err == nil
default:
return 0, false
}
}
// numberToInt64 coerces the numeric kinds database/sql and JSON decode to
// into an int64.
func numberToInt64(value interface{}) (int64, bool) {
switch number := value.(type) {
case int:
return int64(number), true
case int32:
return int64(number), true
case int64:
return number, true
case float64:
return int64(number), true
case float32:
return int64(number), true
case json.Number:
parsed, err := number.Int64()
return parsed, err == nil
case string:
parsed, err := strconv.ParseInt(number, 10, 64)
return parsed, err == nil
default:
return 0, false
}
}
// stringValue renders a value as the empty string for nil, itself for a
// string, and fmt's default otherwise.
func stringValue(value interface{}) string {
if value == nil {
return ""
}
if text, ok := value.(string); ok {
return text
}
return fmt.Sprint(value)
}
+231
View File
@@ -0,0 +1,231 @@
//
// 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 vastbase
import (
"encoding/json"
"testing"
"github.com/lib/pq"
)
// TestNormalizeChunkEncodesPythonStorageFormats checks the physical
// encodings the Python connector's storage formats map to.
func TestNormalizeChunkEncodesPythonStorageFormats(t *testing.T) {
document := map[string]interface{}{
"id": "chunk-1",
"kb_id": []interface{}{"kb-1", "kb-2"},
"important_kwd": []interface{}{"alpha", "beta"},
"source_id": []interface{}{"doc-1"},
"docnm_kwd": "plain text title",
"metadata": map[string]interface{}{"_group_id": "g1"},
"position_int": []interface{}{[]interface{}{int64(1), 2, 3, 4, 5}, []interface{}{int64(6), 7, 8, 9, 10}},
"page_num_int": []interface{}{int64(1), int64(2)},
"q_3_vec": []interface{}{0.1, 0.2, 0.3},
"custom_field": "kept as dynamic column",
}
got, err := normalizeChunk(document)
if err != nil {
t.Fatal(err)
}
if got["kb_id"] != "kb-1" {
t.Fatalf("kb_id list must collapse to first element, got %#v", got["kb_id"])
}
if got["important_kwd"] != "alpha###beta" || got["source_id"] != "doc-1" {
t.Fatalf("keyword lists must ###-join, got %#v / %#v", got["important_kwd"], got["source_id"])
}
if got["docnm_kwd"] != "plain text title" {
t.Fatalf("docnm_kwd stores plain text, got %#v", got["docnm_kwd"])
}
var metadata map[string]interface{}
if err := json.Unmarshal([]byte(got["metadata"].(string)), &metadata); err != nil {
t.Fatal(err)
}
if metadata["_group_id"] != "g1" {
t.Fatalf("metadata map must JSON-encode, got %#v", got["metadata"])
}
position, ok := got["position_int"].(pq.Int64Array)
if !ok || len(position) != 10 || position[0] != 1 || position[9] != 10 {
t.Fatalf("position_int must flatten to int64[], got %#v", got["position_int"])
}
pages, ok := got["page_num_int"].(pq.Int64Array)
if !ok || len(pages) != 2 {
t.Fatalf("page_num_int must bind as int64[], got %#v", got["page_num_int"])
}
if got["q_3_vec"] != "[0.1,0.2,0.3]" {
t.Fatalf("vector text encoding changed: %#v", got["q_3_vec"])
}
if got["custom_field"] != "kept as dynamic column" {
t.Fatalf("unknown fields must pass through for dynamic columns, got %#v", got)
}
}
// TestNormalizeChunkKeepsZeroAndFalsyValues checks that zero-valued fields
// survive normalization instead of being dropped as absent.
func TestNormalizeChunkKeepsZeroAndFalsyValues(t *testing.T) {
got, err := normalizeChunk(map[string]interface{}{
"available_int": 0, "weight_flt": 0.0, "removed_kwd": "",
})
if err != nil {
t.Fatal(err)
}
if got["available_int"] != 0 || got["weight_flt"] != 0.0 || got["removed_kwd"] != "" {
t.Fatalf("zero values must survive encode (DEFAULT '' only fills gaps): %#v", got)
}
}
// TestNormalizeMemoryAliasesAndDefaults checks memory field aliasing and
// default filling.
func TestNormalizeMemoryAliasesAndDefaults(t *testing.T) {
got, err := normalizeMemory(map[string]interface{}{
"id": "memory-1_1", "message_id": "1",
"message_type": "raw", "status": false,
"content": "hello world",
"content_embed": []interface{}{0.25, 0.5},
}, "memory-1")
if err != nil {
t.Fatal(err)
}
if got["message_type_kwd"] != "raw" {
t.Fatalf("message_type must map to message_type_kwd: %#v", got)
}
if got["status_int"] != 0 {
t.Fatalf("status bool must encode as status_int 0/1, got %#v", got["status_int"])
}
if got["content_ltks"] != "hello world" {
t.Fatalf("content must map to content_ltks: %#v", got)
}
if got["q_2_vec"] != "[0.25,0.5]" {
t.Fatalf("content_embed must land in the dimension column: %#v", got["q_2_vec"])
}
if got["tokenized_content_ltks"] == "" {
t.Fatal("tokenized_content_ltks companion must be derived from content")
}
if got["memory_id"] != "memory-1" {
t.Fatalf("memory_id must default to the dataset id, got %#v", got["memory_id"])
}
if got["zone_id"] != 0 {
t.Fatalf("zone_id must default to 0, got %#v", got["zone_id"])
}
if _, present := got["forget_at_flt"]; present {
t.Fatal("memory tables carry no forget_at_flt column")
}
}
// TestNormalizeMemoryStatusDefaultsToOne checks the memory status default.
func TestNormalizeMemoryStatusDefaultsToOne(t *testing.T) {
got, err := normalizeMemory(map[string]interface{}{"id": "m1", "memory_id": "m"}, "m")
if err != nil {
t.Fatal(err)
}
if got["status_int"] != 1 {
t.Fatalf("status must default to 1 (valid), got %#v", got["status_int"])
}
}
// TestNormalizeSkillDerivesTokenFields checks that skill normalization
// derives the tokenized companion fields.
func TestNormalizeSkillDerivesTokenFields(t *testing.T) {
got, err := normalizeSkill(map[string]interface{}{
"skill_id": "s-1", "name": "pdf reader", "content": "opens files",
"q_2_vec": []interface{}{1, 0},
}, "")
if err != nil {
t.Fatal(err)
}
if got["skill_id"] != "s-1" {
t.Fatalf("skill_id must be kept, got %#v", got["skill_id"])
}
if got["name_tks"] == "" || got["content_tks"] == "" {
t.Fatalf("token companion fields must be derived: %#v", got)
}
if _, present := got["tags_tks"]; !present {
t.Fatalf("absent tags must still carry an (empty) token field: %#v", got)
}
if got["q_2_vec"] != "[1,0]" {
t.Fatalf("skill vector encoding changed: %#v", got["q_2_vec"])
}
}
// TestEncodeUpdateValueRoutesByKind checks UPDATE value encoding per row
// kind and column.
func TestEncodeUpdateValueRoutesByKind(t *testing.T) {
vector, err := encodeUpdateValue("chunk", "q_2_vec", []interface{}{0.5, -0.5})
if err != nil || vector != "[0.5,-0.5]" {
t.Fatalf("vector update must encode text, got %#v err %v", vector, err)
}
status, err := encodeUpdateValue("memory", "status_int", true)
if err != nil || status != 1 {
t.Fatalf("memory status update must encode bool as 1, got %#v err %v", status, err)
}
joined, err := encodeUpdateValue("chunk", "important_kwd", []interface{}{"a", "b"})
if err != nil || joined != "a###b" {
t.Fatalf("chunk update must keep ### joining, got %#v err %v", joined, err)
}
raw, err := encodeUpdateValue("memory", "content_ltks", "raw text")
if err != nil || raw != "raw text" {
t.Fatalf("memory text update must pass through, got %#v err %v", raw, err)
}
}
// TestVectorEncodingIsFloat32Precision checks that vector text round-trips
// through float32 precision, matching the Python connector.
func TestVectorEncodingIsFloat32Precision(t *testing.T) {
// Python wrote np.float32 components; values beyond float32 precision
// must round to the float32 rendering so shared tables stay byte-stable.
encoded, err := encodeVector([]interface{}{1.0 / 3.0, 1e-40, 0.1})
if err != nil {
t.Fatal(err)
}
if encoded != "[0.33333334,1e-40,0.1]" {
t.Fatalf("float32 rendering changed: %q", encoded)
}
if _, err := encodeVector("not a vector"); err == nil {
t.Fatal("non-numeric vector must error")
}
}
// TestZeroVectorAndDimensionInference checks the zero-vector placeholder
// text and dimension inference from documents.
func TestZeroVectorAndDimensionInference(t *testing.T) {
if zeroVector(1) != "[0]" || zeroVector(3) != "[0,0,0]" {
t.Fatalf("zeroVector format changed: %q %q", zeroVector(1), zeroVector(3))
}
if d := vectorDimension(map[string]interface{}{"q_1024_vec": []interface{}{0.1}}); d != 1024 {
t.Fatalf("column-name dimension inference = %d", d)
}
if d := vectorDimension(map[string]interface{}{"content_embed": []interface{}{1, 2, 3, 4}}); d != 4 {
t.Fatalf("content_embed dimension inference = %d", d)
}
if d := vectorDimension(map[string]interface{}{"id": "x"}); d != 0 {
t.Fatalf("absent vector must infer 0, got %d", d)
}
}
// TestFieldKeywordClassification checks which columns count as
// keyword-joined fields.
func TestFieldKeywordClassification(t *testing.T) {
for field, want := range map[string]bool{
"source_id": true, "important_kwd": true, "entity_type_kwd": true,
"docnm_kwd": false, "knowledge_graph_kwd": false, "title_tks": false,
"content_with_weight": false, "q_3_vec": false,
} {
if got := fieldKeyword(field); got != want {
t.Fatalf("fieldKeyword(%q) = %v, want %v", field, got, want)
}
}
}
@@ -0,0 +1,194 @@
//
// 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 vastbase
import (
"reflect"
"testing"
"github.com/lib/pq"
)
// These tests lock the byte-level storage contract the Python psycopg2
// connector established on the vb-dev-1 branch. Tables written by either
// implementation must be readable by the other, so the wire formats below are
// frozen: ### joined keyword text, floatvector text rendered at float32
// precision, position_int stored as a flat integer[] regrouped in fives, and
// JSON object columns stored as JSON text.
func TestKeywordRoundTripMatchesPythonFormat(t *testing.T) {
encoded, err := encodeColumnValue("important_kwd", []interface{}{"政策", "检索", "ragflow"})
if err != nil {
t.Fatal(err)
}
if encoded != "政策###检索###ragflow" {
t.Fatalf("python ### join = %#v", encoded)
}
decoded := decodeColumnValue("chunk", "important_kwd", encoded)
if !reflect.DeepEqual(decoded, []string{"政策", "检索", "ragflow"}) {
t.Fatalf("python ### split = %#v", decoded)
}
// A single-element list round-trips without separators.
encoded, _ = encodeColumnValue("source_id", []interface{}{"doc-1"})
decoded = decodeColumnValue("chunk", "source_id", encoded)
if !reflect.DeepEqual(decoded, []string{"doc-1"}) {
t.Fatalf("single element = %#v -> %#v", encoded, decoded)
}
}
// TestVectorRoundTripMatchesPsycopg2Float32 checks float32 precision both
// ways: encode narrows, decode widens back.
func TestVectorRoundTripMatchesPsycopg2Float32(t *testing.T) {
// float32 precision both ways: encode narrows, decode widens back.
encoded, err := encodeVector([]interface{}{0.123456789, 2.5, -0.75})
if err != nil {
t.Fatal(err)
}
if encoded != "[0.12345679,2.5,-0.75]" {
t.Fatalf("float32 text = %q", encoded)
}
decoded := decodeColumnValue("chunk", "q_3_vec", encoded)
numbers, ok := decoded.([]float64)
if !ok || len(numbers) != 3 {
t.Fatalf("float32 round trip = %#v", decoded)
}
// The text form is exact at float32 precision: narrowing the decoded
// float64 back to float32 must reproduce the original components.
for i, original := range []float32{float32(0.123456789), 2.5, -0.75} {
if float32(numbers[i]) != original {
t.Fatalf("component %d: %v != %v", i, float32(numbers[i]), original)
}
}
}
// TestPositionIntRoundTripRegroupsByFive checks that the flattened
// [[page,x1,y1,x2,y2],...] storage regroups into fives on read.
func TestPositionIntRoundTripRegroupsByFive(t *testing.T) {
original := []interface{}{
[]interface{}{int64(0), 12, 34, 56, 78},
[]interface{}{int64(0), 98, 76, 54, 32},
}
stored, err := encodeColumnValue("position_int", original)
if err != nil {
t.Fatal(err)
}
array, ok := stored.(pq.Int64Array)
if !ok || len(array) != 10 {
t.Fatalf("flat storage = %#v", stored)
}
// lib/pq renders int64[] as {0,12,34,...} when the server returns it as
// text into a dynamic scan; decode from that literal.
decoded := decodeColumnValue("chunk", "position_int", "{0,12,34,56,78,0,98,76,54,32}")
want := [][]int64{{0, 12, 34, 56, 78}, {0, 98, 76, 54, 32}}
if !reflect.DeepEqual(decoded, want) {
t.Fatalf("regroup = %#v", decoded)
}
}
// TestIntegerArrayRoundTrip checks plain integer array encode/decode.
func TestIntegerArrayRoundTrip(t *testing.T) {
stored, err := encodeColumnValue("page_num_int", []interface{}{int64(3), int64(1), int64(4)})
if err != nil {
t.Fatal(err)
}
if _, ok := stored.(pq.Int64Array); !ok {
t.Fatalf("int[] binding = %#v", stored)
}
decoded := decodeColumnValue("chunk", "page_num_int", "{3,1,4}")
if !reflect.DeepEqual(decoded, []int64{3, 1, 4}) {
t.Fatalf("int[] decode = %#v", decoded)
}
}
// TestJSONColumnRoundTrip checks JSON column text against map round-trips.
func TestJSONColumnRoundTrip(t *testing.T) {
object := map[string]interface{}{"_group_id": "g-7", "nested": map[string]interface{}{"k": 1.5}}
encoded, err := encodeColumnValue("metadata", object)
if err != nil {
t.Fatal(err)
}
decoded := decodeColumnValue("chunk", "metadata", encoded.(string))
decodedObject, ok := decoded.(map[string]interface{})
if !ok || decodedObject["_group_id"] != "g-7" {
t.Fatalf("json round trip = %#v", decoded)
}
// Unwritten metadata rows hold DEFAULT '' and must pass through so the
// exists-filter semantics stay intact.
if got := decodeColumnValue("chunk", "metadata", ""); got != "" {
t.Fatalf("DEFAULT '' passthrough = %#v", got)
}
}
// TestZeroBackfillVectorFormat checks the zero-vector placeholder format
// backfill writes into stale dimension columns.
func TestZeroBackfillVectorFormat(t *testing.T) {
// Documents missing a vector column the table already has are backfilled
// with zeros; the text form must match encodeVector's rendering.
encoded, err := encodeVector([]interface{}{0.0, 0.0, 0.0, 0.0})
if err != nil {
t.Fatal(err)
}
if zeroVector(4) != encoded {
t.Fatalf("zeroVector(4) = %q, encodeVector = %q", zeroVector(4), encoded)
}
if got := decodeColumnValue("chunk", "q_4_vec", zeroVector(4)); !reflect.DeepEqual(got, []float64{0, 0, 0, 0}) {
t.Fatalf("zero vector decode = %#v", got)
}
}
// TestKbIDListCollapsesToFirst checks that a kb_id list keeps only its
// first element on write.
func TestKbIDListCollapsesToFirst(t *testing.T) {
// The Python writer stored the first kb of the list; shared tables rely
// on that single value for row-level scoping.
encoded, err := encodeColumnValue("kb_id", []interface{}{"kb-1", "kb-2"})
if err != nil {
t.Fatal(err)
}
if encoded != "kb-1" {
t.Fatalf("kb_id = %#v", encoded)
}
if encoded, _ := encodeColumnValue("kb_id", []interface{}{}); encoded != nil {
t.Fatalf("empty kb list = %#v", encoded)
}
}
// TestTextColumnsCarryEmptyStringDefault checks that text columns default
// to empty strings, not NULLs.
func TestTextColumnsCarryEmptyStringDefault(t *testing.T) {
// The mapping's ES-parity contract: unwritten text is '' (not NULL), so
// exists filters can treat '' as absent and must_not as present.
for _, column := range chunkColumns {
if column.name == "position_int" || column.name == "page_num_int" || column.name == "top_int" {
if column.defaultSQL != "" {
t.Fatalf("array column %s must have no DEFAULT, got %q", column.name, column.defaultSQL)
}
continue
}
if column.defaultSQL == "" {
t.Fatalf("column %s lost its DEFAULT clause", column.name)
}
}
for _, column := range memoryColumns {
if column.defaultSQL == "" {
t.Fatalf("memory column %s lost its DEFAULT clause", column.name)
}
}
if metadataColumns[len(metadataColumns)-1].defaultSQL != "'{}'" {
t.Fatalf("meta_fields must DEFAULT '{}', got %q", metadataColumns[len(metadataColumns)-1].defaultSQL)
}
}
+534
View File
@@ -0,0 +1,534 @@
//
// 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 vastbase
import (
"context"
"database/sql"
"fmt"
"sort"
"strconv"
"strings"
"ragflow/internal/engine/types"
)
// InsertChunks writes chunks or memory messages with the Python connector's
// delete-then-insert upsert: legacy tables carry no PRIMARY KEY, so ON
// CONFLICT is not an option. The delete and the grouped inserts share one
// transaction, keeping a failed batch from leaving half-replaced rows.
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
}
kind := tableKind(baseName, datasetID)
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 {
// ChunkStoreExists also reports false while a full-text index DDL
// keeps failing (ensureFullTextIndexes degrades to a warning), so the
// vector size is required only to create a brand-new table; a batch
// without vectors must still flow into an existing one and re-run the
// self-healing CreateChunkStore path.
tablePresent, err := e.tableExists(ctx, baseName)
if err != nil {
return nil, err
}
if !tablePresent && 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); err != nil {
return nil, err
}
}
fieldNames := map[string]bool{}
normalized := make([]map[string]interface{}, 0, len(chunks))
for _, chunk := range chunks {
var document map[string]interface{}
switch kind {
case "memory":
document, err = normalizeMemory(chunk, datasetID)
case "skill":
documentID := stringValue(chunk["skill_id"])
if documentID == "" {
documentID = stringValue(chunk["id"])
}
document, err = normalizeSkill(chunk, documentID)
default:
document, err = normalizeChunk(chunk)
// normalizeChunk returns a nil map on encode errors; checking
// err first keeps the kb_id backfill from panicking on it.
if err == nil && stringValue(document["kb_id"]) == "" {
document["kb_id"] = datasetID
}
}
if err != nil {
return nil, err
}
for field := range document {
fieldNames[field] = true
}
normalized = append(normalized, document)
}
if err := e.ensureDynamicColumns(ctx, baseName, sortedSet(fieldNames)); err != nil {
return nil, err
}
if err := e.backfillVectorColumns(ctx, baseName, normalized); err != nil {
return nil, err
}
identifier := identifierColumn(kind)
ids := make([]string, 0, len(normalized))
for _, document := range normalized {
if id := stringValue(document[identifier]); id != "" {
ids = append(ids, id)
}
}
tx, err := e.db.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
defer tx.Rollback()
if len(ids) > 0 {
// Python-parity IN form ("DELETE ... WHERE id IN (...)"): plain
// placeholders avoid asking the server to infer an array parameter
// type for ANY($n).
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE %s IN (%s)",
quoteIdent(baseName), quoteIdent(identifier), placeholderList(len(ids), 0)),
stringArgs(ids)...); err != nil {
return nil, fmt.Errorf("vastbase: delete before insert in %s: %w", baseName, err)
}
}
for _, group := range groupByColumnSet(normalized) {
if err := insertGroup(ctx, tx, baseName, group); err != nil {
return nil, err
}
}
if err := tx.Commit(); err != nil {
return nil, err
}
return []string{}, nil
}
// backfillVectorColumns writes the zero-vector placeholder for every
// floatvector column the table already has but a document does not carry, the
// Python connector's backfill contract for multi-dimension tables.
func (e *Engine) backfillVectorColumns(ctx context.Context, tableName string, documents []map[string]interface{}) error {
columns, err := e.listTableColumns(ctx, tableName)
if err != nil || len(columns) == 0 {
return err
}
for column := range columns {
match := vectorDimRegex.FindStringSubmatch(column)
if match == nil {
continue
}
dimension, _ := strconv.Atoi(match[1])
zero := zeroVector(dimension)
for _, document := range documents {
if _, ok := document[column]; !ok {
document[column] = zero
}
}
}
return nil
}
// groupByColumnSet buckets documents by their sorted column set so each bucket
// forms one multi-row INSERT.
func groupByColumnSet(documents []map[string]interface{}) [][]map[string]interface{} {
groups := map[string][]map[string]interface{}{}
keys := make([]string, 0)
for _, document := range documents {
key := strings.Join(sortedColumns(document), "\x00")
if _, ok := groups[key]; !ok {
keys = append(keys, key)
}
groups[key] = append(groups[key], document)
}
sort.Strings(keys)
result := make([][]map[string]interface{}, 0, len(keys))
for _, key := range keys {
result = append(result, groups[key])
}
return result
}
// insertGroup writes one same-shape batch of documents as a single
// multi-row INSERT inside the caller's transaction.
func insertGroup(ctx context.Context, tx *sql.Tx, tableName string, documents []map[string]interface{}) error {
columns := sortedColumns(documents[0])
quoted := make([]string, len(columns))
for i, column := range columns {
if err := validateIdentifier(column); err != nil {
return err
}
quoted[i] = quoteIdent(column)
}
args := make([]interface{}, 0, len(columns)*len(documents))
placeholders := make([]string, 0, len(documents))
for _, document := range documents {
row := make([]string, len(columns))
for i, column := range columns {
args = append(args, document[column])
row[i] = fmt.Sprintf("$%d", len(args))
}
placeholders = append(placeholders, "("+strings.Join(row, ",")+")")
}
query := fmt.Sprintf("INSERT INTO %s (%s) VALUES %s",
quoteIdent(tableName), strings.Join(quoted, ", "), strings.Join(placeholders, ","))
if _, err := tx.ExecContext(ctx, query, args...); err != nil {
return fmt.Errorf("vastbase: insert into %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
}
columns, err := e.listTableColumns(ctx, baseName)
if err != nil {
return nil, err
}
condition := map[string]interface{}{identifierColumn(kind): chunkID}
if kind == "memory" {
condition["memory_id"] = datasetIDs
} else if kind == "chunk" {
condition["kb_id"] = datasetIDs
}
whereSQL, args, err := buildFilter(condition, kind, columns)
if err != nil {
return nil, err
}
rows, err := e.queryRows(ctx, "SELECT * FROM "+quoteIdent(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 the shared
// table needs. remove as a field name resets the column to its DEFAULT; remove
// as an object drops values from ###-joined keyword lists via read-modify-write.
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)
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)
}
columns, err := e.listTableColumns(ctx, baseName)
if err != nil {
return err
}
whereSQL, args, err := buildFilter(condition, kind, columns)
if err != nil {
return err
}
setParts := make([]string, 0, len(newValue))
setArgs := make([]interface{}, 0, len(newValue))
removeFields := map[string]interface{}{}
for _, rawKey := range sortedKeys(newValue) {
key := rawKey
if kind == "memory" {
key = mapMemoryField(rawKey)
} else if kind == "skill" && key == "id" {
key = "skill_id"
}
value := newValue[rawKey]
switch key {
case "id":
continue
case "remove":
switch remove := value.(type) {
case string:
if kind == "memory" {
remove = mapMemoryField(remove)
}
if _, known := columns[remove]; !known {
return fmt.Errorf("vastbase: cannot remove unknown column %q", remove)
}
setParts = append(setParts, quoteIdent(remove)+" = DEFAULT")
case map[string]interface{}:
for column, item := range remove {
if _, known := columns[column]; !known {
return fmt.Errorf("vastbase: cannot remove from unknown column %q", column)
}
removeFields[column] = item
}
default:
return fmt.Errorf("vastbase: remove must be a field name or object")
}
default:
if err := validateIdentifier(key); err != nil {
return err
}
encoded, err := encodeUpdateValue(kind, key, value)
if err != nil {
return err
}
setArgs = append(setArgs, encoded)
setParts = append(setParts, fmt.Sprintf("%s = $%d", quoteIdent(key), len(setArgs)))
if kind == "memory" && key == "content_ltks" {
setArgs = append(setArgs, tokenizeMemoryContent(stringValue(value)))
setParts = append(setParts, fmt.Sprintf("%s = $%d", quoteIdent("tokenized_content_ltks"), len(setArgs)))
}
}
}
if err := e.applyKeywordRemovals(ctx, baseName, whereSQL, args, removeFields); err != nil {
return err
}
if len(setParts) == 0 {
return nil
}
query := fmt.Sprintf("UPDATE %s SET %s WHERE %s",
quoteIdent(baseName), strings.Join(setParts, ", "), shiftPlaceholders(whereSQL, len(setArgs)))
setArgs = append(setArgs, args...)
_, err = e.db.ExecContext(ctx, query, setArgs...)
if err != nil {
return fmt.Errorf("vastbase: update %s: %w", baseName, err)
}
return nil
}
// applyKeywordRemovals implements the Python connector's read-modify-write for
// remove objects: keyword columns are ###-joined text, so affected rows are
// read, the value is dropped from the decoded list, and rows landing on the
// same resulting list are updated in one statement per (column, list) pair.
func (e *Engine) applyKeywordRemovals(ctx context.Context, tableName, whereSQL string, args []interface{}, removeFields map[string]interface{}) error {
if len(removeFields) == 0 {
return nil
}
removeColumns := sortedMapKeys(removeFields)
selectColumns := make([]string, 0, len(removeColumns)+1)
for _, column := range removeColumns {
selectColumns = append(selectColumns, quoteIdent(column))
}
selectColumns = append(selectColumns, quoteIdent("id"))
rows, err := e.queryRows(ctx, fmt.Sprintf("SELECT %s FROM %s WHERE %s",
strings.Join(selectColumns, ", "), quoteIdent(tableName), whereSQL), args...)
if err != nil {
return fmt.Errorf("vastbase: read %s for removal: %w", tableName, err)
}
updates := map[string]map[string][]string{}
for _, row := range rows {
id := stringValue(row["id"])
if id == "" {
continue
}
for _, column := range removeColumns {
current, _ := decodeColumnValue("chunk", column, row[column]).([]string)
removeValue := stringValue(removeFields[column])
if !containsString(current, removeValue) {
continue
}
updated := make([]string, 0, len(current))
for _, item := range current {
if item != removeValue {
updated = append(updated, item)
}
}
if updates[column] == nil {
updates[column] = map[string][]string{}
}
updates[column][strings.Join(updated, keywordSeparator)] = append(updates[column][strings.Join(updated, keywordSeparator)], id)
}
}
for _, column := range sortedMapKeys(updates) {
for joined, ids := range updates[column] {
// The filter args keep their original $1..$n numbering (they are
// bound first); the SET value and the id list continue past them.
query := fmt.Sprintf("UPDATE %s SET %s = $%d WHERE %s AND id IN (%s)",
quoteIdent(tableName), quoteIdent(column), len(args)+1, whereSQL,
placeholderList(len(ids), len(args)+1))
execArgs := append(append([]interface{}{}, args...), joined)
execArgs = append(execArgs, stringArgs(ids)...)
if _, err := e.db.ExecContext(ctx, query, execArgs...); err != nil {
return fmt.Errorf("vastbase: remove from %s: %w", tableName, err)
}
}
}
return nil
}
// 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
}
columns, err := e.listTableColumns(ctx, baseName)
if err != nil {
return 0, err
}
whereSQL, args, err := buildFilter(condition, kind, columns)
if err != nil {
return 0, err
}
result, err := e.db.ExecContext(ctx, "DELETE FROM "+quoteIdent(baseName)+" WHERE "+whereSQL, args...)
if err != nil {
return 0, fmt.Errorf("vastbase: delete from %s: %w", baseName, err)
}
return result.RowsAffected()
}
// placeholderList renders "offset+1, offset+2, ..." for an IN clause with the
// given number of bound arguments.
func placeholderList(count, offset int) string {
placeholders := make([]string, count)
for i := range placeholders {
placeholders[i] = fmt.Sprintf("$%d", offset+i+1)
}
return strings.Join(placeholders, ", ")
}
// stringArgs widens a string slice to the interface slice database/sql
// variadic calls need.
func stringArgs(values []string) []interface{} {
args := make([]interface{}, len(values))
for i, value := range values {
args[i] = value
}
return args
}
// tableKind classifies a table row set: skill tables carry skill_id, memory_
// tables hold memory rows, everything else is a chunk table.
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"
}
}
// identifierColumn is the primary-key column name for a row kind.
func identifierColumn(kind string) string {
if kind == "skill" {
return "skill_id"
}
return "id"
}
// mapMemoryField translates a logical memory field to its column name,
// passing unknown fields through.
func mapMemoryField(field string) string {
if mapped, ok := memoryFieldToColumn[field]; ok {
return mapped
}
return field
}
// copyMap shallow-copies a filter map so per-table discriminators can be
// added without aliasing the request's filter.
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
}
// sortedSet lists a set's members in sorted order.
func sortedSet(source map[string]bool) []string {
values := make([]string, 0, len(source))
for value := range source {
values = append(values, value)
}
sort.Strings(values)
return values
}
// sortedMapKeys lists a map's keys in sorted order.
func sortedMapKeys[V any](source map[string]V) []string {
keys := make([]string, 0, len(source))
for key := range source {
keys = append(keys, key)
}
sort.Strings(keys)
return keys
}
+539
View File
@@ -0,0 +1,539 @@
//
// 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 vastbase
import (
"context"
"errors"
"fmt"
"reflect"
"regexp"
"strings"
"testing"
"ragflow/internal/engine/types"
"ragflow/internal/server/config"
"github.com/DATA-DOG/go-sqlmock"
)
const listColumnsQuery = "SELECT column_name, data_type, column_default, is_nullable FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = $1"
// expectExistingChunkStore wires the probes ChunkStoreExists runs against a
// fully provisioned PG-mode chunk table.
func expectExistingChunkStore(mock sqlmock.Sqlmock, table string) {
expectCount(mock, tableCountQuery(), table)
for _, column := range chunkIndexColumns {
expectCount(mock, indexCountQuery(), table, regularIndexName(table, column))
}
expectCount(mock, indexCountQuery(), table, "text_gin_idx_"+table)
}
// expectListColumns stubs the existence probe and column listing round trip
// with name/type pairs.
func expectListColumns(mock sqlmock.Sqlmock, table string, columns ...[2]string) {
expectCount(mock, tableCountQuery(), table)
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
for _, column := range columns {
rows.AddRow(column[0], column[1], nil, "YES")
}
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs(table).WillReturnRows(rows)
}
// TestInsertChunksRunsDeleteThenInsertInOneTransaction checks the replace
// semantics: the delete and the insert share one transaction.
func TestInsertChunksRunsDeleteThenInsertInOneTransaction(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
// ChunkStoreExists: table present with every static index.
expectExistingChunkStore(mock, "t1")
// Existing vector column and graph index: no DDL. ensureDynamicColumns
// also probes the vector column (it is not part of the static schema).
expectCount(mock, columnCountQuery(), "t1", "q_3_vec")
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", "q_3_vec_graph"))
expectCount(mock, columnCountQuery(), "t1", "q_3_vec")
// backfillVectorColumns probes the live column set.
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
[2]string{"content_ltks", "text"}, [2]string{"q_3_vec", "text"},
)
mock.ExpectBegin()
mock.ExpectExec(regexp.QuoteMeta(`DELETE FROM "t1" WHERE "id" IN ($1)`)).
WithArgs("c1").WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(regexp.QuoteMeta(
`INSERT INTO "t1" ("content_ltks", "id", "kb_id", "q_3_vec") VALUES ($1,$2,$3,$4)`)).
WithArgs("hello", "c1", "kb-1", "[0.1,0.2,0.3]").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
inserted, err := engine.InsertChunks(context.Background(), []map[string]interface{}{{
"id": "c1", "content_ltks": "hello", "q_3_vec": []interface{}{0.1, 0.2, 0.3},
}}, "t1", "kb-1")
if err != nil {
t.Fatal(err)
}
if len(inserted) != 0 {
t.Fatalf("InsertChunks returns the Python connector's empty id list, got %#v", inserted)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// A malformed chunk must surface the encode error, not panic: normalizeChunk
// returns a nil map on failure, and the kb_id backfill used to write into it
// before the error check.
func TestInsertChunksReturnsEncodeErrorForMalformedVector(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectExistingChunkStore(mock, "t1")
// vectorDimension reads the column name, so the store side runs before the
// chunk body is ever encoded.
expectCount(mock, columnCountQuery(), "t1", "q_3_vec")
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", "q_3_vec_graph"))
_, err = engine.InsertChunks(context.Background(), []map[string]interface{}{{
"id": "c1", "q_3_vec": "oops",
}}, "t1", "kb-1")
if err == nil || !strings.Contains(err.Error(), "encode q_3_vec") {
t.Fatalf("err = %v", err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestInsertChunksGroupsByColumnSet checks that documents with different
// column sets are inserted as separate grouped statements.
func TestInsertChunksGroupsByColumnSet(t *testing.T) {
documents := []map[string]interface{}{
{"id": "a", "title_tks": "x"},
{"id": "b", "title_tks": "y", "available_int": 1},
{"id": "c", "title_tks": "z"},
}
groups := groupByColumnSet(documents)
if len(groups) != 2 {
t.Fatalf("groups = %d, want 2", len(groups))
}
// Groups are ordered by their sorted column-set key, so the
// available_int group sorts first.
if len(groups[0]) != 1 || len(groups[1]) != 2 {
t.Fatalf("group sizes = %d, %d", len(groups[0]), len(groups[1]))
}
}
// TestInsertChunksRequiresVectorSizeForNewStore checks that creating a
// brand-new chunk store without a vector size fails instead of guessing a
// column.
func TestInsertChunksRequiresVectorSizeForNewStore(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
if _, err := engine.InsertChunks(context.Background(),
[]map[string]interface{}{{"id": "c1", "content_ltks": "x"}}, "t1", "kb-1"); err == nil {
t.Fatal("creating a new store without an inferable vector dimension must fail")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// A batch without vectors must still flow into an existing table even while
// the full-text index DDL keeps failing: ChunkStoreExists reports the store
// missing (it requires the index), but the vector size is only needed to
// create a brand-new table.
func TestInsertChunksVectorlessBatchOnTableWithFailedFullTextIndex(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
// ChunkStoreExists: table and b-tree indexes present, GIN index missing.
expectCount(mock, tableCountQuery(), "t1")
for _, column := range chunkIndexColumns {
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", column))
}
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs("t1", "text_gin_idx_t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
// The table itself exists, so no vector size is inferred; the store is
// re-created to self-heal the missing index.
expectCount(mock, tableCountQuery(), "t1")
// CreateChunkStore: ensureTable and the regular indexes are no-ops.
expectCount(mock, tableCountQuery(), "t1")
for _, column := range chunkIndexColumns {
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", column))
}
// The GIN DDL fails again — a warning, not an error.
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs("t1", "text_gin_idx_t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
expressions := make([]string, 0, len(fullTextFields))
for _, field := range fullTextFields {
expressions = append(expressions, fmt.Sprintf("to_tsvector('cn_tokenizer', %s)", quoteIdent(field)))
}
mock.ExpectExec(regexp.QuoteMeta(fmt.Sprintf(
`CREATE INDEX IF NOT EXISTS "text_gin_idx_t1" ON "t1" USING gin (%s)`,
strings.Join(expressions, ", ")))).WillReturnError(fmt.Errorf("gin extension missing"))
// vectorSize 0: no vector column or index is created.
// backfillVectorColumns lists the live columns (no vector column present).
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"}, [2]string{"content_ltks", "text"},
)
mock.ExpectBegin()
mock.ExpectExec(regexp.QuoteMeta(`DELETE FROM "t1" WHERE "id" IN ($1)`)).
WithArgs("c3").WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectExec(regexp.QuoteMeta(
`INSERT INTO "t1" ("content_ltks", "id", "kb_id") VALUES ($1,$2,$3)`)).
WithArgs("third", "c3", "kb-1").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
if _, err := engine.InsertChunks(context.Background(), []map[string]interface{}{
{"id": "c3", "kb_id": "kb-1", "content_ltks": "third"},
}, "t1", "kb-1"); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestUpdateChunksShiftsWherePlaceholdersAfterSet checks that UPDATE
// placeholder numbering keeps SET arguments ahead of the WHERE clause's.
func TestUpdateChunksShiftsWherePlaceholdersAfterSet(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectCount(mock, tableCountQuery(), "t1")
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"}, [2]string{"title_tks", "text"},
)
mock.ExpectExec(regexp.QuoteMeta(
`UPDATE "t1" SET "title_tks" = $1 WHERE "id" = $2 AND "kb_id" = $3`)).
WithArgs("new title", "c1", "kb-1").WillReturnResult(sqlmock.NewResult(0, 1))
if err := engine.UpdateChunks(context.Background(),
map[string]interface{}{"id": "c1"},
map[string]interface{}{"title_tks": "new title"},
"t1", "kb-1"); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// A write filter whose only row selectors are keyword conditions must keep
// them: dropping compile_kwd/type_kwd/title_kwd (navigation-graph updates)
// would leave the dataset scope alone and hit every row of the dataset.
func TestUpdateAndDeleteKeepKeywordPredicates(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectCount(mock, tableCountQuery(), "t1")
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
[2]string{"title_kwd", "varchar"}, [2]string{"title_tks", "text"},
)
// Filter order follows sortedKeys: kb_id first, then the keyword predicate.
mock.ExpectExec(regexp.QuoteMeta(
`UPDATE "t1" SET "title_tks" = $1 WHERE "kb_id" = $2 AND (('###' || "title_kwd" || '###') LIKE $3 OR ('###' || "title_kwd" || '###') LIKE $4)`)).
WithArgs("x", "kb-1", `%###nav\_cluster###%`, `%###nav\_doc###%`).
WillReturnResult(sqlmock.NewResult(0, 1))
if err := engine.UpdateChunks(context.Background(),
map[string]interface{}{"title_kwd": []string{"nav_cluster", "nav_doc"}},
map[string]interface{}{"title_tks": "x"},
"t1", "kb-1"); err != nil {
t.Fatal(err)
}
expectCount(mock, tableCountQuery(), "t1")
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
[2]string{"compile_kwd", "varchar"}, [2]string{"type_kwd", "varchar"},
)
// sortedKeys puts compile_kwd before kb_id before type_kwd.
mock.ExpectExec(regexp.QuoteMeta(
`DELETE FROM "t1" WHERE (('###' || "compile_kwd" || '###') LIKE $1) AND "kb_id" = $2 AND (('###' || "type_kwd" || '###') LIKE $3)`)).
WithArgs("%###done###%", "kb-1", `%###nav\_cluster###%`).
WillReturnResult(sqlmock.NewResult(0, 1))
if _, err := engine.DeleteChunks(context.Background(),
map[string]interface{}{"compile_kwd": "done", "type_kwd": "nav_cluster"},
"t1", "kb-1"); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestUpdateChunksRemoveStringResetsColumnDefault checks that removing a
// field resets its column to the declared default instead of NULL.
func TestUpdateChunksRemoveStringResetsColumnDefault(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectCount(mock, tableCountQuery(), "t1")
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"}, [2]string{"title_tks", "text"},
)
mock.ExpectExec(regexp.QuoteMeta(
`UPDATE "t1" SET "title_tks" = DEFAULT WHERE "id" = $1 AND "kb_id" = $2`)).
WithArgs("c1", "kb-1").WillReturnResult(sqlmock.NewResult(0, 1))
if err := engine.UpdateChunks(context.Background(),
map[string]interface{}{"id": "c1"},
map[string]interface{}{"remove": "title_tks"},
"t1", "kb-1"); err != nil {
t.Fatal(err)
}
// Removing a column the table lacks is an error, not a silent pass.
expectCount(mock, tableCountQuery(), "t1")
expectListColumns(mock, "t1", [2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"})
if err := engine.UpdateChunks(context.Background(),
map[string]interface{}{"id": "c1"},
map[string]interface{}{"remove": "ghost_col"},
"t1", "kb-1"); err == nil {
t.Fatal("unknown remove column must fail")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestUpdateChunksKeywordRemovalReadsModifiesWrites checks that removing a
// keyword rewrites the joined value through a read-modify-write.
func TestUpdateChunksKeywordRemovalReadsModifiesWrites(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectCount(mock, tableCountQuery(), "t1")
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
[2]string{"important_kwd", "text"},
)
// Read affected rows' current lists.
mock.ExpectQuery(regexp.QuoteMeta(
`SELECT "important_kwd", "id" FROM "t1" WHERE "id" = $1 AND "kb_id" = $2`)).
WithArgs("c1", "kb-1").
WillReturnRows(sqlmock.NewRows([]string{"important_kwd", "id"}).AddRow("alpha###beta", "c1"))
// Write the shortened list for the matching rows. The filter args keep
// their $1/$2 numbering (they are bound first); the SET value continues
// as $3 and the id list after it.
mock.ExpectExec(regexp.QuoteMeta(
`UPDATE "t1" SET "important_kwd" = $3 WHERE "id" = $1 AND "kb_id" = $2 AND id IN ($4)`)).
WithArgs("c1", "kb-1", "alpha", "c1").WillReturnResult(sqlmock.NewResult(0, 1))
if err := engine.UpdateChunks(context.Background(),
map[string]interface{}{"id": "c1"},
map[string]interface{}{"remove": map[string]interface{}{"important_kwd": "beta"}},
"t1", "kb-1"); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestUpdateChunksMemoryContentRewritesTokenizedCompanion checks that a
// memory content update re-derives the tokenized lookup column.
func TestUpdateChunksMemoryContentRewritesTokenizedCompanion(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
expectCount(mock, tableCountQuery(), "memory_t1")
expectListColumns(mock, "memory_t1",
[2]string{"id", "varchar"}, [2]string{"memory_id", "varchar"},
[2]string{"content_ltks", "text"}, [2]string{"tokenized_content_ltks", "text"},
)
mock.ExpectExec(regexp.QuoteMeta(
`UPDATE "memory_t1" SET "content_ltks" = $1, "tokenized_content_ltks" = $2 WHERE "id" = $3 AND "memory_id" = $4`)).
WithArgs("new text", sqlmock.AnyArg(), "m1_1", "mem-1").
WillReturnResult(sqlmock.NewResult(0, 1))
if err := engine.UpdateChunks(context.Background(),
map[string]interface{}{"id": "m1_1"},
map[string]interface{}{"content": "new text"},
"memory_t1", "mem-1"); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestUpdateChunksMissingTableReportsIndexNotFound checks the error shape
// when the named table does not exist.
func TestUpdateChunksMissingTableReportsIndexNotFound(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
err = engine.UpdateChunks(context.Background(),
map[string]interface{}{"id": "c1"}, map[string]interface{}{"title_tks": "x"}, "t1", "kb-1")
if !errors.Is(err, types.ErrIndexNotFound) {
t.Fatalf("err = %v", err)
}
}
// TestDeleteChunksScopesByDataset checks that deletes are scoped to the
// dataset's rows through the kb_id discriminator.
func TestDeleteChunksScopesByDataset(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectCount(mock, tableCountQuery(), "t1")
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"}, [2]string{"doc_id", "varchar"},
)
mock.ExpectExec(regexp.QuoteMeta(
`DELETE FROM "t1" WHERE "doc_id" = $1 AND "kb_id" = $2`)).
WithArgs("doc-1", "kb-1").WillReturnResult(sqlmock.NewResult(0, 4))
deleted, err := engine.DeleteChunks(context.Background(),
map[string]interface{}{"doc_id": "doc-1"}, "t1", "kb-1")
if err != nil || deleted != 4 {
t.Fatalf("deleted = %d err = %v", deleted, err)
}
// Missing tables delete nothing without error.
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
deleted, err = engine.DeleteChunks(context.Background(), nil, "t1", "kb-1")
if err != nil || deleted != 0 {
t.Fatalf("deleted = %d err = %v", deleted, err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestGetChunkScopedQuery pins the GetChunk select and its dataset scoping.
func TestGetChunkScopedQuery(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectCount(mock, tableCountQuery(), "t1")
expectListColumns(mock, "t1",
[2]string{"id", "varchar"}, [2]string{"kb_id", "varchar"},
[2]string{"content_ltks", "text"}, [2]string{"important_kwd", "text"},
)
mock.ExpectQuery(regexp.QuoteMeta(
`SELECT * FROM "t1" WHERE "id" = $1 AND "kb_id" IN ($2) LIMIT 1`)).
WithArgs("c1", "kb-1").
WillReturnRows(sqlmock.NewRows([]string{"id", "content_ltks", "important_kwd"}).
AddRow("c1", "hello", "alpha###beta"))
document, err := engine.GetChunk(context.Background(), "t1", "c1", []string{"kb-1"})
if err != nil {
t.Fatal(err)
}
chunk, ok := document.(map[string]interface{})
if !ok || chunk["id"] != "c1" || chunk["content_ltks"] != "hello" {
t.Fatalf("chunk = %#v", document)
}
if !reflect.DeepEqual(chunk["important_kwd"], []string{"alpha", "beta"}) {
t.Fatalf("keyword decode = %#v", chunk["important_kwd"])
}
// Memory rows report not-found instead of nil.
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("memory_t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
if _, err := engine.GetChunk(context.Background(), "memory_t1", "m1_1", []string{"mem-1"}); !errors.Is(err, types.ErrDocumentNotFound) {
t.Fatalf("err = %v", err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestTableKindRouting checks tableKind classification across skill,
// memory, and chunk tables.
func TestTableKindRouting(t *testing.T) {
cases := []struct {
table string
datasets []string
want string
}{
{"memory_t1", []string{"mem-1"}, "memory"},
{"skill_t1", nil, "skill"},
{"t1", []string{"skill"}, "skill"},
{"ragflow_doc_meta_tenant", nil, "metadata"},
{"t1", []string{"kb-1"}, "chunk"},
}
for _, test := range cases {
if got := tableKind(test.table, test.datasets...); got != test.want {
t.Fatalf("tableKind(%q, %v) = %q, want %q", test.table, test.datasets, got, test.want)
}
}
}
+157
View File
@@ -0,0 +1,157 @@
//
// 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 vastbase
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); 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 := tx.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE skill_id = $1", quoteIdent(indexName)), stringValue(normalized["skill_id"])); err != nil {
return fmt.Errorf("vastbase: delete before skill insert: %w", err)
}
if err := insertGroup(ctx, tx, indexName, []map[string]interface{}{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); err != nil {
return nil, err
}
}
normalized := make([]map[string]interface{}, 0, len(docs))
ids := make([]string, 0, len(docs))
for _, raw := range docs {
document, ok := raw.(map[string]interface{})
if !ok {
continue
}
docID := stringValue(document["skill_id"])
if docID == "" {
docID = stringValue(document["id"])
}
row, normalizeErr := normalizeSkill(document, docID)
if normalizeErr != nil {
return nil, normalizeErr
}
normalized = append(normalized, row)
ids = append(ids, docID)
}
tx, err := e.db.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
defer tx.Rollback()
if len(ids) > 0 {
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE skill_id IN (%s)",
quoteIdent(indexName), placeholderList(len(ids), 0)), stringArgs(ids)...); err != nil {
return nil, fmt.Errorf("vastbase: delete before skill bulk insert: %w", err)
}
}
for _, group := range groupByColumnSet(normalized) {
if err := insertGroup(ctx, tx, indexName, group); err != nil {
return nil, err
}
}
if err := tx.Commit(); err != nil {
return nil, err
}
return map[string]interface{}{"inserted": len(normalized)}, nil
}
// DeleteDocument deletes a skill by primary key; the slash-folded variant
// covers callers that mangle skill ids into path-safe forms.
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 = $1 OR replace(skill_id, '/', '_') = $2", quoteIdent(indexName)), docID, docID)
return err
}
_, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE id = $1", quoteIdent(indexName)), docID)
return err
}
+363
View File
@@ -0,0 +1,363 @@
//
// 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 vastbase
import (
"encoding/json"
"fmt"
"reflect"
"regexp"
"sort"
"strconv"
"strings"
)
var placeholderRegex = regexp.MustCompile(`\$(\d+)`)
// shiftPlaceholders renumbers the $n placeholders of a built clause so it can
// be embedded after clauses that already consumed the first offset args.
func shiftPlaceholders(clause string, offset int) string {
if offset == 0 {
return clause
}
return placeholderRegex.ReplaceAllStringFunc(clause, func(match string) string {
number, _ := strconv.Atoi(match[1:])
return fmt.Sprintf("$%d", number+offset)
})
}
// filterBuilder accumulates WHERE parts with a running $n counter; args order
// is bound to placeholder order by construction.
type filterBuilder struct {
parts []string
args []interface{}
}
// raw appends a pre-built SQL fragment to the conjunction.
func (b *filterBuilder) raw(part string) {
b.parts = append(b.parts, part)
}
// placeholder binds a value and returns its $n reference.
func (b *filterBuilder) placeholder(value interface{}) string {
b.args = append(b.args, value)
return fmt.Sprintf("$%d", len(b.args))
}
// join renders the accumulated parts as one AND chain, with a trivially
// true predicate for the empty case.
func (b *filterBuilder) join() string {
if len(b.parts) == 0 {
return "1=1"
}
return strings.Join(b.parts, " AND ")
}
// buildFilter translates an ES-style condition map into a parameterized WHERE
// clause, porting the Python connector's equivalent_condition_to_str:
//
// - falsy values are dropped (the Python code skipped them too),
// - keyword columns (source_id, *_kwd except docnm_kwd/knowledge_graph_kwd)
// store ###-joined lists, so their term/terms conditions become LIKE
// containment tests on the separator-padded column. Dropping them instead
// would collapse a write filter to its dataset scope and let an update or
// delete hit every row of the dataset,
// - a column missing from the live schema matches no row (1=0), like an
// unmapped ES field; a nil columns map means "unknown schema" and allows
// every column,
// - exists on a text column excludes DEFAULT ” rows, because unwritten text
// columns hold ” rather than NULL; must_not.exists is the mirror image.
//
// Memory tables skip the keyword translation: their keyword-named columns
// (message_type_kwd, source_id) store plain scalars.
//
// Unlike the Python connector, kb_id is NOT skipped: Go tables are shared per
// baseName with kb_id/memory_id as row-level dataset discriminators, so the
// scope condition must reach the SQL.
func buildFilter(condition map[string]interface{}, kind string, columns map[string]columnMeta) (string, []interface{}, error) {
if _, ok := condition["_id"]; ok {
return "", nil, fmt.Errorf("vastbase: _id is not a filterable column, use id")
}
filter := &filterBuilder{}
for _, rawKey := range sortedKeys(condition) {
key := rawKey
if kind == "memory" {
key = mapMemoryField(key)
} else if kind == "skill" && key == "id" {
key = "skill_id"
}
value := condition[rawKey]
if isEmptyFilterValue(value) {
continue
}
if fieldKeyword(key) && kind != "memory" {
filter.appendKeywordTerms(key, value, columns)
continue
}
switch key {
case "exists":
column, ok := value.(string)
if !ok {
continue
}
if kind == "memory" {
column = mapMemoryField(column)
}
filter.appendExists(column, columns)
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)
}
meta, known := columns[column]
if !known {
// A column the table lacks is NULL on every row, so the
// must_not clause holds everywhere.
continue
}
quoted := quoteIdent(column)
if isTextColumn(meta.dataType) {
filter.raw("(" + quoted + " IS NULL OR " + quoted + " = '')")
} else {
filter.raw(quoted + " IS NULL")
}
case "metadata_filtering_conditions":
if err := filter.appendMetadataFilteringConditions(value); err != nil {
return "", nil, err
}
default:
if columns != nil {
if _, known := columns[key]; !known {
filter.raw("1=0")
continue
}
}
if values, ok := interfaceSlice(value); ok {
placeholders := make([]string, 0, len(values))
for _, item := range values {
placeholders = append(placeholders, filter.placeholder(item))
}
filter.raw(quoteIdent(key) + " IN (" + strings.Join(placeholders, ", ") + ")")
} else {
filter.raw(quoteIdent(key) + " = " + filter.placeholder(value))
}
}
}
return filter.join(), filter.args, nil
}
// appendExists renders an exists/does-not-exist filter against a column,
// treating never-stored fields as absent on every row (ES dynamic-mapping
// parity).
func (b *filterBuilder) appendExists(column string, columns map[string]columnMeta) {
meta, known := columns[column]
if !known {
// ES dynamic-mapping parity: a field this table never stored exists
// on no row.
b.raw("1=0")
return
}
quoted := quoteIdent(column)
if isTextColumn(meta.dataType) {
b.raw("(" + quoted + " IS NOT NULL AND " + quoted + " <> '')")
return
}
b.raw(quoted + " IS NOT NULL")
}
// likeEscaper neutralizes LIKE metacharacters so a bound value matches
// literally; backslash is the default LIKE escape character.
var likeEscaper = strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`)
// appendKeywordTerms translates a term/terms condition on a ###-joined keyword
// column into LIKE containment tests: padding the column with separators makes
// first and last list elements match and prevents substring false positives
// ("nav_doc" must not match a stored "nav_doc_2" element). A terms list ORs
// one test per value, matching the ES any-of semantics. The pattern is bound
// as a parameter, so only the value's own LIKE metacharacters need escaping.
func (b *filterBuilder) appendKeywordTerms(column string, value interface{}, columns map[string]columnMeta) {
if columns != nil {
if _, known := columns[column]; !known {
// ES dynamic-mapping parity: a field this table never stored exists
// on no row.
b.raw("1=0")
return
}
}
values := []interface{}{value}
if list, ok := interfaceSlice(value); ok {
values = list
}
padded := "('" + keywordSeparator + "' || " + quoteIdent(column) + " || '" + keywordSeparator + "')"
parts := make([]string, 0, len(values))
for _, item := range values {
text := stringValue(item)
if text == "" {
continue
}
pattern := "%" + keywordSeparator + likeEscaper.Replace(text) + keywordSeparator + "%"
parts = append(parts, padded+" LIKE "+b.placeholder(pattern))
}
if len(parts) == 0 {
// Only empty-string values remain. The decode side drops empty
// elements, so no stored row can ever contain one.
b.raw("1=0")
return
}
b.raw("(" + strings.Join(parts, " OR ") + ")")
}
// metadataJSON reads the metadata text column as jsonb, mapping rows the
// writer left at DEFAULT ” to an empty object so the cast never fails the
// whole query.
const metadataJSON = `CASE WHEN "metadata" LIKE '{%' THEN "metadata"::jsonb ELSE '{}'::jsonb END`
// metadataText builds the jsonb string extraction for one metadata key.
func (b *filterBuilder) metadataText(path string) string {
return "(" + metadataJSON + ") ->> " + b.placeholder(path)
}
// metadataCast wraps the extracted text in a guarded cast: rows whose value
// does not parse as the target type yield NULL (no match) instead of failing
// the query, matching the silent-coercion behavior of the MySQL JSON path.
func (b *filterBuilder) metadataCast(path, castType, pattern string) string {
text := b.metadataText(path)
return "(CASE WHEN " + text + " ~ '" + pattern + "' THEN (" + text + ")::" + castType + " ELSE NULL END)"
}
// appendMetadataFilteringConditions adds the post-filter conditions from
// the metadata filtering option, joined as a conjunction.
func (b *filterBuilder) appendMetadataFilteringConditions(raw interface{}) error {
filter, ok := raw.(map[string]interface{})
if !ok {
return nil
}
conditions, ok := interfaceSlice(filter["conditions"])
if !ok || len(conditions) == 0 {
return nil
}
logic := strings.ToUpper(stringValue(filter["logical_operator"]))
if logic == "" {
logic = "AND"
}
if logic != "AND" && logic != "OR" {
return fmt.Errorf("vastbase: unsupported metadata logical operator: %s", logic)
}
parts := make([]string, 0, len(conditions))
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
}
// A plain top-level key, not a MySQL "$.name" JSON path: the jsonb
// ->>/-> text operand is a literal key lookup, and no metadata object
// carries a key literally named "$.name".
path := name
value := condition["value"]
switch operator {
case "is", "is not":
comparison := "="
if operator == "is not" {
comparison = "<>"
}
parts = append(parts, b.metadataText(path)+" "+comparison+" "+b.placeholder(value))
case "=", "≠", ">", "<", "≥", "≤":
comparison := map[string]string{"=": "=", "≠": "<>", ">": ">", "<": "<", "≥": ">=", "≤": "<="}[operator]
parts = append(parts, b.metadataCast(path, "numeric", `^-?[0-9]+(\.[0-9]+)?([eE][+-]?[0-9]+)?$`)+
" "+comparison+" "+b.placeholder(value))
case "before", "after":
comparison := "<"
if operator == "after" {
comparison = ">"
}
parts = append(parts, b.metadataCast(path, "timestamp", `^[0-9]{4}-[0-9]{2}-[0-9]{2}([ T][0-9]{2}:[0-9]{2}(:[0-9]{2})?)?$`)+
" "+comparison+" "+b.placeholder(value))
case "contains", "not contains":
encoded, err := json.Marshal(value)
if err != nil {
continue
}
part := "(" + metadataJSON + ") -> " + b.placeholder(path) + " @> " + b.placeholder(string(encoded)) + "::jsonb"
if operator == "not contains" {
part = "NOT " + part
}
parts = append(parts, part)
case "start with":
parts = append(parts, b.metadataText(path)+" LIKE "+b.placeholder(stringValue(value)+"%"))
case "end with":
parts = append(parts, b.metadataText(path)+" LIKE "+b.placeholder("%"+stringValue(value)))
case "empty", "not empty":
text := b.metadataText(path)
empty := "(" + text + " IS NULL OR " + text + " = '' OR " + text + " = '[]' OR " + text + " = '{}')"
if operator == "not empty" {
empty = "(" + text + " IS NOT NULL AND " + text + " != '' AND " + text + " != '[]' AND " + text + " != '{}')"
}
parts = append(parts, empty)
}
}
if len(parts) == 0 {
return nil
}
b.raw("(" + strings.Join(parts, " "+logic+" ") + ")")
return nil
}
// isEmptyFilterValue reports whether a filter value means "match nothing":
// nil, empty string/collection, false, or zero.
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()
// Python truthiness also drops numeric zero: the connector's
// `if ... or not v: continue` skipped 0 and 0.0 conditions.
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
return rv.Int() == 0
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
return rv.Uint() == 0
case reflect.Float32, reflect.Float64:
return rv.Float() == 0
}
return false
}
// sortedKeys lists a map's keys in sorted order for deterministic SQL.
func sortedKeys(source map[string]interface{}) []string {
keys := make([]string, 0, len(source))
for key := range source {
keys = append(keys, key)
}
sort.Strings(keys)
return keys
}
+354
View File
@@ -0,0 +1,354 @@
//
// 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 vastbase
import (
"reflect"
"strings"
"testing"
)
// testColumns builds the column metadata fixture the filter tests filter
// against.
func testColumns() map[string]columnMeta {
columns := make(map[string]columnMeta)
for _, column := range chunkColumns {
dataType := column.typeSQL
columns[column.name] = columnMeta{name: column.name, dataType: dataType}
}
return columns
}
// TestBuildFilterBasicsAndPlaceholderOrdering checks simple predicates and
// that placeholder numbering follows sorted field order.
func TestBuildFilterBasicsAndPlaceholderOrdering(t *testing.T) {
columns := testColumns()
sqlText, args, err := buildFilter(map[string]interface{}{
"doc_id": "doc-1",
"kb_id": "kb-9",
"available_int": 0,
}, "chunk", columns)
if err != nil {
t.Fatal(err)
}
// available_int=0 is falsy in the Python semantics and must drop.
if sqlText != `"doc_id" = $1 AND "kb_id" = $2` {
t.Fatalf("filter = %q", sqlText)
}
if !reflect.DeepEqual(args, []interface{}{"doc-1", "kb-9"}) {
t.Fatalf("args = %#v", args)
}
}
// TestBuildFilterListBecomesIN checks that a list value becomes an IN
// predicate with one placeholder per element.
func TestBuildFilterListBecomesIN(t *testing.T) {
sqlText, args, err := buildFilter(map[string]interface{}{
"doc_id": []interface{}{"d1", "d2", "d3"},
}, "chunk", testColumns())
if err != nil {
t.Fatal(err)
}
if sqlText != `"doc_id" IN ($1, $2, $3)` {
t.Fatalf("filter = %q", sqlText)
}
if !reflect.DeepEqual(args, []interface{}{"d1", "d2", "d3"}) {
t.Fatalf("args = %#v", args)
}
}
// TestBuildFilterUnknownColumnMatchesNothing checks that filtering a column
// the table never stored matches no rows.
func TestBuildFilterUnknownColumnMatchesNothing(t *testing.T) {
// ES dynamic mapping parity: a field the table never stored exists on no
// row, so the condition collapses to 1=0 instead of erroring.
sqlText, args, err := buildFilter(map[string]interface{}{"never_stored": "x"}, "chunk", testColumns())
if err != nil {
t.Fatal(err)
}
if sqlText != "1=0" || len(args) != 0 {
t.Fatalf("filter = %q args = %#v", sqlText, args)
}
// A nil columns map means "unknown schema" (callers that cannot probe);
// every column is then allowed through.
sqlText, _, err = buildFilter(map[string]interface{}{"never_stored": "x"}, "chunk", nil)
if err != nil {
t.Fatal(err)
}
if sqlText != `"never_stored" = $1` {
t.Fatalf("permissive filter = %q", sqlText)
}
}
// TestBuildFilterKeywordConditionsMatchJoinedValues checks equality against
// the separator-joined keyword storage, with padding on both sides.
func TestBuildFilterKeywordConditionsMatchJoinedValues(t *testing.T) {
// Keyword columns store ###-joined lists: term conditions translate to
// LIKE containment on the separator-padded column (dropping them would
// collapse a write filter to its dataset scope).
sqlText, args, err := buildFilter(map[string]interface{}{
"source_id": "doc-1",
"important_kwd": "alpha",
"docnm_kwd": "title",
}, "chunk", testColumns())
if err != nil {
t.Fatal(err)
}
want := `"docnm_kwd" = $1 AND (('###' || "important_kwd" || '###') LIKE $2) AND (('###' || "source_id" || '###') LIKE $3)`
if sqlText != want {
t.Fatalf("filter = %q", sqlText)
}
if !reflect.DeepEqual(args, []interface{}{"title", "%###alpha###%", "%###doc-1###%"}) {
t.Fatalf("args = %#v", args)
}
}
// TestBuildFilterKeywordTermsListORsAndEscapes checks that a keyword terms
// list becomes an OR of LIKE predicates with escaped wildcards.
func TestBuildFilterKeywordTermsListORsAndEscapes(t *testing.T) {
sqlText, args, err := buildFilter(map[string]interface{}{
"tag_kwd": []interface{}{"red", "blue%ish", "gr_en"},
}, "chunk", testColumns())
if err != nil {
t.Fatal(err)
}
want := `(('###' || "tag_kwd" || '###') LIKE $1 OR ('###' || "tag_kwd" || '###') LIKE $2 OR ('###' || "tag_kwd" || '###') LIKE $3)`
if sqlText != want {
t.Fatalf("filter = %q", sqlText)
}
// LIKE metacharacters in the values are escaped so they match literally.
if !reflect.DeepEqual(args, []interface{}{"%###red###%", `%###blue\%ish###%`, `%###gr\_en###%`}) {
t.Fatalf("args = %#v", args)
}
// A terms list whose every element is empty matches no row: the decode
// side never yields an empty element.
sqlText, args, err = buildFilter(map[string]interface{}{"tag_kwd": []interface{}{"", ""}}, "chunk", testColumns())
if err != nil || sqlText != "1=0" || len(args) != 0 {
t.Fatalf("empty-element filter = %q args = %#v err %v", sqlText, args, err)
}
}
// TestBuildFilterUnknownKeywordColumnMatchesNothing checks that keyword
// terms against a never-stored column match no rows.
func TestBuildFilterUnknownKeywordColumnMatchesNothing(t *testing.T) {
sqlText, args, err := buildFilter(map[string]interface{}{"ghost_kwd": "x"}, "chunk", testColumns())
if err != nil {
t.Fatal(err)
}
if sqlText != "1=0" || len(args) != 0 {
t.Fatalf("filter = %q args = %#v", sqlText, args)
}
}
// TestBuildFilterExistsParity checks the exists filter against stored and
// never-stored columns.
func TestBuildFilterExistsParity(t *testing.T) {
columns := testColumns()
sqlText, _, err := buildFilter(map[string]interface{}{
"exists": "title_tks",
}, "chunk", columns)
if err != nil {
t.Fatal(err)
}
if sqlText != `("title_tks" IS NOT NULL AND "title_tks" <> '')` {
t.Fatalf("text exists = %q", sqlText)
}
sqlText, _, err = buildFilter(map[string]interface{}{
"exists": "pagerank_fea",
}, "chunk", columns)
if err != nil {
t.Fatal(err)
}
if sqlText != `"pagerank_fea" IS NOT NULL` {
t.Fatalf("scalar exists = %q", sqlText)
}
// exists on a column the table lacks matches nothing.
sqlText, _, _ = buildFilter(map[string]interface{}{"exists": "ghost_field"}, "chunk", columns)
if sqlText != "1=0" {
t.Fatalf("unknown exists = %q", sqlText)
}
}
// TestBuildFilterMustNotExistsMirror checks that must_not exists inverts the
// exists predicate.
func TestBuildFilterMustNotExistsMirror(t *testing.T) {
columns := testColumns()
sqlText, _, err := buildFilter(map[string]interface{}{
"must_not": map[string]interface{}{"exists": "forget_at"},
}, "chunk", columns)
if err != nil {
t.Fatal(err)
}
// forget_at is absent from chunk tables: the clause holds on every row
// and is dropped entirely.
if sqlText != "1=1" {
t.Fatalf("unknown must_not = %q", sqlText)
}
memory := make(map[string]columnMeta)
for _, column := range memoryColumns {
memory[column.name] = columnMeta{name: column.name, dataType: column.typeSQL}
}
sqlText, _, err = buildFilter(map[string]interface{}{
"must_not": map[string]interface{}{"exists": "forget_at"},
}, "memory", memory)
if err != nil {
t.Fatal(err)
}
if sqlText != `("forget_at" IS NULL OR "forget_at" = '')` {
t.Fatalf("text must_not = %q", sqlText)
}
sqlText, _, err = buildFilter(map[string]interface{}{
"must_not": map[string]interface{}{"exists": "status_int"},
}, "memory", memory)
if err != nil {
t.Fatal(err)
}
if sqlText != `"status_int" IS NULL` {
t.Fatalf("scalar must_not = %q", sqlText)
}
}
// TestBuildFilterMemoryFieldMapping checks that logical memory field names
// filter their mapped columns.
func TestBuildFilterMemoryFieldMapping(t *testing.T) {
memory := make(map[string]columnMeta)
for _, column := range memoryColumns {
memory[column.name] = columnMeta{name: column.name, dataType: column.typeSQL}
}
sqlText, args, err := buildFilter(map[string]interface{}{
"message_type": "raw",
"id": "m1_1",
}, "memory", memory)
if err != nil {
t.Fatal(err)
}
// message_type maps to message_type_kwd; memory stores that column as a
// plain scalar, so it keeps plain equality (no ### padding).
if sqlText != `"id" = $1 AND "message_type_kwd" = $2` {
t.Fatalf("memory filter = %q", sqlText)
}
if !reflect.DeepEqual(args, []interface{}{"m1_1", "raw"}) {
t.Fatalf("args = %#v", args)
}
}
// TestBuildFilterIDKeyRejected checks that the reserved "id" key is refused
// rather than treated as a column.
func TestBuildFilterIDKeyRejected(t *testing.T) {
if _, _, err := buildFilter(map[string]interface{}{"_id": "x"}, "chunk", testColumns()); err == nil {
t.Fatal("_id must be rejected in favor of id")
}
}
// TestBuildFilterEmptyAndFalsy checks that empty and falsy filter values
// match nothing instead of matching everything.
func TestBuildFilterEmptyAndFalsy(t *testing.T) {
sqlText, args, err := buildFilter(map[string]interface{}{}, "chunk", testColumns())
if err != nil || sqlText != "1=1" || len(args) != 0 {
t.Fatalf("empty filter = %q %#v err %v", sqlText, args, err)
}
sqlText, _, err = buildFilter(map[string]interface{}{"doc_id": "", "removed_kwd": nil, "tag_kwd": []interface{}{}}, "chunk", testColumns())
if err != nil || sqlText != "1=1" {
t.Fatalf("falsy filter = %q err %v", sqlText, err)
}
}
// TestShiftPlaceholders checks $n renumbering when a fragment joins a
// statement that already consumed arguments.
func TestShiftPlaceholders(t *testing.T) {
if got := shiftPlaceholders(`"a" = $1 AND "b" IN ($2, $3)`, 3); got != `"a" = $4 AND "b" IN ($5, $6)` {
t.Fatalf("shift = %q", got)
}
if got := shiftPlaceholders("$1", 0); got != "$1" {
t.Fatalf("zero offset must be a no-op, got %q", got)
}
}
// TestBuildFilterMetadataConditions checks jsonb operators rendered for the
// metadata post-filter conditions.
func TestBuildFilterMetadataConditions(t *testing.T) {
sqlText, args, err := buildFilter(map[string]interface{}{
"metadata_filtering_conditions": map[string]interface{}{
"logical_operator": "and",
"conditions": []interface{}{
map[string]interface{}{"name": "year", "comparison_operator": ">", "value": 2023},
map[string]interface{}{"name": "author", "comparison_operator": "is", "value": "Ada"},
map[string]interface{}{"name": "title", "comparison_operator": "start with", "value": "RAG"},
},
},
}, "chunk", testColumns())
if err != nil {
t.Fatal(err)
}
want := `(CASE WHEN "metadata" LIKE '{%' THEN "metadata"::jsonb ELSE '{}'::jsonb END) ->> $1`
if !strings.Contains(sqlText, want) {
t.Fatalf("json extraction missing: %q", sqlText)
}
// Guarded numeric cast keeps non-numeric rows as NULL instead of failing.
if !strings.Contains(sqlText, `THEN ((CASE WHEN "metadata" LIKE '{%' THEN "metadata"::jsonb ELSE '{}'::jsonb END) ->> $1)::numeric ELSE NULL END`) {
t.Fatalf("guarded cast missing: %q", sqlText)
}
if !strings.Contains(sqlText, `LIKE $6`) || !strings.Contains(sqlText, `> $2`) || !strings.Contains(sqlText, `= $4`) {
t.Fatalf("operators/args misplaced: %q", sqlText)
}
if len(args) != 6 {
t.Fatalf("args = %#v", args)
}
if args[0] != "year" || args[2] != "author" {
t.Fatalf("json key arguments = %#v", args)
}
}
// TestBuildFilterMetadataContainsUsesJSONBContainment checks that a metadata
// "contains" over an object becomes jsonb containment.
func TestBuildFilterMetadataContainsUsesJSONBContainment(t *testing.T) {
sqlText, args, err := buildFilter(map[string]interface{}{
"metadata_filtering_conditions": map[string]interface{}{
"logical_operator": "or",
"conditions": []interface{}{
map[string]interface{}{"name": "tags", "comparison_operator": "contains", "value": "hot"},
map[string]interface{}{"name": "tags", "comparison_operator": "not contains", "value": "cold"},
},
},
}, "chunk", testColumns())
if err != nil {
t.Fatal(err)
}
if !strings.Contains(sqlText, `-> $1 @> $2::jsonb`) || !strings.Contains(sqlText, `NOT (CASE WHEN`) || !strings.Contains(sqlText, ` OR `) {
t.Fatalf("contains filter = %q", sqlText)
}
if args[1] != `"hot"` {
t.Fatalf("contains value must be JSON-encoded for @>, got %#v", args[1])
}
}
// TestBuildFilterMetadataEmptyChecks checks empty metadata operand handling.
func TestBuildFilterMetadataEmptyChecks(t *testing.T) {
sqlText, _, err := buildFilter(map[string]interface{}{
"metadata_filtering_conditions": map[string]interface{}{
"conditions": []interface{}{
map[string]interface{}{"name": "note", "comparison_operator": "empty"},
},
},
}, "chunk", testColumns())
if err != nil {
t.Fatal(err)
}
if !strings.Contains(sqlText, `IS NULL OR`) || !strings.Contains(sqlText, `= '{}'`) {
t.Fatalf("empty filter = %q", sqlText)
}
}
+188
View File
@@ -0,0 +1,188 @@
//
// 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 vastbase
import (
"context"
"fmt"
"math"
"sort"
"strings"
)
// GetFields projects selected fields out of search chunks, filling missing
// fields with nil — the get_fields contract that lets callers ask for fields
// a table does not carry.
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
}
// GetAggregation counts distinct values of a field across chunks.
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 {
if text, ok := item.(string); 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
}
// GetHighlight marks query keywords in the chunk text client-side; Vastbase
// provides no server-side highlight.
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
}
// GetChunkIDs collects chunk or skill ids from result rows.
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 the source fields for Vastbase.
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
}
// GetScores extracts the per-hit scores from a KNN result envelope.
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
}
// cosineSimilarity computes the cosine of two equally sized vectors,
// yielding 0 for mismatched or empty inputs.
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))
}
+214
View File
@@ -0,0 +1,214 @@
//
// 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 vastbase
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)
}
// normalizeKeywords trims, lowercases, and deduplicates highlight keywords.
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
}
// isMostlyEnglish reports whether Latin letters dominate, which selects the
// word-boundary highlight strategy over CJK substring matching.
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
}
// findPatternMatches collects regex match spans, optionally keeping only
// those with word boundaries on both sides.
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)
}
// findTokenMatches locates keyword spans in text by scanning its tokenized
// form right-to-left, so each token maps to the occurrence nearest its
// neighbors.
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)
}
// hasWordBoundaries reports whether a match starts and ends on non-word
// characters (or the text edges).
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
}
// isWordRune reports whether a rune participates in a word for boundary
// checks.
func isWordRune(r rune) bool {
return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_'
}
// selectMatches picks non-overlapping spans from the candidates, preferring
// earlier starts and longer matches.
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
}
// applyMatches wraps the selected spans in <em> markers.
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()
}
@@ -0,0 +1,315 @@
//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 vastbase
import (
"context"
"os"
"reflect"
"strconv"
"strings"
"testing"
"time"
"ragflow/internal/engine/types"
"ragflow/internal/server/config"
)
// Run against a live Vastbase G100 instance:
//
// RAGFLOW_TEST_VASTBASE_HOST=... go test -tags integration ./internal/engine/vastbase/
//
// DBCOMPATIBILITY selects the server mode under test — the ParadeDB bm25 path
// (B) or the GIN to_tsvector path (PG). B is the default: it is the only mode
// the Python connector verified end to end.
func vbITEnvOr(name, fallback string) string {
if value := strings.TrimSpace(os.Getenv(name)); value != "" {
return value
}
return fallback
}
// TestVastbaseStorageRoundTrip exercises the full storage surface against a
// live Vastbase instance: DDL, writes, searches, and cleanup.
func TestVastbaseStorageRoundTrip(t *testing.T) {
host := vbITEnvOr("RAGFLOW_TEST_VASTBASE_HOST", "")
if host == "" {
t.Skip("RAGFLOW_TEST_VASTBASE_HOST is not set")
}
port, err := strconv.Atoi(vbITEnvOr("RAGFLOW_TEST_VASTBASE_PORT", "5432"))
if err != nil {
t.Fatal(err)
}
compatibility := vbITEnvOr("RAGFLOW_TEST_VASTBASE_DBCOMPATIBILITY", "B")
engine, err := NewEngine(config.VastbaseConfig{
Host: host, Port: port,
User: vbITEnvOr("RAGFLOW_TEST_VASTBASE_USER", "ragflow"),
Password: vbITEnvOr("RAGFLOW_TEST_VASTBASE_PASSWORD", ""),
DBName: vbITEnvOr("RAGFLOW_TEST_VASTBASE_DBNAME", "ragflow"),
DBCompatibility: compatibility,
SSLMode: vbITEnvOr("RAGFLOW_TEST_VASTBASE_SSL_MODE", "disable"),
})
if err != nil {
t.Fatal(err)
}
defer engine.Close()
ctx, cancel := context.WithTimeout(t.Context(), 5*time.Minute)
defer cancel()
if err := engine.Ping(ctx); err != nil {
t.Fatalf("ping vastbase (dbcompatibility %s): %v", compatibility, err)
}
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
table := "vb_it_" + suffix
dataset := "kb-" + suffix
defer engine.DropChunkStore(ctx, table, "")
// --- Chunk store: DDL, dynamic columns, storage-format round trip ---
if err := engine.CreateChunkStore(ctx, table, dataset, 3, "naive"); err != nil {
t.Fatal(err)
}
exists, err := engine.ChunkStoreExists(ctx, table, dataset)
if err != nil || !exists {
t.Fatalf("chunk store exists = %v err = %v", exists, err)
}
if _, err := engine.InsertChunks(ctx, []map[string]interface{}{
{
"id": "c1", "kb_id": dataset, "doc_id": "doc-1",
"content_with_weight": "hello vastbase world", "content_ltks": "hello vastbase world",
"important_kwd": []interface{}{"hello", "vastbase"},
"metadata": map[string]interface{}{"_group_id": "g1"},
"position_int": []interface{}{[]interface{}{int64(1), 2, 3, 4, 5}},
"page_num_int": []interface{}{int64(1)},
"custom_field": "dynamic column value",
"q_3_vec": []interface{}{0.1, 0.2, 0.3},
},
{
"id": "c2", "kb_id": dataset, "doc_id": "doc-1",
"content_with_weight": "second chunk", "content_ltks": "second chunk",
"available_int": 1, "q_3_vec": []interface{}{0.9, 0.8, 0.7},
},
}, table, dataset); err != nil {
t.Fatal(err)
}
row, err := engine.GetChunk(ctx, table, "c1", []string{dataset})
if err != nil {
t.Fatal(err)
}
chunk := row.(map[string]interface{})
if !reflect.DeepEqual(chunk["important_kwd"], []string{"hello", "vastbase"}) {
t.Fatalf("keyword round trip = %#v", chunk["important_kwd"])
}
positions, ok := chunk["position_int"].([][]int64)
if !ok || len(positions) != 1 || len(positions[0]) != 5 || positions[0][0] != 1 {
t.Fatalf("position_int regroup = %#v", chunk["position_int"])
}
if vector, ok := chunk["q_3_vec"].([]float64); !ok || len(vector) != 3 {
t.Fatalf("vector round trip = %#v", chunk["q_3_vec"])
}
if chunk["custom_field"] != "dynamic column value" {
t.Fatalf("dynamic column round trip = %#v", chunk["custom_field"])
}
if groupID, _ := chunk["metadata"].(map[string]interface{}); groupID["_group_id"] != "g1" {
t.Fatalf("metadata JSON round trip = %#v", chunk["metadata"])
}
// --- Four retrieval paths ---
textResult, err := engine.Search(ctx, &types.SearchRequest{
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
SelectFields: []string{"id", "content_ltks"},
MatchExprs: []interface{}{&types.MatchTextExpr{
Fields: []string{"content_ltks"}, MatchingText: "vastbase", TopN: 10,
}},
})
if err != nil {
t.Fatalf("fulltext search: %v", err)
}
if len(textResult.Chunks) != 1 || textResult.Chunks[0]["id"] != "c1" {
t.Fatalf("fulltext hits = %#v", textResult.Chunks)
}
vectorResult, err := engine.Search(ctx, &types.SearchRequest{
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
SelectFields: []string{"id"},
MatchExprs: []interface{}{&types.MatchDenseExpr{
VectorColumnName: "q_3_vec", EmbeddingData: []float64{0.1, 0.2, 0.3},
EmbeddingDataType: "float", TopN: 10,
}},
})
if err != nil {
t.Fatalf("vector search: %v", err)
}
if len(vectorResult.Chunks) != 2 || vectorResult.Chunks[0]["id"] != "c1" {
t.Fatalf("knn order = %#v", vectorResult.Chunks)
}
fusionResult, err := engine.Search(ctx, &types.SearchRequest{
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
SelectFields: []string{"id"},
MatchExprs: []interface{}{
&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "vastbase", TopN: 10},
&types.MatchDenseExpr{VectorColumnName: "q_3_vec", EmbeddingData: []float64{0.1, 0.2, 0.3},
EmbeddingDataType: "float", TopN: 10},
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
FusionParams: map[string]interface{}{"weights": "0.5,0.5"}},
},
})
if err != nil {
t.Fatalf("fusion search: %v", err)
}
if len(fusionResult.Chunks) != 2 || fusionResult.Chunks[0]["id"] != "c1" {
t.Fatalf("fusion order = %#v", fusionResult.Chunks)
}
filterResult, err := engine.Search(ctx, &types.SearchRequest{
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
SelectFields: []string{"id"}, Filter: map[string]interface{}{"doc_id": "doc-1"},
})
if err != nil {
t.Fatalf("filter search: %v", err)
}
if filterResult.Total != 2 || len(filterResult.Chunks) != 2 {
t.Fatalf("filter total = %d chunks = %#v", filterResult.Total, filterResult.Chunks)
}
// --- Update + pagerank + delete ---
if err := engine.UpdateChunks(ctx, map[string]interface{}{"id": "c2"},
map[string]interface{}{"important_kwd": []interface{}{"tagged"}}, table, dataset); err != nil {
t.Fatal(err)
}
if err := engine.AdjustChunkPagerank(ctx, table, "c1", dataset, 1.5, 0, 10); err != nil {
t.Fatal(err)
}
row, err = engine.GetChunk(ctx, table, "c1", []string{dataset})
if err != nil {
t.Fatal(err)
}
if pagerank, _ := numberToFloat(row.(map[string]interface{})["pagerank_fea"]); pagerank != 1.5 {
t.Fatalf("pagerank after adjust = %#v", row.(map[string]interface{})["pagerank_fea"])
}
deleted, err := engine.DeleteChunks(ctx, map[string]interface{}{"id": "c2"}, table, dataset)
if err != nil || deleted != 1 {
t.Fatalf("delete = %d err = %v", deleted, err)
}
// --- Memory store: tsquery fulltext + fusion ---
memoryTable := "memory_vb_it_" + suffix
memoryID := "mem-" + suffix
defer engine.DropChunkStore(ctx, memoryTable, "")
if err := engine.CreateChunkStore(ctx, memoryTable, memoryID, 0, ""); err != nil {
t.Fatal(err)
}
if _, err := engine.InsertChunks(ctx, []map[string]interface{}{
{"id": memoryID + "_1", "content": "remember the vastbase meeting", "status": true,
"content_embed": []interface{}{0.25, 0.5}},
}, memoryTable, memoryID); err != nil {
t.Fatal(err)
}
memoryText, err := engine.Search(ctx, &types.SearchRequest{
IndexNames: []string{memoryTable}, KbIDs: []string{memoryID}, Limit: 10,
SelectFields: []string{"content", "status"},
MatchExprs: []interface{}{&types.MatchTextExpr{MatchingText: "vastbase meeting"}},
})
if err != nil {
t.Fatalf("memory fulltext: %v", err)
}
if len(memoryText.Chunks) != 1 || memoryText.Chunks[0]["content"] != "remember the vastbase meeting" {
t.Fatalf("memory fulltext hits = %#v", memoryText.Chunks)
}
memoryFusion, err := engine.Search(ctx, &types.SearchRequest{
IndexNames: []string{memoryTable}, KbIDs: []string{memoryID}, Limit: 10,
SelectFields: []string{"content"},
MatchExprs: []interface{}{
&types.MatchTextExpr{MatchingText: "vastbase", TopN: 10},
&types.MatchDenseExpr{VectorColumnName: "q_2_vec", EmbeddingData: []float64{0.25, 0.5},
EmbeddingDataType: "float", TopN: 10},
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
FusionParams: map[string]interface{}{"weights": "0.5,0.5"}},
},
})
if err != nil {
t.Fatalf("memory fusion: %v", err)
}
if len(memoryFusion.Chunks) != 1 {
t.Fatalf("memory fusion hits = %#v", memoryFusion.Chunks)
}
// --- Metadata store CRUD ---
tenant := "it" + suffix
defer engine.DropMetadataStore(ctx, tenant)
if err := engine.CreateMetadataStore(ctx, tenant); err != nil {
t.Fatal(err)
}
if _, err := engine.InsertMetadata(ctx, []map[string]interface{}{
{"id": "md-1", "kb_id": dataset, "doc_id": "doc-1",
"meta_fields": map[string]interface{}{"year": 2026, "author": "it"}},
}, tenant); err != nil {
t.Fatal(err)
}
metaResult, err := engine.SearchMetadata(ctx, &types.SearchMetadataRequest{
TenantID: tenant, Limit: 10, Filter: map[string]interface{}{"kb_id": dataset},
})
if err != nil {
t.Fatalf("metadata search: %v", err)
}
if metaResult.Total != 1 || len(metaResult.MetadataRecords) != 1 {
t.Fatalf("metadata search result = %#v", metaResult)
}
if err := engine.UpdateMetadata(ctx, "md-1", dataset,
map[string]interface{}{"year": 2026, "author": "updated"}, tenant); err != nil {
t.Fatal(err)
}
metaResult, err = engine.SearchMetadata(ctx, &types.SearchMetadataRequest{
TenantID: tenant, Limit: 10, Filter: map[string]interface{}{"doc_id": "doc-1"},
})
if err != nil {
t.Fatal(err)
}
fields, _ := metaResult.MetadataRecords[0]["meta_fields"].(map[string]interface{})
if fields["author"] != "updated" {
t.Fatalf("metadata update = %#v", metaResult.MetadataRecords[0])
}
if _, err := engine.DeleteMetadata(ctx, map[string]interface{}{"id": "md-1"}, tenant); err != nil {
t.Fatal(err)
}
// --- Drop deletes dataset rows but keeps the shared table ---
remaining, err := engine.InsertChunks(ctx, []map[string]interface{}{
{"id": "c3", "kb_id": dataset, "content_ltks": "third"},
}, table, dataset)
if err != nil {
t.Fatal(err)
}
_ = remaining
if err := engine.DropChunkStore(ctx, table, dataset); err != nil {
t.Fatal(err)
}
filterResult, err = engine.Search(ctx, &types.SearchRequest{
IndexNames: []string{table}, KbIDs: []string{dataset}, Limit: 10,
SelectFields: []string{"id"},
MatchExprs: []interface{}{&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "third"}},
})
if err != nil {
t.Fatal(err)
}
if len(filterResult.Chunks) != 0 {
t.Fatalf("rows survived dataset drop: %#v", filterResult.Chunks)
}
}
+330
View File
@@ -0,0 +1,330 @@
//
// 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 vastbase
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"sort"
"ragflow/internal/engine/types"
"gorm.io/gorm"
)
// metadataTableName is the per-tenant metadata table name.
func metadataTableName(tenantID string) string { return "ragflow_doc_meta_" + tenantID }
// validatedMetadataTableName maps a tenant to its metadata table name,
// rejecting tenant IDs the identifier pattern refuses.
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 with its kb_id
// index. Tables are per-tenant, so kb_id is a plain attribute column here,
// not a shared-table discriminator.
func (e *Engine) CreateMetadataStore(ctx context.Context, tenantID string) error {
tableName, err := validatedMetadataTableName(tenantID)
if err != nil {
return err
}
if err := e.ensureTable(ctx, tableName, metadataColumns); err != nil {
return err
}
return e.ensureRegularIndex(ctx, tableName, "kb_id")
}
// InsertMetadata stores metadata rows with the same delete-then-insert
// replacement the chunk path uses (id is the primary key).
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
}
}
rows := make([]map[string]interface{}, 0, len(metadata))
ids := make([]string, 0, len(metadata))
for _, document := range metadata {
row := make(map[string]interface{}, len(metadataColumns))
for _, column := range metadataColumns {
value, present := document[column.name]
if !present {
continue
}
if column.name == "meta_fields" {
encoded, err := encodeMetaFields(value)
if err != nil {
return nil, err
}
row[column.name] = encoded
continue
}
encoded, err := encodeColumnValue(column.name, value)
if err != nil {
return nil, err
}
row[column.name] = encoded
}
if stringValue(row["id"]) == "" {
return nil, fmt.Errorf("metadata document is missing id")
}
ids = append(ids, stringValue(row["id"]))
rows = append(rows, row)
}
tx, err := e.db.BeginTx(ctx, nil)
if err != nil {
return nil, err
}
defer tx.Rollback()
if len(ids) > 0 {
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE id IN (%s)",
quoteIdent(tableName), placeholderList(len(ids), 0)), stringArgs(ids)...); err != nil {
return nil, fmt.Errorf("vastbase: delete before metadata insert: %w", err)
}
}
for _, group := range groupByColumnSet(rows) {
if err := insertGroup(ctx, tx, tableName, group); err != nil {
return nil, err
}
}
if err := tx.Commit(); err != nil {
return nil, err
}
return []string{}, nil
}
// encodeMetaFields serializes a metadata value for the meta_fields column,
// passing strings through and JSON-encoding objects.
func encodeMetaFields(value interface{}) (string, error) {
switch typed := value.(type) {
case string:
return typed, nil
case map[string]interface{}:
encoded, err := json.Marshal(typed)
if err != nil {
return "", err
}
return string(encoded), nil
default:
return "{}", nil
}
}
// UpdateMetadata replaces the complete meta_fields JSON object, inserting the
// row when it does not yet exist — 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 := encodeMetaFields(metaFields)
if err != nil {
return err
}
tx, err := e.db.BeginTx(ctx, nil)
if err != nil {
return err
}
defer tx.Rollback()
result, err := tx.ExecContext(ctx, fmt.Sprintf("UPDATE %s SET meta_fields = $1 WHERE id = $2", quoteIdent(tableName)), encoded, docID)
if err != nil {
return fmt.Errorf("vastbase: replace metadata: %w", err)
}
affected, err := result.RowsAffected()
if err != nil {
return fmt.Errorf("vastbase: replace metadata: %w", err)
}
if affected == 0 {
if _, err := tx.ExecContext(ctx, fmt.Sprintf("INSERT INTO %s (id, kb_id, doc_id, meta_fields) VALUES ($1, $2, $3, $4)", quoteIdent(tableName)), docID, datasetID, docID, encoded); err != nil {
return fmt.Errorf("vastbase: replace metadata: %w", err)
}
}
return tx.Commit()
}
// 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
}
columns, err := e.listTableColumns(ctx, tableName)
if err != nil {
return 0, err
}
whereSQL, args, err := buildFilter(condition, "metadata", columns)
if err != nil {
return 0, err
}
result, err := e.db.ExecContext(ctx, "DELETE FROM "+quoteIdent(tableName)+" WHERE "+whereSQL, args...)
if err != nil {
return 0, fmt.Errorf("vastbase: delete metadata: %w", err)
}
return result.RowsAffected()
}
// DeleteMetadataKeys removes selected JSON keys and deletes the row when 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,
fmt.Sprintf("SELECT meta_fields FROM %s WHERE id = $1 AND kb_id = $2 LIMIT 1", quoteIdent(tableName)),
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, fmt.Sprintf("DELETE FROM %s WHERE id = $1 AND kb_id = $2", quoteIdent(tableName)), docID, datasetID)
return err
}
encoded, err := json.Marshal(fields)
if err != nil {
return err
}
_, err = e.db.ExecContext(ctx, fmt.Sprintf("UPDATE %s SET meta_fields = $1 WHERE id = $2 AND kb_id = $3", quoteIdent(tableName)), string(encoded), docID, datasetID)
return err
}
// DropMetadataStore drops a tenant's metadata table.
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 "+quoteIdent(tableName))
return err
}
// MetadataStoreExists reports whether the tenant's metadata table exists.
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 an 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
}
columns, err := e.listTableColumns(ctx, tableName)
if err != nil {
return nil, err
}
output := buildMetadataOutput(req.SelectFields, columns)
whereSQL, args, err := buildFilter(req.Filter, "metadata", columns)
if err != nil {
return nil, err
}
var total int64
if err := e.db.QueryRowContext(ctx,
fmt.Sprintf("SELECT COUNT(id) FROM %s WHERE %s", quoteIdent(tableName), whereSQL),
args...).Scan(&total); err != nil {
return nil, fmt.Errorf("vastbase: count metadata %s: %w", tableName, err)
}
if total == 0 {
return &types.SearchMetadataResult{MetadataRecords: []map[string]interface{}{}, Total: total}, nil
}
orderSQL, err := buildOrderBy(req.OrderBy, "metadata", columns)
if err != nil {
return nil, err
}
query := fmt.Sprintf("SELECT %s FROM %s WHERE %s%s LIMIT %d OFFSET %d",
selectFieldsSQL(output), quoteIdent(tableName), whereSQL, orderSQL, max(req.Offset, 0), positiveOr(req.Limit, 30))
rows, err := e.queryRows(ctx, query, args...)
if err != nil {
return nil, fmt.Errorf("vastbase: search metadata %s: %w", tableName, err)
}
return &types.SearchMetadataResult{MetadataRecords: decodeRows(rows, "metadata"), Total: total}, nil
}
// buildMetadataOutput projects the requested fields onto the live columns; a
// nil field list selects every column.
func buildMetadataOutput(selectFields []string, columns map[string]columnMeta) []string {
if len(selectFields) == 0 {
names := make([]string, 0, len(columns))
for column := range columns {
names = append(names, column)
}
sort.Strings(names)
return names
}
output := make([]string, 0, len(selectFields))
for _, field := range selectFields {
if _, exists := columns[field]; exists && !containsString(output, field) {
output = append(output, field)
}
}
return output
}
// FilterDocIdsByMetaPushdown is not implemented: meta_fields is stored as
// text, so a push-down predicate would cast the whole table to jsonb. The
// caller falls back to in-memory filtering, which is correct and bounded by
// the metadata page size.
func (e *Engine) FilterDocIdsByMetaPushdown(ctx context.Context, sqlDB *gorm.DB, kbIDs []string, conditions []map[string]interface{}, logic string) []string {
return nil
}
+81
View File
@@ -0,0 +1,81 @@
//
// 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 vastbase
import (
"context"
"regexp"
"testing"
"ragflow/internal/server/config"
"github.com/DATA-DOG/go-sqlmock"
)
// An existing row is updated in place: doc_id and create_time survive the
// replace, and no INSERT runs.
func TestUpdateMetadataUpdatesExistingRowInPlace(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
mock.ExpectBegin()
mock.ExpectExec(regexp.QuoteMeta(`UPDATE "ragflow_doc_meta_t1" SET meta_fields = $1 WHERE id = $2`)).
WithArgs(`{"k":"v"}`, "md-1").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
if err := engine.UpdateMetadata(context.Background(), "md-1", "kb-1",
map[string]interface{}{"k": "v"}, "t1"); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// A missing row falls back to INSERT, which must carry doc_id so later
// doc_id-filtered lookups still find the metadata.
func TestUpdateMetadataInsertsMissingRowWithDocID(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
mock.ExpectBegin()
mock.ExpectExec(regexp.QuoteMeta(`UPDATE "ragflow_doc_meta_t1" SET meta_fields = $1 WHERE id = $2`)).
WithArgs(`{"k":"v"}`, "md-2").
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectExec(regexp.QuoteMeta(
`INSERT INTO "ragflow_doc_meta_t1" (id, kb_id, doc_id, meta_fields) VALUES ($1, $2, $3, $4)`)).
WithArgs("md-2", "kb-1", "md-2", `{"k":"v"}`).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
if err := engine.UpdateMetadata(context.Background(), "md-2", "kb-1",
map[string]interface{}{"k": "v"}, "t1"); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
+281
View File
@@ -0,0 +1,281 @@
//
// 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 vastbase
import (
"context"
"database/sql"
"encoding/json"
"strconv"
"strings"
)
// queryRows runs a query and shapes each row into a map keyed by column name.
// NULL columns are omitted, matching the ES _source behavior of not
// returning fields that were never written.
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()
return scanRows(rows)
}
// queryRowsTx is queryRows bound to an open transaction.
func queryRowsTx(ctx context.Context, tx *sql.Tx, query string, args ...interface{}) ([]map[string]interface{}, error) {
rows, err := tx.QueryContext(ctx, query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanRows(rows)
}
// scanRows drains a result set into per-row maps keyed by column name.
func scanRows(rows *sql.Rows) ([]map[string]interface{}, error) {
columns, err := rows.Columns()
if err != nil {
return nil, err
}
var result []map[string]interface{}
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()
}
// decodeLogicalRow converts a physical row into domain values: memory columns
// are renamed to their logical fields, keyword columns are split on ###,
// position_int is regrouped into 5-tuples, integer arrays and JSON text are
// parsed, and floatvector text becomes a float slice (lib/pq returns unknown
// OIDs as text, so decoding rides on the q_(\d+)_vec column-name pattern).
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) {
// Multiple q_N_vec columns coexist after an embedding-model
// change; backfillVectorColumns zero-fills the stale ones, and
// GetChunk's SELECT * reads them all. A zero placeholder must
// not overwrite a real embedding: keep the first non-zero
// vector and skip zeros once content_embed is set.
if text, ok := value.(string); ok {
if vector := parseFloatVectorText(text); !isZeroVector(vector) {
result["content_embed"] = vector
continue
} else if _, set := result["content_embed"]; set {
continue
}
}
column = "content_embed"
} else if mapped, ok := memoryColumnToField[column]; ok {
column = mapped
}
}
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"
}
}
value = decodeColumnValue(kind, rawColumn, value)
result[column] = value
}
return result
}
// decodeColumnValue converts one physical column value to its logical form:
// vector text to float slices, keyword joins back to lists, JSON columns to
// maps.
func decodeColumnValue(kind, column string, value interface{}) interface{} {
if value == nil {
return nil
}
text, isText := value.(string)
switch {
case vectorColumnPattern.MatchString(column):
if isText {
return parseFloatVectorText(text)
}
// Memory keyword-named columns (message_type_kwd, source_id) store plain
// scalars — the Python memory connector read them back without the ###
// split reserved for chunk keyword lists.
case fieldKeyword(column) && kind != "memory":
if isText {
parts := strings.Split(text, keywordSeparator)
keywords := make([]string, 0, len(parts))
for _, part := range parts {
if part != "" {
keywords = append(keywords, part)
}
}
return keywords
}
case column == "position_int":
if isText {
return regroupPositions(parsePgIntArray(text))
}
case integerArrayColumns[column]:
if isText {
return parsePgIntArray(text)
}
case jsonTextColumns[column] || strings.HasSuffix(column, "_feas"):
if isText && strings.HasPrefix(strings.TrimSpace(text), "{") {
var decoded interface{}
if err := json.Unmarshal([]byte(text), &decoded); err == nil {
return decoded
}
}
}
return value
}
// regroupPositions converts the flat integer[] storage back into the original
// [[page, x1, y1, x2, y2], ...] five-element groups.
func regroupPositions(flat []int64) [][]int64 {
groups := make([][]int64, 0, (len(flat)+4)/5)
for start := 0; start < len(flat); start += 5 {
end := start + 5
if end > len(flat) {
end = len(flat)
}
groups = append(groups, flat[start:end])
}
return groups
}
// parseFloatVectorText decodes "[0.1,0.2,...]" floatvector text.
func parseFloatVectorText(text string) []float64 {
text = strings.TrimSpace(text)
text = strings.TrimPrefix(text, "[")
text = strings.TrimSuffix(text, "]")
if text == "" {
return nil
}
parts := strings.Split(text, ",")
values := make([]float64, 0, len(parts))
for _, part := range parts {
part = strings.TrimSpace(part)
if part == "" {
continue
}
number, err := strconv.ParseFloat(part, 64)
if err != nil {
return nil
}
values = append(values, number)
}
return values
}
// isZeroVector reports whether a decoded vector is the all-zero placeholder
// backfillVectorColumns writes into stale dimension columns.
func isZeroVector(values []float64) bool {
for _, value := range values {
if value != 0 {
return false
}
}
return true
}
// parsePgIntArray decodes a PostgreSQL integer[] text literal ("{1,2,3}").
func parsePgIntArray(literal string) []int64 {
parts := parsePgArray(literal)
values := make([]int64, 0, len(parts))
for _, part := range parts {
if part == "" {
continue
}
if number, err := strconv.ParseInt(strings.TrimSpace(part), 10, 64); err == nil {
values = append(values, number)
}
}
return values
}
// parsePgArray decodes a PostgreSQL text-array literal (e.g. {a,"b,c"}) into a
// slice. lib/pq returns array columns as this literal when scanned dynamically.
func parsePgArray(literal string) []string {
if len(literal) < 2 || literal[0] != '{' || literal[len(literal)-1] != '}' {
return nil
}
body := literal[1 : len(literal)-1]
if body == "" {
return []string{}
}
var out []string
var buf strings.Builder
inQuote := false
for i := 0; i < len(body); i++ {
c := body[i]
switch {
case c == '"':
if inQuote && i+1 < len(body) && body[i+1] == '"' {
buf.WriteByte('"')
i++
continue
}
inQuote = !inQuote
case c == '\\' && i+1 < len(body):
buf.WriteByte(body[i+1])
i++
case c == ',' && !inQuote:
out = append(out, buf.String())
buf.Reset()
default:
buf.WriteByte(c)
}
}
out = append(out, buf.String())
return out
}
// containsString reports whether target appears in values.
func containsString(values []string, target string) bool {
for _, value := range values {
if value == target {
return true
}
}
return false
}
+191
View File
@@ -0,0 +1,191 @@
//
// 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 vastbase
import (
"reflect"
"testing"
)
// TestDecodeColumnValueBranches checks per-column-kind decoding: vectors,
// arrays, keyword joins, and JSON.
func TestDecodeColumnValueBranches(t *testing.T) {
vector := decodeColumnValue("chunk", "q_3_vec", "[0.1,0.2,0.3]")
if !reflect.DeepEqual(vector, []float64{0.1, 0.2, 0.3}) {
t.Fatalf("vector text decode = %#v", vector)
}
keywords := decodeColumnValue("chunk", "important_kwd", "alpha###beta###")
if !reflect.DeepEqual(keywords, []string{"alpha", "beta"}) {
t.Fatalf("keyword split must drop empty parts: %#v", keywords)
}
positions := decodeColumnValue("chunk", "position_int", "{1,2,3,4,5,6,7,8,9,10}")
want := [][]int64{{1, 2, 3, 4, 5}, {6, 7, 8, 9, 10}}
if !reflect.DeepEqual(positions, want) {
t.Fatalf("position_int regroup = %#v", positions)
}
pages := decodeColumnValue("chunk", "page_num_int", "{1,2}")
if !reflect.DeepEqual(pages, []int64{1, 2}) {
t.Fatalf("integer array decode = %#v", pages)
}
metadata := decodeColumnValue("chunk", "metadata", `{"_group_id":"g1"}`)
if metadata.(map[string]interface{})["_group_id"] != "g1" {
t.Fatalf("json text decode = %#v", metadata)
}
if got := decodeColumnValue("chunk", "metadata", ""); got != "" {
t.Fatalf("non-json text must pass through, got %#v", got)
}
feas := decodeColumnValue("chunk", "tag_feas", `{"tag_a":0.5}`)
if !reflect.DeepEqual(feas, map[string]interface{}{"tag_a": 0.5}) {
t.Fatalf("*_feas json decode = %#v", feas)
}
if got := decodeColumnValue("chunk", "title_tks", 42); got != 42 {
t.Fatalf("plain values must pass through, got %#v", got)
}
}
// TestDecodeLogicalRowMemoryRenames checks memory column-to-field renaming
// on decoded rows.
func TestDecodeLogicalRowMemoryRenames(t *testing.T) {
decoded := decodeLogicalRow(map[string]interface{}{
"id": "m1_1", "message_type_kwd": "raw", "status_int": int64(1),
"content_ltks": "hi", "q_2_vec": "[0.5,0.25]", "forget_at": "",
}, "memory")
if decoded["message_type"] != "raw" || decoded["content"] != "hi" {
t.Fatalf("memory column renaming = %#v", decoded)
}
// Memory keyword-named columns are scalars: no ### split on read.
scalar := decodeLogicalRow(map[string]interface{}{"source_id": "sess-1"}, "memory")
if scalar["source_id"] != "sess-1" {
t.Fatalf("memory source_id must stay scalar: %#v", scalar["source_id"])
}
if decoded["status"] != true {
t.Fatalf("status_int must decode to bool, got %#v", decoded["status"])
}
if !reflect.DeepEqual(decoded["content_embed"], []float64{0.5, 0.25}) {
t.Fatalf("vector column must decode to content_embed: %#v", decoded["content_embed"])
}
chunk := decodeLogicalRow(map[string]interface{}{"important_kwd": "a###b"}, "chunk")
if !reflect.DeepEqual(chunk["important_kwd"], []string{"a", "b"}) {
t.Fatalf("chunk keyword decode = %#v", chunk)
}
}
// A memory table can carry several q_N_vec columns after an embedding-model
// change; backfillVectorColumns zero-fills the stale ones and GetChunk's
// SELECT * reads them all. Whichever way the map iterates, content_embed must
// be the real embedding, never the zero placeholder.
func TestDecodeLogicalRowMemorySkipsZeroVectorPlaceholders(t *testing.T) {
row := map[string]interface{}{
"id": "m1_1",
"q_2_vec": "[0.5,0.25]",
"q_4_vec": "[0,0,0,0]",
}
for i := 0; i < 50; i++ { // exercise both map iteration orders
decoded := decodeLogicalRow(row, "memory")
if !reflect.DeepEqual(decoded["content_embed"], []float64{0.5, 0.25}) {
t.Fatalf("content_embed = %#v", decoded["content_embed"])
}
}
// All-zero row: a placeholder is retained when no real vector exists
// (its dimension depends on iteration order; only zero-ness is stable).
decoded := decodeLogicalRow(map[string]interface{}{
"q_2_vec": "[0,0]", "q_4_vec": "[0,0,0,0]",
}, "memory")
placeholder, _ := decoded["content_embed"].([]float64)
if !isZeroVector(placeholder) || len(placeholder) == 0 {
t.Fatalf("all-zero content_embed = %#v", decoded["content_embed"])
}
}
// TestRegroupPositionsHandlesRemainders checks regrouping tolerates
// non-multiple-of-five array lengths.
func TestRegroupPositionsHandlesRemainders(t *testing.T) {
// The 5-tuple regroup is a lossy contract locked with the Python writer:
// a partial trailing group is preserved as-is, never padded.
got := regroupPositions([]int64{1, 2, 3, 4, 5, 6})
want := [][]int64{{1, 2, 3, 4, 5}, {6}}
if !reflect.DeepEqual(got, want) {
t.Fatalf("regroup = %#v", got)
}
if got := regroupPositions(nil); len(got) != 0 {
t.Fatalf("empty regroup = %#v", got)
}
}
// TestParseFloatVectorText checks vector text parsing, including malformed
// input.
func TestParseFloatVectorText(t *testing.T) {
if got := parseFloatVectorText("[0.1,0.2]"); !reflect.DeepEqual(got, []float64{0.1, 0.2}) {
t.Fatalf("parse = %#v", got)
}
if got := parseFloatVectorText("[]"); got != nil {
t.Fatalf("empty vector must decode nil, got %#v", got)
}
if got := parseFloatVectorText("not a vector"); got != nil {
t.Fatalf("malformed vector must decode nil, got %#v", got)
}
if got := parseFloatVectorText("[ 0.5 , 1 ]"); !reflect.DeepEqual(got, []float64{0.5, 1}) {
t.Fatalf("whitespace tolerance = %#v", got)
}
}
// TestParsePgArray checks PostgreSQL array-literal parsing.
func TestParsePgArray(t *testing.T) {
if got := parsePgArray("{a,b,c}"); !reflect.DeepEqual(got, []string{"a", "b", "c"}) {
t.Fatalf("simple = %#v", got)
}
if got := parsePgArray(`{a,"b,c",d}`); !reflect.DeepEqual(got, []string{"a", "b,c", "d"}) {
t.Fatalf("quoted comma = %#v", got)
}
if got := parsePgArray(`{"say ""hi"""}`); !reflect.DeepEqual(got, []string{`say "hi"`}) {
t.Fatalf("escaped quote = %#v", got)
}
if got := parsePgArray("{}"); len(got) != 0 {
t.Fatalf("empty = %#v", got)
}
if got := parsePgArray("not-an-array"); got != nil {
t.Fatalf("malformed = %#v", got)
}
}
// TestParsePgIntArray checks integer array-literal parsing.
func TestParsePgIntArray(t *testing.T) {
if got := parsePgIntArray("{1,2,3}"); !reflect.DeepEqual(got, []int64{1, 2, 3}) {
t.Fatalf("parse = %#v", got)
}
if got := parsePgIntArray("{1, x, 3}"); !reflect.DeepEqual(got, []int64{1, 3}) {
t.Fatalf("non-numeric entries must be skipped: %#v", got)
}
}
// TestDecodeLogicalRowToleratesNilColumns checks decoding skips NULL column
// values without KeyError-style panics.
func TestDecodeLogicalRowToleratesNilColumns(t *testing.T) {
// NULL columns are dropped by the queryRows scan layer before decoding
// (the ES _source contract: never-written fields come back absent); any
// nil that still reaches decodeLogicalRow must pass through untouched.
decoded := decodeLogicalRow(map[string]interface{}{
"id": "c1", "title_tks": nil, "content_ltks": "kept",
}, "chunk")
if decoded["title_tks"] != nil {
t.Fatalf("nil must pass through: %#v", decoded)
}
if decoded["content_ltks"] != "kept" {
t.Fatalf("kept column = %#v", decoded)
}
}
+647
View File
@@ -0,0 +1,647 @@
//
// 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 vastbase
import (
"context"
"crypto/md5" // #nosec G501 -- MD5 provides a deterministic identifier checksum, not security.
"database/sql"
"fmt"
"regexp"
"strconv"
"strings"
"sync"
"time"
"ragflow/internal/common"
"ragflow/internal/engine/kvrocks"
"go.uber.org/zap"
)
const (
maxIndexNameLength = 63
indexNameHashLength = 4
indexNameTruncationSpace = 8
)
// columnDefinition is one column of the static table schemas. defaultSQL is a
// rendered SQL literal ("" means no DEFAULT clause). Text columns default to
// ” instead of NULL so that the ES-parity exists-filter can treat an empty
// string as "field absent".
type columnDefinition struct {
name string
typeSQL string
defaultSQL string
}
// columnMeta describes a column as reported by information_schema.
type columnMeta struct {
name string
dataType string
columnDefault sql.NullString
nullable string
}
var (
identifierPattern = regexp.MustCompile(`^[A-Za-z0-9_-]+$`)
vectorDimRegex = regexp.MustCompile(`^q_(\d+)_vec$`)
textTypeRegex = regexp.MustCompile(`(?i)^(character varying|varchar|character|char|text)`)
ddlLocks sync.Map
)
// chunkColumns mirrors conf/vastbase_mapping.json from the Python connector:
// column names are the ES field names verbatim, with PG types.
var chunkColumns = []columnDefinition{
{"id", "varchar(256)", "''"},
{"kb_id", "varchar(256)", "''"},
{"doc_id", "varchar(256)", "''"},
{"doc_type_kwd", "varchar(256)", "''"},
{"docnm_kwd", "text", "''"},
{"title_tks", "text", "''"},
{"title_sm_tks", "text", "''"},
{"content_with_weight", "text", "''"},
{"content_ltks", "text", "''"},
{"content_sm_ltks", "text", "''"},
{"pagerank_fea", "double precision", "0.0"},
{"important_kwd", "text", "''"},
{"important_tks", "text", "''"},
{"question_kwd", "text", "''"},
{"question_tks", "text", "''"},
{"tag_kwd", "varchar(256)", "''"},
{"tag_feas", "text", "''"},
{"available_int", "integer", "1"},
{"create_time", "varchar(32)", "''"},
{"create_timestamp_flt", "double precision", "0.0"},
{"img_id", "varchar(128)", "''"},
{"position_int", "integer[]", ""},
{"page_num_int", "integer[]", ""},
{"top_int", "integer[]", ""},
{"knowledge_graph_kwd", "varchar(256)", "''"},
{"source_id", "text", "''"},
{"entity_kwd", "varchar(256)", "''"},
{"entity_type_kwd", "varchar(256)", "''"},
{"from_entity_kwd", "varchar(256)", "''"},
{"to_entity_kwd", "varchar(256)", "''"},
{"weight_int", "integer", "0"},
{"weight_flt", "double precision", "0.0"},
{"entities_kwd", "text", "''"},
{"rank_flt", "double precision", "0.0"},
{"removed_kwd", "varchar(256)", "'N'"},
{"raptor_kwd", "varchar(32)", "''"},
{"raptor_layer_int", "integer", "0"},
{"metadata", "text", "''"},
{"extra", "text", "''"},
{"chunk_order_int", "integer", "0"},
{"compile_kwd", "varchar(32)", "''"},
{"toc_kwd", "varchar(32)", "''"},
{"mom_id", "varchar(128)", "''"},
{"mom_with_weight", "text", "''"},
}
var memoryColumns = []columnDefinition{
{"id", "varchar(256)", "''"},
{"message_id", "varchar(256)", "''"},
{"message_type_kwd", "varchar(64)", "''"},
{"source_id", "text", "''"},
{"memory_id", "varchar(256)", "''"},
{"user_id", "varchar(256)", "''"},
{"agent_id", "varchar(256)", "''"},
{"session_id", "varchar(256)", "''"},
{"zone_id", "integer", "0"},
{"valid_at", "varchar(64)", "''"},
{"invalid_at", "varchar(64)", "''"},
{"forget_at", "varchar(64)", "''"},
{"status_int", "integer", "1"},
{"content_ltks", "text", "''"},
{"tokenized_content_ltks", "text", "''"},
}
// metadataColumns mirrors conf/doc_meta_vastbase_mapping.json. Tables are
// per-tenant (ragflow_doc_meta_{tenantID}), so kb_id is not a discriminator
// here; it is a plain attribute column.
var metadataColumns = []columnDefinition{
{"id", "varchar(128)", "''"},
{"kb_id", "varchar(128)", "''"},
{"doc_id", "varchar(128)", "''"},
{"create_time", "varchar(32)", "''"},
{"create_timestamp_flt", "double precision", "0.0"},
{"meta_fields", "text", "'{}'"},
}
var skillColumns = []columnDefinition{
{"skill_id", "varchar(256)", "''"},
{"space_id", "varchar(256)", "''"},
{"folder_id", "varchar(256)", "''"},
{"name", "text", "''"},
{"name_tks", "text", "''"},
{"tags", "text", "''"},
{"tags_tks", "text", "''"},
{"description", "text", "''"},
{"description_tks", "text", "''"},
{"content", "text", "''"},
{"content_tks", "text", "''"},
{"version", "varchar(64)", "''"},
{"status", "varchar(64)", "''"},
{"create_time", "bigint", "0"},
{"update_time", "bigint", "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"}
// fullTextFields are the chunk columns covered by the full-text index in both
// compatibility modes.
var fullTextFields = []string{
"title_tks", "title_sm_tks", "important_kwd", "important_tks",
"question_tks", "content_ltks", "content_sm_ltks",
}
var skillFullTextFields = []string{"name_tks", "tags_tks", "description_tks", "content_tks"}
// CreateChunkStore creates or upgrades the shared baseName 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); err != nil {
return err
}
if err := e.ensureFullTextIndexes(ctx, baseName, skillFullTextFields); err != nil {
return err
}
case strings.HasPrefix(baseName, "memory_"):
if err := e.ensureTable(ctx, baseName, memoryColumns); err != nil {
return err
}
for _, field := range memoryIndexColumns {
if err := e.ensureRegularIndex(ctx, baseName, field); err != nil {
return err
}
}
default:
if err := e.ensureTable(ctx, baseName, chunkColumns); err != nil {
return err
}
for _, field := range chunkIndexColumns {
if err := e.ensureRegularIndex(ctx, baseName, field); err != nil {
return err
}
}
if err := e.ensureFullTextIndexes(ctx, baseName, fullTextFields); err != nil {
return err
}
}
return e.ensureVectorColumnAndIndex(ctx, baseName, vectorSize)
}
// ensureTable creates tableName with the given columns under the DDL lock,
// turning an already-present table into a no-op.
func (e *Engine) ensureTable(ctx context.Context, tableName string, columns []columnDefinition) error {
primaryKey := primaryKeyColumn(columns)
return e.withDDLLock(ctx, "vb_create_table_"+tableName, func() (bool, error) {
return e.tableExists(ctx, tableName)
}, func() error {
definitions := make([]string, 0, len(columns)+1)
for _, column := range columns {
definition := quoteIdent(column.name) + " " + column.typeSQL
if column.name == primaryKey {
definition += " NOT NULL PRIMARY KEY"
} else if column.defaultSQL != "" {
definition += " DEFAULT " + column.defaultSQL
}
definitions = append(definitions, definition)
}
query := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)",
quoteIdent(tableName), strings.Join(definitions, ", "))
_, err := e.db.ExecContext(ctx, query)
return err
})
}
// primaryKeyColumn picks the primary key column for a column set: id, or
// skill_id for skill tables.
func primaryKeyColumn(columns []columnDefinition) string {
for _, column := range columns {
switch column.name {
case "id", "skill_id":
return column.name
}
}
return ""
}
// ensureColumn adds one column to a table under the DDL lock when missing.
func (e *Engine) ensureColumn(ctx context.Context, tableName string, column columnDefinition) error {
return e.withDDLLock(ctx, "vb_add_"+column.name+"_"+tableName, func() (bool, error) {
return e.columnExists(ctx, tableName, column.name)
}, func() error {
definition := quoteIdent(column.name) + " " + column.typeSQL
if column.defaultSQL != "" {
definition += " DEFAULT " + column.defaultSQL
}
// No IF NOT EXISTS: Vastbase rejects "ADD COLUMN IF NOT EXISTS" with a
// syntax error (openGauss-derived kernels lack the clause). withDDLLock
// already probes columnExists under the lock and tolerates the
// already-exists race error.
_, err := e.db.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s",
quoteIdent(tableName), definition))
return err
})
}
// ensureDynamicColumns adds columns for chunk fields the static schema does
// not know, mirroring ES dynamic mapping: q_(\d+)_vec becomes floatvector(N),
// anything else becomes text with ” default. Known static columns are
// skipped so their declared types win.
func (e *Engine) ensureDynamicColumns(ctx context.Context, tableName string, fieldNames []string) error {
static := make(map[string]bool, len(chunkColumns)+len(memoryColumns)+len(skillColumns))
for _, column := range chunkColumns {
static[column.name] = true
}
for _, column := range memoryColumns {
static[column.name] = true
}
for _, column := range skillColumns {
static[column.name] = true
}
for _, field := range fieldNames {
if static[field] || !validFieldName(field) {
continue
}
if err := e.ensureColumn(ctx, tableName, dynamicColumnDefinition(field)); err != nil {
return err
}
}
return nil
}
// dynamicColumnDefinition maps an unknown document field to its column type:
// floatvector(N) for q_N_vec names, unbounded text otherwise — normalized
// values are never truncated, so a varchar would reject anything longer than
// its limit and roll back the whole delete+insert batch.
func dynamicColumnDefinition(field string) columnDefinition {
if match := vectorDimRegex.FindStringSubmatch(field); match != nil {
dim, _ := strconv.Atoi(match[1])
return columnDefinition{name: field, typeSQL: fmt.Sprintf("floatvector(%d)", dim)}
}
return columnDefinition{name: field, typeSQL: "text", defaultSQL: "''"}
}
// validFieldName reports whether name matches the identifier pattern every
// table, column, and index component must satisfy.
func validFieldName(field string) bool {
return identifierPattern.MatchString(field)
}
// ensureRegularIndex creates a b-tree index on one column when absent,
// serialized by the DDL lock.
func (e *Engine) ensureRegularIndex(ctx context.Context, tableName, columnName string) error {
indexName := regularIndexName(tableName, columnName)
return e.withDDLLock(ctx, "vb_add_idx_"+tableName+"_"+columnName, func() (bool, error) {
return e.indexExists(ctx, tableName, indexName)
}, func() error {
_, err := e.db.ExecContext(ctx, fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s ON %s (%s)",
quoteIdent(indexName), quoteIdent(tableName), quoteIdent(columnName)))
return err
})
}
// ensureFullTextIndexes creates the full-text structure for the compatibility
// mode. PG mode builds one GIN index over to_tsvector('cn_tokenizer', ...)
// expressions (failure is tolerated — vector search still works); B mode adds
// one "fulltext" index per field, tolerating already-exists.
func (e *Engine) ensureFullTextIndexes(ctx context.Context, tableName string, fields []string) error {
if e.dbCompatibility == compatB {
for _, field := range fields {
indexName := regularIndexName(tableName, field+"_fulltext")
err := e.withDDLLock(ctx, "vb_add_fulltext_idx_"+tableName+"_"+field, func() (bool, error) {
return e.indexExists(ctx, tableName, indexName)
}, func() error {
_, err := e.db.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD INDEX %s USING \"fulltext\" (%s)",
quoteIdent(tableName), quoteIdent(indexName), quoteIdent(field)))
return err
})
if err != nil {
// Python-parity degradation: a failed full-text index only
// costs BM25 search (vector search works without it), so
// ingestion must not be blocked on it. ChunkStoreExists still
// requires the index, so later batches retry and self-heal
// once the server accepts the DDL.
common.Warn("Vastbase full-text index creation failed; vector search works without it",
zap.String("table", tableName), zap.String("field", field), zap.Error(err))
}
}
return nil
}
indexName := "text_gin_idx_" + tableName
exists, err := e.indexExists(ctx, tableName, indexName)
if err != nil || exists {
return err
}
expressions := make([]string, 0, len(fields))
for _, field := range fields {
expressions = append(expressions, fmt.Sprintf("to_tsvector('cn_tokenizer', %s)", quoteIdent(field)))
}
query := fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s ON %s USING gin (%s)",
quoteIdent(indexName), quoteIdent(tableName), strings.Join(expressions, ", "))
if _, err := e.db.ExecContext(ctx, query); err != nil {
common.Warn("Vastbase PG full-text index creation failed; vector search works without it",
zap.String("table", tableName), zap.Error(err))
}
return nil
}
// ensureVectorColumnAndIndex provisions the q_N_vec column for an embedding
// size and its HNSW graph index.
func (e *Engine) ensureVectorColumnAndIndex(ctx context.Context, tableName string, vectorSize int) error {
if vectorSize <= 0 {
return nil
}
columnName := fmt.Sprintf("q_%d_vec", vectorSize)
if err := e.ensureColumn(ctx, tableName, columnDefinition{
name: columnName,
typeSQL: fmt.Sprintf("floatvector(%d)", vectorSize),
}); err != nil {
return err
}
indexName := regularIndexName(tableName, columnName+"_graph")
return e.withDDLLock(ctx, "vb_add_vector_idx_"+tableName+"_"+columnName, func() (bool, error) {
return e.indexExists(ctx, tableName, indexName)
}, func() error {
_, err := e.db.ExecContext(ctx, fmt.Sprintf(
"CREATE INDEX IF NOT EXISTS %s ON %s USING graph_index (%s floatvector_cosine_ops) WITH (m=16, ef_construction=50)",
quoteIdent(indexName), quoteIdent(tableName), quoteIdent(columnName)))
return err
})
}
// regularIndexName builds ix_<table>_<column>, hashing the tail when the
// name would exceed the kernel's index-name length limit.
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
}
// withDDLLock serializes a check-then-act DDL step on one object per
// process: probe under the lock and act only when the probe misses.
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
}
timeout := 60 * time.Second
distributed := kvrocks.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
}
// tableExists probes the current schema for a table by name.
func (e *Engine) tableExists(ctx context.Context, tableName string) (bool, error) {
if !validFieldName(tableName) {
return false, fmt.Errorf("vastbase: invalid table name %q", tableName)
}
var count int
err := e.db.QueryRowContext(ctx,
"SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = $1",
tableName).Scan(&count)
return count > 0, err
}
// columnExists probes a table for a column by name.
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 = current_schema() AND table_name = $1 AND column_name = $2",
tableName, columnName).Scan(&count)
return count > 0, err
}
// indexExists probes pg_indexes for a named index.
func (e *Engine) indexExists(ctx context.Context, tableName, indexName string) (bool, error) {
var count int
err := e.db.QueryRowContext(ctx,
"SELECT COUNT(*) FROM pg_indexes WHERE schemaname = current_schema() AND tablename = $1 AND indexname = $2",
tableName, indexName).Scan(&count)
return count > 0, err
}
// listTableColumns returns the live column set of a table, or nil when the
// table does not exist. Callers use it both to filter SELECT fields down to
// existing columns and to drive the ES-parity exists-filter semantics.
func (e *Engine) listTableColumns(ctx context.Context, tableName string) (map[string]columnMeta, error) {
exists, err := e.tableExists(ctx, tableName)
if err != nil || !exists {
return nil, err
}
rows, err := e.db.QueryContext(ctx,
"SELECT column_name, data_type, column_default, is_nullable FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = $1",
tableName)
if err != nil {
return nil, err
}
defer rows.Close()
columns := make(map[string]columnMeta)
for rows.Next() {
var meta columnMeta
if err := rows.Scan(&meta.name, &meta.dataType, &meta.columnDefault, &meta.nullable); err != nil {
return nil, err
}
columns[meta.name] = meta
}
return columns, rows.Err()
}
// isTextColumn reports whether a column uses a character type, which decides
// the ES-parity empty-string handling in filters.
func isTextColumn(dataType string) bool {
return textTypeRegex.MatchString(strings.TrimSpace(dataType))
}
// ChunkStoreExists checks the shared physical table and its static contract
// (base columns plus the indexes CreateChunkStore would build). Dynamic
// columns are not part of the check.
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 []string
fullText := false
switch kind {
case "memory":
indexColumns = memoryIndexColumns
case "skill":
fullText = true
default:
indexColumns = chunkIndexColumns
fullText = true
}
for _, column := range indexColumns {
exists, err = e.indexExists(ctx, baseName, regularIndexName(baseName, column))
if err != nil || !exists {
return exists, err
}
}
if fullText {
if e.dbCompatibility == compatB {
// Skill tables carry full-text indexes over skillFullTextFields,
// not the chunk fields; probing the chunk set would always miss
// and re-run the whole CreateChunkStore path on every write.
fields := fullTextFields
if kind == "skill" {
fields = skillFullTextFields
}
for _, field := range fields {
exists, err = e.indexExists(ctx, baseName, regularIndexName(baseName, field+"_fulltext"))
if err != nil || !exists {
return exists, err
}
}
} else {
exists, err = e.indexExists(ctx, baseName, "text_gin_idx_"+baseName)
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"
}
result, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE %s = $1",
quoteIdent(baseName), quoteIdent(field)), datasetID)
if err != nil {
return err
}
if affected, err := result.RowsAffected(); err == nil && affected > 0 {
common.Info("Vastbase dropped chunk store rows",
zap.String("table", baseName), zap.String("dataset", datasetID), zap.Int64("rows", affected))
}
return nil
}
_, err := e.db.ExecContext(ctx, "DROP TABLE IF EXISTS "+quoteIdent(baseName))
return err
}
// quoteIdent double-quotes an SQL identifier, escaping embedded quotes.
func quoteIdent(identifier string) string {
return `"` + strings.ReplaceAll(identifier, `"`, `""`) + `"`
}
// validateIdentifier rejects empty names or names outside the identifier
// pattern before they reach any SQL text.
func validateIdentifier(identifier string) error {
if identifier == "" || !identifierPattern.MatchString(identifier) {
return fmt.Errorf("invalid SQL identifier: %q", identifier)
}
return nil
}
// isDuplicateDDLError reports whether a DDL failure is a benign
// already-exists race rather than a real error.
func isDuplicateDDLError(err error) bool {
message := strings.ToLower(err.Error())
return strings.Contains(message, "duplicate") || strings.Contains(message, "already exists")
}
+386
View File
@@ -0,0 +1,386 @@
//
// 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 vastbase
import (
"context"
"database/sql/driver"
"fmt"
"regexp"
"strings"
"testing"
"ragflow/internal/server/config"
"github.com/DATA-DOG/go-sqlmock"
)
// expectCount wires one information_schema/pg_indexes existence probe that
// scans a single integer.
func expectCount(mock sqlmock.Sqlmock, query string, args ...driver.Value) *sqlmock.ExpectedQuery {
return mock.ExpectQuery(regexp.QuoteMeta(query)).WithArgs(args...).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
}
// tableCountQuery is the information_schema probe stub shared by the schema
// tests.
func tableCountQuery() string {
return "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = current_schema() AND table_name = $1"
}
// indexCountQuery is the pg_indexes probe stub shared by the schema tests.
func indexCountQuery() string {
return "SELECT COUNT(*) FROM pg_indexes WHERE schemaname = current_schema() AND tablename = $1 AND indexname = $2"
}
// columnCountQuery is the information_schema column probe stub shared by
// the schema tests.
func columnCountQuery() string {
return "SELECT COUNT(*) FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = $1 AND column_name = $2"
}
// expectIndexProvision expects the withDDLLock triple for one index:
// absence probe, DDL statement, existence confirmation.
func expectIndexProvision(mock sqlmock.Sqlmock, table, indexName, ddl string) {
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs(table, indexName).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(regexp.QuoteMeta(ddl)).WillReturnResult(sqlmock.NewResult(0, 0))
expectCount(mock, indexCountQuery(), table, indexName)
}
// TestCreateMemoryStoreDDLGolden pins the memory table's CREATE TABLE and
// companion DDL.
func TestCreateMemoryStoreDDLGolden(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
createTable := `CREATE TABLE IF NOT EXISTS "memory_t1" (` +
`"id" varchar(256) NOT NULL PRIMARY KEY, "message_id" varchar(256) DEFAULT '', ` +
`"message_type_kwd" varchar(64) DEFAULT '', "source_id" text DEFAULT '', ` +
`"memory_id" varchar(256) DEFAULT '', "user_id" varchar(256) DEFAULT '', ` +
`"agent_id" varchar(256) DEFAULT '', "session_id" varchar(256) DEFAULT '', ` +
`"zone_id" integer DEFAULT 0, "valid_at" varchar(64) DEFAULT '', ` +
`"invalid_at" varchar(64) DEFAULT '', "forget_at" varchar(64) DEFAULT '', ` +
`"status_int" integer DEFAULT 1, "content_ltks" text DEFAULT '', ` +
`"tokenized_content_ltks" text DEFAULT '')`
// ensureTable: absent probe, DDL, existence confirmation.
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("memory_t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(regexp.QuoteMeta(createTable)).WillReturnResult(sqlmock.NewResult(0, 0))
expectCount(mock, tableCountQuery(), "memory_t1")
for _, column := range memoryIndexColumns {
expectIndexProvision(mock, "memory_t1", regularIndexName("memory_t1", column),
fmt.Sprintf(`CREATE INDEX IF NOT EXISTS %s ON "memory_t1" (%s)`,
quoteIdent(regularIndexName("memory_t1", column)), quoteIdent(column)))
}
// vectorSize 0: no vector column or index is created.
if err := engine.CreateChunkStore(context.Background(), "memory_t1", "mem-1", 0, ""); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestCreateChunkStoreBModeProvisionsFullTextPerColumn checks that B mode
// provisions one full-text index per indexed chunk column.
func TestCreateChunkStoreBModeProvisionsFullTextPerColumn(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
// The chunk DDL is asserted by fragments: PK on id, text DEFAULT '',
// array columns without DEFAULT, and the literal column set from the
// Python mapping.
createChunkPattern := strings.Join([]string{
regexp.QuoteMeta(`CREATE TABLE IF NOT EXISTS "t1" ("id" varchar(256) NOT NULL PRIMARY KEY`),
regexp.QuoteMeta(`"docnm_kwd" text DEFAULT ''`),
// Fractional feedback deltas must survive: pagerank_fea is double
// precision, not integer.
regexp.QuoteMeta(`"pagerank_fea" double precision DEFAULT 0.0`),
regexp.QuoteMeta(`"available_int" integer DEFAULT 1`),
regexp.QuoteMeta(`"position_int" integer[]`),
regexp.QuoteMeta(`"metadata" text DEFAULT ''`),
regexp.QuoteMeta(`"mom_with_weight" text DEFAULT '')`),
}, ".*")
mock.ExpectExec(createChunkPattern).WillReturnResult(sqlmock.NewResult(0, 0))
expectCount(mock, tableCountQuery(), "t1")
for _, column := range chunkIndexColumns {
expectIndexProvision(mock, "t1", regularIndexName("t1", column),
fmt.Sprintf(`CREATE INDEX IF NOT EXISTS %s ON "t1" (%s)`,
quoteIdent(regularIndexName("t1", column)), quoteIdent(column)))
}
for _, field := range fullTextFields {
indexName := regularIndexName("t1", field+"_fulltext")
expectIndexProvision(mock, "t1", indexName,
fmt.Sprintf(`ALTER TABLE "t1" ADD INDEX %s USING "fulltext" (%s)`,
quoteIdent(indexName), quoteIdent(field)))
}
if err := engine.CreateChunkStore(context.Background(), "t1", "kb-1", 0, ""); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestCreateChunkStoreBModeDegradesWhenFullTextIndexFails checks that a
// failed full-text index degrades to a warning instead of failing the store.
func TestCreateChunkStoreBModeDegradesWhenFullTextIndexFails(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(regexp.QuoteMeta(`CREATE TABLE IF NOT EXISTS "t1" (`)).
WillReturnResult(sqlmock.NewResult(0, 0))
expectCount(mock, tableCountQuery(), "t1")
for _, column := range chunkIndexColumns {
expectIndexProvision(mock, "t1", regularIndexName("t1", column),
fmt.Sprintf(`CREATE INDEX IF NOT EXISTS %s ON "t1" (%s)`,
quoteIdent(regularIndexName("t1", column)), quoteIdent(column)))
}
// Every full-text index probe reports absence and the DDL fails: the
// Python connector degrades to vector-only search here instead of
// failing the store.
for _, field := range fullTextFields {
indexName := regularIndexName("t1", field+"_fulltext")
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs("t1", indexName).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(regexp.QuoteMeta(fmt.Sprintf(`ALTER TABLE "t1" ADD INDEX %s USING "fulltext"`,
quoteIdent(indexName)))).WillReturnError(fmt.Errorf("fulltext extension missing"))
}
// vectorSize 0: no vector column or index is created; the failed
// full-text DDL must not surface as an error.
if err := engine.CreateChunkStore(context.Background(), "t1", "kb-1", 0, ""); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestEnsureFullTextIndexesPGModeBuildsOneGIN checks that PG mode builds a
// single GIN index over the concatenated text columns.
func TestEnsureFullTextIndexesPGModeBuildsOneGIN(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expressions := make([]string, 0, len(fullTextFields))
for _, field := range fullTextFields {
expressions = append(expressions, fmt.Sprintf("to_tsvector('cn_tokenizer', %s)", quoteIdent(field)))
}
// The PG path probes and creates directly (no withDDLLock, failure only
// warns), so there is no existence confirmation probe afterwards.
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).WithArgs("t1", "text_gin_idx_t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(regexp.QuoteMeta(fmt.Sprintf(
`CREATE INDEX IF NOT EXISTS "text_gin_idx_t1" ON "t1" USING gin (%s)`,
strings.Join(expressions, ", ")))).WillReturnResult(sqlmock.NewResult(0, 0))
if err := engine.ensureFullTextIndexes(context.Background(), "t1", fullTextFields); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestEnsureVectorColumnAndIndexUsesGraphIndex pins the q_N_vec column and
// graph_index DDL.
func TestEnsureVectorColumnAndIndexUsesGraphIndex(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
mock.ExpectQuery(regexp.QuoteMeta(columnCountQuery())).WithArgs("t1", "q_3_vec").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
// No IF NOT EXISTS in the ALTER: Vastbase rejects the clause, and
// withDDLLock's existence probe makes it redundant.
mock.ExpectExec(regexp.QuoteMeta(`ALTER TABLE "t1" ADD COLUMN "q_3_vec" floatvector(3)`)).
WillReturnResult(sqlmock.NewResult(0, 0))
expectCount(mock, columnCountQuery(), "t1", "q_3_vec")
expectIndexProvision(mock, "t1", regularIndexName("t1", "q_3_vec_graph"),
`CREATE INDEX IF NOT EXISTS "ix_t1_q_3_vec_graph" ON "t1" USING graph_index ("q_3_vec" floatvector_cosine_ops) WITH (m=16, ef_construction=50)`)
if err := engine.ensureVectorColumnAndIndex(context.Background(), "t1", 3); err != nil {
t.Fatal(err)
}
if err := engine.ensureVectorColumnAndIndex(context.Background(), "t1", 0); err != nil {
t.Fatal(err) // zero dimension is a no-op
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestDynamicColumnDefinitionTypes checks unknown-field typing: floatvector
// for q_N_vec names, unbounded text otherwise.
func TestDynamicColumnDefinitionTypes(t *testing.T) {
if got := dynamicColumnDefinition("q_1024_vec"); got.typeSQL != "floatvector(1024)" || got.defaultSQL != "" {
t.Fatalf("vector dynamic column = %#v", got)
}
if got := dynamicColumnDefinition("custom_field"); got.typeSQL != "text" || got.defaultSQL != "''" {
t.Fatalf("scalar dynamic column = %#v", got)
}
}
// TestRegularIndexNameTruncatesWithinLimit checks the hashed truncation that
// keeps index names inside the length limit.
func TestRegularIndexNameTruncatesWithinLimit(t *testing.T) {
long := regularIndexName(strings.Repeat("t", 40), strings.Repeat("c", 30))
// Truncation keeps headroom below the 63-char PG limit and appends a
// stable 4-hex digest so distinct long names cannot collide after cut.
if len(long) > maxIndexNameLength {
t.Fatalf("truncated index name length = %d (%q)", len(long), long)
}
if !strings.HasPrefix(long, "ix_") || !regexp.MustCompile(`_[0-9a-f]{4}$`).MatchString(long) {
t.Fatalf("truncated index name = %q", long)
}
if got := regularIndexName("t1", "kb_id"); got != "ix_t1_kb_id" {
t.Fatalf("short index name = %q", got)
}
}
// TestDropChunkStoreDeletesDatasetRowsButKeepsSharedTable checks that
// dropping a chunk store deletes only the dataset's rows in the shared
// table.
func TestDropChunkStoreDeletesDatasetRowsButKeepsSharedTable(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
expectCount(mock, tableCountQuery(), "t1")
mock.ExpectExec(regexp.QuoteMeta(`DELETE FROM "t1" WHERE "kb_id" = $1`)).
WithArgs("kb-1").WillReturnResult(sqlmock.NewResult(0, 7))
if err := engine.DropChunkStore(context.Background(), "t1", "kb-1"); err != nil {
t.Fatal(err)
}
expectCount(mock, tableCountQuery(), "memory_t1")
mock.ExpectExec(regexp.QuoteMeta(`DELETE FROM "memory_t1" WHERE "memory_id" = $1`)).
WithArgs("mem-1").WillReturnResult(sqlmock.NewResult(0, 3))
if err := engine.DropChunkStore(context.Background(), "memory_t1", "mem-1"); err != nil {
t.Fatal(err)
}
// Unscoped and skill drops remove the whole table.
mock.ExpectExec(regexp.QuoteMeta(`DROP TABLE IF EXISTS "t1"`)).WillReturnResult(sqlmock.NewResult(0, 0))
if err := engine.DropChunkStore(context.Background(), "t1", ""); err != nil {
t.Fatal(err)
}
mock.ExpectExec(regexp.QuoteMeta(`DROP TABLE IF EXISTS "skill_t1"`)).WillReturnResult(sqlmock.NewResult(0, 0))
if err := engine.DropChunkStore(context.Background(), "skill_t1", "skill"); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestChunkStoreExistsRequiresStaticIndexes checks that store existence is
// judged by the presence of the static schema indexes, not the data.
func TestChunkStoreExistsRequiresStaticIndexes(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
expectCount(mock, tableCountQuery(), "t1")
for _, column := range chunkIndexColumns {
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", column))
}
for _, field := range fullTextFields {
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", field+"_fulltext"))
}
exists, err := engine.ChunkStoreExists(context.Background(), "t1", "kb-1")
if err != nil || !exists {
t.Fatalf("exists = %v err = %v", exists, err)
}
// One missing full-text index invalidates the store.
expectCount(mock, tableCountQuery(), "t1")
for _, column := range chunkIndexColumns {
expectCount(mock, indexCountQuery(), "t1", regularIndexName("t1", column))
}
mock.ExpectQuery(regexp.QuoteMeta(indexCountQuery())).
WithArgs("t1", regularIndexName("t1", fullTextFields[0]+"_fulltext")).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
exists, err = engine.ChunkStoreExists(context.Background(), "t1", "kb-1")
if err != nil || exists {
t.Fatalf("exists = %v err = %v", exists, err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// In B mode a skill table's full-text contract is skillFullTextFields —
// probing the chunk fields would always miss and re-run CreateChunkStore on
// every skill write.
func TestChunkStoreExistsSkillBModeProbesSkillFullTextFields(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatB}, db)
expectCount(mock, tableCountQuery(), "skill_t1")
for _, field := range skillFullTextFields {
expectCount(mock, indexCountQuery(), "skill_t1", regularIndexName("skill_t1", field+"_fulltext"))
}
exists, err := engine.ChunkStoreExists(context.Background(), "skill_t1", "skill")
if err != nil || !exists {
t.Fatalf("exists = %v err = %v", exists, err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
File diff suppressed because it is too large Load Diff
+918
View File
@@ -0,0 +1,918 @@
//
// 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 vastbase
import (
"context"
"fmt"
"reflect"
"regexp"
"strings"
"testing"
"ragflow/internal/common"
"ragflow/internal/engine/types"
"ragflow/internal/server/config"
"github.com/DATA-DOG/go-sqlmock"
)
// TestMain initializes the shared logger: Search logs every request through
// common.LogSearchRequest, whose level probe dereferences a nil AtomicLevel
// until InitLogger runs.
func TestMain(m *testing.M) {
_ = common.InitLogger("info", common.FileOutput{}, "vastbase_test")
m.Run()
}
// chunkColumnRows is the standard chunk table column list shared by the
// search golden tests.
func chunkColumnRows() []string {
return []string{"id", "kb_id", "content_ltks", "pagerank_fea", "q_3_vec"}
}
// expectChunkColumns stubs the listTableColumns round trip with the standard
// chunk column set.
func expectChunkColumns(mock sqlmock.Sqlmock, table string) {
expectCount(mock, tableCountQuery(), table)
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
for _, column := range chunkColumnRows() {
rows.AddRow(column, "text", nil, "YES")
}
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs(table).WillReturnRows(rows)
}
// expectLexemeProbe stubs the server-side cn_tokenizer parse round trip that
// derives the PG-mode full-text terms.
func expectLexemeProbe(mock sqlmock.Sqlmock, query, rendered string) {
mock.ExpectQuery(regexp.QuoteMeta(`SELECT to_tsvector('cn_tokenizer', $1)::text AS lexemes`)).
WithArgs(query).
WillReturnRows(sqlmock.NewRows([]string{"lexemes"}).AddRow(rendered))
}
// TestFullTextCTESingleColumnGolden pins the single-column bm25 CTE SQL.
func TestFullTextCTESingleColumnGolden(t *testing.T) {
engine := &Engine{dbCompatibility: compatB}
text := &types.MatchTextExpr{
Fields: []string{"content_ltks"},
MatchingText: "hello",
TopN: 100,
ExtraOptions: map[string]interface{}{"minimum_should_match": 0.5},
}
args := &sqlArgs{}
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id", "content_ltks"}, `"kb_id" = $1`,
[]interface{}{"kb-1"}, text, 100, args)
if err != nil {
t.Fatal(err)
}
want := `SELECT "id", "content_ltks", bm25_score AS _score FROM ` +
`(SELECT "id", "content_ltks", bm25_score() AS bm25_score FROM "t1" ` +
`WHERE ("kb_id" = $1) AND ("content_ltks" @~@ $2) ORDER BY bm25_score DESC LIMIT 100) AS ranked`
if got != want {
t.Fatalf("single-column CTE =\n%s", got)
}
operand := "hello @<PARAM:MINIMUM_SHOULD_MATCH=50% PARAM:BOOST=1>@"
if !reflect.DeepEqual(args.values, []interface{}{"kb-1", operand}) {
t.Fatalf("args = %#v", args.values)
}
}
// TestFullTextCTEMultiColumnUnionGolden pins the UNION ALL form the CTE takes
// when several columns carry full-text indexes.
func TestFullTextCTEMultiColumnUnionGolden(t *testing.T) {
engine := &Engine{dbCompatibility: compatB}
text := &types.MatchTextExpr{
Fields: []string{"title_tks^2", "content_ltks"},
MatchingText: "boosted~3 query^2",
TopN: 50,
}
args := &sqlArgs{}
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id"}, `"kb_id" = $1`,
[]interface{}{"kb-1"}, text, 50, args)
if err != nil {
t.Fatal(err)
}
branch := func(field, weight string, filterPart, placeholder string) string {
return `(SELECT "id", bm25_score() AS bm25_score FROM "t1" ` +
`WHERE (` + filterPart + `) AND (` + field + ` @~@ ` + placeholder + `) ORDER BY bm25_score DESC LIMIT 100)`
}
// Each branch re-binds the filter, so placeholders advance per branch; the
// dedup keys on the kind's row identifier.
want := `SELECT "id", "SCORE" AS _score FROM (SELECT DISTINCT ON ("id") "id", bm25_score AS "SCORE" FROM (` +
branch(`"title_tks"`, `2`, `"kb_id" = $1`, `$2`) + " UNION ALL " +
branch(`"content_ltks"`, `1`, `"kb_id" = $3`, `$4`) +
`) AS unioned ORDER BY "id", "SCORE" DESC) AS deduped ORDER BY _score DESC LIMIT 50`
if got != want {
t.Fatalf("multi-column CTE =\n%s\nwant\n%s", got, want)
}
// Boost suffixes move out of the query text into PARAM:BOOST.
if args.values[1] != "boosted query @<PARAM:MINIMUM_SHOULD_MATCH=0% PARAM:BOOST=2>@" {
t.Fatalf("operand 1 = %#v", args.values[1])
}
if args.values[3] != "boosted query @<PARAM:MINIMUM_SHOULD_MATCH=0% PARAM:BOOST=1>@" {
t.Fatalf("operand 2 = %#v", args.values[3])
}
}
// TestFullTextCTEPGModeSingleColumnGolden pins the PG-mode single-column CTE:
// predicates restate the indexed to_tsvector('cn_tokenizer', ...) expression
// against plainto_tsquery, the server-parsed lexemes OR together below 100%
// minimum_should_match, and ts_rank supplies the score.
func TestFullTextCTEPGModeSingleColumnGolden(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
text := &types.MatchTextExpr{
Fields: []string{"content_ltks"},
MatchingText: "hello world",
TopN: 100,
ExtraOptions: map[string]interface{}{"minimum_should_match": 0.5},
}
expectLexemeProbe(mock, "hello world", `'hello':1 'world':2`)
args := &sqlArgs{}
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id", "content_ltks"}, `"kb_id" = $1`,
[]interface{}{"kb-1"}, text, 100, args)
if err != nil {
t.Fatal(err)
}
rank := `GREATEST(ts_rank(to_tsvector('cn_tokenizer', "content_ltks"), plainto_tsquery('cn_tokenizer', $2)), ` +
`ts_rank(to_tsvector('cn_tokenizer', "content_ltks"), plainto_tsquery('cn_tokenizer', $3)))`
match := `to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $2) OR ` +
`to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $3)`
want := `SELECT "id", "content_ltks", (` + rank + ` * 1) AS _score FROM "t1" ` +
`WHERE ("kb_id" = $1) AND (` + match + `) ORDER BY _score DESC LIMIT 100`
if got != want {
t.Fatalf("PG single-column CTE =\n%s\nwant\n%s", got, want)
}
if !reflect.DeepEqual(args.values, []interface{}{"kb-1", "hello", "world"}) {
t.Fatalf("args = %#v", args.values)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestFullTextCTEPGModeMultiColumnSkillGolden pins the PG-mode UNION form and
// the skill table's row identifier: branches deduplicate on skill_id, and a
// 100% minimum_should_match AND-joins the term predicates.
func TestFullTextCTEPGModeMultiColumnSkillGolden(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
text := &types.MatchTextExpr{
Fields: []string{"name_tks^2", "content_tks"},
MatchingText: "deploy",
TopN: 50,
ExtraOptions: map[string]interface{}{"minimum_should_match": "100%"},
}
expectLexemeProbe(mock, "deploy", `'deploy':1`)
args := &sqlArgs{}
got, err := engine.fullTextCTE(context.Background(), "skill_t1", "skill", []string{"skill_id"}, `"space_id" = $1`,
[]interface{}{"sp-1"}, text, 50, args)
if err != nil {
t.Fatal(err)
}
branch := func(field, weight, filterPart, placeholder string) string {
return `(SELECT "skill_id", (GREATEST(ts_rank(to_tsvector('cn_tokenizer', ` + field +
`), plainto_tsquery('cn_tokenizer', ` + placeholder + `))) * ` + weight + `) AS "SCORE" FROM "skill_t1" ` +
`WHERE (` + filterPart + `) AND (to_tsvector('cn_tokenizer', ` + field + `) @@ plainto_tsquery('cn_tokenizer', ` +
placeholder + `)) ORDER BY "SCORE" DESC LIMIT 100)`
}
want := `SELECT "skill_id", "SCORE" AS _score FROM (SELECT DISTINCT ON ("skill_id") "skill_id", "SCORE" FROM (` +
branch(`"name_tks"`, `2`, `"space_id" = $1`, `$2`) + " UNION ALL " +
branch(`"content_tks"`, `1`, `"space_id" = $3`, `$4`) +
`) AS unioned ORDER BY "skill_id", "SCORE" DESC) AS deduped ORDER BY _score DESC LIMIT 50`
if got != want {
t.Fatalf("PG skill multi-column CTE =\n%s\nwant\n%s", got, want)
}
if !reflect.DeepEqual(args.values, []interface{}{"sp-1", "deploy", "sp-1", "deploy"}) {
t.Fatalf("args = %#v", args.values)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestPGFullTextCTECJKQuerySplitsServerLexemes is the CJK regression: an
// unspaced Chinese query parses into several server lexemes, and a
// minimum_should_match below 100% must OR the per-lexeme predicates — with
// whitespace splitting the query stayed one term whose lexemes
// plainto_tsquery AND-joins, so the OR relaxation never engaged.
func TestPGFullTextCTECJKQuerySplitsServerLexemes(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
text := &types.MatchTextExpr{
Fields: []string{"content_ltks"},
MatchingText: "机器学习",
TopN: 100,
}
expectLexemeProbe(mock, "机器学习", `'机器':1 '学习':2`)
args := &sqlArgs{}
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id"}, `"kb_id" = $1`,
[]interface{}{"kb-1"}, text, 100, args)
if err != nil {
t.Fatal(err)
}
match := `to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $2) OR ` +
`to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $3)`
if !strings.Contains(got, match) {
t.Fatalf("per-lexeme OR predicates missing:\n%s", got)
}
if !reflect.DeepEqual(args.values, []interface{}{"kb-1", "机器", "学习"}) {
t.Fatalf("args = %#v", args.values)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestPGFullTextCTEStopwordParseFallsBackToWhitespace checks the fallback
// when the server parse yields no lexemes (stopword-only query): terms fall
// back to whitespace words.
func TestPGFullTextCTEStopwordParseFallsBackToWhitespace(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
text := &types.MatchTextExpr{
Fields: []string{"content_ltks"},
MatchingText: "the of",
TopN: 100,
}
expectLexemeProbe(mock, "the of", "")
args := &sqlArgs{}
got, err := engine.fullTextCTE(context.Background(), "t1", "chunk", []string{"id"}, `"kb_id" = $1`,
[]interface{}{"kb-1"}, text, 100, args)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(got, `@@ plainto_tsquery('cn_tokenizer', $2) OR`) {
t.Fatalf("whitespace fallback predicates missing:\n%s", got)
}
if !reflect.DeepEqual(args.values, []interface{}{"kb-1", "the", "of"}) {
t.Fatalf("args = %#v", args.values)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestParseTSVectorLexemes checks lexeme extraction from a tsvector's text
// rendering: doubled quotes unescape, positions and weights drop, empty and
// duplicate lexemes filter out.
func TestParseTSVectorLexemes(t *testing.T) {
cases := []struct {
rendered string
want []string
}{
{`'机器':1 '学习':2`, []string{"机器", "学习"}},
{`'don''t':2`, []string{"don't"}},
{`'a':1,3 'b':2B`, []string{"a", "b"}},
{`'a':1 'a':2`, []string{"a"}},
{"", []string{}},
}
for _, test := range cases {
if got := parseTSVectorLexemes(test.rendered); !reflect.DeepEqual(got, test.want) {
t.Fatalf("parseTSVectorLexemes(%q) = %#v want %#v", test.rendered, got, test.want)
}
}
}
// TestMinimumShouldMatchCoversAllTerms checks when the PG-mode predicate
// AND-joins terms instead of OR-joining them.
func TestMinimumShouldMatchCoversAllTerms(t *testing.T) {
cases := []struct {
minimumShouldMatch string
terms int
want bool
}{
{"0%", 2, false},
{"50%", 2, false},
{"100%", 2, true},
{"200%", 1, true},
{"75%", 0, false},
{"2", 2, true},
{"1", 2, false},
{"3<-25%", 3, false},
{"nonsense", 2, false},
}
for _, test := range cases {
if got := minimumShouldMatchCoversAllTerms(test.minimumShouldMatch, test.terms); got != test.want {
t.Fatalf("minimumShouldMatchCoversAllTerms(%q, %d) = %v", test.minimumShouldMatch, test.terms, got)
}
}
}
// TestVectorCTEGolden pins the dense-vector CTE SQL, including the bare
// ORDER BY that keeps the graph index eligible.
func TestVectorCTEGolden(t *testing.T) {
dense := &types.MatchDenseExpr{
VectorColumnName: "q_3_vec",
EmbeddingData: []float64{0.1, 0.2, 0.3},
TopN: 10,
ExtraOptions: map[string]interface{}{"similarity": 0.2},
}
build := func(engine *Engine) string {
args := &sqlArgs{}
cte, err := engine.vectorCTE("t1", []string{"id"}, `"kb_id" = $1`,
[]interface{}{"kb-1"}, dense, 10, 0.2, "[0.1,0.2,0.3]", args)
if err != nil {
t.Fatal(err)
}
return cte
}
pg := build(&Engine{dbCompatibility: compatPG})
wantPG := `SELECT * FROM (SELECT "id", (1 - ("q_3_vec" <=> $2)) AS _score FROM "t1" ` +
`WHERE "kb_id" = $1 ORDER BY "q_3_vec" <=> $3 LIMIT 10) AS knn WHERE _score >= $4`
if pg != wantPG {
t.Fatalf("PG vector CTE =\n%s", pg)
}
b := build(&Engine{dbCompatibility: compatB})
if !regexp.MustCompile(`"q_3_vec" <\+> \$2`).MatchString(b) || !regexp.MustCompile(`"q_3_vec" <\+> \$3`).MatchString(b) {
t.Fatalf("B-mode operator missing:\n%s", b)
}
// The plan contract: threshold never inside the KNN scan, ORDER BY is the
// bare column so graph_index stays eligible.
if !regexp.MustCompile(`\) AS knn WHERE _score >= \$4$`).MatchString(b) {
t.Fatalf("threshold placement changed:\n%s", b)
}
}
// TestSearchFusionQueryGolden pins the fused full-text-plus-vector query SQL
// for PG mode: the ts_rank CTE, the vector KNN scan, and the identifier join.
func TestSearchFusionQueryGolden(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectChunkColumns(mock, "t1")
expectChunkColumns(mock, "t1") // searchFusion re-lists columns for the "*" expansion
expectLexemeProbe(mock, "hello", `'hello':1`)
projection := `COALESCE(a."content_ltks", b."content_ltks") AS "content_ltks", ` +
`COALESCE(a."id", b."id") AS "id", COALESCE(a."kb_id", b."kb_id") AS "kb_id", ` +
`COALESCE(a."pagerank_fea", b."pagerank_fea") AS "pagerank_fea", ` +
`COALESCE(a."q_3_vec", b."q_3_vec") AS "q_3_vec"`
fields := `"content_ltks", "id", "kb_id", "pagerank_fea", "q_3_vec"`
fulltextScore := `(GREATEST(ts_rank(to_tsvector('cn_tokenizer', "content_ltks"), plainto_tsquery('cn_tokenizer', $2))) * 1)`
fulltextMatch := `to_tsvector('cn_tokenizer', "content_ltks") @@ plainto_tsquery('cn_tokenizer', $2)`
fulltextCTE := `SELECT ` + fields + `, ` + fulltextScore + ` AS _score FROM "t1" ` +
`WHERE ("kb_id" IN ($1)) AND (` + fulltextMatch + `) ORDER BY _score DESC LIMIT 10`
vectorCTE := `SELECT * FROM (SELECT ` + fields + `, (1 - ("q_3_vec" <=> $4)) AS _score FROM "t1" ` +
`WHERE "kb_id" IN ($3) ORDER BY "q_3_vec" <=> $5 LIMIT 10) AS knn WHERE _score >= $6`
fused := `(COALESCE(a._score, 0) / COALESCE(NULLIF((SELECT MAX(_score) FROM filter_fulltext), 0), 1) * 0.7 + COALESCE(b._score, 0) * 0.3)`
query := `WITH filter_fulltext AS (` + fulltextCTE + `), filter_vector AS (` + vectorCTE + `) ` +
`SELECT ` + projection + `, ` + fused + ` AS _score ` +
`FROM filter_fulltext a FULL OUTER JOIN filter_vector b ON a."id" = b."id" ORDER BY _score DESC LIMIT 10 OFFSET 0`
mock.ExpectQuery(regexp.QuoteMeta(query)).
WithArgs("kb-1", "hello",
"kb-1", "[0.1,0.2,0.3]", "[0.1,0.2,0.3]", 0.2).
WillReturnRows(sqlmock.NewRows([]string{"id", "content_ltks", "kb_id", "pagerank_fea", "q_3_vec", "_score"}).
AddRow("c1", "hello world", "kb-1", 0.5, "[0.1,0.2,0.3]", 0.9))
result, err := engine.Search(context.Background(), &types.SearchRequest{
IndexNames: []string{"t1"},
KbIDs: []string{"kb-1"},
SelectFields: []string{"*"},
Limit: 10,
MatchExprs: []interface{}{
&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "hello", TopN: 10},
&types.MatchDenseExpr{VectorColumnName: "q_3_vec", EmbeddingData: []float64{0.1, 0.2, 0.3},
TopN: 10, ExtraOptions: map[string]interface{}{"similarity": 0.2}},
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
FusionParams: map[string]interface{}{"weights": "0.7,0.3"}},
},
})
if err != nil {
t.Fatal(err)
}
if len(result.Chunks) != 1 || result.Chunks[0]["id"] != "c1" {
t.Fatalf("chunks = %#v", result.Chunks)
}
if score, _ := numberToFloat(result.Chunks[0]["_score"]); score != 0.9 {
t.Fatalf("fused score = %#v", result.Chunks[0]["_score"])
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestSearchSkillFusionUsesSkillIdentifier drives a skill-table hybrid
// search end to end: the fused join and the projected row key must use
// skill_id, not id.
func TestSearchSkillFusionUsesSkillIdentifier(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
skillColumns := []string{"skill_id", "name", "name_tks", "content_tks", "q_2_vec"}
listSkillColumns := func() {
expectCount(mock, tableCountQuery(), "skill_t1")
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
for _, column := range skillColumns {
rows.AddRow(column, "text", nil, "YES")
}
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("skill_t1").WillReturnRows(rows)
}
listSkillColumns()
listSkillColumns() // searchFusion re-lists columns for the "*" expansion
expectLexemeProbe(mock, "deploy", `'deploy':1`)
fields := `"content_tks", "name", "name_tks", "q_2_vec", "skill_id"`
fulltextScore := `(GREATEST(ts_rank(to_tsvector('cn_tokenizer', "name_tks"), plainto_tsquery('cn_tokenizer', $1))) * 1)`
fulltextCTE := `SELECT ` + fields + `, ` + fulltextScore + ` AS _score FROM "skill_t1" ` +
`WHERE (1=1) AND (to_tsvector('cn_tokenizer', "name_tks") @@ plainto_tsquery('cn_tokenizer', $1)) ` +
`ORDER BY _score DESC LIMIT 10`
vectorCTE := `SELECT * FROM (SELECT ` + fields + `, (1 - ("q_2_vec" <=> $2)) AS _score FROM "skill_t1" ` +
`WHERE 1=1 ORDER BY "q_2_vec" <=> $3 LIMIT 10) AS knn WHERE _score >= $4`
projection := `COALESCE(a."content_tks", b."content_tks") AS "content_tks", ` +
`COALESCE(a."name", b."name") AS "name", COALESCE(a."name_tks", b."name_tks") AS "name_tks", ` +
`COALESCE(a."q_2_vec", b."q_2_vec") AS "q_2_vec", COALESCE(a."skill_id", b."skill_id") AS "skill_id"`
fused := `(COALESCE(a._score, 0) / COALESCE(NULLIF((SELECT MAX(_score) FROM filter_fulltext), 0), 1) * 0.5 + COALESCE(b._score, 0) * 0.5)`
query := `WITH filter_fulltext AS (` + fulltextCTE + `), filter_vector AS (` + vectorCTE + `) ` +
`SELECT ` + projection + `, ` + fused + ` AS _score ` +
`FROM filter_fulltext a FULL OUTER JOIN filter_vector b ON a."skill_id" = b."skill_id" ` +
`ORDER BY _score DESC LIMIT 10 OFFSET 0`
mock.ExpectQuery(regexp.QuoteMeta(query)).
WithArgs("deploy", "[0.5,0.25]", "[0.5,0.25]", 0.2).
WillReturnRows(sqlmock.NewRows([]string{"skill_id", "name", "_score"}).AddRow("s1", "deploy skill", 0.8))
result, err := engine.Search(context.Background(), &types.SearchRequest{
IndexNames: []string{"skill_t1"},
SelectFields: []string{"*"},
Limit: 10,
MatchExprs: []interface{}{
&types.MatchTextExpr{Fields: []string{"name_tks"}, MatchingText: "deploy", TopN: 10},
&types.MatchDenseExpr{VectorColumnName: "q_2_vec", EmbeddingData: []float64{0.5, 0.25},
TopN: 10, ExtraOptions: map[string]interface{}{"similarity": 0.2}},
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
FusionParams: map[string]interface{}{"weights": "0.5,0.5"}},
},
})
if err != nil {
t.Fatal(err)
}
if len(result.Chunks) != 1 || result.Chunks[0]["skill_id"] != "s1" {
t.Fatalf("chunks = %#v", result.Chunks)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestSearchFusionWeightDegeneration checks that all-or-nothing fusion
// weights collapse the hybrid path into the single remaining path.
func TestSearchFusionWeightDegeneration(t *testing.T) {
newReq := func(weights string) *types.SearchRequest {
return &types.SearchRequest{
IndexNames: []string{"t1"},
KbIDs: []string{"kb-1"},
SelectFields: []string{"id"},
Limit: 10,
MatchExprs: []interface{}{
&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "hello", TopN: 10},
&types.MatchDenseExpr{VectorColumnName: "q_3_vec", EmbeddingData: []float64{0.1, 0.2, 0.3}, TopN: 10},
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
FusionParams: map[string]interface{}{"weights": weights}},
},
}
}
// Vector weight 1: the text leg is dropped, only the KNN scan runs.
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectChunkColumns(mock, "t1")
mock.ExpectQuery(`^SELECT \* FROM \(SELECT \* FROM \(SELECT .+ <=> \$2.+ <=> \$3 .+\) AS knn WHERE _score >= \$4\) AS candidates ORDER BY _score DESC LIMIT 10 OFFSET 0$`).
WithArgs("kb-1", "[0.1,0.2,0.3]", "[0.1,0.2,0.3]", 0.0).
WillReturnRows(sqlmock.NewRows([]string{"id", "pagerank_fea", "_score"}).AddRow("c1", 0, 0.5))
if _, err := engine.Search(context.Background(), newReq("0,1")); err != nil {
t.Fatal(err)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
// Vector weight 0: the dense leg is dropped, only the full-text scan runs
// (PG mode: the plainto_tsquery predicate).
db2, mock2, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db2.Close()
engine2 := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db2)
expectChunkColumns(mock2, "t1")
expectLexemeProbe(mock2, "hello", `'hello':1`)
mock2.ExpectQuery(regexp.QuoteMeta(`@@ plainto_tsquery('cn_tokenizer', $2)`)).
WithArgs("kb-1", "hello").
WillReturnRows(sqlmock.NewRows([]string{"id", "pagerank_fea", "_score"}).AddRow("c1", 0, 1.5))
if _, err := engine2.Search(context.Background(), newReq("1,0")); err != nil {
t.Fatal(err)
}
if err := mock2.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestSearchMemoryFullTextUsesTSQuery pins the memory keyword path, which
// ranks through to_tsquery instead of the bm25 operator.
func TestSearchMemoryFullTextUsesTSQuery(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectCount(mock, tableCountQuery(), "memory_t1")
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
for _, column := range []string{"id", "memory_id", "content_ltks", "status_int", "forget_at"} {
rows.AddRow(column, "text", nil, "YES")
}
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("memory_t1").WillReturnRows(rows)
// The score expression is built after the match predicate, so the
// ts_rank placeholder ($3) trails the match placeholder ($2); both run
// over the fine-grained tokenized_content_ltks column.
query := `SELECT "content_ltks", "id", ts_rank(to_tsvector('simple', COALESCE(tokenized_content_ltks, '')), to_tsquery('simple', $3)) AS _score ` +
`FROM "memory_t1" WHERE "memory_id" IN ($1) AND ("forget_at" IS NULL OR "forget_at" = '') ` +
`AND to_tsvector('simple', COALESCE(tokenized_content_ltks, '')) @@ to_tsquery('simple', $2) ` +
`ORDER BY _score DESC LIMIT 30 OFFSET 0`
mock.ExpectQuery(regexp.QuoteMeta(query)).
WithArgs("mem-1", "hello & world", "hello & world").
WillReturnRows(sqlmock.NewRows([]string{"content_ltks", "id", "_score"}).AddRow("hello world", "m1_1", 0.1))
result, err := engine.Search(context.Background(), &types.SearchRequest{
IndexNames: []string{"memory_t1"},
KbIDs: []string{"mem-1"},
SelectFields: []string{"content"},
MatchExprs: []interface{}{&types.MatchTextExpr{MatchingText: "hello world"}},
})
if err != nil {
t.Fatal(err)
}
if len(result.Chunks) != 1 || result.Chunks[0]["content"] != "hello world" {
t.Fatalf("chunks = %#v", result.Chunks)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestSearchMemoryFusionCandidatesGolden pins the memory fusion candidate
// query that scores vectors over CTE rows.
func TestSearchMemoryFusionCandidatesGolden(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
expectCount(mock, tableCountQuery(), "memory_t1")
rows := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
for _, column := range []string{"id", "memory_id", "content_ltks", "forget_at", "q_2_vec"} {
rows.AddRow(column, "text", nil, "YES")
}
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("memory_t1").WillReturnRows(rows)
score := `(ts_rank(to_tsvector('simple', COALESCE(tokenized_content_ltks, '')), to_tsquery('simple', $3)) * 0.7 + (1 - ("q_2_vec" <=> $4)) * 0.3)`
inner := `SELECT "content_ltks", "id", "q_2_vec", ` + score + ` AS _score FROM "memory_t1" ` +
`WHERE "memory_id" IN ($1) AND ("forget_at" IS NULL OR "forget_at" = '') ` +
`AND to_tsvector('simple', COALESCE(tokenized_content_ltks, '')) @@ to_tsquery('simple', $2) LIMIT 20`
query := `SELECT * FROM (` + inner + `) AS candidates WHERE (1 - ("q_2_vec" <=> $5)) >= $6 ORDER BY _score DESC LIMIT 10 OFFSET 0`
mock.ExpectQuery(regexp.QuoteMeta(query)).
WithArgs("mem-1", "hello", "hello", "[0.5,0.25]", "[0.5,0.25]", 0.2).
WillReturnRows(sqlmock.NewRows([]string{"content_ltks", "id", "q_2_vec", "_score"}).
AddRow("hello", "m1_1", "[0.5,0.25]", 0.6))
result, err := engine.Search(context.Background(), &types.SearchRequest{
IndexNames: []string{"memory_t1"},
KbIDs: []string{"mem-1"},
SelectFields: []string{"content"},
Limit: 10,
MatchExprs: []interface{}{
&types.MatchTextExpr{MatchingText: "hello", TopN: 10},
&types.MatchDenseExpr{VectorColumnName: "q_2_vec", EmbeddingData: []float64{0.5, 0.25},
TopN: 10, ExtraOptions: map[string]interface{}{"similarity": 0.2}},
&types.FusionExpr{Method: "weighted_sum", TopN: 10,
FusionParams: map[string]interface{}{"weights": "0.7,0.3"}},
},
})
if err != nil {
t.Fatal(err)
}
if len(result.Chunks) != 1 || result.Chunks[0]["content_embed"] == nil {
t.Fatalf("chunks = %#v", result.Chunks)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestSearchFilterOnlyCountsThenPages checks the two-statement shape of the
// filter-only path: one COUNT for the total, one paged select.
func TestSearchFilterOnlyCountsThenPages(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
// Search lists the table's columns before dispatching; filter-only then
// runs the COUNT gate and lists columns once more for ORDER BY validation.
expectCount(mock, tableCountQuery(), "t1")
columns := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
for _, column := range []string{"id", "kb_id", "doc_id", "page_num_int"} {
columns.AddRow(column, "text", nil, "YES")
}
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("t1").WillReturnRows(columns)
mock.ExpectQuery(regexp.QuoteMeta(`SELECT COUNT("id") FROM "t1" WHERE "doc_id" = $1 AND "kb_id" IN ($2)`)).
WithArgs("doc-1", "kb-1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(3))
expectCount(mock, tableCountQuery(), "t1")
orderColumns := sqlmock.NewRows([]string{"column_name", "data_type", "column_default", "is_nullable"})
for _, column := range []string{"id", "kb_id", "doc_id", "page_num_int"} {
orderColumns.AddRow(column, "integer[]", nil, "YES")
}
mock.ExpectQuery(regexp.QuoteMeta(listColumnsQuery)).WithArgs("t1").WillReturnRows(orderColumns)
mock.ExpectQuery(regexp.QuoteMeta(
`SELECT "id" FROM "t1" WHERE "doc_id" = $1 AND "kb_id" IN ($2) `+
`ORDER BY (SELECT AVG(element) FROM unnest("page_num_int") AS element) DESC LIMIT 30 OFFSET 0`)).
WithArgs("doc-1", "kb-1").
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("c1").AddRow("c2"))
result, err := engine.Search(context.Background(), &types.SearchRequest{
IndexNames: []string{"t1"},
KbIDs: []string{"kb-1"},
Filter: map[string]interface{}{"doc_id": "doc-1"},
OrderBy: &types.OrderByExpr{Fields: []types.OrderByField{{Field: "page_num_int", Type: types.SortDesc}}},
})
if err != nil {
t.Fatal(err)
}
if result.Total != 3 || len(result.Chunks) != 2 {
t.Fatalf("total = %d chunks = %#v", result.Total, result.Chunks)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestSearchSkipsMissingTables checks that a search naming a table the
// schema never created skips it instead of failing.
func TestSearchSkipsMissingTables(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
mock.ExpectQuery(regexp.QuoteMeta(tableCountQuery())).WithArgs("t1").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
result, err := engine.Search(context.Background(), &types.SearchRequest{
IndexNames: []string{"t1"},
KbIDs: []string{"kb-1"},
SelectFields: []string{"id"},
MatchExprs: []interface{}{&types.MatchTextExpr{Fields: []string{"content_ltks"}, MatchingText: "x"}},
})
if err != nil {
t.Fatal(err)
}
if len(result.Chunks) != 0 || result.Total != 0 {
t.Fatalf("result = %#v", result)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
// TestMemoryTSQuery checks the memory tsquery builder's term joining and
// escaping.
func TestMemoryTSQuery(t *testing.T) {
cases := []struct {
text *types.MatchTextExpr
want string
}{
{&types.MatchTextExpr{MatchingText: "hello world"}, "hello & world"},
{&types.MatchTextExpr{MatchingText: "boost~3"}, "boost"},
{
&types.MatchTextExpr{
MatchingText: "ignored",
ExtraOptions: map[string]interface{}{"original_query": `user's (input) & more`},
},
"users & input & more",
},
{&types.MatchTextExpr{MatchingText: " "}, ""},
}
for _, test := range cases {
if got := memoryTSQuery(test.text); got != test.want {
t.Fatalf("memoryTSQuery(%#v) = %q, want %q", test.text.MatchingText, got, test.want)
}
}
}
// TestParseFullTextFields checks "field^weight" specification parsing,
// including entries without a weight.
func TestParseFullTextFields(t *testing.T) {
fields, weights := parseFullTextFields(&types.MatchTextExpr{
Fields: []string{"title_tks^10", "content_ltks", "question_tks^0.5"},
})
if !reflect.DeepEqual(fields, []string{"title_tks", "content_ltks", "question_tks"}) {
t.Fatalf("fields = %#v", fields)
}
if !reflect.DeepEqual(weights, []string{"10", "1", "0.5"}) {
t.Fatalf("weights = %#v", weights)
}
}
// TestFusionVectorWeight checks fusion weight parsing and its even-split
// fallbacks.
func TestFusionVectorWeight(t *testing.T) {
cases := []struct {
fusion *types.FusionExpr
want float64
}{
{nil, 0.5},
{&types.FusionExpr{}, 0.5},
{&types.FusionExpr{FusionParams: map[string]interface{}{"weights": "0.7,0.3"}}, 0.3},
{&types.FusionExpr{FusionParams: map[string]interface{}{"weights": "0.7"}}, 0.5},
{&types.FusionExpr{FusionParams: map[string]interface{}{"weights": "0.7,x"}}, 0.5},
{&types.FusionExpr{FusionParams: map[string]interface{}{"weights": " 0.5 , 0.9 "}}, 0.9},
}
for _, test := range cases {
if got := fusionVectorWeight(test.fusion); got != test.want {
t.Fatalf("fusionVectorWeight(%#v) = %v, want %v", test.fusion, got, test.want)
}
}
}
// TestFormatMinimumShouldMatch checks the percent rendering of the
// minimum_should_match option across its accepted types.
func TestFormatMinimumShouldMatch(t *testing.T) {
cases := []struct {
options map[string]interface{}
want string
}{
{nil, "0%"},
{map[string]interface{}{}, "0%"},
{map[string]interface{}{"minimum_should_match": 0.5}, "50%"},
{map[string]interface{}{"minimum_should_match": float32(0.25)}, "25%"},
// A bare count is formatted through the percent helper, matching the
// oceanbase connector's behavior for the same ExtraOptions value.
{map[string]interface{}{"minimum_should_match": 2}, "200%"},
{map[string]interface{}{"minimum_should_match": "75%"}, "75%"},
{map[string]interface{}{"minimum_should_match": ""}, "0%"},
}
for _, test := range cases {
got := formatMinimumShouldMatch(&types.MatchTextExpr{ExtraOptions: test.options})
if got != test.want {
t.Fatalf("formatMinimumShouldMatch(%#v) = %q, want %q", test.options, got, test.want)
}
}
}
// TestBuildSearchOutput checks how requested fields, kinds, and hidden sort
// fields turn into the SQL select list.
func TestBuildSearchOutput(t *testing.T) {
columns := testColumns()
output := buildSearchOutput([]string{"content_ltks", "_score", "ghost_field"}, "chunk", columns, true, nil)
if !reflect.DeepEqual(output, []string{"content_ltks", "id", "pagerank_fea"}) {
t.Fatalf("chunk scored output = %#v", output)
}
plain := buildSearchOutput([]string{"content_ltks"}, "chunk", columns, false, nil)
if !reflect.DeepEqual(plain, []string{"content_ltks", "id"}) {
t.Fatalf("chunk unscored output = %#v", plain)
}
if star := buildSearchOutput([]string{"*"}, "chunk", columns, true, nil); !reflect.DeepEqual(star, []string{"*"}) {
t.Fatalf("star output = %#v", star)
}
memory := make(map[string]columnMeta)
for _, column := range memoryColumns {
memory[column.name] = columnMeta{name: column.name, dataType: column.typeSQL}
}
memory["q_2_vec"] = columnMeta{name: "q_2_vec", dataType: "text"}
memoryOutput := buildSearchOutput([]string{"content", "status", "content_embed"}, "memory", memory, true, nil)
if !reflect.DeepEqual(memoryOutput, []string{"content_ltks", "status_int", "q_2_vec", "id"}) {
t.Fatalf("memory output = %#v", memoryOutput)
}
// Skill rows key on skill_id, not id.
skill := map[string]columnMeta{"skill_id": {}, "name": {}, "content_tks": {}}
skillOutput := buildSearchOutput([]string{"name"}, "skill", skill, true, nil)
if !reflect.DeepEqual(skillOutput, []string{"name", "skill_id"}) {
t.Fatalf("skill output = %#v", skillOutput)
}
}
// TestScoredChunkSortFoldsPagerank checks that merged ranking folds
// pagerank_fea into the sort key.
func TestScoredChunkSortFoldsPagerank(t *testing.T) {
chunks := []map[string]interface{}{
{"id": "low", "_score": 0.8, "pagerank_fea": 0},
{"id": "high", "_score": 0.5, "pagerank_fea": 1.0},
{"id": "mid", "_score": 0.7, "pagerank_fea": 0},
}
sortScoredChunks(chunks)
if chunks[0]["id"] != "high" || chunks[1]["id"] != "low" || chunks[2]["id"] != "mid" {
t.Fatalf("order = %#v", chunks)
}
}
// Multi-table searches page only once, on the merged set: each table's SQL
// must cover the whole candidate range (LIMIT offset+limit OFFSET 0), not
// page with req's OFFSET — otherwise mergeSearchChunks skips offset rows a
// second time and loses page-2 rows that ranked inside a single table.
func TestSearchMultiTableWindowsOnlyAfterMerge(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
for _, table := range []string{"t1", "t2"} {
expectChunkColumns(mock, table)
mock.ExpectQuery(regexp.QuoteMeta(fmt.Sprintf(`SELECT COUNT("id") FROM "%s" WHERE 1=1`, table))).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(5))
// searchFilterOnly re-lists columns before building ORDER BY.
expectChunkColumns(mock, table)
prefix := "a"
if table == "t2" {
prefix = "b"
}
rows := sqlmock.NewRows([]string{"id"})
for i := 1; i <= 5; i++ {
rows.AddRow(fmt.Sprintf("%s%d", prefix, i))
}
mock.ExpectQuery(regexp.QuoteMeta(fmt.Sprintf(
`SELECT "id" FROM "%s" WHERE 1=1 ORDER BY "id" ASC LIMIT 5 OFFSET 0`, table))).
WillReturnRows(rows)
}
result, err := engine.Search(context.Background(), &types.SearchRequest{
IndexNames: []string{"t1", "t2"},
Offset: 2,
Limit: 3,
SelectFields: []string{"id"},
OrderBy: &types.OrderByExpr{Fields: []types.OrderByField{{Field: "id", Type: types.SortAsc}}},
})
if err != nil {
t.Fatal(err)
}
// Global order a1..a5, b1..b5; the page window [2:5) keeps a3..a5 —
// exactly the rows the double-offset path dropped.
got := make([]string, 0, len(result.Chunks))
for _, chunk := range result.Chunks {
got = append(got, chunk["id"].(string))
}
if !reflect.DeepEqual(got, []string{"a3", "a4", "a5"}) {
t.Fatalf("merged page ids = %v", got)
}
if result.Total != 10 {
t.Fatalf("total = %d", result.Total)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
+59
View File
@@ -0,0 +1,59 @@
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package vastbase
import (
"context"
"database/sql"
"fmt"
"regexp"
"strings"
)
var sqlLimitPattern = regexp.MustCompile(`(?i)\blimit\b`)
// RunSQL executes the read-only SQL produced by the chat SQL-retrieval flow.
// The service layer injects the kb_id filter before the SQL reaches the
// engine (shared tables are not auto-scoped here). The SELECT/WITH prefix
// check is a guardrail, not the read-only control: a data-modifying CTE
// (WITH d AS (DELETE FROM ... RETURNING *) SELECT * FROM d) also starts with
// WITH, so the statement runs inside a READ ONLY transaction the engine
// always rolls back — the server itself rejects writes on that path.
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), ";"))
if strings.Contains(normalized, ";") {
return nil, fmt.Errorf("multiple statements are not allowed")
}
lower := strings.ToLower(normalized)
if !strings.HasPrefix(lower, "select ") && !strings.HasPrefix(lower, "with ") {
return nil, fmt.Errorf("only SELECT and WITH statements are allowed")
}
if !sqlLimitPattern.MatchString(normalized) {
normalized += " LIMIT 1024"
}
tx, err := e.db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true})
if err != nil {
return nil, fmt.Errorf("vastbase: begin read-only sql transaction: %w", err)
}
defer tx.Rollback()
return queryRowsTx(ctx, tx, normalized)
}
+160
View File
@@ -0,0 +1,160 @@
//
// 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 vastbase
import (
"context"
"database/sql"
"database/sql/driver"
"fmt"
"regexp"
"testing"
"ragflow/internal/server/config"
"github.com/DATA-DOG/go-sqlmock"
)
// txOptionsConnector opens the sqlmock connection under its registered DSN
// and wraps it so tests can observe the driver.TxOptions database/sql hands
// to BeginTx; go-sqlmock accepts but discards them, so it offers no
// expectation API for transaction options.
type txOptionsConnector struct {
drv driver.Driver
dsn string
got *driver.TxOptions
}
func (c txOptionsConnector) Connect(context.Context) (driver.Conn, error) {
conn, err := c.drv.Open(c.dsn)
if err != nil {
return nil, err
}
return txOptionsConn{Conn: conn, got: c.got}, nil
}
func (c txOptionsConnector) Driver() driver.Driver { return c.drv }
type txOptionsConn struct {
driver.Conn
got *driver.TxOptions
}
func (c txOptionsConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
*c.got = opts
return c.Conn.(driver.ConnBeginTx).BeginTx(ctx, opts)
}
// QueryContext delegates so database/sql keeps using the query path instead
// of falling back to Prepare on the wrapper.
func (c txOptionsConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
return c.Conn.(driver.QueryerContext).QueryContext(ctx, query, args)
}
// newRunSQLEngine wires an Engine around a sqlmock database that records the
// transaction options in got, so tests can prove RunSQL opens its
// transaction READ ONLY.
func newRunSQLEngine(t *testing.T) (*Engine, sqlmock.Sqlmock, *driver.TxOptions) {
t.Helper()
db, mock, err := sqlmock.NewWithDSN(t.Name())
if err != nil {
t.Fatal(err)
}
got := new(driver.TxOptions)
recDB := sql.OpenDB(txOptionsConnector{drv: db.Driver(), dsn: t.Name(), got: got})
t.Cleanup(func() { recDB.Close(); db.Close() })
return newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, recDB), mock, got
}
// TestRunSQLExecutesInReadOnlyRolledBackTransaction checks the admin SQL
// console's read-only, always-rolled-back execution envelope.
func TestRunSQLExecutesInReadOnlyRolledBackTransaction(t *testing.T) {
engine, mock, gotTxOptions := newRunSQLEngine(t)
mock.ExpectBegin()
// No LIMIT in the input: the engine must append one.
mock.ExpectQuery(regexp.QuoteMeta(
`SELECT "id" FROM "t1" WHERE kb_id = 'kb-1' LIMIT 1024`)).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("c1"))
mock.ExpectRollback()
rows, err := engine.RunSQL(context.Background(), "t1",
`SELECT "id" FROM "t1" WHERE kb_id = 'kb-1'`, []string{"kb-1"}, "")
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 || rows[0]["id"] != "c1" {
t.Fatalf("rows = %#v", rows)
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
if !gotTxOptions.ReadOnly {
t.Fatal("RunSQL must open its transaction READ ONLY")
}
}
// TestRunSQLDataModifyingCTEFailsInsideReadOnlyTransaction checks that even
// data-modifying CTEs fail under the read-only envelope.
func TestRunSQLDataModifyingCTEFailsInsideReadOnlyTransaction(t *testing.T) {
engine, mock, gotTxOptions := newRunSQLEngine(t)
// The CTE passes the SELECT/WITH prefix checks; only the READ ONLY
// transaction stops it, with the server rejecting the write.
mock.ExpectBegin()
mock.ExpectQuery("WITH d AS").
WillReturnError(fmt.Errorf("cannot execute DELETE in a read-only transaction"))
mock.ExpectRollback()
if _, err := engine.RunSQL(context.Background(), "t1",
`WITH d AS (DELETE FROM "t1" RETURNING *) SELECT * FROM d`, nil, ""); err == nil {
t.Fatal("data-modifying CTE must fail inside the read-only transaction")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
if !gotTxOptions.ReadOnly {
t.Fatal("RunSQL must open its transaction READ ONLY")
}
}
// TestRunSQLRejectsMalformedInputBeforeTouchingTheDatabase checks statement
// validation happens before any connection is made.
func TestRunSQLRejectsMalformedInputBeforeTouchingTheDatabase(t *testing.T) {
db, mock, err := sqlmock.New()
if err != nil {
t.Fatal(err)
}
defer db.Close()
engine := newEngineWithDB(config.VastbaseConfig{DBCompatibility: compatPG}, db)
if _, err := engine.RunSQL(context.Background(), "bad;table", "SELECT 1", nil, ""); err == nil {
t.Fatal("invalid table names must be rejected")
}
for _, bad := range []string{
"DELETE FROM t1",
"UPDATE t1 SET x = 1",
"SELECT 1; DELETE FROM t1",
} {
if _, err := engine.RunSQL(context.Background(), "t1", bad, nil, ""); err == nil {
t.Fatalf("must reject %q", bad)
}
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Fatal(err)
}
}
+5
View File
@@ -169,6 +169,8 @@ func GetConfig() *config.Config {
return globalConfig
}
// GetAllConfigs collects every backend's connection settings as exportable
// maps, for diagnostics and the admin configuration view.
func GetAllConfigs() ([]map[string]interface{}, error) {
var allConfigs []map[string]interface{}
@@ -200,6 +202,9 @@ func GetAllConfigs() ([]map[string]interface{}, error) {
case "seekdb":
seekDBConfig := globalConfig.GetSeekDBConfig()
allConfigs = append(allConfigs, seekDBConfig.ExportConfigs())
case "vastbase":
vastbaseConfig := globalConfig.GetVastbaseConfig()
allConfigs = append(allConfigs, vastbaseConfig.ExportConfigs())
default:
return nil, fmt.Errorf("not supported doc engine: %s", docEngineType)
}
@@ -29,6 +29,7 @@ type DocEngineConfig struct {
OceanBase OceanBaseConfig `mapstructure:"oceanbase"`
SeekDB OceanBaseConfig `mapstructure:"seekdb"`
SereneDB SereneDBConfig `mapstructure:"serenedb"`
Vastbase VastbaseConfig `mapstructure:"vastbase"`
}
// OceanBaseConfig mirrors the existing oceanbase/seekdb service_conf.yaml
@@ -78,15 +79,81 @@ type SereneDBConfig struct {
SSLMode string `mapstructure:"ssl_mode"`
}
// VastbaseConfig Vastbase G100 configuration. Vastbase is PostgreSQL-wire
// based (database/sql + lib/pq). DBCompatibility selects the compatibility
// mode of the instance: "PG" (pgvector-style <=> cosine distance, GIN
// to_tsvector full-text index) or "B" (Oracle-flavored <+> distance operator,
// per-column ADD INDEX USING "fulltext").
type VastbaseConfig struct {
Host string `mapstructure:"host"`
Port int `mapstructure:"port"`
User string `mapstructure:"user"`
Password string `mapstructure:"password"`
DBName string `mapstructure:"db_name"`
DBCompatibility string `mapstructure:"dbcompatibility"`
// SSLMode is the lib/pq sslmode. Empty is not a plaintext license: the
// engine decides by host — disable on loopback or unix sockets,
// verify-full on every other TCP host. Plaintext to a non-loopback host
// (e.g. the bundled docker network) must set this explicitly.
SSLMode string `mapstructure:"ssl_mode"`
}
// ParseDocEngineConfig reads every document-engine section from the
// service configuration into the doc engine settings.
func (c *Config) ParseDocEngineConfig(v *viper.Viper) error {
c.parseInfinityConfig(v)
c.parseElasticsearchConfig(v)
c.parseOceanBaseConfig(v, "oceanbase", &c.docEngine.OceanBase)
c.parseOceanBaseConfig(v, "seekdb", &c.docEngine.SeekDB)
c.parseSereneDBConfig(v)
c.parseVastbaseConfig(v)
return nil
}
func (c *Config) parseVastbaseConfig(v *viper.Viper) {
// Default Vastbase config. SSLMode stays empty (unset): buildDSN picks
// disable for loopback hosts and verify-full for everything else, so a
// remote Host without an explicit ssl_mode never degrades to cleartext.
c.docEngine.Vastbase.Host = "vastbase"
c.docEngine.Vastbase.Port = 5432
c.docEngine.Vastbase.User = "ragflow"
c.docEngine.Vastbase.Password = "Infini_Rag@123"
c.docEngine.Vastbase.DBName = "ragflow"
c.docEngine.Vastbase.DBCompatibility = "PG"
// Fully qualified parent keys, not v.Sub("vastbase"): Sub drops the root
// AutomaticEnv (env prefix/replacer), and IsSet("vastbase") cannot see
// child environment variables, so an env-only deployment
// (RAGFLOW_VASTBASE_HOST, ... with no YAML section) would silently keep
// every default.
if v.IsSet("vastbase.host") {
c.docEngine.Vastbase.Host = v.GetString("vastbase.host")
}
if v.IsSet("vastbase.port") {
c.docEngine.Vastbase.Port = v.GetInt("vastbase.port")
}
if v.IsSet("vastbase.user") {
c.docEngine.Vastbase.User = v.GetString("vastbase.user")
}
if v.IsSet("vastbase.password") {
c.docEngine.Vastbase.Password = v.GetString("vastbase.password")
}
if v.IsSet("vastbase.db_name") {
c.docEngine.Vastbase.DBName = v.GetString("vastbase.db_name")
}
if v.IsSet("vastbase.dbcompatibility") {
c.docEngine.Vastbase.DBCompatibility = v.GetString("vastbase.dbcompatibility")
}
if v.IsSet("vastbase.ssl_mode") {
c.docEngine.Vastbase.SSLMode = v.GetString("vastbase.ssl_mode")
}
c.docEngine.Vastbase.DBCompatibility = strings.ToUpper(strings.TrimSpace(c.docEngine.Vastbase.DBCompatibility))
if c.docEngine.Vastbase.DBCompatibility != "B" {
c.docEngine.Vastbase.DBCompatibility = "PG"
}
}
func (c *Config) parseSereneDBConfig(v *viper.Viper) {
// Default SereneDB config
c.docEngine.SereneDB.Host = "localhost"
@@ -322,6 +389,23 @@ func (i InfinityConfig) ExportConfigs() map[string]interface{} {
return infinityConfigs
}
// GetSereneDBConfig returns the SereneDB doc engine connection settings.
func (c *Config) GetSereneDBConfig() SereneDBConfig {
return c.docEngine.SereneDB
}
func (c *Config) GetVastbaseConfig() VastbaseConfig {
return c.docEngine.Vastbase
}
func (v VastbaseConfig) ExportConfigs() map[string]interface{} {
return map[string]interface{}{
"host": v.Host,
"port": v.Port,
"user": v.User,
"password": v.Password,
"db_name": v.DBName,
"dbcompatibility": v.DBCompatibility,
"ssl_mode": v.SSLMode,
}
}
+3 -1
View File
@@ -114,6 +114,8 @@ type DefaultSuperUser struct {
Nickname string `mapstructure:"nickname"`
}
// GetEnvironments overlays environment variables on top of the file-based
// settings, filling the runtime environment snapshot.
func (c *Config) GetEnvironments() error {
// Language
@@ -152,7 +154,7 @@ func (c *Config) GetEnvironments() error {
docEngine := common.GetEnvSmall(common.EnvDocEngine)
if docEngine != "" {
switch docEngine {
case "infinity", "elasticsearch", "oceanbase", "seekdb":
case "infinity", "elasticsearch", "oceanbase", "seekdb", "vastbase":
c.environments.DocumentEngineType = docEngine
case "opensearch":
return fmt.Errorf("not implemented: %s", docEngine)
@@ -0,0 +1,135 @@
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package config
import (
"strings"
"testing"
"github.com/spf13/viper"
)
// TestParseVastbaseConfigDefaultsAndOverrides checks the config defaults
// against an empty viper instance and that file values override them.
func TestParseVastbaseConfigDefaultsAndOverrides(t *testing.T) {
t.Setenv("DB_TYPE", "")
v := viper.New()
c := &Config{}
if err := c.ParseDocEngineConfig(v); err != nil {
t.Fatal(err)
}
defaults := c.GetVastbaseConfig()
if defaults.Host != "vastbase" || defaults.Port != 5432 || defaults.User != "ragflow" ||
defaults.DBName != "ragflow" || defaults.DBCompatibility != "PG" || defaults.SSLMode != "" {
t.Fatalf("vastbase defaults = %#v", defaults)
}
v = viper.New()
v.Set("vastbase", map[string]interface{}{
"host": "vb-host", "port": 15432, "user": "vb-user",
"password": "vb-secret", "db_name": "rag_vec",
"dbcompatibility": "b", "ssl_mode": "verify-full",
})
c = &Config{}
if err := c.ParseDocEngineConfig(v); err != nil {
t.Fatal(err)
}
got := c.GetVastbaseConfig()
if got.Host != "vb-host" || got.Port != 15432 || got.User != "vb-user" ||
got.Password != "vb-secret" || got.DBName != "rag_vec" ||
got.DBCompatibility != "B" || got.SSLMode != "verify-full" {
t.Fatalf("vastbase overrides = %#v", got)
}
}
// TestParseVastbaseConfigNormalizesCompatibility checks the accepted
// spellings of the dbcompatibility setting and their canonical B/PG forms.
func TestParseVastbaseConfigNormalizesCompatibility(t *testing.T) {
for _, test := range []struct{ raw, want string }{
{" b ", "B"},
{"pg", "PG"},
{"nonsense", "PG"},
{"", "PG"},
} {
v := viper.New()
if test.raw != "" {
v.Set("vastbase", map[string]interface{}{"dbcompatibility": test.raw})
}
c := &Config{}
if err := c.ParseDocEngineConfig(v); err != nil {
t.Fatal(err)
}
if got := c.GetVastbaseConfig().DBCompatibility; got != test.want {
t.Fatalf("dbcompatibility(%q) = %q, want %q", test.raw, got, test.want)
}
}
}
// Env-only deployments (no service_conf.yaml) configure Vastbase through
// RAGFLOW_VASTBASE_* variables: the parser must read them through the parent
// viper's AutomaticEnv, not a Sub view that drops it.
func TestParseVastbaseConfigEnvironmentOnly(t *testing.T) {
t.Setenv("RAGFLOW_VASTBASE_HOST", "env-host")
t.Setenv("RAGFLOW_VASTBASE_PORT", "25432")
t.Setenv("RAGFLOW_VASTBASE_PASSWORD", "env-secret")
t.Setenv("RAGFLOW_VASTBASE_DBCOMPATIBILITY", "b")
v := viper.New()
v.SetEnvPrefix("RAGFLOW")
v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
v.AutomaticEnv()
c := &Config{}
if err := c.ParseDocEngineConfig(v); err != nil {
t.Fatal(err)
}
got := c.GetVastbaseConfig()
if got.Host != "env-host" || got.Port != 25432 ||
got.Password != "env-secret" || got.DBCompatibility != "B" {
t.Fatalf("vastbase env-only overrides = %#v", got)
}
// Fields without an env override keep their defaults (SSLMode stays
// unset: the engine derives it from the host).
if got.User != "ragflow" || got.DBName != "ragflow" || got.SSLMode != "" {
t.Fatalf("unset vastbase fields must keep defaults: %#v", got)
}
}
// TestVastbaseEnvironmentTypeIsAccepted checks that DOC_ENGINE=vastbase is a
// selectable engine type from the environment.
func TestVastbaseEnvironmentTypeIsAccepted(t *testing.T) {
t.Setenv("DOC_ENGINE", "vastbase")
c := &Config{}
if err := c.GetEnvironments(); err != nil {
t.Fatal(err)
}
if got := c.DocEngineType(); got != "vastbase" {
t.Fatalf("doc engine = %q, want vastbase", got)
}
}
// TestVastbaseExportConfigsKeys checks the exported vastbase configuration
// carries exactly the expected keys.
func TestVastbaseExportConfigsKeys(t *testing.T) {
exported := VastbaseConfig{
Host: "h", Port: 1, User: "u", Password: "p",
DBName: "d", DBCompatibility: "B", SSLMode: "disable",
}.ExportConfigs()
for _, key := range []string{"host", "port", "user", "password", "db_name", "dbcompatibility", "ssl_mode"} {
if _, present := exported[key]; !present {
t.Fatalf("ExportConfigs missing key %q: %#v", key, exported)
}
}
}
+3
View File
@@ -1233,6 +1233,9 @@ func TestExpectedDocNameColumn(t *testing.T) {
if got := expectedDocNameColumn("seekdb"); got != "docnm_kwd" {
t.Errorf("seekdb = %q, want docnm_kwd", got)
}
if got := expectedDocNameColumn("vastbase"); got != "docnm_kwd" {
t.Errorf("vastbase = %q, want docnm_kwd", got)
}
if got := expectedDocNameColumn("elasticsearch"); got != "docnm_kwd" {
t.Errorf("elasticsearch = %q, want docnm_kwd", got)
}
+4 -4
View File
@@ -922,7 +922,7 @@ func (s *MemoryService) DeleteMemory(ctx context.Context, userID, memoryID strin
}
// TODO: Delete associated message index - Implementation pending MessageService
if s.docEngine != nil && engine.IsOceanBaseFamily(s.docEngine.GetType()) {
if s.docEngine != nil && (engine.IsOceanBaseFamily(s.docEngine.GetType()) || engine.IsVastbase(s.docEngine.GetType())) {
if err := s.docEngine.DropChunkStore(ctx, MemoryIndexName(memory.TenantID), memoryID); err != nil {
return fmt.Errorf("delete memory messages: %w", err)
}
@@ -957,9 +957,9 @@ func (s *MemoryService) ForgetMessage(ctx context.Context, userID string, memory
updates := map[string]interface{}{
"forget_at": forgetTime,
}
// OceanBase/SeekDB memory tables contain forget_at but no forget_at_flt.
// Keep the existing companion-field update for other engines.
if !engine.IsOceanBaseFamily(s.docEngine.GetType()) {
// OceanBase/SeekDB/Vastbase memory tables contain forget_at but no
// forget_at_flt. Keep the existing companion-field update for other engines.
if !engine.IsOceanBaseFamily(s.docEngine.GetType()) && !engine.IsVastbase(s.docEngine.GetType()) {
updates["forget_at_flt"] = now.UnixMilli()
}
condition := map[string]interface{}{
+5
View File
@@ -315,6 +315,9 @@ func TestListMemoryFiltersKeepsCanonicalFacetOrder(t *testing.T) {
}
}
// TestForgetMessageKeepsCompanionFieldForNonOceanBaseEngines pins the
// forget_at companion column behavior for engines that store it as a plain
// nullable field.
func TestForgetMessageKeepsCompanionFieldForNonOceanBaseEngines(t *testing.T) {
setupMemoryMessageTestDB(t)
@@ -346,6 +349,8 @@ func TestForgetMessageKeepsCompanionFieldForNonOceanBaseEngines(t *testing.T) {
{engineType: "infinity", wantForgetAtCompanion: true},
{engineType: "oceanbase", wantForgetAtCompanion: false},
{engineType: "seekdb", wantForgetAtCompanion: false},
// Vastbase memory tables carry forget_at but no forget_at_flt column.
{engineType: "vastbase", wantForgetAtCompanion: false},
} {
t.Run(test.engineType, func(t *testing.T) {
docEngine := &memoryMessageDocEngine{engineType: test.engineType}
+8 -2
View File
@@ -447,6 +447,9 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest)
}, nil
}
// scoreSearchResult computes the per-chunk term, vector, and fused
// similarity scores for a search result, reranking through the configured
// model when the engine does not fuse natively.
func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *RetrievalRequest, searchResult *RetrievalSearchResult) ([]float64, []float64, []float64, error) {
// sim = tkWeight*tsim + vtWeight*vsim
vtWeight := *req.VectorSimilarityWeight
@@ -454,6 +457,9 @@ func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *Retrieval
qb := GetQueryBuilder()
useInfinity := engine.GetEngineType() == "infinity"
useOceanBase := engine.IsOceanBaseFamily(s.docEngine.GetType())
// Vastbase fusion returns a single fused score per chunk, so the
// tsim/vsim decomposition below needs the local rerank — same as OB.
useVastbase := engine.IsVastbase(s.docEngine.GetType())
if req.RerankModel != nil && searchResult.Total > 0 {
return RerankByModel(
@@ -487,7 +493,7 @@ func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *Retrieval
return sim, sim, sim, nil
}
if useOceanBase {
if useOceanBase || useVastbase {
sim, tsim, vsim := RerankStandard(
searchResult.Chunks,
nil,
@@ -783,7 +789,7 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque
// Build source with vector column for ES
searchSrc := make([]string, len(searchRequest.SelectFields))
copy(searchSrc, searchRequest.SelectFields)
if engine.GetEngineType() == "elasticsearch" || engine.IsOceanBaseFamily(engine.GetEngineType()) {
if engine.GetEngineType() == "elasticsearch" || engine.IsOceanBaseFamily(engine.GetEngineType()) || engine.IsVastbase(engine.GetEngineType()) {
searchSrc = append(searchSrc, matchDense.VectorColumnName)
}