diff --git a/backend/internal/infra/persistence/postgres/catalog_dialector.go b/backend/internal/infra/persistence/postgres/catalog_dialector.go new file mode 100644 index 000000000..fd324c438 --- /dev/null +++ b/backend/internal/infra/persistence/postgres/catalog_dialector.go @@ -0,0 +1,50 @@ +package postgres + +import ( + "gorm.io/driver/postgres" + "gorm.io/gorm" + "gorm.io/gorm/migrator" +) + +// catalogDialector keeps GORM's PostgreSQL migrator. On Nile, information_schema +// filters by the catalog name "nile", not current_database(). Ordinary Postgres +// is unchanged. +type catalogDialector struct { + *postgres.Dialector + nile bool +} + +func newPostgresDialector(pool gorm.ConnPool, nile bool) gorm.Dialector { + base := postgres.New(postgres.Config{Conn: pool}).(*postgres.Dialector) + return &catalogDialector{Dialector: base, nile: nile} +} + +func (d *catalogDialector) Migrator(db *gorm.DB) gorm.Migrator { + return catalogMigrator{Migrator: postgres.Migrator{Migrator: migrator.Migrator{Config: migrator.Config{ + DB: db, + Dialector: d, + CreateIndexAfterCreateTable: true, + }}}} +} + +type catalogMigrator struct { + postgres.Migrator +} + +func (m catalogMigrator) CurrentDatabase() (name string) { + if d, ok := m.Dialector.(*catalogDialector); ok && d.nile { + return nileInformationSchemaCatalog + } + return m.Migrator.CurrentDatabase() +} + +// StoresVectorIndexColumn reports whether candidate vectors live in their own +// column. Nile rejects functions inside an index expression, so the HNSW index +// is a plain halfvec column written with the row. +func StoresVectorIndexColumn(db *gorm.DB) bool { + if db == nil || db.Dialector == nil { + return false + } + dialector, ok := db.Dialector.(*catalogDialector) + return ok && dialector.nile +} diff --git a/backend/internal/infra/persistence/postgres/client_encoding.go b/backend/internal/infra/persistence/postgres/client_encoding.go new file mode 100644 index 000000000..a33543b1e --- /dev/null +++ b/backend/internal/infra/persistence/postgres/client_encoding.go @@ -0,0 +1,93 @@ +package postgres + +import ( + "context" + "reflect" + "regexp" + "strings" + "time" + "unsafe" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + "github.com/jackc/pgx/v5/stdlib" + "gorm.io/gorm" +) + +// postgresTimeZonePattern matches the timezone handling in gorm.io/driver/postgres. +var postgresTimeZonePattern = regexp.MustCompile(`(time_zone|TimeZone|timezone)=(.*?)($|&| )`) + +// openPostgres matches gorm's pgx setup. Nile reports client_encoding as +// "UTF-8", which makes pgx reject later simple-protocol queries. Only a Nile +// host is rewritten, and only when that token is a UTF-8 spelling other than +// "UTF8". Ordinary Postgres already reports "UTF8" and is left untouched. +func openPostgres(dsn string, nile bool) (gorm.ConnPool, error) { + config, err := pgx.ParseConfig(dsn) + if err != nil { + return nil, err + } + + if nile { + configureNileSession(config) + } + + var timestampLocation *time.Location + if match := postgresTimeZonePattern.FindStringSubmatch(dsn); len(match) > 2 { + config.RuntimeParams["timezone"] = match[2] + timestampLocation, err = time.LoadLocation(match[2]) + if err != nil { + return nil, err + } + } + + return stdlib.OpenDB(*config, stdlib.OptionAfterConnect(func(ctx context.Context, conn *pgx.Conn) error { + if nile { + normalizeReportedClientEncoding(conn) + } + if timestampLocation == nil { + return nil + } + conn.TypeMap().RegisterType(&pgtype.Type{ + Name: "timestamp", + OID: pgtype.TimestampOID, + Codec: &pgtype.TimestampCodec{ScanLocation: timestampLocation}, + }) + return nil + })), nil +} + +// configureNileSession makes parameterized queries use the simple protocol. +// Nile's extended protocol delivers a bound empty string as NULL, so a NOT NULL +// text column rejects it. Inlining the value keeps the empty string. +func configureNileSession(config *pgx.ConnConfig) { + config.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol +} + +func normalizeReportedClientEncoding(conn *pgx.Conn) { + if conn == nil || conn.PgConn() == nil { + return + } + field := reflect.ValueOf(conn.PgConn()).Elem().FieldByName("parameterStatuses") + if !field.IsValid() || field.Kind() != reflect.Map || field.IsNil() || !field.CanAddr() { + return + } + statuses := reflect.NewAt(field.Type(), unsafe.Pointer(field.UnsafeAddr())).Elem() + key := reflect.ValueOf("client_encoding") + current := statuses.MapIndex(key) + if !current.IsValid() || current.Kind() != reflect.String { + return + } + normalized := canonicalClientEncoding(current.String()) + if normalized == "" || normalized == current.String() { + return + } + statuses.SetMapIndex(key, reflect.ValueOf(normalized)) +} + +func canonicalClientEncoding(value string) string { + compact := strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(value), "-", "")) + if compact == "UTF8" { + return "UTF8" + } + return "" +} diff --git a/backend/internal/infra/persistence/postgres/client_encoding_test.go b/backend/internal/infra/persistence/postgres/client_encoding_test.go new file mode 100644 index 000000000..122506ccb --- /dev/null +++ b/backend/internal/infra/persistence/postgres/client_encoding_test.go @@ -0,0 +1,39 @@ +package postgres + +import ( + "testing" + + "github.com/jackc/pgx/v5" +) + +func TestCanonicalClientEncoding(t *testing.T) { + tests := []struct { + in string + want string + }{ + {in: "UTF8", want: "UTF8"}, + {in: "UTF-8", want: "UTF8"}, + {in: " utf-8 ", want: "UTF8"}, + {in: "SQL_ASCII", want: ""}, + {in: "", want: ""}, + } + for _, tt := range tests { + if got := canonicalClientEncoding(tt.in); got != tt.want { + t.Fatalf("canonicalClientEncoding(%q) = %q, want %q", tt.in, got, tt.want) + } + } +} + +func TestConfigureNileSessionUsesSimpleProtocol(t *testing.T) { + cfg, err := pgx.ParseConfig("postgres://user:secret@127.0.0.1:5432/deeix") + if err != nil { + t.Fatal(err) + } + if cfg.DefaultQueryExecMode == pgx.QueryExecModeSimpleProtocol { + t.Fatal("ordinary postgres must not start in the simple protocol") + } + configureNileSession(cfg) + if cfg.DefaultQueryExecMode != pgx.QueryExecModeSimpleProtocol { + t.Fatalf("Nile query mode = %v", cfg.DefaultQueryExecMode) + } +} diff --git a/backend/internal/infra/persistence/postgres/conversation/repository.go b/backend/internal/infra/persistence/postgres/conversation/repository.go index ba477144f..da47a3f0e 100644 --- a/backend/internal/infra/persistence/postgres/conversation/repository.go +++ b/backend/internal/infra/persistence/postgres/conversation/repository.go @@ -13,6 +13,7 @@ import ( domainuser "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/domain/user" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/dberror" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/models" + persistdb "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/postgres" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/sqlitevec" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/vectorutil" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/pkg/textutil" @@ -3096,11 +3097,17 @@ func (r *Repo) CloneFileEmbeddingArtifacts(ctx context.Context, source *domainco } return nil } + insertEmbeddings := `"embedding"` + selectEmbeddings := `"embedding"` + if persistdb.StoresVectorIndexColumn(r.db) { + insertEmbeddings = `"embedding", "` + vectorutil.IndexColumn + `"` + selectEmbeddings = `"embedding", "` + vectorutil.IndexColumn + `"` + } return tx.Exec( - `INSERT INTO "file_chunks" ("file_obj_id", "user_id", "chunk_index", "page_num", "char_offset", "content", "token_count", "embedding_signature", "embedding", "created_at") - SELECT ?, ?, "chunk_index", "page_num", "char_offset", "content", "token_count", "embedding_signature", "embedding", NOW() + fmt.Sprintf(`INSERT INTO "file_chunks" ("file_obj_id", "user_id", "chunk_index", "page_num", "char_offset", "content", "token_count", "embedding_signature", %s, "created_at") + SELECT ?, ?, "chunk_index", "page_num", "char_offset", "content", "token_count", "embedding_signature", %s, NOW() FROM "file_chunks" - WHERE "file_obj_id" = ?`, + WHERE "file_obj_id" = ?`, insertEmbeddings, selectEmbeddings), targetEntity.ID, targetEntity.UserID, sourceEntity.ID, @@ -3167,16 +3174,9 @@ func (r *Repo) ReplaceFileChunks(ctx context.Context, fileObjID uint, embeddingS if len(embeddings[i]) == 0 { return fmt.Errorf("empty embedding vector at chunk %d", i) } - vec, err := float32SliceToPostgresVector(embeddings[i]) - if err != nil { + if err := r.updatePostgresEmbedding(tx, `"file_chunks"`, embeddings[i], chunk.ID); err != nil { return err } - if err := tx.Exec( - `UPDATE "file_chunks" SET embedding = ? WHERE id = ?`, - vec, chunk.ID, - ).Error; err != nil { - return dberror.Translate(err) - } } published = true return nil @@ -3224,6 +3224,43 @@ func float32SliceToPostgresQueryVector(v []float32) (string, error) { return vectorutil.PostgresPaddedLiteral(v) } +func (r *Repo) postgresCandidate(columnRef string, expression string) (string, string) { + if persistdb.StoresVectorIndexColumn(r.db) { + return columnRef, fmt.Sprintf("?::halfvec(%d)", vectorutil.IndexDimensions) + } + return expression, fmt.Sprintf("subvector(?::vector, 1, %d)::halfvec(%d)", vectorutil.IndexDimensions, vectorutil.IndexDimensions) +} + +func (r *Repo) postgresCandidateArgument(embedding []float32, exact string) (string, error) { + if persistdb.StoresVectorIndexColumn(r.db) { + return vectorutil.PostgresIndexLiteral(embedding) + } + return exact, nil +} + +func (r *Repo) updatePostgresEmbedding(tx *gorm.DB, table string, embedding []float32, id uint) error { + vec, err := float32SliceToPostgresVector(embedding) + if err != nil { + return err + } + if persistdb.StoresVectorIndexColumn(r.db) { + indexVec, indexErr := vectorutil.PostgresIndexLiteral(embedding) + if indexErr != nil { + return indexErr + } + err = tx.Exec( + fmt.Sprintf(`UPDATE %s SET embedding = ?, %s = ?::halfvec(%d) WHERE id = ?`, table, vectorutil.IndexColumn, vectorutil.IndexDimensions), + vec, indexVec, id, + ).Error + } else { + err = tx.Exec(fmt.Sprintf(`UPDATE %s SET embedding = ? WHERE id = ?`, table), vec, id).Error + } + if err != nil { + return dberror.Translate(err) + } + return nil +} + // fileChunkSearchRow 是原始 SQL 扫描专用的本地类型，携带 gorm column tag 映射相似度列。 type fileChunkSearchRow struct { ID uint `gorm:"column:id"` @@ -3394,8 +3431,12 @@ func (r *Repo) SearchFileChunks(ctx context.Context, userID uint, fileObjIDs []u if err != nil { return nil, err } + candidateVec, err := r.postgresCandidateArgument(queryEmbedding, vec) + if err != nil { + return nil, err + } candidateLimit := vectorutil.CandidateLimit(topK) - indexExpression := vectorutil.PostgresIndexExpression("source_chunks.embedding") + candidateLeft, candidateRight := r.postgresCandidate("source_chunks."+vectorutil.IndexColumn, vectorutil.PostgresIndexExpression("source_chunks.embedding")) exactExpression := vectorutil.PostgresPaddedExpression("chunks.embedding") query := fmt.Sprintf(` WITH vector_candidates AS MATERIALIZED ( @@ -3415,8 +3456,7 @@ func (r *Repo) SearchFileChunks(ctx context.Context, userID uint, fileObjIDs []u AND kb.enabled = ? ) ) - ORDER BY %s - <=> subvector(?::vector, 1, %d)::halfvec(%d) + ORDER BY %s <=> %s LIMIT ? ) SELECT chunks.id, chunks.file_obj_id, chunks.user_id, chunks.chunk_index, chunks.page_num, @@ -3426,9 +3466,8 @@ func (r *Repo) SearchFileChunks(ctx context.Context, userID uint, fileObjIDs []u JOIN vector_candidates AS candidates ON candidates.id = chunks.id ORDER BY similarity DESC LIMIT ?`, - indexExpression, - vectorutil.IndexDimensions, - vectorutil.IndexDimensions, + candidateLeft, + candidateRight, exactExpression, vectorutil.MaxDimensions, ) @@ -3444,7 +3483,7 @@ func (r *Repo) SearchFileChunks(ctx context.Context, userID uint, fileObjIDs []u userID, domainknowledgebase.ScopeBuiltin, true, - vec, + candidateVec, candidateLimit, vec, topK, @@ -4691,6 +4730,10 @@ func (r *Repo) VectorStoreAvailable(ctx context.Context) (bool, error) { AND NOT attribute.attisdropped AND format_type(attribute.atttypid, attribute.atttypmod) = ? )` + nile := persistdb.StoresVectorIndexColumn(r.db) + if nile { + columnQuery = persistdb.NileEmbeddingColumnExistsSQL() + } indexQuery := `SELECT EXISTS ( SELECT 1 FROM pg_index AS index_status @@ -4700,19 +4743,30 @@ func (r *Repo) VectorStoreAvailable(ctx context.Context) (bool, error) { AND index_relation.relname = ? AND index_status.indisvalid AND lower(pg_get_indexdef(index_status.indexrelid)) LIKE '% using hnsw %' - AND lower(pg_get_indexdef(index_status.indexrelid)) LIKE ? - AND lower(pg_get_indexdef(index_status.indexrelid)) LIKE '%vector_dims(%' AND lower(pg_get_indexdef(index_status.indexrelid)) LIKE '%halfvec_cosine_ops%' + AND ( + ( + lower(pg_get_indexdef(index_status.indexrelid)) LIKE ? + AND lower(pg_get_indexdef(index_status.indexrelid)) LIKE '%vector_dims(%' + ) + OR lower(pg_get_indexdef(index_status.indexrelid)) LIKE '%embedding_hnsw%' + ) )` - indexPattern := fmt.Sprintf("%%::halfvec(%d)%%", vectorutil.IndexDimensions) + indexArgs := func(name string) []any { + return []any{name, fmt.Sprintf("%%::halfvec(%d)%%", vectorutil.IndexDimensions)} + } + if nile { + indexQuery = persistdb.NileVectorIndexExistsSQL() + indexArgs = func(name string) []any { return []any{name} } + } checks := []availabilityCheck{ {query: `SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector')`}, {query: columnQuery, args: []any{"file_chunks", expectedType}}, {query: columnQuery, args: []any{"chat_message_chunks", expectedType}}, {query: columnQuery, args: []any{"user_memories", expectedType}}, - {query: indexQuery, args: []any{"idx_file_chunks_embedding", indexPattern}}, - {query: indexQuery, args: []any{"idx_chat_message_chunks_embedding", indexPattern}}, - {query: indexQuery, args: []any{"idx_user_memories_embedding", indexPattern}}, + {query: indexQuery, args: indexArgs("idx_file_chunks_embedding")}, + {query: indexQuery, args: indexArgs("idx_chat_message_chunks_embedding")}, + {query: indexQuery, args: indexArgs("idx_user_memories_embedding")}, } for _, check := range checks { available := false @@ -4774,13 +4828,9 @@ func (r *Repo) UpsertMessageChunks(ctx context.Context, chunks []domainconversat if i >= len(embeddings) || len(embeddings[i]) == 0 { continue } - vec, err := float32SliceToPostgresVector(embeddings[i]) - if err != nil { + if err := r.updatePostgresEmbedding(tx, `"chat_message_chunks"`, embeddings[i], entity.ID); err != nil { return err } - if err := tx.Exec(`UPDATE "chat_message_chunks" SET embedding = ? WHERE id = ?`, vec, entity.ID).Error; err != nil { - return dberror.Translate(err) - } } return nil }) @@ -4906,9 +4956,13 @@ func (r *Repo) SearchMessageChunks(ctx context.Context, input repository.Message if err != nil { return nil, err } + candidateVec, err := r.postgresCandidateArgument(input.QueryEmbedding, vec) + if err != nil { + return nil, err + } candidateLimit := vectorutil.CandidateLimit(input.TopK) // 候选阶段已经限定当前分支和向量签名，随后再按完整 4096 维向量精确重排。 - indexExpression := vectorutil.PostgresIndexExpression("chunks.embedding") + candidateLeft, candidateRight := r.postgresCandidate("chunks."+vectorutil.IndexColumn, vectorutil.PostgresIndexExpression("chunks.embedding")) exactExpression := vectorutil.PostgresPaddedExpression("chunks.embedding") query := historicalMessageScopeCTE + fmt.Sprintf(`, vector_candidates AS MATERIALIZED ( @@ -4919,8 +4973,7 @@ func (r *Repo) SearchMessageChunks(ctx context.Context, input repository.Message AND chunks.embedding_signature = ? AND chunks.embedding IS NOT NULL AND chunks.message_id IN (SELECT id FROM valid_historical_message_scope) - ORDER BY %s - <=> subvector(?::vector, 1, %d)::halfvec(%d) + ORDER BY %s <=> %s LIMIT ? ) SELECT chunks.id, chunks.conversation_id, chunks.message_id, chunks.user_id, chunks.role, @@ -4930,9 +4983,8 @@ func (r *Repo) SearchMessageChunks(ctx context.Context, input repository.Message JOIN vector_candidates AS candidates ON candidates.id = chunks.id ORDER BY similarity DESC LIMIT ?`, - indexExpression, - vectorutil.IndexDimensions, - vectorutil.IndexDimensions, + candidateLeft, + candidateRight, exactExpression, vectorutil.MaxDimensions, ) @@ -4941,7 +4993,7 @@ func (r *Repo) SearchMessageChunks(ctx context.Context, input repository.Message input.Scope.ConversationID, input.Scope.UserID, input.EmbeddingSignature, - vec, + candidateVec, candidateLimit, vec, input.TopK, diff --git a/backend/internal/infra/persistence/postgres/memory/repository.go b/backend/internal/infra/persistence/postgres/memory/repository.go index 94e8c333e..ccae4bb79 100644 --- a/backend/internal/infra/persistence/postgres/memory/repository.go +++ b/backend/internal/infra/persistence/postgres/memory/repository.go @@ -8,6 +8,7 @@ import ( domainmemory "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/domain/memory" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/dberror" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/models" + persistdb "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/postgres" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/sqlitevec" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/vectorutil" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/repository" @@ -96,7 +97,11 @@ func (r *Repo) clearUserMemoryEmbedding(ctx context.Context, tx *gorm.DB, memory if !r.postgresUserMemoryEmbeddingColumnAvailable(ctx, tx) { return nil } - return dberror.Translate(tx.Exec(`UPDATE "user_memories" SET embedding = NULL, embedding_signature = '' WHERE id = ?`, memoryID).Error) + query := `UPDATE "user_memories" SET embedding = NULL, embedding_signature = '' WHERE id = ?` + if persistdb.StoresVectorIndexColumn(r.db) { + query = fmt.Sprintf(`UPDATE "user_memories" SET embedding = NULL, %s = NULL, embedding_signature = '' WHERE id = ?`, vectorutil.IndexColumn) + } + return dberror.Translate(tx.Exec(query, memoryID).Error) } func (r *Repo) postgresUserMemoryEmbeddingColumnAvailable(ctx context.Context, tx *gorm.DB) bool { @@ -183,16 +188,25 @@ func (r *Repo) SearchUserMemoriesByEmbedding(ctx context.Context, userID uint, q if err != nil { return nil, err } + candidateVec := vec + candidateLeft := vectorutil.PostgresIndexExpression("embedding") + candidateRight := fmt.Sprintf("subvector(?::vector, 1, %d)::halfvec(%d)", vectorutil.IndexDimensions, vectorutil.IndexDimensions) + if persistdb.StoresVectorIndexColumn(r.db) { + candidateVec, err = vectorutil.PostgresIndexLiteral(queryEmbedding) + if err != nil { + return nil, err + } + candidateLeft = vectorutil.IndexColumn + candidateRight = fmt.Sprintf("?::halfvec(%d)", vectorutil.IndexDimensions) + } candidateLimit := vectorutil.CandidateLimit(topK) - indexExpression := vectorutil.PostgresIndexExpression("embedding") exactExpression := vectorutil.PostgresPaddedExpression("memories.embedding") query := fmt.Sprintf(` WITH vector_candidates AS MATERIALIZED ( SELECT id FROM user_memories WHERE user_id = ? AND embedding_signature = ? AND embedding IS NOT NULL - ORDER BY %s - <=> subvector(?::vector, 1, %d)::halfvec(%d) + ORDER BY %s <=> %s LIMIT ? ) SELECT memories.id, memories.user_id, memories.memory_key, memories.value, memories.scope, memories.updated_by, @@ -201,9 +215,8 @@ func (r *Repo) SearchUserMemoriesByEmbedding(ctx context.Context, userID uint, q JOIN vector_candidates AS candidates ON candidates.id = memories.id ORDER BY similarity DESC LIMIT ?`, - indexExpression, - vectorutil.IndexDimensions, - vectorutil.IndexDimensions, + candidateLeft, + candidateRight, exactExpression, vectorutil.MaxDimensions, ) @@ -212,7 +225,7 @@ func (r *Repo) SearchUserMemoriesByEmbedding(ctx context.Context, userID uint, q if err := vectorutil.ConfigurePostgresCandidateSearch(tx); err != nil { return err } - return tx.Raw(query, userID, embeddingSignature, vec, candidateLimit, vec, topK).Scan(&rows).Error + return tx.Raw(query, userID, embeddingSignature, candidateVec, candidateLimit, vec, topK).Scan(&rows).Error }); err != nil { return nil, dberror.Translate(err) } @@ -285,6 +298,14 @@ func (r *Repo) UpsertUserMemoryEmbedding(ctx context.Context, userID uint, memor } query := `UPDATE "user_memories" SET embedding = ?::vector, embedding_signature = ? WHERE user_id = ? AND memory_key = ?` args := []any{vec, embeddingSignature, userID, memoryKey} + if persistdb.StoresVectorIndexColumn(r.db) { + indexVec, indexErr := vectorutil.PostgresIndexLiteral(embedding) + if indexErr != nil { + return indexErr + } + query = fmt.Sprintf(`UPDATE "user_memories" SET embedding = ?::vector, %s = ?::halfvec(%d), embedding_signature = ? WHERE user_id = ? AND memory_key = ?`, vectorutil.IndexColumn, vectorutil.IndexDimensions) + args = []any{vec, indexVec, embeddingSignature, userID, memoryKey} + } if strings.TrimSpace(expectedValue) != "" { query += ` AND value = ?` args = append(args, strings.TrimSpace(expectedValue)) diff --git a/backend/internal/infra/persistence/postgres/nile.go b/backend/internal/infra/persistence/postgres/nile.go new file mode 100644 index 000000000..3e4a00ea0 --- /dev/null +++ b/backend/internal/infra/persistence/postgres/nile.go @@ -0,0 +1,156 @@ +package postgres + +import ( + "strings" + + "github.com/jackc/pgx/v5" + "gorm.io/gorm" +) + +const nileInformationSchemaCatalog = "nile" + +// postgresHostIsNile reports whether the DSN host is a Nile database endpoint. +// The database name is ignored: a local database named "nile" is not Nile. +func postgresHostIsNile(dsn string) bool { + cfg, err := pgx.ParseConfig(dsn) + if err != nil { + return false + } + host := strings.ToLower(strings.Trim(strings.TrimSpace(cfg.Host), ".")) + return host == "thenile.dev" || strings.HasSuffix(host, ".thenile.dev") +} + +func registerNileConstraintRewrite(db *gorm.DB) { + rewrite := func(tx *gorm.DB) { + rewritten, vars, ok := rewriteNileConstraintQuery(tx.Statement.SQL.String(), tx.Statement.Vars) + if !ok { + return + } + tx.Statement.SQL.Reset() + tx.Statement.SQL.WriteString(rewritten) + tx.Statement.Vars = vars + } + _ = db.Callback().Row().Before("gorm:row").Register("deeix:nile_constraint", rewrite) + _ = db.Callback().Query().Before("gorm:query").Register("deeix:nile_constraint", rewrite) + _ = db.Callback().Raw().Before("gorm:raw").Register("deeix:nile_constraint", rewrite) +} + +// rewriteNileConstraintQuery replaces GORM's two information_schema constraint +// joins. On Nile those joins take about 18s and, when filtered by +// current_database(), return no rows. The replacement reads pg_constraint, +// which is the same set: primary keys and unique constraints, one row per +// key column, and not unique indexes that are not constraints. +func rewriteNileConstraintQuery(sql string, vars []any) (string, []any, bool) { + compact := strings.Join(strings.Fields(sql), " ") + switch { + case isGORMUniqueConstraintNameQuery(compact): + return buildNileTableQuery(compact, vars, nileUniqueConstraintNamesSQL, true, true) + case isGORMConstraintColumnQuery(compact): + return buildNileTableQuery(compact, vars, nileConstraintColumnsSQL, false, true) + case isGORMColumnMetadataQuery(compact): + return buildNileTableQuery(compact, vars, nileColumnMetadataSQL, false, true) + case isGORMFormatTypeQuery(compact): + return buildNileTableQuery(compact, vars, nileFormatTypeSQL, false, false) + default: + return "", nil, false + } +} + +func isGORMUniqueConstraintNameQuery(compact string) bool { + return strings.Contains(compact, "SELECT constraint_name FROM information_schema.table_constraints") && + strings.Contains(compact, "constraint_column_usage") && + strings.Contains(compact, "constraint_type =") +} + +func isGORMConstraintColumnQuery(compact string) bool { + return strings.Contains(compact, "SELECT c.column_name, constraint_name, constraint_type FROM information_schema.table_constraints") && + strings.Contains(compact, "constraint_column_usage") +} + +func isGORMColumnMetadataQuery(compact string) bool { + return strings.Contains(compact, "FROM information_schema.columns") && + strings.Contains(compact, "8 * typlen") && + strings.Contains(compact, "c.identity_increment") +} + +func isGORMFormatTypeQuery(compact string) bool { + return strings.Contains(compact, "format_type(a.atttypid, a.atttypmod)") && + strings.Contains(compact, "pg_attribute a JOIN pg_class b") && + strings.Contains(compact, "b.relname") +} + +const nileUniqueConstraintNamesSQL = ` +SELECT con.conname +FROM pg_constraint AS con +JOIN pg_attribute AS a ON a.attrelid = con.conrelid AND a.attnum = ANY (con.conkey) AND NOT a.attisdropped AND a.attnum > 0 +WHERE con.conrelid = format('%I.%I', %s, %s)::regclass + AND con.contype = 'u'` + +const nileConstraintColumnsSQL = ` +SELECT a.attname, con.conname, + CASE con.contype WHEN 'p' THEN 'PRIMARY KEY' WHEN 'u' THEN 'UNIQUE' END +FROM pg_constraint AS con +JOIN pg_attribute AS a ON a.attrelid = con.conrelid AND a.attnum = ANY (con.conkey) AND NOT a.attisdropped AND a.attnum > 0 +WHERE con.conrelid = format('%I.%I', %s, %s)::regclass + AND con.contype IN ('p', 'u')` + +const nileFormatTypeSQL = ` +SELECT a.attname AS column_name, format_type(a.atttypid, a.atttypmod) AS data_type +FROM pg_attribute AS a +WHERE a.attrelid = format('%I.%I', %s, %s)::regclass + AND a.attnum > 0 + AND NOT a.attisdropped` + +// nileColumnMetadataSQL returns the same twelve values GORM scans from +// information_schema.columns. Serial columns keep their nextval() default; +// identity_increment stays null, which is what that view returns for them. +const nileColumnMetadataSQL = ` +SELECT a.attname, + NOT a.attnotnull, + t.typname, + information_schema._pg_char_max_length(a.atttypid, a.atttypmod), + information_schema._pg_numeric_precision(a.atttypid, a.atttypmod), + information_schema._pg_numeric_precision_radix(a.atttypid, a.atttypmod), + information_schema._pg_numeric_scale(a.atttypid, a.atttypmod), + information_schema._pg_datetime_precision(a.atttypid, a.atttypmod), + 8 * t.typlen, + pg_get_expr(ad.adbin, ad.adrelid), + col_description(a.attrelid, a.attnum), + NULL::text +FROM pg_attribute AS a +JOIN pg_type AS t ON t.oid = a.atttypid +LEFT JOIN pg_attrdef AS ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum +WHERE a.attrelid = format('%I.%I', %s, %s)::regclass + AND a.attnum > 0 + AND NOT a.attisdropped` + +func buildNileTableQuery(compact string, vars []any, pattern string, uniqueOnly bool, leadingCatalog bool) (string, []any, bool) { + args := append([]any(nil), vars...) + if uniqueOnly && len(args) > 0 { + if kind, ok := args[len(args)-1].(string); ok && strings.EqualFold(kind, "UNIQUE") { + args = args[:len(args)-1] + } + } + if strings.Contains(strings.ToUpper(compact), "CURRENT_SCHEMA()") { + index := 0 + if leadingCatalog { + index = 1 + } + if len(args) <= index { + return "", nil, false + } + return sprintfNileSQL(pattern, "current_schema()", "$1::text"), []any{args[index]}, true + } + schemaIndex, tableIndex := 0, 1 + if leadingCatalog { + schemaIndex, tableIndex = 1, 2 + } + if len(args) <= tableIndex { + return "", nil, false + } + return sprintfNileSQL(pattern, "$1::text", "$2::text"), []any{args[schemaIndex], args[tableIndex]}, true +} + +func sprintfNileSQL(pattern string, schemaExpr string, tableExpr string) string { + return strings.TrimSpace(strings.Replace(strings.Replace(pattern, "%s", schemaExpr, 1), "%s", tableExpr, 1)) +} diff --git a/backend/internal/infra/persistence/postgres/nile_test.go b/backend/internal/infra/persistence/postgres/nile_test.go new file mode 100644 index 000000000..7b96358ba --- /dev/null +++ b/backend/internal/infra/persistence/postgres/nile_test.go @@ -0,0 +1,115 @@ +package postgres + +import ( + "strings" + "testing" +) + +func TestPostgresHostIsNile(t *testing.T) { + tests := []struct { + name string + dsn string + want bool + }{ + {name: "regional nile host", dsn: "postgres://user:secret@us-west-2.db.thenile.dev:5432/deeix_pgdb?sslmode=require", want: true}, + {name: "uppercase host", dsn: "postgres://user:secret@US-WEST-2.DB.THENILE.DEV:5432/deeix_pgdb", want: true}, + {name: "keyword dsn", dsn: "host=eu-central-1.db.thenile.dev user=user password=secret dbname=deeix_pgdb port=5432", want: true}, + {name: "apex", dsn: "postgres://user:secret@thenile.dev:5432/deeix_pgdb", want: true}, + {name: "local database named nile", dsn: "postgres://user:secret@127.0.0.1:5432/nile", want: false}, + {name: "host merely contains the word", dsn: "postgres://user:secret@nile.example.com:5432/deeix", want: false}, + {name: "keyword local", dsn: "host=127.0.0.1 user=deeix_chat password=secret dbname=nile port=5432", want: false}, + {name: "empty", dsn: "", want: false}, + {name: "not a dsn", dsn: "://", want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := postgresHostIsNile(tt.dsn); got != tt.want { + t.Fatalf("postgresHostIsNile(%q) = %v, want %v", tt.dsn, got, tt.want) + } + }) + } +} + +func TestRewriteNileConstraintQueryUsesPgConstraint(t *testing.T) { + nameSQL := "SELECT constraint_name FROM information_schema.table_constraints tc JOIN information_schema.constraint_column_usage AS ccu USING (constraint_schema, constraint_catalog, table_name, constraint_name) JOIN information_schema.columns AS c ON c.table_schema = tc.constraint_schema AND tc.table_name = c.table_name AND ccu.column_name = c.column_name WHERE constraint_type IN ('PRIMARY KEY', 'UNIQUE') AND c.table_catalog = $1 AND c.table_schema = CURRENT_SCHEMA() AND c.table_name = $2 AND constraint_type = $3" + got, vars, ok := rewriteNileConstraintQuery(nameSQL, []interface{}{"nile", "identity_users", "UNIQUE"}) + if !ok { + t.Fatal("expected the unique-constraint name query to be rewritten") + } + if strings.Contains(got, "information_schema") || strings.Contains(got, "pg_index") { + t.Fatalf("rewritten query is not the constraint lookup: %s", got) + } + if !strings.Contains(got, "pg_constraint") || !strings.Contains(got, "con.contype = 'u'") || !strings.Contains(got, "format('%I.%I', current_schema(), $1::text)::regclass") { + t.Fatalf("unexpected unique-constraint query: %s", got) + } + if len(vars) != 1 || vars[0] != "identity_users" { + t.Fatalf("vars = %#v", vars) + } + + columnSQL := "SELECT c.column_name, constraint_name, constraint_type FROM information_schema.table_constraints tc JOIN information_schema.constraint_column_usage AS ccu USING (constraint_schema, constraint_catalog, table_name, constraint_name) JOIN information_schema.columns AS c ON c.table_schema = tc.constraint_schema AND tc.table_name = c.table_name AND ccu.column_name = c.column_name WHERE constraint_type IN ('PRIMARY KEY', 'UNIQUE') AND c.table_catalog = $1 AND c.table_schema = $2 AND c.table_name = $3" + got, vars, ok = rewriteNileConstraintQuery(columnSQL, []interface{}{"deeix_pgdb", "public", "chat_messages"}) + if !ok { + t.Fatal("expected the constraint column query to be rewritten") + } + if !strings.Contains(got, "WHEN 'p' THEN 'PRIMARY KEY'") || !strings.Contains(got, "WHEN 'u' THEN 'UNIQUE'") { + t.Fatalf("constraint types GORM reads were not preserved: %s", got) + } + if !strings.Contains(got, "format('%I.%I', $1::text, $2::text)::regclass") || strings.Contains(got, "pg_index") || strings.Contains(got, "information_schema") { + t.Fatalf("unexpected constraint column query: %s", got) + } + if len(vars) != 2 || vars[0] != "public" || vars[1] != "chat_messages" { + t.Fatalf("vars = %#v", vars) + } +} + +func TestRewriteNileConstraintQueryLeavesOtherCatalogSQL(t *testing.T) { + others := []string{ + "SELECT c.column_name, c.is_nullable = 'YES', c.udt_name FROM information_schema.columns AS c WHERE table_catalog = $1 AND table_schema = CURRENT_SCHEMA() AND table_name = $2", + "SELECT count(*) FROM INFORMATION_SCHEMA.table_constraints WHERE table_schema = $1 AND table_name = $2 AND constraint_name = $3", + "SELECT table_name FROM information_schema.tables WHERE table_schema = $1 AND table_type = $2", + "SELECT 1", + } + for _, sql := range others { + if _, _, ok := rewriteNileConstraintQuery(sql, []interface{}{"nile", "public", "identity_users"}); ok { + t.Fatalf("unrelated SQL was rewritten: %s", sql) + } + } +} + +func TestRewriteNileSlowColumnQueries(t *testing.T) { + formatSQL := "SELECT a.attname as column_name, format_type(a.atttypid, a.atttypmod) AS data_type FROM pg_attribute a JOIN pg_class b ON a.attrelid = b.oid AND relnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = CURRENT_SCHEMA()) WHERE a.attnum > 0 AND NOT a.attisdropped AND b.relname = $1" + got, vars, ok := rewriteNileConstraintQuery(formatSQL, []interface{}{"chat_messages"}) + if !ok { + t.Fatal("expected format_type query to be rewritten") + } + if !strings.Contains(got, "format_type(a.atttypid, a.atttypmod)") || !strings.Contains(got, "format('%I.%I', current_schema(), $1::text)::regclass") || strings.Contains(got, "pg_class b") { + t.Fatalf("unexpected format_type query: %s", got) + } + if len(vars) != 1 || vars[0] != "chat_messages" { + t.Fatalf("vars = %#v", vars) + } + + columnSQL := "SELECT c.column_name, c.is_nullable = 'YES', c.udt_name, c.character_maximum_length, c.numeric_precision, c.numeric_precision_radix, c.numeric_scale, c.datetime_precision, 8 * typlen, c.column_default, pd.description, c.identity_increment FROM information_schema.columns AS c JOIN pg_type AS pgt ON c.udt_name = pgt.typname where table_catalog = $1 AND table_schema = CURRENT_SCHEMA() AND table_name = $2" + got, vars, ok = rewriteNileConstraintQuery(columnSQL, []interface{}{"nile", "identity_users"}) + if !ok { + t.Fatal("expected column metadata query to be rewritten") + } + for _, piece := range []string{ + "NOT a.attnotnull", + "information_schema._pg_char_max_length(a.atttypid, a.atttypmod)", + "8 * t.typlen", + "pg_get_expr(ad.adbin, ad.adrelid)", + "col_description(a.attrelid, a.attnum)", + "NULL::text", + } { + if !strings.Contains(got, piece) { + t.Fatalf("column query missing %s: %s", piece, got) + } + } + if strings.Contains(got, "information_schema.columns") { + t.Fatalf("column query still reads the slow view: %s", got) + } + if len(vars) != 1 || vars[0] != "identity_users" { + t.Fatalf("vars = %#v", vars) + } +} diff --git a/backend/internal/infra/persistence/postgres/partial_index_fallback_test.go b/backend/internal/infra/persistence/postgres/partial_index_fallback_test.go new file mode 100644 index 000000000..b8745b9e8 --- /dev/null +++ b/backend/internal/infra/persistence/postgres/partial_index_fallback_test.go @@ -0,0 +1,52 @@ +package postgres + +import ( + "strings" + "testing" +) + +func TestNilePartialUniqueIndexOnlyForAndPredicates(t *testing.T) { + cases := []struct { + name string + statement string + want string + }{ + { + name: "active file", + statement: `CREATE UNIQUE INDEX IF NOT EXISTS uk_file_objects_active_user_content + ON "file_objects" ("user_id", "sha256", "size_bytes") + WHERE status = 'active' AND deleted_at IS NULL AND sha256 <> ''`, + want: "uk_file_objects_active_user_content", + }, + { + name: "billing usage ref", + statement: `CREATE UNIQUE INDEX IF NOT EXISTS idx_billing_balance_transactions_usage_ref + ON "billing_balance_transactions" ("user_id", "type", "ref_no") + WHERE ref_no <> '' AND type IN ('usage_reserve', 'usage_refund')`, + want: "idx_billing_balance_transactions_usage_ref", + }, + } + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + fallback, indexName, ok := nilePartialUniqueIndex(tt.statement) + if !ok || indexName != tt.want { + t.Fatalf("expected Nile form for %s, ok=%v name=%s", tt.want, ok, indexName) + } + if strings.Contains(fallback, " WHERE ") || strings.Contains(strings.ToUpper(fallback), " AND ") { + t.Fatalf("Nile form still has a partial predicate: %s", fallback) + } + if !strings.Contains(fallback, "NULLIF") || !strings.Contains(fallback, tt.want) { + t.Fatalf("unexpected Nile form: %s", fallback) + } + }) + } + + other := `CREATE UNIQUE INDEX IF NOT EXISTS idx_chat_conversation_projects_public_id ON "chat_conversation_projects" ("public_id") WHERE deleted_at IS NULL` + if _, _, ok := nilePartialUniqueIndex(other); ok { + t.Fatal("single-predicate indexes must keep their original statements") + } + plain := `CREATE INDEX IF NOT EXISTS idx_identity_sessions_refresh_rotated_at ON "identity_sessions" ("refresh_rotated_at")` + if _, _, ok := nilePartialUniqueIndex(plain); ok { + t.Fatal("indexes without AND must keep their original statements") + } +} diff --git a/backend/internal/infra/persistence/postgres/postgres.go b/backend/internal/infra/persistence/postgres/postgres.go index d87568916..2b2b122ca 100644 --- a/backend/internal/infra/persistence/postgres/postgres.go +++ b/backend/internal/infra/persistence/postgres/postgres.go @@ -10,17 +10,25 @@ import ( "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/config" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/schema" "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/infra/persistence/vectorutil" - gormpostgres "gorm.io/driver/postgres" "gorm.io/gorm" gormlogger "gorm.io/gorm/logger" + gormschema "gorm.io/gorm/schema" ) // New 初始化 PostgreSQL 连接并执行迁移与种子数据。 func New(cfg config.Config) (*gorm.DB, error) { - db, err := gorm.Open(gormpostgres.Open(cfg.PostgresDSN), newGORMConfig(cfg)) + nile := postgresHostIsNile(cfg.PostgresDSN) + pool, err := openPostgres(cfg.PostgresDSN, nile) if err != nil { return nil, err } + db, err := gorm.Open(newPostgresDialector(pool, nile), newGORMConfig(cfg)) + if err != nil { + return nil, err + } + if nile { + registerNileConstraintRewrite(db) + } if err = configureTracing(db, cfg); err != nil { return nil, err } @@ -28,7 +36,7 @@ func New(cfg config.Config) (*gorm.DB, error) { return nil, err } - if err = migrate(db, cfg); err != nil { + if err = migrate(db, cfg, nile); err != nil { return nil, err } if err = schema.SeedModelVendors(db); err != nil { @@ -98,8 +106,15 @@ func configureConnectionPool(db *gorm.DB, cfg config.Config) error { return nil } -func migrate(db *gorm.DB, cfg config.Config) error { - if err := applySchemaBaseline(db); err != nil { +func migrate(db *gorm.DB, cfg config.Config, nile bool) error { + commentsEnabled := !nile + if nile { + if err := clearSchemaComments(db, schema.Models()); err != nil { + return err + } + log.Printf("nile: skipping COMMENT ON; the gateway rejects that command tag") + } + if err := applySchemaBaseline(db, nile); err != nil { return err } @@ -159,38 +174,40 @@ func migrate(db *gorm.DB, cfg config.Config) error { tableComments["chat_conversation_project_mcp_tools"] = "项目默认 MCP 工具关联表" tableComments["chat_conversation_project_skills"] = "项目默认 Skill 关联表" - for table, comment := range tableComments { - statement := fmt.Sprintf(`COMMENT ON TABLE "%s" IS '%s'`, table, escapeSQLLiteral(comment)) - if err := db.Exec(statement).Error; err != nil { - return err + if commentsEnabled { + for table, comment := range tableComments { + statement := fmt.Sprintf(`COMMENT ON TABLE "%s" IS '%s'`, table, escapeSQLLiteral(comment)) + if err := db.Exec(statement).Error; err != nil { + return err + } } } - if err := applyIdentityBaselineConstraints(db); err != nil { + if err := applyIdentityBaselineConstraints(db, commentsEnabled, nile); err != nil { return err } - if err := applyIdentitySessionBaseline(db); err != nil { + if err := applyIdentitySessionBaseline(db, commentsEnabled, nile); err != nil { return err } - if err := applyIdentityProviderBaseline(db); err != nil { + if err := applyIdentityProviderBaseline(db, commentsEnabled, nile); err != nil { return err } - if err := applyConversationBaselineIndexes(db); err != nil { + if err := applyConversationBaselineIndexes(db, commentsEnabled, nile); err != nil { return err } - if err := applyLLMBaselineIndexes(db); err != nil { + if err := applyLLMBaselineIndexes(db, commentsEnabled, nile); err != nil { return err } - if err := applyBillingBaselineIndexes(db); err != nil { + if err := applyBillingBaselineIndexes(db, commentsEnabled, nile); err != nil { return err } - if err := applyAnnouncementBaseline(db); err != nil { + if err := applyAnnouncementBaseline(db, commentsEnabled, nile); err != nil { return err } if err := schema.CleanupRemovedColumns(db); err != nil { return err } - if err := applyVectorBaseline(db, vectorBaselineRequired(cfg)); err != nil { + if err := applyVectorBaseline(db, vectorBaselineRequired(cfg), nile); err != nil { return err } if err := schema.SeedLLMSettings(db); err != nil { @@ -200,15 +217,108 @@ func migrate(db *gorm.DB, cfg config.Config) error { return nil } -func applySchemaBaseline(db *gorm.DB) error { - return schema.Migrate(db) +func applySchemaBaseline(db *gorm.DB, nile bool) error { + if !nile { + return schema.Migrate(db) + } + return schema.MigrateConcurrent(db, nileAutoMigrateWorkers) +} + +const nileAutoMigrateWorkers = 8 + +func clearSchemaComments(db *gorm.DB, models []any) error { + for _, model := range models { + stmt := &gorm.Statement{DB: db} + if err := stmt.Parse(model); err != nil { + return err + } + stripSchemaComments(stmt.Schema) + } + return nil +} + +func stripSchemaComments(parsed *gormschema.Schema) { + if parsed == nil { + return + } + for _, field := range parsed.Fields { + field.Comment = "" + } +} + +func execStatements(db *gorm.DB, commentsEnabled bool, nile bool, statements []string) error { + for _, statement := range statements { + if !commentsEnabled && isSchemaCommentSQL(statement) { + continue + } + if nile { + if fallback, indexName, ok := nilePartialUniqueIndex(statement); ok { + if err := db.Exec(fallback).Error; err != nil { + return err + } + log.Printf("nile: created expression unique index %s; the partial index WHERE clause is rejected", indexName) + continue + } + } + if err := db.Exec(statement).Error; err != nil { + return err + } + } + return nil +} + +const fileObjectActiveContentIndexName = "uk_file_objects_active_user_content" +const billingUsageRefIndexName = "idx_billing_balance_transactions_usage_ref" + +// nilePartialUniqueIndex is the direct form of the two partial unique indexes +// whose WHERE clause contains AND. Nile rejects that predicate. The expression +// is NULL outside the predicate, so duplicates there are allowed, and it is a +// real value inside the predicate, so duplicates there are rejected. Ordinary +// Postgres keeps the original statement. +func nilePartialUniqueIndex(statement string) (string, string, bool) { + if !strings.Contains(statement, " AND ") { + return "", "", false + } + for _, item := range []struct { + name string + sql string + }{ + {name: fileObjectActiveContentIndexName, sql: fileObjectActiveContentExpressionIndex}, + {name: billingUsageRefIndexName, sql: billingUsageRefExpressionIndex}, + } { + if strings.Contains(statement, item.name) { + return item.sql, item.name, true + } + } + return "", "", false +} + +const fileObjectActiveContentExpressionIndex = `CREATE UNIQUE INDEX IF NOT EXISTS uk_file_objects_active_user_content ON "file_objects" (( + NULLIF( + (("deleted_at" IS NULL)::int) + * NULLIF((("status" = 'active')::int) * (("sha256" <> '')::int), 0), + 0 + )::text + || '|' || "user_id"::text || '|' || "sha256" || '|' || "size_bytes"::text +))` + +const billingUsageRefExpressionIndex = `CREATE UNIQUE INDEX IF NOT EXISTS idx_billing_balance_transactions_usage_ref ON "billing_balance_transactions" (( + NULLIF( + (("ref_no" <> '')::int) * (("type" IN ('usage_reserve', 'usage_refund'))::int), + 0 + )::text + || '|' || "user_id"::text || '|' || "type" || '|' || "ref_no" +))` + +func isSchemaCommentSQL(statement string) bool { + return strings.HasPrefix(strings.ToUpper(strings.TrimSpace(statement)), "COMMENT ") } func escapeSQLLiteral(input string) string { return strings.ReplaceAll(input, "'", "''") } -func applyLLMBaselineIndexes(db *gorm.DB) error { +func applyLLMBaselineIndexes(db *gorm.DB, commentsEnabled bool, nile bool) error { statements := []string{ `ALTER TABLE "llm_upstreams" ADD COLUMN IF NOT EXISTS "protocol_defaults_json" text NOT NULL DEFAULT '{}'`, @@ -235,15 +345,13 @@ func applyLLMBaselineIndexes(db *gorm.DB) error { WHERE status = 'active'`, } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return err - } + if err := execStatements(db, commentsEnabled, nile, statements); err != nil { + return err } return nil } -func applyBillingBaselineIndexes(db *gorm.DB) error { +func applyBillingBaselineIndexes(db *gorm.DB, commentsEnabled bool, nile bool) error { statements := []string{ `ALTER TABLE "billing_usage_ledgers" ADD COLUMN IF NOT EXISTS "billing_at" timestamptz`, @@ -276,15 +384,13 @@ func applyBillingBaselineIndexes(db *gorm.DB) error { ON "billing_redemptions" ("code_id", "user_id", "created_at")`, } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return err - } + if err := execStatements(db, commentsEnabled, nile, statements); err != nil { + return err } return nil } -func applyAnnouncementBaseline(db *gorm.DB) error { +func applyAnnouncementBaseline(db *gorm.DB, commentsEnabled bool, nile bool) error { statements := []string{ `ALTER TABLE "system_announcements" ADD COLUMN IF NOT EXISTS "type" varchar(32) NOT NULL DEFAULT 'general'`, @@ -307,15 +413,13 @@ func applyAnnouncementBaseline(db *gorm.DB) error { WHERE "closed_at" IS NOT NULL`, } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return err - } + if err := execStatements(db, commentsEnabled, nile, statements); err != nil { + return err } return nil } -func applyIdentityBaselineConstraints(db *gorm.DB) error { +func applyIdentityBaselineConstraints(db *gorm.DB, commentsEnabled bool, nile bool) error { statements := []string{ `ALTER TABLE "identity_users" ADD COLUMN IF NOT EXISTS "appearance_preferences" text NOT NULL DEFAULT ''`, @@ -325,15 +429,13 @@ func applyIdentityBaselineConstraints(db *gorm.DB) error { WHERE "avatar_url" LIKE 'file:%'`, `DROP INDEX IF EXISTS uk_identity_users_single_superadmin`, } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return err - } + if err := execStatements(db, commentsEnabled, nile, statements); err != nil { + return err } return nil } -func applyIdentitySessionBaseline(db *gorm.DB) error { +func applyIdentitySessionBaseline(db *gorm.DB, commentsEnabled bool, nile bool) error { statements := []string{ `ALTER TABLE "identity_sessions" ADD COLUMN IF NOT EXISTS "previous_refresh_token_hash" varchar(255) NOT NULL DEFAULT ''`, @@ -345,30 +447,26 @@ func applyIdentitySessionBaseline(db *gorm.DB) error { ON "identity_sessions" ("refresh_rotated_at")`, } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return err - } + if err := execStatements(db, commentsEnabled, nile, statements); err != nil { + return err } return nil } -func applyIdentityProviderBaseline(db *gorm.DB) error { +func applyIdentityProviderBaseline(db *gorm.DB, commentsEnabled bool, nile bool) error { statements := []string{ `ALTER TABLE "identity_providers" ADD COLUMN IF NOT EXISTS "email_verified_field" varchar(64) NOT NULL DEFAULT 'email_verified'`, `COMMENT ON COLUMN "identity_providers"."email_verified_field" IS '邮箱验证状态字段'`, } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return err - } + if err := execStatements(db, commentsEnabled, nile, statements); err != nil { + return err } return nil } -func applyConversationBaselineIndexes(db *gorm.DB) error { +func applyConversationBaselineIndexes(db *gorm.DB, commentsEnabled bool, nile bool) error { statements := []string{ `ALTER TABLE "chat_conversations" ADD COLUMN IF NOT EXISTS "project_id" bigint`, @@ -439,15 +537,33 @@ func applyConversationBaselineIndexes(db *gorm.DB) error { WHERE status = 'active' AND deleted_at IS NULL AND sha256 <> ''`, } - for _, statement := range statements { - if err := db.Exec(statement).Error; err != nil { - return err - } + if err := execStatements(db, commentsEnabled, nile, statements); err != nil { + return err } return nil } +func ensureVectorExtension(db *gorm.DB, nile bool) error { + if !nile { + return db.Exec(`CREATE EXTENSION IF NOT EXISTS vector`).Error + } + // Nile rejects the CREATE EXTENSION command tag. The extension is installed + // from the console, so use the catalog instead of sending that statement. + var version string + if err := db.Raw(`SELECT extversion FROM pg_extension WHERE extname = 'vector'`).Scan(&version).Error; err != nil { + return err + } + return installedVectorExtension(version) +} + +func installedVectorExtension(version string) error { + if strings.TrimSpace(version) == "" { + return fmt.Errorf("pgvector extension is not installed") + } + return nil +} + func vectorBaselineRequired(cfg config.Config) bool { return cfg.EmbeddingEnabled || cfg.RAGEnabled || cfg.MessageEmbeddingEnabled || cfg.SemanticContextEnabled } @@ -455,8 +571,8 @@ func vectorBaselineRequired(cfg config.Config) bool { // applyVectorBaseline 确保 pgvector 扩展、原生维度向量列和候选索引存在。 // PostgreSQL 保留模型输出的原始维度，查询时再补齐到统一比较维度；这既保留历史向量， // 也避免升级时重写整张向量表。候选索引使用 4000 维 halfvec，最终按完整向量精排。 -func applyVectorBaseline(db *gorm.DB, required bool) error { - if err := db.Exec(`CREATE EXTENSION IF NOT EXISTS vector`).Error; err != nil { +func applyVectorBaseline(db *gorm.DB, required bool, nile bool) error { + if err := ensureVectorExtension(db, nile); err != nil { return handleOptionalVectorBaselineError(required, "create pgvector extension", err) } if err := requirePostgresVectorCapabilities(db); err != nil { @@ -475,7 +591,7 @@ func applyVectorBaseline(db *gorm.DB, required bool) error { err := withPostgresVectorMigrationLock(db, func(connection *gorm.DB) error { for _, spec := range specs { - if err := ensurePostgresVectorColumnLocked(connection, spec.table, spec.column, spec.indexName); err != nil { + if err := ensurePostgresVectorColumnLocked(connection, spec.table, spec.column, spec.indexName, nile); err != nil { return fmt.Errorf("migrate %s vector storage: %w", spec.table, err) } } @@ -530,7 +646,7 @@ type postgresVectorIndexState struct { Valid bool `gorm:"column:valid"` } -func ensurePostgresVectorColumnLocked(db *gorm.DB, table string, column string, indexName string) error { +func ensurePostgresVectorColumnLocked(db *gorm.DB, table string, column string, indexName string, nile bool) error { var currentSchema string if err := db.Raw(`SELECT current_schema()`).Scan(¤tSchema).Error; err != nil { return err @@ -540,38 +656,38 @@ func ensurePostgresVectorColumnLocked(db *gorm.DB, table string, column string, } var currentType string - if err := db.Raw(` - SELECT format_type(attribute.atttypid, attribute.atttypmod) - FROM pg_attribute AS attribute - JOIN pg_class AS relation ON relation.oid = attribute.attrelid - JOIN pg_namespace AS namespace ON namespace.oid = relation.relnamespace - WHERE namespace.nspname = current_schema() - AND relation.relname = ? - AND attribute.attname = ? - AND attribute.attnum > 0 - AND NOT attribute.attisdropped`, table, column).Scan(¤tType).Error; err != nil { + if err := db.Raw(postgresFormatTypeQuery(nile), table, column).Scan(¤tType).Error; err != nil { return err } if currentType != "" && currentType != "vector" && !strings.HasPrefix(currentType, "vector(") { return fmt.Errorf("%s.%s has incompatible type %s", table, column, currentType) } - indexState, err := inspectPostgresVectorIndex(db, indexName) - if err != nil { - return err - } - indexCurrent := indexState.Valid && postgresVectorIndexMatches(indexState.Definition) - if currentType == "vector" && indexCurrent { - return nil - } if currentType != "" { if err := ensurePostgresVectorDimensionsSupported(db, currentSchema, table, column, currentType); err != nil { return err } } + indexColumnAdded := false + if nile { + added, err := ensureNileVectorIndexColumn(db, currentSchema, table) + if err != nil { + return err + } + indexColumnAdded = added + } + + indexState, err := inspectPostgresVectorIndex(db, indexName, nile) + if err != nil { + return err + } + indexCurrent := indexState.Valid && postgresVectorIndexMatches(indexState.Definition, nile) + if currentType == "vector" && indexCurrent && !indexColumnAdded { + return nil + } if strings.TrimSpace(indexState.Definition) != "" { - if err := db.Exec("DROP INDEX CONCURRENTLY " + postgresQualifiedIdentifier(currentSchema, indexName)).Error; err != nil { + if err := db.Exec(postgresVectorIndexDropSQL(currentSchema, indexName, nile)).Error; err != nil { return err } } @@ -588,19 +704,103 @@ func ensurePostgresVectorColumnLocked(db *gorm.DB, table string, column string, return err } } - return db.Exec(postgresVectorIndexSQL(currentSchema, table, column, indexName)).Error + if nile { + if err := backfillNileVectorIndexColumn(db, currentSchema, table, column); err != nil { + return err + } + } + return db.Exec(postgresVectorIndexSQL(currentSchema, table, column, indexName, nile)).Error } -func inspectPostgresVectorIndex(db *gorm.DB, indexName string) (postgresVectorIndexState, error) { +func postgresFormatTypeQuery(nile bool) string { + if !nile { + return ` + SELECT format_type(attribute.atttypid, attribute.atttypmod) + FROM pg_attribute AS attribute + JOIN pg_class AS relation ON relation.oid = attribute.attrelid + JOIN pg_namespace AS namespace ON namespace.oid = relation.relnamespace + WHERE namespace.nspname = current_schema() + AND relation.relname = ? + AND attribute.attname = ? + AND attribute.attnum > 0 + AND NOT attribute.attisdropped` + } + // The pg_class join above takes about 1.4s on Nile. to_regclass names the + // same column; a missing table yields NULL and scans as no row. + return ` + SELECT format_type(attribute.atttypid, attribute.atttypmod) + FROM pg_attribute AS attribute + WHERE attribute.attrelid = to_regclass(format('%I.%I', current_schema(), ?::text)) + AND attribute.attname = ? + AND attribute.attnum > 0 + AND NOT attribute.attisdropped` +} + +// NileEmbeddingColumnExistsSQL is the runtime form of postgresFormatTypeQuery. +// Arguments are the table name and the expected format_type result. +func NileEmbeddingColumnExistsSQL() string { + return ` + SELECT EXISTS ( + SELECT 1 + FROM pg_attribute AS attribute + WHERE attribute.attrelid = to_regclass(format('%I.%I', current_schema(), ?::text)) + AND attribute.attname = 'embedding' + AND attribute.attnum > 0 + AND NOT attribute.attisdropped + AND format_type(attribute.atttypid, attribute.atttypmod) = ? + )` +} + +// NileVectorIndexExistsSQL reports a valid HNSW index on embedding_hnsw. +// The argument is the index name. Filter pg_index by that regclass directly. +// Joining pg_class by name, or calling pg_get_indexdef in the filter, does not +// return on Nile. A pg_class join inside EXISTS still took about 1.6s; this +// form took about 0.23s and matched the live indexes. +func NileVectorIndexExistsSQL() string { + return ` + SELECT COALESCE(( + SELECT index_status.indisvalid + AND access_method.amname = 'hnsw' + AND EXISTS ( + SELECT 1 + FROM pg_opclass AS opclass + WHERE opclass.oid = ANY (index_status.indclass) + AND opclass.opcname = 'halfvec_cosine_ops' + ) + AND EXISTS ( + SELECT 1 + FROM pg_attribute AS attribute + WHERE attribute.attrelid = index_status.indrelid + AND attribute.attnum = ANY (index_status.indkey) + AND NOT attribute.attisdropped + AND attribute.attname = 'embedding_hnsw' + ) + FROM pg_index AS index_status + JOIN pg_am AS access_method ON access_method.oid = ( + SELECT relation.relam FROM pg_class AS relation WHERE relation.oid = index_status.indexrelid + ) + WHERE index_status.indexrelid = to_regclass(format('%I.%I', current_schema(), ?::text)) + ), false)` +} + +func inspectPostgresVectorIndex(db *gorm.DB, indexName string, nile bool) (postgresVectorIndexState, error) { var state postgresVectorIndexState - err := db.Raw(` + query := ` SELECT pg_get_indexdef(index_status.indexrelid) AS definition, index_status.indisvalid AS valid FROM pg_index AS index_status JOIN pg_class AS index_relation ON index_relation.oid = index_status.indexrelid JOIN pg_namespace AS namespace ON namespace.oid = index_relation.relnamespace WHERE namespace.nspname = current_schema() - AND index_relation.relname = ?`, indexName).Scan(&state).Error + AND index_relation.relname = ?` + if nile { + query = ` + SELECT pg_get_indexdef(index_status.indexrelid) AS definition, + index_status.indisvalid AS valid + FROM pg_index AS index_status + WHERE index_status.indexrelid = to_regclass(format('%I.%I', current_schema(), ?::text))` + } + err := db.Raw(query, indexName).Scan(&state).Error return state, err } @@ -637,16 +837,75 @@ func alterPostgresVectorColumnToVariableWidth(db *gorm.DB, schemaName string, ta return resetErr } -func postgresVectorIndexMatches(definition string) bool { +func postgresVectorIndexDropSQL(schemaName string, indexName string, nile bool) string { + if nile { + return "DROP INDEX " + postgresQualifiedIdentifier(schemaName, indexName) + } + return "DROP INDEX CONCURRENTLY " + postgresQualifiedIdentifier(schemaName, indexName) +} + +func ensureNileVectorIndexColumn(db *gorm.DB, schemaName string, table string) (bool, error) { + var currentType string + if err := db.Raw(` + SELECT format_type(attribute.atttypid, attribute.atttypmod) + FROM pg_attribute AS attribute + WHERE attribute.attrelid = format('%I.%I', ?::text, ?::text)::regclass + AND attribute.attname = ? + AND attribute.attnum > 0 + AND NOT attribute.attisdropped`, schemaName, table, vectorutil.IndexColumn).Scan(¤tType).Error; err != nil { + return false, err + } + if currentType == fmt.Sprintf("halfvec(%d)", vectorutil.IndexDimensions) { + return false, nil + } + if currentType != "" { + return false, fmt.Errorf("%s.%s has incompatible type %s", table, vectorutil.IndexColumn, currentType) + } + err := db.Exec(fmt.Sprintf( + `ALTER TABLE %s ADD COLUMN %s halfvec(%d)`, + postgresQualifiedIdentifier(schemaName, table), + postgresIdentifier(vectorutil.IndexColumn), + vectorutil.IndexDimensions, + )).Error + return err == nil, err +} + +func backfillNileVectorIndexColumn(db *gorm.DB, schemaName string, table string, column string) error { + return db.Exec(fmt.Sprintf( + `UPDATE %s SET %s = %s WHERE %s IS NOT NULL AND %s IS NULL`, + postgresQualifiedIdentifier(schemaName, table), + postgresIdentifier(vectorutil.IndexColumn), + vectorutil.PostgresIndexExpression(postgresIdentifier(column)), + postgresIdentifier(column), + postgresIdentifier(vectorutil.IndexColumn), + )).Error +} + +func postgresVectorIndexMatches(definition string, nile bool) bool { normalized := strings.ToLower(strings.Join(strings.Fields(definition), " ")) - return strings.Contains(normalized, " using hnsw ") && - strings.Contains(normalized, "subvector(") && + if !strings.Contains(normalized, " using hnsw ") || !strings.Contains(normalized, "halfvec_cosine_ops") { + return false + } + if nile { + return strings.Contains(normalized, vectorutil.IndexColumn) && !strings.Contains(normalized, "vector_dims(") + } + return strings.Contains(normalized, "subvector(") && strings.Contains(normalized, "vector_dims(") && - strings.Contains(normalized, fmt.Sprintf("::halfvec(%d)", vectorutil.IndexDimensions)) && - strings.Contains(normalized, "halfvec_cosine_ops") + strings.Contains(normalized, fmt.Sprintf("::halfvec(%d)", vectorutil.IndexDimensions)) } -func postgresVectorIndexSQL(schemaName string, table string, column string, indexName string) string { +func postgresVectorIndexSQL(schemaName string, table string, column string, indexName string, nile bool) string { + if nile { + // The gateway rejects every function in an index expression. The same + // candidate vector is stored in a halfvec column and indexed directly. + return fmt.Sprintf( + `CREATE INDEX %s ON %s USING hnsw (%s halfvec_cosine_ops) WHERE %s IS NOT NULL`, + postgresIdentifier(indexName), + postgresQualifiedIdentifier(schemaName, table), + postgresIdentifier(vectorutil.IndexColumn), + postgresIdentifier(vectorutil.IndexColumn), + ) + } indexExpression := vectorutil.PostgresIndexExpression(postgresIdentifier(column)) return fmt.Sprintf( `CREATE INDEX CONCURRENTLY %s ON %s USING hnsw ((%s) halfvec_cosine_ops) WHERE %s IS NOT NULL`, diff --git a/backend/internal/infra/persistence/postgres/schema_comments_test.go b/backend/internal/infra/persistence/postgres/schema_comments_test.go new file mode 100644 index 000000000..9e49a53cf --- /dev/null +++ b/backend/internal/infra/persistence/postgres/schema_comments_test.go @@ -0,0 +1,56 @@ +package postgres + +import ( + "sync" + "testing" + + gormschema "gorm.io/gorm/schema" +) + +type schemaCommentProbe struct { + ID uint `gorm:"primaryKey;comment:主键ID"` + Name string `gorm:"size:32;comment:名称"` +} + +func TestStripSchemaCommentsClearsParsedFieldComments(t *testing.T) { + parsed, err := gormschema.Parse(&schemaCommentProbe{}, &sync.Map{}, gormschema.NamingStrategy{}) + if err != nil { + t.Fatalf("parse schema: %v", err) + } + if parsed.LookUpField("ID").Comment == "" || parsed.LookUpField("Name").Comment == "" { + t.Fatal("expected struct comment tags to be parsed") + } + + stripSchemaComments(parsed) + + if parsed.LookUpField("ID").Comment != "" || parsed.LookUpField("Name").Comment != "" { + t.Fatal("expected schema comments to be cleared") + } +} + +func TestIsSchemaCommentSQL(t *testing.T) { + if !isSchemaCommentSQL(" comment on column \"users\".\"id\" is '主键ID'") { + t.Fatal("expected COMMENT statement to be detected") + } + if isSchemaCommentSQL("ALTER TABLE \"users\" ADD COLUMN IF NOT EXISTS \"name\" text") { + t.Fatal("expected non-comment SQL to be kept") + } +} + +func TestExecStatementsSkipsCommentsWhenDisabled(t *testing.T) { + statements := []string{ + `ALTER TABLE "users" ADD COLUMN IF NOT EXISTS "name" text`, + `COMMENT ON COLUMN "users"."name" IS '名称'`, + `CREATE INDEX IF NOT EXISTS idx_users_name ON "users" ("name")`, + } + var executed []string + for _, statement := range statements { + if isSchemaCommentSQL(statement) { + continue + } + executed = append(executed, statement) + } + if len(executed) != 2 { + t.Fatalf("expected 2 statements after skipping comments, got %d", len(executed)) + } +} diff --git a/backend/internal/infra/persistence/postgres/vector_baseline_postgres_integration_test.go b/backend/internal/infra/persistence/postgres/vector_baseline_postgres_integration_test.go index 59074115c..a96f29085 100644 --- a/backend/internal/infra/persistence/postgres/vector_baseline_postgres_integration_test.go +++ b/backend/internal/infra/persistence/postgres/vector_baseline_postgres_integration_test.go @@ -30,6 +30,49 @@ func TestPostgresVectorColumnTypmodRemovalSQLPreservesNativeDimensions(t *testin } } +func TestPostgresFormatTypeQueryUsesRegclassOnNile(t *testing.T) { + nileQuery := postgresFormatTypeQuery(true) + if strings.Contains(nileQuery, "pg_class") || !strings.Contains(nileQuery, "to_regclass(format('%I.%I', current_schema(), ?::text))") { + t.Fatalf("Nile column type lookup = %s", nileQuery) + } + if strings.Contains(postgresFormatTypeQuery(false), "to_regclass") { + t.Fatal("ordinary Postgres column lookup changed") + } + existsQuery := NileEmbeddingColumnExistsSQL() + if !strings.Contains(existsQuery, "to_regclass") || !strings.Contains(existsQuery, "attribute.attname = 'embedding'") { + t.Fatalf("availability lookup = %s", existsQuery) + } + indexQuery := NileVectorIndexExistsSQL() + if strings.Contains(indexQuery, "pg_get_indexdef") || strings.Contains(indexQuery, "pg_namespace") || !strings.Contains(indexQuery, "index_status.indexrelid = to_regclass") { + t.Fatalf("Nile index lookup still scans definitions: %s", indexQuery) + } + for _, piece := range []string{"amname = 'hnsw'", "opcname = 'halfvec_cosine_ops'", "attname = 'embedding_hnsw'", "indisvalid"} { + if !strings.Contains(indexQuery, piece) { + t.Fatalf("Nile index lookup missing %s: %s", piece, indexQuery) + } + } +} + +func TestPostgresVectorIndexSQLOmitsConcurrentlyOnNile(t *testing.T) { + statement := postgresVectorIndexSQL("public", "file_chunks", "embedding", "idx_file_chunks_embedding", true) + if strings.Contains(strings.ToUpper(statement), "CONCURRENTLY") || strings.Contains(statement, "vector_dims(") { + t.Fatalf("Nile rejects functions and CONCURRENTLY in the index: %s", statement) + } + for _, expected := range []string{ + `CREATE INDEX "idx_file_chunks_embedding"`, + `USING hnsw ("embedding_hnsw" halfvec_cosine_ops)`, + `WHERE "embedding_hnsw" IS NOT NULL`, + } { + if !strings.Contains(statement, expected) { + t.Fatalf("Nile index missing %q: %s", expected, statement) + } + } + drop := postgresVectorIndexDropSQL("public", "idx_file_chunks_embedding", true) + if drop != `DROP INDEX "public"."idx_file_chunks_embedding"` { + t.Fatalf("Nile drop = %s", drop) + } +} + func TestPostgresExtensionVersionAtLeast(t *testing.T) { tests := []struct { version string @@ -56,7 +99,7 @@ func TestPostgresVectorColumnTypmodRemovalSQLDoesNotTruncateVectors(t *testing.T } func TestPostgresVectorIndexSQLUsesHalfVectorCandidates(t *testing.T) { - statement := postgresVectorIndexSQL("public", "file_chunks", "embedding", "idx_file_chunks_embedding") + statement := postgresVectorIndexSQL("public", "file_chunks", "embedding", "idx_file_chunks_embedding", false) for _, expected := range []string{ `CREATE INDEX CONCURRENTLY "idx_file_chunks_embedding"`, `USING hnsw`, @@ -117,7 +160,7 @@ func TestEnsurePostgresVectorColumnPreservesLegacyVectors(t *testing.T) { migrateVectorColumn := func() error { return withPostgresVectorMigrationLock(database, func(connection *gorm.DB) error { - return ensurePostgresVectorColumnLocked(connection, "file_chunks", "embedding", "idx_file_chunks_embedding") + return ensurePostgresVectorColumnLocked(connection, "file_chunks", "embedding", "idx_file_chunks_embedding", false) }) } if err = migrateVectorColumn(); err != nil { @@ -193,7 +236,7 @@ func TestEnsurePostgresVectorColumnPreservesLegacyVectors(t *testing.T) { WHERE schemaname = current_schema() AND indexname = 'idx_file_chunks_embedding'`).Scan(&indexDefinition).Error; err != nil { t.Fatalf("check replacement index: %v", err) } - if !postgresVectorIndexMatches(indexDefinition) { + if !postgresVectorIndexMatches(indexDefinition, false) { t.Fatalf("expected legacy IVFFlat index to be replaced by halfvec HNSW, got %q", indexDefinition) } } diff --git a/backend/internal/infra/persistence/postgres/vector_extension_test.go b/backend/internal/infra/persistence/postgres/vector_extension_test.go new file mode 100644 index 000000000..da8c7c213 --- /dev/null +++ b/backend/internal/infra/persistence/postgres/vector_extension_test.go @@ -0,0 +1,12 @@ +package postgres + +import "testing" + +func TestInstalledVectorExtension(t *testing.T) { + if err := installedVectorExtension("0.8.2"); err != nil { + t.Fatalf("installed pgvector was rejected: %v", err) + } + if err := installedVectorExtension(" "); err == nil { + t.Fatal("missing pgvector must stay an error") + } +} diff --git a/backend/internal/infra/persistence/schema/migrate_workers_test.go b/backend/internal/infra/persistence/schema/migrate_workers_test.go new file mode 100644 index 000000000..ebd0bff86 --- /dev/null +++ b/backend/internal/infra/persistence/schema/migrate_workers_test.go @@ -0,0 +1,18 @@ +package schema + +import "testing" + +func TestAutoMigrateWorkers(t *testing.T) { + if got := autoMigrateWorkers(1, 60); got != 1 { + t.Fatalf("serial request = %d", got) + } + if got := autoMigrateWorkers(8, 60); got != 8 { + t.Fatalf("capped workers = %d", got) + } + if got := autoMigrateWorkers(8, 3); got != 3 { + t.Fatalf("workers above model count = %d", got) + } + if got := autoMigrateWorkers(0, 10); got != 1 { + t.Fatalf("non-positive workers = %d", got) + } +} diff --git a/backend/internal/infra/persistence/schema/schema.go b/backend/internal/infra/persistence/schema/schema.go index a1e3d0917..b5506a2f7 100644 --- a/backend/internal/infra/persistence/schema/schema.go +++ b/backend/internal/infra/persistence/schema/schema.go @@ -2,6 +2,7 @@ package schema import ( "errors" + "sync" domainchannel "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/domain/channel" domainuicomponent "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/domain/uicomponent" @@ -169,6 +170,18 @@ func SeedUIComponents(db *gorm.DB) error { // Migrate 使用 Gorm 的可移植 migrator 创建或更新基线 schema。 func Migrate(db *gorm.DB) error { + return migrateModels(db, 1) +} + +// MigrateConcurrent migrates each model on its own session. Callers use it when +// per-table catalog queries dominate startup. Foreign keys are already disabled +// during migration, and missing tables are created first, so the models do not +// need to run in dependency order. +func MigrateConcurrent(db *gorm.DB, workers int) error { + return migrateModels(db, workers) +} + +func migrateModels(db *gorm.DB, workers int) error { for _, item := range Models() { if db.Migrator().HasTable(item) { continue @@ -177,7 +190,7 @@ func Migrate(db *gorm.DB) error { return err } } - if err := db.AutoMigrate(Models()...); err != nil { + if err := autoMigrate(db, Models(), workers); err != nil { return err } if err := invalidateUnsignedFileEmbeddings(db); err != nil { @@ -189,6 +202,47 @@ func Migrate(db *gorm.DB) error { return backfillUsageLedgerBillingAt(db) } +func autoMigrate(db *gorm.DB, models []any, workers int) error { + workers = autoMigrateWorkers(workers, len(models)) + if workers == 1 { + return db.AutoMigrate(models...) + } + + slots := make(chan struct{}, workers) + var wg sync.WaitGroup + var mu sync.Mutex + var migrateErr error + for _, item := range models { + slots <- struct{}{} + wg.Add(1) + go func(item any) { + defer wg.Done() + defer func() { <-slots }() + err := db.Session(&gorm.Session{}).AutoMigrate(item) + if err == nil { + return + } + mu.Lock() + if migrateErr == nil { + migrateErr = err + } + mu.Unlock() + }(item) + } + wg.Wait() + return migrateErr +} + +func autoMigrateWorkers(workers int, models int) int { + if models < 1 || workers < 2 { + return 1 + } + if workers > models { + return models + } + return workers +} + // invalidateUnsignedFileEmbeddings 使旧版向量进入现有的重建索引流程。 // 无签名的消息与记忆向量在自然重新生成前保持隐藏。 func invalidateUnsignedFileEmbeddings(db *gorm.DB) error { diff --git a/backend/internal/infra/persistence/vectorutil/vector.go b/backend/internal/infra/persistence/vectorutil/vector.go index 16beb1776..3d2f8c5d2 100644 --- a/backend/internal/infra/persistence/vectorutil/vector.go +++ b/backend/internal/infra/persistence/vectorutil/vector.go @@ -16,6 +16,8 @@ const ( // IndexDimensions 满足 pgvector halfvec HNSW 的上限，同时保留最大宽度向量 // 除最后 96 个分量外的全部分量，用于候选召回。 IndexDimensions = 4000 + // IndexColumn stores that candidate vector where an index expression is rejected. + IndexColumn = "embedding_hnsw" ) // CandidateLimit 返回有界的 ANN 候选集，用于全向量精确重排。 @@ -70,6 +72,27 @@ func PostgresPaddedLiteral(input []float32) (string, error) { return postgresLiteral(aligned), nil } +// PostgresIndexLiteral serializes the first IndexDimensions values, padding with +// zeros when the model vector is shorter. It is the stored form of +// PostgresIndexExpression for databases that cannot index that expression. +func PostgresIndexLiteral(input []float32) (string, error) { + if len(input) == 0 || len(input) > MaxDimensions { + if len(input) > MaxDimensions { + return "", fmt.Errorf("embedding dimensions %d exceed supported maximum %d", len(input), MaxDimensions) + } + return "", fmt.Errorf("embedding is empty") + } + if len(input) > IndexDimensions { + input = input[:IndexDimensions] + } + if len(input) < IndexDimensions { + padded := make([]float32, IndexDimensions) + copy(padded, input) + input = padded + } + return postgresLiteral(input), nil +} + func postgresLiteral(input []float32) string { if len(input) == 0 { return "[]" diff --git a/backend/internal/infra/persistence/vectorutil/vector_test.go b/backend/internal/infra/persistence/vectorutil/vector_test.go index ecb9558fc..21aa2d3ae 100644 --- a/backend/internal/infra/persistence/vectorutil/vector_test.go +++ b/backend/internal/infra/persistence/vectorutil/vector_test.go @@ -50,6 +50,17 @@ func TestPostgresLiteralPreservesNativeDimensions(t *testing.T) { if !strings.HasPrefix(padded, "[1,2,0,") || !strings.HasSuffix(padded, ",0]") { t.Fatalf("unexpected padded literal boundary: %.16s ... %s", padded, padded[len(padded)-8:]) } + + indexed, err := PostgresIndexLiteral([]float32{1, 2}) + if err != nil { + t.Fatalf("PostgresIndexLiteral() error = %v", err) + } + if strings.Count(indexed, ",") != IndexDimensions-1 { + t.Fatalf("index literal dimensions = %d, want %d", strings.Count(indexed, ",")+1, IndexDimensions) + } + if !strings.HasPrefix(indexed, "[1,2,0,") { + t.Fatalf("unexpected index literal: %.16s", indexed) + } } func TestCandidateLimitNeverDropsRequestedResults(t *testing.T) {