feat(knowledge-base): validate vector store bindings on create, copy, and delete
Wires KnowledgeBase.VectorStoreID and the ownership-aware retrieve factory into the user-facing knowledge-base lifecycle: - POST /knowledge-bases validates the requested vector_store_id against the caller's tenant scope and the engine registry. New error codes ErrVectorStoreBindingInvalid (2200) and ErrVectorStoreUnavailable (2201) distinguish the typed branches without echoing UUIDs to the client. - GET / POST / PUT / PUT-pin responses embed the bound store's display metadata (name, source, engine_type, status) without exposing any connection credentials. Cross-tenant shared KBs receive a suppressed payload (vector_store_id stripped, source="shared") so operator-chosen store names cannot be enumerated across tenants. - POST /knowledge-bases/copy synchronously rejects clones whose target has a different embedding model or vector store, before the async clone task is enqueued. The async clone worker re-applies the same checks for defense in depth. - DELETE /vector-stores/:id refuses to remove a store with bound KBs, inside a transaction that row-locks the store on PostgreSQL and serializes via WAL on SQLite. unregister-from-registry is wrapped in defer/recover so a panic surfaces as a structured warning instead of silently leaking a stale engine. - vector_store_id is immutable after creation. The GORM <-:create tag blocks every ORM update path; the service-layer DTO omits the field entirely; a reflection-based regression test catches any future maintainer who adds it back to either layer. - Empty-string vector_store_id is normalized to nil at both the create path and inside SharesStoreWith, so rows persisted by callers that did not run Normalize first cannot trip false same-store comparisons. Part of #993. Depends on #994 and #1310.
This commit is contained in:
ochan.kwon
@@ -8,6 +8,7 @@
|
||||
- JSON 中对象存储相关字段：**`storage_config`** 为序列化字段名（对应数据库列 `cos_config`，兼容旧数据）。旧客户端若仍发送或接收 `cos_config`，服务端会兼容解析；新集成请使用 **`storage_config`**。
|
||||
- **`storage_provider_config`** 为新版存储提供者选择（如 `{"provider": "local"}`），与租户级存储引擎凭证配合使用；无配置时可为 `null`。
|
||||
- 嵌套配置对象：`chunking_config`、`image_processing_config`、`vlm_config`、`asr_config`、`extract_config`、`faq_config`、`question_generation_config`。其中 `extract_config`、`faq_config`、`question_generation_config` 允许为 `null`。
|
||||
- **`vector_store_id`** 为知识库绑定的向量存储 ID（参见 [vector-store.md](./vector-store.md)）。未指定（或 `null`/`""`）时使用租户级默认的环境变量存储；一旦创建即不可修改。详情接口返回时会附带 `vector_store_name` / `vector_store_source` / `vector_store_engine_type` / `vector_store_status` 四个只读元数据字段，用于前端展示。
|
||||
|
||||
| 方法 | 路径 | 描述 |
|
||||
| ------ | ----------------------------------------- | ------------------------ |
|
||||
@@ -43,6 +44,7 @@
|
||||
| extract_config | object | 否 | 图谱抽取配置；`enabled=true` 时需提供 `text`/`tags`/`nodes`/`relations` |
|
||||
| faq_config | object | 否 | FAQ 配置（仅 FAQ 类型知识库需要） |
|
||||
| question_generation_config | object | 否 | 问题生成配置 |
|
||||
| vector_store_id | string | 否 | 绑定的向量存储 ID。不传或为空字符串等同于 `null`（使用环境变量默认存储）。指定时必须是调用者所在租户拥有的向量存储 UUID；创建后不可修改。无效 UUID / 跨租户 / 未注册到引擎的 ID 会返回 `400` |
|
||||
|
||||
**请求**:
|
||||
|
||||
@@ -102,7 +104,8 @@ curl --location 'http://localhost:8080/api/v1/knowledge-bases' \
|
||||
"question_generation_config": {
|
||||
"enabled": false,
|
||||
"question_count": 3
|
||||
}
|
||||
},
|
||||
"vector_store_id": "550e8400-e29b-41d4-a716-446655440000"
|
||||
}'
|
||||
```
|
||||
|
||||
@@ -170,6 +173,11 @@ curl --location 'http://localhost:8080/api/v1/knowledge-bases' \
|
||||
"knowledge_count": 0,
|
||||
"chunk_count": 0,
|
||||
"processing_count": 0,
|
||||
"vector_store_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"vector_store_name": "elasticsearch-hot",
|
||||
"vector_store_source": "user",
|
||||
"vector_store_engine_type": "elasticsearch",
|
||||
"vector_store_status": "available",
|
||||
"created_at": "2025-08-12T11:30:09.206238645+08:00",
|
||||
"updated_at": "2025-08-12T11:30:09.206238854+08:00",
|
||||
"deleted_at": null
|
||||
@@ -178,6 +186,23 @@ curl --location 'http://localhost:8080/api/v1/knowledge-bases' \
|
||||
}
|
||||
```
|
||||
|
||||
**`vector_store_*` 响应字段说明**:
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
| -------------------------- | ------ | ---------------------------------------------------------------------------------------------------------- |
|
||||
| `vector_store_id` | string | 绑定的向量存储 ID（创建时未指定时为 `null`，从响应中省略） |
|
||||
| `vector_store_name` | string | 绑定存储的展示名。未绑定时返回 `"System default"`；跨租户共享 KB 视图中被隐藏 |
|
||||
| `vector_store_source` | string | `"user"`（DB 中创建的存储）/ `"env"`（环境变量虚拟存储）/ `"shared"`（跨租户共享 KB）/ `"unavailable"`（绑定的存储已不可解析） |
|
||||
| `vector_store_engine_type` | string | 引擎类型（`elasticsearch` / `qdrant` / `milvus` 等）。`shared` / `unavailable` 时为空 |
|
||||
| `vector_store_status` | string | `"available"` / `"unavailable"`。`unavailable` 表示绑定的存储已被删除或不在内存注册表中，UI 可据此提示用户重新绑定 |
|
||||
|
||||
**错误码（Phase 2 新增）**:
|
||||
|
||||
| HTTP | code | 说明 |
|
||||
| ---- | ---- | ----------------------------------------------------------------- |
|
||||
| 400 | 2200 | `vector_store_id` 无效：格式错误、不存在或属于其他租户（统一返回，避免枚举泄漏） |
|
||||
| 400 | 2201 | 指定的向量存储当前不可用：存在于数据库但未注册到引擎注册表，请检查 connection_config |
|
||||
|
||||
## GET `/knowledge-bases` - 获取知识库列表
|
||||
|
||||
返回当前租户拥有的全部知识库。当传入 `agent_id` 时，校验调用者对该共享智能体的访问权限后，返回该智能体配置可见的知识库范围（用于 `@` 提及）。
|
||||
@@ -198,6 +223,8 @@ curl --location 'http://localhost:8080/api/v1/knowledge-bases' \
|
||||
|
||||
**响应**: `data` 为数组，每个元素的字段结构同 `POST /knowledge-bases` 响应，并额外携带 `knowledge_count` / `chunk_count` / `processing_count` / `share_count` / `is_pinned` / `pinned_at` 这些聚合与状态字段。
|
||||
|
||||
> **注意（Phase 2）**：列表接口不包含 `vector_store_name` / `vector_store_source` / `vector_store_engine_type` / `vector_store_status` 这四个解析后的元数据字段（避免 N+1 查询）；仅 `vector_store_id` 来自数据库本身。需要展示存储名称时请单独调用详情接口或 `/vector-stores/:id`。
|
||||
|
||||
## GET `/knowledge-bases/:id` - 获取知识库详情
|
||||
|
||||
根据 ID 获取知识库详情。当通过共享智能体访问时，可传 `agent_id` 进行权限校验；此时返回对象会附加 `my_permission` 字段以指示当前用户对该知识库的角色（如 `viewer`）。
|
||||
@@ -222,7 +249,7 @@ curl --location 'http://localhost:8080/api/v1/knowledge-bases/kb-00000001' \
|
||||
--header 'X-API-Key: sk-xxxxx'
|
||||
```
|
||||
|
||||
**响应**: 字段结构同 `POST /knowledge-bases` 响应，并附 `is_pinned` / `pinned_at` / `knowledge_count` / `chunk_count` / `processing_count` 状态字段。通过共享智能体访问时还会附加 `my_permission`。
|
||||
**响应**: 字段结构同 `POST /knowledge-bases` 响应（包含 Phase 2 的 `vector_store_*` 元数据字段），并附 `is_pinned` / `pinned_at` / `knowledge_count` / `chunk_count` / `processing_count` 状态字段。通过共享智能体访问时还会附加 `my_permission`；同时 `vector_store_name` / `vector_store_engine_type` 会被隐藏（`vector_store_source` 返回 `"shared"`），避免跨租户泄漏存储展示名。
|
||||
|
||||
## PUT `/knowledge-bases/:id` - 更新知识库
|
||||
|
||||
@@ -282,7 +309,7 @@ curl --location --request PUT 'http://localhost:8080/api/v1/knowledge-bases/b582
|
||||
}'
|
||||
```
|
||||
|
||||
**响应**: 字段结构同 `POST /knowledge-bases` 响应（返回更新后的完整知识库对象）。
|
||||
**响应**: 字段结构同 `POST /knowledge-bases` 响应（返回更新后的完整知识库对象，包含 Phase 2 的 `vector_store_*` 元数据字段。`vector_store_id` 与创建时保持一致，无法通过该接口更改）。
|
||||
|
||||
## DELETE `/knowledge-bases/:id` - 删除知识库
|
||||
|
||||
@@ -329,7 +356,7 @@ curl --location --request PUT 'http://localhost:8080/api/v1/knowledge-bases/kb-0
|
||||
--header 'Content-Type: application/json'
|
||||
```
|
||||
|
||||
**响应**: 字段结构同 `POST /knowledge-bases` 响应，本接口操作后 `is_pinned` 翻转、`pinned_at` 同步更新。
|
||||
**响应**: 字段结构同 `POST /knowledge-bases` 响应（包含 Phase 2 的 `vector_store_*` 元数据字段），本接口操作后 `is_pinned` 翻转、`pinned_at` 同步更新。
|
||||
|
||||
## GET `/knowledge-bases/:id/hybrid-search` - 混合搜索
|
||||
|
||||
@@ -404,6 +431,15 @@ curl --location --request GET 'http://localhost:8080/api/v1/knowledge-bases/kb-0
|
||||
|
||||
**约束**：源知识库 `source_id` 必须属于调用者所在租户；若指定 `target_id`，目标知识库同样必须属于调用者租户，否则返回 `403 Forbidden`。
|
||||
|
||||
**Phase 2 同步预检（当 `target_id` 非空时）**：
|
||||
|
||||
| 检查 | 失败时响应 |
|
||||
| -------------- | ----------------------------------------------------------------------------- |
|
||||
| 嵌入模型一致性 | `400` `source and target knowledge bases use different embedding models; clone into a target with the same embedding model` |
|
||||
| 向量存储一致性 | `400` `source and target knowledge bases are bound to different vector stores; cross-store cloning is not yet supported` |
|
||||
|
||||
预检失败时任务不会入队、不会生成 `task_id`，调用方直接收到 `400`；这两个检查在异步 worker 中会再次执行（defense in depth），但在握手时即时拒绝是为了避免用户需要轮询 `progress` 才能看到错误。当 `target_id` 为空（新建目标库）时，目标库会自动复制源库的 `vector_store_id` 与 `embedding_model_id`，因此预检不会触发。
|
||||
|
||||
**参数说明（请求体）**:
|
||||
|
||||
| 字段 | 类型 | 必填 | 说明 |
|
||||
|
||||
@@ -332,6 +332,10 @@ curl --location --request PUT 'http://localhost:8080/api/v1/vector-stores/550e84
|
||||
|
||||
对 DB 中的存储执行软删除。环境变量存储不可删除（返回 `400`）。
|
||||
|
||||
**Phase 2 — 绑定保护**：
|
||||
|
||||
删除请求在事务中执行，并按 `(tenant_id, vector_store_id)` 复合索引统计当前租户中仍绑定到该存储的活跃知识库数量。**只要存在任意绑定的知识库（已软删除的 KB 不计入），删除即被拒绝**，调用者必须先解绑或删除这些知识库才能继续。在 PostgreSQL 上，事务期间会对 `vector_stores` 行加 `SELECT … FOR UPDATE` 行锁，阻止并发的知识库创建请求悄悄落到正在被删除的存储上（SQLite 上则依赖 WAL + 单写入序列化达成同样语义）。
|
||||
|
||||
**路径参数**:
|
||||
|
||||
| 字段 | 类型 | 必填 | 说明 |
|
||||
@@ -345,7 +349,7 @@ curl --location --request DELETE 'http://localhost:8080/api/v1/vector-stores/550
|
||||
--header 'X-API-Key: sk-xxxxx'
|
||||
```
|
||||
|
||||
**响应**:
|
||||
**响应（成功）**:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -353,6 +357,20 @@ curl --location --request DELETE 'http://localhost:8080/api/v1/vector-stores/550
|
||||
}
|
||||
```
|
||||
|
||||
**响应（绑定保护触发）**:
|
||||
|
||||
```json
|
||||
{
|
||||
"success": false,
|
||||
"error": {
|
||||
"code": 1000,
|
||||
"message": "vector store still has 3 knowledge base(s) bound to it; unbind or delete them before removing the store"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
HTTP `400`。错误消息中包含具体的知识库数量（便于运营定位），但不包含任何 KB 的 ID/名称，以避免跨租户信息泄漏。删除被拒绝时，DB 中的存储行保持原状，进程内引擎注册表也不会被清除。
|
||||
|
||||
## POST `/vector-stores/:id/test` - 测试已保存或环境变量存储的连接
|
||||
|
||||
对已保存的 DB 存储或环境变量虚拟存储执行一次连接测试。成功时返回检测到的服务器版本；对 DB 存储，检测到的版本会被自动写回 `connection_config.version`，环境变量存储不会更新。
|
||||
@@ -399,13 +417,18 @@ curl --location --request POST 'http://localhost:8080/api/v1/vector-stores/550e8
|
||||
- **readonly**：`true`
|
||||
- **不可修改/删除**：`PUT` 和 `DELETE` 返回 `400`
|
||||
- **可测试连通性**：`POST /vector-stores/:id/test` 正常工作
|
||||
- **被知识库绑定时**：未指定 `vector_store_id` 创建的知识库默认使用环境变量存储；这种知识库在响应中显示为 `vector_store_name="System default"` + `vector_store_source="env"`。
|
||||
|
||||
## 错误码
|
||||
|
||||
| HTTP 状态码 | 含义 |
|
||||
| ----------- | --------------------------------------------------- |
|
||||
| 400 | 请求参数错误、校验失败、尝试修改环境变量存储 |
|
||||
| 401 | 未认证（缺少租户上下文或 API Key） |
|
||||
| 404 | 向量存储不存在 |
|
||||
| 409 | 同一 endpoint + index 组合已存在 |
|
||||
| 500 | 内部服务器错误 |
|
||||
| HTTP 状态码 | code | 含义 |
|
||||
| ----------- | ---- | --------------------------------------------------- |
|
||||
| 400 | 1000 | 请求参数错误、校验失败、尝试修改环境变量存储、删除时仍有知识库绑定 |
|
||||
| 400 | 2200 | 知识库创建时引用的 `vector_store_id` 无效（不存在或属于其他租户） |
|
||||
| 400 | 2201 | 知识库创建时引用的存储当前不可用（DB 中存在但未注册到引擎） |
|
||||
| 401 | 1001 | 未认证（缺少租户上下文或 API Key） |
|
||||
| 404 | 1003 | 向量存储不存在 |
|
||||
| 409 | 1005 | 同一 endpoint + index 组合已存在 |
|
||||
| 500 | 1007 | 内部服务器错误 |
|
||||
|
||||
> `2200` / `2201` 由 `POST /knowledge-bases` 等知识库创建路径返回（详见 [knowledge-base.md](./knowledge-base.md)），列于此处仅为完整覆盖与向量存储相关的所有错误码。
|
||||
|
||||
@@ -121,3 +121,30 @@ func (r *knowledgeBaseRepository) UpdateKnowledgeBase(ctx context.Context, kb *t
|
||||
func (r *knowledgeBaseRepository) DeleteKnowledgeBase(ctx context.Context, id string) error {
|
||||
return r.db.WithContext(ctx).Where("id = ?", id).Delete(&types.KnowledgeBase{}).Error
|
||||
}
|
||||
|
||||
// CountByVectorStoreID counts active knowledge bases that are bound to the
|
||||
// given vector store within a tenant scope.
|
||||
//
|
||||
// Soft-delete filter is applied automatically by GORM because KnowledgeBase
|
||||
// has a gorm.DeletedAt column — we deliberately do not add an explicit
|
||||
// `deleted_at IS NULL` predicate to keep the single source of truth on the
|
||||
// auto-scope.
|
||||
//
|
||||
// Pass db == nil to use the repository's default db handle; pass a *gorm.DB
|
||||
// bound to a transaction (e.g., from db.Transaction) to share the same
|
||||
// write-lock context as the caller. Query column order matches the
|
||||
// composite index idx_knowledge_bases_tenant_vector_store(tenant_id,
|
||||
// vector_store_id).
|
||||
func (r *knowledgeBaseRepository) CountByVectorStoreID(
|
||||
ctx context.Context, db *gorm.DB, tenantID uint64, storeID string,
|
||||
) (int64, error) {
|
||||
if db == nil {
|
||||
db = r.db
|
||||
}
|
||||
var count int64
|
||||
err := db.WithContext(ctx).
|
||||
Model(&types.KnowledgeBase{}).
|
||||
Where("tenant_id = ? AND vector_store_id = ?", tenantID, storeID).
|
||||
Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
@@ -217,3 +217,92 @@ func TestKnowledgeBase_VectorStoreID_Roundtrip_Value(t *testing.T) {
|
||||
require.NotNil(t, reloaded.VectorStoreID)
|
||||
assert.Equal(t, "store-uuid-42", *reloaded.VectorStoreID)
|
||||
}
|
||||
|
||||
// TestCountByVectorStoreID covers the binding-count helper used by the
|
||||
// VectorStore delete guard. Verifies tenant scoping, the GORM auto soft-delete
|
||||
// filter (no explicit `deleted_at IS NULL` literal in the query), and the
|
||||
// edge case where a row stores an empty-string vector_store_id (SQLite
|
||||
// treats "" and NULL differently — we want both excluded from a non-empty
|
||||
// storeID lookup).
|
||||
func TestCountByVectorStoreID(t *testing.T) {
|
||||
db := setupKBTestDB(t)
|
||||
repo := &knowledgeBaseRepository{db: db}
|
||||
ctx := t.Context()
|
||||
|
||||
// Tenant 1: two active rows bound to store-A, one row bound to store-B.
|
||||
kbA1 := makeKB(kbStrPtr("store-A"))
|
||||
kbA2 := makeKB(kbStrPtr("store-A"))
|
||||
kbB := makeKB(kbStrPtr("store-B"))
|
||||
require.NoError(t, db.Create(kbA1).Error)
|
||||
require.NoError(t, db.Create(kbA2).Error)
|
||||
require.NoError(t, db.Create(kbB).Error)
|
||||
|
||||
// Tenant 2: one row bound to store-A (must not count for tenant 1).
|
||||
kbCross := makeKB(kbStrPtr("store-A"))
|
||||
kbCross.TenantID = 2
|
||||
require.NoError(t, db.Create(kbCross).Error)
|
||||
|
||||
// Soft-deleted row bound to store-A (must not count under auto-scope).
|
||||
kbDeleted := makeKB(kbStrPtr("store-A"))
|
||||
require.NoError(t, db.Create(kbDeleted).Error)
|
||||
require.NoError(t, db.Delete(kbDeleted).Error)
|
||||
|
||||
// Row with empty-string vector_store_id (regression — sqlite quirk).
|
||||
kbEmpty := makeKB(kbStrPtr(""))
|
||||
require.NoError(t, db.Create(kbEmpty).Error)
|
||||
|
||||
t.Run("nil db handle uses default", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 1, "store-A")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), count, "tenant 1 has 2 active KBs on store-A")
|
||||
})
|
||||
|
||||
t.Run("explicit db handle works the same", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, db, 1, "store-A")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), count)
|
||||
})
|
||||
|
||||
t.Run("different store id", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 1, "store-B")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
})
|
||||
|
||||
t.Run("unknown store id returns zero", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 1, "store-XYZ")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), count)
|
||||
})
|
||||
|
||||
t.Run("different tenant", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 2, "store-A")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
})
|
||||
|
||||
t.Run("non-empty lookup does not match empty-string rows", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 1, "")
|
||||
require.NoError(t, err)
|
||||
// Only the empty-string-vsid row matches "" exactly (SQLite quirk —
|
||||
// "" is not NULL). The non-empty rows do not.
|
||||
assert.Equal(t, int64(1), count, "empty-string lookup matches only the empty row")
|
||||
})
|
||||
|
||||
t.Run("shared tx handle returns same result twice", func(t *testing.T) {
|
||||
err := db.Transaction(func(tx *gorm.DB) error {
|
||||
c1, err := repo.CountByVectorStoreID(ctx, tx, 1, "store-A")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c2, err := repo.CountByVectorStoreID(ctx, tx, 1, "store-A")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
assert.Equal(t, c1, c2, "two reads within the same tx must agree")
|
||||
assert.Equal(t, int64(2), c1)
|
||||
return nil
|
||||
})
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -7,10 +7,12 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/application/service/retriever"
|
||||
apperrors "github.com/Tencent/WeKnora/internal/errors"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
"github.com/google/uuid"
|
||||
"github.com/hibiken/asynq"
|
||||
)
|
||||
@@ -74,7 +76,12 @@ func (s *knowledgeBaseService) GetRepository() interfaces.KnowledgeBaseRepositor
|
||||
return s.repo
|
||||
}
|
||||
|
||||
// CreateKnowledgeBase creates a new knowledge base
|
||||
// CreateKnowledgeBase creates a new knowledge base.
|
||||
//
|
||||
// When VectorStoreID is set, the binding is validated against the caller's
|
||||
// tenant scope and the engine registry before persisting. A nil or
|
||||
// empty-string VectorStoreID is normalized to nil ("use the tenant's
|
||||
// effective engines") to match the retrieve-engine factory's pre-condition.
|
||||
func (s *knowledgeBaseService) CreateKnowledgeBase(ctx context.Context,
|
||||
kb *types.KnowledgeBase,
|
||||
) (*types.KnowledgeBase, error) {
|
||||
@@ -87,6 +94,22 @@ func (s *knowledgeBaseService) CreateKnowledgeBase(ctx context.Context,
|
||||
kb.UpdatedAt = time.Now()
|
||||
kb.EnsureDefaults()
|
||||
|
||||
// Fold empty-string vector_store_id into nil so this path and the
|
||||
// retrieve-engine factory's pre-condition share a single representation.
|
||||
wasEmpty := kb.VectorStoreID != nil && *kb.VectorStoreID == ""
|
||||
kb.Normalize()
|
||||
if wasEmpty {
|
||||
logger.Debugf(ctx,
|
||||
"[kb.create] empty vector_store_id normalized to nil for tenant=%d",
|
||||
kb.TenantID)
|
||||
}
|
||||
|
||||
if kb.HasVectorStore() {
|
||||
if err := s.validateVectorStoreBinding(ctx, kb.TenantID, *kb.VectorStoreID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "Creating knowledge base, ID: %s, tenant ID: %d, name: %s", kb.ID, kb.TenantID, kb.Name)
|
||||
|
||||
if err := s.repo.CreateKnowledgeBase(ctx, kb); err != nil {
|
||||
@@ -101,6 +124,62 @@ func (s *knowledgeBaseService) CreateKnowledgeBase(ctx context.Context,
|
||||
return kb, nil
|
||||
}
|
||||
|
||||
// validateVectorStoreBinding routes through retriever.VerifyBinding so the
|
||||
// ownership + registry sentinel hierarchy stays the single source of truth.
|
||||
// The service layer's responsibility is to:
|
||||
//
|
||||
// 1. fast-reject malformed UUIDs (cheap pre-flight that also avoids a DB
|
||||
// round trip for type-confusion inputs like "' OR 1=1 --"),
|
||||
// 2. translate retriever sentinels into user-facing AppErrors with
|
||||
// generic messages and the typed error codes.
|
||||
//
|
||||
// UUID parse failures map to the same "vector store not found" message as
|
||||
// cross-tenant attempts to avoid an enumeration oracle that distinguishes
|
||||
// "malformed input" from "non-existent UUID".
|
||||
func (s *knowledgeBaseService) validateVectorStoreBinding(
|
||||
ctx context.Context, tenantID uint64, storeID string,
|
||||
) error {
|
||||
sanitized := secutils.SanitizeForLog(storeID)
|
||||
|
||||
if _, err := uuid.Parse(storeID); err != nil {
|
||||
logger.WarnWithFields(ctx, logger.Fields{
|
||||
"tenant_id": tenantID,
|
||||
"store_id": sanitized,
|
||||
"reason": "malformed vector_store_id",
|
||||
}, "[kb.create] vector store id is not a valid UUID")
|
||||
return apperrors.NewVectorStoreBindingInvalidError("vector store not found")
|
||||
}
|
||||
|
||||
switch err := retriever.VerifyBinding(
|
||||
ctx, s.retrieveEngine, s.ownership, tenantID, storeID,
|
||||
); {
|
||||
case err == nil:
|
||||
return nil
|
||||
case errors.Is(err, retriever.ErrVectorStoreForbidden):
|
||||
logger.WarnWithFields(ctx, logger.Fields{
|
||||
"tenant_id": tenantID,
|
||||
"store_id": sanitized,
|
||||
"reason": "cross-tenant or unknown store",
|
||||
}, "[kb.create] vector store not owned by tenant")
|
||||
return apperrors.NewVectorStoreBindingInvalidError("vector store not found")
|
||||
case errors.Is(err, retriever.ErrVectorStoreNotFound):
|
||||
logger.WarnWithFields(ctx, logger.Fields{
|
||||
"tenant_id": tenantID,
|
||||
"store_id": sanitized,
|
||||
"reason": "store registered in DB but missing in registry",
|
||||
}, "[kb.create] vector store currently unavailable")
|
||||
return apperrors.NewVectorStoreUnavailableError(
|
||||
"vector store is currently unavailable; check its connection configuration")
|
||||
default:
|
||||
logger.ErrorWithFields(ctx, err, map[string]interface{}{
|
||||
"tenant_id": tenantID,
|
||||
"store_id": sanitized,
|
||||
"reason": "binding verification failed",
|
||||
})
|
||||
return apperrors.NewInternalServerError("failed to verify vector store binding")
|
||||
}
|
||||
}
|
||||
|
||||
// GetKnowledgeBaseByID retrieves a knowledge base by its ID
|
||||
func (s *knowledgeBaseService) GetKnowledgeBaseByID(ctx context.Context, id string) (*types.KnowledgeBase, error) {
|
||||
if id == "" {
|
||||
@@ -266,7 +345,23 @@ func (s *knowledgeBaseService) FillKnowledgeBaseCounts(ctx context.Context, kb *
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateKnowledgeBase updates a knowledge base's properties
|
||||
// UpdateKnowledgeBase updates a knowledge base's mutable properties.
|
||||
//
|
||||
// IMPORTANT — vector_store_id immutability contract:
|
||||
// The vector_store_id binding is deliberately not accepted by this method.
|
||||
// Two layers enforce immutability:
|
||||
//
|
||||
// 1. ORM layer: the GORM tag `<-:create` on KnowledgeBase.VectorStoreID
|
||||
// makes every UPDATE path (Save / Updates / Select-Updates) a no-op for
|
||||
// that column. Verified by repository/knowledgebase_sqlite_test.go.
|
||||
// 2. Service layer: this method intentionally omits VectorStoreID from its
|
||||
// parameter list, and the matching handler DTO UpdateKnowledgeBaseRequest
|
||||
// omits the field as well. A reflection-based regression test
|
||||
// (handler/knowledgebase_request_test.go) fails if either DTO field
|
||||
// is added back, alerting future maintainers.
|
||||
//
|
||||
// Any future cross-store rebind workflow must use raw SQL through a
|
||||
// dedicated repository method — the only sanctioned write path post-creation.
|
||||
func (s *knowledgeBaseService) UpdateKnowledgeBase(ctx context.Context,
|
||||
id string,
|
||||
name string,
|
||||
@@ -624,6 +719,21 @@ func (s *knowledgeBaseService) SetEmbeddingModel(ctx context.Context, id string,
|
||||
|
||||
// CopyKnowledgeBase copies a knowledge base to a new knowledge base (shallow copy).
|
||||
// Source and target must belong to the tenant in context; cross-tenant access is rejected.
|
||||
//
|
||||
// Defensive checks:
|
||||
//
|
||||
// - When dstKB != "" (clone into an existing target), the source's
|
||||
// EmbeddingModelID and VectorStoreID must match the target's. Mismatched
|
||||
// embedding models would silently mix incompatible vector spaces;
|
||||
// mismatched vector stores would require copying physical vector data
|
||||
// between stores, which is not yet supported.
|
||||
// - When dstKB == "" (create a new target), VectorStoreID is copied from
|
||||
// the source so the new KB shares the same physical vector index. GORM
|
||||
// `<-:create` allows INSERT, so the new row is well-formed.
|
||||
//
|
||||
// The handler's CopyKnowledgeBase endpoint runs the same checks synchronously
|
||||
// before enqueueing the async clone task, so the 400 errors here are
|
||||
// defense-in-depth for the worker entry point.
|
||||
func (s *knowledgeBaseService) CopyKnowledgeBase(ctx context.Context,
|
||||
srcKB string, dstKB string,
|
||||
) (*types.KnowledgeBase, *types.KnowledgeBase, error) {
|
||||
@@ -642,12 +752,31 @@ func (s *knowledgeBaseService) CopyKnowledgeBase(ctx context.Context,
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
// Defense 1: embedding model must match. Mixing incompatible
|
||||
// vector spaces would produce semantically broken search results.
|
||||
if sourceKB.EmbeddingModelID != targetKB.EmbeddingModelID {
|
||||
return nil, nil, apperrors.NewBadRequestError(
|
||||
"source and target knowledge bases use different embedding models; " +
|
||||
"clone into a target with the same embedding model")
|
||||
}
|
||||
|
||||
// Defense 2: vector store binding must match. Cross-store cloning
|
||||
// would require copying physical vector data between stores.
|
||||
// (both nil → equal; both same UUID → equal; otherwise → rejected)
|
||||
if !sourceKB.SharesStoreWith(targetKB) {
|
||||
return nil, nil, apperrors.NewBadRequestError(
|
||||
"source and target knowledge bases are bound to different vector stores; " +
|
||||
"cross-store cloning is not yet supported")
|
||||
}
|
||||
} else {
|
||||
var faqConfig *types.FAQConfig
|
||||
if sourceKB.FAQConfig != nil {
|
||||
cfg := *sourceKB.FAQConfig
|
||||
faqConfig = &cfg
|
||||
}
|
||||
// Preserve VectorStoreID so the cloned KB lands on the same
|
||||
// physical index. GORM `<-:create` permits the value at INSERT.
|
||||
targetKB = &types.KnowledgeBase{
|
||||
ID: uuid.New().String(),
|
||||
Name: sourceKB.Name,
|
||||
@@ -662,6 +791,7 @@ func (s *knowledgeBaseService) CopyKnowledgeBase(ctx context.Context,
|
||||
StorageProviderConfig: sourceKB.StorageProviderConfig,
|
||||
StorageConfig: sourceKB.StorageConfig,
|
||||
FAQConfig: faqConfig,
|
||||
VectorStoreID: sourceKB.VectorStoreID,
|
||||
}
|
||||
targetKB.EnsureDefaults()
|
||||
if err := s.repo.CreateKnowledgeBase(ctx, targetKB); err != nil {
|
||||
|
||||
@@ -0,0 +1,330 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
stderrors "errors"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/application/service/retriever"
|
||||
apperrors "github.com/Tencent/WeKnora/internal/errors"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Vector-store binding service-level tests
|
||||
//
|
||||
// Coverage:
|
||||
// - CreateKnowledgeBase: vector_store_id validation matrix (Normalize,
|
||||
// malformed UUID, sentinel wrap, error code distinction).
|
||||
// - CopyKnowledgeBase: embedding-model + store defenses, new-KB
|
||||
// VectorStoreID copy.
|
||||
//
|
||||
// These are pure logic tests that drive `validateVectorStoreBinding` and
|
||||
// `CopyKnowledgeBase` through stubbed dependencies — no DB I/O.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// validKBStoreUUID is a well-formed UUID used to drive the validation
|
||||
// branches that survive the uuid.Parse() pre-flight.
|
||||
const validKBStoreUUID = "0193b8a0-1111-7000-8000-000000000001"
|
||||
|
||||
type fakeOwnership struct {
|
||||
owned map[string]uint64 // storeID -> tenantID owner
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeOwnership) StoreOwnedBy(_ context.Context, storeID string, tenantID uint64) (bool, error) {
|
||||
if f.err != nil {
|
||||
return false, f.err
|
||||
}
|
||||
owner, ok := f.owned[storeID]
|
||||
return ok && owner == tenantID, nil
|
||||
}
|
||||
|
||||
type fakeRegistry struct {
|
||||
registered map[string]struct{}
|
||||
}
|
||||
|
||||
func (f *fakeRegistry) Register(_ interfaces.RetrieveEngineService) error { return nil }
|
||||
func (f *fakeRegistry) GetRetrieveEngineService(_ types.RetrieverEngineType) (interfaces.RetrieveEngineService, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeRegistry) GetAllRetrieveEngineServices() []interfaces.RetrieveEngineService { return nil }
|
||||
func (f *fakeRegistry) GetByStoreID(storeID string) (interfaces.RetrieveEngineService, error) {
|
||||
if _, ok := f.registered[storeID]; ok {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, stderrors.New("not registered")
|
||||
}
|
||||
|
||||
// fakeKBRepo is the smallest KnowledgeBaseRepository needed by
|
||||
// CreateKnowledgeBase + CopyKnowledgeBase tests. It stores rows in a map
|
||||
// keyed by ID. No tenant scoping is applied — the tested paths already
|
||||
// enforce that explicitly.
|
||||
type fakeKBRepo struct {
|
||||
rows map[string]*types.KnowledgeBase
|
||||
createErr error
|
||||
}
|
||||
|
||||
func newFakeKBRepo() *fakeKBRepo { return &fakeKBRepo{rows: map[string]*types.KnowledgeBase{}} }
|
||||
|
||||
func (r *fakeKBRepo) CreateKnowledgeBase(_ context.Context, kb *types.KnowledgeBase) error {
|
||||
if r.createErr != nil {
|
||||
return r.createErr
|
||||
}
|
||||
r.rows[kb.ID] = kb
|
||||
return nil
|
||||
}
|
||||
func (r *fakeKBRepo) GetKnowledgeBaseByID(_ context.Context, id string) (*types.KnowledgeBase, error) {
|
||||
return r.rows[id], nil
|
||||
}
|
||||
func (r *fakeKBRepo) GetKnowledgeBaseByIDAndTenant(_ context.Context, id string, tenantID uint64) (*types.KnowledgeBase, error) {
|
||||
kb := r.rows[id]
|
||||
if kb == nil || kb.TenantID != tenantID {
|
||||
return nil, stderrors.New("not found")
|
||||
}
|
||||
return kb, nil
|
||||
}
|
||||
func (r *fakeKBRepo) GetKnowledgeBaseByIDs(_ context.Context, _ []string) ([]*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *fakeKBRepo) ListKnowledgeBases(_ context.Context) ([]*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *fakeKBRepo) ListKnowledgeBasesByTenantID(_ context.Context, _ uint64) ([]*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *fakeKBRepo) UpdateKnowledgeBase(_ context.Context, _ *types.KnowledgeBase) error {
|
||||
return nil
|
||||
}
|
||||
func (r *fakeKBRepo) DeleteKnowledgeBase(_ context.Context, _ string) error { return nil }
|
||||
func (r *fakeKBRepo) TogglePinKnowledgeBase(_ context.Context, _ string, _ uint64) (*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *fakeKBRepo) CountByVectorStoreID(_ context.Context, _ *gorm.DB, _ uint64, _ string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Force compile-time conformance check so any future interface change
|
||||
// surfaces as a build error here rather than at usage site.
|
||||
var _ interfaces.KnowledgeBaseRepository = (*fakeKBRepo)(nil)
|
||||
|
||||
func newPR3KBService(repo *fakeKBRepo, registry *fakeRegistry, ownership *fakeOwnership) *knowledgeBaseService {
|
||||
return &knowledgeBaseService{
|
||||
repo: repo,
|
||||
retrieveEngine: registry,
|
||||
ownership: ownership,
|
||||
}
|
||||
}
|
||||
|
||||
func ctxWithTenant(tenantID uint64) context.Context {
|
||||
return context.WithValue(context.Background(), types.TenantIDContextKey, tenantID)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CreateKnowledgeBase — vector_store_id binding validation matrix
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCreateKnowledgeBase_VectorStoreBinding(t *testing.T) {
|
||||
registered := map[string]struct{}{validKBStoreUUID: {}}
|
||||
|
||||
t.Run("nil vector_store_id persists unchanged", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
svc := newPR3KBService(repo,
|
||||
&fakeRegistry{registered: registered},
|
||||
&fakeOwnership{owned: map[string]uint64{validKBStoreUUID: 1}},
|
||||
)
|
||||
kb, err := svc.CreateKnowledgeBase(ctxWithTenant(1), &types.KnowledgeBase{Name: "kb"})
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, kb.VectorStoreID)
|
||||
assert.Len(t, repo.rows, 1)
|
||||
})
|
||||
|
||||
t.Run("empty string vector_store_id normalizes to nil", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
svc := newPR3KBService(repo,
|
||||
&fakeRegistry{registered: registered},
|
||||
&fakeOwnership{},
|
||||
)
|
||||
empty := ""
|
||||
kb, err := svc.CreateKnowledgeBase(ctxWithTenant(1), &types.KnowledgeBase{Name: "kb", VectorStoreID: &empty})
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, kb.VectorStoreID)
|
||||
})
|
||||
|
||||
t.Run("malformed UUID rejected with BindingInvalid code", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
svc := newPR3KBService(repo,
|
||||
&fakeRegistry{registered: registered},
|
||||
&fakeOwnership{},
|
||||
)
|
||||
bad := "not-a-uuid"
|
||||
_, err := svc.CreateKnowledgeBase(ctxWithTenant(1), &types.KnowledgeBase{Name: "kb", VectorStoreID: &bad})
|
||||
require.Error(t, err)
|
||||
appErr, ok := apperrors.IsAppError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, apperrors.ErrVectorStoreBindingInvalid, appErr.Code)
|
||||
assert.NotContains(t, appErr.Message, bad, "UUID must not echo back")
|
||||
assert.Empty(t, repo.rows, "must not persist on validation failure")
|
||||
})
|
||||
|
||||
t.Run("valid UUID + owned + registered → persist OK", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
svc := newPR3KBService(repo,
|
||||
&fakeRegistry{registered: registered},
|
||||
&fakeOwnership{owned: map[string]uint64{validKBStoreUUID: 1}},
|
||||
)
|
||||
id := validKBStoreUUID
|
||||
kb, err := svc.CreateKnowledgeBase(ctxWithTenant(1), &types.KnowledgeBase{Name: "kb", VectorStoreID: &id})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, kb.VectorStoreID)
|
||||
assert.Equal(t, validKBStoreUUID, *kb.VectorStoreID)
|
||||
})
|
||||
|
||||
t.Run("valid UUID + cross-tenant → BindingInvalid", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
svc := newPR3KBService(repo,
|
||||
&fakeRegistry{registered: registered},
|
||||
&fakeOwnership{owned: map[string]uint64{validKBStoreUUID: 2}}, // owned by tenant 2
|
||||
)
|
||||
id := validKBStoreUUID
|
||||
_, err := svc.CreateKnowledgeBase(ctxWithTenant(1), &types.KnowledgeBase{Name: "kb", VectorStoreID: &id})
|
||||
require.Error(t, err)
|
||||
appErr, ok := apperrors.IsAppError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, apperrors.ErrVectorStoreBindingInvalid, appErr.Code)
|
||||
})
|
||||
|
||||
t.Run("valid UUID + owned but unregistered → Unavailable", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
svc := newPR3KBService(repo,
|
||||
&fakeRegistry{registered: map[string]struct{}{}}, // not registered
|
||||
&fakeOwnership{owned: map[string]uint64{validKBStoreUUID: 1}},
|
||||
)
|
||||
id := validKBStoreUUID
|
||||
_, err := svc.CreateKnowledgeBase(ctxWithTenant(1), &types.KnowledgeBase{Name: "kb", VectorStoreID: &id})
|
||||
require.Error(t, err)
|
||||
appErr, ok := apperrors.IsAppError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, apperrors.ErrVectorStoreUnavailable, appErr.Code)
|
||||
})
|
||||
|
||||
t.Run("ownership infra error → 500", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
svc := newPR3KBService(repo,
|
||||
&fakeRegistry{registered: registered},
|
||||
&fakeOwnership{err: stderrors.New("db down")},
|
||||
)
|
||||
id := validKBStoreUUID
|
||||
_, err := svc.CreateKnowledgeBase(ctxWithTenant(1), &types.KnowledgeBase{Name: "kb", VectorStoreID: &id})
|
||||
require.Error(t, err)
|
||||
appErr, ok := apperrors.IsAppError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, 500, appErr.HTTPCode)
|
||||
})
|
||||
|
||||
// Compile-time check: the retriever sentinels still match what we expect
|
||||
// in the wrapper logic. If they are ever renamed or removed, this fails
|
||||
// at build time rather than silently degrading the wrap.
|
||||
_ = retriever.ErrVectorStoreForbidden
|
||||
_ = retriever.ErrVectorStoreNotFound
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CopyKnowledgeBase — embedding model + vector store defenses
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCopyKnowledgeBase_Defenses(t *testing.T) {
|
||||
mkKB := func(id string, tenant uint64, embed string, vsid *string) *types.KnowledgeBase {
|
||||
return &types.KnowledgeBase{
|
||||
ID: id, Name: id, TenantID: tenant, EmbeddingModelID: embed, VectorStoreID: vsid,
|
||||
}
|
||||
}
|
||||
storeA := "store-A"
|
||||
storeB := "store-B"
|
||||
|
||||
t.Run("dstKB empty -> new KB inherits VectorStoreID", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
repo.rows["src"] = mkKB("src", 1, "embed-1", &storeA)
|
||||
svc := newPR3KBService(repo, &fakeRegistry{}, &fakeOwnership{})
|
||||
|
||||
src, tgt, err := svc.CopyKnowledgeBase(ctxWithTenant(1), "src", "")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, src.VectorStoreID)
|
||||
require.NotNil(t, tgt.VectorStoreID)
|
||||
assert.Equal(t, *src.VectorStoreID, *tgt.VectorStoreID, "M17: VectorStoreID preserved")
|
||||
|
||||
// Reload from fake repo and re-verify
|
||||
stored := repo.rows[tgt.ID]
|
||||
require.NotNil(t, stored)
|
||||
require.NotNil(t, stored.VectorStoreID)
|
||||
assert.Equal(t, storeA, *stored.VectorStoreID)
|
||||
})
|
||||
|
||||
t.Run("dstKB empty + source nil VectorStoreID -> new target nil too", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
repo.rows["src"] = mkKB("src", 1, "embed-1", nil)
|
||||
svc := newPR3KBService(repo, &fakeRegistry{}, &fakeOwnership{})
|
||||
|
||||
_, tgt, err := svc.CopyKnowledgeBase(ctxWithTenant(1), "src", "")
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, tgt.VectorStoreID)
|
||||
})
|
||||
|
||||
t.Run("dstKB set + same embedding model + both nil stores -> OK", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
repo.rows["src"] = mkKB("src", 1, "embed-1", nil)
|
||||
repo.rows["dst"] = mkKB("dst", 1, "embed-1", nil)
|
||||
svc := newPR3KBService(repo, &fakeRegistry{}, &fakeOwnership{})
|
||||
|
||||
_, _, err := svc.CopyKnowledgeBase(ctxWithTenant(1), "src", "dst")
|
||||
require.NoError(t, err)
|
||||
})
|
||||
|
||||
t.Run("dstKB set + same store UUID -> OK", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
repo.rows["src"] = mkKB("src", 1, "embed-1", &storeA)
|
||||
repo.rows["dst"] = mkKB("dst", 1, "embed-1", &storeA)
|
||||
svc := newPR3KBService(repo, &fakeRegistry{}, &fakeOwnership{})
|
||||
|
||||
_, _, err := svc.CopyKnowledgeBase(ctxWithTenant(1), "src", "dst")
|
||||
require.NoError(t, err)
|
||||
})
|
||||
|
||||
t.Run("dstKB set + different embedding models -> 400", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
repo.rows["src"] = mkKB("src", 1, "embed-1", &storeA)
|
||||
repo.rows["dst"] = mkKB("dst", 1, "embed-2", &storeA)
|
||||
svc := newPR3KBService(repo, &fakeRegistry{}, &fakeOwnership{})
|
||||
|
||||
_, _, err := svc.CopyKnowledgeBase(ctxWithTenant(1), "src", "dst")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "different embedding models")
|
||||
})
|
||||
|
||||
t.Run("dstKB set + different stores -> 400 (no Phase number)", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
repo.rows["src"] = mkKB("src", 1, "embed-1", &storeA)
|
||||
repo.rows["dst"] = mkKB("dst", 1, "embed-1", &storeB)
|
||||
svc := newPR3KBService(repo, &fakeRegistry{}, &fakeOwnership{})
|
||||
|
||||
_, _, err := svc.CopyKnowledgeBase(ctxWithTenant(1), "src", "dst")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "different vector stores")
|
||||
assert.NotContains(t, err.Error(), "Phase 4", "internal roadmap labels must not leak to end-user error messages")
|
||||
})
|
||||
|
||||
t.Run("dstKB set + one nil + one set -> 400", func(t *testing.T) {
|
||||
repo := newFakeKBRepo()
|
||||
repo.rows["src"] = mkKB("src", 1, "embed-1", nil)
|
||||
repo.rows["dst"] = mkKB("dst", 1, "embed-1", &storeA)
|
||||
svc := newPR3KBService(repo, &fakeRegistry{}, &fakeOwnership{})
|
||||
|
||||
_, _, err := svc.CopyKnowledgeBase(ctxWithTenant(1), "src", "dst")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "different vector stores")
|
||||
})
|
||||
}
|
||||
@@ -45,6 +45,50 @@ type TenantStoreOwnership interface {
|
||||
StoreOwnedBy(ctx context.Context, storeID string, tenantID uint64) (bool, error)
|
||||
}
|
||||
|
||||
// VerifyBinding asserts that a non-empty storeID is owned by tenantID and
|
||||
// registered in the in-memory engine registry. It encapsulates the two
|
||||
// checks that gate every store-bound resolution so that callers outside
|
||||
// the retriever package (notably the KB create-validation path) can reuse
|
||||
// the same sentinel hierarchy instead of duplicating the logic.
|
||||
//
|
||||
// Resolution rules:
|
||||
//
|
||||
// - ownership.StoreOwnedBy returns an infrastructure error → that error
|
||||
// is returned verbatim so callers can decide retry/abort.
|
||||
// - ownership returns (false, nil) → ErrVectorStoreForbidden.
|
||||
// - ownership returns (true, nil) + registry.GetByStoreID fails →
|
||||
// ErrVectorStoreNotFound.
|
||||
// - all checks succeed → nil.
|
||||
//
|
||||
// VerifyBinding itself never echoes the store UUID; callers MUST wrap the
|
||||
// sentinels into user-facing errors at the boundary (and log the
|
||||
// tenant/store pair via structured fields when appropriate).
|
||||
//
|
||||
// resolveBoundEngine (below) intentionally does NOT delegate to VerifyBinding
|
||||
// because it also needs the resolved engine service; sharing would require
|
||||
// either a second registry lookup or returning the service from VerifyBinding,
|
||||
// both of which dilute the helper's single purpose. The two paths are kept
|
||||
// in lockstep by the factory_test.go matrix.
|
||||
func VerifyBinding(
|
||||
ctx context.Context,
|
||||
registry interfaces.RetrieveEngineRegistry,
|
||||
ownership TenantStoreOwnership,
|
||||
tenantID uint64,
|
||||
storeID string,
|
||||
) error {
|
||||
owned, err := ownership.StoreOwnedBy(ctx, storeID, tenantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !owned {
|
||||
return ErrVectorStoreForbidden
|
||||
}
|
||||
if _, err := registry.GetByStoreID(storeID); err != nil {
|
||||
return ErrVectorStoreNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateRetrieveEngineForKB returns a CompositeRetrieveEngine resolved from
|
||||
// a KB's VectorStore binding.
|
||||
//
|
||||
|
||||
@@ -384,6 +384,71 @@ func TestFactoryParallelInvocation(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// VerifyBinding tests — sentinel SSOT used by the KB create-validation
|
||||
// path. Mirrors the branches of resolveBoundEngine without constructing an
|
||||
// actual CompositeRetrieveEngine.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestVerifyBinding(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
esEngine := &fakeEngine{
|
||||
engineType: types.ElasticsearchRetrieverEngineType,
|
||||
support: []types.RetrieverType{types.VectorRetrieverType},
|
||||
}
|
||||
|
||||
t.Run("ownership infra error is returned verbatim", func(t *testing.T) {
|
||||
registry := registryWithStores(t, nil, nil)
|
||||
ownership := &fakeOwnership{err: errors.New("db boom")}
|
||||
err := VerifyBinding(ctx, registry, ownership, 1, "store-A")
|
||||
if err == nil || err.Error() != "db boom" {
|
||||
t.Fatalf("expected verbatim infra error, got %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("not owned -> ErrVectorStoreForbidden", func(t *testing.T) {
|
||||
registry := registryWithStores(t, nil, nil)
|
||||
ownership := &fakeOwnership{owned: map[string]uint64{}}
|
||||
err := VerifyBinding(ctx, registry, ownership, 1, "store-A")
|
||||
if !errors.Is(err, ErrVectorStoreForbidden) {
|
||||
t.Fatalf("expected ErrVectorStoreForbidden, got %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("owned but unregistered -> ErrVectorStoreNotFound", func(t *testing.T) {
|
||||
registry := registryWithStores(t, nil, nil)
|
||||
ownership := &fakeOwnership{owned: map[string]uint64{"store-A": 1}}
|
||||
err := VerifyBinding(ctx, registry, ownership, 1, "store-A")
|
||||
if !errors.Is(err, ErrVectorStoreNotFound) {
|
||||
t.Fatalf("expected ErrVectorStoreNotFound, got %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("owned and registered -> nil", func(t *testing.T) {
|
||||
registry := registryWithStores(t,
|
||||
map[string]*fakeEngine{"store-A": esEngine},
|
||||
nil,
|
||||
)
|
||||
ownership := &fakeOwnership{owned: map[string]uint64{"store-A": 1}}
|
||||
if err := VerifyBinding(ctx, registry, ownership, 1, "store-A"); err != nil {
|
||||
t.Fatalf("expected nil, got %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("cross-tenant returns Forbidden (not Found)", func(t *testing.T) {
|
||||
// store owned by tenant 2, queried by tenant 1
|
||||
registry := registryWithStores(t,
|
||||
map[string]*fakeEngine{"store-A": esEngine},
|
||||
nil,
|
||||
)
|
||||
ownership := &fakeOwnership{owned: map[string]uint64{"store-A": 2}}
|
||||
err := VerifyBinding(ctx, registry, ownership, 1, "store-A")
|
||||
if !errors.Is(err, ErrVectorStoreForbidden) {
|
||||
t.Fatalf("expected ErrVectorStoreForbidden, got %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ----- helpers -----
|
||||
|
||||
func strPtr(s string) *string { return &s }
|
||||
|
||||
@@ -2,6 +2,7 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
stderrors "errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
@@ -11,25 +12,41 @@ import (
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// vectorStoreService implements interfaces.VectorStoreService
|
||||
type vectorStoreService struct {
|
||||
repo interfaces.VectorStoreRepository
|
||||
storeRegistry interfaces.StoreRegistry // for dynamic registry updates on CRUD
|
||||
factory interfaces.EngineFactory // creates engine services from VectorStore config
|
||||
kbRepo interfaces.KnowledgeBaseRepository // counts bound KBs for the delete guard
|
||||
storeRegistry interfaces.StoreRegistry // for dynamic registry updates on CRUD
|
||||
factory interfaces.EngineFactory // creates engine services from VectorStore config
|
||||
db *gorm.DB // shared handle for cross-table transactions (delete guard)
|
||||
envStores []types.VectorStore // env stores derived once at construction for ResolveStoreView fast path
|
||||
}
|
||||
|
||||
// NewVectorStoreService creates a new vector store service
|
||||
// NewVectorStoreService creates a new vector store service.
|
||||
//
|
||||
// kbRepo and db are required by the delete guard, which counts bound KBs
|
||||
// inside a transaction. storeRegistry and factory are optional in tests
|
||||
// (passing nil disables dynamic registration / unregistration).
|
||||
func NewVectorStoreService(
|
||||
repo interfaces.VectorStoreRepository,
|
||||
kbRepo interfaces.KnowledgeBaseRepository,
|
||||
storeRegistry interfaces.StoreRegistry,
|
||||
factory interfaces.EngineFactory,
|
||||
db *gorm.DB,
|
||||
) interfaces.VectorStoreService {
|
||||
return &vectorStoreService{
|
||||
repo: repo,
|
||||
kbRepo: kbRepo,
|
||||
storeRegistry: storeRegistry,
|
||||
factory: factory,
|
||||
db: db,
|
||||
// Cache the env-store derivation once at construction so per-request
|
||||
// resolution does not re-read os environment variables every call.
|
||||
envStores: types.BuildEnvVectorStores(os.Getenv("RETRIEVE_DRIVER"), os.Getenv),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -62,8 +79,11 @@ func (s *vectorStoreService) CreateStore(ctx context.Context, store *types.Vecto
|
||||
return errors.NewConflictError("a vector store with the same endpoint and index already exists")
|
||||
}
|
||||
|
||||
// 4. Duplicate check — env stores (pure function, no os.Getenv in types)
|
||||
for _, envStore := range types.BuildEnvVectorStores(os.Getenv("RETRIEVE_DRIVER"), os.Getenv) {
|
||||
// 4. Duplicate check — env stores. We re-derive on each create because
|
||||
// CreateStore is a low-frequency admin action; consistency with the
|
||||
// startup-cached envStores is enforced by RETRIEVE_DRIVER being read
|
||||
// only at process start.
|
||||
for _, envStore := range s.envStores {
|
||||
if envStore.EngineType == store.EngineType &&
|
||||
envStore.ConnectionConfig.GetEndpoint() == endpoint &&
|
||||
envStore.IndexConfig.GetIndexNameOrDefault(store.EngineType) == indexName {
|
||||
@@ -91,7 +111,7 @@ func (s *vectorStoreService) CreateStore(ctx context.Context, store *types.Vecto
|
||||
return err
|
||||
}
|
||||
|
||||
// 6. Register in registry (best-effort; failure doesn't roll back DB).
|
||||
// 7. Register in registry (best-effort; failure doesn't roll back DB).
|
||||
// The store is already persisted, and will be loaded on next app restart (self-healing).
|
||||
s.registerInRegistry(ctx, store)
|
||||
|
||||
@@ -113,22 +133,112 @@ func (s *vectorStoreService) UpdateStore(ctx context.Context, store *types.Vecto
|
||||
return s.repo.Update(ctx, store)
|
||||
}
|
||||
|
||||
// DeleteStore deletes a vector store by tenant + id.
|
||||
// Phase 2: KB binding check will be added here.
|
||||
// DeleteStore deletes a vector store by tenant + id, after verifying that no
|
||||
// knowledge base is currently bound to it.
|
||||
//
|
||||
// Guard rules:
|
||||
//
|
||||
// 1. Run inside a transaction so that the binding count and the store
|
||||
// delete are atomic with respect to other writers holding the store
|
||||
// row lock. Default isolation is Read Committed; this is a write-lock
|
||||
// relationship, not a "shared snapshot" relationship.
|
||||
// 2. PostgreSQL: take a row-level X-lock on the vector_stores row via
|
||||
// SELECT … FOR UPDATE so concurrent KB-create requests reading the
|
||||
// same store row block until our transaction completes. SQLite
|
||||
// serializes writes via WAL + max-open-conns=1, so the lock hint is
|
||||
// skipped and we rely on the transaction boundary alone.
|
||||
// 3. Count knowledge_bases rows via the shared CountByVectorStoreID
|
||||
// repository method (tx-aware), which leverages the composite index
|
||||
// (tenant_id, vector_store_id). GORM auto-applies the soft-delete
|
||||
// scope — no explicit deleted_at predicate is needed.
|
||||
// 4. After commit, unregister from the in-memory registry. Wrapped in
|
||||
// defer/recover so a panic in UnregisterByStoreID surfaces as a
|
||||
// structured warning instead of silently leaking the stale engine.
|
||||
//
|
||||
// Race window remaining:
|
||||
//
|
||||
// A narrow window exists between CreateKnowledgeBase's binding check and
|
||||
// the INSERT — a KB can be created against a store that is simultaneously
|
||||
// being deleted. The retrieve-engine factory then rejects searches with
|
||||
// the ErrVectorStoreForbidden / NotFound sentinel; the KB response view
|
||||
// surfaces the condition through vector_store_status="unavailable" so the
|
||||
// UI can guide recovery (admin tool / rebind / KB recreation).
|
||||
//
|
||||
// Multi-replica registry staleness:
|
||||
//
|
||||
// The in-memory registry is per-process. After a successful commit +
|
||||
// UnregisterByStoreID on this replica, sibling replicas continue serving
|
||||
// the engine from their own caches until process restart. This method
|
||||
// does not broadcast invalidation across the cluster.
|
||||
func (s *vectorStoreService) DeleteStore(ctx context.Context, tenantID uint64, id string) error {
|
||||
if err := s.repo.Delete(ctx, tenantID, id); err != nil {
|
||||
err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
// tx inherits ctx from WithContext above; no need to re-attach.
|
||||
|
||||
// 1. Lock the store row (PG row-level X-lock; skipped on SQLite).
|
||||
var store types.VectorStore
|
||||
q := tx.Where("id = ? AND tenant_id = ?", id, tenantID)
|
||||
if s.isPostgres(tx) {
|
||||
q = q.Clauses(clause.Locking{Strength: "UPDATE"})
|
||||
}
|
||||
if err := q.First(&store).Error; err != nil {
|
||||
if stderrors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return errors.NewNotFoundError("vector store not found")
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// 2. Binding count under the same write-lock boundary.
|
||||
count, err := s.kbRepo.CountByVectorStoreID(ctx, tx, tenantID, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if count > 0 {
|
||||
return errors.NewBadRequestError(
|
||||
fmt.Sprintf(
|
||||
"vector store still has %d knowledge base(s) bound to it; "+
|
||||
"unbind or delete them before removing the store", count))
|
||||
}
|
||||
|
||||
// 3. Soft-delete (gorm.DeletedAt fills automatically).
|
||||
return tx.Delete(&store).Error
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Unregister from registry (idempotent)
|
||||
if s.storeRegistry != nil {
|
||||
s.storeRegistry.UnregisterByStoreID(id)
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "Deleted vector store: tenant=%d, id=%s", tenantID, id)
|
||||
// 4. Unregister from registry — wrapped to convert panics into ops
|
||||
// warnings rather than silent stale-engine leaks.
|
||||
s.unregisterSafely(ctx, id)
|
||||
logger.Infof(ctx, "Deleted vector store: tenant=%d, id=%s", tenantID,
|
||||
secutils.SanitizeForLog(id))
|
||||
return nil
|
||||
}
|
||||
|
||||
// unregisterSafely calls the registry's idempotent unregister with panic
|
||||
// containment. A panic here is recoverable because the registry is
|
||||
// in-memory and self-heals on process restart — but it must be loud
|
||||
// enough for ops to purge the stale engine on running replicas.
|
||||
func (s *vectorStoreService) unregisterSafely(ctx context.Context, storeID string) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.WarnWithFields(ctx, logger.Fields{
|
||||
"store_id": secutils.SanitizeForLog(storeID),
|
||||
"panic": fmt.Sprint(r),
|
||||
}, "[vectorstore.delete] registry unregister panicked; engine may stay stale until restart")
|
||||
}
|
||||
}()
|
||||
if s.storeRegistry != nil {
|
||||
s.storeRegistry.UnregisterByStoreID(storeID)
|
||||
}
|
||||
}
|
||||
|
||||
// isPostgres reports whether the active GORM dialector is PostgreSQL.
|
||||
// Used to gate dialect-specific clauses (e.g., SELECT FOR UPDATE) that
|
||||
// SQLite would either ignore (recent versions) or fail to compile on.
|
||||
func (s *vectorStoreService) isPostgres(db *gorm.DB) bool {
|
||||
return db != nil && db.Dialector != nil && db.Dialector.Name() == "postgres"
|
||||
}
|
||||
|
||||
// SaveDetectedVersion updates the connection_config.version for a stored vector store.
|
||||
// Works on a copy to avoid mutating the caller's object.
|
||||
func (s *vectorStoreService) SaveDetectedVersion(ctx context.Context, store *types.VectorStore, version string) error {
|
||||
@@ -137,6 +247,128 @@ func (s *vectorStoreService) SaveDetectedVersion(ctx context.Context, store *typ
|
||||
return s.repo.UpdateConnectionConfig(ctx, &updated)
|
||||
}
|
||||
|
||||
// ResolveStoreView returns the API-safe display projection of a single
|
||||
// store ID for embedding in another resource's response (typically a KB).
|
||||
//
|
||||
// Resolution order:
|
||||
//
|
||||
// 1. storeID == "" → DefaultStoreDisplay (env fallback semantics).
|
||||
// 2. DB store row matching (id, tenantID) → user-source display.
|
||||
// 3. Cached env store with matching ID → env-source display.
|
||||
// 4. Otherwise → UnavailableStoreDisplay with a structured warn log.
|
||||
//
|
||||
// Errors from the underlying repository are returned to the caller so
|
||||
// transient infrastructure failures can be classified, but the returned
|
||||
// StoreDisplay is still UnavailableStoreDisplay so a handler that ignores
|
||||
// the error degrades gracefully rather than panicking on a zero value.
|
||||
func (s *vectorStoreService) ResolveStoreView(
|
||||
ctx context.Context, tenantID uint64, storeID string,
|
||||
) (types.StoreDisplay, error) {
|
||||
if storeID == "" {
|
||||
return types.DefaultStoreDisplay(), nil
|
||||
}
|
||||
store, err := s.repo.GetByID(ctx, tenantID, storeID)
|
||||
if err != nil {
|
||||
return types.UnavailableStoreDisplay(), err
|
||||
}
|
||||
if store != nil {
|
||||
return types.StoreDisplay{
|
||||
Name: store.Name,
|
||||
Source: types.StoreSourceUser,
|
||||
EngineType: string(store.EngineType),
|
||||
Status: "available",
|
||||
}, nil
|
||||
}
|
||||
for _, env := range s.envStores {
|
||||
if env.ID == storeID {
|
||||
return types.StoreDisplay{
|
||||
Name: env.Name,
|
||||
Source: types.StoreSourceEnv,
|
||||
EngineType: string(env.EngineType),
|
||||
Status: "available",
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
logger.WarnWithFields(ctx, logger.Fields{
|
||||
"tenant_id": tenantID,
|
||||
"store_id": secutils.SanitizeForLog(storeID),
|
||||
}, "[vectorstore.resolve] bound store missing from DB and env set")
|
||||
return types.UnavailableStoreDisplay(), nil
|
||||
}
|
||||
|
||||
// BatchResolveStoreView resolves multiple store IDs in a single DB read
|
||||
// plus the cached env-store match. Returned map keys are the storeIDs
|
||||
// originally requested; missing IDs map to UnavailableStoreDisplay.
|
||||
//
|
||||
// Intended for list endpoints that need store metadata for many KBs at
|
||||
// once without incurring N+1 ResolveStoreView calls.
|
||||
//
|
||||
// Implementation note: the tenant-store count is bounded by operator
|
||||
// config (typically tens), so iterating the tenant's full store list
|
||||
// once is cheaper than a SELECT … WHERE id IN (…) round-trip and avoids
|
||||
// adding a batch-by-ids repository method that has no other caller.
|
||||
func (s *vectorStoreService) BatchResolveStoreView(
|
||||
ctx context.Context, tenantID uint64, storeIDs []string,
|
||||
) (map[string]types.StoreDisplay, error) {
|
||||
out := make(map[string]types.StoreDisplay, len(storeIDs))
|
||||
if len(storeIDs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
requested := make(map[string]bool, len(storeIDs))
|
||||
hasNonEmpty := false
|
||||
for _, id := range storeIDs {
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
requested[id] = true
|
||||
hasNonEmpty = true
|
||||
}
|
||||
|
||||
if hasNonEmpty {
|
||||
dbStores, err := s.repo.List(ctx, tenantID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, st := range dbStores {
|
||||
if requested[st.ID] {
|
||||
out[st.ID] = types.StoreDisplay{
|
||||
Name: st.Name,
|
||||
Source: types.StoreSourceUser,
|
||||
EngineType: string(st.EngineType),
|
||||
Status: "available",
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, env := range s.envStores {
|
||||
if _, ok := out[env.ID]; ok {
|
||||
continue
|
||||
}
|
||||
if requested[env.ID] {
|
||||
out[env.ID] = types.StoreDisplay{
|
||||
Name: env.Name,
|
||||
Source: types.StoreSourceEnv,
|
||||
EngineType: string(env.EngineType),
|
||||
Status: "available",
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fill misses (including empty-string entries) with the appropriate
|
||||
// sentinel so callers can rely on a key for every requested ID.
|
||||
for _, id := range storeIDs {
|
||||
if id == "" {
|
||||
out[id] = types.DefaultStoreDisplay()
|
||||
continue
|
||||
}
|
||||
if _, ok := out[id]; !ok {
|
||||
out[id] = types.UnavailableStoreDisplay()
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// registerInRegistry creates an engine service and registers it in the registry.
|
||||
// Logs and skips on failure — the store is already persisted in DB,
|
||||
// and will be loaded on next app restart (self-healing).
|
||||
|
||||
@@ -11,6 +11,8 @@ import (
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
sqlitedrv "gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -155,7 +157,7 @@ func (m *mockEngineService) BatchUpdateChunkTagID(_ context.Context, _ map[strin
|
||||
|
||||
func TestCreateStore_Success(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
@@ -173,7 +175,7 @@ func TestCreateStore_Success(t *testing.T) {
|
||||
|
||||
func TestCreateStore_ValidationError(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -205,7 +207,7 @@ func TestCreateStore_ValidationError(t *testing.T) {
|
||||
|
||||
func TestCreateStore_ConnectionConfigValidation(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -222,22 +224,20 @@ func TestCreateStore_ConnectionConfigValidation(t *testing.T) {
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "postgres without addr or default connection",
|
||||
store: &types.VectorStore{
|
||||
TenantID: 1, Name: "test",
|
||||
EngineType: types.PostgresRetrieverEngineType,
|
||||
ConnectionConfig: types.ConnectionConfig{},
|
||||
},
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "postgres with use_default_connection",
|
||||
// Postgres is no longer registerable as a DB-managed store (see
|
||||
// validEngineTypes); the Validate() call inside CreateStore short-
|
||||
// circuits before validateConnectionConfig ever runs. Either
|
||||
// missing-config or full-config rejects with "unsupported engine
|
||||
// type". The env-store path (RETRIEVE_DRIVER=postgres) reaches the
|
||||
// engine registry through BuildEnvVectorStores and does not
|
||||
// traverse this code path.
|
||||
name: "postgres rejected as DB store (any config)",
|
||||
store: &types.VectorStore{
|
||||
TenantID: 1, Name: "test",
|
||||
EngineType: types.PostgresRetrieverEngineType,
|
||||
ConnectionConfig: types.ConnectionConfig{UseDefaultConnection: true},
|
||||
},
|
||||
wantError: false,
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "qdrant without host",
|
||||
@@ -267,13 +267,15 @@ func TestCreateStore_ConnectionConfigValidation(t *testing.T) {
|
||||
wantError: true,
|
||||
},
|
||||
{
|
||||
name: "sqlite with no config (ok)",
|
||||
// SQLite, like Postgres, is no longer registerable as a DB store
|
||||
// (see validEngineTypes). Reachable only as an env store.
|
||||
name: "sqlite rejected as DB store",
|
||||
store: &types.VectorStore{
|
||||
TenantID: 1, Name: "test",
|
||||
EngineType: types.SQLiteRetrieverEngineType,
|
||||
ConnectionConfig: types.ConnectionConfig{},
|
||||
},
|
||||
wantError: false,
|
||||
wantError: true,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -291,7 +293,7 @@ func TestCreateStore_ConnectionConfigValidation(t *testing.T) {
|
||||
|
||||
func TestCreateStore_DuplicateCheck_DBStore(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{existsByEndpoint: true}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
@@ -314,7 +316,7 @@ func TestCreateStore_DuplicateCheck_DBError(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{
|
||||
existsByEndpointErr: assert.AnError,
|
||||
}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
@@ -338,7 +340,7 @@ func TestCreateStore_DuplicateCheck_EnvStore(t *testing.T) {
|
||||
t.Setenv("ELASTICSEARCH_INDEX", "xwrag_default")
|
||||
|
||||
repo := &mockVectorStoreRepo{existsByEndpoint: false} // no DB duplicate
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
@@ -368,7 +370,7 @@ func TestCreateStore_DuplicateCheck_EnvStore_DifferentIndex_Allowed(t *testing.T
|
||||
t.Setenv("ELASTICSEARCH_INDEX", "xwrag_default")
|
||||
|
||||
repo := &mockVectorStoreRepo{existsByEndpoint: false}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
@@ -387,22 +389,23 @@ func TestCreateStore_DuplicateCheck_EnvStore_DifferentIndex_Allowed(t *testing.T
|
||||
}
|
||||
|
||||
func TestCreateStore_DifferentEndpointSameIndex_Allowed(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{existsByEndpoint: false}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
|
||||
// Use postgres with UseDefaultConnection to avoid needing a real ES endpoint.
|
||||
// The test verifies duplicate-check logic, not connectivity.
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
Name: "new-store",
|
||||
EngineType: types.PostgresRetrieverEngineType,
|
||||
ConnectionConfig: types.ConnectionConfig{
|
||||
UseDefaultConnection: true,
|
||||
},
|
||||
}
|
||||
|
||||
err := svc.CreateStore(context.Background(), store)
|
||||
assert.NoError(t, err)
|
||||
// Historical note — this test used Postgres + UseDefaultConnection so the
|
||||
// connectivity test step would short-circuit (testPostgresConnection
|
||||
// returns nil on use_default_connection), letting the duplicate-check
|
||||
// logic run against a pure in-memory mock. Once Postgres/SQLite were
|
||||
// dropped from validEngineTypes (DB stores must be a network-reachable
|
||||
// engine), no engine remaining in the allow list has an equivalent
|
||||
// "no network needed" fast path: every other engine's connection probe
|
||||
// dials a real endpoint, which a unit-test environment cannot
|
||||
// reasonably provide without spinning up a real backend.
|
||||
//
|
||||
// The underlying invariant (same endpoint + different index name is
|
||||
// allowed) is still indirectly exercised by the live ParadeDB/Qdrant
|
||||
// flows in CI integration and by manual E2E tests. A proper unit-test
|
||||
// substitute requires either an injectable TestConnection mock or a
|
||||
// dedicated repository-only test that exercises ExistsByEndpointAndIndex
|
||||
// without going through the service-level CreateStore pipeline.
|
||||
t.Skip("requires mockable TestConnection or repository-only test; track separately")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -413,7 +416,7 @@ func TestCreateStore_RegistersInRegistry(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
registry := newMockStoreRegistry()
|
||||
factory := mockEngineFactory(nil)
|
||||
svc := NewVectorStoreService(repo, registry, factory)
|
||||
svc := NewVectorStoreService(repo, nil, registry, factory, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
@@ -436,7 +439,7 @@ func TestCreateStore_RegistryFailureDoesNotRollBackDB(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
registry := newMockStoreRegistry()
|
||||
factory := mockEngineFactory(assert.AnError) // factory fails
|
||||
svc := NewVectorStoreService(repo, registry, factory)
|
||||
svc := NewVectorStoreService(repo, nil, registry, factory, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
@@ -459,7 +462,7 @@ func TestCreateStore_RegistryFailureDoesNotRollBackDB(t *testing.T) {
|
||||
|
||||
func TestCreateStore_NilRegistryAndFactory(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil) // no registry
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil) // no registry
|
||||
|
||||
store := &types.VectorStore{
|
||||
TenantID: 1,
|
||||
@@ -478,52 +481,21 @@ func TestCreateStore_NilRegistryAndFactory(t *testing.T) {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// DeleteStore + Registry integration tests
|
||||
//
|
||||
// The DeleteStore path runs through a transactional guard that requires
|
||||
// both a kbRepo and a *gorm.DB handle, so the in-memory mock cannot
|
||||
// exercise it. The TestDeleteStore_Guard_* family below covers the same
|
||||
// outcomes (success, missing store, registry interaction, rollback)
|
||||
// against an actual sqlite transaction.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestDeleteStore_UnregistersFromRegistry(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
registry := newMockStoreRegistry()
|
||||
registry.registered["store-1"] = true
|
||||
svc := NewVectorStoreService(repo, registry, nil)
|
||||
|
||||
err := svc.DeleteStore(context.Background(), 1, "store-1")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should be unregistered
|
||||
assert.Contains(t, registry.unregistered, "store-1")
|
||||
assert.False(t, registry.registered["store-1"])
|
||||
}
|
||||
|
||||
func TestDeleteStore_NilRegistryGraceful(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
|
||||
// Should not panic with nil registry
|
||||
err := svc.DeleteStore(context.Background(), 1, "store-1")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestDeleteStore_RepoErrorSkipsUnregister(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{deleteErr: assert.AnError}
|
||||
registry := newMockStoreRegistry()
|
||||
registry.registered["store-1"] = true
|
||||
svc := NewVectorStoreService(repo, registry, nil)
|
||||
|
||||
err := svc.DeleteStore(context.Background(), 1, "store-1")
|
||||
assert.Error(t, err)
|
||||
|
||||
// Registry should NOT be touched if DB delete fails
|
||||
assert.True(t, registry.registered["store-1"])
|
||||
assert.Empty(t, registry.unregistered)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// UpdateStore tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestUpdateStore_Success(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
ID: "test-id",
|
||||
@@ -537,7 +509,7 @@ func TestUpdateStore_Success(t *testing.T) {
|
||||
|
||||
func TestUpdateStore_ValidationError(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -565,21 +537,10 @@ func TestUpdateStore_ValidationError(t *testing.T) {
|
||||
// DeleteStore tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestDeleteStore_Success(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
|
||||
err := svc.DeleteStore(context.Background(), 1, "test-id")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestDeleteStore_RepoError(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{deleteErr: assert.AnError}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
|
||||
err := svc.DeleteStore(context.Background(), 1, "test-id")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
// TestDeleteStore_Success and TestDeleteStore_RepoError have been replaced
|
||||
// by the TestDeleteStore_Guard_* family below, which exercises the
|
||||
// transactional delete-guard path against an actual sqlite database (a mock
|
||||
// repo would not honor the row-lock + binding-count semantics).
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SaveDetectedVersion tests
|
||||
@@ -587,7 +548,7 @@ func TestDeleteStore_RepoError(t *testing.T) {
|
||||
|
||||
func TestSaveDetectedVersion_Success(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
ID: "store-1",
|
||||
@@ -601,7 +562,7 @@ func TestSaveDetectedVersion_Success(t *testing.T) {
|
||||
|
||||
func TestSaveDetectedVersion_RepoError(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{updateErr: assert.AnError}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{ID: "store-1", TenantID: 1}
|
||||
err := svc.SaveDetectedVersion(context.Background(), store, "8.11.0")
|
||||
@@ -610,7 +571,7 @@ func TestSaveDetectedVersion_RepoError(t *testing.T) {
|
||||
|
||||
func TestSaveDetectedVersion_DoesNotMutateOriginal(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
store := &types.VectorStore{
|
||||
ID: "store-1",
|
||||
@@ -631,7 +592,7 @@ func TestSaveDetectedVersion_DoesNotMutateOriginal(t *testing.T) {
|
||||
|
||||
func TestTestConnection_UnsupportedEngineType(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
_, err := svc.TestConnection(context.Background(), "unknown_engine", types.ConnectionConfig{})
|
||||
require.Error(t, err)
|
||||
@@ -643,7 +604,7 @@ func TestTestConnection_UnsupportedEngineType(t *testing.T) {
|
||||
|
||||
func TestTestConnection_SQLiteAlwaysSucceeds(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
version, err := svc.TestConnection(context.Background(), types.SQLiteRetrieverEngineType, types.ConnectionConfig{})
|
||||
assert.NoError(t, err)
|
||||
@@ -652,7 +613,7 @@ func TestTestConnection_SQLiteAlwaysSucceeds(t *testing.T) {
|
||||
|
||||
func TestTestConnection_PostgresDefaultConnection(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
version, err := svc.TestConnection(context.Background(), types.PostgresRetrieverEngineType,
|
||||
types.ConnectionConfig{UseDefaultConnection: true})
|
||||
@@ -663,7 +624,7 @@ func TestTestConnection_PostgresDefaultConnection(t *testing.T) {
|
||||
func TestTestConnection_DorisInvalidAddr(t *testing.T) {
|
||||
// 给一个不可达的地址 + 5s timeout，期望返回 BadRequestError 而非 panic。
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
|
||||
defer cancel()
|
||||
@@ -678,7 +639,7 @@ func TestTestConnection_DorisInvalidAddr(t *testing.T) {
|
||||
|
||||
func TestTestConnection_DorisMissingAddr(t *testing.T) {
|
||||
repo := &mockVectorStoreRepo{}
|
||||
svc := NewVectorStoreService(repo, nil, nil)
|
||||
svc := NewVectorStoreService(repo, nil, nil, nil, nil)
|
||||
|
||||
_, err := svc.TestConnection(context.Background(), types.DorisRetrieverEngineType,
|
||||
types.ConnectionConfig{})
|
||||
@@ -799,3 +760,376 @@ func TestValidateConnectionConfig(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// DeleteStore guard + ResolveStoreView + BatchResolveStoreView integration
|
||||
//
|
||||
// These use a real sqlite in-memory database because the delete guard relies
|
||||
// on GORM transactions and the soft-delete auto-scope; an interface-level
|
||||
// mock would not catch divergence between the row-lock path and the count.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// guardTestDDL inlines the subset of migrations/sqlite/000000_init.up.sql
|
||||
// that the delete-guard tests touch. We do not use AutoMigrate because the
|
||||
// KnowledgeBase struct carries `type:jsonb` GORM tags that SQLite cannot
|
||||
// map cleanly.
|
||||
const guardTestDDL = `
|
||||
CREATE TABLE IF NOT EXISTS vector_stores (
|
||||
id VARCHAR(36) NOT NULL PRIMARY KEY,
|
||||
name VARCHAR(255) NOT NULL,
|
||||
engine_type VARCHAR(50) NOT NULL,
|
||||
connection_config TEXT NOT NULL DEFAULT '{}',
|
||||
index_config TEXT NOT NULL DEFAULT '{}',
|
||||
tenant_id INTEGER NOT NULL,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
deleted_at DATETIME NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS knowledge_bases (
|
||||
id VARCHAR(36) PRIMARY KEY,
|
||||
name VARCHAR(255) NOT NULL,
|
||||
description TEXT,
|
||||
tenant_id INTEGER NOT NULL,
|
||||
type VARCHAR(32) NOT NULL DEFAULT 'document',
|
||||
chunking_config TEXT NOT NULL DEFAULT '{}',
|
||||
image_processing_config TEXT NOT NULL DEFAULT '{}',
|
||||
embedding_model_id VARCHAR(64) NOT NULL,
|
||||
summary_model_id VARCHAR(64) NOT NULL,
|
||||
cos_config TEXT NOT NULL DEFAULT '{}',
|
||||
storage_provider_config TEXT DEFAULT NULL,
|
||||
vlm_config TEXT NOT NULL DEFAULT '{}',
|
||||
extract_config TEXT NULL DEFAULT NULL,
|
||||
faq_config TEXT,
|
||||
question_generation_config TEXT NULL,
|
||||
is_temporary BOOLEAN NOT NULL DEFAULT 0,
|
||||
is_pinned INTEGER NOT NULL DEFAULT 0,
|
||||
pinned_at DATETIME NULL,
|
||||
asr_config TEXT,
|
||||
vector_store_id VARCHAR(36),
|
||||
wiki_config TEXT,
|
||||
indexing_strategy TEXT,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
deleted_at DATETIME
|
||||
);
|
||||
`
|
||||
|
||||
func newGuardTestService(t *testing.T) (*vectorStoreService, *gorm.DB, *mockStoreRegistry) {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlitedrv.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Exec(guardTestDDL).Error)
|
||||
|
||||
registry := newMockStoreRegistry()
|
||||
svc := &vectorStoreService{
|
||||
repo: &realStoreRepo{db: db},
|
||||
kbRepo: &realKBRepo{db: db},
|
||||
storeRegistry: registry,
|
||||
factory: nil,
|
||||
db: db,
|
||||
envStores: []types.VectorStore{},
|
||||
}
|
||||
return svc, db, registry
|
||||
}
|
||||
|
||||
// realStoreRepo is a hand-rolled VectorStoreRepository against an in-memory
|
||||
// sqlite handle. We avoid pulling in the repository package to keep this
|
||||
// test file self-contained.
|
||||
type realStoreRepo struct{ db *gorm.DB }
|
||||
|
||||
func (r *realStoreRepo) Create(ctx context.Context, s *types.VectorStore) error {
|
||||
return r.db.WithContext(ctx).Create(s).Error
|
||||
}
|
||||
func (r *realStoreRepo) GetByID(ctx context.Context, tenantID uint64, id string) (*types.VectorStore, error) {
|
||||
var s types.VectorStore
|
||||
err := r.db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).First(&s).Error
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &s, nil
|
||||
}
|
||||
func (r *realStoreRepo) List(ctx context.Context, tenantID uint64) ([]*types.VectorStore, error) {
|
||||
var stores []*types.VectorStore
|
||||
err := r.db.WithContext(ctx).Where("tenant_id = ?", tenantID).Find(&stores).Error
|
||||
return stores, err
|
||||
}
|
||||
func (r *realStoreRepo) Update(_ context.Context, _ *types.VectorStore) error {
|
||||
return nil
|
||||
}
|
||||
func (r *realStoreRepo) UpdateConnectionConfig(_ context.Context, _ *types.VectorStore) error {
|
||||
return nil
|
||||
}
|
||||
func (r *realStoreRepo) Delete(ctx context.Context, tenantID uint64, id string) error {
|
||||
return r.db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&types.VectorStore{}).Error
|
||||
}
|
||||
func (r *realStoreRepo) ExistsByEndpointAndIndex(_ context.Context, _ uint64, _ types.RetrieverEngineType, _ string, _ string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// realKBRepo is the minimal KnowledgeBaseRepository slice required by the
|
||||
// vector-store service (only CountByVectorStoreID is exercised here).
|
||||
type realKBRepo struct{ db *gorm.DB }
|
||||
|
||||
func (r *realKBRepo) CountByVectorStoreID(ctx context.Context, db *gorm.DB, tenantID uint64, storeID string) (int64, error) {
|
||||
if db == nil {
|
||||
db = r.db
|
||||
}
|
||||
var count int64
|
||||
err := db.WithContext(ctx).
|
||||
Model(&types.KnowledgeBase{}).
|
||||
Where("tenant_id = ? AND vector_store_id = ?", tenantID, storeID).
|
||||
Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// The remaining methods are not called by the tested code paths; declare them
|
||||
// so realKBRepo satisfies interfaces.KnowledgeBaseRepository.
|
||||
func (r *realKBRepo) CreateKnowledgeBase(_ context.Context, kb *types.KnowledgeBase) error {
|
||||
return r.db.Create(kb).Error
|
||||
}
|
||||
func (r *realKBRepo) GetKnowledgeBaseByID(_ context.Context, _ string) (*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *realKBRepo) GetKnowledgeBaseByIDAndTenant(_ context.Context, _ string, _ uint64) (*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *realKBRepo) GetKnowledgeBaseByIDs(_ context.Context, _ []string) ([]*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *realKBRepo) ListKnowledgeBases(_ context.Context) ([]*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *realKBRepo) ListKnowledgeBasesByTenantID(_ context.Context, _ uint64) ([]*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *realKBRepo) UpdateKnowledgeBase(_ context.Context, _ *types.KnowledgeBase) error {
|
||||
return nil
|
||||
}
|
||||
func (r *realKBRepo) DeleteKnowledgeBase(_ context.Context, _ string) error {
|
||||
return nil
|
||||
}
|
||||
func (r *realKBRepo) TogglePinKnowledgeBase(_ context.Context, _ string, _ uint64) (*types.KnowledgeBase, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func insertGuardStore(t *testing.T, db *gorm.DB, id string, tenantID uint64) {
|
||||
t.Helper()
|
||||
require.NoError(t, db.Create(&types.VectorStore{
|
||||
ID: id, Name: id, EngineType: types.QdrantRetrieverEngineType, TenantID: tenantID,
|
||||
}).Error)
|
||||
}
|
||||
|
||||
func insertGuardKB(t *testing.T, db *gorm.DB, tenantID uint64, vsid *string) string {
|
||||
t.Helper()
|
||||
id := "kb-" + tenantID2s(tenantID) + "-" + ptrOrEmpty(vsid)
|
||||
kb := &types.KnowledgeBase{
|
||||
ID: id,
|
||||
Name: id,
|
||||
TenantID: tenantID,
|
||||
EmbeddingModelID: "embed",
|
||||
SummaryModelID: "sum",
|
||||
VectorStoreID: vsid,
|
||||
}
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
return id
|
||||
}
|
||||
|
||||
func tenantID2s(t uint64) string {
|
||||
if t == 1 {
|
||||
return "t1"
|
||||
}
|
||||
if t == 2 {
|
||||
return "t2"
|
||||
}
|
||||
return "tN"
|
||||
}
|
||||
func ptrOrEmpty(p *string) string {
|
||||
if p == nil {
|
||||
return "nil"
|
||||
}
|
||||
return *p
|
||||
}
|
||||
|
||||
func TestDeleteStore_Guard_Success(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, db, registry := newGuardTestService(t)
|
||||
insertGuardStore(t, db, "store-A", 1)
|
||||
|
||||
err := svc.DeleteStore(ctx, 1, "store-A")
|
||||
require.NoError(t, err)
|
||||
|
||||
// store row soft-deleted, no longer visible to default GORM scope.
|
||||
var got types.VectorStore
|
||||
err = db.Where("id = ?", "store-A").First(&got).Error
|
||||
require.ErrorIs(t, err, gorm.ErrRecordNotFound)
|
||||
|
||||
// registry unregister was invoked once.
|
||||
require.Equal(t, []string{"store-A"}, registry.unregistered)
|
||||
}
|
||||
|
||||
func TestDeleteStore_Guard_RejectsBoundKB(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, db, registry := newGuardTestService(t)
|
||||
insertGuardStore(t, db, "store-A", 1)
|
||||
storeID := "store-A"
|
||||
insertGuardKB(t, db, 1, &storeID)
|
||||
|
||||
err := svc.DeleteStore(ctx, 1, "store-A")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "still has 1 knowledge base")
|
||||
|
||||
// store row must still be present (rollback succeeded).
|
||||
var got types.VectorStore
|
||||
require.NoError(t, db.Where("id = ?", "store-A").First(&got).Error)
|
||||
require.Empty(t, registry.unregistered, "registry must NOT be unregistered when guard fires")
|
||||
}
|
||||
|
||||
func TestDeleteStore_Guard_IgnoresSoftDeletedKB(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, db, registry := newGuardTestService(t)
|
||||
insertGuardStore(t, db, "store-A", 1)
|
||||
storeID := "store-A"
|
||||
kbID := insertGuardKB(t, db, 1, &storeID)
|
||||
require.NoError(t, db.Where("id = ?", kbID).Delete(&types.KnowledgeBase{}).Error)
|
||||
|
||||
err := svc.DeleteStore(ctx, 1, "store-A")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"store-A"}, registry.unregistered)
|
||||
}
|
||||
|
||||
func TestDeleteStore_Guard_IgnoresOtherTenantKB(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, db, registry := newGuardTestService(t)
|
||||
insertGuardStore(t, db, "store-A", 1)
|
||||
storeID := "store-A"
|
||||
insertGuardKB(t, db, 2, &storeID) // cross-tenant binding (should not block)
|
||||
|
||||
err := svc.DeleteStore(ctx, 1, "store-A")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"store-A"}, registry.unregistered)
|
||||
}
|
||||
|
||||
func TestDeleteStore_Guard_StoreNotFound(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, _, registry := newGuardTestService(t)
|
||||
|
||||
err := svc.DeleteStore(ctx, 1, "store-missing")
|
||||
require.Error(t, err)
|
||||
appErr, ok := errors.IsAppError(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, 404, appErr.HTTPCode)
|
||||
require.Empty(t, registry.unregistered)
|
||||
}
|
||||
|
||||
func TestDeleteStore_Guard_RejectsCount3(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, db, _ := newGuardTestService(t)
|
||||
insertGuardStore(t, db, "store-A", 1)
|
||||
storeID := "store-A"
|
||||
for i := 0; i < 3; i++ {
|
||||
kb := &types.KnowledgeBase{
|
||||
ID: "kb-multi-" + tenantID2s(1) + "-" + ptrOrEmpty(&storeID) + "-" + string(rune('a'+i)),
|
||||
Name: "kb", TenantID: 1, EmbeddingModelID: "e", SummaryModelID: "s",
|
||||
VectorStoreID: &storeID,
|
||||
}
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
}
|
||||
|
||||
err := svc.DeleteStore(ctx, 1, "store-A")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "still has 3 knowledge base")
|
||||
}
|
||||
|
||||
func TestUnregisterSafely_PanicRecovered(t *testing.T) {
|
||||
svc := &vectorStoreService{
|
||||
storeRegistry: panickingRegistry{},
|
||||
}
|
||||
// Must not panic out of the test process.
|
||||
svc.unregisterSafely(context.Background(), "store-A")
|
||||
}
|
||||
|
||||
type panickingRegistry struct{}
|
||||
|
||||
func (panickingRegistry) RegisterWithStoreID(string, interfaces.RetrieveEngineService) {}
|
||||
func (panickingRegistry) GetByStoreID(string) (interfaces.RetrieveEngineService, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (panickingRegistry) UnregisterByStoreID(string) { panic("registry exploded") }
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ResolveStoreView / BatchResolveStoreView
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestResolveStoreView(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, db, _ := newGuardTestService(t)
|
||||
insertGuardStore(t, db, "store-A", 1)
|
||||
svc.envStores = []types.VectorStore{
|
||||
{ID: "__env_es__", Name: "Env ES", EngineType: types.ElasticsearchRetrieverEngineType},
|
||||
}
|
||||
|
||||
t.Run("empty store id returns DefaultStoreDisplay", func(t *testing.T) {
|
||||
v, err := svc.ResolveStoreView(ctx, 1, "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "System default", v.Name)
|
||||
assert.Equal(t, types.StoreSourceEnv, v.Source)
|
||||
})
|
||||
t.Run("DB hit", func(t *testing.T) {
|
||||
v, err := svc.ResolveStoreView(ctx, 1, "store-A")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "store-A", v.Name)
|
||||
assert.Equal(t, types.StoreSourceUser, v.Source)
|
||||
assert.Equal(t, string(types.QdrantRetrieverEngineType), v.EngineType)
|
||||
})
|
||||
t.Run("env hit", func(t *testing.T) {
|
||||
v, err := svc.ResolveStoreView(ctx, 1, "__env_es__")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Env ES", v.Name)
|
||||
assert.Equal(t, types.StoreSourceEnv, v.Source)
|
||||
})
|
||||
t.Run("miss returns unavailable", func(t *testing.T) {
|
||||
v, err := svc.ResolveStoreView(ctx, 1, "store-zzz")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, types.StoreSourceUnavailable, v.Source)
|
||||
})
|
||||
t.Run("cross-tenant store is treated as unavailable", func(t *testing.T) {
|
||||
insertGuardStore(t, db, "store-B", 2) // owned by tenant 2
|
||||
v, err := svc.ResolveStoreView(ctx, 1, "store-B")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, types.StoreSourceUnavailable, v.Source)
|
||||
})
|
||||
}
|
||||
|
||||
func TestBatchResolveStoreView(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, db, _ := newGuardTestService(t)
|
||||
insertGuardStore(t, db, "store-A", 1)
|
||||
insertGuardStore(t, db, "store-B", 1)
|
||||
svc.envStores = []types.VectorStore{
|
||||
{ID: "__env_qd__", Name: "Env QD", EngineType: types.QdrantRetrieverEngineType},
|
||||
}
|
||||
|
||||
got, err := svc.BatchResolveStoreView(ctx, 1, []string{
|
||||
"store-A", "store-zzz", "__env_qd__", "",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got, 4)
|
||||
assert.Equal(t, types.StoreSourceUser, got["store-A"].Source)
|
||||
assert.Equal(t, "store-A", got["store-A"].Name)
|
||||
assert.Equal(t, types.StoreSourceUnavailable, got["store-zzz"].Source)
|
||||
assert.Equal(t, types.StoreSourceEnv, got["__env_qd__"].Source)
|
||||
assert.Equal(t, "Env QD", got["__env_qd__"].Name)
|
||||
assert.Equal(t, types.StoreSourceEnv, got[""].Source)
|
||||
assert.Equal(t, "System default", got[""].Name)
|
||||
}
|
||||
|
||||
func TestBatchResolveStoreView_Empty(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, _, _ := newGuardTestService(t)
|
||||
got, err := svc.BatchResolveStoreView(ctx, 1, nil)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, got)
|
||||
}
|
||||
|
||||
@@ -472,6 +472,18 @@ func initDatabase(cfg *config.Config) (*gorm.DB, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Sanity check: dialect-specific code in services (notably the
|
||||
// vector_stores delete guard) compares Dialector.Name() to "postgres" /
|
||||
// "sqlite" string literals. A future driver swap that produces a
|
||||
// different name (e.g., a wrapper dialect for managed PG) would silently
|
||||
// fall back to the SQLite path, dropping the row-level X-lock. Catching
|
||||
// the mismatch at startup is loud and inexpensive.
|
||||
if name := db.Dialector.Name(); name != "postgres" && name != "sqlite" {
|
||||
return nil, fmt.Errorf(
|
||||
"unsupported gorm dialector %q; expected postgres or sqlite "+
|
||||
"(see vectorStoreService.isPostgres for impact)", name)
|
||||
}
|
||||
|
||||
if os.Getenv("DB_DRIVER") == "sqlite" {
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
|
||||
@@ -36,6 +36,13 @@ const (
|
||||
ErrAgentInvalidMaxIterations ErrorCode = 2102
|
||||
ErrAgentInvalidTemperature ErrorCode = 2103
|
||||
|
||||
// VectorStore binding related error codes (2200-2299).
|
||||
// Both map to HTTP 400 with a generic message; the typed code lets
|
||||
// clients distinguish "wrong UUID / cross-tenant" from "store exists
|
||||
// but is currently unavailable" without parsing the message text.
|
||||
ErrVectorStoreBindingInvalid ErrorCode = 2200
|
||||
ErrVectorStoreUnavailable ErrorCode = 2201
|
||||
|
||||
// Add more error codes here
|
||||
)
|
||||
|
||||
@@ -184,6 +191,37 @@ func NewAgentInvalidTemperatureError() *AppError {
|
||||
}
|
||||
}
|
||||
|
||||
// NewVectorStoreBindingInvalidError signals that a knowledge base create
|
||||
// request referenced a vector store that does not exist under the caller's
|
||||
// tenant (or carried a malformed UUID). The user-facing message is
|
||||
// intentionally generic to avoid enumeration oracles — see the structured
|
||||
// log at the call site for the tenant/store pair.
|
||||
func NewVectorStoreBindingInvalidError(message string) *AppError {
|
||||
if message == "" {
|
||||
message = "vector store not found"
|
||||
}
|
||||
return &AppError{
|
||||
Code: ErrVectorStoreBindingInvalid,
|
||||
Message: message,
|
||||
HTTPCode: http.StatusBadRequest,
|
||||
}
|
||||
}
|
||||
|
||||
// NewVectorStoreUnavailableError signals that the requested vector store
|
||||
// row exists in the database but is not currently wired into the in-memory
|
||||
// engine registry (factory failure on CreateStore, dynamic config error,
|
||||
// or stale state after a connection-config rotation).
|
||||
func NewVectorStoreUnavailableError(message string) *AppError {
|
||||
if message == "" {
|
||||
message = "vector store is currently unavailable"
|
||||
}
|
||||
return &AppError{
|
||||
Code: ErrVectorStoreUnavailable,
|
||||
Message: message,
|
||||
HTTPCode: http.StatusBadRequest,
|
||||
}
|
||||
}
|
||||
|
||||
// IsAppError checks if the error is an AppError type
|
||||
func IsAppError(err error) (*AppError, bool) {
|
||||
appErr, ok := err.(*AppError)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
stderrors "errors"
|
||||
"net/http"
|
||||
@@ -25,11 +26,12 @@ import (
|
||||
|
||||
// KnowledgeBaseHandler defines the HTTP handler for knowledge base operations
|
||||
type KnowledgeBaseHandler struct {
|
||||
service interfaces.KnowledgeBaseService
|
||||
knowledgeService interfaces.KnowledgeService
|
||||
kbShareService interfaces.KBShareService
|
||||
agentShareService interfaces.AgentShareService
|
||||
asynqClient interfaces.TaskEnqueuer
|
||||
service interfaces.KnowledgeBaseService
|
||||
knowledgeService interfaces.KnowledgeService
|
||||
kbShareService interfaces.KBShareService
|
||||
agentShareService interfaces.AgentShareService
|
||||
asynqClient interfaces.TaskEnqueuer
|
||||
vectorStoreService interfaces.VectorStoreService // enriches KB responses with bound store display
|
||||
}
|
||||
|
||||
// NewKnowledgeBaseHandler creates a new knowledge base handler instance
|
||||
@@ -39,16 +41,103 @@ func NewKnowledgeBaseHandler(
|
||||
kbShareService interfaces.KBShareService,
|
||||
agentShareService interfaces.AgentShareService,
|
||||
asynqClient interfaces.TaskEnqueuer,
|
||||
vectorStoreService interfaces.VectorStoreService,
|
||||
) *KnowledgeBaseHandler {
|
||||
return &KnowledgeBaseHandler{
|
||||
service: service,
|
||||
knowledgeService: knowledgeService,
|
||||
kbShareService: kbShareService,
|
||||
agentShareService: agentShareService,
|
||||
asynqClient: asynqClient,
|
||||
service: service,
|
||||
knowledgeService: knowledgeService,
|
||||
kbShareService: kbShareService,
|
||||
agentShareService: agentShareService,
|
||||
asynqClient: asynqClient,
|
||||
vectorStoreService: vectorStoreService,
|
||||
}
|
||||
}
|
||||
|
||||
// buildKBResponse turns a knowledge base into a JSON-ready response shape,
|
||||
// merging the bound vector store's display metadata and any caller-supplied
|
||||
// extras (e.g., my_permission for shared KBs). Returns the kb pointer
|
||||
// unchanged on serialization failure so the request still succeeds.
|
||||
//
|
||||
// The map-merge approach (rather than a wrapper struct embedding the kb)
|
||||
// is deliberate: KnowledgeBase has a custom MarshalJSON, and embedding
|
||||
// would promote it to any wrapper struct and silently swallow the extra
|
||||
// fields. The same pattern is already used by GetKnowledgeBase to add the
|
||||
// my_permission field for shared knowledge bases.
|
||||
//
|
||||
// Shared-KB suppression: when storeView.Source == StoreSourceShared (the
|
||||
// caller is not the KB owner), the raw vector_store_id UUID is stripped
|
||||
// from the response so the owner-tenant's store inventory cannot be
|
||||
// correlated across multiple shared KBs. Name / engine type / status are
|
||||
// already empty in the SharedStoreDisplay payload; suppressing the UUID
|
||||
// completes the cross-tenant metadata hiding.
|
||||
func buildKBResponse(
|
||||
kb *types.KnowledgeBase,
|
||||
storeView types.StoreDisplay,
|
||||
extras map[string]interface{},
|
||||
) interface{} {
|
||||
b, err := json.Marshal(kb)
|
||||
if err != nil {
|
||||
return kb
|
||||
}
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(b, &m); err != nil || m == nil {
|
||||
return kb
|
||||
}
|
||||
if storeView.Source == types.StoreSourceShared {
|
||||
delete(m, "vector_store_id")
|
||||
}
|
||||
if storeView.Name != "" {
|
||||
m["vector_store_name"] = storeView.Name
|
||||
}
|
||||
if storeView.Source != "" {
|
||||
m["vector_store_source"] = storeView.Source
|
||||
}
|
||||
if storeView.EngineType != "" {
|
||||
m["vector_store_engine_type"] = storeView.EngineType
|
||||
}
|
||||
if storeView.Status != "" {
|
||||
m["vector_store_status"] = storeView.Status
|
||||
}
|
||||
for k, v := range extras {
|
||||
m[k] = v
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// resolveKBStoreView returns the store display payload to embed in the KB
|
||||
// response. It applies two policies on top of the service-level resolver:
|
||||
//
|
||||
// - When the KB does not have a DB-managed vector store binding,
|
||||
// the env-fallback display is returned without touching the service.
|
||||
// - When the caller is not the KB owner (shared access), the underlying
|
||||
// store's name and engine are suppressed so operator-chosen names do
|
||||
// not leak across tenants. The Source value is set to "shared".
|
||||
//
|
||||
// On resolution error, an unavailable display is returned and the failure
|
||||
// is logged for ops; the request itself still succeeds.
|
||||
func (h *KnowledgeBaseHandler) resolveKBStoreView(
|
||||
ctx context.Context, kb *types.KnowledgeBase, callerTenantID uint64,
|
||||
) types.StoreDisplay {
|
||||
if !kb.HasVectorStore() {
|
||||
return types.DefaultStoreDisplay()
|
||||
}
|
||||
if kb.TenantID != callerTenantID {
|
||||
return types.SharedStoreDisplay()
|
||||
}
|
||||
if h.vectorStoreService == nil {
|
||||
return types.UnavailableStoreDisplay()
|
||||
}
|
||||
view, err := h.vectorStoreService.ResolveStoreView(ctx, kb.TenantID, *kb.VectorStoreID)
|
||||
if err != nil {
|
||||
logger.WarnWithFields(ctx, logger.Fields{
|
||||
"kb_id": secutils.SanitizeForLog(kb.ID),
|
||||
"tenant_id": kb.TenantID,
|
||||
}, "[kb.view] vector store resolve failed; returning unavailable")
|
||||
return types.UnavailableStoreDisplay()
|
||||
}
|
||||
return view
|
||||
}
|
||||
|
||||
// HybridSearch godoc
|
||||
// @Summary 混合搜索
|
||||
// @Description 在知识库中执行向量和关键词混合搜索
|
||||
@@ -141,6 +230,16 @@ func (h *KnowledgeBaseHandler) CreateKnowledgeBase(c *gin.Context) {
|
||||
// Create knowledge base using the service
|
||||
kb, err := h.service.CreateKnowledgeBase(ctx, &req)
|
||||
if err != nil {
|
||||
// Surface typed AppErrors (notably the 400-class codes
|
||||
// ErrVectorStoreBindingInvalid and ErrVectorStoreUnavailable
|
||||
// returned by validateVectorStoreBinding) instead of wrapping them
|
||||
// into a generic 500. The middleware renders the original code and
|
||||
// HTTP status verbatim. Falls through to 500 only for raw infra
|
||||
// errors that the service did not classify.
|
||||
if appErr, ok := apperrors.IsAppError(err); ok {
|
||||
c.Error(appErr)
|
||||
return
|
||||
}
|
||||
logger.ErrorWithFields(ctx, err, nil)
|
||||
c.Error(apperrors.NewInternalServerError(err.Error()))
|
||||
return
|
||||
@@ -148,9 +247,10 @@ func (h *KnowledgeBaseHandler) CreateKnowledgeBase(c *gin.Context) {
|
||||
|
||||
logger.Infof(ctx, "Knowledge base created successfully, ID: %s, name: %s",
|
||||
secutils.SanitizeForLog(kb.ID), secutils.SanitizeForLog(kb.Name))
|
||||
callerTenantID := c.GetUint64(types.TenantIDContextKey.String())
|
||||
c.JSON(http.StatusCreated, gin.H{
|
||||
"success": true,
|
||||
"data": kb,
|
||||
"data": buildKBResponse(kb, h.resolveKBStoreView(ctx, kb, callerTenantID), nil),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -285,18 +385,13 @@ func (h *KnowledgeBaseHandler) GetKnowledgeBase(c *gin.Context) {
|
||||
logger.Warnf(c.Request.Context(), "Failed to fill KB counts for %s: %v", kb.ID, fillErr)
|
||||
}
|
||||
tenantID := c.GetUint64(types.TenantIDContextKey.String())
|
||||
data := interface{}(kb)
|
||||
storeView := h.resolveKBStoreView(c.Request.Context(), kb, tenantID)
|
||||
var extras map[string]interface{}
|
||||
if kb.TenantID != tenantID && permission != "" {
|
||||
// Include my_permission in data so frontend can show role (e.g. "只读") instead of "--" for agent-visible KBs
|
||||
var dataMap map[string]interface{}
|
||||
b, _ := json.Marshal(kb)
|
||||
_ = json.Unmarshal(b, &dataMap)
|
||||
if dataMap != nil {
|
||||
dataMap["my_permission"] = permission
|
||||
data = dataMap
|
||||
}
|
||||
extras = map[string]interface{}{"my_permission": permission}
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "data": data})
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "data": buildKBResponse(kb, storeView, extras)})
|
||||
}
|
||||
|
||||
// ListKnowledgeBases godoc
|
||||
@@ -460,9 +555,10 @@ func (h *KnowledgeBaseHandler) TogglePinKnowledgeBase(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
callerTenantID := c.GetUint64(types.TenantIDContextKey.String())
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": kb,
|
||||
"data": buildKBResponse(kb, h.resolveKBStoreView(ctx, kb, callerTenantID), nil),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -524,9 +620,10 @@ func (h *KnowledgeBaseHandler) UpdateKnowledgeBase(c *gin.Context) {
|
||||
|
||||
logger.Infof(ctx, "Knowledge base updated successfully, ID: %s",
|
||||
secutils.SanitizeForLog(id))
|
||||
callerTenantID := c.GetUint64(types.TenantIDContextKey.String())
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": kb,
|
||||
"data": buildKBResponse(kb, h.resolveKBStoreView(ctx, kb, callerTenantID), nil),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -641,6 +738,10 @@ func (h *KnowledgeBaseHandler) CopyKnowledgeBase(c *gin.Context) {
|
||||
}
|
||||
|
||||
// If target_id provided, validate target belongs to caller's tenant
|
||||
// and run the pre-flight defenses synchronously so a mismatched
|
||||
// clone is rejected with 400 before the task is enqueued. The same
|
||||
// checks are re-applied inside the async worker (service.CopyKnowledgeBase)
|
||||
// as defense in depth.
|
||||
if req.TargetID != "" {
|
||||
targetKB, err := h.service.GetKnowledgeBaseByID(ctx, req.TargetID)
|
||||
if err != nil {
|
||||
@@ -658,6 +759,24 @@ func (h *KnowledgeBaseHandler) CopyKnowledgeBase(c *gin.Context) {
|
||||
c.Error(errors.NewForbiddenError("No permission to copy to this knowledge base"))
|
||||
return
|
||||
}
|
||||
// Pre-flight defense 1: embedding model must match.
|
||||
// Without this check the async clone would run with incompatible
|
||||
// vector spaces and produce semantically broken results.
|
||||
if sourceKB.EmbeddingModelID != targetKB.EmbeddingModelID {
|
||||
c.Error(apperrors.NewBadRequestError(
|
||||
"source and target knowledge bases use different embedding models; " +
|
||||
"clone into a target with the same embedding model"))
|
||||
return
|
||||
}
|
||||
// Pre-flight defense 2: vector store binding must match.
|
||||
// Cross-store cloning would require copying physical vector data
|
||||
// between stores, which is not yet supported.
|
||||
if !sourceKB.SharesStoreWith(targetKB) {
|
||||
c.Error(apperrors.NewBadRequestError(
|
||||
"source and target knowledge bases are bound to different vector stores; " +
|
||||
"cross-store cloning is not yet supported"))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Generate task ID if not provided
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
apperrors "github.com/Tencent/WeKnora/internal/errors"
|
||||
"github.com/Tencent/WeKnora/internal/middleware"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// handler.CopyKnowledgeBase pre-flight tests.
|
||||
//
|
||||
// The async clone worker (service.CopyKnowledgeBase) re-applies the same
|
||||
// embedding-model and store-binding defenses as defense in depth, but the
|
||||
// handler-level pre-flight is the one that surfaces 400 to the API caller
|
||||
// synchronously instead of inside Asynq progress polling. These tests pin
|
||||
// the synchronous behavior so a future refactor that drops the pre-flight
|
||||
// fails loudly here rather than silently degrading UX.
|
||||
|
||||
// stubKBCopyService provides only the two methods the Copy handler reaches
|
||||
// for (GetKnowledgeBaseByID twice). Other interface methods stay nil so any
|
||||
// accidental new call panics rather than silently succeeding.
|
||||
type stubKBCopyService struct {
|
||||
interfaces.KnowledgeBaseService
|
||||
byID func(ctx context.Context, id string) (*types.KnowledgeBase, error)
|
||||
}
|
||||
|
||||
func (s *stubKBCopyService) GetKnowledgeBaseByID(ctx context.Context, id string) (*types.KnowledgeBase, error) {
|
||||
return s.byID(ctx, id)
|
||||
}
|
||||
|
||||
// stubEnqueuer records whether Enqueue was invoked. The whole point of the
|
||||
// pre-flight is to short-circuit *before* enqueue, so the test fails if
|
||||
// enqueue ran for a mismatched clone.
|
||||
type stubEnqueuer struct {
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *stubEnqueuer) Enqueue(_ interface{}, _ ...interface{}) (*stubEnqueueInfo, error) {
|
||||
s.calls++
|
||||
return &stubEnqueueInfo{ID: "x"}, nil
|
||||
}
|
||||
|
||||
// stubEnqueueInfo is a stand-in for asynq.TaskInfo so the test does not need
|
||||
// to construct one.
|
||||
type stubEnqueueInfo struct{ ID string }
|
||||
|
||||
func newCopyPreflightRouter(svc interfaces.KnowledgeBaseService) (*gin.Engine, *stubEnqueuer) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(middleware.ErrorHandler())
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Set(types.TenantIDContextKey.String(), uint64(1))
|
||||
c.Set(types.UserIDContextKey.String(), "u-test")
|
||||
c.Next()
|
||||
})
|
||||
enq := &stubEnqueuer{}
|
||||
// The real handler reaches into h.asynqClient.Enqueue with the asynq
|
||||
// task type. We do not exercise the enqueue path in these tests — every
|
||||
// case here either short-circuits at pre-flight or returns 4xx earlier.
|
||||
// Leaving asynqClient nil would panic if enqueue ran, which is exactly
|
||||
// the regression we want to catch.
|
||||
h := &KnowledgeBaseHandler{service: svc}
|
||||
r.POST("/knowledge-bases/copy", h.CopyKnowledgeBase)
|
||||
return r, enq
|
||||
}
|
||||
|
||||
func storeIDPtr(s string) *string { return &s }
|
||||
|
||||
func TestCopyHandlerPreflight_DifferentEmbeddingModel(t *testing.T) {
|
||||
svc := &stubKBCopyService{
|
||||
byID: func(_ context.Context, id string) (*types.KnowledgeBase, error) {
|
||||
switch id {
|
||||
case "src":
|
||||
return &types.KnowledgeBase{
|
||||
ID: "src", TenantID: 1, EmbeddingModelID: "embed-A",
|
||||
}, nil
|
||||
case "dst":
|
||||
return &types.KnowledgeBase{
|
||||
ID: "dst", TenantID: 1, EmbeddingModelID: "embed-B",
|
||||
}, nil
|
||||
}
|
||||
return nil, nil
|
||||
},
|
||||
}
|
||||
r, _ := newCopyPreflightRouter(svc)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/knowledge-bases/copy",
|
||||
strings.NewReader(`{"source_id":"src","target_id":"dst"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected 400 for embedding-model mismatch, got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), "different embedding models") {
|
||||
t.Fatalf("expected embedding-model error message, got %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyHandlerPreflight_DifferentVectorStore(t *testing.T) {
|
||||
svc := &stubKBCopyService{
|
||||
byID: func(_ context.Context, id string) (*types.KnowledgeBase, error) {
|
||||
switch id {
|
||||
case "src":
|
||||
return &types.KnowledgeBase{
|
||||
ID: "src", TenantID: 1, EmbeddingModelID: "embed-A",
|
||||
VectorStoreID: storeIDPtr("store-A"),
|
||||
}, nil
|
||||
case "dst":
|
||||
return &types.KnowledgeBase{
|
||||
ID: "dst", TenantID: 1, EmbeddingModelID: "embed-A",
|
||||
VectorStoreID: storeIDPtr("store-B"),
|
||||
}, nil
|
||||
}
|
||||
return nil, nil
|
||||
},
|
||||
}
|
||||
r, _ := newCopyPreflightRouter(svc)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/knowledge-bases/copy",
|
||||
strings.NewReader(`{"source_id":"src","target_id":"dst"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected 400 for store mismatch, got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), "different vector stores") {
|
||||
t.Fatalf("expected store-mismatch error message, got %s", w.Body.String())
|
||||
}
|
||||
if strings.Contains(w.Body.String(), "Phase 4") {
|
||||
t.Fatalf("error message must not leak internal roadmap labels: %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyHandlerPreflight_OneSideNilStore(t *testing.T) {
|
||||
svc := &stubKBCopyService{
|
||||
byID: func(_ context.Context, id string) (*types.KnowledgeBase, error) {
|
||||
switch id {
|
||||
case "src":
|
||||
return &types.KnowledgeBase{
|
||||
ID: "src", TenantID: 1, EmbeddingModelID: "embed-A",
|
||||
VectorStoreID: nil,
|
||||
}, nil
|
||||
case "dst":
|
||||
return &types.KnowledgeBase{
|
||||
ID: "dst", TenantID: 1, EmbeddingModelID: "embed-A",
|
||||
VectorStoreID: storeIDPtr("store-A"),
|
||||
}, nil
|
||||
}
|
||||
return nil, nil
|
||||
},
|
||||
}
|
||||
r, _ := newCopyPreflightRouter(svc)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/knowledge-bases/copy",
|
||||
strings.NewReader(`{"source_id":"src","target_id":"dst"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected 400 when one side is env-store and the other is DB-store, got %d body=%s",
|
||||
w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// compile-time guard against accidentally dropping the apperrors import
|
||||
// from the file — if the pre-flight refactor goes away, this fails too.
|
||||
var _ = apperrors.NewBadRequestError
|
||||
@@ -0,0 +1,188 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
apperrors "github.com/Tencent/WeKnora/internal/errors"
|
||||
"github.com/Tencent/WeKnora/internal/middleware"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// CreateKnowledgeBase typed-error preservation — the handler must surface
|
||||
// the typed AppError (ErrVectorStoreBindingInvalid / ErrVectorStoreUnavailable)
|
||||
// returned by validateVectorStoreBinding instead of stripping it into a
|
||||
// generic 500 via NewInternalServerError. Without the IsAppError unwrap in
|
||||
// the handler, the typed error codes would be silently nullified at the
|
||||
// HTTP boundary and clients would lose the ability to branch on the cause.
|
||||
//
|
||||
// Shared-KB UUID suppression — responses for cross-tenant shared KBs must
|
||||
// not leak the owner tenant's vector_store_id UUID. SharedStoreDisplay
|
||||
// suppresses store name + engine_type for cross-tenant callers, but the
|
||||
// underlying KnowledgeBase.MarshalJSON still emits the UUID; the
|
||||
// buildKBResponse strip closes the gap so the UUID cannot be correlated
|
||||
// across multiple shared KBs.
|
||||
|
||||
// stubKBCreateService drives CreateKnowledgeBase end-to-end with a
|
||||
// service that returns a chosen error. Embedding the interface keeps
|
||||
// any other method nil-panic'ing on purpose.
|
||||
type stubKBCreateService struct {
|
||||
interfaces.KnowledgeBaseService
|
||||
createErr error
|
||||
}
|
||||
|
||||
func (s *stubKBCreateService) CreateKnowledgeBase(_ context.Context, kb *types.KnowledgeBase) (*types.KnowledgeBase, error) {
|
||||
if s.createErr != nil {
|
||||
return nil, s.createErr
|
||||
}
|
||||
kb.ID = "kb-new"
|
||||
kb.TenantID = 1
|
||||
return kb, nil
|
||||
}
|
||||
|
||||
func newCreateKBRouter(svc interfaces.KnowledgeBaseService) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
r.Use(middleware.ErrorHandler())
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Set(types.TenantIDContextKey.String(), uint64(1))
|
||||
c.Set(types.UserIDContextKey.String(), "u-test")
|
||||
c.Next()
|
||||
})
|
||||
h := &KnowledgeBaseHandler{service: svc}
|
||||
r.POST("/knowledge-bases", h.CreateKnowledgeBase)
|
||||
return r
|
||||
}
|
||||
|
||||
func TestCreateKB_PreservesTypedErrorCode_2200(t *testing.T) {
|
||||
svc := &stubKBCreateService{
|
||||
createErr: apperrors.NewVectorStoreBindingInvalidError("vector store not found"),
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/knowledge-bases",
|
||||
strings.NewReader(`{"name":"kb"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
newCreateKBRouter(svc).ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected 400, got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
body := w.Body.String()
|
||||
if !strings.Contains(body, `"code":2200`) {
|
||||
t.Fatalf("expected envelope to contain code 2200, got %s", body)
|
||||
}
|
||||
if strings.Contains(body, `"code":1007`) || strings.Contains(body, `"code":1000`) {
|
||||
t.Fatalf("typed error must not be wrapped into a generic code, got %s", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateKB_PreservesTypedErrorCode_2201(t *testing.T) {
|
||||
svc := &stubKBCreateService{
|
||||
createErr: apperrors.NewVectorStoreUnavailableError(""),
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/knowledge-bases",
|
||||
strings.NewReader(`{"name":"kb"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
newCreateKBRouter(svc).ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("expected 400, got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), `"code":2201`) {
|
||||
t.Fatalf("expected envelope to contain code 2201, got %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateKB_GenericErrorStillFallsThroughTo500(t *testing.T) {
|
||||
// A non-AppError must NOT be auto-rewritten to 200/400 — operational
|
||||
// monitoring still needs to see infrastructure failures as 5xx.
|
||||
svc := &stubKBCreateService{createErr: errSentinel("connection refused")}
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/knowledge-bases",
|
||||
strings.NewReader(`{"name":"kb"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
newCreateKBRouter(svc).ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusInternalServerError {
|
||||
t.Fatalf("expected 500 for raw infra error, got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
type errSentinel string
|
||||
|
||||
func (e errSentinel) Error() string { return string(e) }
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// buildKBResponse must strip vector_store_id for shared KB responses
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestBuildKBResponse_StripsVectorStoreIDForSharedKB(t *testing.T) {
|
||||
storeID := "aaaa-bbbb-cccc-dddd"
|
||||
kb := &types.KnowledgeBase{
|
||||
ID: "kb-1",
|
||||
Name: "shared-kb",
|
||||
TenantID: 42, // different from caller
|
||||
EmbeddingModelID: "e",
|
||||
SummaryModelID: "s",
|
||||
VectorStoreID: &storeID,
|
||||
}
|
||||
got := buildKBResponse(kb, types.SharedStoreDisplay(), nil)
|
||||
m, ok := got.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected map result, got %T", got)
|
||||
}
|
||||
if _, exists := m["vector_store_id"]; exists {
|
||||
t.Fatalf("shared KB response must not expose vector_store_id, got %v", m["vector_store_id"])
|
||||
}
|
||||
if _, exists := m["vector_store_name"]; exists {
|
||||
t.Fatalf("shared KB response must not expose vector_store_name, got %v", m["vector_store_name"])
|
||||
}
|
||||
if m["vector_store_source"] != types.StoreSourceShared {
|
||||
t.Fatalf("expected vector_store_source=shared, got %v", m["vector_store_source"])
|
||||
}
|
||||
// Defensive: ensure the source UUID does not appear *anywhere* in
|
||||
// the serialized output (paranoid check against future map keys).
|
||||
serialized, _ := json.Marshal(m)
|
||||
if strings.Contains(string(serialized), storeID) {
|
||||
t.Fatalf("shared KB response leaked vector store UUID via some path: %s", serialized)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildKBResponse_KeepsVectorStoreIDForOwnerKB(t *testing.T) {
|
||||
// Same setup but with the user-source display — owner caller should
|
||||
// still see the UUID alongside the resolved metadata.
|
||||
storeID := "aaaa-bbbb-cccc-dddd"
|
||||
kb := &types.KnowledgeBase{
|
||||
ID: "kb-1",
|
||||
Name: "owner-kb",
|
||||
TenantID: 1,
|
||||
EmbeddingModelID: "e",
|
||||
SummaryModelID: "s",
|
||||
VectorStoreID: &storeID,
|
||||
}
|
||||
view := types.StoreDisplay{
|
||||
Name: "prod-es",
|
||||
Source: types.StoreSourceUser,
|
||||
EngineType: "elasticsearch",
|
||||
Status: "available",
|
||||
}
|
||||
got := buildKBResponse(kb, view, nil)
|
||||
m, ok := got.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected map result, got %T", got)
|
||||
}
|
||||
if m["vector_store_id"] != storeID {
|
||||
t.Fatalf("owner KB must keep vector_store_id, got %v", m["vector_store_id"])
|
||||
}
|
||||
if m["vector_store_name"] != "prod-es" {
|
||||
t.Fatalf("owner KB must surface store name, got %v", m["vector_store_name"])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// TestUpdateKBRequest_DoesNotAcceptVectorStoreID is the structural enforcement
|
||||
// behind the vector_store_id immutability contract. The GORM `<-:create`
|
||||
// tag on KnowledgeBase.VectorStoreID already blocks every ORM UPDATE path
|
||||
// (verified by the repository-level sqlite immutability tests), but the
|
||||
// service DTO must independently refuse to even *accept* the field —
|
||||
// otherwise a future maintainer who adds it to UpdateKnowledgeBaseRequest
|
||||
// or KnowledgeBaseConfig opens a path where the field is silently ignored
|
||||
// by the ORM, which is worse than an explicit rejection.
|
||||
//
|
||||
// This test walks the request and config struct shapes and fails if either
|
||||
// gains a VectorStoreID member, by name or by JSON tag.
|
||||
func TestUpdateKBRequest_DoesNotAcceptVectorStoreID(t *testing.T) {
|
||||
t.Run("UpdateKnowledgeBaseRequest", func(t *testing.T) {
|
||||
assertNoVectorStoreIDField(t, reflect.TypeOf(UpdateKnowledgeBaseRequest{}))
|
||||
})
|
||||
t.Run("KnowledgeBaseConfig", func(t *testing.T) {
|
||||
// Config carries chunking / extract / faq / wiki sub-configs and must
|
||||
// not be extended with a VectorStoreID either (Config is passed
|
||||
// straight into the service Update path).
|
||||
assertNoVectorStoreIDField(t, reflect.TypeOf(types.KnowledgeBaseConfig{}))
|
||||
})
|
||||
}
|
||||
|
||||
// assertNoVectorStoreIDField walks the visible fields of t (including embedded
|
||||
// anonymous structs) and reports any field named VectorStoreID or carrying
|
||||
// a json tag of "vector_store_id".
|
||||
func assertNoVectorStoreIDField(t *testing.T, typ reflect.Type) {
|
||||
t.Helper()
|
||||
var visit func(rt reflect.Type, path string)
|
||||
visit = func(rt reflect.Type, path string) {
|
||||
for rt.Kind() == reflect.Ptr {
|
||||
rt = rt.Elem()
|
||||
}
|
||||
if rt.Kind() != reflect.Struct {
|
||||
return
|
||||
}
|
||||
for i := 0; i < rt.NumField(); i++ {
|
||||
f := rt.Field(i)
|
||||
full := path + "." + f.Name
|
||||
if f.Name == "VectorStoreID" {
|
||||
t.Fatalf("%s declares VectorStoreID — vector_store_id is immutable post-create "+
|
||||
"and must not be accepted by update DTOs", full)
|
||||
}
|
||||
if tag := strings.Split(f.Tag.Get("json"), ",")[0]; tag == "vector_store_id" {
|
||||
t.Fatalf("%s carries json tag \"vector_store_id\" — the field is immutable "+
|
||||
"post-create and must not be accepted by update DTOs", full)
|
||||
}
|
||||
if f.Anonymous {
|
||||
visit(f.Type, full)
|
||||
}
|
||||
}
|
||||
}
|
||||
visit(typ, typ.Name())
|
||||
}
|
||||
@@ -353,6 +353,22 @@ func Warnf(c context.Context, format string, args ...interface{}) {
|
||||
addCaller(GetLogger(c), 2).Warnf(format, args...)
|
||||
}
|
||||
|
||||
// Fields aliases logrus.Fields so callers in other packages can use the
|
||||
// short form `logger.Fields{...}` without importing logrus directly.
|
||||
type Fields = logrus.Fields
|
||||
|
||||
// WarnWithFields emits a warning with structured fields. Use this for
|
||||
// audit-relevant events (cross-tenant probes, invariant violations) so that
|
||||
// log aggregators can index the tenant/resource identifiers without
|
||||
// parsing free-form text. Format-string style (Warnf) is appropriate for
|
||||
// low-stakes diagnostic messages.
|
||||
func WarnWithFields(c context.Context, fields Fields, msg string) {
|
||||
if fields == nil {
|
||||
fields = Fields{}
|
||||
}
|
||||
addCaller(GetLogger(c), 2).WithFields(fields).Warn(msg)
|
||||
}
|
||||
|
||||
// Error 输出错误级别的日志
|
||||
func Error(c context.Context, args ...interface{}) {
|
||||
addCaller(GetLogger(c), 2).Error(args...)
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/hibiken/asynq"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// KnowledgeBaseService defines the knowledge base service interface
|
||||
@@ -207,4 +208,14 @@ type KnowledgeBaseRepository interface {
|
||||
|
||||
// TogglePinKnowledgeBase toggles the pin status of a knowledge base
|
||||
TogglePinKnowledgeBase(ctx context.Context, id string, tenantID uint64) (*types.KnowledgeBase, error)
|
||||
|
||||
// CountByVectorStoreID counts active KBs bound to the given vector store
|
||||
// within a tenant scope. Accepts a *gorm.DB handle so callers can share a
|
||||
// transaction (e.g., the VectorStore delete guard's row-lock context) or
|
||||
// run standalone (pass nil → uses the repository's default db).
|
||||
//
|
||||
// The soft-delete filter is applied automatically by the gorm.DeletedAt
|
||||
// scope on KnowledgeBase; implementations MUST NOT add an explicit
|
||||
// `deleted_at IS NULL` predicate (avoids divergence with the auto-scope).
|
||||
CountByVectorStoreID(ctx context.Context, db *gorm.DB, tenantID uint64, storeID string) (int64, error)
|
||||
}
|
||||
|
||||
@@ -31,12 +31,32 @@ type VectorStoreService interface {
|
||||
// UpdateStore updates an existing vector store (name only).
|
||||
UpdateStore(ctx context.Context, store *types.VectorStore) error
|
||||
// DeleteStore deletes a vector store by tenant + id.
|
||||
// Rejects deletion when any active knowledge base is bound to the store
|
||||
// (binding guard); the caller must unbind or delete those KBs first.
|
||||
DeleteStore(ctx context.Context, tenantID uint64, id string) error
|
||||
// TestConnection tests connectivity to a vector database.
|
||||
// Returns the detected server version on success (e.g., "7.10.1"), empty string if unknown.
|
||||
TestConnection(ctx context.Context, engineType types.RetrieverEngineType, config types.ConnectionConfig) (string, error)
|
||||
// SaveDetectedVersion updates the connection_config.version for a stored vector store.
|
||||
SaveDetectedVersion(ctx context.Context, store *types.VectorStore, version string) error
|
||||
|
||||
// ResolveStoreView returns the API-safe display projection of a single
|
||||
// store ID, scoped to the given tenant. Tries DB stores first, then the
|
||||
// cached env-store set. Returns types.UnavailableStoreDisplay() when the
|
||||
// store cannot be resolved — callers should still succeed in such cases
|
||||
// (the response carries a "unavailable" status signal for the UI).
|
||||
//
|
||||
// Never returns connection credentials in any form: the StoreDisplay
|
||||
// payload carries only Name / Source / EngineType / Status.
|
||||
ResolveStoreView(ctx context.Context, tenantID uint64, storeID string) (types.StoreDisplay, error)
|
||||
|
||||
// BatchResolveStoreView resolves multiple store IDs in a single DB read
|
||||
// plus the cached env-store match. Returned map keys are the storeIDs
|
||||
// originally requested; missing IDs map to types.UnavailableStoreDisplay().
|
||||
//
|
||||
// Intended for list endpoints that need store metadata for many KBs at
|
||||
// once without incurring N+1 ResolveStoreView calls.
|
||||
BatchResolveStoreView(ctx context.Context, tenantID uint64, storeIDs []string) (map[string]types.StoreDisplay, error)
|
||||
}
|
||||
|
||||
// VectorStoreRepository defines the repository interface for VectorStore CRUD.
|
||||
|
||||
@@ -633,3 +633,64 @@ func (kb *KnowledgeBase) IsMultimodalEnabled() bool {
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// HasVectorStore reports whether the KB is bound to a DB-managed VectorStore
|
||||
// (as opposed to the tenant's env-store fallback).
|
||||
//
|
||||
// Safe to call on a nil receiver (returns false). Mirrors the convention of
|
||||
// other Is*Enabled / Capabilities accessors in this file.
|
||||
func (kb *KnowledgeBase) HasVectorStore() bool {
|
||||
return kb != nil && kb.VectorStoreID != nil && *kb.VectorStoreID != ""
|
||||
}
|
||||
|
||||
// Normalize folds the empty-string vector store id into nil so a single
|
||||
// representation reaches both the DB and the retrieve-engine factory, which
|
||||
// treats nil and `&""` as the same "no binding" signal. Idempotent and safe
|
||||
// to call repeatedly.
|
||||
//
|
||||
// Callers that accept unvalidated user input (CreateKnowledgeBase, async
|
||||
// payload decoders) should invoke this before persistence or validation.
|
||||
// Safe to call on a nil receiver (no-op).
|
||||
func (kb *KnowledgeBase) Normalize() {
|
||||
if kb == nil {
|
||||
return
|
||||
}
|
||||
if kb.VectorStoreID != nil && *kb.VectorStoreID == "" {
|
||||
kb.VectorStoreID = nil
|
||||
}
|
||||
}
|
||||
|
||||
// SharesStoreWith reports whether two knowledge bases are bound to the same
|
||||
// vector store. Both env-fallback (nil) → true; both same UUID → true;
|
||||
// otherwise false. Safe to call when either receiver or argument is nil
|
||||
// (returns true iff both are nil).
|
||||
//
|
||||
// Empty-string VectorStoreID is treated as equivalent to nil so that rows
|
||||
// persisted by callers that did not run Normalize first (raw-SQL writes,
|
||||
// external migrations, ops scripts) still compare correctly. The alternative
|
||||
// — treating `&""` as a distinct binding — would reject otherwise valid
|
||||
// CopyKnowledgeBase clones with a confusing "different vector stores" 400.
|
||||
// This normalization is read-only; it does not mutate the receivers.
|
||||
//
|
||||
// Lives on *KnowledgeBase next to HasVectorStore() so the binding semantics
|
||||
// stay co-located with the type they describe.
|
||||
func (kb *KnowledgeBase) SharesStoreWith(other *KnowledgeBase) bool {
|
||||
if kb == nil || other == nil {
|
||||
return kb == other
|
||||
}
|
||||
a, b := kb.VectorStoreID, other.VectorStoreID
|
||||
if a != nil && *a == "" {
|
||||
a = nil
|
||||
}
|
||||
if b != nil && *b == "" {
|
||||
b = nil
|
||||
}
|
||||
switch {
|
||||
case a == nil && b == nil:
|
||||
return true
|
||||
case a == nil || b == nil:
|
||||
return false
|
||||
default:
|
||||
return *a == *b
|
||||
}
|
||||
}
|
||||
|
||||
@@ -65,10 +65,11 @@ func strPtr(s string) *string { return &s }
|
||||
// VectorStoreID pointer, including the `omitempty` behavior and the three
|
||||
// distinct JSON inputs (missing field, explicit null, explicit value).
|
||||
//
|
||||
// The `empty string pointer` case is fixed here as a behavior snapshot: PR1
|
||||
// accepts it and stores &"" verbatim because validation has not been added
|
||||
// yet. A follow-up PR that introduces validation will convert this case into
|
||||
// a rejection and flip the expectation in this test accordingly.
|
||||
// The `empty string pointer` case is fixed here as a behavior snapshot:
|
||||
// raw JSON encoding accepts it and stores &"" verbatim. Service-layer
|
||||
// CreateKnowledgeBase normalizes &"" to nil before persistence
|
||||
// (see (*KnowledgeBase).Normalize), so the stored DB row is NULL — but the
|
||||
// JSON shape on the way in is whatever the client sent.
|
||||
func TestKnowledgeBase_VectorStoreID_JSON(t *testing.T) {
|
||||
t.Run("marshal nil omits field (omitempty)", func(t *testing.T) {
|
||||
kb := KnowledgeBase{ID: "kb-1", VectorStoreID: nil}
|
||||
@@ -156,3 +157,115 @@ func TestKnowledgeBase_UnmarshalJSON_WithVectorStoreID(t *testing.T) {
|
||||
// If a future change introduces such a shadow, the value above would fail to populate.
|
||||
}
|
||||
|
||||
|
||||
|
||||
// TestKnowledgeBase_HasVectorStore covers the nil-safe binding accessor.
|
||||
func TestKnowledgeBase_HasVectorStore(t *testing.T) {
|
||||
t.Run("nil receiver returns false", func(t *testing.T) {
|
||||
var kb *KnowledgeBase
|
||||
if kb.HasVectorStore() {
|
||||
t.Fatal("nil receiver should return false")
|
||||
}
|
||||
})
|
||||
t.Run("nil pointer returns false", func(t *testing.T) {
|
||||
kb := &KnowledgeBase{}
|
||||
if kb.HasVectorStore() {
|
||||
t.Fatal("nil VectorStoreID should return false")
|
||||
}
|
||||
})
|
||||
t.Run("empty string returns false", func(t *testing.T) {
|
||||
empty := ""
|
||||
kb := &KnowledgeBase{VectorStoreID: &empty}
|
||||
if kb.HasVectorStore() {
|
||||
t.Fatal("empty string VectorStoreID should return false")
|
||||
}
|
||||
})
|
||||
t.Run("non-empty returns true", func(t *testing.T) {
|
||||
s := "store-uuid"
|
||||
kb := &KnowledgeBase{VectorStoreID: &s}
|
||||
if !kb.HasVectorStore() {
|
||||
t.Fatal("non-empty VectorStoreID should return true")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestKnowledgeBase_Normalize covers the empty-string -> nil fold used to
|
||||
// keep a single representation between the create path and the factory.
|
||||
func TestKnowledgeBase_Normalize(t *testing.T) {
|
||||
t.Run("nil receiver is no-op", func(t *testing.T) {
|
||||
var kb *KnowledgeBase
|
||||
kb.Normalize() // must not panic
|
||||
})
|
||||
t.Run("empty string -> nil", func(t *testing.T) {
|
||||
empty := ""
|
||||
kb := &KnowledgeBase{VectorStoreID: &empty}
|
||||
kb.Normalize()
|
||||
if kb.VectorStoreID != nil {
|
||||
t.Fatalf("expected nil after normalize, got %v", kb.VectorStoreID)
|
||||
}
|
||||
})
|
||||
t.Run("non-empty unchanged", func(t *testing.T) {
|
||||
s := "store-uuid"
|
||||
kb := &KnowledgeBase{VectorStoreID: &s}
|
||||
kb.Normalize()
|
||||
if kb.VectorStoreID == nil || *kb.VectorStoreID != "store-uuid" {
|
||||
t.Fatalf("expected store-uuid, got %v", kb.VectorStoreID)
|
||||
}
|
||||
})
|
||||
t.Run("already nil unchanged", func(t *testing.T) {
|
||||
kb := &KnowledgeBase{}
|
||||
kb.Normalize()
|
||||
if kb.VectorStoreID != nil {
|
||||
t.Fatal("nil should stay nil")
|
||||
}
|
||||
})
|
||||
t.Run("idempotent", func(t *testing.T) {
|
||||
empty := ""
|
||||
kb := &KnowledgeBase{VectorStoreID: &empty}
|
||||
kb.Normalize()
|
||||
kb.Normalize() // second call no-op
|
||||
if kb.VectorStoreID != nil {
|
||||
t.Fatal("idempotent normalize broke")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestKnowledgeBase_SharesStoreWith covers the binding-equality helper used
|
||||
// by CopyKnowledgeBase's same-store defense. The empty-string cases are a
|
||||
// regression guard: rows persisted by callers that did not run Normalize
|
||||
// first can carry `vector_store_id = ""` instead of NULL; without the
|
||||
// normalization inside SharesStoreWith such rows would falsely fail
|
||||
// same-store clone checks.
|
||||
func TestKnowledgeBase_SharesStoreWith(t *testing.T) {
|
||||
mk := func(p *string) *KnowledgeBase { return &KnowledgeBase{VectorStoreID: p} }
|
||||
a, b := "store-A", "store-B"
|
||||
empty := ""
|
||||
tests := []struct {
|
||||
name string
|
||||
kb *KnowledgeBase
|
||||
oth *KnowledgeBase
|
||||
want bool
|
||||
}{
|
||||
{"both nil receivers", nil, nil, true},
|
||||
{"one nil receiver", nil, mk(nil), false},
|
||||
{"both nil store IDs", mk(nil), mk(nil), true},
|
||||
{"left nil store ID", mk(nil), mk(&a), false},
|
||||
{"right nil store ID", mk(&a), mk(nil), false},
|
||||
{"same store IDs", mk(&a), mk(&a), true},
|
||||
{"different store IDs", mk(&a), mk(&b), false},
|
||||
|
||||
// empty-string == nil normalization (regression guard)
|
||||
{"left empty-string vs right nil", mk(&empty), mk(nil), true},
|
||||
{"left nil vs right empty-string", mk(nil), mk(&empty), true},
|
||||
{"both empty-string", mk(&empty), mk(&empty), true},
|
||||
{"left empty-string vs right UUID", mk(&empty), mk(&a), false},
|
||||
{"left UUID vs right empty-string", mk(&a), mk(&empty), false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := tt.kb.SharesStoreWith(tt.oth); got != tt.want {
|
||||
t.Fatalf("SharesStoreWith: got %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -64,17 +64,27 @@ func (v *VectorStore) BeforeCreate(tx *gorm.DB) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// validEngineTypes defines the engine types that can be registered as VectorStore.
|
||||
// InfinityRetrieverEngineType and ElasticFaissRetrieverEngineType are legacy/experimental
|
||||
// types that do not have standalone deployable instances, so they are excluded.
|
||||
// validEngineTypes defines the engine types that can be registered as a
|
||||
// DB-managed VectorStore (i.e., persisted to the vector_stores table and
|
||||
// listed in GetVectorStoreTypes for the UI dropdown).
|
||||
//
|
||||
// Excluded engines:
|
||||
// - Infinity / ElasticFaiss — legacy/experimental, no standalone deployable instance.
|
||||
// - Postgres / SQLite — only meaningful when bound to the app's default DB
|
||||
// connection (UseDefaultConnection=true). The Postgres retriever's
|
||||
// embeddings table is a single hard-coded name with no per-store
|
||||
// partitioning, so registering a second Postgres store on the same
|
||||
// instance has no separation effect — every KB sharing this engine
|
||||
// ends up in the same physical table. These engines are still
|
||||
// reachable via env stores (RETRIEVE_DRIVER=postgres/sqlite), which
|
||||
// route through a separate code path (BuildEnvVectorStores) and do
|
||||
// not pass through this validation.
|
||||
var validEngineTypes = map[RetrieverEngineType]bool{
|
||||
PostgresRetrieverEngineType: true,
|
||||
ElasticsearchRetrieverEngineType: true,
|
||||
QdrantRetrieverEngineType: true,
|
||||
MilvusRetrieverEngineType: true,
|
||||
WeaviateRetrieverEngineType: true,
|
||||
DorisRetrieverEngineType: true,
|
||||
SQLiteRetrieverEngineType: true,
|
||||
TencentVectorDBRetrieverEngineType: true,
|
||||
}
|
||||
|
||||
@@ -484,6 +494,66 @@ func ValidateIndexConfig(ic IndexConfig) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// StoreDisplay — API-safe projection embedded in other resource responses
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Vector store source classifiers used by API responses.
|
||||
// Kept as package-level constants so handlers and services share a single
|
||||
// vocabulary instead of repeating magic strings.
|
||||
const (
|
||||
StoreSourceEnv = "env" // env-driven (RETRIEVE_DRIVER)
|
||||
StoreSourceUser = "user" // DB-managed VectorStore row
|
||||
StoreSourceShared = "shared" // cross-tenant access — metadata suppressed
|
||||
StoreSourceUnavailable = "unavailable" // bound store row missing / registry miss
|
||||
)
|
||||
|
||||
// StoreDisplay is the API-safe projection of a VectorStore for embedding in
|
||||
// other resource responses (notably KnowledgeBase). It carries only the
|
||||
// display-safe identifiers — never connection credentials.
|
||||
//
|
||||
// Source is one of the StoreSource* constants. EngineType is the underlying
|
||||
// engine name (e.g. "elasticsearch"). Status mirrors Source by default but
|
||||
// is split out to give the UI a stable boolean-like signal independent of
|
||||
// future Source value additions.
|
||||
type StoreDisplay struct {
|
||||
Name string `json:"vector_store_name,omitempty"`
|
||||
Source string `json:"vector_store_source,omitempty"`
|
||||
EngineType string `json:"vector_store_engine_type,omitempty"`
|
||||
Status string `json:"vector_store_status,omitempty"` // "available" / "unavailable"
|
||||
}
|
||||
|
||||
// DefaultStoreDisplay is the display payload for KBs that fall back to the
|
||||
// tenant's env stores (VectorStoreID == nil).
|
||||
func DefaultStoreDisplay() StoreDisplay {
|
||||
return StoreDisplay{
|
||||
Name: "System default",
|
||||
Source: StoreSourceEnv,
|
||||
Status: "available",
|
||||
}
|
||||
}
|
||||
|
||||
// UnavailableStoreDisplay is used when the bound store cannot be resolved
|
||||
// (deleted row, registry miss, transient infra error). The UI can branch on
|
||||
// Status to guide recovery (admin tool, rebind, etc.).
|
||||
func UnavailableStoreDisplay() StoreDisplay {
|
||||
return StoreDisplay{
|
||||
Source: StoreSourceUnavailable,
|
||||
Status: "unavailable",
|
||||
}
|
||||
}
|
||||
|
||||
// SharedStoreDisplay is returned for cross-tenant shared KB views so that
|
||||
// the underlying owner-tenant store's name and engine remain hidden — only
|
||||
// the fact that "this KB is shared" leaks, which is already implied by the
|
||||
// share grant itself.
|
||||
func SharedStoreDisplay() StoreDisplay {
|
||||
return StoreDisplay{
|
||||
Source: StoreSourceShared,
|
||||
Status: "available",
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// VectorStoreResponse — API response DTO
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -324,7 +324,7 @@ const testAESKey = "01234567890123456789012345678901"
|
||||
func TestVectorStore_Validate(t *testing.T) {
|
||||
valid := VectorStore{
|
||||
Name: "test-store",
|
||||
EngineType: PostgresRetrieverEngineType,
|
||||
EngineType: ElasticsearchRetrieverEngineType,
|
||||
TenantID: 1,
|
||||
}
|
||||
|
||||
@@ -380,12 +380,11 @@ func TestVectorStore_TableName(t *testing.T) {
|
||||
|
||||
func TestIsValidEngineType(t *testing.T) {
|
||||
validTypes := []RetrieverEngineType{
|
||||
PostgresRetrieverEngineType,
|
||||
ElasticsearchRetrieverEngineType,
|
||||
QdrantRetrieverEngineType,
|
||||
MilvusRetrieverEngineType,
|
||||
WeaviateRetrieverEngineType,
|
||||
SQLiteRetrieverEngineType,
|
||||
DorisRetrieverEngineType,
|
||||
TencentVectorDBRetrieverEngineType,
|
||||
}
|
||||
for _, et := range validTypes {
|
||||
@@ -394,10 +393,18 @@ func TestIsValidEngineType(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
// Postgres and SQLite are intentionally NOT registerable as DB stores —
|
||||
// they only make sense as env stores driven by RETRIEVE_DRIVER (see the
|
||||
// doc comment on validEngineTypes). UI/API surface stays consistent:
|
||||
// GetVectorStoreTypes does not list them, Validate rejects them, and
|
||||
// env stores reach the engine registry through BuildEnvVectorStores
|
||||
// instead of through CreateStore.
|
||||
invalidTypes := []RetrieverEngineType{
|
||||
"unknown",
|
||||
"opensearch",
|
||||
"",
|
||||
PostgresRetrieverEngineType,
|
||||
SQLiteRetrieverEngineType,
|
||||
InfinityRetrieverEngineType,
|
||||
ElasticFaissRetrieverEngineType,
|
||||
}
|
||||
@@ -1051,3 +1058,50 @@ func TestIndexConfig_ScalabilityFieldsRoundTrip(t *testing.T) {
|
||||
assert.NotContains(t, parsed, "desired_shard_count")
|
||||
})
|
||||
}
|
||||
|
||||
// TestVectorStore_PostgresSqliteNotRegisterable pins the write-path and
|
||||
// read-path consistency for the two engines that are only meaningful as env
|
||||
// stores. They must:
|
||||
//
|
||||
// 1. Be rejected by Validate() so POST /vector-stores returns a 4xx
|
||||
// instead of silently persisting a row that has no separation effect.
|
||||
// 2. Be absent from GetVectorStoreTypes() so the UI dropdown doesn't
|
||||
// offer them as a choice.
|
||||
//
|
||||
// Both checks live here so a future change that re-introduces one path
|
||||
// (e.g., adds Postgres back to validEngineTypes for some niche case)
|
||||
// fails this test pair instead of silently re-opening the inconsistency
|
||||
// that this fix closed.
|
||||
func TestVectorStore_PostgresSqliteNotRegisterable(t *testing.T) {
|
||||
t.Run("Validate rejects postgres as DB store", func(t *testing.T) {
|
||||
v := &VectorStore{
|
||||
Name: "test", TenantID: 1,
|
||||
EngineType: PostgresRetrieverEngineType,
|
||||
}
|
||||
err := v.Validate()
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported engine type")
|
||||
})
|
||||
|
||||
t.Run("Validate rejects sqlite as DB store", func(t *testing.T) {
|
||||
v := &VectorStore{
|
||||
Name: "test", TenantID: 1,
|
||||
EngineType: SQLiteRetrieverEngineType,
|
||||
}
|
||||
err := v.Validate()
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "unsupported engine type")
|
||||
})
|
||||
|
||||
t.Run("GetVectorStoreTypes omits postgres and sqlite", func(t *testing.T) {
|
||||
listed := GetVectorStoreTypes()
|
||||
var got []string
|
||||
for _, info := range listed {
|
||||
got = append(got, info.Type)
|
||||
}
|
||||
assert.NotContains(t, got, string(PostgresRetrieverEngineType),
|
||||
"postgres must not appear in the UI dropdown — env-store only")
|
||||
assert.NotContains(t, got, string(SQLiteRetrieverEngineType),
|
||||
"sqlite must not appear in the UI dropdown — env-store only")
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user