diff --git a/api/db/services/doc_metadata_service.py b/api/db/services/doc_metadata_service.py index 0b38643a36..f9cfefeaeb 100644 --- a/api/db/services/doc_metadata_service.py +++ b/api/db/services/doc_metadata_service.py @@ -484,7 +484,13 @@ class DocMetadataService: # Index exists - check if document exists try: - doc_exists = settings.docStoreConn.get(doc_id, index_name, [kb_id]) + # Doc-meta tables are per-tenant, not per-kb; ``kb_id`` is + # unused by ES/OB and the Infinity connector special-cases + # ``ragflow_doc_meta_`` to ignore it as well. Pass ``[]`` + # so the Infinity connector doesn't try to build a + # non-existent ``ragflow_doc_meta__`` table + # name. See also: ``_get_doc_meta_index_name`` above. + doc_exists = settings.docStoreConn.get(doc_id, index_name, []) if doc_exists: # Document exists - replace meta_fields entirely. # Using update with a `doc` body would deep-merge the meta_fields @@ -555,10 +561,15 @@ class DocMetadataService: # Try to get the metadata to confirm it exists before deleting # This is more efficient than attempting delete on non-existent records try: + # Doc-meta tables are per-tenant, not per-kb, so there is no + # kb suffix to query by. Pass an empty list — the Infinity + # connector (and OB) special-cases ``ragflow_doc_meta_`` + # indexes and queries the table directly; ES relies on the + # empty-``kb_id`` filter it already injects below. existing_metadata = settings.docStoreConn.get( doc_id, index_name, - [""], # Empty list for metadata tables + [], ) logging.debug(f"[METADATA DELETE] Get result: {existing_metadata is not None}") if not existing_metadata: @@ -698,13 +709,13 @@ class DocMetadataService: return {} # Extract fields - doc_obj = doc tenant_id = doc.knowledgebase.tenant_id - kb_id = doc_obj.kb_id index_name = cls._get_doc_meta_index_name(tenant_id) - # Try to get metadata from ES/Infinity - metadata_doc = settings.docStoreConn.get(doc_id, index_name, [kb_id]) + # Try to get metadata from ES/Infinity. Doc-meta tables are + # per-tenant, not per-kb; pass ``[]`` to avoid the Infinity + # connector building a non-existent ``_`` table name. + metadata_doc = settings.docStoreConn.get(doc_id, index_name, []) if metadata_doc: # Extract and unflatten metadata diff --git a/common/doc_store/infinity_conn_base.py b/common/doc_store/infinity_conn_base.py index 0484468860..26520a601b 100644 --- a/common/doc_store/infinity_conn_base.py +++ b/common/doc_store/infinity_conn_base.py @@ -14,24 +14,25 @@ # limitations under the License. # +import json import logging import os import random import re -import json import time from abc import abstractmethod from typing import Callable, TypeVar import infinity -from infinity.common import ConflictType -from infinity.index import IndexInfo, IndexType -from infinity.errors import ErrorCode import pandas as pd -from common.file_utils import get_project_base_directory -from rag.nlp import is_english +from infinity.common import ConflictType +from infinity.errors import ErrorCode +from infinity.index import IndexInfo, IndexType + from common import settings from common.doc_store.doc_store_base import DocStoreConnection, MatchExpr, OrderByExpr +from common.file_utils import get_project_base_directory +from rag.nlp import is_english # Concurrent CREATE/DROP TABLE on the same Infinity instance can race on # Infinity's RocksDB-backed catalog counters (e.g. ``db|1|next_table_id``). @@ -274,7 +275,7 @@ class InfinityConnectionBase(DocStoreConnection): return lst return sep.join(lst) - def equivalent_condition_to_str(self, condition: dict, table_instance=None) -> str | None: + def equivalent_condition_to_str(self, condition: dict, table_instance=None, is_delete: bool = False) -> str | None: assert "_id" not in condition columns = {} if table_instance: @@ -314,12 +315,48 @@ class InfinityConnectionBase(DocStoreConnection): "rechunked_from_chunk_ids", }: values = v if isinstance(v, list) else [v] - json_conditions = [] + # The same JSON-list columns were migrated from `varchar` to + # `json` in #17288. Pre-#17288 chunk tables in the wild still + # have these as `varchar` with a `whitespace-#` analyzer and + # store the data as a `###`-joined string (e.g. + # ``doc1###doc2``). ``json_contains`` on such a column returns + # 3030 ``json_contains(Varchar, Varchar) not found``, so fall + # back to ``filter_fulltext`` with the bare item value when + # the column is Varchar. New tables with the JSON schema use + # ``json_contains`` directly. + col_type = "" + if columns: + col_type = (columns.get(k, ("",))[0] or "").lower() + is_json_col = "json" in col_type + col_present = bool(columns) and k in columns + logger = getattr(self, "logger", None) or logging.getLogger(__name__) + if is_json_col: + logger.debug("INFINITY filter: using json_contains for JSON column %s", k) + elif col_present and "char" in col_type: + logger.debug("INFINITY filter: using filter_fulltext fallback for Varchar column %s", k) + else: + logger.debug("INFINITY filter: skipping predicate for unmapped/non-string column %s", k) + list_conditions = [] for item in values: - literal = json.dumps(item, ensure_ascii=False).replace("'", "''") - json_conditions.append(f"json_contains({k}, '{literal}')") - if json_conditions: - cond.append("(" + " or ".join(json_conditions) + ")") + if is_json_col: + # ``json_contains`` accepts any JSON-encodable value. + literal = json.dumps(item, ensure_ascii=False).replace("'", "''") + list_conditions.append(f"json_contains({k}, '{literal}')") + elif col_present and "char" in col_type and isinstance(item, str): + # Legacy Varchar column: bare item matches a token + # under the `whitespace-#` analyzer for the old + # `###`-joined encoding. Numeric / other non-string + # values were not meaningfully searchable against the + # legacy encoding, so skip them rather than emit a + # query that returns nothing. + escaped = item.replace("'", "''") + list_conditions.append(f"filter_fulltext('{self.convert_matching_field(k)}', '{escaped}')") + elif is_delete: + raise ValueError(f"Cannot build delete predicate for column '{k}' (type='{col_type}') with value {item!r}") + if list_conditions: + cond.append("(" + " or ".join(list_conditions) + ")") + elif is_delete: + raise ValueError(f"No valid delete predicate could be generated for column '{k}'") elif k in {"compile_kwd", "type_kwd", "parent_kwd"}: values = v if isinstance(v, list) else [v] exact_conditions = [] @@ -680,7 +717,11 @@ class InfinityConnectionBase(DocStoreConnection): except Exception: self.logger.warning(f"Skipped deleting from table {table_name} since the table doesn't exist.") return 0 - filter = self.equivalent_condition_to_str(condition, table_instance) + filter = self.equivalent_condition_to_str(condition, table_instance, is_delete=True) + if condition and (not filter or filter == "1=1"): + msg = f"INFINITY delete aborted: non-empty condition produced an unconstrained filter on table {table_name}." + self.logger.error(msg) + raise ValueError(msg) self.logger.debug(f"INFINITY delete table {table_name}, filter {filter}.") res = table_instance.delete(filter) return res.deleted_rows diff --git a/conf/infinity_mapping.json b/conf/infinity_mapping.json index d6bbc72aee..2ab32c2bc1 100644 --- a/conf/infinity_mapping.json +++ b/conf/infinity_mapping.json @@ -47,6 +47,7 @@ "compile_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, "plan_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, "scope_kwd": {"type": "varchar", "default": "doc", "analyzer": "whitespace-#"}, + "deleted_doc_id": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, "source_chunk_ids": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, "source_doc_ids": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, "compilation_template_ids": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, diff --git a/internal/engine/elasticsearch/chunk.go b/internal/engine/elasticsearch/chunk.go index c4dbe77b12..ae98fa1371 100644 --- a/internal/engine/elasticsearch/chunk.go +++ b/internal/engine/elasticsearch/chunk.go @@ -885,6 +885,9 @@ func (e *Engine) DeleteChunks(ctx context.Context, condition map[string]interfac // Build the query var qry map[string]interface{} if len(filterClauses) == 0 && len(mustClauses) == 0 && len(mustNotClauses) == 0 { + if len(condition) > 0 { + return 0, fmt.Errorf("ES delete aborted: non-empty condition yielded match_all query on index %s", fullIndexName) + } qry = map[string]interface{}{"match_all": map[string]interface{}{}} } else { boolMap := map[string]interface{}{} diff --git a/internal/engine/infinity/chunk.go b/internal/engine/infinity/chunk.go index 4c5264266c..6d1cecb243 100644 --- a/internal/engine/infinity/chunk.go +++ b/internal/engine/infinity/chunk.go @@ -686,6 +686,10 @@ func (e *Engine) DeleteChunks(ctx context.Context, condition map[string]interfac // Build filter from condition filter := buildFilterFromCondition(condition, clmns) + if len(condition) > 0 && (filter == "" || filter == "1=1") { + return 0, fmt.Errorf("INFINITY delete aborted: non-empty condition yielded unconstrained filter on table %s", tableName) + } + delResp, err := table.Delete(filter) if err != nil { return 0, fmt.Errorf("failed to delete: %w", err) diff --git a/internal/engine/infinity/metadata.go b/internal/engine/infinity/metadata.go index 833325e394..696c836ee0 100644 --- a/internal/engine/infinity/metadata.go +++ b/internal/engine/infinity/metadata.go @@ -386,6 +386,10 @@ func (e *Engine) deleteMetadataWithTable(table *infinity.Table, condition map[st // Build filter from condition filter := buildFilterFromCondition(condition, clmns) + if len(condition) > 0 && (filter == "" || filter == "1=1") { + return 0, fmt.Errorf("INFINITY delete aborted: non-empty condition yielded unconstrained filter") + } + delResp, err := table.Delete(filter) if err != nil { return 0, fmt.Errorf("failed to delete metadata: %w", err) diff --git a/internal/engine/infinity/sql_test.go b/internal/engine/infinity/sql_test.go index 755b124791..98993a7599 100644 --- a/internal/engine/infinity/sql_test.go +++ b/internal/engine/infinity/sql_test.go @@ -379,3 +379,24 @@ func TestLoadFieldMapping_EmptyNameDefaultsToInfinityMappingJSON(t *testing.T) { t.Errorf("empty name + no file should yield empty maps; got a2a=%v r2a=%v", a2a, r2a) } } + +func TestBuildFilterFromCondition_UnconstrainedFilter(t *testing.T) { + clmns := map[string]struct { + Type string + Default interface{} + }{ + "id": {"Varchar", ""}, + } + // empty condition yields "1=1" + if got := buildFilterFromCondition(map[string]interface{}{}, clmns); got != "1=1" { + t.Errorf("empty condition: got %q, want '1=1'", got) + } + // condition with nil or empty string values yields "1=1" + cond := map[string]interface{}{ + "source_id": "", + "nil_field": nil, + } + if got := buildFilterFromCondition(cond, clmns); got != "1=1" { + t.Errorf("non-empty condition with blank values: got %q, want '1=1'", got) + } +} diff --git a/rag/svr/task_executor_refactor/dataset_structure_merger.py b/rag/svr/task_executor_refactor/dataset_structure_merger.py index 9cdc0cfdff..ab1ee77812 100644 --- a/rag/svr/task_executor_refactor/dataset_structure_merger.py +++ b/rag/svr/task_executor_refactor/dataset_structure_merger.py @@ -42,12 +42,16 @@ import xxhash from common import settings from common.misc_utils import thread_pool_exec -from rag.nlp import search from rag.advanced_rag.knowlege_compile._common import ( encode as _encode, - tokenize_for_search as _tokenize_for_search, +) +from rag.advanced_rag.knowlege_compile._common import ( stable_row_id as _stable_row_id, ) +from rag.advanced_rag.knowlege_compile._common import ( + tokenize_for_search as _tokenize_for_search, +) +from rag.nlp import search from rag.svr.task_executor_refactor.task_context import TaskContext # --------------------------------------------------------------------------- @@ -215,14 +219,16 @@ async def _disabled_doc_ids(kb_id: str) -> set[str]: # --------------------------------------------------------------------------- # Document deletion tracking (sync, called from DocumentService.remove_document) # --------------------------------------------------------------------------- +_UPGRADED_TABLES: set[tuple[str, str]] = set() def record_doc_deletion(tenant_id: str, kb_id: str, doc_id: str) -> None: - """Record a document deletion for incremental ghost cleanup. + """Write a deletion-marker row into the chunk table. - Creates a lightweight ES meta row so that the next incremental merge build - can look up which docs were deleted and clean up orphaned dataset-level - entity rows (those whose *only* source document was the deleted doc). + The row is consumed by :func:`_cleanup_deleted_docs` during the next + incremental merge build to strip the deleted doc's ID from candidate + entity rows and purge ghost entities (those whose *only* source document + was the deleted doc). The deleted doc ID is stored in **deleted_doc_id** (not ``doc_id``) so the row survives the ``doc_id``-based chunk sweep that runs in @@ -236,6 +242,18 @@ def record_doc_deletion(tenant_id: str, kb_id: str, doc_id: str) -> None: index = search.index_name(tenant_id) if not settings.docStoreConn.index_exist(index, kb_id): return + # Chunk tables created before #17685 don't have a `deleted_doc_id` + # column; the next insert() would fail with 3013. Upgrade the table + # in place if needed. New tables pick the column up from + # conf/infinity_mapping.json automatically. + if (index, kb_id) not in _UPGRADED_TABLES: + ensure_fn = getattr(settings.docStoreConn, "ensure_columns", None) + if callable(ensure_fn): + ensure_fn( + index, + kb_id, + {"deleted_doc_id": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}}, + ) row = { "id": f"{_DELETION_META_KWD}:{kb_id}:{doc_id}", "kb_id": kb_id, @@ -245,8 +263,10 @@ def record_doc_deletion(tenant_id: str, kb_id: str, doc_id: str) -> None: "create_timestamp_flt": datetime.datetime.now().timestamp(), } settings.docStoreConn.insert([row], index, kb_id) + _UPGRADED_TABLES.add((index, kb_id)) except Exception: logging.exception("structure_merge: failed to record doc deletion kb=%s doc=%s", kb_id, doc_id) + raise # --------------------------------------------------------------------------- @@ -835,11 +855,10 @@ async def run_structure_merge(ctx: TaskContext) -> None: # Resolve embedding model (required for re-embedding dataset rows) embd_mdl = None try: - from api.db.services.llm_service import LLMBundle from api.apps.services.dataset_api_service import resolve_model_config - from common.constants import LLMType - from api.db.services.knowledgebase_service import KnowledgebaseService + from api.db.services.llm_service import LLMBundle + from common.constants import LLMType ok, kb = KnowledgebaseService.get_by_id(ctx.kb_id) if ok and kb: diff --git a/rag/utils/infinity_conn.py b/rag/utils/infinity_conn.py index 33b8cb8250..a869a5c20d 100644 --- a/rag/utils/infinity_conn.py +++ b/rag/utils/infinity_conn.py @@ -331,19 +331,37 @@ class InfinityConnection(InfinityConnectionBase): self.connPool.release_conn(inf_conn) def get(self, chunk_id: str, index_name: str, knowledgebase_ids: list[str]) -> dict | None: + # Doc-meta tables are per-tenant, not per-kb: they have no `_kb_id` + # suffix. Match the special-casing used by index_exist/insert/delete + # in InfinityConnectionBase so callers can pass either a chunk + # index (``ragflow_``) or a doc-meta index + # (``ragflow_doc_meta_``) without us logging a bogus + # "blank knowledgebase_ids" warning. + is_meta_table = index_name.startswith("ragflow_doc_meta_") + + # Validate the per-kb list BEFORE acquiring a connection — the + # blank-list case is a caller bug and shouldn't burn a connection + # from the pool. For meta tables the list is unused, so an empty + # list is fine. + if not is_meta_table: + if not knowledgebase_ids: + self.logger.warning("INFINITY get called with empty knowledgebase_ids for index %s", index_name) + return None + kb_table_names = [f"{index_name}_{kb_id}" for kb_id in knowledgebase_ids if kb_id] + if not kb_table_names: + self.logger.warning("INFINITY get has only blank knowledgebase_ids for index %s", index_name) + return None + inf_conn = self.connPool.get_conn() try: db_instance = inf_conn.get_database(self.dbName) df_list = list() assert isinstance(knowledgebase_ids, list) table_list = list() - if not knowledgebase_ids: - self.logger.warning("INFINITY get called with empty knowledgebase_ids for index %s", index_name) - return None - table_names_to_search = [f"{index_name}_{kb_id}" for kb_id in knowledgebase_ids if kb_id] - if not table_names_to_search: - self.logger.warning("INFINITY get has only blank knowledgebase_ids for index %s", index_name) - return None + if is_meta_table: + table_names_to_search = [index_name] + else: + table_names_to_search = kb_table_names for table_name in table_names_to_search: table_list.append(table_name) try: @@ -379,6 +397,58 @@ class InfinityConnection(InfinityConnectionBase): chunk["id"] = chunk_id return chunk + def ensure_columns(self, index_name: str, knowledgebase_id: str, column_defs: dict) -> None: + """Make sure the per-KB chunk table carries the given columns. + + Infinity's ``add_columns`` is idempotent for already-present columns, + so this is safe to call repeatedly. Used by callers that write new + marker rows referencing columns which were introduced after the + chunk-table schema was last updated (e.g. ``deleted_doc_id``, added + in #17685). For per-tenant doc-meta tables, pass ``knowledgebase_id`` + as the empty string. + + Logs and swallows any failure — this is a best-effort upgrade helper, + not a hard requirement of the calling write path. New tables are + created with the current ``conf/infinity_mapping.json`` schema so + callers will not need to invoke this in the steady state. + """ + if index_name.startswith("ragflow_doc_meta_"): + table_name = index_name + else: + table_name = f"{index_name}_{knowledgebase_id}" if knowledgebase_id else None + if not table_name: + return + inf_conn = self.connPool.get_conn() + try: + db_instance = inf_conn.get_database(self.dbName) + try: + table_instance = db_instance.get_table(table_name) + except InfinityException as e: + # src/common/status.cppm, kTableNotExist = 3022 + if e.error_code != ErrorCode.TABLE_NOT_EXIST: + raise + # Table doesn't exist yet — the next insert() will create it + # with the current schema, so we have nothing to upgrade. + return + existing = {n for n, *_ in table_instance.show_columns().rows()} + missing = {c: d for c, d in column_defs.items() if c not in existing} + if not missing: + return + self.logger.info( + "INFINITY adding %d missing column(s) [%s] to %s", + len(missing), + ", ".join(sorted(missing)), + table_name, + ) + table_instance.add_columns(missing) + except Exception: + self.logger.exception( + "INFINITY failed to upgrade columns on %s; the next insert() may fail", + table_name, + ) + finally: + self.connPool.release_conn(inf_conn) + def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None, refresh: str | bool = "wait_for") -> list[str]: """ # Save input to file to test inserting from file in GO diff --git a/rag/utils/ob_conn.py b/rag/utils/ob_conn.py index 3ad0d07e6f..6020c2c312 100644 --- a/rag/utils/ob_conn.py +++ b/rag/utils/ob_conn.py @@ -52,6 +52,7 @@ column_chunk_data = Column("chunk_data", JSON, nullable=True, comment="table par column_raptor_kwd = Column("raptor_kwd", String(256), nullable=True, comment="RAPTOR summary marker") column_raptor_layer_int = Column("raptor_layer_int", Integer, nullable=True, comment="RAPTOR summary layer") column_n_hop_with_weight = Column("n_hop_with_weight", LONGTEXT, nullable=True, comment="JSON-encoded n-hop neighbour paths and weights for a graph entity") +column_deleted_doc_id = Column("deleted_doc_id", String(256), nullable=True, index=True, comment="marker for incremental structure-merge ghost cleanup (#17685)") column_definitions: list[Column] = [ Column("id", String(256), primary_key=True, comment="chunk id"), @@ -98,6 +99,7 @@ column_definitions: list[Column] = [ column_order_id, column_group_id, column_mom_id, + column_deleted_doc_id, ] column_names: list[str] = [col.name for col in column_definitions] @@ -141,6 +143,7 @@ EXTRA_COLUMNS: list[Column] = [ column_raptor_kwd, column_raptor_layer_int, column_n_hop_with_weight, + column_deleted_doc_id, ] diff --git a/test/unit_test/api/db/services/test_doc_metadata_service.py b/test/unit_test/api/db/services/test_doc_metadata_service.py new file mode 100644 index 0000000000..65c4b1d785 --- /dev/null +++ b/test/unit_test/api/db/services/test_doc_metadata_service.py @@ -0,0 +1,77 @@ +# +# Copyright 2025 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Unit tests for ``DocMetadataService`` metadata lookup and deletion methods.""" + +from unittest.mock import MagicMock, patch + +import pytest + +pytestmark = pytest.mark.p2 + +from api.db.services.doc_metadata_service import DocMetadataService + + +class TestDocMetadataServiceConnectorGetCalls: + """Verify that ``DocMetadataService`` passes ``[]`` as knowledgebase_ids + when interacting with ``docStoreConn.get`` for per-tenant metadata indexes, + avoiding the creation of non-existent ``_`` table names.""" + + def test_delete_document_metadata_passes_empty_kb_list_to_get(self): + mock_doc_store = MagicMock() + mock_doc_store.index_exist.return_value = True + mock_doc_store.get.return_value = None # No metadata found + + with ( + patch("api.db.services.doc_metadata_service.settings") as mock_settings, + patch("api.db.db_models.DB.connect"), + patch("api.db.db_models.DB.connection_context"), + ): + mock_settings.docStoreConn = mock_doc_store + + res = DocMetadataService.delete_document_metadata("doc_789", "kb_123", tenant_id="tenant_456") + + assert res is True + mock_doc_store.get.assert_called_once_with( + "doc_789", + "ragflow_doc_meta_tenant_456", + [], + ) + + def test_get_document_metadata_passes_empty_kb_list_to_get(self): + mock_doc_store = MagicMock() + mock_doc_store.get.return_value = {"id": "doc_789", "meta_fields": {"author": "alice"}} + + mock_doc = MagicMock() + mock_doc.knowledgebase.tenant_id = "tenant_456" + mock_doc.kb_id = "kb_123" + + with ( + patch("api.db.services.doc_metadata_service.settings") as mock_settings, + patch("api.db.services.doc_metadata_service.Document") as mock_document_model, + patch("api.db.db_models.DB.connect"), + patch("api.db.db_models.DB.connection_context"), + ): + mock_settings.docStoreConn = mock_doc_store + mock_document_model.select.return_value.join.return_value.where.return_value.first.return_value = mock_doc + + result = DocMetadataService.get_document_metadata("doc_789") + + mock_doc_store.get.assert_called_once_with( + "doc_789", + "ragflow_doc_meta_tenant_456", + [], + ) + assert result == {"author": "alice"} diff --git a/test/unit_test/common/test_infinity_condition.py b/test/unit_test/common/test_infinity_condition.py new file mode 100644 index 0000000000..fd777e1841 --- /dev/null +++ b/test/unit_test/common/test_infinity_condition.py @@ -0,0 +1,414 @@ +# +# Copyright 2025 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Unit tests for the Infinity ``equivalent_condition_to_str`` branch that +handles the migrated JSON-list columns +(``source_doc_ids``/``source_chunk_ids``/``compilation_template_ids``/ +``doc_ids_kwd``/``entity_names_kwd``/``outlinks_kwd``/ +``related_kb_pages_kwd``/``rechunked_from_chunk_ids``). + +These columns were migrated from ``varchar`` (``whitespace-#`` analyzer, +``###``-joined encoding) to ``json`` in #17288, and then exposed through a +``json_contains`` filter in ``InfinityConnectionBase.equivalent_condition_to_str``. + +The pre-#17288 chunk tables still in the wild have these columns as +``varchar``. ``json_contains`` against a Varchar column returns +``3030 json_contains(Varchar, Varchar) not found``, so the translator must +fall back to a ``filter_fulltext`` query that matches the legacy encoding. +This module pins that behavior down (#17685). + +Run with: python -m pytest test/unit_test/common/test_infinity_condition.py -v +""" + +from unittest.mock import MagicMock, patch + +import pytest + +pytestmark = pytest.mark.p2 + + +# ``common.doc_store.infinity_conn_base`` is loaded via ``common.settings``, +# which in turn imports the rag- and memory-side Infinity connectors. We +# pre-load ``common.settings`` first so the partial-module circular import +# in the rag/memory side is already resolved by the time we reach the base +# class. +import common.settings # noqa: F401 +from rag.utils import infinity_conn as rag_infinity_conn + + +# ``InfinityConnection`` is wrapped by ``@common.decorator.singleton``, +# which replaces the class object with a factory function. Reach into the +# closure to recover the undecorated class so we can call its methods +# without dialing Infinity. +def _resolve_infinity_class(): + factory = rag_infinity_conn.InfinityConnection + closure_vars = factory.__closure__ + assert closure_vars, "singleton factory has no closure" + for cell in closure_vars: + cls = cell.cell_contents + if isinstance(cls, type): + return cls + raise RuntimeError("could not recover InfinityConnection from singleton closure") + + +_InfinityConnection = _resolve_infinity_class() + + +# --------------------------------------------------------------------------- +# Lightweight stand-in for ``infinity.remote_thrift.table.RemoteTable``. +# Captures the column metadata that ``equivalent_condition_to_str`` reads via +# ``table_instance.show_columns().rows()``. +# --------------------------------------------------------------------------- + + +class _FakeInfinityTable: + """Minimal mock of an Infinity table: returns a fixed column list.""" + + def __init__(self, columns): + # ``columns`` is a dict of name -> (type_string, default). + self._columns = dict(columns) + + def show_columns(self): + class _Resp: + def __init__(self, rows): + self._rows = rows + + def rows(self): + return self._rows + + return _Resp([(n, ty, de, "") for n, (ty, de) in self._columns.items()]) + + +# --------------------------------------------------------------------------- +# The translator is a static-ish helper that does not actually touch Infinity +# at runtime; we instantiate the base class only for ``convert_matching_field`` +# and the column-typing helpers. +# --------------------------------------------------------------------------- + + +def _translate(condition, columns): + """Run ``equivalent_condition_to_str`` against the supplied schema.""" + # ``equivalent_condition_to_str`` does not touch the connection, so we + # can skip ``__init__`` (which would otherwise try to dial Infinity via + # the singleton decorator). + return _InfinityConnection.equivalent_condition_to_str( + _InfinityConnection.__new__(_InfinityConnection), + dict(condition), + table_instance=_FakeInfinityTable(columns), + ) + + +_JSON_COLS = { + # New schema (since #17288) + "source_doc_ids": ("Json", "[]"), + "source_chunk_ids": ("Json", "[]"), + "compilation_template_ids": ("Json", "[]"), + "doc_ids_kwd": ("Json", "[]"), + "entity_names_kwd": ("Json", "[]"), + "outlinks_kwd": ("Json", "[]"), + "related_kb_pages_kwd": ("Json", "[]"), + "rechunked_from_chunk_ids": ("Json", "[]"), +} + +_VARCHAR_COLS = { + # Legacy schema (pre-#17288) — Varchar with a ``###``-joined encoding + "source_doc_ids": ("Varchar", ""), + "source_chunk_ids": ("Varchar", ""), + "compilation_template_ids": ("Varchar", ""), + "doc_ids_kwd": ("Varchar", ""), + "entity_names_kwd": ("Varchar", ""), + "outlinks_kwd": ("Varchar", ""), + "related_kb_pages_kwd": ("Varchar", ""), + "rechunked_from_chunk_ids": ("Varchar", ""), +} + + +# --------------------------------------------------------------------------- +# JSON (post-#17288) columns +# --------------------------------------------------------------------------- + + +class TestJsonColumnsUseJsonContains: + """New tables (post-#17288) have these columns as Json and must use + ``json_contains`` with a JSON-serialized literal.""" + + @pytest.mark.parametrize( + "col", + [ + "source_doc_ids", + "source_chunk_ids", + "compilation_template_ids", + "doc_ids_kwd", + "entity_names_kwd", + "outlinks_kwd", + "related_kb_pages_kwd", + "rechunked_from_chunk_ids", + ], + ) + def test_string_value_uses_json_contains(self, col): + result = _translate({col: ["doc-1"]}, {col: ("Json", "[]")}) + assert result is not None + # The JSON literal for a string is the quoted form. + assert f"json_contains({col}, '\"doc-1\"')" in result + + def test_list_of_strings_joined_with_or(self): + result = _translate( + {"source_doc_ids": ["doc-1", "doc-2"]}, + {"source_doc_ids": ("Json", "[]")}, + ) + assert result is not None + assert "json_contains(source_doc_ids, '\"doc-1\"')" in result + assert "json_contains(source_doc_ids, '\"doc-2\"')" in result + assert " or " in result + + def test_numeric_value_uses_unquoted_literal(self): + result = _translate( + {"doc_ids_kwd": [42, 99]}, + {"doc_ids_kwd": ("Json", "[]")}, + ) + assert result is not None + # ``json.dumps(42) == '42'`` (no surrounding quotes). + assert "json_contains(doc_ids_kwd, '42')" in result + assert "json_contains(doc_ids_kwd, '99')" in result + + def test_apostrophe_in_value_is_escaped(self): + result = _translate( + {"source_doc_ids": ["o'brien"]}, + {"source_doc_ids": ("Json", "[]")}, + ) + assert result is not None + # ``json.dumps("o'brien")`` -> ``"o'brien"``; the single quote inside + # is doubled to keep the surrounding SQL literal valid. + assert "json_contains(source_doc_ids, '\"o''brien\"')" in result + + +# --------------------------------------------------------------------------- +# Legacy Varchar (pre-#17288) columns +# --------------------------------------------------------------------------- + + +class TestLegacyVarcharColumnsUseFilterFulltext: + """Pre-#17288 chunk tables store these columns as Varchar with a + ``###``-joined encoding. ``json_contains`` returns + ``3030 json_contains(Varchar, Varchar) not found`` on them, so the + translator must fall back to ``filter_fulltext`` with the bare item + value (the ``whitespace-#`` analyzer tokenizes the ``###``-joined + string into the individual values).""" + + @pytest.mark.parametrize( + "col", + [ + "source_doc_ids", + "source_chunk_ids", + "compilation_template_ids", + "doc_ids_kwd", + "entity_names_kwd", + "outlinks_kwd", + "related_kb_pages_kwd", + "rechunked_from_chunk_ids", + ], + ) + def test_uses_filter_fulltext_with_bare_value(self, col): + result = _translate({col: ["doc-1"]}, {col: ("Varchar", "")}) + assert result is not None + # Bare value, NOT the JSON-serialized literal. ``filter_fulltext`` + # takes a quoted column name and a quoted value. + assert f"filter_fulltext('{col}', 'doc-1')" in result + # The buggy ``json_contains`` form must NOT be emitted. + assert "json_contains" not in result + + def test_list_of_strings_joined_with_or(self): + result = _translate( + {"source_doc_ids": ["doc-1", "doc-2"]}, + {"source_doc_ids": ("Varchar", "")}, + ) + assert result is not None + assert "filter_fulltext('source_doc_ids', 'doc-1')" in result + assert "filter_fulltext('source_doc_ids', 'doc-2')" in result + assert " or " in result + + def test_apostrophe_in_value_is_escaped(self): + result = _translate( + {"source_doc_ids": ["o'brien"]}, + {"source_doc_ids": ("Varchar", "")}, + ) + assert result is not None + assert "filter_fulltext('source_doc_ids', 'o''brien')" in result + + def test_numeric_value_is_skipped_on_legacy_varchar(self): + """Pre-#17288 the ``###``-joined encoding could not represent a + numeric value in a searchable way — emitting a query would just + return nothing, so we skip non-string items rather than emit a + query that lies to the caller.""" + result = _translate( + {"doc_ids_kwd": [42]}, + {"doc_ids_kwd": ("Varchar", "")}, + ) + # No predicate should be emitted, so the empty condition yields + # the ``1=1`` default. + assert result == "1=1" + + +# --------------------------------------------------------------------------- +# Unknown / missing columns +# --------------------------------------------------------------------------- + + +class TestUnknownColumnsAreSkipped: + """If the condition references one of the JSON-list columns but the + table doesn't carry it (or carries it under an unknown type), the + translator must skip the predicate rather than emit a query that + Infinity would reject. The remaining conditions (or ``1=1``) keep the + request valid.""" + + def test_json_list_column_missing_from_schema_is_skipped(self): + result = _translate( + { + "source_doc_ids": ["doc-1"], + # ``source_chunk_ids`` is in the JSON-list set but not in the + # supplied table schema — we cannot tell its type, so we + # skip the predicate rather than risk a + # ``json_contains(Varchar, Varchar) not found`` (#17685). + "source_chunk_ids": ["chunk-x"], + }, + {"source_doc_ids": ("Json", "[]")}, + ) + assert result is not None + assert "json_contains(source_doc_ids, '\"doc-1\"')" in result + # The unknown-type column contributes nothing. + assert "source_chunk_ids" not in result + + def test_only_unknown_type_column_yields_one_equals_one(self): + result = _translate( + {"source_doc_ids": ["doc-1"]}, + # No columns at all — we cannot tell the type, so skip. + {}, + ) + assert result == "1=1" + + def test_no_table_metadata_skips_json_predicate(self): + """``table_instance=None`` means we have no column metadata. We must + not fabricate column types — the predicate is skipped to avoid a + query that Infinity would reject.""" + + result = _InfinityConnection.equivalent_condition_to_str( + _InfinityConnection.__new__(_InfinityConnection), + {"source_doc_ids": ["doc-1"]}, + table_instance=None, + ) + # The Json predicate is gated on having seen the column as Json. With + # no metadata we skip the predicate rather than risk the legacy + # Varchar ``json_contains`` failure. + assert result == "1=1" + + +# --------------------------------------------------------------------------- +# Other behavior (smoke) +# --------------------------------------------------------------------------- + + +class TestOtherConditionBranches: + """Confirm we didn't accidentally regress the non-JSON-list branches.""" + + def test_available_int(self): + result = _translate({"available_int": 1}, {}) + assert result == "available_int=1" + + def test_compile_kwd_string(self): + result = _translate({"compile_kwd": ["entity"]}, {}) + assert result == "(compile_kwd='entity')" + + def test_compile_kwd_multi(self): + result = _translate( + {"compile_kwd": ["entity", "relation"]}, + {}, + ) + assert "compile_kwd='entity'" in result + assert "compile_kwd='relation'" in result + assert " or " in result + + def test_kb_id_varchar(self): + result = _translate({"kb_id": "kb-1"}, {"kb_id": ("Varchar", "")}) + assert result == "kb_id='kb-1'" + + +class TestDeleteSafety: + """``delete()`` must abort and raise ValueError if a non-empty condition generates + an unconstrained filter ('1=1') or unmapped predicate to prevent accidental table truncation.""" + + def test_delete_raises_when_condition_yields_unconstrained_filter(self): + inst = _InfinityConnection.__new__(_InfinityConnection) + inst.dbName = "default_db" + inst.logger = MagicMock() + inst.connPool = MagicMock() + + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + # Empty schema -> equivalent_condition_to_str yields "1=1" + table.show_columns.return_value.rows.return_value = [] + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(inst.connPool, "get_conn", return_value=inf_conn), patch.object(inst.connPool, "release_conn"): + with pytest.raises(ValueError, match="Cannot build delete predicate|unconstrained filter"): + inst.delete({"source_doc_ids": ["doc-1"]}, "ragflow_tenant", "kb-1") + + # Must NOT call table.delete() + table.delete.assert_not_called() + + def test_delete_raises_value_error_for_unmapped_delete_predicate(self): + inst = _InfinityConnection.__new__(_InfinityConnection) + inst.dbName = "default_db" + inst.logger = MagicMock() + inst.connPool = MagicMock() + + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + table.show_columns.return_value.rows.return_value = [("other_col", "Varchar", "", "")] + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(inst.connPool, "get_conn", return_value=inf_conn), patch.object(inst.connPool, "release_conn"): + with pytest.raises(ValueError, match="Cannot build delete predicate"): + inst.delete({"source_doc_ids": ["doc-1"]}, "ragflow_tenant", "kb-1") + + table.delete.assert_not_called() + + def test_legacy_varchar_chunk_table_delete_by_source_doc_ids(self): + inst = _InfinityConnection.__new__(_InfinityConnection) + inst.dbName = "default_db" + inst.logger = MagicMock() + inst.connPool = MagicMock() + + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + # Schema with legacy Varchar column + table.show_columns.return_value.rows.return_value = [("source_doc_ids", "Varchar", "", "")] + table.delete.return_value = MagicMock(deleted_rows=5) + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(inst.connPool, "get_conn", return_value=inf_conn), patch.object(inst.connPool, "release_conn"): + deleted = inst.delete({"source_doc_ids": ["doc-123"]}, "ragflow_tenant", "kb-1") + + assert deleted == 5 + table.delete.assert_called_once_with("(filter_fulltext('source_doc_ids', 'doc-123'))") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_dataset_structure_merger.py b/test/unit_test/rag/svr/task_executor_refactor/test_dataset_structure_merger.py new file mode 100644 index 0000000000..7709a48784 --- /dev/null +++ b/test/unit_test/rag/svr/task_executor_refactor/test_dataset_structure_merger.py @@ -0,0 +1,132 @@ +# +# Copyright 2025 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Unit tests for ``record_doc_deletion`` in ``dataset_structure_merger.py``.""" + +from unittest.mock import MagicMock, patch + +import pytest + +pytestmark = pytest.mark.p2 + +from rag.svr.task_executor_refactor.dataset_structure_merger import _UPGRADED_TABLES, record_doc_deletion + + +class TestRecordDocDeletion: + def setup_method(self): + _UPGRADED_TABLES.clear() + + def test_record_doc_deletion_invokes_ensure_columns_and_insert(self): + tenant_id = "tenant_123" + kb_id = "kb_456" + doc_id = "doc_789" + + mock_conn = MagicMock() + mock_conn.index_exist.return_value = True + mock_conn.ensure_columns = MagicMock() + + with patch("rag.svr.task_executor_refactor.dataset_structure_merger.settings") as mock_settings, patch("rag.svr.task_executor_refactor.dataset_structure_merger.search") as mock_search: + mock_settings.docStoreConn = mock_conn + mock_search.index_name.return_value = "ragflow_tenant_123" + + record_doc_deletion(tenant_id, kb_id, doc_id) + + mock_conn.index_exist.assert_called_once_with("ragflow_tenant_123", kb_id) + mock_conn.ensure_columns.assert_called_once_with( + "ragflow_tenant_123", + kb_id, + {"deleted_doc_id": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}}, + ) + mock_conn.insert.assert_called_once() + inserted_rows, index, kb = mock_conn.insert.call_args[0] + assert index == "ragflow_tenant_123" + assert kb == kb_id + assert len(inserted_rows) == 1 + assert inserted_rows[0]["deleted_doc_id"] == doc_id + assert inserted_rows[0]["kb_id"] == kb_id + + def test_record_doc_deletion_caches_upgraded_tables_per_kb(self): + tenant_id = "tenant_123" + kb_id1 = "kb_456" + kb_id2 = "kb_789" + + mock_conn = MagicMock() + mock_conn.index_exist.return_value = True + mock_conn.ensure_columns = MagicMock() + + with patch("rag.svr.task_executor_refactor.dataset_structure_merger.settings") as mock_settings, patch("rag.svr.task_executor_refactor.dataset_structure_merger.search") as mock_search: + mock_settings.docStoreConn = mock_conn + mock_search.index_name.return_value = "ragflow_tenant_123" + + record_doc_deletion(tenant_id, kb_id1, "doc_1") + record_doc_deletion(tenant_id, kb_id1, "doc_2") + record_doc_deletion(tenant_id, kb_id2, "doc_3") + + # ensure_columns is called twice: once for kb_456, once for kb_789 + assert mock_conn.ensure_columns.call_count == 2 + assert mock_conn.insert.call_count == 3 + inserted_doc_ids = [call.args[0][0]["deleted_doc_id"] for call in mock_conn.insert.call_args_list] + assert inserted_doc_ids == ["doc_1", "doc_2", "doc_3"] + + def test_record_doc_deletion_retries_ensure_columns_on_insert_failure(self): + tenant_id = "tenant_123" + kb_id = "kb_456" + + mock_conn = MagicMock() + mock_conn.index_exist.return_value = True + mock_conn.ensure_columns = MagicMock() + # First insert fails, second succeeds + mock_conn.insert.side_effect = [RuntimeError("insert failed"), None] + + with patch("rag.svr.task_executor_refactor.dataset_structure_merger.settings") as mock_settings, patch("rag.svr.task_executor_refactor.dataset_structure_merger.search") as mock_search: + mock_settings.docStoreConn = mock_conn + mock_search.index_name.return_value = "ragflow_tenant_123" + + with pytest.raises(RuntimeError, match="insert failed"): + record_doc_deletion(tenant_id, kb_id, "doc_1") + assert (("ragflow_tenant_123", kb_id)) not in _UPGRADED_TABLES + + record_doc_deletion(tenant_id, kb_id, "doc_2") + assert (("ragflow_tenant_123", kb_id)) in _UPGRADED_TABLES + + # ensure_columns was retried on the second attempt + assert mock_conn.ensure_columns.call_count == 2 + assert mock_conn.insert.call_count == 2 + + def test_record_doc_deletion_when_index_not_exist(self): + mock_conn = MagicMock() + mock_conn.index_exist.return_value = False + mock_conn.ensure_columns = MagicMock() + + with patch("rag.svr.task_executor_refactor.dataset_structure_merger.settings") as mock_settings, patch("rag.svr.task_executor_refactor.dataset_structure_merger.search") as mock_search: + mock_settings.docStoreConn = mock_conn + mock_search.index_name.return_value = "ragflow_tenant_123" + + record_doc_deletion("tenant_123", "kb_456", "doc_789") + + mock_conn.ensure_columns.assert_not_called() + mock_conn.insert.assert_not_called() + + def test_record_doc_deletion_when_ensure_columns_missing(self): + mock_conn = MagicMock(spec=["index_exist", "insert"]) + mock_conn.index_exist.return_value = True + + with patch("rag.svr.task_executor_refactor.dataset_structure_merger.settings") as mock_settings, patch("rag.svr.task_executor_refactor.dataset_structure_merger.search") as mock_search: + mock_settings.docStoreConn = mock_conn + mock_search.index_name.return_value = "ragflow_tenant_123" + + record_doc_deletion("tenant_123", "kb_456", "doc_789") + + mock_conn.insert.assert_called_once() diff --git a/test/unit_test/rag/utils/test_infinity_conn_helpers.py b/test/unit_test/rag/utils/test_infinity_conn_helpers.py new file mode 100644 index 0000000000..1c7f368356 --- /dev/null +++ b/test/unit_test/rag/utils/test_infinity_conn_helpers.py @@ -0,0 +1,305 @@ +# +# Copyright 2025 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Unit tests for the two ``InfinityConnection`` helpers added in #17685: + +- ``get`` special-cases the ``ragflow_doc_meta_`` index family so callers + can pass an empty ``knowledgebase_ids`` (the meta table is per-tenant and + has no ``_kb_id`` suffix). +- ``ensure_columns`` upgrades pre-existing chunk tables in place to add + columns that were introduced after the table was created (e.g. + ``deleted_doc_id``). + +The methods need a live Infinity connection, so we patch the connection +pool and exercise the routing logic. Run with:: + + python -m pytest test/unit_test/rag/utils/test_infinity_conn_helpers.py -v +""" + +import logging +from unittest.mock import MagicMock, patch + +import pytest + +pytestmark = pytest.mark.p2 + + +import common.settings # noqa: F401 -- see test_infinity_condition for the why +from rag.utils import infinity_conn as rag_infinity_conn + + +def _resolve_infinity_class(): + """The class is wrapped by ``@common.decorator.singleton``; recover the + underlying class object from the wrapper's closure (same trick as in + test_infinity_condition).""" + factory = rag_infinity_conn.InfinityConnection + for cell in factory.__closure__ or (): + cls = cell.cell_contents + if isinstance(cls, type): + return cls + raise RuntimeError("could not recover InfinityConnection from singleton closure") + + +_Inf = _resolve_infinity_class() + + +# --------------------------------------------------------------------------- +# ``InfinityConnection.get`` — doc-meta special case +# --------------------------------------------------------------------------- + + +class TestGetHandlesDocMetaIndex: + """``get`` must treat ``ragflow_doc_meta_`` indexes as per-tenant, not + per-kb. The pre-#17685 code built a ``_`` table name and + either logged a "blank knowledgebase_ids" warning (when ``[""]`` was + passed) or a "Table not found" warning (when a real ``kb_id`` was + passed), without ever actually querying the table.""" + + def _new_conn(self): + """Build a bare ``InfinityConnection`` that does not touch the real + connection pool. ``__new__`` skips ``__init__``; we hand-roll the + attributes the methods under test read.""" + inst = _Inf.__new__(_Inf) + inst.dbName = "default_db" + inst.logger = logging.getLogger("test.infinity_conn_helpers") + inst.connPool = MagicMock() + return inst + + def test_meta_index_with_empty_kb_ids_uses_index_name_as_table(self): + conn = self._new_conn() + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + empty_df = MagicMock() + empty_df.empty = True + empty_df.columns.tolist.return_value = ["id"] + table.output.return_value.filter.return_value.to_df.return_value = (empty_df, 0) + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(conn.connPool, "get_conn", return_value=inf_conn), patch.object(conn.connPool, "release_conn") as release: + # ``[""]`` (or any list of blanks) used to trigger the + # "blank knowledgebase_ids" warning. The fixed code treats + # meta tables by index name and silently returns ``None`` for + # the (now nonexistent) row. + result = conn.get( + "doc-1", + "ragflow_doc_meta_tenant-A", + [""], + ) + + assert result is None + # The table is queried by index name, not by ``index_name + "_" + kb_id``. + db.get_table.assert_called_once_with("ragflow_doc_meta_tenant-A") + release.assert_called_once_with(inf_conn) + + def test_meta_index_with_real_kb_id_still_uses_index_name(self): + """Existing callers that pass ``[kb_id]`` for a meta index used to + build the non-existent ``_`` table and log a warning. + The fix routes them to the meta table directly.""" + conn = self._new_conn() + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + empty_df = MagicMock() + empty_df.empty = True + empty_df.columns.tolist.return_value = ["id"] + table.output.return_value.filter.return_value.to_df.return_value = (empty_df, 0) + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(conn.connPool, "get_conn", return_value=inf_conn): + conn.get("doc-1", "ragflow_doc_meta_tenant-A", ["kb-1"]) + + db.get_table.assert_called_once_with("ragflow_doc_meta_tenant-A") + + def test_non_meta_index_with_empty_kb_ids_still_returns_none(self): + """The pre-existing behavior for chunk indexes is unchanged: an + empty ``knowledgebase_ids`` is a programmer error and the call + short-circuits to ``None``.""" + conn = self._new_conn() + inf_conn = MagicMock() + with patch.object(conn.connPool, "get_conn", return_value=inf_conn): + result = conn.get("chunk-1", "ragflow_tenant-A", []) + + assert result is None + inf_conn.get_database.assert_not_called() + + def test_non_meta_index_with_blank_kb_ids_still_returns_none(self): + conn = self._new_conn() + inf_conn = MagicMock() + with patch.object(conn.connPool, "get_conn", return_value=inf_conn): + result = conn.get("chunk-1", "ragflow_tenant-A", ["", ""]) + + assert result is None + inf_conn.get_database.assert_not_called() + + +# --------------------------------------------------------------------------- +# ``InfinityConnection.ensure_columns`` +# --------------------------------------------------------------------------- + + +class TestEnsureColumns: + """``ensure_columns`` upgrades chunk tables in place with columns that + are not yet present (e.g. ``deleted_doc_id`` from #17685). The method + is idempotent and silent on already-present columns.""" + + def _new_conn(self): + inst = _Inf.__new__(_Inf) + inst.dbName = "default_db" + inst.logger = logging.getLogger("test.infinity_conn_helpers") + inst.connPool = MagicMock() + return inst + + def test_adds_missing_column_via_add_columns(self): + conn = self._new_conn() + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + # ``deleted_doc_id`` is missing; ``kb_id`` is present. + table.show_columns.return_value.rows.return_value = [ + ("id", "Varchar", "", ""), + ("kb_id", "Varchar", "", ""), + ] + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(conn.connPool, "get_conn", return_value=inf_conn), patch.object(conn.connPool, "release_conn") as release: + conn.ensure_columns( + "ragflow_tenant-A", + "kb-1", + {"deleted_doc_id": {"type": "varchar", "default": ""}}, + ) + + # Only the missing column is passed to ``add_columns``. + table.add_columns.assert_called_once_with({"deleted_doc_id": {"type": "varchar", "default": ""}}) + release.assert_called_once_with(inf_conn) + + def test_skips_when_all_columns_present(self): + conn = self._new_conn() + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + table.show_columns.return_value.rows.return_value = [ + ("id", "Varchar", "", ""), + ("deleted_doc_id", "Varchar", "", ""), + ] + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(conn.connPool, "get_conn", return_value=inf_conn), patch.object(conn.connPool, "release_conn") as release: + conn.ensure_columns( + "ragflow_tenant-A", + "kb-1", + {"deleted_doc_id": {"type": "varchar", "default": ""}}, + ) + + table.add_columns.assert_not_called() + release.assert_called_once_with(inf_conn) + + def test_silent_when_table_missing(self): + """A missing table means the next ``insert()`` will create it with + the current schema; ``ensure_columns`` must not interfere.""" + from infinity.common import InfinityException + + conn = self._new_conn() + inf_conn = MagicMock() + db = MagicMock() + db.get_table.side_effect = InfinityException(3022, "table missing") + inf_conn.get_database.return_value = db + + with patch.object(conn.connPool, "get_conn", return_value=inf_conn), patch.object(conn.connPool, "release_conn") as release: + conn.ensure_columns( + "ragflow_tenant-A", + "kb-1", + {"deleted_doc_id": {"type": "varchar", "default": ""}}, + ) + + release.assert_called_once_with(inf_conn) + + def test_logs_exception_when_other_infinity_error(self): + """Non-TABLE_NOT_EXIST Infinity exceptions are re-raised internally + and caught/logged by the outer exception handler.""" + from infinity.common import InfinityException + + conn = self._new_conn() + inf_conn = MagicMock() + db = MagicMock() + db.get_table.side_effect = InfinityException(3000, "catalog corrupted") + inf_conn.get_database.return_value = db + + with patch.object(conn.connPool, "get_conn", return_value=inf_conn), patch.object(conn.connPool, "release_conn") as release: + conn.ensure_columns( + "ragflow_tenant-A", + "kb-1", + {"deleted_doc_id": {"type": "varchar", "default": ""}}, + ) + + release.assert_called_once_with(inf_conn) + + def test_meta_table_uses_index_name_directly(self): + conn = self._new_conn() + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + table.show_columns.return_value.rows.return_value = [ + ("id", "Varchar", "", ""), + ("kb_id", "Varchar", "", ""), + ("meta_fields", "Json", "{}", ""), + ] + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(conn.connPool, "get_conn", return_value=inf_conn): + # Doc-meta tables pass an empty ``knowledgebase_id``; the helper + # should still resolve the table by index name only. + conn.ensure_columns( + "ragflow_doc_meta_tenant-A", + "", + {"new_col": {"type": "varchar", "default": ""}}, + ) + + db.get_table.assert_called_once_with("ragflow_doc_meta_tenant-A") + table.add_columns.assert_called_once_with({"new_col": {"type": "varchar", "default": ""}}) + + def test_swallows_add_columns_failure(self): + """``add_columns`` is best-effort; a failure must not propagate so + the caller's own write path can still proceed (and log a + structured warning via the connector).""" + conn = self._new_conn() + inf_conn = MagicMock() + db = MagicMock() + table = MagicMock() + table.show_columns.return_value.rows.return_value = [ + ("id", "Varchar", "", ""), + ] + table.add_columns.side_effect = RuntimeError("boom") + db.get_table.return_value = table + inf_conn.get_database.return_value = db + + with patch.object(conn.connPool, "get_conn", return_value=inf_conn), patch.object(conn.connPool, "release_conn") as release: + # Must not raise. + conn.ensure_columns( + "ragflow_tenant-A", + "kb-1", + {"deleted_doc_id": {"type": "varchar", "default": ""}}, + ) + + release.assert_called_once_with(inf_conn) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])