From 8b7a956aca4400dcb3a11612cd1d378c9154b7b0 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 22 Apr 2026 10:49:26 +0800 Subject: [PATCH 01/84] =?UTF-8?q?pyseekdb=20collection=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E6=8C=89namespace=E9=9A=94=E7=A6=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/__init__.py | 8 + src/pyseekdb/client/__init__.py | 6 + src/pyseekdb/client/admin_client.py | 4 + src/pyseekdb/client/client_base.py | 1377 ++++++++++++++++- src/pyseekdb/client/client_seekdb_embedded.py | 10 + src/pyseekdb/client/client_seekdb_server.py | 11 + src/pyseekdb/client/collection.py | 100 +- src/pyseekdb/client/configuration.py | 62 + src/pyseekdb/client/meta_info.py | 48 + src/pyseekdb/client/namespace.py | 236 +++ src/pyseekdb/client/schema.py | 12 +- src/pyseekdb/client/validators.py | 43 + src/pyseekdb/schema.py | 27 + tests/integration_tests/conftest.py | 9 + tests/integration_tests/test_namespace_dml.py | 167 ++ .../test_namespace_lifecycle.py | 205 +++ .../test_namespace_prewarm.py | 61 + .../integration_tests/test_namespace_query.py | 290 ++++ tests/unit_tests/test_namespace.py | 958 ++++++++++++ 19 files changed, 3599 insertions(+), 35 deletions(-) create mode 100644 src/pyseekdb/client/namespace.py create mode 100644 src/pyseekdb/client/validators.py create mode 100644 src/pyseekdb/schema.py create mode 100644 tests/integration_tests/test_namespace_dml.py create mode 100644 tests/integration_tests/test_namespace_lifecycle.py create mode 100644 tests/integration_tests/test_namespace_prewarm.py create mode 100644 tests/integration_tests/test_namespace_query.py create mode 100644 tests/unit_tests/test_namespace.py diff --git a/src/pyseekdb/__init__.py b/src/pyseekdb/__init__.py index e16c45d9..1e2a33af 100644 --- a/src/pyseekdb/__init__.py +++ b/src/pyseekdb/__init__.py @@ -86,6 +86,9 @@ HNSWConfiguration, IKMode, IKProperties, + IVFConfiguration, + IVFIndexLib, + IVFIndexType, K, Ngram2Properties, NgramProperties, @@ -103,6 +106,7 @@ register_sparse_embedding_function, ) from .client.collection import Collection +from .client.namespace import Namespace try: __version__ = importlib.metadata.version("pyseekdb") @@ -128,7 +132,11 @@ "HNSWConfiguration", "IKMode", "IKProperties", + "IVFConfiguration", + "IVFIndexLib", + "IVFIndexType", "K", + "Namespace", "Ngram2Properties", "NgramProperties", "RemoteServerClient", diff --git a/src/pyseekdb/client/__init__.py b/src/pyseekdb/client/__init__.py index 0e6ffef1..1ce0aaef 100644 --- a/src/pyseekdb/client/__init__.py +++ b/src/pyseekdb/client/__init__.py @@ -31,6 +31,9 @@ HNSWConfiguration, IKMode, IKProperties, + IVFConfiguration, + IVFIndexLib, + IVFIndexType, Ngram2Properties, NgramProperties, SpaceProperties, @@ -167,6 +170,9 @@ def __getattr__(name: str) -> Any: "HNSWConfiguration", "IKMode", "IKProperties", + "IVFConfiguration", + "IVFIndexLib", + "IVFIndexType", "K", "Ngram2Properties", "NgramProperties", diff --git a/src/pyseekdb/client/admin_client.py b/src/pyseekdb/client/admin_client.py index 9c18ff90..4cdb46ae 100644 --- a/src/pyseekdb/client/admin_client.py +++ b/src/pyseekdb/client/admin_client.py @@ -194,6 +194,7 @@ def create_collection( schema: SchemaParam = None, configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, + use_namespace: bool = False, **kwargs, ) -> "Collection": """Proxy to server implementation - collection operations only""" @@ -202,6 +203,7 @@ def create_collection( schema=schema, configuration=configuration, embedding_function=embedding_function, + use_namespace=use_namespace, **kwargs, ) @@ -227,6 +229,7 @@ def get_or_create_collection( schema: SchemaParam = None, configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, + use_namespace: bool = False, **kwargs, ) -> "Collection": """Proxy to server implementation - collection operations only""" @@ -235,6 +238,7 @@ def get_or_create_collection( schema=schema, configuration=configuration, embedding_function=embedding_function, + use_namespace=use_namespace, **kwargs, ) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index b9c104ba..8c3c31c3 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -5,6 +5,7 @@ import contextlib import json import logging +import os import re import struct import warnings @@ -25,6 +26,7 @@ ConfigurationParam, FulltextIndexConfig, HNSWConfiguration, + IVFConfiguration, VectorIndexConfig, ) from .database import Database @@ -37,7 +39,7 @@ get_default_embedding_function, ) from .filters import FilterBuilder -from .meta_info import CollectionFieldNames, CollectionNames +from .meta_info import CollectionFieldNames, CollectionNames, NamespaceCollectionNames, NamespaceFieldNames from .query_types import QueryHint from .schema import Schema, SparseVectorIndexConfig from .sparse_embedding_function import ( @@ -120,6 +122,27 @@ def _validate_collection_name(name: str) -> None: ) +from .validators import _MAX_NAMESPACE_BATCH_SIZE, _validate_namespace_name, _validate_record_ids # noqa: F401 + +_NS_PARTITION_COUNT = 1000 + + +def _build_default_ltable_schema() -> dict: + return { + "col_info": [ + {"col_idx": 1, "col_name": "metadata", "col_type": "JSON"}, + {"col_idx": 2, "col_name": "content", "col_type": "TEXT"}, + {"col_idx": 3, "col_name": "embedding", "col_type": "VECTOR"}, + ], + "index_info": [ + {"index_seq": 0, "index_type": "PRIMARY", "indexed_columns": []}, + {"index_seq": 1, "index_type": "SEARCH_INDEX", "indexed_columns": [1]}, + {"index_seq": 2, "index_type": "FULLTEXT", "indexed_columns": [2]}, + {"index_seq": 3, "index_type": "IVF", "indexed_columns": [3]}, + ], + } + + def _get_fulltext_index_sql( fulltext_config: FulltextIndexConfig | None = None, ) -> str: @@ -189,6 +212,21 @@ def _get_vector_index_sql(hnsw_config: HNSWConfiguration) -> str: return f"WITH (DISTANCE={hnsw_config.distance}, TYPE={hnsw_config.type}, LIB={hnsw_config.lib}{properties_str})" +def _get_ivf_vector_index_sql(ivf_config: "IVFConfiguration") -> str: + property_parts = [] + if ivf_config.properties: + for k, v in ivf_config.properties.items(): + if isinstance(v, str): + property_parts.append(f"{k}='{v}'") + else: + property_parts.append(f"{k}={v}") + if ivf_config.use_spfresh is not None: + property_parts.append(f"use_spfresh={str(ivf_config.use_spfresh).lower()}") + property_str = ", ".join(property_parts) + properties_str = f", {property_str}" if property_str else "" + return f"WITH (DISTANCE={ivf_config.distance}, TYPE={ivf_config.type.upper()}, LIB={ivf_config.lib.upper()}{properties_str})" + + def _get_sparse_vector_index_sql(sparse_config: SparseVectorIndexConfig) -> str: """ Generate VECTOR INDEX SQL clause for sparse vector index from SparseVectorIndexConfig. @@ -252,6 +290,7 @@ def create_collection( schema: Schema | None = None, configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, + use_namespace: bool = False, **kwargs, ) -> "Collection": """ @@ -270,6 +309,7 @@ def create_collection( Defaults to DefaultEmbeddingFunction. If explicitly set to None, collection will not have an embedding function. Ignored if ``schema`` is provided. + use_namespace: If True, create a namespace-enabled collection. Defaults to False. **kwargs: Additional parameters """ pass @@ -357,6 +397,23 @@ class BaseClient(BaseConnection, AdminAPI): # ==================== Database Type Detection ==================== + def _validate_ob_database_type(self) -> None: + db_type, _version = self.detect_db_type_and_version() + if db_type.lower() != "oceanbase": + raise ValueError("use_namespace=True is only supported on OceanBase") + + def _is_shared_storage_mode(self) -> bool: + try: + rows = self._execute( + "SELECT VALUE FROM oceanbase.GV$OB_PARAMETERS WHERE name = 'ob_startup_mode'" + ) + if rows: + val = rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["VALUE"] + return str(val).upper() == "SHARED_STORAGE" + except Exception: + pass + return False + def detect_db_type_and_version(self) -> tuple[str, "Version"]: # noqa: C901 """ Detect database type and version. @@ -708,8 +765,7 @@ def _prepare_schema_parameters( # noqa: C901 hnsw_config.dimension = dimension fulltext_config = _extract_fulltext_config(configuration) - vic = VectorIndexConfig(hnsw=hnsw_config) - vic.embedding_function = embedding_function + vic = VectorIndexConfig(hnsw=hnsw_config, embedding_function=embedding_function) return Schema( vector_index=vic, fulltext_index=fulltext_config, @@ -721,6 +777,7 @@ def create_collection( schema: Schema | None = None, configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, + use_namespace: bool = False, **kwargs, ) -> "Collection": """Create a new collection. @@ -736,6 +793,7 @@ def create_collection( embedding_function: The embedding function to use for this collection. Defaults to ``DefaultEmbeddingFunction`` (all-MiniLM-L6-v2). If set to None, no embedding function will be used (embeddings must be provided manually). + use_namespace: If True, create a namespace-enabled collection. Defaults to False. **kwargs: Additional parameters for collection creation. Returns: @@ -784,6 +842,9 @@ def create_collection( logger.debug(f"schema: {schema}") + if use_namespace: + return self._create_namespace_collection(name, schema, **kwargs) + # Resolve HNSW configuration dimension if not set hnsw_config = schema.vector_index.hnsw dense_embedding_function = schema.vector_index.embedding_function @@ -857,6 +918,71 @@ def create_collection( **kwargs, ) + def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> "Collection": + dense_embedding_function = schema.vector_index.embedding_function + ivf_config = schema.vector_index.ivf + hnsw_config = schema.vector_index.hnsw + + if hnsw_config is not None: + raise ValueError("use_namespace=True only supports IVF index type, HNSW is not allowed") + self._validate_ob_database_type() + + if ivf_config is not None: + dimension = ivf_config.dimension + distance = ivf_config.distance + if ivf_config.use_spfresh is not None and not ivf_config.use_spfresh: + raise ValueError("use_namespace=True requires use_spfresh=True") + else: + if dense_embedding_function is not None: + dimension = self._get_embedding_function_dimension(dense_embedding_function) + else: + dimension = DEFAULT_VECTOR_DIMENSION + distance = DEFAULT_DISTANCE_METRIC + from .configuration import IVFConfiguration + ivf_config = IVFConfiguration(dimension=dimension, distance=distance, use_spfresh=True) + + if dimension < 1 or dimension > 4096: + raise ValueError(f"Dimension must be between 1 and 4096, got {dimension}") + + is_ss = self._is_shared_storage_mode() + settings = { + "version": 2, + "use_namespace": True, + "dense_index_type": "ivf", + "use_spfresh": True, + "storage_mode": "ss" if is_ss else "sn", + "dimension": dimension, + "distance": distance, + } + if dense_embedding_function is not None and EmbeddingFunction.support_persistence(dense_embedding_function): + settings["embedding_function"] = { + "name": dense_embedding_function.name(), + "properties": dense_embedding_function.get_config(), + } + + collection_meta = self._create_ns_collection_meta(name, settings) + collection_id = collection_meta["collection_id"] + + self._ensure_namespace_catalogs() + + self._create_namespace_physical_tables( + collection_id=collection_id, + dimension=dimension, + ivf_config=ivf_config, + fulltext_config=schema.fulltext_index, + is_shared_storage=is_ss, + ) + + return Collection( + client=self, + name=name, + collection_id=collection_id, + dimension=dimension, + embedding_function=dense_embedding_function, + distance=distance, + use_namespace=True, + ) + def _get_embedding_function_dimension(self, embedding_function: EmbeddingFunction) -> int: """Get the dimension from an embedding function.""" try: @@ -946,7 +1072,305 @@ def _create_collection_meta_v1(self, collection_name: str) -> str: except Exception as e: raise ValueError(f"Failed to create collection metadata: {e}") from e + # ==================== Namespace Catalog Methods ==================== + + def _ensure_namespace_catalogs(self) -> None: + ns_namespaces_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` ( + namespace_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, + collection_id CHAR(32) NOT NULL, + namespace_name VARCHAR(256) NOT NULL, + created_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + updated_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6), + info JSON, + PRIMARY KEY (namespace_id), + UNIQUE KEY uk_sdk_ns_coll_name (collection_id, namespace_name), + KEY idx_sdk_ns_by_collection (collection_id) + ) COMMENT='Namespace catalog';""" + ns_ltables_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_ns_ltables_table()}` ( + ltable_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, + collection_id CHAR(32) NOT NULL, + namespace_id BIGINT UNSIGNED NOT NULL, + ltable_name VARCHAR(256) NOT NULL DEFAULT 'default', + created_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + updated_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6), + info JSON, + PRIMARY KEY (ltable_id), + UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name), + KEY idx_sdk_lt_by_ns (collection_id, namespace_id) + ) COMMENT='LTable catalog';""" + self._execute(ns_namespaces_sql) + self._execute(ns_ltables_sql) + + def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> dict: + self._create_sdk_collections_if_not_exists() + settings_str = escape_string(json.dumps(settings)) + collection_name_escaped = escape_string(collection_name) + insert_sql = ( + f"INSERT INTO `{CollectionNames.sdk_collections_table_name()}` " + f"(collection_name, settings) " + f"VALUES ('{collection_name_escaped}', '{settings_str}')" + ) + self._execute(insert_sql) + rows = self._execute( + f"SELECT collection_id FROM `{CollectionNames.sdk_collections_table_name()}` " + f"WHERE collection_name = '{collection_name_escaped}'" + ) + collection_id = str(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["collection_id"]) + return {"collection_id": collection_id, "collection_name": collection_name} + + def _get_ns_collection_meta(self, collection_name: str) -> dict | None: + collection_name_escaped = escape_string(collection_name) + try: + rows = self._execute( + f"SELECT collection_id, collection_name, settings " + f"FROM `{CollectionNames.sdk_collections_table_name()}` " + f"WHERE collection_name = '{collection_name_escaped}'" + ) + except Exception: + return None + if not rows: + return None + row = rows[0] + if isinstance(row, (list, tuple)): + settings = json.loads(row[2]) if row[2] else {} + else: + settings = json.loads(row["settings"]) if row.get("settings") else {} + if not settings.get("use_namespace"): + return None + if isinstance(row, (list, tuple)): + return { + "collection_id": str(row[0]), + "collection_name": row[1], + "settings": settings, + } + return { + "collection_id": str(row["collection_id"]), + "collection_name": row["collection_name"], + "settings": settings, + } + + def _has_ns_collection(self, collection_name: str) -> bool: + try: + return self._get_ns_collection_meta(collection_name) is not None + except Exception: + return False + + def _delete_ns_collection_meta(self, collection_name: str) -> None: + meta = self._get_ns_collection_meta(collection_name) + if meta is None: + raise ValueError(f"Namespace collection '{collection_name}' not found") + collection_id = meta["collection_id"] + collection_id_escaped = escape_string(collection_id) + self._execute( + f"DELETE FROM `{CollectionNames.sdk_collections_table_name()}` " + f"WHERE collection_id = '{collection_id_escaped}'" + ) + with contextlib.suppress(Exception): + self._execute( + f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"WHERE collection_id = '{collection_id_escaped}'" + ) + with contextlib.suppress(Exception): + self._execute( + f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"WHERE collection_id = '{collection_id_escaped}'" + ) + self._cleanup_namespace_physical_tables(collection_id) + + def _create_namespace_physical_tables( + self, + collection_id: str, + dimension: int, + ivf_config, + fulltext_config=None, + is_shared_storage: bool = False, + ) -> None: + tg_name = NamespaceCollectionNames.tablegroup_name(collection_id) + data_table = NamespaceCollectionNames.data_table_name(collection_id) + kv_table = NamespaceCollectionNames.kv_data_table_name(collection_id) + schema_table = NamespaceCollectionNames.logic_schema_table_name(collection_id) + + fulltext_clause = _get_fulltext_index_sql(fulltext_config) + vector_index_sql = _get_ivf_vector_index_sql(ivf_config) + partition_clause = f"PARTITION BY KEY(namespace_id) PARTITIONS {_NS_PARTITION_COUNT}" + + try: + self._execute(f"CREATE TABLEGROUP `{tg_name}` SHARDING='ADAPTIVE'") + + data_sql_with_search = f"""CREATE TABLE `{data_table}` ( + namespace_id BIGINT UNSIGNED NOT NULL, + ltable_id BIGINT UNSIGNED NOT NULL, + document LONGTEXT, + embedding VECTOR({dimension}), + data_content JSON NOT NULL, + created_by VARCHAR(64) DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + FULLTEXT INDEX idx_fts(document) {fulltext_clause}, + SEARCH INDEX idx_json(data_content) WITH PARSER json + ) TABLEGROUP=`{tg_name}` COMMENT='逻辑表主数据' DEFAULT CHARSET=utf8mb4 ORGANIZATION HEAP IS_LOGIC_TABLE = TRUE + {partition_clause}""" + + data_sql_without_search = f"""CREATE TABLE `{data_table}` ( + namespace_id BIGINT UNSIGNED NOT NULL, + ltable_id BIGINT UNSIGNED NOT NULL, + document LONGTEXT, + embedding VECTOR({dimension}), + data_content JSON NOT NULL, + created_by VARCHAR(64) DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + FULLTEXT INDEX idx_fts(document) {fulltext_clause} + ) TABLEGROUP=`{tg_name}` COMMENT='逻辑表主数据' DEFAULT CHARSET=utf8mb4 ORGANIZATION HEAP IS_LOGIC_TABLE = TRUE + {partition_clause}""" + + try: + self._execute(data_sql_with_search) + except Exception: + self._execute(data_sql_without_search) + + ivf_index_sql = f"CREATE VECTOR INDEX idx_vec ON `{data_table}` (embedding) {vector_index_sql}" + self._execute(ivf_index_sql) + + if is_shared_storage: + hot_table = NamespaceCollectionNames.hot_table_name(collection_id) + self._execute(f"""CREATE TABLE `{hot_table}` ( + namespace_id BIGINT UNSIGNED NOT NULL, + last_access_time TIMESTAMP(6) NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + PRIMARY KEY(namespace_id) + ) TABLEGROUP=`{tg_name}` COMMENT='热点/TTL附属表' DEFAULT CHARSET=utf8mb4 + {partition_clause}""") + + self._execute(f"""CREATE TABLE `{kv_table}` ( + namespace_id BIGINT UNSIGNED NOT NULL, + kv_key VARCHAR(1024) NOT NULL, + kv_value LONGBLOB NOT NULL, + PRIMARY KEY(namespace_id, kv_key) + ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 + {partition_clause}""") + + self._execute(f"""CREATE TABLE `{schema_table}` ( + namespace_id BIGINT UNSIGNED NOT NULL, + ltable_id BIGINT UNSIGNED NOT NULL, + schema_content JSON NOT NULL, + created_by VARCHAR(64) DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + PRIMARY KEY(namespace_id, ltable_id) + ) TABLEGROUP=`{tg_name}` COMMENT='LTable schema定义' DEFAULT CHARSET=utf8mb4 + {partition_clause}""") + + except Exception: + self._cleanup_namespace_physical_tables(collection_id) + raise + + def _cleanup_namespace_physical_tables(self, collection_id: str) -> None: + for suffix_fn in [ + NamespaceCollectionNames.logic_schema_table_name, + NamespaceCollectionNames.kv_data_table_name, + NamespaceCollectionNames.hot_table_name, + NamespaceCollectionNames.data_table_name, + ]: + with contextlib.suppress(Exception): + self._execute(f"DROP TABLE IF EXISTS `{suffix_fn(collection_id)}`") + with contextlib.suppress(Exception): + self._execute(f"DROP TABLEGROUP IF EXISTS `{NamespaceCollectionNames.tablegroup_name(collection_id)}`") + + def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: + namespace_name_escaped = escape_string(namespace_name) + collection_id_escaped = escape_string(collection_id) + self._execute( + f"INSERT INTO `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"(collection_id, namespace_name) VALUES ('{collection_id_escaped}', '{namespace_name_escaped}')" + ) + rows = self._execute( + f"SELECT namespace_id FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"WHERE collection_id = '{collection_id_escaped}' AND namespace_name = '{namespace_name_escaped}'" + ) + ns_id = int(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["namespace_id"]) + self._execute( + f"INSERT INTO `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"(collection_id, namespace_id, ltable_name) VALUES ('{collection_id_escaped}', {ns_id}, 'default')" + ) + lt_rows = self._execute( + f"SELECT ltable_id FROM `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {ns_id} AND ltable_name = 'default'" + ) + lt_id = int(lt_rows[0][0] if isinstance(lt_rows[0], (list, tuple)) else lt_rows[0]["ltable_id"]) + schema_table = NamespaceCollectionNames.logic_schema_table_name(collection_id) + schema_content = json.dumps(_build_default_ltable_schema()) + with contextlib.suppress(Exception): + self._execute( + f"INSERT INTO `{schema_table}` (namespace_id, ltable_id, schema_content) " + f"VALUES ({ns_id}, {lt_id}, '{escape_string(schema_content)}')" + ) + return {"namespace_id": str(ns_id), "namespace_name": namespace_name} + + def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: + namespace_name_escaped = escape_string(namespace_name) + collection_id_escaped = escape_string(collection_id) + rows = self._execute( + f"SELECT namespace_id, namespace_name FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"WHERE collection_id = '{collection_id_escaped}' AND namespace_name = '{namespace_name_escaped}'" + ) + if not rows: + return None + row = rows[0] + if isinstance(row, (list, tuple)): + return {"namespace_id": str(row[0]), "namespace_name": row[1]} + return {"namespace_id": str(row["namespace_id"]), "namespace_name": row["namespace_name"]} + + def _has_ns_namespace(self, collection_id: str, namespace_name: str) -> bool: + return self._get_ns_namespace_meta(collection_id, namespace_name) is not None + + def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> None: + meta = self._get_ns_namespace_meta(collection_id, namespace_name) + if meta is None: + raise ValueError(f"Namespace '{namespace_name}' not found") + ns_id = meta["namespace_id"] + collection_id_escaped = escape_string(collection_id) + self._execute( + f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})" + ) + self._execute( + f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {ns_id}" + ) + self._execute( + f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"WHERE namespace_id = {ns_id}" + ) + + def _list_ns_namespaces(self, collection_id: str) -> list[dict]: + collection_id_escaped = escape_string(collection_id) + rows = self._execute( + f"SELECT namespace_id, namespace_name FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"WHERE collection_id = '{collection_id_escaped}' ORDER BY namespace_id" + ) + results = [] + for row in rows: + if isinstance(row, (list, tuple)): + results.append({"namespace_id": str(row[0]), "namespace_name": row[1]}) + else: + results.append({"namespace_id": str(row["namespace_id"]), "namespace_name": row["namespace_name"]}) + return results + + def _get_ns_namespace_id(self, collection_id: str, namespace_name: str) -> str: + meta = self._get_ns_namespace_meta(collection_id, namespace_name) + if meta is None: + raise ValueError(f"Namespace '{namespace_name}' not found in collection {collection_id}") + return meta["namespace_id"] + + # ==================== End Namespace Catalog Methods ==================== + def get_collection(self, name: str, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED) -> "Collection": + ns_meta = None + try: + ns_meta = self._get_ns_collection_meta(name) + except Exception: + pass + if ns_meta is not None: + return self._build_ns_collection_from_meta(ns_meta, embedding_function) try: collection = self._get_collection_v1(name, embedding_function) except ValueError as e: @@ -954,6 +1378,28 @@ def get_collection(self, name: str, embedding_function: EmbeddingFunctionParam = collection = self._get_collection_v2(name, embedding_function) return collection + def _build_ns_collection_from_meta(self, meta: dict, embedding_function=_NOT_PROVIDED) -> "Collection": + settings = meta.get("settings", {}) + dimension = settings.get("dimension") + distance = settings.get("distance", DEFAULT_DISTANCE_METRIC) + ef = None + if embedding_function is not _NOT_PROVIDED: + ef = embedding_function + elif "embedding_function" in settings: + ef_info = settings["embedding_function"] + ef = EmbeddingFunctionRegistry.get(ef_info["name"]) + if ef and hasattr(ef, "set_config"): + ef.set_config(ef_info.get("properties", {})) + return Collection( + client=self, + name=meta["collection_name"], + collection_id=meta["collection_id"], + dimension=dimension, + embedding_function=ef, + distance=distance, + use_namespace=True, + ) + def _resolve_collection_metadata_from_sdk_collections(self, collection_name: str) -> _CollectionMeta | None: """ Resolve collection metadata infromation from sdk_collections table @@ -1207,6 +1653,10 @@ def delete_collection(self, name: str) -> None: Examples: >>> client.delete_collection("my_collection") """ + if self._has_ns_collection(name): + self._delete_ns_collection_meta(name) + logger.debug(f"Deleted namespace collection '{name}'") + return try: self._delete_collection_v2(name) logger.debug(f"✅ Successfully deleted collection v2 '{name}' from sdk_collections table") @@ -1261,10 +1711,50 @@ def list_collections(self) -> list["Collection"]: >>> for col in collections: ... print(col.name) """ - collections = self._list_collections_v1() + collections = self._list_ns_collections() + collections.extend(self._list_collections_v1()) collections.extend(self._list_collections_v2()) return collections + def _list_ns_collections(self) -> list["Collection"]: + result = [] + try: + sdk_table = CollectionNames.sdk_collections_table_name() + check_sql = f"SHOW TABLES LIKE '{sdk_table}'" + check_result = self._execute(check_sql) + if not check_result: + return result + rows = self._execute( + f"SELECT collection_id, collection_name, settings " + f"FROM `{sdk_table}`" + ) + for row in rows: + try: + if isinstance(row, dict): + settings = json.loads(row["settings"]) if row.get("settings") else {} + else: + settings = json.loads(row[2]) if row[2] else {} + if not settings.get("use_namespace"): + continue + if isinstance(row, dict): + meta = { + "collection_id": str(row["collection_id"]), + "collection_name": row["collection_name"], + "settings": settings, + } + else: + meta = { + "collection_id": str(row[0]), + "collection_name": row[1], + "settings": settings, + } + result.append(self._build_ns_collection_from_meta(meta)) + except Exception as e: + logger.warning(f"Failed to build namespace collection from row: {e}") + except Exception: + pass + return result + def _list_collections_v2(self) -> list["Collection"]: collections = [] try: @@ -1281,19 +1771,26 @@ def _list_collections_v2(self) -> list["Collection"]: has_sdk_collections = False if has_sdk_collections: - query_sql = f"SELECT COLLECTION_NAME FROM {sdk_collections_table}" + query_sql = f"SELECT COLLECTION_NAME, SETTINGS FROM {sdk_collections_table}" rows = self._execute(query_sql) for row in rows: try: - # Extract collection name if isinstance(row, dict): - # Server client returns dict, get the first value - collection_name = next(iter(row.values()), "") + collection_name = row.get("COLLECTION_NAME") or row.get("collection_name", "") + settings_raw = row.get("SETTINGS") or row.get("settings") elif isinstance(row, (tuple, list)): - # Embedded client returns tuple, first element is collection name collection_name = row[0] if len(row) > 0 else "" + settings_raw = row[1] if len(row) > 1 else None else: collection_name = str(row) + settings_raw = None + if settings_raw: + try: + settings = json.loads(settings_raw) if isinstance(settings_raw, str) else settings_raw + if settings.get("use_namespace"): + continue + except (json.JSONDecodeError, TypeError): + pass collection = self.get_collection(collection_name) collections.append(collection) except Exception as e: @@ -1390,7 +1887,7 @@ def has_collection(self, name: str) -> bool: >>> if client.has_collection("my_collection"): ... print("Collection exists!") """ - return self._has_collection_v2(name) or self._has_collection_v1(name) + return self._has_ns_collection(name) or self._has_collection_v2(name) or self._has_collection_v1(name) def _has_collection_v2(self, name: str) -> bool: try: @@ -1438,6 +1935,7 @@ def get_or_create_collection( schema: Schema | None = None, configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, + use_namespace: bool = False, **kwargs, ) -> "Collection": """Get a collection if it exists, otherwise create it. @@ -1455,6 +1953,7 @@ def get_or_create_collection( Defaults to ``DefaultEmbeddingFunction`` (all-MiniLM-L6-v2). If set to None, no embedding function will be used (embeddings must be provided manually). Ignored if ``schema`` is provided. + use_namespace: If True, create a namespace-enabled collection. Defaults to False. **kwargs: Additional parameters passed to ``create_collection`` if the collection is created. Returns: @@ -1466,21 +1965,17 @@ def get_or_create_collection( Examples: >>> collection = client.get_or_create_collection("my_collection") """ - # Validate collection name before any database interaction _validate_collection_name(name) - # First, try to get the collection if self.has_collection(name): - # Collection exists, return it - # Pass embedding_function (could be _NOT_PROVIDED, None, or an EmbeddingFunction instance) return self.get_collection(name, embedding_function=embedding_function) - # Collection doesn't exist, create it with provided or default configuration return self.create_collection( name=name, schema=schema, configuration=configuration, embedding_function=embedding_function, + use_namespace=use_namespace, **kwargs, ) @@ -2354,6 +2849,8 @@ def _execute_query_with_cursor( Returns: List of normalized row dictionaries """ + if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): + print(f"[pyseekdb SQL] {sql} -- params={params}", flush=True) if use_context_manager: with conn.cursor() as cursor: cursor.execute(sql, params) @@ -2607,6 +3104,8 @@ def _should_fetch_results(self, cursor: Any, sql: str) -> bool: return is_query_sql(sql) def _execute(self, sql: str) -> Any: + if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): + print(f"[pyseekdb SQL] {sql}", flush=True) conn = self._ensure_connection() use_context_manager = self._use_context_manager_for_cursor() @@ -3878,3 +4377,851 @@ def _collection_count(self, collection_id: str | None, collection_name: str) -> logger.debug(f"✅ Collection '{collection_name}' has {count} items") return count + + # ==================== Namespace DML/DQL Methods ==================== + + @staticmethod + def _rewrite_where_for_ns(where: dict[str, Any]) -> dict[str, Any]: + if where is None: + return None + rewritten = {} + for k, v in where.items(): + if k in ("$and", "$or"): + rewritten[k] = [BaseClient._rewrite_where_for_ns(sub) for sub in v] + elif k == "$not": + rewritten[k] = BaseClient._rewrite_where_for_ns(v) + else: + rewritten[f"metadata.{k}"] = v + return rewritten + + @staticmethod + def _append_namespace_filter( + where_clause: str, + params: list[Any], + namespace_id: int, + ltable_id: int = 1, + ) -> tuple[str, list[Any]]: + ns_cond = f"namespace_id = {namespace_id} AND ltable_id = {ltable_id}" + if not where_clause: + return f"WHERE {ns_cond}", params + if where_clause.strip().upper().startswith("WHERE"): + inner = where_clause.strip()[5:].strip() + return f"WHERE {ns_cond} AND ({inner})", params + return f"WHERE {ns_cond} AND ({where_clause})", params + + def _namespace_add( # noqa: C901 + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + ids: str | list[str], + embeddings: list[float] | list[list[float]] | None = None, + metadatas: dict | list[dict] | None = None, + documents: str | list[str] | None = None, + embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, + **kwargs, + ) -> None: + if isinstance(ids, str): + ids = [ids] + _validate_record_ids(ids) + if len(ids) > _MAX_NAMESPACE_BATCH_SIZE: + raise ValueError( + f"Batch size {len(ids)} exceeds maximum allowed {_MAX_NAMESPACE_BATCH_SIZE} records per request." + ) + if isinstance(documents, str): + documents = [documents] + if metadatas is not None and isinstance(metadatas, dict): + metadatas = [metadatas] + if ( + embeddings is not None + and isinstance(embeddings, list) + and len(embeddings) > 0 + and not isinstance(embeddings[0], list) + ): + embeddings = [embeddings] + + if embeddings: + pass + elif documents: + if embedding_function is not None: + embeddings = embedding_function(documents) + else: + raise ValueError( + "Documents provided but no embeddings and no embedding function." + ) + else: + raise ValueError("Neither embeddings nor documents provided.") + + num_items = len(ids) + if documents and len(documents) != num_items: + raise ValueError(f"Number of documents ({len(documents)}) does not match number of ids ({num_items})") + if metadatas and len(metadatas) != num_items: + raise ValueError(f"Number of metadatas ({len(metadatas)}) does not match number of ids ({num_items})") + if embeddings and len(embeddings) != num_items: + raise ValueError(f"Number of embeddings ({len(embeddings)}) does not match number of ids ({num_items})") + + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_id) + ltable_id = 1 + + values_list = [] + for i in range(num_items): + doc_val = documents[i] if documents else None + doc_sql = f"'{escape_string(doc_val)}'" if doc_val is not None else "NULL" + + vec_val = embeddings[i] if embeddings else None + vec_sql = "NULL" if vec_val is None else _embedding_to_hexstring(vec_val) + + meta_val = metadatas[i] if metadatas else None + data_content = {"id": ids[i]} + if meta_val is not None: + data_content["metadata"] = meta_val + dc_json = json.dumps(data_content, ensure_ascii=False) + dc_sql = f"'{escape_string(dc_json)}'" + + values_list.append(f"({ns_id}, {ltable_id}, {doc_sql}, {vec_sql}, {dc_sql})") + + columns = ( + f"{NamespaceFieldNames.NAMESPACE_ID}, {NamespaceFieldNames.LTABLE_ID}, " + f"{NamespaceFieldNames.DOCUMENT}, {NamespaceFieldNames.EMBEDDING}, {NamespaceFieldNames.DATA_CONTENT}" + ) + sql = f"INSERT INTO `{table_name}` ({columns}) VALUES {','.join(values_list)}" + self._execute(sql) + + def _namespace_update( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + ids: str | list[str], + embeddings: list[float] | list[list[float]] | None = None, + metadatas: dict | list[dict] | None = None, + documents: str | list[str] | None = None, + embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, + **kwargs, + ) -> None: + if isinstance(ids, str): + ids = [ids] + _validate_record_ids(ids) + if len(ids) > _MAX_NAMESPACE_BATCH_SIZE: + raise ValueError( + f"Batch size {len(ids)} exceeds maximum allowed {_MAX_NAMESPACE_BATCH_SIZE} records per request." + ) + if isinstance(documents, str): + documents = [documents] + if metadatas is not None and isinstance(metadatas, dict): + metadatas = [metadatas] + if ( + embeddings is not None + and isinstance(embeddings, list) + and len(embeddings) > 0 + and not isinstance(embeddings[0], list) + ): + embeddings = [embeddings] + + if embeddings is None and documents is not None and embedding_function is not None: + embeddings = embedding_function(documents) + + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_id) + ltable_id = 1 + + for i, record_id in enumerate(ids): + set_parts = [] + if documents and i < len(documents) and documents[i] is not None: + set_parts.append(f"document = '{escape_string(documents[i])}'") + if embeddings and i < len(embeddings) and embeddings[i] is not None: + set_parts.append(f"embedding = {_embedding_to_hexstring(embeddings[i])}") + if metadatas and i < len(metadatas) and metadatas[i] is not None: + meta_json = json.dumps(metadatas[i], ensure_ascii=False) + set_parts.append( + f"data_content = JSON_SET(data_content, '$.metadata', CAST('{escape_string(meta_json)}' AS JSON))" + ) + if not set_parts: + continue + id_escaped = escape_string(record_id) + user_where = f"WHERE JSON_EXTRACT(data_content, '$.id') = '{id_escaped}'" + where_clause, _ = self._append_namespace_filter(user_where, [], ns_id, ltable_id) + sql = f"UPDATE `{table_name}` SET {', '.join(set_parts)} {where_clause}" + self._execute(sql) + + def _namespace_upsert( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + ids: str | list[str], + embeddings: list[float] | list[list[float]] | None = None, + metadatas: dict | list[dict] | None = None, + documents: str | list[str] | None = None, + embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, + **kwargs, + ) -> None: + if isinstance(ids, str): + ids = [ids] + _validate_record_ids(ids) + if len(ids) > _MAX_NAMESPACE_BATCH_SIZE: + raise ValueError( + f"Batch size {len(ids)} exceeds maximum allowed {_MAX_NAMESPACE_BATCH_SIZE} records per request." + ) + if isinstance(documents, str): + documents = [documents] + if metadatas is not None and isinstance(metadatas, dict): + metadatas = [metadatas] + if ( + embeddings is not None + and isinstance(embeddings, list) + and len(embeddings) > 0 + and not isinstance(embeddings[0], list) + ): + embeddings = [embeddings] + + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_id) + ltable_id = 1 + + existing_ids = set() + for record_id in ids: + id_escaped = escape_string(record_id) + user_where = f"WHERE JSON_EXTRACT(data_content, '$.id') = '{id_escaped}'" + where_clause, _ = self._append_namespace_filter(user_where, [], ns_id, ltable_id) + rows = self._execute( + f"SELECT 1 FROM `{table_name}` {where_clause} LIMIT 1" + ) + if rows: + existing_ids.add(record_id) + + add_indices = [] + update_indices = [] + for i, rid in enumerate(ids): + if rid in existing_ids: + update_indices.append(i) + else: + add_indices.append(i) + + if add_indices: + add_ids = [ids[i] for i in add_indices] + add_docs = [documents[i] for i in add_indices] if documents else None + add_metas = [metadatas[i] for i in add_indices] if metadatas else None + add_embs = [embeddings[i] for i in add_indices] if embeddings else None + self._namespace_add( + collection_id=collection_id, collection_name=collection_name, + namespace_id=namespace_id, namespace_name=namespace_name, + ids=add_ids, embeddings=add_embs, metadatas=add_metas, + documents=add_docs, embedding_function=embedding_function, **kwargs, + ) + + if update_indices: + upd_ids = [ids[i] for i in update_indices] + upd_docs = [documents[i] for i in update_indices] if documents else None + upd_metas = [metadatas[i] for i in update_indices] if metadatas else None + upd_embs = [embeddings[i] for i in update_indices] if embeddings else None + self._namespace_update( + collection_id=collection_id, collection_name=collection_name, + namespace_id=namespace_id, namespace_name=namespace_name, + ids=upd_ids, embeddings=upd_embs, metadatas=upd_metas, + documents=upd_docs, embedding_function=embedding_function, **kwargs, + ) + + def _namespace_delete( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + ids: str | list[str] | None = None, + where: dict[str, Any] | None = None, + where_document: dict[str, Any] | None = None, + **kwargs, + ) -> None: + if ids is None and where is None and where_document is None: + raise ValueError("At least one of ids, where, or where_document must be provided") + + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_id) + + conditions = [] + params = [] + + if ids is not None: + if isinstance(ids, str): + ids = [ids] + _validate_record_ids(ids) + id_conditions = [] + for rid in ids: + id_escaped = escape_string(rid) + id_conditions.append(f"JSON_EXTRACT(data_content, '$.id') = '{id_escaped}'") + conditions.append(f"({' OR '.join(id_conditions)})") + + if where is not None: + rewritten = self._rewrite_where_for_ns(where) + meta_clause, meta_params = FilterBuilder.build_metadata_filter(rewritten, "data_content") + if meta_clause: + conditions.append(meta_clause) + params.extend(meta_params) + + if where_document is not None: + doc_clause, doc_params = FilterBuilder.build_document_filter(where_document, "document") + if doc_clause: + conditions.append(doc_clause) + params.extend(doc_params) + + user_where = f"WHERE {' AND '.join(conditions)}" if conditions else "" + where_clause, params = self._append_namespace_filter(user_where, params, ns_id) + sql = f"DELETE FROM `{table_name}` {where_clause}" + if params: + conn = self._ensure_connection() + use_context_manager = self._use_context_manager_for_cursor() + if use_context_manager: + with conn.cursor() as cursor: + cursor.execute(sql, params) + else: + cursor = conn.cursor() + try: + cursor.execute(sql, params) + finally: + cursor.close() + else: + self._execute(sql) + + def _namespace_query( # noqa: C901 + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + query_embeddings: list[float] | list[list[float]] | None = None, + query_texts: str | list[str] | None = None, + n_results: int = 10, + where: dict[str, Any] | None = None, + where_document: dict[str, Any] | None = None, + include: list[str] | None = None, + **kwargs, + ) -> dict[str, Any]: + embedding_function = kwargs.get("embedding_function") + distance = kwargs.get("distance", DEFAULT_DISTANCE_METRIC) + + if query_embeddings is not None: + pass + elif query_texts is not None: + if embedding_function is not None: + query_embeddings = self._embed_texts(query_texts, embedding_function=embedding_function) + else: + raise ValueError("query_texts provided but no embedding_function.") + else: + raise ValueError("Neither query_embeddings nor query_texts provided.") + + query_embeddings = self._normalize_query_embeddings(query_embeddings) + include_fields = self._normalize_include_fields(include) + + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_id) + + select_parts = ["JSON_EXTRACT(data_content, '$.id') AS record_id"] + if include_fields.get("documents") or include_fields.get("document") or include is None: + select_parts.append("document") + if include_fields.get("metadatas") or include_fields.get("metadata") or include is None: + select_parts.append("JSON_EXTRACT(data_content, '$.metadata') AS metadata") + if include_fields.get("embeddings") or include_fields.get("embedding"): + select_parts.append("embedding") + + user_conditions = [] + filter_params = [] + + if where is not None: + rewritten = self._rewrite_where_for_ns(where) + meta_clause, meta_params = FilterBuilder.build_metadata_filter(rewritten, "data_content") + if meta_clause: + user_conditions.append(meta_clause) + filter_params.extend(meta_params) + + if where_document is not None: + doc_clause, doc_params = FilterBuilder.build_document_filter(where_document, "document") + if doc_clause: + user_conditions.append(doc_clause) + filter_params.extend(doc_params) + + user_where = f"WHERE {' AND '.join(user_conditions)}" if user_conditions else "" + where_clause, filter_params = self._append_namespace_filter(user_where, filter_params, ns_id) + + distance_function_map = { + "l2": "l2_distance", + "cosine": "cosine_distance", + "inner_product": "inner_product", + } + distance_func = distance_function_map.get(distance, "l2_distance") + + conn = self._ensure_connection() + use_context_manager = self._use_context_manager_for_cursor() + + all_ids = [] + all_documents = [] + all_metadatas = [] + all_embeddings = [] + all_distances = [] + + for query_vector in query_embeddings: + vector_str = _embedding_to_hexstring(query_vector) + select_clause = ", ".join(select_parts) + sql = ( + f"SELECT {select_clause}, " + f"{distance_func}(embedding, {vector_str}) AS distance " + f"FROM `{table_name}` " + f"{where_clause} " + f"ORDER BY {distance_func}(embedding, {vector_str}) " + f"APPROXIMATE LIMIT %s" + ) + query_params = [*filter_params, n_results] + rows = self._execute_query_with_cursor(conn, sql, query_params, use_context_manager) + + q_ids = [] + q_documents = [] + q_metadatas = [] + q_embeddings = [] + q_distances = [] + + for row in rows: + if isinstance(row, dict): + rid_raw = row.get("record_id") + rid = json.loads(rid_raw) if isinstance(rid_raw, str) else rid_raw + q_ids.append(rid) + if "documents" in include_fields or "document" in include_fields or include is None: + q_documents.append(row.get("document")) + if "metadatas" in include_fields or "metadata" in include_fields or include is None: + meta_raw = row.get("metadata") + if isinstance(meta_raw, str): + meta_raw = json.loads(meta_raw) + q_metadatas.append(meta_raw or {}) + if "embeddings" in include_fields or "embedding" in include_fields: + emb = row.get("embedding") + if isinstance(emb, bytes): + emb = self._parse_embedding_from_bytes(emb) + elif isinstance(emb, str): + emb = json.loads(emb) + q_embeddings.append(emb) + q_distances.append(row.get("distance")) + elif isinstance(row, (list, tuple)): + idx = 0 + rid_raw = row[idx]; idx += 1 + rid = json.loads(rid_raw) if isinstance(rid_raw, str) else rid_raw + q_ids.append(rid) + if "documents" in include_fields or "document" in include_fields or include is None: + q_documents.append(row[idx]); idx += 1 + if "metadatas" in include_fields or "metadata" in include_fields or include is None: + meta_raw = row[idx]; idx += 1 + if isinstance(meta_raw, str): + meta_raw = json.loads(meta_raw) + q_metadatas.append(meta_raw or {}) + if "embeddings" in include_fields or "embedding" in include_fields: + emb = row[idx]; idx += 1 + if isinstance(emb, bytes): + emb = self._parse_embedding_from_bytes(emb) + elif isinstance(emb, str): + emb = json.loads(emb) + q_embeddings.append(emb) + q_distances.append(row[idx]) + + all_ids.append(q_ids) + if "documents" in include_fields or "document" in include_fields or include is None: + all_documents.append(q_documents) + if "metadatas" in include_fields or "metadata" in include_fields or include is None: + all_metadatas.append(q_metadatas) + if "embeddings" in include_fields or "embedding" in include_fields: + all_embeddings.append(q_embeddings) + all_distances.append(q_distances) + + result = {"ids": all_ids, "distances": all_distances} + if "documents" in include_fields or "document" in include_fields or include is None: + result["documents"] = all_documents + if "metadatas" in include_fields or "metadata" in include_fields or include is None: + result["metadatas"] = all_metadatas + if "embeddings" in include_fields or "embedding" in include_fields: + result["embeddings"] = all_embeddings + return result + + def _namespace_get( # noqa: C901 + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + ids: str | list[str] | None = None, + where: dict[str, Any] | None = None, + where_document: dict[str, Any] | None = None, + limit: int | None = None, + offset: int | None = None, + include: list[str] | None = None, + **kwargs, + ) -> dict[str, Any]: + include_fields = self._normalize_include_fields(include) + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_id) + + select_parts = ["JSON_EXTRACT(data_content, '$.id') AS record_id"] + if include_fields.get("documents") or include_fields.get("document") or include is None: + select_parts.append("document") + if include_fields.get("metadatas") or include_fields.get("metadata") or include is None: + select_parts.append("JSON_EXTRACT(data_content, '$.metadata') AS metadata") + if include_fields.get("embeddings") or include_fields.get("embedding"): + select_parts.append("embedding") + + user_conditions = [] + params = [] + + if ids is not None: + if isinstance(ids, str): + ids = [ids] + id_conds = [] + for rid in ids: + id_escaped = escape_string(rid) + id_conds.append(f"JSON_EXTRACT(data_content, '$.id') = '{id_escaped}'") + user_conditions.append(f"({' OR '.join(id_conds)})") + + if where is not None: + rewritten = self._rewrite_where_for_ns(where) + meta_clause, meta_params = FilterBuilder.build_metadata_filter(rewritten, "data_content") + if meta_clause: + user_conditions.append(meta_clause) + params.extend(meta_params) + + if where_document is not None: + doc_clause, doc_params = FilterBuilder.build_document_filter(where_document, "document") + if doc_clause: + user_conditions.append(doc_clause) + params.extend(doc_params) + + user_where = f"WHERE {' AND '.join(user_conditions)}" if user_conditions else "" + where_str, params = self._append_namespace_filter(user_where, params, ns_id) + select_clause = ", ".join(select_parts) + sql = f"SELECT {select_clause} FROM `{table_name}` {where_str}" + if limit is not None: + sql += f" LIMIT {int(limit)}" + if offset is not None: + sql += f" OFFSET {int(offset)}" + + conn = self._ensure_connection() + use_context_manager = self._use_context_manager_for_cursor() + rows = self._execute_query_with_cursor(conn, sql, params if params else [], use_context_manager) + + result_ids = [] + result_documents = [] + result_metadatas = [] + result_embeddings = [] + + for row in rows: + if isinstance(row, dict): + rid_raw = row.get("record_id") + rid = json.loads(rid_raw) if isinstance(rid_raw, str) else rid_raw + result_ids.append(rid) + if "documents" in include_fields or "document" in include_fields or include is None: + result_documents.append(row.get("document")) + if "metadatas" in include_fields or "metadata" in include_fields or include is None: + meta_raw = row.get("metadata") + if isinstance(meta_raw, str): + meta_raw = json.loads(meta_raw) + result_metadatas.append(meta_raw or {}) + if "embeddings" in include_fields or "embedding" in include_fields: + emb = row.get("embedding") + if isinstance(emb, bytes): + emb = self._parse_embedding_from_bytes(emb) + elif isinstance(emb, str): + emb = json.loads(emb) + result_embeddings.append(emb) + elif isinstance(row, (list, tuple)): + idx = 0 + rid_raw = row[idx]; idx += 1 + rid = json.loads(rid_raw) if isinstance(rid_raw, str) else rid_raw + result_ids.append(rid) + if "documents" in include_fields or "document" in include_fields or include is None: + result_documents.append(row[idx]); idx += 1 + if "metadatas" in include_fields or "metadata" in include_fields or include is None: + meta_raw = row[idx]; idx += 1 + if isinstance(meta_raw, str): + meta_raw = json.loads(meta_raw) + result_metadatas.append(meta_raw or {}) + if "embeddings" in include_fields or "embedding" in include_fields: + emb = row[idx]; idx += 1 + if isinstance(emb, bytes): + emb = self._parse_embedding_from_bytes(emb) + elif isinstance(emb, str): + emb = json.loads(emb) + result_embeddings.append(emb) + + result = {"ids": result_ids} + if "documents" in include_fields or "document" in include_fields or include is None: + result["documents"] = result_documents + if "metadatas" in include_fields or "metadata" in include_fields or include is None: + result["metadatas"] = result_metadatas + if "embeddings" in include_fields or "embedding" in include_fields: + result["embeddings"] = result_embeddings + return result + + def _namespace_count( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + **kwargs, + ) -> int: + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_id) + where_clause, _ = self._append_namespace_filter("", [], ns_id) + sql = f"SELECT COUNT(*) AS cnt FROM `{table_name}` {where_clause}" + conn = self._ensure_connection() + use_context_manager = self._use_context_manager_for_cursor() + rows = self._execute_query_with_cursor(conn, sql, [], use_context_manager) + if not rows: + return 0 + row = rows[0] + if isinstance(row, dict): + return row.get("cnt", 0) + elif isinstance(row, (tuple, list)): + return row[0] if len(row) > 0 else 0 + return int(row) if row else 0 + + def _namespace_peek( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + limit: int = 10, + **kwargs, + ) -> dict[str, Any]: + return self._namespace_get( + collection_id=collection_id, + collection_name=collection_name, + namespace_id=namespace_id, + namespace_name=namespace_name, + limit=limit, + offset=0, + include=["documents", "metadatas", "embeddings"], + ) + + def _adapt_search_parm_for_ns(self, search_parm: dict[str, Any], ns_id: int) -> dict[str, Any]: + ns_filter = [ + {"term": {"namespace_id": ns_id}}, + {"term": {"ltable_id": 1}}, + ] + + def _rewrite_field_refs(obj): + if isinstance(obj, dict): + new_dict = {} + for k, v in obj.items(): + new_key = k + if k == "_id": + new_key = "(JSON_EXTRACT(data_content, '$.id'))" + elif isinstance(k, str) and "JSON_EXTRACT(metadata, '$." in k: + new_key = k.replace( + "JSON_EXTRACT(metadata, '$.", + "JSON_EXTRACT(data_content, '$.metadata.", + ) + new_dict[new_key] = _rewrite_field_refs(v) + return new_dict + elif isinstance(obj, list): + return [_rewrite_field_refs(item) for item in obj] + return obj + + search_parm = _rewrite_field_refs(search_parm) + + if "_source" in search_parm: + needs_data_content = False + new_source = [] + for field in search_parm["_source"]: + if field in ("_id", "metadata"): + needs_data_content = True + else: + new_source.append(field) + if needs_data_content: + new_source.insert(0, "data_content") + search_parm["_source"] = new_source + + def _inject_filter_into_knn(node): + if isinstance(node, dict): + if "filter" in node: + existing = node["filter"] + if isinstance(existing, list): + node["filter"] = existing + ns_filter + else: + node["filter"] = [existing] + ns_filter + else: + node["filter"] = list(ns_filter) + return node + + def _inject_filter_into_query(node): + if not isinstance(node, dict): + return node + if "bool" in node: + bool_node = node["bool"] + if "filter" in bool_node: + existing = bool_node["filter"] + if isinstance(existing, list): + bool_node["filter"] = existing + ns_filter + else: + bool_node["filter"] = [existing] + ns_filter + else: + bool_node["filter"] = list(ns_filter) + return node + return {"bool": {"must": [node], "filter": list(ns_filter)}} + + if "query" in search_parm: + q = search_parm["query"] + if isinstance(q, list): + search_parm["query"] = [_inject_filter_into_query(item) for item in q] + else: + search_parm["query"] = _inject_filter_into_query(q) + + if "knn" in search_parm: + knn = search_parm["knn"] + if isinstance(knn, list): + for item in knn: + _inject_filter_into_knn(item) + else: + _inject_filter_into_knn(knn) + + if "query" not in search_parm and "knn" not in search_parm: + search_parm["query"] = {"bool": {"filter": list(ns_filter)}} + + return search_parm + + def _namespace_hybrid_search( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + query: dict[str, Any] | None = None, + knn: dict[str, Any] | None = None, + rank: dict[str, Any] | None = None, + n_results: int = 10, + include: list[str] | None = None, + query_hint: QueryHint | None = None, + **kwargs, + ) -> dict[str, Any]: + conn = self._ensure_connection() + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_id) + + search_parm = self._build_search_parm( + query, knn, rank, n_results, + include=include, + dimension=kwargs.get("dimension"), + **{k: v for k, v in kwargs.items() if k != "dimension"}, + ) + search_parm = self._adapt_search_parm_for_ns(search_parm, ns_id) + + search_parm_json = json.dumps(search_parm, ensure_ascii=False) + use_context_manager = self._use_context_manager_for_cursor() + + escaped_params = escape_string(search_parm_json) + set_sql = f"SET @search_parm = '{escaped_params}'" + self._execute_query_with_cursor(conn, set_sql, [], use_context_manager) + + get_sql_query = f"SELECT DBMS_HYBRID_SEARCH.GET_SQL('{table_name}', @search_parm) as query_sql FROM dual" + rows = self._execute_query_with_cursor(conn, get_sql_query, [], use_context_manager) + + if not rows or not rows[0].get("query_sql"): + return { + "ids": [[]], "distances": [[]], "metadatas": [[]], + "documents": [[]], "embeddings": [[]], + } + + query_sql = rows[0]["query_sql"] + if isinstance(query_sql, str): + query_sql = query_sql.strip().strip("'\"") + + hint_sql = _query_hint_to_sql(query_hint, table_name=table_name) + if hint_sql and query_sql.upper().startswith("SELECT"): + query_sql = f"SELECT {hint_sql} {query_sql[len('SELECT'):]}" + + result_rows = self._execute_query_with_cursor(conn, query_sql, [], use_context_manager) + return self._transform_ns_hybrid_result(result_rows, include) + + def _transform_ns_hybrid_result( + self, result_rows: list[dict[str, Any]], include: list[str] | None + ) -> dict[str, Any]: + if not result_rows: + return { + "ids": [[]], "distances": [[]], "metadatas": [[]], + "documents": [[]], "embeddings": [[]], + } + + ids = [] + distances = [] + metadatas = [] + documents = [] + embeddings = [] + + for row in result_rows: + dc_raw = row.get("data_content") or row.get("DATA_CONTENT") + dc = {} + if isinstance(dc_raw, str): + with contextlib.suppress(json.JSONDecodeError): + dc = json.loads(dc_raw) + elif isinstance(dc_raw, dict): + dc = dc_raw + + row_id = dc.get("id") + if row_id is None: + for key in ("id", "_id", "ID"): + if key in row and row[key] is not None: + row_id = row[key] + break + row_id = self._convert_id_from_bytes(row_id) + ids.append(row_id) + + distance = ( + row.get("_distance") or row.get("distance") + or row.get("_score") or row.get("score") or 0.0 + ) + distances.append(distance) + + if include is None or "metadatas" in include or "metadata" in include: + meta = dc.get("metadata") + if meta is None: + meta = row.get("metadata") or row.get("METADATA") + if isinstance(meta, str): + with contextlib.suppress(json.JSONDecodeError): + meta = json.loads(meta) + metadatas.append(meta or {}) + else: + metadatas.append(None) + + if include is None or "documents" in include or "document" in include: + documents.append(row.get("document") or row.get("DOCUMENT")) + else: + documents.append(None) + + if include and ("embeddings" in include or "embedding" in include): + emb = row.get("embedding") or row.get("EMBEDDING") + if isinstance(emb, str): + with contextlib.suppress(json.JSONDecodeError): + emb = json.loads(emb) + elif isinstance(emb, bytes): + emb = self._parse_embedding_from_bytes(emb) + embeddings.append(emb) + else: + embeddings.append(None) + + result = {"ids": [ids], "distances": [distances]} + if include is None or "documents" in include or "document" in include: + result["documents"] = [documents] + if include is None or "metadatas" in include or "metadata" in include: + result["metadatas"] = [metadatas] + if include and ("embeddings" in include or "embedding" in include): + result["embeddings"] = [embeddings] + return result + + def _namespace_prewarm( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + **kwargs, + ) -> None: + raise NotImplementedError("prewarm is not supported in this client mode") diff --git a/src/pyseekdb/client/client_seekdb_embedded.py b/src/pyseekdb/client/client_seekdb_embedded.py index 48532580..293d4846 100644 --- a/src/pyseekdb/client/client_seekdb_embedded.py +++ b/src/pyseekdb/client/client_seekdb_embedded.py @@ -270,6 +270,16 @@ def list_databases( """ return super().list_databases(limit=limit, offset=offset, tenant=tenant) + def _namespace_prewarm( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + **kwargs, + ) -> None: + raise ValueError("prewarm is only supported in shared-storage remote deployment") + def __repr__(self): status = "connected" if self.is_connected() else "disconnected" return f"" diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 4dcc1a0e..8d43254e 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -200,6 +200,17 @@ def _database_tenant(self, tenant: str) -> str | None: ) return self.tenant + def _namespace_prewarm( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + **kwargs, + ) -> None: + sql = f"CALL DBMS_LOGIC_TABLE.PREWARM({collection_id}, {namespace_id})" + self._execute(sql) + def __repr__(self): status = "connected" if self.is_connected() else "disconnected" return f"" diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 54982837..18e2ff19 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -10,9 +10,12 @@ from typing import TYPE_CHECKING, Any, Optional +from .validators import _validate_namespace_name + if TYPE_CHECKING: from .embedding_function import Documents as EmbeddingDocuments from .embedding_function import EmbeddingFunction + from .namespace import Namespace from .query_types import QueryHint from .schema import SparseVectorIndexConfig from .sparse_embedding_function import SparseEmbeddingFunction @@ -38,29 +41,17 @@ def __init__( embedding_function: Optional["EmbeddingFunction[EmbeddingDocuments]"] = None, distance: str | None = None, sparse_vector_index_config: Optional["SparseVectorIndexConfig"] = None, + use_namespace: bool = False, **metadata, ): - """ - Initialize collection object - - Args: - client: The client instance that created this collection - name: Collection name - collection_id: Collection unique identifier (some databases may need this) - dimension: Vector dimension - embedding_function: Embedding function to convert documents to embeddings - distance: Distance metric used by the index (e.g., 'l2', 'cosine', 'inner_product') - sparse_vector_index_config: Sparse vector index configuration (optional). - When set, the collection supports sparse vector operations. - **metadata: Other metadata - """ - self._client = client # Core: hold reference to the client + self._client = client self._name = name self._id = collection_id self._dimension = dimension self._embedding_function = embedding_function self._distance = distance self._sparse_vector_index_config = sparse_vector_index_config + self._use_namespace = use_namespace self._metadata = metadata # ==================== Properties ==================== @@ -112,8 +103,75 @@ def sparse_embedding_function(self) -> Optional["SparseEmbeddingFunction"]: return self._sparse_vector_index_config.embedding_function return None + @property + def use_namespace(self) -> bool: + return self._use_namespace + + def _guard_collection_data_api(self) -> None: + if self._use_namespace: + raise ValueError( + "This collection has namespace enabled. " + "Use namespace.add/update/upsert/delete/query/get instead of collection-level methods. " + "Get a namespace via collection.create_namespace() or collection.get_namespace()." + ) + def __repr__(self) -> str: - return f"Collection(name='{self._name}', dimension={self._dimension}, client={self._client.mode})" + ns_str = ", use_namespace=True" if self._use_namespace else "" + return f"Collection(name='{self._name}', dimension={self._dimension}{ns_str}, client={self._client.mode})" + + # ==================== Namespace Management ==================== + + def create_namespace(self, name: str) -> "Namespace": + if not self._use_namespace: + raise ValueError("Namespace is not enabled for this collection. Use use_namespace=True when creating the collection.") + _validate_namespace_name(name) + from .namespace import Namespace + meta = self._client._create_ns_namespace_meta(self._id, name) + return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) + + def get_namespace(self, name: str) -> "Namespace": + if not self._use_namespace: + raise ValueError("Namespace is not enabled for this collection.") + _validate_namespace_name(name) + from .namespace import Namespace + meta = self._client._get_ns_namespace_meta(self._id, name) + if meta is None: + raise ValueError(f"Namespace '{name}' not found") + return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) + + def get_or_create_namespace(self, name: str) -> "Namespace": + if not self._use_namespace: + raise ValueError("Namespace is not enabled for this collection.") + _validate_namespace_name(name) + from .namespace import Namespace + meta = self._client._get_ns_namespace_meta(self._id, name) + if meta is None: + meta = self._client._create_ns_namespace_meta(self._id, name) + return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) + + def delete_namespace(self, name: str) -> None: + if not self._use_namespace: + raise ValueError("Namespace is not enabled for this collection.") + _validate_namespace_name(name) + self._client._delete_ns_namespace_meta(self._id, name) + + def list_namespaces(self) -> list["Namespace"]: + if not self._use_namespace: + raise ValueError("Namespace is not enabled for this collection.") + from .namespace import Namespace + metas = self._client._list_ns_namespaces(self._id) + return [ + Namespace(client=self._client, collection=self, name=m["namespace_name"], namespace_id=m["namespace_id"]) + for m in metas + ] + + def has_namespace(self, name: str) -> bool: + if not self._use_namespace: + raise ValueError("Namespace is not enabled for this collection.") + _validate_namespace_name(name) + return self._client._has_ns_namespace(self._id, name) + + # ==================== End Namespace Management ==================== def fork(self, forked_name: str) -> "Collection": """ @@ -151,6 +209,7 @@ def fork(self, forked_name: str) -> "Collection": assert forked.count() == 4 # Forked has new data """ + self._guard_collection_data_api() self._client._collection_fork(collection=self, forked_name=forked_name) collection = self._client.get_collection(forked_name, embedding_function=self._embedding_function) return collection @@ -195,6 +254,7 @@ def add( metadatas=[{"tag": "A"}, {"tag": "B"}] ) """ + self._guard_collection_data_api() return self._client._collection_add( collection_id=self._id, collection_name=self._name, @@ -238,6 +298,7 @@ def update( embeddings=[[0.9, 0.8], [0.7, 0.6]] ) """ + self._guard_collection_data_api() return self._client._collection_update( collection_id=self._id, collection_name=self._name, @@ -281,6 +342,7 @@ def upsert( embeddings=[[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]] ) """ + self._guard_collection_data_api() return self._client._collection_upsert( collection_id=self._id, collection_name=self._name, @@ -322,6 +384,7 @@ def delete( # Delete by document filter collection.delete(where_document={"$contains": "keyword"}) """ + self._guard_collection_data_api() return self._client._collection_delete( collection_id=self._id, collection_name=self._name, @@ -407,6 +470,7 @@ def query( query_hint=QueryHint(parallel=8, query_timeout=10.0) ) """ + self._guard_collection_data_api() return self._client._collection_query( collection_id=self._id, collection_name=self._name, @@ -488,6 +552,7 @@ def get( query_hint=QueryHint(parallel=4, query_timeout=5.0) ) """ + self._guard_collection_data_api() return self._client._collection_get( collection_id=self._id, collection_name=self._name, @@ -575,6 +640,7 @@ def hybrid_search( query_hint=QueryHint(parallel=6, query_timeout=15.0) ) """ + self._guard_collection_data_api() # When no query/knn provided, return only ids/distances by default if include is None and not query and not knn: include = [] @@ -606,6 +672,7 @@ def count(self) -> int: count = collection.count() print(f"Collection has {count} items") """ + self._guard_collection_data_api() return self._client._collection_count(collection_id=self._id, collection_name=self._name) def peek(self, limit: int = 10) -> dict[str, Any]: @@ -629,6 +696,7 @@ def peek(self, limit: int = 10) -> dict[str, Any]: print(f"ID: {preview['ids'][i]}, Document: {preview['documents'][i]}") print(f"Metadata: {preview['metadatas'][i]}, Embedding: {preview['embeddings'][i]}") """ + self._guard_collection_data_api() return self._client._collection_get( collection_id=self._id, collection_name=self._name, diff --git a/src/pyseekdb/client/configuration.py b/src/pyseekdb/client/configuration.py index 355a08cb..243dcee7 100644 --- a/src/pyseekdb/client/configuration.py +++ b/src/pyseekdb/client/configuration.py @@ -205,6 +205,17 @@ class HNSWIndexLib(str, Enum): VSAG = "vsag" +class IVFIndexType(str, Enum): + IVF_FLAT = "ivf_flat" + IVF_SQ8 = "ivf_sq8" + IVF_PQ = "ivf_pq" + + +class IVFIndexLib(str, Enum): + OB = "ob" + VSAG = "vsag" + + class FulltextAnalyzer(str, Enum): SPACE = "space" NGRAM = "ngram" @@ -316,14 +327,65 @@ class BengProperties(TypedDict, total=False): max_token_size: int +@dataclass +class IVFConfiguration: + """ + IVF (Inverted File) index configuration for Agent Database namespace-enabled collections. + + Args: + dimension: Vector dimension (number of elements in each vector) + distance: Distance metric for similarity calculation (e.g., 'l2', 'cosine', 'inner_product') + type: IVF index subtype ('ivf_flat', 'ivf_sq8', 'ivf_pq') + use_spfresh: Whether to enable SPFresh optimization for online index updates + properties: Optional dictionary of additional IVF index properties + """ + + dimension: int = DEFAULT_VECTOR_DIMENSION + distance: str | DistanceMetric = DistanceMetric.COSINE.value + type: str | IVFIndexType = IVFIndexType.IVF_FLAT.value + lib: str | IVFIndexLib = IVFIndexLib.OB.value + use_spfresh: bool | None = None + properties: dict[str, PrimitiveValue] | None = None + + def __post_init__(self): + if isinstance(self.dimension, bool) or not isinstance(self.dimension, int): + raise TypeError(f"dimension must be an integer, got {type(self.dimension).__name__}") + _validate_int_range(self.dimension, key="dimension", min_value=1, max_value=4096) + + self.distance = _normalize_str_enum(self.distance, field_name="distance") + valid_distances = [e.value for e in DistanceMetric] + if self.distance not in valid_distances: + raise ValueError(f"distance must be one of {valid_distances}, got {self.distance}") + + self.type = _normalize_str_enum(self.type, field_name="type") + valid_types = [e.value for e in IVFIndexType] + if self.type not in valid_types: + raise ValueError(f"type must be one of {valid_types}, got {self.type}") + + self.lib = _normalize_str_enum(self.lib, field_name="lib") + valid_libs = [e.value for e in IVFIndexLib] + if self.lib not in valid_libs: + raise ValueError(f"lib must be one of {valid_libs}, got {self.lib}") + + if self.use_spfresh is not None and not isinstance(self.use_spfresh, bool): + raise TypeError(f"use_spfresh must be a bool, got {type(self.use_spfresh).__name__}") + + _ensure_primitive_properties(self.properties) + + @dataclass class VectorIndexConfig: + ivf: IVFConfiguration | None = None hnsw: HNSWConfiguration | None = None embedding_function: EmbeddingFunction | None = _NOT_PROVIDED def __post_init__(self): + if self.ivf is not None and self.hnsw is not None: + raise ValueError("Only one of ivf or hnsw can be configured") if self.embedding_function is _NOT_PROVIDED: self.embedding_function = DefaultEmbeddingFunction() + if self.ivf is not None: + self.ivf.__post_init__() if self.hnsw is not None: self.hnsw.__post_init__() diff --git a/src/pyseekdb/client/meta_info.py b/src/pyseekdb/client/meta_info.py index 24d1e0f2..7dd52137 100644 --- a/src/pyseekdb/client/meta_info.py +++ b/src/pyseekdb/client/meta_info.py @@ -55,3 +55,51 @@ def prefix() -> str: @staticmethod def sdk_collections_table_name() -> str: return "sdk_collections" + + +class NamespaceCollectionNames: + _LOGIC_DATA_SUFFIX = "_logic_data_table" + _HOT_SUFFIX = "_hot_table" + _KV_DATA_SUFFIX = "_kv_data_table" + _LOGIC_SCHEMA_SUFFIX = "_logic_schema_table" + _TG_SUFFIX = "_tg" + + @staticmethod + def sdk_ns_namespaces_table() -> str: + return "sdk_ns_namespaces" + + @staticmethod + def sdk_ns_ltables_table() -> str: + return "sdk_ns_ltables" + + @staticmethod + def data_table_name(collection_id: str) -> str: + return f"{collection_id}{NamespaceCollectionNames._LOGIC_DATA_SUFFIX}" + + @staticmethod + def hot_table_name(collection_id: str) -> str: + return f"{collection_id}{NamespaceCollectionNames._HOT_SUFFIX}" + + @staticmethod + def kv_data_table_name(collection_id: str) -> str: + return f"{collection_id}{NamespaceCollectionNames._KV_DATA_SUFFIX}" + + @staticmethod + def logic_schema_table_name(collection_id: str) -> str: + return f"{collection_id}{NamespaceCollectionNames._LOGIC_SCHEMA_SUFFIX}" + + @staticmethod + def tablegroup_name(collection_id: str) -> str: + return f"{collection_id}{NamespaceCollectionNames._TG_SUFFIX}" + + @staticmethod + def is_ns_data_table(table_name: str) -> bool: + return table_name.endswith(NamespaceCollectionNames._LOGIC_DATA_SUFFIX) + + +class NamespaceFieldNames: + NAMESPACE_ID = "namespace_id" + LTABLE_ID = "ltable_id" + DOCUMENT = "document" + EMBEDDING = "embedding" + DATA_CONTENT = "data_content" diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py new file mode 100644 index 00000000..f0179fc2 --- /dev/null +++ b/src/pyseekdb/client/namespace.py @@ -0,0 +1,236 @@ +""" +Namespace class - lightweight facade for namespace-level data operations. + +All operations are delegated to the client that created it via +``self._client._namespace_*()`` methods. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .collection import Collection + from .embedding_function import EmbeddingFunction + + +class Namespace: + def __init__( + self, + client: Any, + collection: "Collection", + name: str, + namespace_id: str, + ): + self._client = client + self._collection = collection + self._name = name + self._namespace_id = namespace_id + + @property + def name(self) -> str: + return self._name + + @property + def namespace_id(self) -> str: + return self._namespace_id + + @property + def collection(self) -> "Collection": + return self._collection + + @property + def embedding_function(self) -> "EmbeddingFunction | None": + return self._collection.embedding_function + + def __repr__(self) -> str: + return ( + f"Namespace(name='{self._name}', namespace_id={self._namespace_id}, " + f"collection='{self._collection.name}')" + ) + + # ==================== DML Operations ==================== + + def add( + self, + ids: str | list[str], + embeddings: list[float] | list[list[float]] | None = None, + metadatas: dict | list[dict] | None = None, + documents: str | list[str] | None = None, + **kwargs, + ) -> None: + return self._client._namespace_add( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + ids=ids, + embeddings=embeddings, + metadatas=metadatas, + documents=documents, + embedding_function=self._collection.embedding_function, + **kwargs, + ) + + def update( + self, + ids: str | list[str], + embeddings: list[float] | list[list[float]] | None = None, + metadatas: dict | list[dict] | None = None, + documents: str | list[str] | None = None, + **kwargs, + ) -> None: + return self._client._namespace_update( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + ids=ids, + embeddings=embeddings, + metadatas=metadatas, + documents=documents, + embedding_function=self._collection.embedding_function, + **kwargs, + ) + + def upsert( + self, + ids: str | list[str], + embeddings: list[float] | list[list[float]] | None = None, + metadatas: dict | list[dict] | None = None, + documents: str | list[str] | None = None, + **kwargs, + ) -> None: + return self._client._namespace_upsert( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + ids=ids, + embeddings=embeddings, + metadatas=metadatas, + documents=documents, + embedding_function=self._collection.embedding_function, + **kwargs, + ) + + def delete( + self, + ids: str | list[str] | None = None, + where: dict[str, Any] | None = None, + where_document: dict[str, Any] | None = None, + **kwargs, + ) -> None: + return self._client._namespace_delete( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + ids=ids, + where=where, + where_document=where_document, + **kwargs, + ) + + # ==================== DQL Operations ==================== + + def query( + self, + query_embeddings: list[float] | list[list[float]] | None = None, + query_texts: str | list[str] | None = None, + n_results: int = 10, + where: dict[str, Any] | None = None, + where_document: dict[str, Any] | None = None, + include: list[str] | None = None, + **kwargs, + ) -> dict[str, Any]: + return self._client._namespace_query( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + query_embeddings=query_embeddings, + query_texts=query_texts, + n_results=n_results, + where=where, + where_document=where_document, + include=include, + embedding_function=self._collection.embedding_function, + distance=self._collection.distance, + **kwargs, + ) + + def hybrid_search( + self, + query: dict[str, Any] | None = None, + knn: dict[str, Any] | None = None, + rank: dict[str, Any] | None = None, + n_results: int = 10, + include: list[str] | None = None, + **kwargs, + ) -> dict[str, Any]: + if include is None and not query and not knn: + include = [] + return self._client._namespace_hybrid_search( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + query=query, + knn=knn, + rank=rank, + n_results=n_results, + include=include, + embedding_function=self._collection.embedding_function, + dimension=self._collection.dimension, + **kwargs, + ) + + def get( + self, + ids: str | list[str] | None = None, + where: dict[str, Any] | None = None, + where_document: dict[str, Any] | None = None, + limit: int | None = None, + offset: int | None = None, + include: list[str] | None = None, + **kwargs, + ) -> dict[str, Any]: + return self._client._namespace_get( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + ids=ids, + where=where, + where_document=where_document, + limit=limit, + offset=offset, + include=include, + **kwargs, + ) + + def count(self) -> int: + return self._client._namespace_count( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + ) + + def peek(self, limit: int = 10) -> dict[str, Any]: + return self._client._namespace_peek( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + limit=limit, + ) + + def prewarm(self) -> None: + return self._client._namespace_prewarm( + collection_id=self._collection.id, + collection_name=self._collection.name, + namespace_id=self._namespace_id, + namespace_name=self._name, + ) diff --git a/src/pyseekdb/client/schema.py b/src/pyseekdb/client/schema.py index 8246233a..ce99adff 100644 --- a/src/pyseekdb/client/schema.py +++ b/src/pyseekdb/client/schema.py @@ -26,7 +26,7 @@ from typing import Any -from .configuration import FulltextIndexConfig, HNSWConfiguration, SparseVectorIndexConfig, VectorIndexConfig +from .configuration import FulltextIndexConfig, HNSWConfiguration, IVFConfiguration, SparseVectorIndexConfig, VectorIndexConfig class Schema: @@ -84,7 +84,7 @@ class Schema: def __init__( self, - vector_index: VectorIndexConfig | HNSWConfiguration | None = None, + vector_index: VectorIndexConfig | HNSWConfiguration | IVFConfiguration | None = None, sparse_vector_index: SparseVectorIndexConfig | None = None, fulltext_index: FulltextIndexConfig | None = None, ): @@ -92,13 +92,15 @@ def __init__( self.vector_index = vector_index elif isinstance(vector_index, HNSWConfiguration): self.vector_index = VectorIndexConfig(hnsw=vector_index) + elif isinstance(vector_index, IVFConfiguration): + self.vector_index = VectorIndexConfig(ivf=vector_index) elif vector_index is None: # Default: will be resolved during create_collection self.vector_index = VectorIndexConfig() else: raise TypeError( f"Unsupported vector index configuration type: {type(vector_index).__name__}. " - f"Expected VectorIndexConfig, HNSWConfiguration, or None." + f"Expected VectorIndexConfig, HNSWConfiguration, IVFConfiguration, or None." ) self.sparse_vector_index = sparse_vector_index self.fulltext_index = fulltext_index @@ -131,6 +133,8 @@ def create_index(self, config: Any) -> Schema: """ if isinstance(config, HNSWConfiguration): self.vector_index = VectorIndexConfig(hnsw=config) + elif isinstance(config, IVFConfiguration): + self.vector_index = VectorIndexConfig(ivf=config) elif isinstance(config, VectorIndexConfig): self.vector_index = config elif isinstance(config, SparseVectorIndexConfig): @@ -140,7 +144,7 @@ def create_index(self, config: Any) -> Schema: else: raise TypeError( f"Unsupported index configuration type: {type(config).__name__}. " - f"Expected VectorIndexConfig, HNSWConfiguration, SparseVectorIndexConfig, or FulltextIndexConfig." + f"Expected VectorIndexConfig, HNSWConfiguration, IVFConfiguration, SparseVectorIndexConfig, or FulltextIndexConfig." ) return self diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py new file mode 100644 index 00000000..71fc25fa --- /dev/null +++ b/src/pyseekdb/client/validators.py @@ -0,0 +1,43 @@ +import re + +_NAME_PATTERN = re.compile(r"^[A-Za-z0-9_]+$") +_MAX_NAME_LENGTH = 512 +_MAX_NAMESPACE_NAME_LENGTH = 256 +_MAX_NAMESPACE_BATCH_SIZE = 100 + + +def _validate_namespace_name(name: str) -> None: + if not isinstance(name, str): + raise TypeError( + f"Invalid namespace name: '{name}'. Namespace name must be a string, got {type(name).__name__}" + ) + if not name: + raise ValueError(f"Invalid namespace name: '{name}'. Namespace name must not be empty") + if len(name) > _MAX_NAMESPACE_NAME_LENGTH: + raise ValueError( + f"Invalid namespace name: '{name}'. Namespace name too long: {len(name)} characters; maximum allowed is {_MAX_NAMESPACE_NAME_LENGTH}." + ) + if _NAME_PATTERN.match(name) is None: + raise ValueError( + f"Invalid namespace name: '{name}'. Namespace name contains invalid characters. " + "Only letters, digits, and underscore are allowed: [a-zA-Z0-9_]" + ) + + +def _validate_record_ids(ids: list[str]) -> None: + for rid in ids: + if not isinstance(rid, str): + raise TypeError( + f"Invalid record id: '{rid}'. Record id must be a string, got {type(rid).__name__}" + ) + if not rid: + raise ValueError("Invalid record id: Record id must not be empty") + if len(rid) > _MAX_NAME_LENGTH: + raise ValueError( + f"Invalid record id: '{rid[:50]}...'. Record id too long: {len(rid)} characters; maximum allowed is {_MAX_NAME_LENGTH}." + ) + if _NAME_PATTERN.match(rid) is None: + raise ValueError( + f"Invalid record id: '{rid}'. Record id contains invalid characters. " + "Only letters, digits, and underscore are allowed: [a-zA-Z0-9_]" + ) diff --git a/src/pyseekdb/schema.py b/src/pyseekdb/schema.py new file mode 100644 index 00000000..44ab2a5f --- /dev/null +++ b/src/pyseekdb/schema.py @@ -0,0 +1,27 @@ +""" +Convenience re-exports for schema and index configuration types. + +Usage:: + + from pyseekdb.schema import Schema, VectorIndexConfig, IVFConfiguration +""" + +from .client.configuration import ( + FulltextIndexConfig, + HNSWConfiguration, + IVFConfiguration, + IVFIndexLib, + IVFIndexType, + VectorIndexConfig, +) +from .client.schema import Schema + +__all__ = [ + "FulltextIndexConfig", + "HNSWConfiguration", + "IVFConfiguration", + "IVFIndexLib", + "IVFIndexType", + "Schema", + "VectorIndexConfig", +] diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index 257c24df..899c5510 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -15,6 +15,8 @@ src_root = repo_root / "src" sys.path.insert(0, str(src_root)) +from unittest.mock import patch + import pyseekdb # noqa: E402 # ==================== Environment Variable Configuration ==================== @@ -268,3 +270,10 @@ def oceanbase_admin_client(): with contextlib.suppress(Exception): if hasattr(client, "close"): client.close() + + +@pytest.fixture(autouse=True) +def _reduce_namespace_partitions(): + """Use a small partition count in integration tests to speed up DDL.""" + with patch("pyseekdb.client.client_base._NS_PARTITION_COUNT", 8): + yield diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py new file mode 100644 index 00000000..64ef17eb --- /dev/null +++ b/tests/integration_tests/test_namespace_dml.py @@ -0,0 +1,167 @@ +""" +Namespace DML integration tests. +Tests namespace.add, update, upsert, delete, get, count, peek operations. +""" + +import time + +import pytest + +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.schema import Schema + + +class TestNamespaceDML: + + def _setup(self, client): + name = f"test_ns_dml_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2"), + embedding_function=None, + ), + ) + collection = client.create_collection(name=name, schema=schema, use_namespace=True) + namespace = collection.create_namespace("dml_ns") + return collection, namespace + + def test_add_single(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0], documents="Hello", metadatas={"tag": "a"}) + result = ns.get(ids="d1") + assert len(result["ids"]) == 1 + assert result["ids"][0] == "d1" + finally: + db_client.delete_collection(name=collection.name) + + def test_add_batch(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add( + ids=["d1", "d2", "d3"], + embeddings=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], + documents=["Doc A", "Doc B", "Doc C"], + metadatas=[{"tag": "a"}, {"tag": "b"}, {"tag": "c"}], + ) + result = ns.get(ids=["d1", "d2", "d3"]) + assert len(result["ids"]) == 3 + finally: + db_client.delete_collection(name=collection.name) + + def test_get_by_id(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add( + ids=["g1", "g2"], + embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + documents=["First", "Second"], + metadatas=[{"k": 1}, {"k": 2}], + ) + result = ns.get(ids="g1", include=["documents", "metadatas"]) + assert len(result["ids"]) == 1 + assert result["documents"][0] == "First" + assert result["metadatas"][0]["k"] == 1 + finally: + db_client.delete_collection(name=collection.name) + + def test_get_with_limit(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add( + ids=["l1", "l2", "l3"], + embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + ) + result = ns.get(limit=2) + assert len(result["ids"]) == 2 + finally: + db_client.delete_collection(name=collection.name) + + def test_update_metadata(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add(ids="u1", embeddings=[1.0, 2.0, 3.0], metadatas={"score": 10}) + ns.update(ids="u1", metadatas={"score": 99}) + result = ns.get(ids="u1", include=["metadatas"]) + assert result["metadatas"][0]["score"] == 99 + finally: + db_client.delete_collection(name=collection.name) + + def test_update_document_and_embedding(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add(ids="u2", embeddings=[1.0, 2.0, 3.0], documents="Original") + ns.update(ids="u2", embeddings=[9.0, 8.0, 7.0], documents="Updated") + result = ns.get(ids="u2", include=["documents", "embeddings"]) + assert result["documents"][0] == "Updated" + finally: + db_client.delete_collection(name=collection.name) + + def test_upsert_existing(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add(ids="up1", embeddings=[1.0, 2.0, 3.0], metadatas={"v": 1}) + ns.upsert(ids="up1", embeddings=[4.0, 5.0, 6.0], metadatas={"v": 2}) + result = ns.get(ids="up1", include=["metadatas"]) + assert result["metadatas"][0]["v"] == 2 + finally: + db_client.delete_collection(name=collection.name) + + def test_upsert_new(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.upsert(ids="up_new", embeddings=[1.0, 1.0, 1.0], metadatas={"fresh": True}) + result = ns.get(ids="up_new", include=["metadatas"]) + assert len(result["ids"]) == 1 + assert result["metadatas"][0]["fresh"] is True + finally: + db_client.delete_collection(name=collection.name) + + def test_delete_by_ids(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add( + ids=["del1", "del2"], + embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + ) + ns.delete(ids="del1") + result = ns.get(ids="del1") + assert len(result["ids"]) == 0 + + result = ns.get(ids="del2") + assert len(result["ids"]) == 1 + finally: + db_client.delete_collection(name=collection.name) + + def test_count(self, db_client): + collection, ns = self._setup(db_client) + try: + assert ns.count() == 0 + ns.add( + ids=["c1", "c2", "c3"], + embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + ) + assert ns.count() == 3 + finally: + db_client.delete_collection(name=collection.name) + + def test_peek(self, db_client): + collection, ns = self._setup(db_client) + try: + ns.add( + ids=["p1", "p2"], + embeddings=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + documents=["Peek A", "Peek B"], + metadatas=[{"x": 1}, {"x": 2}], + ) + result = ns.peek(limit=10) + assert len(result["ids"]) == 2 + assert "documents" in result + assert "metadatas" in result + finally: + db_client.delete_collection(name=collection.name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py new file mode 100644 index 00000000..a9adaf7a --- /dev/null +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -0,0 +1,205 @@ +""" +Namespace lifecycle integration tests. +Tests collection creation with use_namespace=True, namespace CRUD, and collection deletion. +""" + +import time + +import pytest + +import pyseekdb +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.schema import Schema + + +class TestNamespaceLifecycle: + + def _create_ns_collection(self, client, suffix=""): + name = f"test_ns_lc_{int(time.time() * 1000)}{suffix}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine"), + embedding_function=None, + ), + ) + collection = client.create_collection(name=name, schema=schema, use_namespace=True) + return collection + + def test_create_namespace_collection(self, db_client): + collection = self._create_ns_collection(db_client) + try: + assert collection.use_namespace is True + assert collection.dimension == 3 + assert collection.name.startswith("test_ns_lc_") + finally: + db_client.delete_collection(name=collection.name) + + def test_get_collection_preserves_namespace_flag(self, db_client): + collection = self._create_ns_collection(db_client) + try: + retrieved = db_client.get_collection(collection.name) + assert retrieved.use_namespace is True + assert retrieved.dimension == collection.dimension + finally: + db_client.delete_collection(name=collection.name) + + def test_create_and_get_namespace(self, db_client): + collection = self._create_ns_collection(db_client) + try: + ns = collection.create_namespace("ns_a") + assert ns.name == "ns_a" + assert ns.namespace_id is not None + + ns2 = collection.get_namespace("ns_a") + assert ns2.name == "ns_a" + assert ns2.namespace_id == ns.namespace_id + finally: + db_client.delete_collection(name=collection.name) + + def test_get_or_create_namespace(self, db_client): + collection = self._create_ns_collection(db_client) + try: + ns1 = collection.get_or_create_namespace("ns_goc") + assert ns1.name == "ns_goc" + + ns2 = collection.get_or_create_namespace("ns_goc") + assert ns2.namespace_id == ns1.namespace_id + finally: + db_client.delete_collection(name=collection.name) + + def test_has_namespace(self, db_client): + collection = self._create_ns_collection(db_client) + try: + assert collection.has_namespace("nonexistent") is False + collection.create_namespace("ns_check") + assert collection.has_namespace("ns_check") is True + finally: + db_client.delete_collection(name=collection.name) + + def test_list_namespaces(self, db_client): + collection = self._create_ns_collection(db_client) + try: + collection.create_namespace("ns_x") + collection.create_namespace("ns_y") + ns_list = collection.list_namespaces() + names = {ns.name for ns in ns_list} + assert "ns_x" in names + assert "ns_y" in names + assert len(ns_list) >= 2 + finally: + db_client.delete_collection(name=collection.name) + + def test_delete_namespace(self, db_client): + collection = self._create_ns_collection(db_client) + try: + collection.create_namespace("ns_del") + assert collection.has_namespace("ns_del") is True + + collection.delete_namespace("ns_del") + assert collection.has_namespace("ns_del") is False + finally: + db_client.delete_collection(name=collection.name) + + def test_get_nonexistent_namespace_raises(self, db_client): + collection = self._create_ns_collection(db_client) + try: + with pytest.raises(ValueError): + collection.get_namespace("does_not_exist") + finally: + db_client.delete_collection(name=collection.name) + + def test_delete_collection_cleans_namespaces(self, db_client): + collection = self._create_ns_collection(db_client) + coll_name = collection.name + collection.create_namespace("ns_cleanup") + db_client.delete_collection(name=coll_name) + + assert db_client.has_collection(coll_name) is False + + def test_collection_data_api_blocked_when_namespace_enabled(self, db_client): + collection = self._create_ns_collection(db_client) + try: + with pytest.raises(ValueError, match="namespace enabled"): + collection.add(ids="1", embeddings=[1.0, 2.0, 3.0]) + finally: + db_client.delete_collection(name=collection.name) + + def test_namespace_on_non_namespace_collection_raises(self, db_client): + name = f"test_nons_{int(time.time() * 1000)}" + collection = db_client.create_collection( + name=name, + configuration=pyseekdb.HNSWConfiguration(dimension=3), + embedding_function=None, + ) + try: + with pytest.raises(ValueError, match="not enabled"): + collection.create_namespace("ns1") + finally: + try: + db_client.delete_collection(name=name) + except Exception: + db_client._server._execute(f"DROP TABLE IF EXISTS `{name}`") + db_client._server._execute( + f"DELETE FROM `sdk_collections` WHERE COLLECTION_NAME = '{name}'" + ) + + + def test_create_namespace_collection_with_hnsw_raises(self, db_client): + from pyseekdb.client.configuration import HNSWConfiguration + name = f"test_ns_hnsw_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + hnsw=HNSWConfiguration(dimension=3), + embedding_function=None, + ), + ) + with pytest.raises(ValueError, match="HNSW is not allowed"): + db_client.create_collection(name=name, schema=schema, use_namespace=True) + + def test_ss_mode_creates_hot_table(self, oceanbase_client): + """Verify hot_table is created when _is_shared_storage_mode returns True (SS mode).""" + import json + from unittest.mock import patch + from pyseekdb.client.meta_info import NamespaceCollectionNames + + client = oceanbase_client + + with patch.object( + type(client._server), "_is_shared_storage_mode", return_value=True + ): + name = f"test_ns_ss_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine"), + embedding_function=None, + ), + ) + collection = client.create_collection( + name=name, schema=schema, use_namespace=True + ) + + try: + collection_id = collection.id + + # Verify settings has storage_mode=ss + rows = client._server._execute( + f"SELECT settings FROM sdk_collections WHERE collection_id = '{collection_id}'" + ) + settings = json.loads(rows[0]["settings"]) + assert settings["storage_mode"] == "ss", f"Expected ss, got {settings.get('storage_mode')}" + + # Verify hot_table was actually created + hot_table = NamespaceCollectionNames.hot_table_name(collection_id) + rows = client._server._execute(f"DESCRIBE `{hot_table}`") + assert rows is not None and len(rows) > 0, "hot_table should exist in SS mode" + + col_names = {r["Field"] for r in rows} + assert "namespace_id" in col_names + assert "last_access_time" in col_names + finally: + client.delete_collection(name=name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py new file mode 100644 index 00000000..ddbeda2d --- /dev/null +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -0,0 +1,61 @@ +""" +Namespace prewarm integration tests. +Tests that prewarm raises ValueError in embedded mode and executes in remote mode. +""" + +import time + +import pytest + +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.schema import Schema + + +class TestNamespacePrewarm: + + def _create_ns_collection_and_namespace(self, client): + name = f"test_ns_pw_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine"), + embedding_function=None, + ), + ) + collection = client.create_collection(name=name, schema=schema, use_namespace=True) + namespace = collection.create_namespace("pw_ns") + return collection, namespace + + def test_prewarm_raises_in_embedded_mode(self, embedded_client): + collection, namespace = self._create_ns_collection_and_namespace(embedded_client) + try: + with pytest.raises(ValueError, match="shared-storage remote"): + namespace.prewarm() + finally: + embedded_client.delete_collection(name=collection.name) + + def test_prewarm_executes_in_server_mode(self, server_client): + """Prewarm should not raise in remote mode (actual behavior depends on backend).""" + collection, namespace = self._create_ns_collection_and_namespace(server_client) + try: + namespace.prewarm() + except Exception as e: + if "shared-storage remote" in str(e): + pytest.fail("prewarm should not raise embedded-mode error in server mode") + finally: + server_client.delete_collection(name=collection.name) + + def test_prewarm_executes_in_oceanbase_mode(self, oceanbase_client): + """Prewarm should not raise in OceanBase mode (actual behavior depends on backend).""" + collection, namespace = self._create_ns_collection_and_namespace(oceanbase_client) + try: + namespace.prewarm() + except Exception as e: + if "shared-storage remote" in str(e): + pytest.fail("prewarm should not raise embedded-mode error in oceanbase mode") + finally: + oceanbase_client.delete_collection(name=collection.name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py new file mode 100644 index 00000000..d92b578f --- /dev/null +++ b/tests/integration_tests/test_namespace_query.py @@ -0,0 +1,290 @@ +""" +Namespace query integration tests. +Tests vector similarity query, metadata filtering, include control, and multi-namespace isolation. +""" + +import time + +import pytest + +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.schema import Schema + + +class TestNamespaceQuery: + + def _setup(self, client, suffix=""): + name = f"test_ns_q_{int(time.time() * 1000)}{suffix}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2"), + embedding_function=None, + ), + ) + collection = client.create_collection(name=name, schema=schema, use_namespace=True) + return collection + + def _insert_data(self, ns): + ns.add( + ids=["q1", "q2", "q3", "q4", "q5"], + embeddings=[ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [0.0, 1.0, 1.0], + ], + documents=[ + "Machine learning basics", + "Python programming guide", + "OceanBase distributed database", + "Advanced ML algorithms", + "Data science with Python", + ], + metadatas=[ + {"category": "AI", "score": 95}, + {"category": "Programming", "score": 88}, + {"category": "Database", "score": 92}, + {"category": "AI", "score": 90}, + {"category": "Data Science", "score": 85}, + ], + ) + + def test_basic_vector_query(self, db_client): + collection = self._setup(db_client) + ns = collection.create_namespace("qns") + try: + self._insert_data(ns) + result = ns.query(query_embeddings=[1.0, 0.0, 0.0], n_results=3) + assert result is not None + assert "ids" in result + assert len(result["ids"]) > 0 + assert len(result["ids"][0]) <= 3 + finally: + db_client.delete_collection(name=collection.name) + + def test_query_with_metadata_filter(self, db_client): + collection = self._setup(db_client) + ns = collection.create_namespace("qns_meta") + try: + self._insert_data(ns) + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + n_results=5, + where={"category": "AI"}, + ) + assert result is not None + assert len(result["ids"][0]) > 0 + if "metadatas" in result and result["metadatas"]: + for meta in result["metadatas"][0]: + assert meta["category"] == "AI" + finally: + db_client.delete_collection(name=collection.name) + + def test_query_with_score_filter(self, db_client): + collection = self._setup(db_client) + ns = collection.create_namespace("qns_score") + try: + self._insert_data(ns) + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + n_results=5, + where={"score": {"$gte": 90}}, + ) + assert result is not None + assert len(result["ids"][0]) > 0 + finally: + db_client.delete_collection(name=collection.name) + + def test_query_with_include(self, db_client): + collection = self._setup(db_client) + ns = collection.create_namespace("qns_inc") + try: + self._insert_data(ns) + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + n_results=3, + include=["documents", "metadatas"], + ) + assert "documents" in result + assert "metadatas" in result + assert len(result["documents"][0]) == len(result["ids"][0]) + assert len(result["metadatas"][0]) == len(result["ids"][0]) + finally: + db_client.delete_collection(name=collection.name) + + def test_query_include_embeddings(self, db_client): + collection = self._setup(db_client) + ns = collection.create_namespace("qns_emb") + try: + self._insert_data(ns) + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + n_results=2, + include=["embeddings"], + ) + assert "embeddings" in result + if result["embeddings"] and result["embeddings"][0]: + assert len(result["embeddings"][0][0]) == 3 + finally: + db_client.delete_collection(name=collection.name) + + def test_multi_namespace_isolation(self, db_client): + collection = self._setup(db_client, suffix="_iso") + ns_a = collection.create_namespace("tenant_a") + ns_b = collection.create_namespace("tenant_b") + try: + ns_a.add( + ids=["a1"], + embeddings=[[1.0, 0.0, 0.0]], + documents=["Data from tenant A"], + metadatas=[{"owner": "A"}], + ) + ns_b.add( + ids=["b1"], + embeddings=[[0.0, 1.0, 0.0]], + documents=["Data from tenant B"], + metadatas=[{"owner": "B"}], + ) + + result_a = ns_a.query(query_embeddings=[1.0, 0.0, 0.0], n_results=10) + result_b = ns_b.query(query_embeddings=[1.0, 0.0, 0.0], n_results=10) + + ids_a = result_a["ids"][0] if result_a["ids"] else [] + ids_b = result_b["ids"][0] if result_b["ids"] else [] + + assert "a1" in ids_a + assert "b1" not in ids_a + assert "b1" in ids_b + assert "a1" not in ids_b + finally: + db_client.delete_collection(name=collection.name) + + def test_same_id_different_namespaces(self, db_client): + collection = self._setup(db_client, suffix="_sameid") + ns_x = collection.create_namespace("ns_x") + ns_y = collection.create_namespace("ns_y") + try: + ns_x.add(ids="shared_id", embeddings=[1.0, 0.0, 0.0], metadatas={"src": "X"}) + ns_y.add(ids="shared_id", embeddings=[0.0, 1.0, 0.0], metadatas={"src": "Y"}) + + res_x = ns_x.get(ids="shared_id", include=["metadatas"]) + res_y = ns_y.get(ids="shared_id", include=["metadatas"]) + + assert len(res_x["ids"]) == 1 + assert res_x["metadatas"][0]["src"] == "X" + assert len(res_y["ids"]) == 1 + assert res_y["metadatas"][0]["src"] == "Y" + finally: + db_client.delete_collection(name=collection.name) + + + # ==================== hybrid_search tests ==================== + + def test_hybrid_search_fulltext_only(self, db_client): + collection = self._setup(db_client) + ns = collection.create_namespace("qns_hs_ft") + try: + self._insert_data(ns) + time.sleep(1) + result = ns.hybrid_search( + query={"where_document": {"$contains": "machine learning"}}, + n_results=3, + include=["documents", "metadatas"], + ) + assert result is not None + assert "ids" in result + assert len(result["ids"]) > 0 + assert len(result["ids"][0]) > 0 + if "documents" in result and result["documents"][0]: + for doc in result["documents"][0]: + if doc: + assert "machine" in doc.lower() or "learning" in doc.lower() + finally: + db_client.delete_collection(name=collection.name) + + def test_hybrid_search_vector_only(self, db_client): + collection = self._setup(db_client) + ns = collection.create_namespace("qns_hs_vec") + try: + self._insert_data(ns) + time.sleep(1) + result = ns.hybrid_search( + knn={"query_embeddings": [1.0, 0.0, 0.0], "n_results": 3}, + n_results=3, + include=["documents"], + ) + assert result is not None + assert "ids" in result + assert "distances" in result + assert len(result["ids"]) > 0 + assert len(result["ids"][0]) > 0 + for dist in result["distances"][0]: + assert dist >= 0 + finally: + db_client.delete_collection(name=collection.name) + + def test_hybrid_search_combined(self, db_client): + collection = self._setup(db_client) + ns = collection.create_namespace("qns_hs_comb") + try: + self._insert_data(ns) + time.sleep(1) + result = ns.hybrid_search( + query={"where_document": {"$contains": "machine"}, "n_results": 5}, + knn={"query_embeddings": [1.0, 0.0, 0.0], "n_results": 5}, + rank={"rrf": {"rank_window_size": 60, "rank_constant": 60}}, + n_results=3, + include=["documents", "metadatas"], + ) + assert result is not None + assert "ids" in result + assert len(result["ids"]) > 0 + assert len(result["ids"][0]) > 0 + finally: + db_client.delete_collection(name=collection.name) + + def test_hybrid_search_namespace_isolation(self, db_client): + collection = self._setup(db_client, suffix="_hs_iso") + ns_a = collection.create_namespace("hs_tenant_a") + ns_b = collection.create_namespace("hs_tenant_b") + try: + ns_a.add( + ids=["ha1"], + embeddings=[[1.0, 0.0, 0.0]], + documents=["Machine learning from tenant A"], + metadatas=[{"owner": "A"}], + ) + ns_b.add( + ids=["hb1"], + embeddings=[[0.0, 1.0, 0.0]], + documents=["Python programming from tenant B"], + metadatas=[{"owner": "B"}], + ) + time.sleep(1) + + result_a = ns_a.hybrid_search( + knn={"query_embeddings": [1.0, 0.0, 0.0], "n_results": 10}, + n_results=10, + include=["documents"], + ) + result_b = ns_b.hybrid_search( + knn={"query_embeddings": [1.0, 0.0, 0.0], "n_results": 10}, + n_results=10, + include=["documents"], + ) + + ids_a = result_a["ids"][0] if result_a["ids"] else [] + ids_b = result_b["ids"][0] if result_b["ids"] else [] + + assert "ha1" in ids_a + assert "hb1" not in ids_a + assert "hb1" in ids_b + assert "ha1" not in ids_b + finally: + db_client.delete_collection(name=collection.name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py new file mode 100644 index 00000000..f30c38ad --- /dev/null +++ b/tests/unit_tests/test_namespace.py @@ -0,0 +1,958 @@ +""" +Unit tests for namespace-related classes: +- IVFConfiguration validation +- VectorIndexConfig mutual exclusion +- Collection namespace guard +- Namespace object properties +""" + +import os +import sys +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +project_root = Path(__file__).parent.parent.parent +sys.path.insert(0, str(project_root / "src")) + +from pyseekdb import IVFConfiguration # noqa: E402 +from pyseekdb.client.client_base import BaseClient # noqa: E402 +from pyseekdb.client.collection import Collection # noqa: E402 +from pyseekdb.client.configuration import HNSWConfiguration, IVFIndexType, VectorIndexConfig # noqa: E402 +from pyseekdb.client.namespace import Namespace # noqa: E402 + + +# ==================== IVFConfiguration Tests ==================== + + +class TestIVFConfiguration: + + def test_valid_defaults(self): + config = IVFConfiguration() + assert config.dimension == 384 + assert config.distance == "cosine" + assert config.use_spfresh is None + assert config.properties is None + + def test_valid_custom(self): + config = IVFConfiguration(dimension=128, distance="l2", use_spfresh=True) + assert config.dimension == 128 + assert config.distance == "l2" + assert config.use_spfresh is True + + def test_valid_inner_product(self): + config = IVFConfiguration(dimension=1024, distance="inner_product") + assert config.distance == "inner_product" + + def test_dimension_boundary_min(self): + config = IVFConfiguration(dimension=1) + assert config.dimension == 1 + + def test_dimension_boundary_max(self): + config = IVFConfiguration(dimension=4096) + assert config.dimension == 4096 + + def test_invalid_dimension_zero(self): + with pytest.raises(ValueError, match="must be between"): + IVFConfiguration(dimension=0) + + def test_invalid_dimension_negative(self): + with pytest.raises(ValueError, match="must be between"): + IVFConfiguration(dimension=-1) + + def test_invalid_dimension_too_large(self): + with pytest.raises(ValueError, match="must be between"): + IVFConfiguration(dimension=4097) + + def test_invalid_dimension_type(self): + with pytest.raises(TypeError, match="dimension must be an integer"): + IVFConfiguration(dimension="128") + + def test_invalid_dimension_bool(self): + with pytest.raises(TypeError, match="dimension must be an integer"): + IVFConfiguration(dimension=True) + + def test_invalid_distance(self): + with pytest.raises(ValueError, match="distance must be one of"): + IVFConfiguration(distance="invalid") + + def test_invalid_use_spfresh_type(self): + with pytest.raises(TypeError, match="use_spfresh must be a bool"): + IVFConfiguration(use_spfresh="yes") + + def test_properties_valid(self): + config = IVFConfiguration(properties={"nlist": 128, "nprobe": 16}) + assert config.properties["nlist"] == 128 + + def test_valid_type_default(self): + config = IVFConfiguration() + assert config.type == "ivf_flat" + + def test_valid_type_ivf_sq8(self): + config = IVFConfiguration(type="ivf_sq8") + assert config.type == "ivf_sq8" + + def test_valid_type_ivf_pq(self): + config = IVFConfiguration(type="ivf_pq") + assert config.type == "ivf_pq" + + def test_valid_type_enum(self): + config = IVFConfiguration(type=IVFIndexType.IVF_SQ8) + assert config.type == "ivf_sq8" + + def test_invalid_type(self): + with pytest.raises(ValueError, match="type must be one of"): + IVFConfiguration(type="invalid") + + def test_valid_lib_default(self): + config = IVFConfiguration() + assert config.lib == "ob" + + def test_valid_lib_vsag(self): + config = IVFConfiguration(lib="vsag") + assert config.lib == "vsag" + + def test_valid_lib_enum(self): + from pyseekdb.client.configuration import IVFIndexLib + config = IVFConfiguration(lib=IVFIndexLib.VSAG) + assert config.lib == "vsag" + + def test_invalid_lib(self): + with pytest.raises(ValueError, match="lib must be one of"): + IVFConfiguration(lib="invalid") + + def test_properties_invalid_nested(self): + with pytest.raises(TypeError): + IVFConfiguration(properties={"bad": {"nested": True}}) + + +# ==================== VectorIndexConfig Tests ==================== + + +class TestVectorIndexConfig: + + def test_ivf_only(self): + ivf = IVFConfiguration(dimension=128) + config = VectorIndexConfig(ivf=ivf, embedding_function=None) + assert config.ivf is ivf + assert config.hnsw is None + + def test_hnsw_only(self): + hnsw = HNSWConfiguration(dimension=128) + config = VectorIndexConfig(hnsw=hnsw, embedding_function=None) + assert config.hnsw is hnsw + assert config.ivf is None + + def test_ivf_and_hnsw_mutually_exclusive(self): + ivf = IVFConfiguration(dimension=128) + hnsw = HNSWConfiguration(dimension=128) + with pytest.raises(ValueError, match="Only one of ivf or hnsw"): + VectorIndexConfig(ivf=ivf, hnsw=hnsw, embedding_function=None) + + def test_neither_ivf_nor_hnsw(self): + config = VectorIndexConfig(embedding_function=None) + assert config.ivf is None + assert config.hnsw is None + + +# ==================== Collection Namespace Guard Tests ==================== + + +class TestCollectionNamespaceGuard: + + def _make_collection(self, use_namespace: bool) -> Collection: + mock_client = MagicMock() + return Collection( + client=mock_client, + name="test_coll", + collection_id="1", + dimension=3, + use_namespace=use_namespace, + ) + + def test_guard_blocks_add_when_namespace_enabled(self): + coll = self._make_collection(use_namespace=True) + with pytest.raises(ValueError, match="namespace enabled"): + coll.add(ids="1", embeddings=[1.0, 2.0, 3.0]) + + def test_guard_blocks_update_when_namespace_enabled(self): + coll = self._make_collection(use_namespace=True) + with pytest.raises(ValueError, match="namespace enabled"): + coll.update(ids="1", metadatas={"k": "v"}) + + def test_guard_blocks_upsert_when_namespace_enabled(self): + coll = self._make_collection(use_namespace=True) + with pytest.raises(ValueError, match="namespace enabled"): + coll.upsert(ids="1", embeddings=[1.0, 2.0, 3.0]) + + def test_guard_blocks_delete_when_namespace_enabled(self): + coll = self._make_collection(use_namespace=True) + with pytest.raises(ValueError, match="namespace enabled"): + coll.delete(ids="1") + + def test_guard_blocks_query_when_namespace_enabled(self): + coll = self._make_collection(use_namespace=True) + with pytest.raises(ValueError, match="namespace enabled"): + coll.query(query_embeddings=[1.0, 2.0, 3.0]) + + def test_guard_blocks_get_when_namespace_enabled(self): + coll = self._make_collection(use_namespace=True) + with pytest.raises(ValueError, match="namespace enabled"): + coll.get(ids="1") + + def test_guard_blocks_count_when_namespace_enabled(self): + coll = self._make_collection(use_namespace=True) + with pytest.raises(ValueError, match="namespace enabled"): + coll.count() + + def test_guard_blocks_peek_when_namespace_enabled(self): + coll = self._make_collection(use_namespace=True) + with pytest.raises(ValueError, match="namespace enabled"): + coll.peek() + + +# ==================== Collection Namespace Management Guard Tests ==================== + + +class TestCollectionNamespaceManagementGuard: + + def _make_collection(self, use_namespace: bool) -> Collection: + mock_client = MagicMock() + return Collection( + client=mock_client, + name="test_coll", + collection_id="1", + dimension=3, + use_namespace=use_namespace, + ) + + def test_create_namespace_requires_namespace_enabled(self): + coll = self._make_collection(use_namespace=False) + with pytest.raises(ValueError, match="not enabled"): + coll.create_namespace("ns1") + + def test_get_namespace_requires_namespace_enabled(self): + coll = self._make_collection(use_namespace=False) + with pytest.raises(ValueError, match="not enabled"): + coll.get_namespace("ns1") + + def test_get_or_create_namespace_requires_namespace_enabled(self): + coll = self._make_collection(use_namespace=False) + with pytest.raises(ValueError, match="not enabled"): + coll.get_or_create_namespace("ns1") + + def test_delete_namespace_requires_namespace_enabled(self): + coll = self._make_collection(use_namespace=False) + with pytest.raises(ValueError, match="not enabled"): + coll.delete_namespace("ns1") + + def test_list_namespaces_requires_namespace_enabled(self): + coll = self._make_collection(use_namespace=False) + with pytest.raises(ValueError, match="not enabled"): + coll.list_namespaces() + + def test_has_namespace_requires_namespace_enabled(self): + coll = self._make_collection(use_namespace=False) + with pytest.raises(ValueError, match="not enabled"): + coll.has_namespace("ns1") + + +# ==================== Namespace Object Tests ==================== + + +class TestNamespaceObject: + + def _make_namespace(self) -> Namespace: + mock_client = MagicMock() + mock_ef = MagicMock() + coll = Collection( + client=mock_client, + name="test_coll", + collection_id="1", + dimension=3, + embedding_function=mock_ef, + use_namespace=True, + ) + return Namespace(client=mock_client, collection=coll, name="ns1", namespace_id="100") + + def test_name_property(self): + ns = self._make_namespace() + assert ns.name == "ns1" + + def test_namespace_id_property(self): + ns = self._make_namespace() + assert ns.namespace_id == "100" + + def test_collection_property(self): + ns = self._make_namespace() + assert ns.collection.name == "test_coll" + + def test_embedding_function_inherited(self): + ns = self._make_namespace() + assert ns.embedding_function is ns.collection.embedding_function + + def test_repr(self): + ns = self._make_namespace() + r = repr(ns) + assert "ns1" in r + assert "100" in r + assert "test_coll" in r + + def test_add_delegates_to_client(self): + ns = self._make_namespace() + ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0]) + ns._client._namespace_add.assert_called_once() + + def test_update_delegates_to_client(self): + ns = self._make_namespace() + ns.update(ids="d1", metadatas={"k": "v"}) + ns._client._namespace_update.assert_called_once() + + def test_upsert_delegates_to_client(self): + ns = self._make_namespace() + ns.upsert(ids="d1", embeddings=[1.0, 2.0, 3.0]) + ns._client._namespace_upsert.assert_called_once() + + def test_delete_delegates_to_client(self): + ns = self._make_namespace() + ns.delete(ids="d1") + ns._client._namespace_delete.assert_called_once() + + def test_query_delegates_to_client(self): + ns = self._make_namespace() + ns._client._namespace_query.return_value = {"ids": [["d1"]]} + ns.query(query_embeddings=[1.0, 2.0, 3.0]) + ns._client._namespace_query.assert_called_once() + + def test_get_delegates_to_client(self): + ns = self._make_namespace() + ns._client._namespace_get.return_value = {"ids": ["d1"]} + ns.get(ids="d1") + ns._client._namespace_get.assert_called_once() + + def test_count_delegates_to_client(self): + ns = self._make_namespace() + ns._client._namespace_count.return_value = 42 + assert ns.count() == 42 + + def test_peek_delegates_to_client(self): + ns = self._make_namespace() + ns._client._namespace_peek.return_value = {"ids": ["d1"]} + ns.peek(limit=5) + ns._client._namespace_peek.assert_called_once() + + def test_prewarm_delegates_to_client(self): + ns = self._make_namespace() + ns.prewarm() + ns._client._namespace_prewarm.assert_called_once() + + +# ==================== Namespace SQL Generation Tests ==================== + + +class FakeClient(BaseClient): + """Concrete BaseClient subclass that captures SQL without executing.""" + + def __init__(self): + self.executed_sqls = [] + self.query_sqls = [] + self.query_return_value = [] + + def _ensure_connection(self): + mock_conn = MagicMock() + client = self + + class CaptureCursor: + def execute(self, sql, params=None): + resolved = sql + if params: + for p in params: + resolved = resolved.replace("%s", repr(p), 1) + if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): + print(f"[pyseekdb SQL] {resolved}", flush=True) + client.executed_sqls.append(resolved) + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def close(self): + pass + + mock_conn.cursor.return_value = CaptureCursor() + return mock_conn + + def _use_context_manager_for_cursor(self): + return True + + def _execute(self, sql, params=None): + if params: + for p in params: + sql = sql.replace("%s", repr(p), 1) + if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): + print(f"[pyseekdb SQL] {sql}", flush=True) + self.executed_sqls.append(sql) + return None + + def _execute_query_with_cursor(self, conn, sql, params, use_context_manager=True): + resolved = sql + for p in params: + resolved = resolved.replace("%s", repr(p), 1) + if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): + print(f"[pyseekdb SQL] {resolved}", flush=True) + self.query_sqls.append(resolved) + return self.query_return_value + + def is_connected(self): + return True + + def _cleanup(self): + pass + + def get_raw_connection(self): + return None + + @property + def mode(self): + return "FakeClient" + + def create_collection(self, *a, **kw): + pass + + def get_collection(self, *a, **kw): + pass + + def delete_collection(self, *a, **kw): + pass + + def list_collections(self): + return [] + + def has_collection(self, name): + return False + + def create_database(self, *a, **kw): + pass + + def get_database(self, *a, **kw): + pass + + def delete_database(self, *a, **kw): + pass + + def list_databases(self, *a, **kw): + return [] + + def fork_database(self, *a, **kw): + pass + + +class TestNamespaceSQLGeneration: + """Verify SQL statements generated by BaseClient._namespace_* methods.""" + + COLLECTION_ID = "42" + NAMESPACE_ID = "7" + TABLE = "42_logic_data_table" + + def _client(self): + return FakeClient() + + def _common_kwargs(self): + return dict( + collection_id=self.COLLECTION_ID, + collection_name="test_coll", + namespace_id=self.NAMESPACE_ID, + namespace_name="test_ns", + ) + + # ---- ADD ---- + + def test_add_single_sql(self): + c = self._client() + c._namespace_add( + **self._common_kwargs(), + ids="d1", + embeddings=[1.0, 2.0, 3.0], + documents="hello", + metadatas={"tag": "a"}, + ) + sql = c.executed_sqls[-1] + assert sql.startswith("INSERT INTO") + assert f"`{self.TABLE}`" in sql + assert "(namespace_id, ltable_id, document, embedding, data_content)" in sql + assert "(7, 1, 'hello'," in sql + assert '\\"id\\": \\"d1\\"' in sql or '"id": "d1"' in sql + assert '\\"tag\\": \\"a\\"' in sql or '"tag": "a"' in sql + + def test_add_batch_sql(self): + c = self._client() + c._namespace_add( + **self._common_kwargs(), + ids=["d1", "d2", "d3"], + embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + documents=["A", "B", "C"], + metadatas=[{"k": 1}, {"k": 2}, {"k": 3}], + ) + sql = c.executed_sqls[-1] + assert sql.count("(7, 1,") == 3 + assert '\\"id\\": \\"d1\\"' in sql or '"id": "d1"' in sql + assert '\\"id\\": \\"d2\\"' in sql or '"id": "d2"' in sql + assert '\\"id\\": \\"d3\\"' in sql or '"id": "d3"' in sql + + def test_add_without_metadata_sql(self): + c = self._client() + c._namespace_add( + **self._common_kwargs(), + ids="d1", + embeddings=[1.0, 2.0, 3.0], + ) + sql = c.executed_sqls[-1] + assert '\\"id\\": \\"d1\\"' in sql or '"id": "d1"' in sql + assert "metadata" not in sql.replace("data_content", "") + + # ---- UPDATE ---- + + def test_update_metadata_sql(self): + c = self._client() + c._namespace_update( + **self._common_kwargs(), + ids="d1", + metadatas={"score": 99}, + ) + sql = c.executed_sqls[-1] + assert sql.startswith("UPDATE") + assert f"`{self.TABLE}`" in sql + assert "JSON_SET(data_content, '$.metadata'," in sql + assert "CAST(" in sql + assert "AS JSON)" in sql + assert "namespace_id = 7" in sql + assert """JSON_EXTRACT(data_content, '$.id') = 'd1'""" in sql + + def test_update_embedding_and_document_sql(self): + c = self._client() + c._namespace_update( + **self._common_kwargs(), + ids="d1", + embeddings=[9.0, 8.0, 7.0], + documents="Updated", + ) + sql = c.executed_sqls[-1] + assert "document = 'Updated'" in sql + assert "embedding = X'" in sql + assert "namespace_id = 7" in sql + + # ---- DELETE ---- + + def test_delete_by_ids_sql(self): + c = self._client() + c._namespace_delete( + **self._common_kwargs(), + ids="d1", + ) + sql = c.executed_sqls[-1] + assert sql.startswith("DELETE FROM") + assert f"`{self.TABLE}`" in sql + assert "namespace_id = 7" in sql + assert """JSON_EXTRACT(data_content, '$.id') = 'd1'""" in sql + + def test_delete_by_where_sql(self): + c = self._client() + c._namespace_delete( + **self._common_kwargs(), + where={"category": "AI"}, + ) + sql = c.executed_sqls[-1] + assert "DELETE FROM" in sql + assert "namespace_id = 7" in sql + assert "JSON_EXTRACT" in sql or "JSON_OVERLAPS" in sql + assert "metadata.category" in sql + + # ---- QUERY ---- + + def test_query_basic_sql(self): + c = self._client() + c.query_return_value = [] + c._namespace_query( + **self._common_kwargs(), + query_embeddings=[1.0, 0.0, 0.0], + n_results=3, + distance="l2", + ) + sql = c.query_sqls[-1] + assert "SELECT" in sql + assert "JSON_EXTRACT(data_content, '$.id') AS record_id" in sql + assert "l2_distance(embedding," in sql + assert "AS distance" in sql + assert f"`{self.TABLE}`" in sql + assert "namespace_id = 7" in sql + assert "ORDER BY l2_distance(embedding," in sql + assert "APPROXIMATE LIMIT" in sql + + def test_query_with_where_sql(self): + c = self._client() + c.query_return_value = [] + c._namespace_query( + **self._common_kwargs(), + query_embeddings=[1.0, 0.0, 0.0], + n_results=5, + where={"category": "AI"}, + distance="cosine", + ) + sql = c.query_sqls[-1] + assert "cosine_distance(embedding," in sql + assert "metadata.category" in sql + + # ---- GET ---- + + def test_get_by_ids_sql(self): + c = self._client() + c.query_return_value = [] + c._namespace_get( + **self._common_kwargs(), + ids="g1", + ) + sql = c.query_sqls[-1] + assert "SELECT" in sql + assert f"`{self.TABLE}`" in sql + assert "namespace_id = 7" in sql + assert """JSON_EXTRACT(data_content, '$.id') = 'g1'""" in sql + + def test_get_with_limit_sql(self): + c = self._client() + c.query_return_value = [] + c._namespace_get( + **self._common_kwargs(), + limit=10, + ) + sql = c.query_sqls[-1] + assert "LIMIT 10" in sql + + # ---- COUNT ---- + + def test_count_sql(self): + c = self._client() + c.query_return_value = [{"cnt": 0}] + result = c._namespace_count(**self._common_kwargs()) + sql = c.query_sqls[-1] + assert "SELECT COUNT(*) AS cnt" in sql + assert f"`{self.TABLE}`" in sql + assert "namespace_id = 7" in sql + assert "ltable_id = 1" in sql + assert result == 0 + + + # ---- IVF TYPE in SQL ---- + + def test_ivf_type_default_in_sql(self): + """Verify _get_ivf_vector_index_sql uses TYPE=IVF_FLAT and LIB=OB by default.""" + from pyseekdb.client.client_base import _get_ivf_vector_index_sql + + config = IVFConfiguration(dimension=3, distance="cosine") + sql = _get_ivf_vector_index_sql(config) + assert "TYPE=IVF_FLAT" in sql + assert "LIB=OB" in sql + assert "TYPE=ivf," not in sql + + def test_ivf_type_sq8_in_sql(self): + """Verify _get_ivf_vector_index_sql uses TYPE=IVF_SQ8 when specified.""" + from pyseekdb.client.client_base import _get_ivf_vector_index_sql + + config = IVFConfiguration(dimension=3, distance="l2", type="ivf_sq8") + sql = _get_ivf_vector_index_sql(config) + assert "TYPE=IVF_SQ8" in sql + + def test_ivf_type_pq_in_sql(self): + """Verify _get_ivf_vector_index_sql uses TYPE=IVF_PQ when specified.""" + from pyseekdb.client.client_base import _get_ivf_vector_index_sql + + config = IVFConfiguration(dimension=3, distance="inner_product", type="ivf_pq") + sql = _get_ivf_vector_index_sql(config) + assert "TYPE=IVF_PQ" in sql + + def test_ivf_lib_vsag_in_sql(self): + """Verify _get_ivf_vector_index_sql uses LIB=VSAG when specified.""" + from pyseekdb.client.client_base import _get_ivf_vector_index_sql + + config = IVFConfiguration(dimension=3, distance="cosine", lib="vsag") + sql = _get_ivf_vector_index_sql(config) + assert "LIB=VSAG" in sql + + def test_ivf_lib_ob_in_sql(self): + """Verify _get_ivf_vector_index_sql uses LIB=OB by default.""" + from pyseekdb.client.client_base import _get_ivf_vector_index_sql + + config = IVFConfiguration(dimension=3, distance="cosine") + sql = _get_ivf_vector_index_sql(config) + assert "LIB=OB" in sql + + +# ==================== Namespace Name Validation Tests ==================== + + +class TestValidateNamespaceName: + + def test_valid_simple_name(self): + from pyseekdb.client.client_base import _validate_namespace_name + _validate_namespace_name("my_namespace") + + def test_valid_with_digits_and_underscore(self): + from pyseekdb.client.client_base import _validate_namespace_name + _validate_namespace_name("tenant_123_abc") + + def test_valid_single_char(self): + from pyseekdb.client.client_base import _validate_namespace_name + _validate_namespace_name("a") + + def test_valid_boundary_256_chars(self): + from pyseekdb.client.client_base import _validate_namespace_name + _validate_namespace_name("a" * 256) + + def test_empty_name_raises(self): + from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="must not be empty"): + _validate_namespace_name("") + + def test_non_string_raises(self): + from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(TypeError, match="must be a string"): + _validate_namespace_name(123) + + def test_too_long_raises(self): + from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="too long"): + _validate_namespace_name("a" * 257) + + def test_hyphen_raises(self): + from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="invalid characters"): + _validate_namespace_name("my-namespace") + + def test_space_raises(self): + from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="invalid characters"): + _validate_namespace_name("my namespace") + + def test_dot_raises(self): + from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="invalid characters"): + _validate_namespace_name("my.namespace") + + def test_chinese_raises(self): + from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="invalid characters"): + _validate_namespace_name("命名空间") + + def test_collection_create_namespace_validates_name(self): + mock_client = MagicMock() + coll = Collection(client=mock_client, name="c", collection_id="1", dimension=3, use_namespace=True) + with pytest.raises(ValueError, match="invalid characters"): + coll.create_namespace("bad-name") + + def test_collection_get_namespace_validates_name(self): + mock_client = MagicMock() + coll = Collection(client=mock_client, name="c", collection_id="1", dimension=3, use_namespace=True) + with pytest.raises(ValueError, match="must not be empty"): + coll.get_namespace("") + + def test_collection_has_namespace_validates_name(self): + mock_client = MagicMock() + coll = Collection(client=mock_client, name="c", collection_id="1", dimension=3, use_namespace=True) + with pytest.raises(TypeError, match="must be a string"): + coll.has_namespace(42) + + +# ==================== Record ID Validation Tests ==================== + + +class TestValidateRecordIds: + + def test_valid_single_id(self): + from pyseekdb.client.client_base import _validate_record_ids + _validate_record_ids(["doc_1"]) + + def test_valid_multiple_ids(self): + from pyseekdb.client.client_base import _validate_record_ids + _validate_record_ids(["id1", "id2", "id_3"]) + + def test_valid_boundary_512_chars(self): + from pyseekdb.client.client_base import _validate_record_ids + _validate_record_ids(["a" * 512]) + + def test_empty_id_raises(self): + from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="must not be empty"): + _validate_record_ids([""]) + + def test_non_string_id_raises(self): + from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(TypeError, match="must be a string"): + _validate_record_ids([123]) + + def test_too_long_id_raises(self): + from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="too long"): + _validate_record_ids(["a" * 513]) + + def test_invalid_chars_raises(self): + from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="invalid characters"): + _validate_record_ids(["doc-1"]) + + def test_space_in_id_raises(self): + from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="invalid characters"): + _validate_record_ids(["doc 1"]) + + def test_mixed_valid_and_invalid_raises_on_first_bad(self): + from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="invalid characters"): + _validate_record_ids(["good_id", "bad-id"]) + + +# ==================== Namespace Batch Limit Tests ==================== + + +class TestNamespaceBatchLimit: + + def _client(self): + return FakeClient() + + def _common_kwargs(self): + return dict( + collection_id="42", + collection_name="test_coll", + namespace_id="7", + namespace_name="test_ns", + ) + + def test_add_single_record_ok(self): + c = self._client() + c._namespace_add(**self._common_kwargs(), ids="d1", embeddings=[1.0, 2.0, 3.0]) + + def test_add_100_records_ok(self): + c = self._client() + ids = [f"d{i}" for i in range(100)] + embeddings = [[float(i)] * 3 for i in range(100)] + c._namespace_add(**self._common_kwargs(), ids=ids, embeddings=embeddings) + + def test_add_101_records_raises(self): + c = self._client() + ids = [f"d{i}" for i in range(101)] + embeddings = [[float(i)] * 3 for i in range(101)] + with pytest.raises(ValueError, match="exceeds maximum allowed 100"): + c._namespace_add(**self._common_kwargs(), ids=ids, embeddings=embeddings) + + def test_update_101_records_raises(self): + c = self._client() + ids = [f"d{i}" for i in range(101)] + with pytest.raises(ValueError, match="exceeds maximum allowed 100"): + c._namespace_update(**self._common_kwargs(), ids=ids, metadatas=[{"k": "v"}] * 101) + + def test_upsert_101_records_raises(self): + c = self._client() + ids = [f"d{i}" for i in range(101)] + embeddings = [[float(i)] * 3 for i in range(101)] + with pytest.raises(ValueError, match="exceeds maximum allowed 100"): + c._namespace_upsert(**self._common_kwargs(), ids=ids, embeddings=embeddings) + + def test_add_invalid_id_raises(self): + c = self._client() + with pytest.raises(ValueError, match="invalid characters"): + c._namespace_add(**self._common_kwargs(), ids="bad-id", embeddings=[1.0, 2.0, 3.0]) + + def test_delete_invalid_id_raises(self): + c = self._client() + with pytest.raises(ValueError, match="invalid characters"): + c._namespace_delete(**self._common_kwargs(), ids="bad-id") + + +# ==================== Physical Table Names Tests ==================== + + +class TestPhysicalTableNames: + + def test_data_table_name(self): + from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.data_table_name("my_coll") == "my_coll_logic_data_table" + + def test_hot_table_name(self): + from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.hot_table_name("my_coll") == "my_coll_hot_table" + + def test_kv_data_table_name(self): + from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.kv_data_table_name("my_coll") == "my_coll_kv_data_table" + + def test_logic_schema_table_name(self): + from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.logic_schema_table_name("my_coll") == "my_coll_logic_schema_table" + + def test_tablegroup_name(self): + from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.tablegroup_name("my_coll") == "my_coll_tg" + + def test_is_ns_data_table_true(self): + from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.is_ns_data_table("my_coll_logic_data_table") is True + + def test_is_ns_data_table_false(self): + from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.is_ns_data_table("my_coll_hot_table") is False + + +# ==================== UseNamespace Validation Tests ==================== + + +class TestUseNamespaceValidation: + + def test_hnsw_raises(self): + c = FakeClient() + from pyseekdb.client.schema import Schema + from pyseekdb.client.configuration import VectorIndexConfig, HNSWConfiguration + hnsw = HNSWConfiguration(dimension=128) + schema = Schema(vector_index=VectorIndexConfig(hnsw=hnsw, embedding_function=None)) + with pytest.raises(ValueError, match="HNSW is not allowed"): + c._create_namespace_collection("test", schema) + + def test_ob_type_validation(self): + c = FakeClient() + c.detect_db_type_and_version = MagicMock(return_value=("mysql", "8.0")) + from pyseekdb.client.schema import Schema + from pyseekdb.client.configuration import VectorIndexConfig + schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, use_spfresh=True), embedding_function=None)) + with pytest.raises(ValueError, match="only supported on OceanBase"): + c._create_namespace_collection("test", schema) + + def test_use_spfresh_false_raises(self): + c = FakeClient() + c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", "4.3")) + from pyseekdb.client.schema import Schema + from pyseekdb.client.configuration import VectorIndexConfig + schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, use_spfresh=False), embedding_function=None)) + with pytest.raises(ValueError, match="requires use_spfresh=True"): + c._create_namespace_collection("test", schema) + + +# ==================== Delete Namespace Uses Kernel ==================== + + +class TestDeleteNamespaceUsesKernel: + + def test_delete_namespace_calls_dbms_logic_table(self): + c = FakeClient() + c._execute = MagicMock(side_effect=[ + [{"namespace_id": 10, "namespace_name": "ns1"}], + None, + None, + None, + ]) + c._delete_ns_namespace_meta("abc123", "ns1") + calls = [str(call) for call in c._execute.call_args_list] + assert any("DBMS_LOGIC_TABLE.DROP_NAMESPACE" in s for s in calls) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From c4897dde23a3e9d4c717af9748b1f39ee942fdf9 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 29 Apr 2026 14:19:24 +0800 Subject: [PATCH 02/84] =?UTF-8?q?pyseekdb=E6=94=AF=E6=8C=81=E8=AE=BE?= =?UTF-8?q?=E7=BD=AEnamespace=5Fid=E7=9B=B8=E5=85=B3=E7=9A=84session?= =?UTF-8?q?=E7=BC=93=E5=AD=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 24 +++++- .../test_namespace_session_vars.py | 82 +++++++++++++++++++ 2 files changed, 104 insertions(+), 2 deletions(-) create mode 100644 tests/integration_tests/test_namespace_session_vars.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 8c3c31c3..09436217 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1116,6 +1116,7 @@ def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> di f"WHERE collection_name = '{collection_name_escaped}'" ) collection_id = str(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["collection_id"]) + self._set_session_ns_context(collection_id=collection_id) return {"collection_id": collection_id, "collection_name": collection_name} def _get_ns_collection_meta(self, collection_name: str) -> dict | None: @@ -1276,6 +1277,19 @@ def _cleanup_namespace_physical_tables(self, collection_id: str) -> None: with contextlib.suppress(Exception): self._execute(f"DROP TABLEGROUP IF EXISTS `{NamespaceCollectionNames.tablegroup_name(collection_id)}`") + def _set_session_ns_context( + self, + collection_id: str | None = None, + namespace_id: int | None = None, + ltable_id: int | None = None, + ) -> None: + if collection_id is not None: + self._execute(f"SET @collection_id = '{escape_string(str(collection_id))}'") + if namespace_id is not None: + self._execute(f"SET @namespace_id = {int(namespace_id)}") + if ltable_id is not None: + self._execute(f"SET @ltable_id = {int(ltable_id)}") + def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) @@ -1304,6 +1318,7 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> f"INSERT INTO `{schema_table}` (namespace_id, ltable_id, schema_content) " f"VALUES ({ns_id}, {lt_id}, '{escape_string(schema_content)}')" ) + self._set_session_ns_context(namespace_id=ns_id, ltable_id=lt_id) return {"namespace_id": str(ns_id), "namespace_name": namespace_name} def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: @@ -1317,8 +1332,13 @@ def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dic return None row = rows[0] if isinstance(row, (list, tuple)): - return {"namespace_id": str(row[0]), "namespace_name": row[1]} - return {"namespace_id": str(row["namespace_id"]), "namespace_name": row["namespace_name"]} + ns_id = str(row[0]) + ns_name = row[1] + else: + ns_id = str(row["namespace_id"]) + ns_name = row["namespace_name"] + self._set_session_ns_context(namespace_id=int(ns_id)) + return {"namespace_id": ns_id, "namespace_name": ns_name} def _has_ns_namespace(self, collection_id: str, namespace_name: str) -> bool: return self._get_ns_namespace_meta(collection_id, namespace_name) is not None diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py new file mode 100644 index 00000000..7eb7e47d --- /dev/null +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -0,0 +1,82 @@ +""" +Integration tests for namespace session variable propagation. +Verifies that collection_id, namespace_id, and ltable_id are set +as OceanBase session user variables (@collection_id, @namespace_id, @ltable_id) +during namespace collection and namespace creation/retrieval. +""" + +import time + +import pytest + +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.schema import Schema + + +class TestNamespaceSessionVars: + + def _create_ns_collection(self, client): + name = f"test_ns_sessvar_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2"), + embedding_function=None, + ), + ) + return client.create_collection(name=name, schema=schema, use_namespace=True) + + def _query_session_vars(self, client): + rows = client._server._execute( + "SELECT @collection_id AS cid, @namespace_id AS nsid, @ltable_id AS ltid" + ) + row = rows[0] + if isinstance(row, (list, tuple)): + return {"cid": row[0], "nsid": row[1], "ltid": row[2]} + return {"cid": row.get("cid"), "nsid": row.get("nsid"), "ltid": row.get("ltid")} + + def test_session_collection_id_after_create(self, db_client): + collection = self._create_ns_collection(db_client) + try: + vars_ = self._query_session_vars(db_client) + assert vars_["cid"] is not None, "@collection_id should be set after create_collection" + assert str(vars_["cid"]) == str(collection.id) + finally: + db_client.delete_collection(name=collection.name) + + def test_session_ns_ids_after_create_namespace(self, db_client): + collection = self._create_ns_collection(db_client) + try: + ns = collection.create_namespace("sess_ns") + vars_ = self._query_session_vars(db_client) + assert vars_["nsid"] is not None, "@namespace_id should be set after create_namespace" + assert int(vars_["nsid"]) == int(ns.namespace_id) + assert vars_["ltid"] is not None, "@ltable_id should be set after create_namespace" + assert int(vars_["ltid"]) > 0, "@ltable_id should be a positive integer" + finally: + db_client.delete_collection(name=collection.name) + + def test_session_ns_id_after_get_namespace(self, db_client): + collection = self._create_ns_collection(db_client) + try: + ns1 = collection.create_namespace("sess_get_ns") + ns1_id = int(ns1.namespace_id) + + ns2 = collection.create_namespace("sess_get_ns2") + ns2_id = int(ns2.namespace_id) + + # After creating ns2, session should have ns2's id + vars_ = self._query_session_vars(db_client) + assert int(vars_["nsid"]) == ns2_id + + # Now get_namespace(ns1) should update session to ns1's id + ns1_again = collection.get_namespace("sess_get_ns") + vars_ = self._query_session_vars(db_client) + assert int(vars_["nsid"]) == ns1_id + assert int(ns1_again.namespace_id) == ns1_id + finally: + db_client.delete_collection(name=collection.name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) From da41d77c8bc82a0d96d9a2bf34e0b1373d012a8d Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 30 Apr 2026 11:11:38 +0800 Subject: [PATCH 03/84] =?UTF-8?q?=E4=BB=A3=E7=A0=81=E6=A3=80=E8=A7=86?= =?UTF-8?q?=E9=97=AE=E9=A2=98=E4=BF=AE=E6=94=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 185 ++++++++++++++---- .../test_namespace_lifecycle.py | 24 +++ .../test_namespace_session_vars.py | 9 +- tests/unit_tests/test_namespace.py | 53 ++++- 4 files changed, 220 insertions(+), 51 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 09436217..830d8d3a 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -965,13 +965,22 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " self._ensure_namespace_catalogs() - self._create_namespace_physical_tables( - collection_id=collection_id, - dimension=dimension, - ivf_config=ivf_config, - fulltext_config=schema.fulltext_index, - is_shared_storage=is_ss, - ) + try: + self._create_namespace_physical_tables( + collection_id=collection_id, + dimension=dimension, + ivf_config=ivf_config, + fulltext_config=schema.fulltext_index, + is_shared_storage=is_ss, + ) + except Exception: + with contextlib.suppress(Exception): + collection_id_escaped = escape_string(collection_id) + self._execute( + f"DELETE FROM `{CollectionNames.sdk_collections_table_name()}` " + f"WHERE collection_id = '{collection_id_escaped}'" + ) + raise return Collection( client=self, @@ -1324,9 +1333,18 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) + ns_table = NamespaceCollectionNames.sdk_ns_namespaces_table() + lt_table = NamespaceCollectionNames.sdk_ns_ltables_table() rows = self._execute( - f"SELECT namespace_id, namespace_name FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " - f"WHERE collection_id = '{collection_id_escaped}' AND namespace_name = '{namespace_name_escaped}'" + f"SELECT n.namespace_id AS namespace_id, n.namespace_name AS namespace_name, " + f"l.ltable_id AS ltable_id " + f"FROM `{ns_table}` n " + f"LEFT JOIN `{lt_table}` l " + f"ON l.collection_id = n.collection_id " + f"AND l.namespace_id = n.namespace_id " + f"AND l.ltable_name = 'default' " + f"WHERE n.collection_id = '{collection_id_escaped}' " + f"AND n.namespace_name = '{namespace_name_escaped}'" ) if not rows: return None @@ -1334,11 +1352,17 @@ def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dic if isinstance(row, (list, tuple)): ns_id = str(row[0]) ns_name = row[1] + lt_raw = row[2] if len(row) > 2 else None else: ns_id = str(row["namespace_id"]) ns_name = row["namespace_name"] - self._set_session_ns_context(namespace_id=int(ns_id)) - return {"namespace_id": ns_id, "namespace_name": ns_name} + lt_raw = row.get("ltable_id") + lt_id = int(lt_raw) if lt_raw is not None else None + self._set_session_ns_context(namespace_id=int(ns_id), ltable_id=lt_id) + meta: dict = {"namespace_id": ns_id, "namespace_name": ns_name} + if lt_id is not None: + meta["ltable_id"] = str(lt_id) + return meta def _has_ns_namespace(self, collection_id: str, namespace_name: str) -> bool: return self._get_ns_namespace_meta(collection_id, namespace_name) is not None @@ -1348,7 +1372,14 @@ def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> if meta is None: raise ValueError(f"Namespace '{namespace_name}' not found") ns_id = meta["namespace_id"] + lt_id_raw = meta.get("ltable_id") + lt_id = int(lt_id_raw) if lt_id_raw is not None else None collection_id_escaped = escape_string(collection_id) + # Ensure DBMS_LOGIC_TABLE.DROP_NAMESPACE sees the right session context; + # otherwise it may read stale @collection_id / @ltable_id from a previous call. + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(ns_id), ltable_id=lt_id + ) self._execute( f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})" ) @@ -4401,7 +4432,7 @@ def _collection_count(self, collection_id: str | None, collection_name: str) -> # ==================== Namespace DML/DQL Methods ==================== @staticmethod - def _rewrite_where_for_ns(where: dict[str, Any]) -> dict[str, Any]: + def _rewrite_where_for_ns(where: dict[str, Any] | None) -> dict[str, Any] | None: if where is None: return None rewritten = {} @@ -4548,24 +4579,96 @@ def _namespace_update( ns_id = int(namespace_id) ltable_id = 1 + id_expr = "JSON_EXTRACT(data_content, '$.id')" + active_ids = [] for i, record_id in enumerate(ids): - set_parts = [] - if documents and i < len(documents) and documents[i] is not None: - set_parts.append(f"document = '{escape_string(documents[i])}'") + has_update = ( + (documents and i < len(documents) and documents[i] is not None) + or (embeddings and i < len(embeddings) and embeddings[i] is not None) + or (metadatas and i < len(metadatas) and metadatas[i] is not None) + ) + if has_update: + active_ids.append((i, record_id)) + + if not active_ids: + return + + conn = self._ensure_connection() + use_context_manager = self._use_context_manager_for_cursor() + + # VECTOR columns do not support CASE-WHEN in OceanBase, so embedding + # updates are executed per-row while document/metadata use batch CASE-WHEN. + emb_ids = [] + for i, record_id in active_ids: if embeddings and i < len(embeddings) and embeddings[i] is not None: - set_parts.append(f"embedding = {_embedding_to_hexstring(embeddings[i])}") + emb_ids.append((i, record_id)) + + if emb_ids: + for i, record_id in emb_ids: + vec_sql = _embedding_to_hexstring(embeddings[i]) + sql = ( + f"UPDATE `{table_name}` SET embedding = {vec_sql} " + f"WHERE namespace_id = {ns_id} AND ltable_id = {ltable_id} " + f"AND {id_expr} = %s" + ) + if use_context_manager: + with conn.cursor() as cursor: + cursor.execute(sql, [record_id]) + else: + cursor = conn.cursor() + try: + cursor.execute(sql, [record_id]) + finally: + cursor.close() + + doc_case_parts = [] + meta_case_parts = [] + params = [] + has_doc = False + has_meta = False + batch_ids = [] + + for i, record_id in active_ids: + if documents and i < len(documents) and documents[i] is not None: + has_doc = True + doc_case_parts.append(f"WHEN {id_expr} = %s THEN %s") + params.extend([record_id, documents[i]]) + if record_id not in [rid for rid in batch_ids]: + batch_ids.append(record_id) if metadatas and i < len(metadatas) and metadatas[i] is not None: + has_meta = True meta_json = json.dumps(metadatas[i], ensure_ascii=False) - set_parts.append( - f"data_content = JSON_SET(data_content, '$.metadata', CAST('{escape_string(meta_json)}' AS JSON))" + meta_case_parts.append( + f"WHEN {id_expr} = %s THEN JSON_SET(data_content, '$.metadata', CAST(%s AS JSON))" ) - if not set_parts: - continue - id_escaped = escape_string(record_id) - user_where = f"WHERE JSON_EXTRACT(data_content, '$.id') = '{id_escaped}'" - where_clause, _ = self._append_namespace_filter(user_where, [], ns_id, ltable_id) - sql = f"UPDATE `{table_name}` SET {', '.join(set_parts)} {where_clause}" - self._execute(sql) + params.extend([record_id, meta_json]) + if record_id not in batch_ids: + batch_ids.append(record_id) + + if has_doc or has_meta: + set_clauses = [] + if has_doc: + set_clauses.append(f"document = CASE {' '.join(doc_case_parts)} ELSE document END") + if has_meta: + set_clauses.append(f"data_content = CASE {' '.join(meta_case_parts)} ELSE data_content END") + + id_placeholders = ", ".join(["%s"] * len(batch_ids)) + params.extend(batch_ids) + + sql = ( + f"UPDATE `{table_name}` SET {', '.join(set_clauses)} " + f"WHERE namespace_id = {ns_id} AND ltable_id = {ltable_id} " + f"AND {id_expr} IN ({id_placeholders})" + ) + if use_context_manager: + with conn.cursor() as cursor: + cursor.execute(sql, params) + else: + cursor = conn.cursor() + try: + cursor.execute(sql, params) + finally: + cursor.close() def _namespace_upsert( self, @@ -4604,15 +4707,21 @@ def _namespace_upsert( ltable_id = 1 existing_ids = set() - for record_id in ids: - id_escaped = escape_string(record_id) - user_where = f"WHERE JSON_EXTRACT(data_content, '$.id') = '{id_escaped}'" - where_clause, _ = self._append_namespace_filter(user_where, [], ns_id, ltable_id) - rows = self._execute( - f"SELECT 1 FROM `{table_name}` {where_clause} LIMIT 1" - ) - if rows: - existing_ids.add(record_id) + id_expr = "JSON_EXTRACT(data_content, '$.id')" + id_placeholders = ", ".join(["%s"] * len(ids)) + check_sql = ( + f"SELECT {id_expr} AS rid FROM `{table_name}` " + f"WHERE namespace_id = {ns_id} AND ltable_id = {ltable_id} " + f"AND {id_expr} IN ({id_placeholders})" + ) + conn = self._ensure_connection() + use_ctx = self._use_context_manager_for_cursor() + rows = self._execute_query_with_cursor(conn, check_sql, list(ids), use_ctx) + for row in rows: + rid_raw = row.get("rid") if isinstance(row, dict) else row[0] + rid = json.loads(rid_raw) if isinstance(rid_raw, str) else rid_raw + if rid is not None: + existing_ids.add(str(rid)) add_indices = [] update_indices = [] @@ -4670,11 +4779,9 @@ def _namespace_delete( if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) - id_conditions = [] - for rid in ids: - id_escaped = escape_string(rid) - id_conditions.append(f"JSON_EXTRACT(data_content, '$.id') = '{id_escaped}'") - conditions.append(f"({' OR '.join(id_conditions)})") + id_placeholders = " OR ".join(["JSON_EXTRACT(data_content, '$.id') = %s"] * len(ids)) + conditions.append(f"({id_placeholders})") + params.extend(ids) if where is not None: rewritten = self._rewrite_where_for_ns(where) diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index a9adaf7a..384ae968 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -13,6 +13,21 @@ from pyseekdb.client.schema import Schema +# delete_namespace currently calls into the OceanBase PL package +# `DBMS_LOGIC_TABLE.DROP_NAMESPACE`, which is missing on fresh deployments. +# Catch the resulting (1049 "Unknown database 'DBMS_LOGIC_TABLE'") so xfail can +# match it precisely; once the package is installed, this test should XPASS and +# we can drop the marker. +try: + from pymysql.err import OperationalError as _PyMysqlOperationalError +except Exception: # pragma: no cover + _PyMysqlOperationalError = Exception +_DELETE_NS_XFAIL_RAISES: tuple[type[BaseException], ...] = ( + _PyMysqlOperationalError, + Exception, +) + + class TestNamespaceLifecycle: def _create_ns_collection(self, client, suffix=""): @@ -90,6 +105,15 @@ def test_list_namespaces(self, db_client): finally: db_client.delete_collection(name=collection.name) + @pytest.mark.xfail( + strict=False, + raises=_DELETE_NS_XFAIL_RAISES, + reason=( + "delete_namespace requires the OceanBase PL package " + "DBMS_LOGIC_TABLE.DROP_NAMESPACE, which is not installed on fresh " + "deployments. Test will XPASS once the package is available." + ), + ) def test_delete_namespace(self, db_client): collection = self._create_ns_collection(db_client) try: diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py index 7eb7e47d..bc7081ad 100644 --- a/tests/integration_tests/test_namespace_session_vars.py +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -61,18 +61,23 @@ def test_session_ns_id_after_get_namespace(self, db_client): try: ns1 = collection.create_namespace("sess_get_ns") ns1_id = int(ns1.namespace_id) + ns1_ltid = int(self._query_session_vars(db_client)["ltid"]) ns2 = collection.create_namespace("sess_get_ns2") ns2_id = int(ns2.namespace_id) + ns2_ltid = int(self._query_session_vars(db_client)["ltid"]) - # After creating ns2, session should have ns2's id + # After creating ns2, session should have ns2's id and its own @ltable_id vars_ = self._query_session_vars(db_client) assert int(vars_["nsid"]) == ns2_id + assert int(vars_["ltid"]) == ns2_ltid - # Now get_namespace(ns1) should update session to ns1's id + # Now get_namespace(ns1) should update session to ns1's id AND ns1's @ltable_id; + # otherwise kernel-side ops would observe ns2's stale @ltable_id while ns1 is in use. ns1_again = collection.get_namespace("sess_get_ns") vars_ = self._query_session_vars(db_client) assert int(vars_["nsid"]) == ns1_id + assert int(vars_["ltid"]) == ns1_ltid assert int(ns1_again.namespace_id) == ns1_id finally: db_client.delete_collection(name=collection.name) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index f30c38ad..678c68f6 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -525,11 +525,12 @@ def test_update_metadata_sql(self): sql = c.executed_sqls[-1] assert sql.startswith("UPDATE") assert f"`{self.TABLE}`" in sql + assert "CASE" in sql assert "JSON_SET(data_content, '$.metadata'," in sql assert "CAST(" in sql assert "AS JSON)" in sql assert "namespace_id = 7" in sql - assert """JSON_EXTRACT(data_content, '$.id') = 'd1'""" in sql + assert "'d1'" in sql def test_update_embedding_and_document_sql(self): c = self._client() @@ -539,10 +540,13 @@ def test_update_embedding_and_document_sql(self): embeddings=[9.0, 8.0, 7.0], documents="Updated", ) - sql = c.executed_sqls[-1] - assert "document = 'Updated'" in sql - assert "embedding = X'" in sql - assert "namespace_id = 7" in sql + emb_sql = c.executed_sqls[-2] + assert "embedding = X'" in emb_sql + assert "namespace_id = 7" in emb_sql + doc_sql = c.executed_sqls[-1] + assert "CASE" in doc_sql + assert "'Updated'" in doc_sql + assert "namespace_id = 7" in doc_sql # ---- DELETE ---- @@ -556,7 +560,8 @@ def test_delete_by_ids_sql(self): assert sql.startswith("DELETE FROM") assert f"`{self.TABLE}`" in sql assert "namespace_id = 7" in sql - assert """JSON_EXTRACT(data_content, '$.id') = 'd1'""" in sql + assert "JSON_EXTRACT(data_content, '$.id') = " in sql + assert "'d1'" in sql def test_delete_by_where_sql(self): c = self._client() @@ -944,14 +949,42 @@ class TestDeleteNamespaceUsesKernel: def test_delete_namespace_calls_dbms_logic_table(self): c = FakeClient() c._execute = MagicMock(side_effect=[ - [{"namespace_id": 10, "namespace_name": "ns1"}], - None, - None, - None, + [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], + None, # _get_ns_namespace_meta -> _set_session_ns_context SET @namespace_id + None, # _get_ns_namespace_meta -> _set_session_ns_context SET @ltable_id + None, # _delete_ns_namespace_meta -> SET @collection_id + None, # _delete_ns_namespace_meta -> SET @namespace_id + None, # _delete_ns_namespace_meta -> SET @ltable_id + None, # CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE + None, # DELETE FROM sdk_ns_ltables + None, # DELETE FROM sdk_ns_namespaces ]) c._delete_ns_namespace_meta("abc123", "ns1") calls = [str(call) for call in c._execute.call_args_list] assert any("DBMS_LOGIC_TABLE.DROP_NAMESPACE" in s for s in calls) + # session context must be set BEFORE DROP_NAMESPACE so the kernel sees the + # right @collection_id / @namespace_id / @ltable_id for this call. + drop_idx = next(i for i, s in enumerate(calls) if "DBMS_LOGIC_TABLE.DROP_NAMESPACE" in s) + before_drop = " | ".join(calls[:drop_idx]) + assert "SET @collection_id" in before_drop + assert "SET @namespace_id" in before_drop + assert "SET @ltable_id" in before_drop + + def test_get_ns_namespace_meta_sets_ltable_id(self): + """get_namespace must populate @ltable_id along with @namespace_id.""" + c = FakeClient() + c._execute = MagicMock(side_effect=[ + [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], + None, # SET @namespace_id + None, # SET @ltable_id + ]) + meta = c._get_ns_namespace_meta("abc123", "ns1") + assert meta == {"namespace_id": "10", "namespace_name": "ns1", "ltable_id": "7"} + calls = [str(call) for call in c._execute.call_args_list] + # JOIN against sdk_ns_ltables so we can resolve the default ltable in one round trip. + assert any("sdk_ns_ltables" in s for s in calls) + assert any("SET @namespace_id" in s for s in calls) + assert any("SET @ltable_id" in s for s in calls) if __name__ == "__main__": From 5f0e1a92f691ee939f0049bee21f8ffe97b8b073 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 12 May 2026 15:54:16 +0800 Subject: [PATCH 04/84] =?UTF-8?q?delete=20namespace=E8=B0=83=E7=94=A8ocean?= =?UTF-8?q?base=E7=9A=84=E5=86=85=E6=A0=B8DBMS=5FLOGIC=5FTABLE=E5=8C=85?= =?UTF-8?q?=EF=BC=8C=E5=86=85=E6=A0=B8=E8=87=AA=E8=A1=8C=E7=AE=A1=E7=90=86?= =?UTF-8?q?=E5=85=83=E6=95=B0=E6=8D=AE=E8=A1=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 8 ------- src/pyseekdb/client/client_seekdb_server.py | 2 +- .../test_namespace_lifecycle.py | 23 ------------------- tests/unit_tests/test_namespace.py | 2 -- 4 files changed, 1 insertion(+), 34 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 830d8d3a..7e430377 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1383,14 +1383,6 @@ def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> self._execute( f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})" ) - self._execute( - f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " - f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {ns_id}" - ) - self._execute( - f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " - f"WHERE namespace_id = {ns_id}" - ) def _list_ns_namespaces(self, collection_id: str) -> list[dict]: collection_id_escaped = escape_string(collection_id) diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 8d43254e..5b625fb6 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -208,7 +208,7 @@ def _namespace_prewarm( namespace_name: str, **kwargs, ) -> None: - sql = f"CALL DBMS_LOGIC_TABLE.PREWARM({collection_id}, {namespace_id})" + sql = f"CALL DBMS_LOGIC_TABLE.PREWARM('{collection_id}', {namespace_id})" self._execute(sql) def __repr__(self): diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 384ae968..2909ba77 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -13,20 +13,6 @@ from pyseekdb.client.schema import Schema -# delete_namespace currently calls into the OceanBase PL package -# `DBMS_LOGIC_TABLE.DROP_NAMESPACE`, which is missing on fresh deployments. -# Catch the resulting (1049 "Unknown database 'DBMS_LOGIC_TABLE'") so xfail can -# match it precisely; once the package is installed, this test should XPASS and -# we can drop the marker. -try: - from pymysql.err import OperationalError as _PyMysqlOperationalError -except Exception: # pragma: no cover - _PyMysqlOperationalError = Exception -_DELETE_NS_XFAIL_RAISES: tuple[type[BaseException], ...] = ( - _PyMysqlOperationalError, - Exception, -) - class TestNamespaceLifecycle: @@ -105,15 +91,6 @@ def test_list_namespaces(self, db_client): finally: db_client.delete_collection(name=collection.name) - @pytest.mark.xfail( - strict=False, - raises=_DELETE_NS_XFAIL_RAISES, - reason=( - "delete_namespace requires the OceanBase PL package " - "DBMS_LOGIC_TABLE.DROP_NAMESPACE, which is not installed on fresh " - "deployments. Test will XPASS once the package is available." - ), - ) def test_delete_namespace(self, db_client): collection = self._create_ns_collection(db_client) try: diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 678c68f6..7b4d851f 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -956,8 +956,6 @@ def test_delete_namespace_calls_dbms_logic_table(self): None, # _delete_ns_namespace_meta -> SET @namespace_id None, # _delete_ns_namespace_meta -> SET @ltable_id None, # CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE - None, # DELETE FROM sdk_ns_ltables - None, # DELETE FROM sdk_ns_namespaces ]) c._delete_ns_namespace_meta("abc123", "ns1") calls = [str(call) for call in c._execute.call_args_list] From f7f6b1cbc2574dc1e89ea77e3d533b294b5feb36 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 13 May 2026 15:12:17 +0800 Subject: [PATCH 05/84] =?UTF-8?q?pyseekdb=E4=BF=AE=E6=94=B9ob=E7=9A=84sear?= =?UTF-8?q?ch=C2=A0index=20with=C2=A0parser?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 21 +++------------------ 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7e430377..825cb1ef 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1207,7 +1207,7 @@ def _create_namespace_physical_tables( try: self._execute(f"CREATE TABLEGROUP `{tg_name}` SHARDING='ADAPTIVE'") - data_sql_with_search = f"""CREATE TABLE `{data_table}` ( + data_sql = f"""CREATE TABLE `{data_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, ltable_id BIGINT UNSIGNED NOT NULL, document LONGTEXT, @@ -1216,26 +1216,11 @@ def _create_namespace_physical_tables( created_by VARCHAR(64) DEFAULT '', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FULLTEXT INDEX idx_fts(document) {fulltext_clause}, - SEARCH INDEX idx_json(data_content) WITH PARSER json + SEARCH INDEX idx_json(data_content) ) TABLEGROUP=`{tg_name}` COMMENT='逻辑表主数据' DEFAULT CHARSET=utf8mb4 ORGANIZATION HEAP IS_LOGIC_TABLE = TRUE {partition_clause}""" - data_sql_without_search = f"""CREATE TABLE `{data_table}` ( - namespace_id BIGINT UNSIGNED NOT NULL, - ltable_id BIGINT UNSIGNED NOT NULL, - document LONGTEXT, - embedding VECTOR({dimension}), - data_content JSON NOT NULL, - created_by VARCHAR(64) DEFAULT '', - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - FULLTEXT INDEX idx_fts(document) {fulltext_clause} - ) TABLEGROUP=`{tg_name}` COMMENT='逻辑表主数据' DEFAULT CHARSET=utf8mb4 ORGANIZATION HEAP IS_LOGIC_TABLE = TRUE - {partition_clause}""" - - try: - self._execute(data_sql_with_search) - except Exception: - self._execute(data_sql_without_search) + self._execute(data_sql) ivf_index_sql = f"CREATE VECTOR INDEX idx_vec ON `{data_table}` (embedding) {vector_index_sql}" self._execute(ivf_index_sql) From ca1a06992ecdab7affb53e0eede60b1856923831 Mon Sep 17 00:00:00 2001 From: ly435438 Date: Wed, 13 May 2026 17:45:21 +0800 Subject: [PATCH 06/84] =?UTF-8?q?=E6=96=B0=E5=A2=9E=20namespace=20?= =?UTF-8?q?=E7=BB=9F=E8=AE=A1=20catalog=20=E8=A1=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 20 +++++++++++ src/pyseekdb/client/meta_info.py | 4 +++ tests/unit_tests/test_namespace.py | 54 ++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7e430377..9571916c 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1107,8 +1107,23 @@ def _ensure_namespace_catalogs(self) -> None: UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name), KEY idx_sdk_lt_by_ns (collection_id, namespace_id) ) COMMENT='LTable catalog';""" + namespaces_stats_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_namespaces_stats_table()}` ( + collection_id CHAR(32) NOT NULL COMMENT 'collection id', + namespace_id BIGINT UNSIGNED NOT NULL COMMENT 'namespace internal id', + ltable_id BIGINT UNSIGNED NOT NULL COMMENT 'logic table internal id, 0 means namespace summary', + estimated_rows BIGINT NOT NULL DEFAULT 0 COMMENT 'estimated row count', + average_row_size BIGINT NOT NULL DEFAULT 0 COMMENT 'average row size in bytes', + row_limit BIGINT NOT NULL DEFAULT -1 COMMENT 'row count limit, -1 means unlimited', + size_limit BIGINT NOT NULL DEFAULT -1 COMMENT 'storage size limit in bytes, -1 means unlimited', + last_estimate_time TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) COMMENT 'last estimate time', + included_index BOOL NOT NULL DEFAULT FALSE COMMENT 'whether stats include index data', + PRIMARY KEY (namespace_id, ltable_id, included_index), + KEY idx_sdk_ns_stat_by_collection (collection_id) + ) COMMENT='Logic table row count and storage size statistics' DEFAULT CHARSET=utf8mb4 + PARTITION BY KEY(namespace_id) PARTITIONS 1000;""" self._execute(ns_namespaces_sql) self._execute(ns_ltables_sql) + self._execute(namespaces_stats_sql) def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> dict: self._create_sdk_collections_if_not_exists() @@ -1185,6 +1200,11 @@ def _delete_ns_collection_meta(self, collection_name: str) -> None: f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " f"WHERE collection_id = '{collection_id_escaped}'" ) + with contextlib.suppress(Exception): + self._execute( + f"DELETE FROM `{NamespaceCollectionNames.sdk_namespaces_stats_table()}` " + f"WHERE collection_id = '{collection_id_escaped}'" + ) self._cleanup_namespace_physical_tables(collection_id) def _create_namespace_physical_tables( diff --git a/src/pyseekdb/client/meta_info.py b/src/pyseekdb/client/meta_info.py index 7dd52137..53be60e5 100644 --- a/src/pyseekdb/client/meta_info.py +++ b/src/pyseekdb/client/meta_info.py @@ -72,6 +72,10 @@ def sdk_ns_namespaces_table() -> str: def sdk_ns_ltables_table() -> str: return "sdk_ns_ltables" + @staticmethod + def sdk_namespaces_stats_table() -> str: + return "sdk_namespaces_stats" + @staticmethod def data_table_name(collection_id: str) -> str: return f"{collection_id}{NamespaceCollectionNames._LOGIC_DATA_SUFFIX}" diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 7b4d851f..0db30bed 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -899,6 +899,12 @@ def test_tablegroup_name(self): from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.tablegroup_name("my_coll") == "my_coll_tg" + def test_namespace_catalog_table_names(self): + from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.sdk_ns_namespaces_table() == "sdk_ns_namespaces" + assert NamespaceCollectionNames.sdk_ns_ltables_table() == "sdk_ns_ltables" + assert NamespaceCollectionNames.sdk_namespaces_stats_table() == "sdk_namespaces_stats" + def test_is_ns_data_table_true(self): from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.is_ns_data_table("my_coll_logic_data_table") is True @@ -908,6 +914,54 @@ def test_is_ns_data_table_false(self): assert NamespaceCollectionNames.is_ns_data_table("my_coll_hot_table") is False +# ==================== Namespace Catalog Tests ==================== + + +class TestNamespaceCatalogs: + + def test_ensure_namespace_catalogs_creates_all_catalog_tables(self): + c = FakeClient() + c._ensure_namespace_catalogs() + + sql = "\n".join(c.executed_sqls) + assert "CREATE TABLE IF NOT EXISTS `sdk_ns_namespaces`" in sql + assert "PRIMARY KEY (namespace_id)" in sql + assert "UNIQUE KEY uk_sdk_ns_coll_name (collection_id, namespace_name)" in sql + + assert "CREATE TABLE IF NOT EXISTS `sdk_ns_ltables`" in sql + assert "PRIMARY KEY (ltable_id)" in sql + assert "UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name)" in sql + + assert "CREATE TABLE IF NOT EXISTS `sdk_namespaces_stats`" in sql + assert "collection_id CHAR(32) NOT NULL" in sql + assert "namespace_id BIGINT UNSIGNED NOT NULL" in sql + assert "ltable_id BIGINT UNSIGNED NOT NULL" in sql + assert "included_index BOOL" in sql + assert "PRIMARY KEY (namespace_id, ltable_id, included_index)" in sql + assert "PARTITION BY KEY(namespace_id) PARTITIONS 1000" in sql + + def test_ensure_namespace_catalogs_creates_catalog_tables_in_order(self): + c = FakeClient() + c._ensure_namespace_catalogs() + + assert len(c.executed_sqls) == 3 + assert "`sdk_ns_namespaces`" in c.executed_sqls[0] + assert "`sdk_ns_ltables`" in c.executed_sqls[1] + assert "`sdk_namespaces_stats`" in c.executed_sqls[2] + + def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): + c = FakeClient() + c._get_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) + c._cleanup_namespace_physical_tables = MagicMock() + c._execute = MagicMock() + + c._delete_ns_collection_meta("coll") + + calls = [str(call) for call in c._execute.call_args_list] + assert any("DELETE FROM `sdk_namespaces_stats`" in s for s in calls) + assert any("WHERE collection_id = 'abc123'" in s for s in calls) + + # ==================== UseNamespace Validation Tests ==================== From 0622ed9f6b818ced5961cbba263c5c417a30eb50 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 14 May 2026 10:33:49 +0800 Subject: [PATCH 07/84] =?UTF-8?q?prewarm=E9=9B=86=E6=88=90=E6=B5=8B?= =?UTF-8?q?=E8=AF=95=E7=94=A8=E4=BE=8B=E5=8A=A0=E5=BC=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../test_namespace_prewarm.py | 57 ++++++++++++++----- 1 file changed, 44 insertions(+), 13 deletions(-) diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index ddbeda2d..fd8056af 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -1,6 +1,8 @@ """ Namespace prewarm integration tests. -Tests that prewarm raises ValueError in embedded mode and executes in remote mode. + +- embedded mode: prewarm raises ValueError +- oceanbase SS mode: prewarm inserts hot_table record, repeat prewarm updates last_access_time """ import time @@ -9,11 +11,20 @@ from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.meta_info import NamespaceCollectionNames from pyseekdb.client.schema import Schema class TestNamespacePrewarm: + @pytest.fixture(autouse=True) + def _reduce_namespace_partitions(self): + """Override conftest: use default 1000 partitions for prewarm realism.""" + # TODO: 1000 partitions causes ~18min DDL; using 8 for now until DDL is faster + from unittest.mock import patch + with patch("pyseekdb.client.client_base._NS_PARTITION_COUNT", 8): + yield + def _create_ns_collection_and_namespace(self, client): name = f"test_ns_pw_{int(time.time() * 1000)}" schema = Schema( @@ -26,6 +37,14 @@ def _create_ns_collection_and_namespace(self, client): namespace = collection.create_namespace("pw_ns") return collection, namespace + def _query_hot_table(self, client, collection, namespace_id): + coll_id = collection.id + hot_table = NamespaceCollectionNames.hot_table_name(coll_id) + rows = client._server._execute( + f"SELECT namespace_id, last_access_time FROM `{hot_table}` WHERE namespace_id = {namespace_id}" + ) + return rows + def test_prewarm_raises_in_embedded_mode(self, embedded_client): collection, namespace = self._create_ns_collection_and_namespace(embedded_client) try: @@ -34,25 +53,37 @@ def test_prewarm_raises_in_embedded_mode(self, embedded_client): finally: embedded_client.delete_collection(name=collection.name) - def test_prewarm_executes_in_server_mode(self, server_client): - """Prewarm should not raise in remote mode (actual behavior depends on backend).""" - collection, namespace = self._create_ns_collection_and_namespace(server_client) + def test_prewarm_inserts_hot_table_record(self, oceanbase_client): + """First prewarm should insert a record into hot_table.""" + collection, namespace = self._create_ns_collection_and_namespace(oceanbase_client) try: + rows_before = self._query_hot_table(oceanbase_client, collection, namespace._namespace_id) + assert len(rows_before) == 0, "hot_table should be empty before prewarm" + namespace.prewarm() - except Exception as e: - if "shared-storage remote" in str(e): - pytest.fail("prewarm should not raise embedded-mode error in server mode") + + rows_after = self._query_hot_table(oceanbase_client, collection, namespace._namespace_id) + assert len(rows_after) == 1, "hot_table should have exactly one record after prewarm" finally: - server_client.delete_collection(name=collection.name) + oceanbase_client.delete_collection(name=collection.name) - def test_prewarm_executes_in_oceanbase_mode(self, oceanbase_client): - """Prewarm should not raise in OceanBase mode (actual behavior depends on backend).""" + def test_prewarm_repeat_updates_access_time(self, oceanbase_client): + """Second prewarm should update last_access_time (not insert a new row).""" collection, namespace = self._create_ns_collection_and_namespace(oceanbase_client) try: namespace.prewarm() - except Exception as e: - if "shared-storage remote" in str(e): - pytest.fail("prewarm should not raise embedded-mode error in oceanbase mode") + rows1 = self._query_hot_table(oceanbase_client, collection, namespace._namespace_id) + assert len(rows1) == 1 + ts1 = rows1[0]["last_access_time"] if isinstance(rows1[0], dict) else rows1[0][1] + + time.sleep(2) + + namespace.prewarm() + rows2 = self._query_hot_table(oceanbase_client, collection, namespace._namespace_id) + assert len(rows2) == 1, "repeat prewarm should not create a second row" + ts2 = rows2[0]["last_access_time"] if isinstance(rows2[0], dict) else rows2[0][1] + + assert ts2 > ts1, f"last_access_time should increase after repeat prewarm: {ts1} -> {ts2}" finally: oceanbase_client.delete_collection(name=collection.name) From 5eb9b44f569b54c26924809da6581c67ab6a916c Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 14 May 2026 11:45:06 +0800 Subject: [PATCH 08/84] =?UTF-8?q?=E6=8F=90=E4=BE=9Bset=5Fnamespace=5Fparti?= =?UTF-8?q?tion=5Fcount=E5=8A=A8=E6=80=81=E8=AE=BE=E7=BD=AE=E5=88=86?= =?UTF-8?q?=E5=8C=BA=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/__init__.py | 4 ++++ src/pyseekdb/client/__init__.py | 4 +++- src/pyseekdb/client/client_base.py | 11 +++++++++ tests/integration_tests/conftest.py | 9 +++---- .../test_namespace_lifecycle.py | 24 +++++++++++++++++++ .../test_namespace_prewarm.py | 19 +++++---------- 6 files changed, 53 insertions(+), 18 deletions(-) diff --git a/src/pyseekdb/__init__.py b/src/pyseekdb/__init__.py index 1e2a33af..da587421 100644 --- a/src/pyseekdb/__init__.py +++ b/src/pyseekdb/__init__.py @@ -102,8 +102,10 @@ VectorIndexConfig, Version, get_default_embedding_function, + get_namespace_partition_count, register_embedding_function, register_sparse_embedding_function, + set_namespace_partition_count, ) from .client.collection import Collection from .client.namespace import Namespace @@ -149,6 +151,8 @@ "VectorIndexConfig", "Version", "get_default_embedding_function", + "get_namespace_partition_count", "register_embedding_function", "register_sparse_embedding_function", + "set_namespace_partition_count", ] diff --git a/src/pyseekdb/client/__init__.py b/src/pyseekdb/client/__init__.py index 1ce0aaef..0fc90630 100644 --- a/src/pyseekdb/client/__init__.py +++ b/src/pyseekdb/client/__init__.py @@ -22,7 +22,7 @@ from .admin_client import AdminAPI, _AdminClientProxy, _ClientProxy from .base_connection import BaseConnection -from .client_base import BaseClient, ClientAPI +from .client_base import BaseClient, ClientAPI, get_namespace_partition_count, set_namespace_partition_count from .client_seekdb_server import RemoteServerClient from .configuration import ( BengProperties, @@ -188,8 +188,10 @@ def __getattr__(name: str) -> Any: "VectorIndexConfig", "Version", "get_default_embedding_function", + "get_namespace_partition_count", "register_embedding_function", "register_sparse_embedding_function", + "set_namespace_partition_count", ] diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index f254f47c..90d1094e 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -127,6 +127,17 @@ def _validate_collection_name(name: str) -> None: _NS_PARTITION_COUNT = 1000 +def set_namespace_partition_count(n: int) -> None: + global _NS_PARTITION_COUNT + if n < 1: + raise ValueError("namespace partition count must be >= 1") + _NS_PARTITION_COUNT = n + + +def get_namespace_partition_count() -> int: + return _NS_PARTITION_COUNT + + def _build_default_ltable_schema() -> dict: return { "col_info": [ diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index 899c5510..9f930cdc 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -15,8 +15,6 @@ src_root = repo_root / "src" sys.path.insert(0, str(src_root)) -from unittest.mock import patch - import pyseekdb # noqa: E402 # ==================== Environment Variable Configuration ==================== @@ -275,5 +273,8 @@ def oceanbase_admin_client(): @pytest.fixture(autouse=True) def _reduce_namespace_partitions(): """Use a small partition count in integration tests to speed up DDL.""" - with patch("pyseekdb.client.client_base._NS_PARTITION_COUNT", 8): - yield + from pyseekdb import get_namespace_partition_count, set_namespace_partition_count + original = get_namespace_partition_count() + set_namespace_partition_count(8) + yield + set_namespace_partition_count(original) diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 2909ba77..3049e856 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -202,5 +202,29 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): client.delete_collection(name=name) + def test_custom_namespace_partition_count(self, oceanbase_client): + """Verify set_namespace_partition_count controls the PARTITIONS clause.""" + from pyseekdb import get_namespace_partition_count, set_namespace_partition_count + from pyseekdb.client.meta_info import NamespaceCollectionNames + + original = get_namespace_partition_count() + try: + set_namespace_partition_count(4) + collection = self._create_ns_collection(oceanbase_client, suffix="_pc") + try: + data_table = NamespaceCollectionNames.data_table_name(collection.id) + rows = oceanbase_client._server._execute(f"SHOW CREATE TABLE `{data_table}`") + create_sql = rows[0].get("Create Table", "") if rows else "" + import re + partitions = re.findall(r"partition `p\d+`", create_sql) + assert len(partitions) == 4, ( + f"Expected 4 partitions, found {len(partitions)}: {create_sql[-300:]}" + ) + finally: + oceanbase_client.delete_collection(name=collection.name) + finally: + set_namespace_partition_count(original) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index fd8056af..4c2f5e3c 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -19,11 +19,12 @@ class TestNamespacePrewarm: @pytest.fixture(autouse=True) def _reduce_namespace_partitions(self): - """Override conftest: use default 1000 partitions for prewarm realism.""" - # TODO: 1000 partitions causes ~18min DDL; using 8 for now until DDL is faster - from unittest.mock import patch - with patch("pyseekdb.client.client_base._NS_PARTITION_COUNT", 8): - yield + """Override conftest: use 8 partitions for prewarm until DDL is faster.""" + from pyseekdb import get_namespace_partition_count, set_namespace_partition_count + original = get_namespace_partition_count() + set_namespace_partition_count(8) + yield + set_namespace_partition_count(original) def _create_ns_collection_and_namespace(self, client): name = f"test_ns_pw_{int(time.time() * 1000)}" @@ -45,14 +46,6 @@ def _query_hot_table(self, client, collection, namespace_id): ) return rows - def test_prewarm_raises_in_embedded_mode(self, embedded_client): - collection, namespace = self._create_ns_collection_and_namespace(embedded_client) - try: - with pytest.raises(ValueError, match="shared-storage remote"): - namespace.prewarm() - finally: - embedded_client.delete_collection(name=collection.name) - def test_prewarm_inserts_hot_table_record(self, oceanbase_client): """First prewarm should insert a record into hot_table.""" collection, namespace = self._create_ns_collection_and_namespace(oceanbase_client) From 2d901a4f2ee93c2644a5c045e6ef5c7704b2ae0e Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 15 May 2026 11:27:39 +0800 Subject: [PATCH 09/84] =?UTF-8?q?=E5=B1=8F=E8=94=BDhybrid=5Fsearch?= =?UTF-8?q?=E7=94=A8search=5Fpara=E8=A7=A3=E6=9E=90=E6=88=90sql=EF=BC=8C?= =?UTF-8?q?=E7=9B=B4=E6=8E=A5=E5=90=91=E5=86=85=E6=A0=B8=E8=BE=93=E5=85=A5?= =?UTF-8?q?hybrid=5Fsearch=E6=8E=A5=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 26 ++++++++++---------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 90d1094e..9626f7f1 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -5249,31 +5249,25 @@ def _namespace_hybrid_search( ) search_parm = self._adapt_search_parm_for_ns(search_parm, ns_id) + search_parm.pop("_source", None) + search_parm_json = json.dumps(search_parm, ensure_ascii=False) use_context_manager = self._use_context_manager_for_cursor() - escaped_params = escape_string(search_parm_json) - set_sql = f"SET @search_parm = '{escaped_params}'" - self._execute_query_with_cursor(conn, set_sql, [], use_context_manager) + escaped_params = search_parm_json.replace("'", "''") - get_sql_query = f"SELECT DBMS_HYBRID_SEARCH.GET_SQL('{table_name}', @search_parm) as query_sql FROM dual" - rows = self._execute_query_with_cursor(conn, get_sql_query, [], use_context_manager) + hint_sql = _query_hint_to_sql(query_hint, table_name=table_name) or "" + hybrid_sql = ( + f"SELECT {hint_sql + ' ' if hint_sql else ''}* " + f"FROM hybrid_search(TABLE `{table_name}`, '{escaped_params}')" + ) - if not rows or not rows[0].get("query_sql"): + result_rows = self._execute_query_with_cursor(conn, hybrid_sql, [], use_context_manager) + if not result_rows: return { "ids": [[]], "distances": [[]], "metadatas": [[]], "documents": [[]], "embeddings": [[]], } - - query_sql = rows[0]["query_sql"] - if isinstance(query_sql, str): - query_sql = query_sql.strip().strip("'\"") - - hint_sql = _query_hint_to_sql(query_hint, table_name=table_name) - if hint_sql and query_sql.upper().startswith("SELECT"): - query_sql = f"SELECT {hint_sql} {query_sql[len('SELECT'):]}" - - result_rows = self._execute_query_with_cursor(conn, query_sql, [], use_context_manager) return self._transform_ns_hybrid_result(result_rows, include) def _transform_ns_hybrid_result( From 750f58178d48d3d13487d831049639c3c46cc56c Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 15 May 2026 17:08:03 +0800 Subject: [PATCH 10/84] =?UTF-8?q?=E6=B7=B7=E6=90=9C=E6=94=AF=E6=8C=81data?= =?UTF-8?q?=5Fcontent=E8=BF=87=E6=BB=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 18 +++-- tests/integration_tests/Untitled | 1 + .../integration_tests/test_namespace_query.py | 66 +++++++++++++++++++ 3 files changed, 78 insertions(+), 7 deletions(-) create mode 100644 tests/integration_tests/Untitled diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 9626f7f1..32111e19 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -4043,7 +4043,10 @@ def _build_metadata_filter_for_search_parm(self, where: dict[str, Any] | None) - def _build_search_parm_field_name(self, key: str) -> str: """ Build field name used in search_parm filters. - Supports special "#id" to refer to the primary key column directly. + + Uses ``JSON_EXTRACT``-wrapped keys for ``DBMS_HYBRID_SEARCH.GET_SQL`` (collection path). + Namespace ``hybrid_search(TABLE ...)`` rewrites these to ``data_content.metadata.*`` DSL + keys in ``_adapt_search_parm_for_ns``. """ if key == "#id" or key == CollectionFieldNames.ID: return CollectionFieldNames.ID @@ -5149,12 +5152,13 @@ def _rewrite_field_refs(obj): for k, v in obj.items(): new_key = k if k == "_id": - new_key = "(JSON_EXTRACT(data_content, '$.id'))" - elif isinstance(k, str) and "JSON_EXTRACT(metadata, '$." in k: - new_key = k.replace( - "JSON_EXTRACT(metadata, '$.", - "JSON_EXTRACT(data_content, '$.metadata.", - ) + new_key = "data_content.id" + elif isinstance(k, str): + prefix = "(JSON_EXTRACT(metadata, '$." + suffix = "'))" + if k.startswith(prefix) and k.endswith(suffix) and len(k) > len(prefix) + len(suffix): + inner = k[len(prefix) : -len(suffix)] + new_key = f"data_content.metadata.{inner}" new_dict[new_key] = _rewrite_field_refs(v) return new_dict elif isinstance(obj, list): diff --git a/tests/integration_tests/Untitled b/tests/integration_tests/Untitled new file mode 100644 index 00000000..4a0ec5c6 --- /dev/null +++ b/tests/integration_tests/Untitled @@ -0,0 +1 @@ +test_collection_data_api_blocked_when_namespace_enabled[ \ No newline at end of file diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index d92b578f..575f69fb 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -285,6 +285,72 @@ def test_hybrid_search_namespace_isolation(self, db_client): finally: db_client.delete_collection(name=collection.name) + def test_hybrid_search_with_data_content_filter(self, db_client): + """ + Namespace hybrid_search applies metadata / id filters on JSON under ``data_content`` + (SDK maps ``metadata.*`` / ``_id`` to ``data_content.metadata.*`` / ``data_content.id`` + for OceanBase hybrid_search field syntax). + """ + collection = self._setup(db_client) + ns = collection.create_namespace("qns_hs_dc") + try: + self._insert_data(ns) + time.sleep(1) + + # Full-text + filter on data_content.metadata.category + result_ft = ns.hybrid_search( + query={ + "where_document": {"$contains": "Python"}, + "where": {"category": "Programming"}, + "n_results": 10, + }, + n_results=10, + include=["documents", "metadatas"], + ) + ids_ft = result_ft["ids"][0] if result_ft["ids"] else [] + assert ids_ft == ["q2"], ids_ft + if result_ft.get("metadatas") and result_ft["metadatas"][0]: + for meta in result_ft["metadatas"][0]: + if meta: + assert meta["category"] == "Programming" + finally: + db_client.delete_collection(name=collection.name) + + def test_hybrid_search_with_data_content_filter_knn_branch(self, db_client, request): + """ + KNN + metadata / id filters on ``data_content`` (same field mapping as full-text branch). + + OceanBase ``hybrid_search`` KNN currently requires an HNSW vector index; namespace + logical tables are created with IVF only, so this branch is skipped for ``[oceanbase]``. + """ + if "oceanbase" in request.node.nodeid: + pytest.skip("OceanBase hybrid_search KNN requires HNSW; namespace collections use IVF only.") + collection = self._setup(db_client) + ns = collection.create_namespace("qns_hs_dc_knn") + try: + self._insert_data(ns) + time.sleep(1) + result_knn = ns.hybrid_search( + knn={ + "query_embeddings": [1.0, 0.0, 0.0], + "where": { + "$and": [ + {"score": {"$gte": 90}}, + {"#id": "q1"}, + ] + }, + "n_results": 10, + }, + n_results=10, + include=["metadatas"], + ) + ids_knn = result_knn["ids"][0] if result_knn["ids"] else [] + assert ids_knn == ["q1"], ids_knn + if result_knn.get("metadatas") and result_knn["metadatas"][0]: + assert result_knn["metadatas"][0][0]["score"] >= 90 + finally: + db_client.delete_collection(name=collection.name) + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) From 0c62edc4ad6819058464905b54626b70229f51bc Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 15 May 2026 17:38:19 +0800 Subject: [PATCH 11/84] =?UTF-8?q?kv=20table=E5=88=B0key=20=E6=94=B9?= =?UTF-8?q?=E4=B8=BAvarbinary?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 32111e19..7ace1d4b 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1269,7 +1269,7 @@ def _create_namespace_physical_tables( self._execute(f"""CREATE TABLE `{kv_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, - kv_key VARCHAR(1024) NOT NULL, + kv_key VARBINARY(1024) NOT NULL, kv_value LONGBLOB NOT NULL, PRIMARY KEY(namespace_id, kv_key) ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 From 765dceed7c6df8b7c7bbf006b98ea3b52f46b25e Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 18 May 2026 13:57:14 +0800 Subject: [PATCH 12/84] =?UTF-8?q?namespace.add,get=E7=AD=89=E6=93=8D?= =?UTF-8?q?=E4=BD=9C=E6=98=BE=E7=A4=BA=E8=AE=BE=E7=BD=AEsession?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 28 +++++++++++++++++++-- src/pyseekdb/client/client_seekdb_server.py | 3 +++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7ace1d4b..f526c277 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -124,7 +124,7 @@ def _validate_collection_name(name: str) -> None: from .validators import _MAX_NAMESPACE_BATCH_SIZE, _validate_namespace_name, _validate_record_ids # noqa: F401 -_NS_PARTITION_COUNT = 1000 +_NS_PARTITION_COUNT = 8 def set_namespace_partition_count(n: int) -> None: @@ -1131,7 +1131,7 @@ def _ensure_namespace_catalogs(self) -> None: PRIMARY KEY (namespace_id, ltable_id, included_index), KEY idx_sdk_ns_stat_by_collection (collection_id) ) COMMENT='Logic table row count and storage size statistics' DEFAULT CHARSET=utf8mb4 - PARTITION BY KEY(namespace_id) PARTITIONS 1000;""" + PARTITION BY KEY(namespace_id) PARTITIONS {_NS_PARTITION_COUNT};""" self._execute(ns_namespaces_sql) self._execute(ns_ltables_sql) self._execute(namespaces_stats_sql) @@ -4484,6 +4484,9 @@ def _namespace_add( # noqa: C901 embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) @@ -4564,6 +4567,9 @@ def _namespace_update( embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) @@ -4694,6 +4700,9 @@ def _namespace_upsert( embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) @@ -4777,6 +4786,9 @@ def _namespace_delete( where_document: dict[str, Any] | None = None, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) if ids is None and where is None and where_document is None: raise ValueError("At least one of ids, where, or where_document must be provided") @@ -4839,6 +4851,9 @@ def _namespace_query( # noqa: C901 include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) embedding_function = kwargs.get("embedding_function") distance = kwargs.get("distance", DEFAULT_DISTANCE_METRIC) @@ -4994,6 +5009,9 @@ def _namespace_get( # noqa: C901 include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) include_fields = self._normalize_include_fields(include) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) @@ -5105,6 +5123,9 @@ def _namespace_count( namespace_name: str, **kwargs, ) -> int: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) where_clause, _ = self._append_namespace_filter("", [], ns_id) @@ -5241,6 +5262,9 @@ def _namespace_hybrid_search( query_hint: QueryHint | None = None, **kwargs, ) -> dict[str, Any]: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) conn = self._ensure_connection() table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 5b625fb6..5b204f87 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -208,6 +208,9 @@ def _namespace_prewarm( namespace_name: str, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) sql = f"CALL DBMS_LOGIC_TABLE.PREWARM('{collection_id}', {namespace_id})" self._execute(sql) From 941a0c9917d98d1c1ee30a675aed313d37c39ef0 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 18 May 2026 13:57:14 +0800 Subject: [PATCH 13/84] =?UTF-8?q?namespace.add,get=E7=AD=89=E6=93=8D?= =?UTF-8?q?=E4=BD=9C=E6=98=BE=E7=A4=BA=E8=AE=BE=E7=BD=AEsession?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 26 ++++++++++++++++++++- src/pyseekdb/client/client_seekdb_server.py | 3 +++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7ace1d4b..fadef14b 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1131,7 +1131,7 @@ def _ensure_namespace_catalogs(self) -> None: PRIMARY KEY (namespace_id, ltable_id, included_index), KEY idx_sdk_ns_stat_by_collection (collection_id) ) COMMENT='Logic table row count and storage size statistics' DEFAULT CHARSET=utf8mb4 - PARTITION BY KEY(namespace_id) PARTITIONS 1000;""" + PARTITION BY KEY(namespace_id) PARTITIONS {_NS_PARTITION_COUNT};""" self._execute(ns_namespaces_sql) self._execute(ns_ltables_sql) self._execute(namespaces_stats_sql) @@ -4484,6 +4484,9 @@ def _namespace_add( # noqa: C901 embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) @@ -4564,6 +4567,9 @@ def _namespace_update( embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) @@ -4694,6 +4700,9 @@ def _namespace_upsert( embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) @@ -4777,6 +4786,9 @@ def _namespace_delete( where_document: dict[str, Any] | None = None, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) if ids is None and where is None and where_document is None: raise ValueError("At least one of ids, where, or where_document must be provided") @@ -4839,6 +4851,9 @@ def _namespace_query( # noqa: C901 include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) embedding_function = kwargs.get("embedding_function") distance = kwargs.get("distance", DEFAULT_DISTANCE_METRIC) @@ -4994,6 +5009,9 @@ def _namespace_get( # noqa: C901 include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) include_fields = self._normalize_include_fields(include) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) @@ -5105,6 +5123,9 @@ def _namespace_count( namespace_name: str, **kwargs, ) -> int: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) where_clause, _ = self._append_namespace_filter("", [], ns_id) @@ -5241,6 +5262,9 @@ def _namespace_hybrid_search( query_hint: QueryHint | None = None, **kwargs, ) -> dict[str, Any]: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) conn = self._ensure_connection() table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 5b625fb6..5b204f87 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -208,6 +208,9 @@ def _namespace_prewarm( namespace_name: str, **kwargs, ) -> None: + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + ) sql = f"CALL DBMS_LOGIC_TABLE.PREWARM('{collection_id}', {namespace_id})" self._execute(sql) From 4ee60df5acec57168af2488b330d6e7c42255de5 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 18 May 2026 20:51:02 +0800 Subject: [PATCH 14/84] =?UTF-8?q?pyseekdb=E6=96=B0=E5=A2=9Edbms=5Flogic=5F?= =?UTF-8?q?table.drop=5Fnamespace=E7=94=A8=E4=BE=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../test_namespace_drop_validation.py | 379 ++++++++++++++++++ 1 file changed, 379 insertions(+) create mode 100644 tests/integration_tests/test_namespace_drop_validation.py diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py new file mode 100644 index 00000000..74045610 --- /dev/null +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -0,0 +1,379 @@ +import threading +import time + +import pytest + +import pyseekdb +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.meta_info import NamespaceCollectionNames +from pyseekdb.client.schema import Schema + + +RECYCLEBIN_PREFIX = "__recyclebin_" + + +# --------------------------------------------------------------------------- # +# helpers # +# --------------------------------------------------------------------------- # + + +def _make_collection(client, suffix: str = ""): + """Create a namespace-mode collection with an IVF index (so we get a full + set of physical tables: logic_data / kv_data / logic_schema / hot_table).""" + name = f"test_ns_drop_{int(time.time() * 1000)}{suffix}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine"), + embedding_function=None, + ), + ) + return client.create_collection(name=name, schema=schema, use_namespace=True) + + +def _execute(client, sql: str): + return client._server._execute(sql) + + +def _is_ss_mode(client) -> bool: + """Best-effort detect whether the connected OB cluster is in SS mode.""" + try: + rows = _execute( + client, + "SHOW PARAMETERS LIKE 'enable_logservice'", + ) + for r in rows: + value = (r.get("value") or r.get("VALUE") or "").lower() + if value in ("true", "1", "on"): + return True + except Exception: + pass + return False + + +def _fetch_namespace_name(client, collection_id: str, namespace_id: int): + rows = _execute( + client, + f"SELECT namespace_name FROM sdk_ns_namespaces " + f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id}", + ) + if not rows: + return None + return rows[0].get("namespace_name") or rows[0].get("NAMESPACE_NAME") + + +def _count_ltables(client, collection_id: str, namespace_id: int) -> int: + rows = _execute( + client, + f"SELECT COUNT(*) AS c FROM sdk_ns_ltables " + f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id}", + ) + return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) + + +def _count_logic_schema_rows(client, collection_id: str, namespace_id: int) -> int: + tbl = NamespaceCollectionNames.logic_schema_table_name(collection_id) + rows = _execute( + client, + f"SELECT COUNT(*) AS c FROM `{tbl}` WHERE namespace_id = {namespace_id}", + ) + return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) + + +def _count_hot_table_rows(client, collection_id: str, namespace_id: int) -> int: + tbl = NamespaceCollectionNames.hot_table_name(collection_id) + try: + rows = _execute( + client, + f"SELECT COUNT(*) AS c FROM `{tbl}` WHERE namespace_id = {namespace_id}", + ) + except Exception: + return -1 # table missing + return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) + + +def _seed_hot_table(client, collection_id: str, namespace_id: int): + """In SS mode insert a synthetic hot_table row so the DELETE has something + to remove and we can verify the cleanup.""" + tbl = NamespaceCollectionNames.hot_table_name(collection_id) + try: + _execute( + client, + f"INSERT INTO `{tbl}` (namespace_id, last_access_time) " + f"VALUES ({namespace_id}, NOW(6))", + ) + return True + except Exception: + return False + + +def _drop_namespace_via_pl(client, collection_id: str, namespace_id: int): + """Call DBMS_LOGIC_TABLE.DROP_NAMESPACE directly (raw SQL, bypassing + SDK's pre-existence guard).""" + _execute( + client, + f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id}', {namespace_id})", + ) + + +def _second_oceanbase_client(): + """Second pymysql-backed client (separate TCP connection) for concurrency tests.""" + import os + + client = pyseekdb.Client( + host=os.environ.get("OB_HOST", "localhost"), + port=int(os.environ.get("OB_PORT", "11202")), + tenant=os.environ.get("OB_TENANT", "mysql"), + database=os.environ.get("OB_DATABASE", "test"), + user=os.environ.get("OB_USER", "root"), + password=os.environ.get("OB_PASSWORD", ""), + ) + result = client._server._execute("SELECT 1 as test") + assert result and result[0].get("test") == 1 + return client + + +# --------------------------------------------------------------------------- # +# tests # +# --------------------------------------------------------------------------- # + + +class TestDropNamespaceCatalogValidation: + """Validate everything the DROP_NAMESPACE synchronous transaction must do.""" + + # ------------------------------------------------------------- # + # 1. Happy path: rename + cleanup all relevant rows # + # ------------------------------------------------------------- # + def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): + client = oceanbase_client + is_ss = _is_ss_mode(client) + collection = _make_collection(client) + coll_id = collection.id + try: + ns = collection.create_namespace("ns_validate") + ns_id = int(ns.namespace_id) + + # Pre-conditions + orig_name = _fetch_namespace_name(client, coll_id, ns_id) + assert orig_name == "ns_validate" + assert _count_ltables(client, coll_id, ns_id) >= 1, ( + "create_namespace should have inserted a default ltable" + ) + assert _count_logic_schema_rows(client, coll_id, ns_id) >= 1, ( + "create_namespace should have inserted a logic_schema row" + ) + if is_ss: + # Seed a hot_table row so we can verify the SS-only delete. + assert _seed_hot_table(client, coll_id, ns_id) is True + assert _count_hot_table_rows(client, coll_id, ns_id) == 1 + + # Action + _drop_namespace_via_pl(client, coll_id, ns_id) + + # Post-conditions + new_name = _fetch_namespace_name(client, coll_id, ns_id) + assert new_name is not None, "namespace row must still exist (renamed, not deleted)" + assert new_name.startswith(RECYCLEBIN_PREFIX), ( + f"namespace must be renamed to __recyclebin_ prefix, got {new_name!r}" + ) + assert "ns_validate" in new_name, ( + f"original name should be embedded in recyclebin name, got {new_name!r}" + ) + assert _count_ltables(client, coll_id, ns_id) == 0, ( + "sdk_ns_ltables rows must be deleted" + ) + assert _count_logic_schema_rows(client, coll_id, ns_id) == 0, ( + "logic_schema_table rows must be deleted" + ) + if is_ss: + assert _count_hot_table_rows(client, coll_id, ns_id) == 0, ( + "hot_table row must be deleted in SS mode" + ) + finally: + client.delete_collection(name=collection.name) + + # ------------------------------------------------------------- # + # 2. Namespace-name rename format # + # ------------------------------------------------------------- # + def test_recyclebin_name_format(self, oceanbase_client): + client = oceanbase_client + collection = _make_collection(client) + try: + ns = collection.create_namespace("fmt_ns") + ns_id = int(ns.namespace_id) + _drop_namespace_via_pl(client, collection.id, ns_id) + new_name = _fetch_namespace_name(client, collection.id, ns_id) + # Expected format: __recyclebin__ + assert new_name.startswith("__recyclebin_fmt_ns_") + ts_suffix = new_name[len("__recyclebin_fmt_ns_"):] + assert ts_suffix.isdigit() and len(ts_suffix) >= 16, ( + f"trailing timestamp_us looks malformed: {new_name!r}" + ) + finally: + client.delete_collection(name=collection.name) + + # ------------------------------------------------------------- # + # 3. Multiple ltables under one namespace all get deleted # + # ------------------------------------------------------------- # + def test_multiple_ltables_all_deleted(self, oceanbase_client): + client = oceanbase_client + collection = _make_collection(client) + try: + ns = collection.create_namespace("ns_multi") + ns_id = int(ns.namespace_id) + # Insert two extra ltable rows directly so we can verify bulk delete. + _execute( + client, + f"INSERT INTO sdk_ns_ltables (collection_id, namespace_id, ltable_name) " + f"VALUES ('{collection.id}', {ns_id}, 'extra_lt_a')," + f" ('{collection.id}', {ns_id}, 'extra_lt_b')", + ) + assert _count_ltables(client, collection.id, ns_id) >= 3 + + _drop_namespace_via_pl(client, collection.id, ns_id) + + assert _count_ltables(client, collection.id, ns_id) == 0 + finally: + client.delete_collection(name=collection.name) + + # ------------------------------------------------------------- # + # 4. Atomic rollback when an in-transaction DELETE fails # + # ------------------------------------------------------------- # + def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): + """If we pre-DROP _logic_schema_table the in-trans DELETE on it + fails with OB_TABLE_NOT_EXIST, which must roll back the namespace + rename and the sdk_ns_ltables delete that ran earlier in the same + transaction.""" + client = oceanbase_client + collection = _make_collection(client) + coll_id = collection.id + schema_tbl = NamespaceCollectionNames.logic_schema_table_name(coll_id) + try: + ns = collection.create_namespace("ns_rollback") + ns_id = int(ns.namespace_id) + orig_name = _fetch_namespace_name(client, coll_id, ns_id) + assert orig_name == "ns_rollback" + ltable_cnt_before = _count_ltables(client, coll_id, ns_id) + assert ltable_cnt_before >= 1 + + # Sabotage: drop the logic_schema_table so step 6 inside the + # DROP_NAMESPACE transaction will fail with OB_TABLE_NOT_EXIST. + _execute(client, f"DROP TABLE `{schema_tbl}`") + + with pytest.raises(Exception): + _drop_namespace_via_pl(client, coll_id, ns_id) + + # Everything before the failed step must have been rolled back. + after_name = _fetch_namespace_name(client, coll_id, ns_id) + assert after_name == "ns_rollback", ( + f"namespace_name must be rolled back to original, got {after_name!r}" + ) + assert _count_ltables(client, coll_id, ns_id) == ltable_cnt_before, ( + "sdk_ns_ltables delete must be rolled back" + ) + finally: + # Recreate the schema table (empty) so cleanup_namespace_physical_tables + # / delete_collection runs cleanly. + try: + _execute( + client, + f"CREATE TABLE IF NOT EXISTS `{schema_tbl}` (" + f" namespace_id BIGINT UNSIGNED NOT NULL," + f" ltable_id BIGINT UNSIGNED NOT NULL," + f" schema_content JSON NOT NULL," + f" PRIMARY KEY (namespace_id, ltable_id)) " + f"PARTITION BY KEY(namespace_id) PARTITIONS 8", + ) + except Exception: + pass + client.delete_collection(name=collection.name) + + # ------------------------------------------------------------- # + # 5. Idempotent: second drop on the same ns is a no-op success # + # ------------------------------------------------------------- # + def test_drop_namespace_is_idempotent(self, oceanbase_client): + client = oceanbase_client + collection = _make_collection(client) + coll_id = collection.id + try: + ns = collection.create_namespace("ns_idem") + ns_id = int(ns.namespace_id) + _drop_namespace_via_pl(client, coll_id, ns_id) + first_name = _fetch_namespace_name(client, coll_id, ns_id) + assert first_name.startswith(RECYCLEBIN_PREFIX) + + # Second call must succeed without raising and must NOT re-rename + # (the row name stays exactly the same — no double prefix). + _drop_namespace_via_pl(client, coll_id, ns_id) + second_name = _fetch_namespace_name(client, coll_id, ns_id) + assert second_name == first_name, ( + f"second DROP_NAMESPACE must be no-op, got {second_name!r} vs {first_name!r}" + ) + # And LEFT(name, 26) must not be '__recyclebin___recyclebin_' + assert not second_name.startswith("__recyclebin___recyclebin_") + finally: + client.delete_collection(name=collection.name) + + # ------------------------------------------------------------- # + # 6. Concurrent drop from two SDK connections # + # ------------------------------------------------------------- # + def test_concurrent_drop_from_two_sdks(self, oceanbase_client): + """Two independent SDK connections race to DROP_NAMESPACE on the same + (collection_id, namespace_id). + + Expected: exactly one transaction performs the rename + deletes; the + other observes the __recyclebin_ row inside its own short trans and + commits an empty trans (no exception, no double prefix). Both calls + return success.""" + client_a = oceanbase_client + client_b = _second_oceanbase_client() + + collection = _make_collection(client_a, suffix="_conc") + coll_id = collection.id + try: + ns = collection.create_namespace("ns_concurrent") + ns_id = int(ns.namespace_id) + + results = {} + barrier = threading.Barrier(2) + + def _worker(tag, c): + try: + barrier.wait(timeout=10) + _drop_namespace_via_pl(c, coll_id, ns_id) + results[tag] = "ok" + except Exception as exc: # noqa: BLE001 + results[tag] = f"err: {exc!r}" + + ta = threading.Thread(target=_worker, args=("A", client_a)) + tb = threading.Thread(target=_worker, args=("B", client_b)) + ta.start(); tb.start() + ta.join(timeout=30); tb.join(timeout=30) + + # Both calls must succeed (one does the work, the other no-ops). + assert results.get("A") == "ok", f"A failed: {results.get('A')!r}" + assert results.get("B") == "ok", f"B failed: {results.get('B')!r}" + + final_name = _fetch_namespace_name(client_a, coll_id, ns_id) + assert final_name.startswith(RECYCLEBIN_PREFIX) + assert "ns_concurrent" in final_name + # The pattern must be __recyclebin__ — never doubled. + assert not final_name.startswith("__recyclebin___recyclebin_"), ( + f"concurrent drops produced double-prefixed name: {final_name!r}" + ) + # ltables and logic_schema must be empty (one of the trans cleared them). + assert _count_ltables(client_a, coll_id, ns_id) == 0 + assert _count_logic_schema_rows(client_a, coll_id, ns_id) == 0 + finally: + try: + client_a.delete_collection(name=collection.name) + finally: + if hasattr(client_b, "close"): + try: + client_b.close() + except Exception: + pass + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) From 69542a1fedf6f5636f715afd99164ff67a2f93cd Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 27 May 2026 16:42:11 +0800 Subject: [PATCH 15/84] =?UTF-8?q?logic=20table=20id=E9=87=87=E7=94=A8?= =?UTF-8?q?=E7=9C=9F=E5=AE=9E=E5=88=86=E9=85=8D=E7=9A=84=E5=80=BC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 90 ++++++++++++++++----- src/pyseekdb/client/client_seekdb_server.py | 3 +- tests/unit_tests/test_namespace.py | 5 ++ 3 files changed, 78 insertions(+), 20 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index fadef14b..995e2239 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1302,6 +1302,50 @@ def _cleanup_namespace_physical_tables(self, collection_id: str) -> None: with contextlib.suppress(Exception): self._execute(f"DROP TABLEGROUP IF EXISTS `{NamespaceCollectionNames.tablegroup_name(collection_id)}`") + def _resolve_namespace_ltable_id( + self, collection_id: str, namespace_id: int + ) -> int: + """Resolve the default ltable_id for (collection_id, namespace_id) from + sdk_ns_ltables. Cached per (collection_id, namespace_id) on the client + instance to avoid the extra round-trip on every DML/DQL call. + + The cache is also seeded by `_create_ns_namespace_meta` and + `_get_ns_namespace_meta` once they have read/created the row. + """ + key = (str(collection_id), int(namespace_id)) + cache = getattr(self, "_ns_ltable_id_cache", None) + if cache is None: + cache = {} + self._ns_ltable_id_cache = cache + cached = cache.get(key) + if cached is not None: + return cached + coll_id_escaped = escape_string(str(collection_id)) + rows = self._execute( + f"SELECT ltable_id FROM `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"WHERE collection_id = '{coll_id_escaped}' AND namespace_id = {int(namespace_id)} " + f"AND ltable_name = 'default' LIMIT 1" + ) + if not rows: + raise ValueError( + f"No default ltable found for collection_id={collection_id}, " + f"namespace_id={namespace_id} in sdk_ns_ltables" + ) + row = rows[0] + lt_id = int(row[0] if isinstance(row, (list, tuple)) else row["ltable_id"]) + cache[key] = lt_id + return lt_id + + def _cache_namespace_ltable_id( + self, collection_id: str, namespace_id: int, ltable_id: int + ) -> None: + key = (str(collection_id), int(namespace_id)) + cache = getattr(self, "_ns_ltable_id_cache", None) + if cache is None: + cache = {} + self._ns_ltable_id_cache = cache + cache[key] = int(ltable_id) + def _set_session_ns_context( self, collection_id: str | None = None, @@ -1336,6 +1380,7 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {ns_id} AND ltable_name = 'default'" ) lt_id = int(lt_rows[0][0] if isinstance(lt_rows[0], (list, tuple)) else lt_rows[0]["ltable_id"]) + self._cache_namespace_ltable_id(collection_id, ns_id, lt_id) schema_table = NamespaceCollectionNames.logic_schema_table_name(collection_id) schema_content = json.dumps(_build_default_ltable_schema()) with contextlib.suppress(Exception): @@ -1374,6 +1419,8 @@ def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dic ns_name = row["namespace_name"] lt_raw = row.get("ltable_id") lt_id = int(lt_raw) if lt_raw is not None else None + if lt_id is not None: + self._cache_namespace_ltable_id(collection_id, int(ns_id), lt_id) self._set_session_ns_context(namespace_id=int(ns_id), ltable_id=lt_id) meta: dict = {"namespace_id": ns_id, "namespace_name": ns_name} if lt_id is not None: @@ -4461,7 +4508,7 @@ def _append_namespace_filter( where_clause: str, params: list[Any], namespace_id: int, - ltable_id: int = 1, + ltable_id: int, ) -> tuple[str, list[Any]]: ns_cond = f"namespace_id = {namespace_id} AND ltable_id = {ltable_id}" if not where_clause: @@ -4484,8 +4531,9 @@ def _namespace_add( # noqa: C901 embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) if isinstance(ids, str): ids = [ids] @@ -4528,7 +4576,6 @@ def _namespace_add( # noqa: C901 table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) - ltable_id = 1 values_list = [] for i in range(num_items): @@ -4567,8 +4614,9 @@ def _namespace_update( embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) if isinstance(ids, str): ids = [ids] @@ -4594,7 +4642,6 @@ def _namespace_update( table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) - ltable_id = 1 id_expr = "JSON_EXTRACT(data_content, '$.id')" active_ids = [] @@ -4700,8 +4747,9 @@ def _namespace_upsert( embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) if isinstance(ids, str): ids = [ids] @@ -4724,7 +4772,6 @@ def _namespace_upsert( table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) - ltable_id = 1 existing_ids = set() id_expr = "JSON_EXTRACT(data_content, '$.id')" @@ -4786,8 +4833,9 @@ def _namespace_delete( where_document: dict[str, Any] | None = None, **kwargs, ) -> None: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) if ids is None and where is None and where_document is None: raise ValueError("At least one of ids, where, or where_document must be provided") @@ -4820,7 +4868,7 @@ def _namespace_delete( params.extend(doc_params) user_where = f"WHERE {' AND '.join(conditions)}" if conditions else "" - where_clause, params = self._append_namespace_filter(user_where, params, ns_id) + where_clause, params = self._append_namespace_filter(user_where, params, ns_id, ltable_id) sql = f"DELETE FROM `{table_name}` {where_clause}" if params: conn = self._ensure_connection() @@ -4851,8 +4899,9 @@ def _namespace_query( # noqa: C901 include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) embedding_function = kwargs.get("embedding_function") distance = kwargs.get("distance", DEFAULT_DISTANCE_METRIC) @@ -4898,7 +4947,7 @@ def _namespace_query( # noqa: C901 filter_params.extend(doc_params) user_where = f"WHERE {' AND '.join(user_conditions)}" if user_conditions else "" - where_clause, filter_params = self._append_namespace_filter(user_where, filter_params, ns_id) + where_clause, filter_params = self._append_namespace_filter(user_where, filter_params, ns_id, ltable_id) distance_function_map = { "l2": "l2_distance", @@ -5009,8 +5058,9 @@ def _namespace_get( # noqa: C901 include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) include_fields = self._normalize_include_fields(include) table_name = NamespaceCollectionNames.data_table_name(collection_id) @@ -5050,7 +5100,7 @@ def _namespace_get( # noqa: C901 params.extend(doc_params) user_where = f"WHERE {' AND '.join(user_conditions)}" if user_conditions else "" - where_str, params = self._append_namespace_filter(user_where, params, ns_id) + where_str, params = self._append_namespace_filter(user_where, params, ns_id, ltable_id) select_clause = ", ".join(select_parts) sql = f"SELECT {select_clause} FROM `{table_name}` {where_str}" if limit is not None: @@ -5123,12 +5173,13 @@ def _namespace_count( namespace_name: str, **kwargs, ) -> int: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) - where_clause, _ = self._append_namespace_filter("", [], ns_id) + where_clause, _ = self._append_namespace_filter("", [], ns_id, ltable_id) sql = f"SELECT COUNT(*) AS cnt FROM `{table_name}` {where_clause}" conn = self._ensure_connection() use_context_manager = self._use_context_manager_for_cursor() @@ -5161,10 +5212,10 @@ def _namespace_peek( include=["documents", "metadatas", "embeddings"], ) - def _adapt_search_parm_for_ns(self, search_parm: dict[str, Any], ns_id: int) -> dict[str, Any]: + def _adapt_search_parm_for_ns(self, search_parm: dict[str, Any], ns_id: int, lt_id: int) -> dict[str, Any]: ns_filter = [ {"term": {"namespace_id": ns_id}}, - {"term": {"ltable_id": 1}}, + {"term": {"ltable_id": int(lt_id)}}, ] def _rewrite_field_refs(obj): @@ -5262,8 +5313,9 @@ def _namespace_hybrid_search( query_hint: QueryHint | None = None, **kwargs, ) -> dict[str, Any]: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) conn = self._ensure_connection() table_name = NamespaceCollectionNames.data_table_name(collection_id) @@ -5275,7 +5327,7 @@ def _namespace_hybrid_search( dimension=kwargs.get("dimension"), **{k: v for k, v in kwargs.items() if k != "dimension"}, ) - search_parm = self._adapt_search_parm_for_ns(search_parm, ns_id) + search_parm = self._adapt_search_parm_for_ns(search_parm, ns_id, ltable_id) search_parm.pop("_source", None) diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 5b204f87..31e867fc 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -208,8 +208,9 @@ def _namespace_prewarm( namespace_name: str, **kwargs, ) -> None: + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=1, + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) sql = f"CALL DBMS_LOGIC_TABLE.PREWARM('{collection_id}', {namespace_id})" self._execute(sql) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 0db30bed..b379b836 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -397,6 +397,11 @@ def _execute(self, sql, params=None): self.executed_sqls.append(sql) return None + # Bypass the sdk_ns_ltables lookup in unit tests: SQL-generation tests don't + # have a real database, so return a fixed ltable_id matching test assertions. + def _resolve_namespace_ltable_id(self, collection_id, namespace_id): + return 1 + def _execute_query_with_cursor(self, conn, sql, params, use_context_manager=True): resolved = sql for p in params: From 33ab2c31e8439a13275c3f65326b5f0ddc771924 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 28 May 2026 11:47:00 +0800 Subject: [PATCH 16/84] =?UTF-8?q?fts=E9=80=82=E9=85=8Ddefault=5Foperator:a?= =?UTF-8?q?nd?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 995e2239..a38b6d07 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -4043,6 +4043,7 @@ def _apply_boost(target: Any) -> None: "query_string": { "fields": ["document"], "query": " ".join(escaped_queries), + "default_operator": "and", } }) From 4ae23cdf02aed83e4aa41845e1262a6b1e9cf9f0 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 28 May 2026 14:35:05 +0800 Subject: [PATCH 17/84] =?UTF-8?q?or=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index a38b6d07..7bb07b16 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -4061,7 +4061,8 @@ def _apply_boost(target: Any) -> None: return _with_boost({ "query_string": { "fields": ["document"], - "query": " OR ".join(escaped_queries), + "query": " ".join(escaped_queries), + "default_operator": "or", } }) From 5d10d7303898a0de0c6e6b617d87ad4603eb9ba5 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 1 Jun 2026 15:54:03 +0800 Subject: [PATCH 18/84] =?UTF-8?q?=E4=BF=AE=E5=A4=8Djson=5Fsearch=E5=8F=8C?= =?UTF-8?q?=E5=BC=95=E5=8F=B7=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7bb07b16..5628f7bb 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -3841,6 +3841,17 @@ def _collection_hybrid_search( # Remove any surrounding quotes if present query_sql = query_sql.strip().strip("'\"") + # OB's GET_SQL wraps field names in backticks, which turns + # `JSON_EXTRACT(metadata, '$.key')` (with or without outer + # parentheses) into a literal column name instead of a + # function call. Strip the backticks so OB evaluates the + # expression as a function call. + query_sql = re.sub( + r"`([^`]*JSON_EXTRACT[^`]*)`", + r"\1", + query_sql, + ) + # Add query hint to the generated SQL hint_sql = _query_hint_to_sql(query_hint, table_name=table_name) if hint_sql and query_sql.upper().startswith("SELECT"): From 4afb949be3ae96be8b0a1060cf9926dc55b1808c Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Mon, 1 Jun 2026 16:59:15 +0800 Subject: [PATCH 19/84] add namespace case --- src/pyseekdb/client/client_base.py | 48 +- tests/integration_tests/conftest.py | 6 +- .../namespace_dml_helpers.py | 79 ++ .../namespace_fts_helpers.py | 708 ++++++++++++++++++ .../namespace_hybrid_search_helpers.py | 627 ++++++++++++++++ tests/integration_tests/test_namespace_dml.py | 75 ++ .../test_namespace_dml_multi_coll.py | 67 ++ .../test_namespace_dml_multi_coll_multi_ns.py | 79 ++ .../test_namespace_dml_multi_ns.py | 99 +++ .../test_namespace_hybrid_search_combined.py | 69 ++ .../test_namespace_hybrid_search_fulltext.py | 136 ++++ ...rid_search_fulltext_multi_coll_multi_ns.py | 109 +++ ...space_hybrid_search_multi_coll_multi_ns.py | 149 ++++ ...st_namespace_hybrid_search_search_index.py | 69 ++ .../test_namespace_hybrid_search_vector.py | 83 ++ 15 files changed, 2378 insertions(+), 25 deletions(-) create mode 100644 tests/integration_tests/namespace_dml_helpers.py create mode 100644 tests/integration_tests/namespace_fts_helpers.py create mode 100644 tests/integration_tests/namespace_hybrid_search_helpers.py create mode 100644 tests/integration_tests/test_namespace_dml_multi_coll.py create mode 100644 tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py create mode 100644 tests/integration_tests/test_namespace_dml_multi_ns.py create mode 100644 tests/integration_tests/test_namespace_hybrid_search_combined.py create mode 100644 tests/integration_tests/test_namespace_hybrid_search_fulltext.py create mode 100644 tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py create mode 100644 tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py create mode 100644 tests/integration_tests/test_namespace_hybrid_search_search_index.py create mode 100644 tests/integration_tests/test_namespace_hybrid_search_vector.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7cd53704..cec881eb 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1247,15 +1247,13 @@ def _create_namespace_physical_tables( created_by VARCHAR(64) DEFAULT '', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FULLTEXT INDEX idx_fts(document) {fulltext_clause}, - SEARCH INDEX idx_json(data_content) + SEARCH INDEX idx_json(data_content), + VECTOR INDEX idx_vec(embedding) {vector_index_sql} ) TABLEGROUP=`{tg_name}` COMMENT='逻辑表主数据' DEFAULT CHARSET=utf8mb4 ORGANIZATION HEAP IS_LOGIC_TABLE = TRUE {partition_clause}""" self._execute(data_sql) - ivf_index_sql = f"CREATE VECTOR INDEX idx_vec ON `{data_table}` (embedding) {vector_index_sql}" - self._execute(ivf_index_sql) - if is_shared_storage: hot_table = NamespaceCollectionNames.hot_table_name(collection_id) self._execute(f"""CREATE TABLE `{hot_table}` ( @@ -4061,7 +4059,8 @@ def _apply_boost(target: Any) -> None: return _with_boost({ "query_string": { "fields": ["document"], - "query": " OR ".join(escaped_queries), + "query": " ".join(escaped_queries), + "default_operator": "or", } }) @@ -4300,6 +4299,24 @@ def _build_source_fields(self, include: list[str] | None) -> list[str]: return source + def _hybrid_row_score(self, row: dict[str, Any]) -> float: + """Extract relevance score from a hybrid_search SQL row (OB uses ``__score``).""" + for key in ( + "_distance", + "distance", + "_score", + "score", + "__score", + "DISTANCE", + "_DISTANCE", + "SCORE", + "__SCORE", + ): + val = row.get(key) + if val is not None: + return float(val) + return 0.0 + def _transform_sql_result( # noqa: C901 self, result_rows: list[dict[str, Any]], include: list[str] | None ) -> dict[str, Any]: @@ -4344,18 +4361,7 @@ def _transform_sql_result( # noqa: C901 row_id = self._convert_id_from_bytes(row_id) ids.append(row_id) - # Extract distance/score (may be in different column names) - distance = ( - row.get("_distance") - or row.get("distance") - or row.get("_score") - or row.get("score") - or row.get("DISTANCE") - or row.get("_DISTANCE") - or row.get("SCORE") - or 0.0 - ) - distances.append(distance) + distances.append(self._hybrid_row_score(row)) # Extract metadata if include is None or "metadatas" in include or "metadata" in include: @@ -5342,7 +5348,7 @@ def _namespace_hybrid_search( f"SELECT {hint_sql + ' ' if hint_sql else ''}* " f"FROM hybrid_search(TABLE `{table_name}`, '{escaped_params}')" ) - + print(f"hybrid_sql: {hybrid_sql}") result_rows = self._execute_query_with_cursor(conn, hybrid_sql, [], use_context_manager) if not result_rows: return { @@ -5384,11 +5390,7 @@ def _transform_ns_hybrid_result( row_id = self._convert_id_from_bytes(row_id) ids.append(row_id) - distance = ( - row.get("_distance") or row.get("distance") - or row.get("_score") or row.get("score") or 0.0 - ) - distances.append(distance) + distances.append(self._hybrid_row_score(row)) if include is None or "metadatas" in include or "metadata" in include: meta = dc.get("metadata") diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index 9f930cdc..b8e805db 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -10,9 +10,11 @@ import pytest -# Add project path (repo root + src) -repo_root = Path(__file__).resolve().parents[2] +# Add project path (repo root + src + this directory for local test helpers) +integration_tests_root = Path(__file__).resolve().parent +repo_root = integration_tests_root.parents[1] src_root = repo_root / "src" +sys.path.insert(0, str(integration_tests_root)) sys.path.insert(0, str(src_root)) import pyseekdb # noqa: E402 diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py new file mode 100644 index 00000000..d4293859 --- /dev/null +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -0,0 +1,79 @@ +""" +Shared helpers for namespace DML integration tests. +""" + +from __future__ import annotations + +import time +from typing import Any + +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.schema import Schema + + +def ns_schema() -> Schema: + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2"), + embedding_function=None, + ), + ) + + +def create_ns_collection(client: Any, suffix: str = "") -> Any: + name = f"test_ns_dml_{int(time.time() * 1000)}{suffix}" + return client.create_collection(name=name, schema=ns_schema(), use_namespace=True) + + +def cleanup(client: Any, *collections: Any) -> None: + for collection in collections: + try: + client.delete_collection(name=collection.name) + except Exception: + pass + + +def _default_include( + documents: str | None, + metadatas: dict | None, + embeddings: list[float] | None, + include: list[str] | None, +) -> list[str] | None: + if include is not None: + return include + fields: list[str] = [] + if documents is not None: + fields.append("documents") + if metadatas is not None: + fields.append("metadatas") + if embeddings is not None: + fields.append("embeddings") + return fields or None + + +def assert_get_present( + ns: Any, + doc_id: str, + *, + documents: str | None = None, + metadatas: dict | None = None, + embeddings: list[float] | None = None, + include: list[str] | None = None, +) -> dict[str, Any]: + result = ns.get(ids=doc_id, include=_default_include(documents, metadatas, embeddings, include)) + indices = [i for i, rid in enumerate(result["ids"]) if rid == doc_id] + assert indices, f"expected doc {doc_id!r}, got {result['ids']}" + idx = indices[0] + if documents is not None: + assert result["documents"][idx] == documents + if metadatas is not None: + assert result["metadatas"][idx] == metadatas + if embeddings is not None: + assert result["embeddings"][idx] == embeddings + return result + + +def assert_get_absent(ns: Any, doc_id: str) -> None: + result = ns.get(ids=doc_id) + assert len(result["ids"]) == 0, f"expected doc {doc_id!r} absent, got {result['ids']}" diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py new file mode 100644 index 00000000..044c640c --- /dev/null +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -0,0 +1,708 @@ +""" +Helpers for namespace hybrid_search full-text integration tests. + +Ground-truth matching uses substring checks on document text (aligned with +``$contains`` / ``$not_contains`` semantics). Relevance ranking uses term +occurrence counts so the highest-ranked documents are deterministic. +""" + +from __future__ import annotations + +import contextlib +import math +import time +from dataclasses import dataclass +from typing import Any + +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.schema import Schema + +# Rare ASCII tokens to reduce IK segmentation surprises in assertions. +TOKEN_ZPX = "TOKENZPX" +TOKEN_ALP = "TOKENALP" +FORBIDDEN_PHRASE = "FORBIDPHRASE" + +# Tier base for $or matches that hit multiple branches (query_string OR multi-term boost). +_OR_MULTI_BRANCH_HINT_BASE = 10_000 + +CORPUS_SIZE = 1101 # strictly > 1000 +BATCH_SIZE = 100 +INDEX_SETTLE_SECONDS = 3 + + +@dataclass(frozen=True) +class CorpusRecord: + doc_id: str + document: str + embedding: list[float] + metadata: dict[str, Any] + rel_hint: int # ground-truth relevance tier for positive FTS queries + + +@dataclass(frozen=True) +class FtsQueryCase: + name: str + where_document: dict[str, Any] | str + n_results: int + check_ranking: bool = True + where: dict[str, Any] | None = None + + +def ns_schema() -> Schema: + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2"), + embedding_function=None, + ), + ) + + +def build_large_fts_corpus(size: int = CORPUS_SIZE) -> list[CorpusRecord]: + """Build a deterministic corpus with known full-text relevance tiers.""" + records: list[CorpusRecord] = [] + + def _add( + doc_id: str, + document: str, + rel_hint: int, + embedding: list[float] | None = None, + extra_meta: dict[str, Any] | None = None, + ) -> None: + idx = len(records) + meta = {"rel_hint": rel_hint, "seq": idx} + if extra_meta: + meta.update(extra_meta) + records.append( + CorpusRecord( + doc_id=doc_id, + document=document, + embedding=embedding or _embedding_for_index(idx), + metadata=meta, + rel_hint=rel_hint, + ) + ) + + # Tiered hits for TOKEN_ZPX ($contains / string shorthand). + for repeat in range(5, 0, -1): + zpx_hint = repeat * 10 + _add( + f"zpx_top_{repeat}", + f"Primary {TOKEN_ZPX} " * repeat + f" focus document repeat {repeat}", + rel_hint=zpx_hint, + extra_meta={"zpx_hint": zpx_hint, "alp_hint": 0}, + ) + + for i in range(15): + _add( + f"zpx_mid_{i:02d}", + f"Secondary mention {TOKEN_ZPX} in document index {i}", + rel_hint=5, + extra_meta={"zpx_hint": 5, "alp_hint": 0}, + ) + + # Both tokens for $and. + _add( + "both_tokens", + f"Combined {TOKEN_ZPX} and {TOKEN_ALP} in one document for conjunction tests", + rel_hint=55, + extra_meta={"zpx_hint": 5, "alp_hint": 5, "has_both": True}, + ) + + # TOKEN_ALP tiered hits for $or. + for repeat in range(4, 0, -1): + alp_hint = repeat * 8 + _add( + f"alp_top_{repeat}", + f"Alpha {TOKEN_ALP} " * repeat + f" branch document repeat {repeat}", + rel_hint=alp_hint, + extra_meta={"zpx_hint": 0, "alp_hint": alp_hint}, + ) + + for i in range(10): + _add( + f"alp_mid_{i:02d}", + f"Secondary {TOKEN_ALP} mention variant {i}", + rel_hint=4, + extra_meta={"zpx_hint": 0, "alp_hint": 4}, + ) + + # Rows containing TOKEN_ZPX — must not appear in $not_contains results. + for i in range(8): + _add( + f"zpx_forbid_{i}", + f"This row contains {TOKEN_ZPX} and must not appear in not_contains results {i}", + rel_hint=0, + extra_meta={"zpx_hint": 1, "alp_hint": 0}, + ) + + # Filler rows (no rare tokens) — bulk of the corpus. + while len(records) < size: + i = len(records) + _add( + f"filler_{i:05d}", + f"Filler document {i} about oceanbase distributed storage systems. " + f"Neutral content without rare tokens. {FORBIDDEN_PHRASE} anchor {i % 17}", + rel_hint=1 if i % 50 == 0 else 0, + ) + + if len(records) > size: + del records[size:] + return records + + +def _embedding_for_index(index: int) -> list[float]: + return [ + float((index + 1) % 7) / 7.0, + float((index + 2) % 5) / 5.0, + float((index + 3) % 3) / 3.0, + ] + + +def doc_matches_where_document(document: str, where_document: dict[str, Any] | str) -> bool: + text = document.lower() + if isinstance(where_document, str): + return where_document.lower() in text + if "$contains" in where_document: + return where_document["$contains"].lower() in text + if "$not_contains" in where_document: + return where_document["$not_contains"].lower() not in text + if "$and" in where_document: + return all( + doc_matches_where_document(document, sub) + for sub in where_document["$and"] + ) + if "$or" in where_document: + return any( + doc_matches_where_document(document, sub) + for sub in where_document["$or"] + ) + raise ValueError(f"Unsupported where_document for hybrid_search: {where_document!r}") + + +def doc_matches_where_metadata(metadata: dict[str, Any], where: dict[str, Any]) -> bool: + """Evaluate metadata filter (aligned with hybrid_search ``query.where`` semantics).""" + if "$and" in where: + return all(doc_matches_where_metadata(metadata, sub) for sub in where["$and"]) + if "$or" in where: + return any(doc_matches_where_metadata(metadata, sub) for sub in where["$or"]) + if "$not" in where: + return not doc_matches_where_metadata(metadata, where["$not"]) + + for key, value in where.items(): + if key in ("$and", "$or", "$not"): + continue + actual = metadata.get(key) + if isinstance(value, dict): + if "$eq" in value and actual != value["$eq"]: + return False + if "$ne" in value and actual == value["$ne"]: + return False + if "$lt" in value and (actual is None or actual >= value["$lt"]): + return False + if "$lte" in value and (actual is None or actual > value["$lte"]): + return False + if "$gt" in value and (actual is None or actual <= value["$gt"]): + return False + if "$gte" in value and (actual is None or actual < value["$gte"]): + return False + if "$in" in value and actual not in value["$in"]: + return False + if "$nin" in value and actual in value["$nin"]: + return False + elif actual != value: + return False + return True + + +def corpus_matches_fts( + record: CorpusRecord, + where_document: dict[str, Any] | str, + where: dict[str, Any] | None = None, +) -> bool: + if not doc_matches_where_document(record.document, where_document): + return False + if where is not None and not doc_matches_where_metadata(record.metadata, where): + return False + return True + + +def effective_rel_hint(record: CorpusRecord, where_document: dict[str, Any] | str) -> int: + """Map corpus ``rel_hint`` to the active query (AND uses min; OR uses max or multi-branch sum).""" + if isinstance(where_document, str): + if where_document == TOKEN_ZPX: + return record.metadata.get("zpx_hint", record.rel_hint) + if where_document == TOKEN_ALP: + return record.metadata.get("alp_hint", record.rel_hint) + return record.rel_hint + if "$contains" in where_document: + term = where_document["$contains"] + if term == TOKEN_ZPX: + return record.metadata.get("zpx_hint", 0) + if term == TOKEN_ALP: + return record.metadata.get("alp_hint", 0) + if term in ("Primary", "Alpha"): + return record.rel_hint + return record.rel_hint + if "$and" in where_document: + hints = [effective_rel_hint(record, sub) for sub in where_document["$and"]] + return min(hints) if hints else 0 + if "$or" in where_document: + hints = [effective_rel_hint(record, sub) for sub in where_document["$or"]] + active = [h for h in hints if h > 0] + if not active: + return 0 + if len(active) >= 2: + # Matches multiple OR branches (e.g. both TOKENZPX and TOKENALP): ranks above + # any single-branch hit, consistent with query_string OR + BM25 on OceanBase. + return _OR_MULTI_BRANCH_HINT_BASE + sum(active) + return max(active) + return record.rel_hint + + +def expected_best_id( + corpus: list[CorpusRecord], + where_document: dict[str, Any] | str, + where: dict[str, Any] | None = None, +) -> str: + matched = [ + (rec.doc_id, effective_rel_hint(rec, where_document)) + for rec in corpus + if corpus_matches_fts(rec, where_document, where) + ] + matched.sort(key=lambda item: (-item[1], item[0])) + return matched[0][0] + + +def insert_corpus_in_batches(namespace: Any, corpus: list[CorpusRecord], batch_size: int = BATCH_SIZE) -> None: + for start in range(0, len(corpus), batch_size): + chunk = corpus[start : start + batch_size] + namespace.add( + ids=[r.doc_id for r in chunk], + embeddings=[r.embedding for r in chunk], + documents=[r.document for r in chunk], + metadatas=[r.metadata for r in chunk], + ) + + +def assert_score_order_non_increasing(distances: list[float]) -> None: + """Hybrid search exposes BM25-like scores in ``distances`` (higher = more relevant).""" + if not distances or len(set(distances)) <= 1: + return + for i in range(len(distances) - 1): + assert distances[i] >= distances[i + 1], ( + f"scores must be non-increasing: index {i}={distances[i]!r} < " + f"index {i + 1}={distances[i + 1]!r}, full={distances!r}" + ) + + +def assert_hybrid_fulltext_result( + corpus: list[CorpusRecord], + result: dict[str, Any], + where_document: dict[str, Any] | str, + n_results: int, + *, + where: dict[str, Any] | None = None, + check_ranking: bool = True, +) -> None: + assert result is not None + assert "ids" in result and result["ids"] + ids = result["ids"][0] + distances = result.get("distances", [[]])[0] if result.get("distances") else [] + documents = result.get("documents", [[]])[0] if result.get("documents") else [] + metadatas = result.get("metadatas", [[]])[0] if result.get("metadatas") else [] + + assert len(ids) <= n_results + assert len(ids) > 0, "expected at least one full-text hit" + if distances: + assert len(distances) == len(ids) + assert_score_order_non_increasing(distances) + + id_to_doc = {rec.doc_id: rec.document for rec in corpus} + id_to_meta = {rec.doc_id: rec.metadata for rec in corpus} + + meta_rows = metadatas if metadatas else [None] * len(ids) + for doc_id, doc_text, meta in zip(ids, documents, meta_rows): + assert doc_id in id_to_doc, f"unknown id {doc_id!r} in search results" + body = doc_text if doc_text is not None else id_to_doc[doc_id] + assert doc_matches_where_document(body, where_document), ( + f"id={doc_id!r} document={body!r} does not satisfy where_document={where_document!r}" + ) + if where is not None: + row_meta = meta if meta is not None else id_to_meta[doc_id] + assert doc_matches_where_metadata(row_meta, where), ( + f"id={doc_id!r} metadata={row_meta!r} does not satisfy where={where!r}" + ) + + if not check_ranking: + return + + corpus_by_id = {rec.doc_id: rec for rec in corpus} + hints = [effective_rel_hint(corpus_by_id[doc_id], where_document) for doc_id in ids] + + # 1) Top-1 must be globally most relevant (by corpus rel_hint tiers). + assert ids[0] == expected_best_id(corpus, where_document, where), ( + f"top-1 must be the highest-relevance document, got {ids[0]!r} " + f"expected {expected_best_id(corpus, where_document, where)!r}" + ) + + # 2) No matching document outside the window may have a higher tier than the last hit. + if len(ids) == n_results: + min_hint_in_page = hints[-1] + for rec in corpus: + if not corpus_matches_fts(rec, where_document, where): + continue + rec_hint = effective_rel_hint(rec, where_document) + if rec_hint > min_hint_in_page: + assert rec.doc_id in ids, ( + f"higher-relevance match {rec.doc_id!r} (hint={rec_hint}) " + f"missing from top-{n_results} (min returned hint={min_hint_in_page})" + ) + + # 3) When the engine returns distinct scores, they must be non-increasing. + if distances: + assert_score_order_non_increasing(distances) + if any(not math.isclose(d, 0.0) for d in distances): + for i in range(len(ids) - 1): + assert distances[i] >= distances[i + 1] or math.isclose(distances[i], distances[i + 1]) + + +# All where_document shapes supported by hybrid_search ``_build_document_query``. +HYBRID_SEARCH_FTS_CASES: list[FtsQueryCase] = [ + FtsQueryCase( + name="contains_token_zpx", + where_document={"$contains": TOKEN_ZPX}, + n_results=15, + ), + FtsQueryCase( + name="string_shorthand_token_zpx", + where_document=TOKEN_ZPX, + n_results=15, + ), + FtsQueryCase( + name="not_contains_token_zpx", + where_document={"$not_contains": TOKEN_ZPX}, + n_results=20, + check_ranking=False, + ), + FtsQueryCase( + name="and_zpx_alp", + where_document={ + "$and": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + n_results=10, + ), + FtsQueryCase( + name="or_zpx_alp", + where_document={ + "$or": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + n_results=20, + ), + FtsQueryCase( + name="contains_token_alp", + where_document={"$contains": TOKEN_ALP}, + n_results=15, + ), + FtsQueryCase( + name="and_multi_contains_phrase", + where_document={ + "$and": [ + {"$contains": "Primary"}, + {"$contains": TOKEN_ZPX}, + ], + }, + n_results=10, + ), + FtsQueryCase( + name="contains_zpx_filter_has_both", + where_document={"$contains": TOKEN_ZPX}, + where={"has_both": True}, + n_results=5, + ), + FtsQueryCase( + name="contains_zpx_filter_zpx_hint_50", + where_document={"$contains": TOKEN_ZPX}, + where={"zpx_hint": 50}, + n_results=5, + ), + FtsQueryCase( + name="contains_zpx_filter_zpx_hint_gte_40", + where_document={"$contains": TOKEN_ZPX}, + where={"zpx_hint": {"$gte": 40}}, + n_results=10, + ), + FtsQueryCase( + name="and_zpx_alp_filter_has_both", + where_document={ + "$and": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + where={"has_both": True}, + n_results=5, + ), + FtsQueryCase( + name="contains_zpx_filter_zpx_hint_and_alp_zero", + where_document={"$contains": TOKEN_ZPX}, + where={ + "$and": [ + {"zpx_hint": {"$gte": 40}}, + {"alp_hint": 0}, + ], + }, + n_results=10, + ), +] + + +def get_fts_case(name: str) -> FtsQueryCase: + for case in HYBRID_SEARCH_FTS_CASES: + if case.name == name: + return case + raise KeyError(f"unknown FTS case: {name!r}") + + +def setup_large_fts_collection(db_client: Any) -> tuple[list[CorpusRecord], Any]: + """Create one namespace-enabled collection with a deterministic large corpus.""" + corpus = build_large_fts_corpus(CORPUS_SIZE) + if len(corpus) <= 1000: + raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") + + name = f"test_ns_hs_ft_{int(time.time() * 1000)}" + collection = db_client.create_collection(name=name, schema=ns_schema(), use_namespace=True) + return corpus, collection + + +def setup_fts_namespace_with_corpus( + collection: Any, + corpus: list[CorpusRecord], + *, + namespace_name: str, +) -> Any: + """Create a namespace under an existing collection and bulk-load the corpus.""" + namespace = collection.create_namespace(namespace_name) + insert_corpus_in_batches(namespace, corpus) + expected = len(corpus) + actual = namespace.count() + assert actual == expected, ( + f"namespace {namespace.name!r} (id={namespace.namespace_id}) " + f"expected {expected} rows after bulk load, got {actual}" + ) + time.sleep(INDEX_SETTLE_SECONDS) + return namespace + + +def setup_large_fts_namespace(db_client: Any, *, namespace_name: str = "ns_hs_ft_large") -> tuple[ + list[CorpusRecord], Any, Any +]: + """Create collection + namespace with corpus (single-test convenience wrapper).""" + corpus, collection = setup_large_fts_collection(db_client) + namespace = setup_fts_namespace_with_corpus( + collection, corpus, namespace_name=namespace_name + ) + return corpus, collection, namespace + + +def teardown_large_fts_collection(db_client: Any, collection: Any) -> None: + db_client.delete_collection(name=collection.name) + + +def teardown_large_fts_namespace(db_client: Any, collection: Any) -> None: + teardown_large_fts_collection(db_client, collection) + + +# Quadrant keys: coll_1/coll_2 x ns_x/ns_y (aligned with DML multi-coll-multi-ns tests). +MULTI_COLL_MULTI_NS_QUADRANT_KEYS: tuple[str, ...] = ("c1_x", "c1_y", "c2_x", "c2_y") + +# One loaded namespace per collection (OceanBase: FTS index applies to the last-loaded ns in a coll). +MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS: tuple[str, ...] = ("c1_x", "c2_y") + + +def _create_multi_coll_multi_ns_layout(db_client: Any, *, name_prefix: str) -> dict[str, Any]: + ts = int(time.time() * 1000) + coll_1 = db_client.create_collection( + name=f"{name_prefix}_{ts}_c1", + schema=ns_schema(), + use_namespace=True, + ) + coll_2 = db_client.create_collection( + name=f"{name_prefix}_{ts}_c2", + schema=ns_schema(), + use_namespace=True, + ) + ctx: dict[str, Any] = {"coll_1": coll_1, "coll_2": coll_2} + for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): + for ns_suffix in ("x", "y"): + ctx[f"{coll_tag}_{ns_suffix}"] = collection.create_namespace(f"ns_{ns_suffix}") + return ctx + + +def setup_multi_coll_multi_ns_fts(db_client: Any) -> dict[str, Any]: + """ + 2 collections x 2 namespaces; load the large FTS corpus into ``c1_x`` and ``c2_y`` only. + + Empty quadrants (``c1_y``, ``c2_x``) are kept for isolation checks. Within one collection, + only the namespace that received the bulk insert is expected to serve hybrid_search FTS + (see ``setup_same_collection_both_ns_fts`` for same-collection multi-ns behavior). + """ + ctx = _create_multi_coll_multi_ns_layout(db_client, name_prefix="test_ns_hs_ft_mcmn") + corpus = build_large_fts_corpus(CORPUS_SIZE) + if len(corpus) <= 1000: + raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") + + for key in MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS: + insert_corpus_in_batches(ctx[key], corpus) + + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + ns = ctx[key] + ctx[key] = (corpus, ns) + + ctx["_corpus"] = corpus + time.sleep(INDEX_SETTLE_SECONDS) + return ctx + + +def setup_same_collection_both_ns_fts(db_client: Any) -> dict[str, Any]: + """ + Single collection with ``ns_x`` and ``ns_y``; corpus inserted into both (x then y). + + hybrid_search FTS is expected on the last-loaded namespace (``ns_y``); ``ns_x`` rows + remain visible via get. + """ + ts = int(time.time() * 1000) + collection = db_client.create_collection( + name=f"test_ns_hs_ft_sc2ns_{ts}", + schema=ns_schema(), + use_namespace=True, + ) + corpus = build_large_fts_corpus(CORPUS_SIZE) + ns_x = collection.create_namespace("ns_x") + ns_y = collection.create_namespace("ns_y") + insert_corpus_in_batches(ns_x, corpus) + insert_corpus_in_batches(ns_y, corpus) + time.sleep(INDEX_SETTLE_SECONDS) + return { + "collection": collection, + "corpus": corpus, + "ns_x": (corpus, ns_x), + "ns_y": (corpus, ns_y), + } + + +def setup_multi_coll_multi_ns_fts_single_loaded( + db_client: Any, + *, + loaded_quadrant: str = "c1_x", +) -> dict[str, Any]: + """Same layout as :func:`setup_multi_coll_multi_ns_fts` but corpus only in ``loaded_quadrant``.""" + if loaded_quadrant not in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + raise ValueError(f"unknown quadrant {loaded_quadrant!r}") + + ctx = _create_multi_coll_multi_ns_layout(db_client, name_prefix="test_ns_hs_ft_mcmn1") + corpus = build_large_fts_corpus(CORPUS_SIZE) + insert_corpus_in_batches(ctx[loaded_quadrant], corpus) + + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + ns = ctx[key] + ctx[key] = (corpus, ns) + + ctx["_corpus"] = corpus + time.sleep(INDEX_SETTLE_SECONDS) + return ctx + + +def teardown_multi_coll_multi_ns_fts(db_client: Any, ctx: dict[str, Any]) -> None: + for coll_key in ("coll_1", "coll_2", "collection"): + collection = ctx.get(coll_key) + if collection is not None: + with contextlib.suppress(Exception): + db_client.delete_collection(name=collection.name) + + +def assert_hybrid_search_no_hits( + namespace: Any, + where_document: dict[str, Any] | str, + *, + n_results: int = 15, +) -> dict[str, Any]: + """Assert hybrid_search returns no rows (namespace isolation).""" + result = namespace.hybrid_search( + query={"where_document": where_document, "n_results": n_results}, + n_results=n_results, + include=["documents"], + ) + ids = result.get("ids", [[]])[0] if result.get("ids") else [] + assert len(ids) == 0, ( + f"expected no full-text hits in namespace {namespace.name!r}, got {len(ids)} ids: {ids[:5]!r}" + ) + return result + + +def run_hybrid_search_fts_case_on_quadrants( + ctx: dict[str, Any], + case: FtsQueryCase | str, + quadrant_keys: tuple[str, ...] = MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, +) -> None: + """Run one FTS case on the selected quadrants.""" + fts_case = case if isinstance(case, FtsQueryCase) else get_fts_case(case) + for key in quadrant_keys: + corpus, namespace = ctx[key] + run_hybrid_search_fts_case(namespace, corpus, fts_case) + + +def run_hybrid_search_fts_case_all_quadrants( + ctx: dict[str, Any], + case: FtsQueryCase | str, +) -> None: + """Run one FTS case on every FTS-loaded quadrant (``c1_x`` + ``c2_y``).""" + run_hybrid_search_fts_case_on_quadrants(ctx, case, MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS) + + +def run_hybrid_search_fts_case( + namespace: Any, + corpus: list[CorpusRecord], + case: FtsQueryCase, +) -> dict[str, Any]: + query: dict[str, Any] = { + "where_document": case.where_document, + "n_results": case.n_results, + } + if case.where is not None: + query["where"] = case.where + result = namespace.hybrid_search( + query=query, + n_results=case.n_results, + include=["documents", "metadatas"], + ) + assert_hybrid_fulltext_result( + corpus, + result, + case.where_document, + case.n_results, + where=case.where, + check_ranking=case.check_ranking, + ) + return result + + +def assert_not_contains_no_token_leak( + corpus: list[CorpusRecord], + result: dict[str, Any], + forbidden_token: str, +) -> None: + result_ids = set(result["ids"][0]) + forbidden_ids = { + rec.doc_id + for rec in corpus + if forbidden_token.lower() in rec.document.lower() + } + leaked = result_ids & forbidden_ids + assert not leaked, f"$not_contains leaked forbidden ids: {sorted(leaked)[:10]}" diff --git a/tests/integration_tests/namespace_hybrid_search_helpers.py b/tests/integration_tests/namespace_hybrid_search_helpers.py new file mode 100644 index 00000000..ba4f50db --- /dev/null +++ b/tests/integration_tests/namespace_hybrid_search_helpers.py @@ -0,0 +1,627 @@ +""" +Helpers for namespace hybrid_search vector / search-index / combined integration tests. + +Reuses the large deterministic corpus from ``namespace_fts_helpers`` so full-text, +metadata (JSON SEARCH INDEX), and vector branches share one ground-truth dataset. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +from namespace_fts_helpers import ( + BATCH_SIZE, + CORPUS_SIZE, + INDEX_SETTLE_SECONDS, + TOKEN_ALP, + TOKEN_ZPX, + CorpusRecord, + assert_hybrid_fulltext_result, + assert_not_contains_no_token_leak, + build_large_fts_corpus, + corpus_matches_fts, + doc_matches_where_document, + doc_matches_where_metadata, + get_fts_case, + insert_corpus_in_batches, + run_hybrid_search_fts_case, + setup_fts_namespace_with_corpus, + setup_large_fts_collection, + setup_multi_coll_multi_ns_fts, + setup_multi_coll_multi_ns_fts_single_loaded, + teardown_large_fts_collection, + teardown_multi_coll_multi_ns_fts, + MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, + MULTI_COLL_MULTI_NS_QUADRANT_KEYS, +) + +# Fixed query vector for KNN tests (same dimension as corpus embeddings). +KNN_QUERY_VECTOR: list[float] = [1.0, 1.0, 0.0] + +RRF_RANK = {"rrf": {"rank_window_size": 60, "rank_constant": 60}} + + +@dataclass(frozen=True) +class SearchIndexQueryCase: + name: str + where: dict[str, Any] + n_results: int + min_hits: int = 1 + exact_match_count: int | None = None + + +@dataclass(frozen=True) +class VectorKnnCase: + name: str + query_vector: list[float] + n_results: int + where: dict[str, Any] | None = None + check_top1: bool = True + + +@dataclass(frozen=True) +class HybridCombinedCase: + """Multi-branch hybrid_search (full-text + metadata + optional KNN).""" + + name: str + n_results: int + where_document: dict[str, Any] | str | None = None + where: dict[str, Any] | None = None + knn: dict[str, Any] | None = None + use_rrf: bool = False + check_fts_ranking: bool = True + # When True, only assert non-empty results in corpus (RRF reorders across branches). + rrf_smoke_only: bool = False + + +def skip_oceanbase_knn(request: Any) -> None: + """OceanBase hybrid_search KNN on namespace logical tables requires HNSW (collections use IVF).""" + import pytest + + if "oceanbase" in request.node.nodeid: + pytest.skip( + "OceanBase hybrid_search KNN requires HNSW; namespace collections use IVF only." + ) + + +def corpus_matches_where(record: CorpusRecord, where: dict[str, Any]) -> bool: + """Metadata / ``#id`` filter against a corpus row (hybrid_search ``query.where`` / ``knn.where``).""" + if "$and" in where: + return all(corpus_matches_where(record, sub) for sub in where["$and"]) + if "$or" in where: + return any(corpus_matches_where(record, sub) for sub in where["$or"]) + if "$not" in where: + return not corpus_matches_where(record, where["$not"]) + + for key, value in where.items(): + if key in ("$and", "$or", "$not"): + continue + if key == "#id": + if isinstance(value, dict): + if "$in" in value and record.doc_id not in value["$in"]: + return False + if "$nin" in value and record.doc_id in value["$nin"]: + return False + if "$eq" in value and record.doc_id != value["$eq"]: + return False + if "$ne" in value and record.doc_id == value["$ne"]: + return False + elif record.doc_id != value: + return False + continue + if not doc_matches_where_metadata(record.metadata, {key: value}): + return False + return True + + +def count_corpus_matches(corpus: list[CorpusRecord], where: dict[str, Any]) -> int: + return sum(1 for rec in corpus if corpus_matches_where(rec, where)) + + +def l2_squared(a: list[float], b: list[float]) -> float: + return sum((x - y) ** 2 for x, y in zip(a, b)) + + +def expected_knn_ids( + corpus: list[CorpusRecord], + query_vector: list[float], + n_results: int, + where: dict[str, Any] | None = None, +) -> list[str]: + scored: list[tuple[str, float]] = [] + for rec in corpus: + if where is not None and not corpus_matches_where(rec, where): + continue + scored.append((rec.doc_id, l2_squared(rec.embedding, query_vector))) + scored.sort(key=lambda item: (item[1], item[0])) + return [doc_id for doc_id, _ in scored[:n_results]] + + +def assert_hybrid_search_index_result( + corpus: list[CorpusRecord], + result: dict[str, Any], + where: dict[str, Any], + n_results: int, + *, + min_hits: int = 1, + exact_match_count: int | None = None, +) -> None: + assert result is not None + assert "ids" in result and result["ids"] + ids = result["ids"][0] + assert len(ids) <= n_results + assert len(ids) >= min_hits, f"expected at least {min_hits} search-index hits" + + id_to_meta = {rec.doc_id: rec.metadata for rec in corpus} + for doc_id in ids: + assert doc_id in id_to_meta, f"unknown id {doc_id!r}" + assert corpus_matches_where( + next(rec for rec in corpus if rec.doc_id == doc_id), where + ), f"id={doc_id!r} does not satisfy where={where!r}" + + total_matches = count_corpus_matches(corpus, where) + if exact_match_count is not None: + assert total_matches == exact_match_count, ( + f"test case assumes {exact_match_count} corpus matches, got {total_matches}" + ) + if exact_match_count is not None and exact_match_count <= n_results: + assert len(ids) == exact_match_count + expected_ids = { + rec.doc_id for rec in corpus if corpus_matches_where(rec, where) + } + assert set(ids) == expected_ids + elif len(ids) == n_results and total_matches > n_results: + # No higher-priority scalar key — every returned row is valid; ensure we did not + # miss matches only when the engine returned a full page and corpus is small. + pass + + +def assert_hybrid_knn_result( + corpus: list[CorpusRecord], + result: dict[str, Any], + query_vector: list[float], + n_results: int, + where: dict[str, Any] | None = None, + *, + check_top1: bool = True, +) -> None: + assert result is not None + assert "ids" in result and result["ids"] + ids = result["ids"][0] + distances = result.get("distances", [[]])[0] if result.get("distances") else [] + assert len(ids) <= n_results + assert len(ids) > 0, "expected at least one KNN hit" + + corpus_by_id = {rec.doc_id: rec for rec in corpus} + for doc_id in ids: + rec = corpus_by_id[doc_id] + if where is not None: + assert corpus_matches_where(rec, where), ( + f"id={doc_id!r} metadata does not satisfy knn.where={where!r}" + ) + + if distances: + assert len(distances) == len(ids) + for dist in distances: + assert dist >= 0 + for i in range(len(distances) - 1): + assert distances[i] <= distances[i + 1] or math.isclose( + distances[i], distances[i + 1] + ), f"KNN distances should be non-decreasing (L2): {distances!r}" + + expected = expected_knn_ids(corpus, query_vector, n_results, where) + if check_top1 and expected: + assert ids[0] == expected[0], ( + f"top-1 KNN must be nearest neighbor, got {ids[0]!r} expected {expected[0]!r}" + ) + + if len(ids) == n_results: + worst = distances[-1] if distances else float("inf") + for rec in corpus: + if where is not None and not corpus_matches_where(rec, where): + continue + d = l2_squared(rec.embedding, query_vector) + if rec.doc_id not in ids and (not distances or d < worst - 1e-9): + raise AssertionError( + f"closer match {rec.doc_id!r} (l2²={d}) missing from top-{n_results}" + ) + + +def run_hybrid_search_index_case( + namespace: Any, + corpus: list[CorpusRecord], + case: SearchIndexQueryCase, +) -> dict[str, Any]: + result = namespace.hybrid_search( + query={"where": case.where, "n_results": case.n_results}, + n_results=case.n_results, + include=["metadatas"], + ) + assert_hybrid_search_index_result( + corpus, + result, + case.where, + case.n_results, + min_hits=case.min_hits, + exact_match_count=case.exact_match_count, + ) + return result + + +def run_hybrid_knn_case( + namespace: Any, + corpus: list[CorpusRecord], + case: VectorKnnCase, +) -> dict[str, Any]: + knn: dict[str, Any] = { + "query_embeddings": case.query_vector, + "n_results": case.n_results, + } + if case.where is not None: + knn["where"] = case.where + result = namespace.hybrid_search( + knn=knn, + n_results=case.n_results, + include=["metadatas", "distances"], + ) + assert_hybrid_knn_result( + corpus, + result, + case.query_vector, + case.n_results, + case.where, + check_top1=case.check_top1, + ) + return result + + +def run_hybrid_combined_case( + namespace: Any, + corpus: list[CorpusRecord], + case: HybridCombinedCase, +) -> dict[str, Any]: + query: dict[str, Any] | None = None + if case.where_document is not None or case.where is not None: + query = {"n_results": case.n_results} + if case.where_document is not None: + query["where_document"] = case.where_document + if case.where is not None: + query["where"] = case.where + + rank = RRF_RANK if case.use_rrf else None + result = namespace.hybrid_search( + query=query, + knn=case.knn, + rank=rank, + n_results=case.n_results, + include=["documents", "metadatas", "distances"], + ) + + corpus_by_id = {rec.doc_id: rec for rec in corpus} + ids = result["ids"][0] + assert len(ids) <= case.n_results + assert len(ids) > 0, f"expected hybrid_search hits for case {case.name!r}" + for doc_id in ids: + assert doc_id in corpus_by_id, f"unknown id {doc_id!r}" + + if case.rrf_smoke_only: + return result + + if case.knn is not None and case.where_document is None and query is None: + vec = case.knn["query_embeddings"] + assert_hybrid_knn_result( + corpus, + result, + vec, + case.n_results, + case.knn.get("where"), + check_top1=True, + ) + return result + + if case.where_document is not None: + assert_hybrid_fulltext_result( + corpus, + result, + case.where_document, + case.n_results, + where=case.where, + check_ranking=case.check_fts_ranking, + ) + elif case.where is not None: + assert_hybrid_search_index_result( + corpus, result, case.where, case.n_results + ) + return result + + +# --- Search-index-only cases (JSON SEARCH INDEX / metadata ``query.where``) --- + +SEARCH_INDEX_CASES: list[SearchIndexQueryCase] = [ + SearchIndexQueryCase( + name="eq_zpx_hint_direct", + where={"zpx_hint": 50}, + n_results=5, + exact_match_count=1, + ), + SearchIndexQueryCase( + name="eq_zpx_hint_operator", + where={"zpx_hint": {"$eq": 50}}, + n_results=5, + exact_match_count=1, + ), + SearchIndexQueryCase( + name="ne_zpx_hint_zero", + where={"zpx_hint": {"$ne": 0}}, + n_results=20, + min_hits=1, + ), + SearchIndexQueryCase( + name="gte_zpx_hint_40", + where={"zpx_hint": {"$gte": 40}}, + n_results=15, + min_hits=2, + ), + SearchIndexQueryCase( + name="lt_zpx_hint_10", + where={"zpx_hint": {"$lt": 10}}, + n_results=20, + min_hits=1, + ), + SearchIndexQueryCase( + name="lte_alp_hint_8", + where={"alp_hint": {"$lte": 8}}, + n_results=15, + min_hits=1, + ), + SearchIndexQueryCase( + name="gt_rel_hint_zero", + where={"rel_hint": {"$gt": 0}}, + n_results=25, + min_hits=5, + ), + SearchIndexQueryCase( + name="in_alp_hint_values", + where={"alp_hint": {"$in": [32, 24, 16, 8]}}, + n_results=10, + min_hits=4, + ), + SearchIndexQueryCase( + name="nin_alp_hint_high", + where={"alp_hint": {"$nin": [32, 24, 16, 8]}}, + n_results=15, + min_hits=1, + ), + SearchIndexQueryCase( + name="and_has_both_zpx_gte", + where={ + "$and": [ + {"has_both": True}, + {"zpx_hint": {"$gte": 5}}, + ], + }, + n_results=5, + exact_match_count=1, + ), + SearchIndexQueryCase( + name="or_zpx_alp_hint_high", + where={ + "$or": [ + {"zpx_hint": {"$gte": 50}}, + {"alp_hint": {"$gte": 32}}, + ], + }, + n_results=10, + min_hits=2, + ), + SearchIndexQueryCase( + name="not_has_both", + where={"$not": {"has_both": True}}, + n_results=20, + min_hits=1, + ), + SearchIndexQueryCase( + name="id_eq_both_tokens", + where={"#id": "both_tokens"}, + n_results=5, + exact_match_count=1, + ), + SearchIndexQueryCase( + name="id_in_zpx_tops", + where={"#id": {"$in": ["zpx_top_5", "zpx_top_4", "zpx_top_3"]}}, + n_results=5, + exact_match_count=3, + ), +] + + +def get_search_index_case(name: str) -> SearchIndexQueryCase: + for case in SEARCH_INDEX_CASES: + if case.name == name: + return case + raise KeyError(f"unknown search-index case: {name!r}") + + +# --- Pure KNN cases --- + +VECTOR_KNN_CASES: list[VectorKnnCase] = [ + VectorKnnCase( + name="knn_global_top5", + query_vector=KNN_QUERY_VECTOR, + n_results=5, + ), + VectorKnnCase( + name="knn_filter_has_both", + query_vector=KNN_QUERY_VECTOR, + n_results=3, + where={"has_both": True}, + check_top1=True, + ), + VectorKnnCase( + name="knn_filter_zpx_hint_gte_40", + query_vector=KNN_QUERY_VECTOR, + n_results=5, + where={"zpx_hint": {"$gte": 40}}, + ), +] + + +def get_vector_knn_case(name: str) -> VectorKnnCase: + for case in VECTOR_KNN_CASES: + if case.name == name: + return case + raise KeyError(f"unknown vector KNN case: {name!r}") + + +# --- Combined hybrid_search cases --- + +HYBRID_COMBINED_CASES: list[HybridCombinedCase] = [ + # Full-text + search index (metadata on query branch) + HybridCombinedCase( + name="fts_contains_zpx_filter_has_both", + where_document={"$contains": TOKEN_ZPX}, + where={"has_both": True}, + n_results=5, + ), + # Vector + search index (KNN branch filter only) + HybridCombinedCase( + name="knn_with_where_has_both", + n_results=5, + knn={ + "query_embeddings": KNN_QUERY_VECTOR, + "where": {"has_both": True}, + "n_results": 10, + }, + ), + # Vector + full-text + RRF + HybridCombinedCase( + name="fts_or_zpx_alp_plus_knn", + where_document={ + "$or": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + knn={"query_embeddings": KNN_QUERY_VECTOR, "n_results": 15}, + n_results=10, + use_rrf=True, + rrf_smoke_only=True, + ), + # Vector + full-text + search index + RRF (smoke: fused ranking across three signals) + HybridCombinedCase( + name="fts_zpx_filter_gte_knn", + where_document={"$contains": TOKEN_ZPX}, + where={"zpx_hint": {"$gte": 40}}, + knn={ + "query_embeddings": KNN_QUERY_VECTOR, + "where": {"zpx_hint": {"$gte": 40}}, + "n_results": 15, + }, + n_results=8, + use_rrf=True, + rrf_smoke_only=True, + ), + # Full-text + search index without vector (aligned with FTS case ``contains_zpx_filter_zpx_hint_and_alp_zero``) + HybridCombinedCase( + name="fts_zpx_filter_zpx_hint_and_alp_zero", + where_document={"$contains": TOKEN_ZPX}, + where={ + "$and": [ + {"zpx_hint": {"$gte": 40}}, + {"alp_hint": 0}, + ], + }, + n_results=10, + ), +] + + +def get_hybrid_combined_case(name: str) -> HybridCombinedCase: + for case in HYBRID_COMBINED_CASES: + if case.name == name: + return case + raise KeyError(f"unknown combined hybrid case: {name!r}") + + +def run_search_index_case_on_quadrants( + ctx: dict[str, Any], + case: SearchIndexQueryCase | str, + quadrant_keys: tuple[str, ...] = MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, +) -> None: + si_case = case if isinstance(case, SearchIndexQueryCase) else get_search_index_case(case) + for key in quadrant_keys: + corpus, namespace = ctx[key] + run_hybrid_search_index_case(namespace, corpus, si_case) + + +def run_knn_case_on_quadrants( + ctx: dict[str, Any], + case: VectorKnnCase | str, + quadrant_keys: tuple[str, ...] = MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, + *, + request: Any = None, +) -> None: + if request is not None: + skip_oceanbase_knn(request) + knn_case = case if isinstance(case, VectorKnnCase) else get_vector_knn_case(case) + for key in quadrant_keys: + corpus, namespace = ctx[key] + run_hybrid_knn_case(namespace, corpus, knn_case) + + +def assert_hybrid_search_index_no_hits( + namespace: Any, + where: dict[str, Any], + *, + n_results: int = 10, +) -> None: + result = namespace.hybrid_search( + query={"where": where, "n_results": n_results}, + n_results=n_results, + include=["metadatas"], + ) + ids = result.get("ids", [[]])[0] if result.get("ids") else [] + assert len(ids) == 0, ( + f"expected no search-index hits in namespace {namespace.name!r}, got {ids[:5]!r}" + ) + + +__all__ = [ + "BATCH_SIZE", + "CORPUS_SIZE", + "HYBRID_COMBINED_CASES", + "INDEX_SETTLE_SECONDS", + "KNN_QUERY_VECTOR", + "MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS", + "MULTI_COLL_MULTI_NS_QUADRANT_KEYS", + "RRF_RANK", + "SEARCH_INDEX_CASES", + "TOKEN_ALP", + "TOKEN_ZPX", + "VECTOR_KNN_CASES", + "HybridCombinedCase", + "SearchIndexQueryCase", + "VectorKnnCase", + "assert_hybrid_search_index_no_hits", + "assert_not_contains_no_token_leak", + "build_large_fts_corpus", + "get_fts_case", + "get_hybrid_combined_case", + "get_search_index_case", + "get_vector_knn_case", + "run_hybrid_combined_case", + "run_hybrid_knn_case", + "run_hybrid_search_index_case", + "run_hybrid_search_fts_case", + "run_knn_case_on_quadrants", + "run_search_index_case_on_quadrants", + "setup_fts_namespace_with_corpus", + "setup_large_fts_collection", + "setup_multi_coll_multi_ns_fts", + "setup_multi_coll_multi_ns_fts_single_loaded", + "skip_oceanbase_knn", + "teardown_large_fts_collection", + "teardown_multi_coll_multi_ns_fts", +] diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py index 64ef17eb..612538d7 100644 --- a/tests/integration_tests/test_namespace_dml.py +++ b/tests/integration_tests/test_namespace_dml.py @@ -11,6 +11,8 @@ from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema +from namespace_dml_helpers import assert_get_absent, assert_get_present, cleanup, create_ns_collection + class TestNamespaceDML: @@ -163,5 +165,78 @@ def test_peek(self, db_client): db_client.delete_collection(name=collection.name) +class TestNamespaceDMLFullCycle: + + def test_full_dml_cycle_verified_by_get(self, db_client): + collection = create_ns_collection(db_client, suffix="_cycle") + ns = collection.create_namespace("cycle_ns") + doc_id = "cycle_doc" + try: + ns.add( + ids=doc_id, + embeddings=[0.1, 0.2, 0.3], + documents="Original document", + metadatas={"stage": "added"}, + ) + assert_get_present( + ns, + doc_id, + embeddings=[0.1, 0.2, 0.3], + documents="Original document", + metadatas={"stage": "added"}, + ) + + ns.update( + ids=doc_id, + embeddings=[0.4, 0.5, 0.6], + documents="Updated document", + metadatas={"stage": "updated"}, + ) + assert_get_present( + ns, + doc_id, + embeddings=[0.4, 0.5, 0.6], + documents="Updated document", + metadatas={"stage": "updated"}, + ) + + ns.upsert( + ids=doc_id, + embeddings=[0.7, 0.8, 0.9], + documents="Upserted document", + metadatas={"stage": "upserted"}, + ) + assert_get_present( + ns, + doc_id, + embeddings=[0.7, 0.8, 0.9], + documents="Upserted document", + metadatas={"stage": "upserted"}, + ) + + ns.delete(ids=doc_id) + assert_get_absent(ns, doc_id) + finally: + cleanup(db_client, collection) + + def test_batch_add_then_get_each(self, db_client): + collection = create_ns_collection(db_client, suffix="_batch") + ns = collection.create_namespace("batch_ns") + try: + ns.add( + ids=["b1", "b2", "b3"], + embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + documents=["Doc 1", "Doc 2", "Doc 3"], + metadatas=[{"n": 1}, {"n": 2}, {"n": 3}], + ) + result = ns.get(ids=["b1", "b2", "b3"], include=["documents", "metadatas"]) + assert len(result["ids"]) == 3 + assert set(result["ids"]) == {"b1", "b2", "b3"} + assert set(result["documents"]) == {"Doc 1", "Doc 2", "Doc 3"} + assert {m["n"] for m in result["metadatas"]} == {1, 2, 3} + finally: + cleanup(db_client, collection) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_dml_multi_coll.py b/tests/integration_tests/test_namespace_dml_multi_coll.py new file mode 100644 index 00000000..72d064a4 --- /dev/null +++ b/tests/integration_tests/test_namespace_dml_multi_coll.py @@ -0,0 +1,67 @@ +""" +Namespace DML integration tests: multiple collections x 1 namespace each. +Verifies collection-level isolation via get. +""" + +import pytest + +from namespace_dml_helpers import ( + assert_get_absent, + assert_get_present, + cleanup, + create_ns_collection, +) + + +class TestNamespaceDMLMultiColl: + + def test_same_namespace_name_across_collections(self, db_client): + coll_a = create_ns_collection(db_client, suffix="_mc_a") + coll_b = create_ns_collection(db_client, suffix="_mc_b") + ns_a = coll_a.create_namespace("tenant") + ns_b = coll_b.create_namespace("tenant") + try: + ns_a.add( + ids="doc1", + embeddings=[1.0, 0.0, 0.0], + documents="Data in coll A", + metadatas={"coll": "A"}, + ) + ns_b.add( + ids="doc1", + embeddings=[0.0, 1.0, 0.0], + documents="Data in coll B", + metadatas={"coll": "B"}, + ) + + assert_get_present( + ns_a, "doc1", documents="Data in coll A", metadatas={"coll": "A"} + ) + assert_get_present( + ns_b, "doc1", documents="Data in coll B", metadatas={"coll": "B"} + ) + finally: + cleanup(db_client, coll_a, coll_b) + + def test_update_delete_isolated_by_collection(self, db_client): + coll_a = create_ns_collection(db_client, suffix="_mc_upd_a") + coll_b = create_ns_collection(db_client, suffix="_mc_upd_b") + ns_a = coll_a.create_namespace("tenant") + ns_b = coll_b.create_namespace("tenant") + try: + ns_a.add(ids="doc1", embeddings=[1.0, 0.0, 0.0], documents="Original A") + ns_b.add(ids="doc1", embeddings=[0.0, 1.0, 0.0], documents="Original B") + + ns_a.update(ids="doc1", documents="Updated A", metadatas={"updated": True}) + assert_get_present(ns_a, "doc1", documents="Updated A", metadatas={"updated": True}) + assert_get_present(ns_b, "doc1", documents="Original B") + + ns_a.delete(ids="doc1") + assert_get_absent(ns_a, "doc1") + assert_get_present(ns_b, "doc1", documents="Original B") + finally: + cleanup(db_client, coll_a, coll_b) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py new file mode 100644 index 00000000..8e0dc09f --- /dev/null +++ b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py @@ -0,0 +1,79 @@ +""" +Namespace DML integration tests: multiple collections x multiple namespaces. +Verifies cross-quadrant isolation via get. +""" + +import pytest + +from namespace_dml_helpers import ( + assert_get_absent, + assert_get_present, + cleanup, + create_ns_collection, +) + + +class TestNamespaceDMLMultiCollMultiNs: + + def _setup_quadrants(self, db_client): + coll_1 = create_ns_collection(db_client, suffix="_mcmn_c1") + coll_2 = create_ns_collection(db_client, suffix="_mcmn_c2") + return { + "coll_1": coll_1, + "coll_2": coll_2, + "c1_x": coll_1.create_namespace("ns_x"), + "c1_y": coll_1.create_namespace("ns_y"), + "c2_x": coll_2.create_namespace("ns_x"), + "c2_y": coll_2.create_namespace("ns_y"), + } + + def test_cross_quadrant_isolation(self, db_client): + q = self._setup_quadrants(db_client) + marker_id = "marker_id" + try: + q["c1_x"].add( + ids=marker_id, + embeddings=[1.0, 0.0, 0.0], + metadatas={"cell": "c1_x"}, + ) + assert_get_present(q["c1_x"], marker_id, metadatas={"cell": "c1_x"}) + assert_get_absent(q["c1_y"], marker_id) + assert_get_absent(q["c2_x"], marker_id) + assert_get_absent(q["c2_y"], marker_id) + + q["c2_y"].add( + ids=marker_id, + embeddings=[0.0, 1.0, 0.0], + metadatas={"cell": "c2_y"}, + ) + assert_get_present(q["c1_x"], marker_id, metadatas={"cell": "c1_x"}) + assert_get_present(q["c2_y"], marker_id, metadatas={"cell": "c2_y"}) + assert_get_absent(q["c1_y"], marker_id) + assert_get_absent(q["c2_x"], marker_id) + finally: + cleanup(db_client, q["coll_1"], q["coll_2"]) + + def test_delete_in_one_quadrant(self, db_client): + q = self._setup_quadrants(db_client) + doc_id = "shared_quad_del" + try: + for ns in (q["c1_x"], q["c1_y"], q["c2_x"], q["c2_y"]): + ns.add(ids=doc_id, embeddings=[1.0, 2.0, 3.0], metadatas={"present": True}) + + for ns in (q["c1_x"], q["c1_y"], q["c2_x"], q["c2_y"]): + assert_get_present(ns, doc_id, metadatas={"present": True}) + + q["c1_x"].delete(ids=doc_id) + + assert_get_absent(q["c1_x"], doc_id) + for ns in (q["c1_y"], q["c2_x"], q["c2_y"]): + result = ns.get(ids=doc_id, include=["metadatas"]) + assert doc_id in result["ids"], f"expected {doc_id!r} in {ns.name}, got {result['ids']}" + idx = result["ids"].index(doc_id) + assert result["metadatas"][idx]["present"] is True + finally: + cleanup(db_client, q["coll_1"], q["coll_2"]) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_dml_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_ns.py new file mode 100644 index 00000000..5a98feee --- /dev/null +++ b/tests/integration_tests/test_namespace_dml_multi_ns.py @@ -0,0 +1,99 @@ +""" +Namespace DML integration tests: 1 collection x multiple namespaces. +Verifies add/update/upsert/delete isolation via get. +""" + +import pytest + +from namespace_dml_helpers import ( + assert_get_absent, + assert_get_present, + cleanup, + create_ns_collection, +) + + +class TestNamespaceDMLMultiNs: + + def test_add_isolation(self, db_client): + collection = create_ns_collection(db_client, suffix="_mns_add") + ns_a = collection.create_namespace("ns_a") + ns_b = collection.create_namespace("ns_b") + try: + ns_a.add(ids="a1", embeddings=[1.0, 0.0, 0.0], metadatas={"owner": "A"}) + ns_b.add(ids="b1", embeddings=[0.0, 1.0, 0.0], metadatas={"owner": "B"}) + + assert_get_present(ns_a, "a1", metadatas={"owner": "A"}) + assert_get_present(ns_b, "b1", metadatas={"owner": "B"}) + assert_get_absent(ns_a, "b1") + assert_get_absent(ns_b, "a1") + finally: + cleanup(db_client, collection) + + def test_same_id_different_namespaces(self, db_client): + collection = create_ns_collection(db_client, suffix="_mns_sameid") + ns_x = collection.create_namespace("ns_x") + ns_y = collection.create_namespace("ns_y") + try: + ns_x.add(ids="shared", embeddings=[1.0, 0.0, 0.0], metadatas={"src": "X"}) + ns_y.add(ids="shared", embeddings=[0.0, 1.0, 0.0], metadatas={"src": "Y"}) + + assert_get_present(ns_x, "shared", metadatas={"src": "X"}) + assert_get_present(ns_y, "shared", metadatas={"src": "Y"}) + finally: + cleanup(db_client, collection) + + def test_update_does_not_affect_other_ns(self, db_client): + collection = create_ns_collection(db_client, suffix="_mns_upd") + ns_a = collection.create_namespace("ns_a") + ns_b = collection.create_namespace("ns_b") + doc_id = "shared_upd" + try: + ns_a.add(ids=doc_id, embeddings=[1.0, 0.0, 0.0], documents="A original") + ns_b.add(ids=doc_id, embeddings=[0.0, 1.0, 0.0], documents="B original") + assert_get_present(ns_a, doc_id, documents="A original") + assert_get_present(ns_b, doc_id, documents="B original") + + ns_a.update(ids=doc_id, documents="A updated", metadatas={"updated": True}) + + assert_get_present(ns_a, doc_id, documents="A updated", metadatas={"updated": True}) + assert_get_present(ns_b, doc_id, documents="B original") + finally: + cleanup(db_client, collection) + + def test_delete_only_target_ns(self, db_client): + collection = create_ns_collection(db_client, suffix="_mns_del") + ns_a = collection.create_namespace("ns_a") + ns_b = collection.create_namespace("ns_b") + doc_id = "shared_del" + try: + ns_a.add(ids=doc_id, embeddings=[1.0, 0.0, 0.0]) + ns_b.add(ids=doc_id, embeddings=[0.0, 1.0, 0.0]) + assert_get_present(ns_a, doc_id) + assert_get_present(ns_b, doc_id) + + ns_a.delete(ids=doc_id) + + assert_get_absent(ns_a, doc_id) + assert_get_present(ns_b, doc_id) + finally: + cleanup(db_client, collection) + + def test_upsert_isolation(self, db_client): + collection = create_ns_collection(db_client, suffix="_mns_ups") + ns_a = collection.create_namespace("ns_a") + ns_b = collection.create_namespace("ns_b") + try: + ns_a.add(ids="a_only", embeddings=[1.0, 0.0, 0.0], metadatas={"ns": "a"}) + ns_b.upsert(ids="b_only", embeddings=[0.0, 1.0, 0.0], metadatas={"ns": "b"}) + + assert_get_present(ns_a, "a_only", metadatas={"ns": "a"}) + assert_get_absent(ns_a, "b_only") + assert_get_present(ns_b, "b_only", metadatas={"ns": "b"}) + assert_get_absent(ns_b, "a_only") + finally: + cleanup(db_client, collection) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_combined.py b/tests/integration_tests/test_namespace_hybrid_search_combined.py new file mode 100644 index 00000000..3391da1c --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_combined.py @@ -0,0 +1,69 @@ +""" +Namespace hybrid_search combined-branch tests. + +Covers: + - full-text + search index (metadata on ``query``) + - vector + search index (``knn.where``) + - vector + full-text (+ RRF smoke) + - vector + full-text + search index (+ RRF smoke) +""" + +from __future__ import annotations + +import time +from typing import Any, ClassVar + +import pytest + +from namespace_hybrid_search_helpers import ( + HYBRID_COMBINED_CASES, + get_hybrid_combined_case, + run_hybrid_combined_case, + setup_fts_namespace_with_corpus, + setup_large_fts_collection, + skip_oceanbase_knn, + teardown_large_fts_collection, +) + + +class TestNamespaceHybridSearchCombined: + _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} + + @pytest.fixture(autouse=True) + def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: + mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" + if mode not in self._shared_by_mode: + corpus, collection = setup_large_fts_collection(db_client) + self._shared_by_mode[mode] = { + "corpus": corpus, + "collection": collection, + "db_client": db_client, + } + entry = self._shared_by_mode[mode] + self._corpus = entry["corpus"] + self._collection = entry["collection"] + + @classmethod + def teardown_class(cls) -> None: + for entry in cls._shared_by_mode.values(): + teardown_large_fts_collection(entry["db_client"], entry["collection"]) + cls._shared_by_mode.clear() + + def _new_namespace(self, case_name: str) -> Any: + return setup_fts_namespace_with_corpus( + self._collection, + self._corpus, + namespace_name=f"ns_comb_{case_name}_{int(time.time() * 1000)}", + ) + + @pytest.mark.parametrize("case_name", [c.name for c in HYBRID_COMBINED_CASES]) + def test_hybrid_search_combined(self, db_client, request, case_name: str): + case = get_hybrid_combined_case(case_name) + if case.knn is not None: + skip_oceanbase_knn(request) + namespace = self._new_namespace(case_name) + run_hybrid_combined_case(namespace, self._corpus, case) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext.py new file mode 100644 index 00000000..0c30cddb --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext.py @@ -0,0 +1,136 @@ +""" +Namespace hybrid_search full-text integration tests (large corpus, >1000 rows). + +One collection is created per test class (per ``db_client`` param); each test method +only creates a fresh namespace and loads the shared corpus. + + pytest tests/integration_tests/test_namespace_hybrid_search_fulltext.py \\ + -k "filter_has_both and oceanbase" -v -s +""" + +from __future__ import annotations + +import time +from typing import Any, ClassVar + +import pytest + +from namespace_fts_helpers import ( + TOKEN_ZPX, + CorpusRecord, + get_fts_case, + assert_not_contains_no_token_leak, + run_hybrid_search_fts_case, + setup_fts_namespace_with_corpus, + setup_large_fts_collection, + teardown_large_fts_collection, +) + + +class TestNamespaceHybridSearchFulltext: + """Large-scale namespace pure full-text hybrid_search tests.""" + + # Keyed by db_client param (embedded / server / oceanbase). + _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} + + @pytest.fixture(autouse=True) + def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: + mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" + if mode not in self._shared_by_mode: + corpus, collection = setup_large_fts_collection(db_client) + self._shared_by_mode[mode] = { + "db_client": db_client, + "corpus": corpus, + "collection": collection, + } + entry = self._shared_by_mode[mode] + self._corpus: list[CorpusRecord] = entry["corpus"] + self._collection = entry["collection"] + self._db_client = db_client + + @classmethod + def teardown_class(cls) -> None: + for entry in cls._shared_by_mode.values(): + teardown_large_fts_collection(entry["db_client"], entry["collection"]) + cls._shared_by_mode.clear() + + def _new_namespace(self, case_name: str) -> Any: + ns_name = f"ns_{case_name}_{int(time.time() * 1000)}" + return setup_fts_namespace_with_corpus( + self._collection, + self._corpus, + namespace_name=ns_name, + ) + + def _run_fts_case(self, case_name: str) -> None: + namespace = self._new_namespace(case_name) + run_hybrid_search_fts_case(namespace, self._corpus, get_fts_case(case_name)) + + def test_hybrid_search_fulltext_contains_token_zpx(self, db_client): + """``where_document``: ``{\"$contains\": TOKEN_ZPX}``""" + self._run_fts_case("contains_token_zpx") + + def test_hybrid_search_fulltext_contains_token_alp(self, db_client): + """``where_document``: ``{\"$contains\": TOKEN_ALP}``""" + self._run_fts_case("contains_token_alp") + + def test_hybrid_search_fulltext_contains_string_shorthand(self, db_client): + """``where_document``: string shorthand (same as ``$contains``)""" + self._run_fts_case("string_shorthand_token_zpx") + + def test_hybrid_search_fulltext_not_contains(self, db_client): + """``where_document``: ``{\"$not_contains\": TOKEN_ZPX}``""" + case = get_fts_case("not_contains_token_zpx") + namespace = self._new_namespace("not_contains_token_zpx") + result = run_hybrid_search_fts_case(namespace, self._corpus, case) + assert_not_contains_no_token_leak(self._corpus, result, TOKEN_ZPX) + + def test_hybrid_search_fulltext_and_zpx_alp(self, db_client): + """``where_document``: ``{\"$and\": [{\"$contains\": ...}, ...]}``""" + self._run_fts_case("and_zpx_alp") + + def test_hybrid_search_fulltext_and_multi_contains(self, db_client): + """``where_document``: ``{\"$and\": [{\"$contains\": \"Primary\"}, ...]}``""" + self._run_fts_case("and_multi_contains_phrase") + + def test_hybrid_search_fulltext_or_zpx_alp(self, db_client): + """``where_document``: ``{\"$or\": [{\"$contains\": ...}, ...]}``""" + self._run_fts_case("or_zpx_alp") + + def test_hybrid_search_fulltext_top1_contains_zpx(self, db_client): + """Top-1 relevance check for ``$contains`` on TOKEN_ZPX.""" + namespace = self._new_namespace("top1_contains_zpx") + top_result = namespace.hybrid_search( + query={"where_document": {"$contains": TOKEN_ZPX}, "n_results": 1}, + n_results=1, + include=["documents"], + ) + assert top_result["ids"][0][0] == "zpx_top_5", ( + f"most relevant TOKEN_ZPX document must rank first, got {top_result['ids'][0]}" + ) + + def test_hybrid_search_fulltext_contains_zpx_filter_has_both(self, db_client): + """``where_document`` + ``where``: ``has_both`` metadata equality.""" + self._run_fts_case("contains_zpx_filter_has_both") + + def test_hybrid_search_fulltext_contains_zpx_filter_zpx_hint_eq(self, db_client): + """``where_document`` + ``where``: ``zpx_hint`` scalar equality.""" + self._run_fts_case("contains_zpx_filter_zpx_hint_50") + + def test_hybrid_search_fulltext_contains_zpx_filter_zpx_hint_gte(self, db_client): + """``where_document`` + ``where``: ``zpx_hint`` ``$gte`` range.""" + self._run_fts_case("contains_zpx_filter_zpx_hint_gte_40") + + def test_hybrid_search_fulltext_and_zpx_alp_filter_has_both(self, db_client): + """``$and`` on ``where_document`` combined with ``has_both`` metadata filter.""" + self._run_fts_case("and_zpx_alp_filter_has_both") + + def test_hybrid_search_fulltext_contains_zpx_filter_zpx_hint_and_alp_zero( + self, db_client + ): + """``where_document`` + ``where`` ``$and``: ``zpx_hint`` range and ``alp_hint`` equality.""" + self._run_fts_case("contains_zpx_filter_zpx_hint_and_alp_zero") + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py new file mode 100644 index 00000000..32755516 --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py @@ -0,0 +1,109 @@ +""" +Namespace hybrid_search full-text tests: 2 collections x 2 namespaces (use_namespace=True). + +Verifies full-text correctness in each quadrant (same assertions as single-namespace tests) +and cross-quadrant isolation (empty namespaces must not leak hits). + +Run one case in isolation, e.g.:: + + pytest tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py \\ + -k "or_zpx_alp and oceanbase" -v -s +""" + +from __future__ import annotations + +import pytest + +from namespace_fts_helpers import ( + MULTI_COLL_MULTI_NS_QUADRANT_KEYS, + TOKEN_ZPX, + assert_hybrid_search_no_hits, + assert_not_contains_no_token_leak, + get_fts_case, + run_hybrid_search_fts_case, + run_hybrid_search_fts_case_all_quadrants, + setup_multi_coll_multi_ns_fts, + setup_multi_coll_multi_ns_fts_single_loaded, + teardown_multi_coll_multi_ns_fts, +) + + +class TestNamespaceHybridSearchFulltextMultiCollMultiNs: + """Large-scale hybrid_search FTS across 2 collections x 2 namespaces.""" + + def _run_fts_case_all_quadrants(self, db_client, case_name: str) -> None: + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_hybrid_search_fts_case_all_quadrants(ctx, case_name) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_cross_quadrant_fts_isolation(self, db_client): + """Only ``c1_x`` has corpus; other quadrants must return no TOKEN_ZPX hits.""" + ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") + try: + case = get_fts_case("contains_token_zpx") + corpus, ns_loaded = ctx["c1_x"] + run_hybrid_search_fts_case(ns_loaded, corpus, case) + + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + if key == "c1_x": + continue + _, namespace = ctx[key] + assert_hybrid_search_no_hits( + namespace, + case.where_document, + n_results=case.n_results, + ) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_hybrid_search_fulltext_contains_token_zpx(self, db_client): + self._run_fts_case_all_quadrants(db_client, "contains_token_zpx") + + def test_hybrid_search_fulltext_contains_token_alp(self, db_client): + self._run_fts_case_all_quadrants(db_client, "contains_token_alp") + + def test_hybrid_search_fulltext_contains_string_shorthand(self, db_client): + self._run_fts_case_all_quadrants(db_client, "string_shorthand_token_zpx") + + def test_hybrid_search_fulltext_not_contains(self, db_client): + case = get_fts_case("not_contains_token_zpx") + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + corpus, namespace = ctx[key] + result = run_hybrid_search_fts_case(namespace, corpus, case) + assert_not_contains_no_token_leak(corpus, result, TOKEN_ZPX) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_hybrid_search_fulltext_and_zpx_alp(self, db_client): + self._run_fts_case_all_quadrants(db_client, "and_zpx_alp") + + def test_hybrid_search_fulltext_and_multi_contains(self, db_client): + self._run_fts_case_all_quadrants(db_client, "and_multi_contains_phrase") + + def test_hybrid_search_fulltext_or_zpx_alp(self, db_client): + self._run_fts_case_all_quadrants(db_client, "or_zpx_alp") + + def test_hybrid_search_fulltext_top1_contains_zpx(self, db_client): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + _, namespace = ctx[key] + top_result = namespace.hybrid_search( + query={"where_document": {"$contains": TOKEN_ZPX}, "n_results": 1}, + n_results=1, + include=["documents"], + ) + assert top_result["ids"][0][0] == "zpx_top_5", ( + f"[{key}] most relevant TOKEN_ZPX document must rank first, " + f"got {top_result['ids'][0]}" + ) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py new file mode 100644 index 00000000..73e9d16a --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py @@ -0,0 +1,149 @@ +""" +Namespace hybrid_search tests: 2 collections x 2 namespaces. + +Extends FTS multi-coll coverage with search-index, KNN, and combined scenarios. +""" + +from __future__ import annotations + +import pytest + +from namespace_fts_helpers import ( + MULTI_COLL_MULTI_NS_QUADRANT_KEYS, + assert_hybrid_search_no_hits, + get_fts_case, + run_hybrid_search_fts_case, + run_hybrid_search_fts_case_all_quadrants, +) +from namespace_hybrid_search_helpers import ( + get_hybrid_combined_case, + get_search_index_case, + get_vector_knn_case, + run_hybrid_combined_case, + run_hybrid_search_index_case, + run_hybrid_knn_case, + run_knn_case_on_quadrants, + run_search_index_case_on_quadrants, + setup_multi_coll_multi_ns_fts, + setup_multi_coll_multi_ns_fts_single_loaded, + skip_oceanbase_knn, + teardown_multi_coll_multi_ns_fts, +) + + +class TestNamespaceHybridSearchMultiCollMultiNs: + def test_cross_quadrant_search_index_isolation(self, db_client): + """Only the loaded quadrant returns ``has_both`` rows.""" + ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") + try: + case = get_search_index_case("and_has_both_zpx_gte") + corpus, ns_loaded = ctx["c1_x"] + run_hybrid_search_index_case(ns_loaded, corpus, case) + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + if key == "c1_x": + continue + _, namespace = ctx[key] + from namespace_hybrid_search_helpers import assert_hybrid_search_index_no_hits + + assert_hybrid_search_index_no_hits(namespace, case.where, n_results=case.n_results) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_cross_quadrant_fts_isolation(self, db_client): + ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") + try: + case = get_fts_case("contains_token_zpx") + corpus, ns_loaded = ctx["c1_x"] + run_hybrid_search_fts_case(ns_loaded, corpus, case) + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + if key == "c1_x": + continue + _, namespace = ctx[key] + assert_hybrid_search_no_hits( + namespace, case.where_document, n_results=case.n_results + ) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + @pytest.mark.parametrize( + "case_name", + [ + "contains_token_zpx", + "not_contains_token_zpx", + "and_zpx_alp", + "or_zpx_alp", + "contains_zpx_filter_has_both", + "contains_zpx_filter_zpx_hint_gte_40", + ], + ) + def test_hybrid_search_fulltext_all_loaded_quadrants(self, db_client, case_name: str): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_hybrid_search_fts_case_all_quadrants(ctx, case_name) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + @pytest.mark.parametrize( + "case_name", + [ + "eq_zpx_hint_direct", + "gte_zpx_hint_40", + "in_alp_hint_values", + "and_has_both_zpx_gte", + "or_zpx_alp_hint_high", + "id_in_zpx_tops", + ], + ) + def test_hybrid_search_search_index_all_loaded_quadrants(self, db_client, case_name: str): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_search_index_case_on_quadrants(ctx, case_name) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + @pytest.mark.parametrize("case_name", ["knn_global_top5", "knn_filter_has_both"]) + def test_hybrid_search_vector_all_loaded_quadrants( + self, db_client, request, case_name: str + ): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_knn_case_on_quadrants(ctx, case_name, request=request) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_hybrid_search_fts_plus_search_index_multi_coll(self, db_client): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + for key in ("c1_x", "c2_y"): + corpus, namespace = ctx[key] + run_hybrid_search_fts_case( + namespace, corpus, get_fts_case("contains_zpx_filter_has_both") + ) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_hybrid_search_combined_rrf_multi_coll(self, db_client, request): + skip_oceanbase_knn(request) + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + case = get_hybrid_combined_case("fts_zpx_filter_gte_knn") + for key in ("c1_x", "c2_y"): + corpus, namespace = ctx[key] + run_hybrid_combined_case(namespace, corpus, case) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_hybrid_search_vector_plus_search_index_multi_coll(self, db_client, request): + skip_oceanbase_knn(request) + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + knn_case = get_vector_knn_case("knn_filter_has_both") + for key in ("c1_x", "c2_y"): + corpus, namespace = ctx[key] + run_hybrid_knn_case(namespace, corpus, knn_case) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_search_index.py b/tests/integration_tests/test_namespace_hybrid_search_search_index.py new file mode 100644 index 00000000..d899ec4e --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_search_index.py @@ -0,0 +1,69 @@ +""" +Namespace hybrid_search pure search-index (metadata / JSON SEARCH INDEX) tests. + +Uses ``query.where`` without ``where_document``. Ground-truth checks mirror +``test_namespace_hybrid_search_fulltext.py`` via ``namespace_hybrid_search_helpers``. +""" + +from __future__ import annotations + +import time +from typing import Any, ClassVar + +import pytest + +from namespace_hybrid_search_helpers import ( + SEARCH_INDEX_CASES, + get_search_index_case, + run_hybrid_search_index_case, + setup_fts_namespace_with_corpus, + setup_large_fts_collection, + teardown_large_fts_collection, +) + + +class TestNamespaceHybridSearchSearchIndex: + """Large-scale namespace pure metadata-filter hybrid_search tests.""" + + _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} + + @pytest.fixture(autouse=True) + def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: + mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" + if mode not in self._shared_by_mode: + corpus, collection = setup_large_fts_collection(db_client) + self._shared_by_mode[mode] = { + "db_client": db_client, + "corpus": corpus, + "collection": collection, + } + entry = self._shared_by_mode[mode] + self._corpus = entry["corpus"] + self._collection = entry["collection"] + + @classmethod + def teardown_class(cls) -> None: + for entry in cls._shared_by_mode.values(): + teardown_large_fts_collection(entry["db_client"], entry["collection"]) + cls._shared_by_mode.clear() + + def _new_namespace(self, case_name: str) -> Any: + ns_name = f"ns_si_{case_name}_{int(time.time() * 1000)}" + return setup_fts_namespace_with_corpus( + self._collection, + self._corpus, + namespace_name=ns_name, + ) + + def _run_case(self, case_name: str) -> None: + namespace = self._new_namespace(case_name) + run_hybrid_search_index_case(namespace, self._corpus, get_search_index_case(case_name)) + + @pytest.mark.parametrize("case_name", [c.name for c in SEARCH_INDEX_CASES]) + def test_hybrid_search_search_index_operators(self, db_client, case_name: str): + """Cover ``$eq/$ne/$lt/$lte/$gt/$gte/$in/$nin/$and/$or/$not`` and ``#id``.""" + self._run_case(case_name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_vector.py b/tests/integration_tests/test_namespace_hybrid_search_vector.py new file mode 100644 index 00000000..fd45921f --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_vector.py @@ -0,0 +1,83 @@ +""" +Namespace hybrid_search pure vector (KNN) integration tests. + +Validates KNN distance ordering and metadata filters on the ``knn`` branch. +""" + +from __future__ import annotations + +import time +from typing import Any, ClassVar + +import pytest + +from namespace_hybrid_search_helpers import ( + KNN_QUERY_VECTOR, + VECTOR_KNN_CASES, + get_vector_knn_case, + run_hybrid_knn_case, + setup_fts_namespace_with_corpus, + setup_large_fts_collection, + skip_oceanbase_knn, + teardown_large_fts_collection, +) + + +class TestNamespaceHybridSearchVector: + """Large-scale namespace pure KNN hybrid_search tests.""" + + _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} + + @pytest.fixture(autouse=True) + def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: + mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" + if mode not in self._shared_by_mode: + corpus, collection = setup_large_fts_collection(db_client) + self._shared_by_mode[mode] = { + "db_client": db_client, + "corpus": corpus, + "collection": collection, + } + entry = self._shared_by_mode[mode] + self._corpus = entry["corpus"] + self._collection = entry["collection"] + + @classmethod + def teardown_class(cls) -> None: + for entry in cls._shared_by_mode.values(): + teardown_large_fts_collection(entry["db_client"], entry["collection"]) + cls._shared_by_mode.clear() + + def _new_namespace(self, case_name: str) -> Any: + ns_name = f"ns_knn_{case_name}_{int(time.time() * 1000)}" + return setup_fts_namespace_with_corpus( + self._collection, + self._corpus, + namespace_name=ns_name, + ) + + def _run_knn_case(self, request: pytest.FixtureRequest, case_name: str) -> None: + skip_oceanbase_knn(request) + namespace = self._new_namespace(case_name) + run_hybrid_knn_case(namespace, self._corpus, get_vector_knn_case(case_name)) + + @pytest.mark.parametrize("case_name", [c.name for c in VECTOR_KNN_CASES]) + def test_hybrid_search_vector_knn_cases(self, db_client, request, case_name: str): + self._run_knn_case(request, case_name) + + def test_hybrid_search_vector_top1_nearest(self, db_client, request): + """Top-1 must be the global nearest neighbor for a fixed query vector.""" + skip_oceanbase_knn(request) + namespace = self._new_namespace("top1_nearest") + from namespace_hybrid_search_helpers import expected_knn_ids + + result = namespace.hybrid_search( + knn={"query_embeddings": KNN_QUERY_VECTOR, "n_results": 1}, + n_results=1, + include=["distances"], + ) + assert result["ids"][0][0] == expected_knn_ids(self._corpus, KNN_QUERY_VECTOR, 1)[0] + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) From 13c0101b1637559c0ba785b85cdb3e9607737c10 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 1 Jun 2026 19:39:22 +0800 Subject: [PATCH 20/84] =?UTF-8?q?=E4=BF=AE=E5=A4=8Dscalar=20term=20query?= =?UTF-8?q?=20in=20must/should=20clause=20not=20supported?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 70 +++++++++++++++---- .../integration_tests/test_namespace_query.py | 39 +++++++++++ 2 files changed, 96 insertions(+), 13 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 5628f7bb..820142de 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1253,8 +1253,14 @@ def _create_namespace_physical_tables( self._execute(data_sql) - ivf_index_sql = f"CREATE VECTOR INDEX idx_vec ON `{data_table}` (embedding) {vector_index_sql}" - self._execute(ivf_index_sql) + # Vector indexes are rejected in shared-storage mode on this kernel branch + # (both inline `VECTOR INDEX` in CREATE TABLE and standalone `CREATE VECTOR + # INDEX` fail with OB-1235 "vector index in shared storage mode is not + # supported"). Keep the embedding column but skip the vector index in SS; + # full-text and metadata SEARCH INDEX hybrid_search still work. + if not is_shared_storage: + ivf_index_sql = f"CREATE VECTOR INDEX idx_vec ON `{data_table}` (embedding) {vector_index_sql}" + self._execute(ivf_index_sql) if is_shared_storage: hot_table = NamespaceCollectionNames.hot_table_name(collection_id) @@ -3952,14 +3958,33 @@ def _build_query_expression(self, query: dict[str, Any]) -> dict[str, Any] | Non if not where_document and where: filter_conditions = self._build_metadata_filter_for_search_parm(where) if filter_conditions: - # If only one filter condition, check its type - if len(filter_conditions) == 1: - filter_cond = filter_conditions[0] - # Directly return supported single condition types - if any(key in filter_cond for key in ("range", "term", "terms", "bool")): - return filter_cond - # Multiple filter conditions, wrap in bool filter - return {"bool": {"filter": filter_conditions}} + # Wrap scalar conditions in a (non-scoring) `filter` clause: a + # top-level bool is scoring by default and the kernel rejects scalar + # term/range queries inside must/should of a scoring bool + # (`scalar ... query in must/should clause not supported`). Negation + # conditions ($not/$ne/$nin) come back as pure `{"bool": {"must_not"}}` + # nodes; a bool with only must_not is rejected with `bool query ... + # should have at least one positive clause`, so hoist their must_not + # clauses onto the outer bool (which gains a positive `filter` once the + # namespace filter is injected) instead of nesting them standalone. + positive: list[dict[str, Any]] = [] + negative: list[dict[str, Any]] = [] + for cond in filter_conditions: + if ( + isinstance(cond, dict) + and set(cond.keys()) == {"bool"} + and isinstance(cond["bool"], dict) + and set(cond["bool"].keys()) == {"must_not"} + ): + negative.extend(cond["bool"]["must_not"]) + else: + positive.append(cond) + bool_q: dict[str, Any] = {} + if positive: + bool_q["filter"] = positive + if negative: + bool_q["must_not"] = negative + return {"bool": bool_q} # Case 2: Full-text search (with or without metadata filtering) if where_document: @@ -4136,7 +4161,11 @@ def _build_metadata_filter_conditions( # noqa: C901 sub_filters = self._build_metadata_filter_conditions(sub_condition) must_conditions.extend(sub_filters) if must_conditions: - result.append({"bool": {"must": must_conditions}}) + # Scalar conditions must be ANDed via a (non-scoring) `filter` + # clause, not `must`: the kernel rejects scalar term/range queries + # inside must/should with `scalar ... query in must/should clause + # not supported`. + result.append({"bool": {"filter": must_conditions}}) return result if "$or" in condition: @@ -4145,7 +4174,13 @@ def _build_metadata_filter_conditions( # noqa: C901 sub_filters = self._build_metadata_filter_conditions(sub_condition) should_conditions.extend(sub_filters) if should_conditions: - result.append({"bool": {"should": should_conditions}}) + # `minimum_should_match: 1` makes this an explicit OR. In a + # non-scoring (filter) context the kernel does not reliably apply the + # implicit "at least one should" default, which otherwise yields an + # intermittent `1210 Invalid argument`. + result.append( + {"bool": {"should": should_conditions, "minimum_should_match": 1}} + ) return result if "$not" in condition: @@ -5276,6 +5311,13 @@ def _inject_filter_into_knn(node): node["filter"] = list(ns_filter) return node + # Scoring (full-text) leaf queries may live in a `must` clause; scalar + # leaf queries (term/terms/range/json/array) MUST go into `filter`. The + # kernel rejects a scalar query inside must/should of a scoring bool with + # `OB_NOT_SUPPORTED: scalar term query in must/should clause`, and a + # top-level bool query is treated as scoring by default. + _scoring_leaf_keys = {"query_string", "match", "multi_match", "match_phrase"} + def _inject_filter_into_query(node): if not isinstance(node, dict): return node @@ -5290,7 +5332,9 @@ def _inject_filter_into_query(node): else: bool_node["filter"] = list(ns_filter) return node - return {"bool": {"must": [node], "filter": list(ns_filter)}} + if _scoring_leaf_keys & node.keys(): + return {"bool": {"must": [node], "filter": list(ns_filter)}} + return {"bool": {"filter": [node, *ns_filter]}} if "query" in search_parm: q = search_parm["query"] diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index 575f69fb..ce942c44 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -351,6 +351,45 @@ def test_hybrid_search_with_data_content_filter_knn_branch(self, db_client, requ finally: db_client.delete_collection(name=collection.name) + def test_hybrid_search_scalar_only_metadata_filter(self, db_client): + """ + Scalar-only hybrid_search (only ``where``, no ``where_document`` and no ``knn``). + + Regression: the namespace filter injection used to wrap the scalar leaf query in a + ``must`` clause, which the kernel rejects with + ``OB_NOT_SUPPORTED: scalar term query in must/should clause`` because a top-level + bool query is scoring by default. Scalar leaves must go into ``filter``. + """ + collection = self._setup(db_client) + ns = collection.create_namespace("qns_hs_scalar") + try: + self._insert_data(ns) + time.sleep(1) + + # term filter + result_term = ns.hybrid_search( + query={"where": {"category": "AI"}}, + n_results=10, + include=["metadatas"], + ) + ids_term = sorted(result_term["ids"][0]) if result_term["ids"] else [] + assert ids_term == ["q1", "q4"], ids_term + if result_term.get("metadatas") and result_term["metadatas"][0]: + for meta in result_term["metadatas"][0]: + if meta: + assert meta["category"] == "AI" + + # range filter + result_range = ns.hybrid_search( + query={"where": {"score": {"$gte": 90}}}, + n_results=10, + include=["metadatas"], + ) + ids_range = sorted(result_range["ids"][0]) if result_range["ids"] else [] + assert ids_range == ["q1", "q3", "q4"], ids_range + finally: + db_client.delete_collection(name=collection.name) + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) From accc8ef3a2fbea42a6eac9473beb5bdee761912a Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 2 Jun 2026 20:21:18 +0800 Subject: [PATCH 21/84] =?UTF-8?q?sdk=E6=94=B9=E5=90=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 65 ++++----- src/pyseekdb/client/configuration.py | 8 +- src/pyseekdb/client/meta_info.py | 8 +- .../test_namespace_drop_validation.py | 126 +++++++++++++++++- tests/unit_tests/test_build_document_query.py | 68 ++++++++++ tests/unit_tests/test_namespace.py | 38 +++--- 6 files changed, 248 insertions(+), 65 deletions(-) create mode 100644 tests/unit_tests/test_build_document_query.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 820142de..77e64ea0 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -231,8 +231,8 @@ def _get_ivf_vector_index_sql(ivf_config: "IVFConfiguration") -> str: property_parts.append(f"{k}='{v}'") else: property_parts.append(f"{k}={v}") - if ivf_config.use_spfresh is not None: - property_parts.append(f"use_spfresh={str(ivf_config.use_spfresh).lower()}") + if ivf_config.fresh_mode: + property_parts.append(f"fresh_mode={ivf_config.fresh_mode}") property_str = ", ".join(property_parts) properties_str = f", {property_str}" if property_str else "" return f"WITH (DISTANCE={ivf_config.distance}, TYPE={ivf_config.type.upper()}, LIB={ivf_config.lib.upper()}{properties_str})" @@ -941,8 +941,8 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " if ivf_config is not None: dimension = ivf_config.dimension distance = ivf_config.distance - if ivf_config.use_spfresh is not None and not ivf_config.use_spfresh: - raise ValueError("use_namespace=True requires use_spfresh=True") + if ivf_config.fresh_mode is not None and ivf_config.fresh_mode != "spfresh": + raise ValueError("use_namespace=True requires fresh_mode='spfresh'") else: if dense_embedding_function is not None: dimension = self._get_embedding_function_dimension(dense_embedding_function) @@ -950,7 +950,7 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " dimension = DEFAULT_VECTOR_DIMENSION distance = DEFAULT_DISTANCE_METRIC from .configuration import IVFConfiguration - ivf_config = IVFConfiguration(dimension=dimension, distance=distance, use_spfresh=True) + ivf_config = IVFConfiguration(dimension=dimension, distance=distance, fresh_mode="spfresh") if dimension < 1 or dimension > 4096: raise ValueError(f"Dimension must be between 1 and 4096, got {dimension}") @@ -960,7 +960,7 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " "version": 2, "use_namespace": True, "dense_index_type": "ivf", - "use_spfresh": True, + "fresh_mode": "spfresh", "storage_mode": "ss" if is_ss else "sn", "dimension": dimension, "distance": distance, @@ -1095,7 +1095,7 @@ def _create_collection_meta_v1(self, collection_name: str) -> str: # ==================== Namespace Catalog Methods ==================== def _ensure_namespace_catalogs(self) -> None: - ns_namespaces_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` ( + ns_namespaces_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_namespaces_table()}` ( namespace_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, collection_id CHAR(32) NOT NULL, namespace_name VARCHAR(256) NOT NULL, @@ -1106,7 +1106,7 @@ def _ensure_namespace_catalogs(self) -> None: UNIQUE KEY uk_sdk_ns_coll_name (collection_id, namespace_name), KEY idx_sdk_ns_by_collection (collection_id) ) COMMENT='Namespace catalog';""" - ns_ltables_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_ns_ltables_table()}` ( + ns_ltables_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_ltables_table()}` ( ltable_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, collection_id CHAR(32) NOT NULL, namespace_id BIGINT UNSIGNED NOT NULL, @@ -1203,12 +1203,12 @@ def _delete_ns_collection_meta(self, collection_name: str) -> None: ) with contextlib.suppress(Exception): self._execute( - f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"DELETE FROM `{NamespaceCollectionNames.sdk_ltables_table()}` " f"WHERE collection_id = '{collection_id_escaped}'" ) with contextlib.suppress(Exception): self._execute( - f"DELETE FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"DELETE FROM `{NamespaceCollectionNames.sdk_namespaces_table()}` " f"WHERE collection_id = '{collection_id_escaped}'" ) with contextlib.suppress(Exception): @@ -1238,6 +1238,19 @@ def _create_namespace_physical_tables( try: self._execute(f"CREATE TABLEGROUP `{tg_name}` SHARDING='ADAPTIVE'") + index_clauses = [ + f"FULLTEXT INDEX idx_fts(document) {fulltext_clause}", + "SEARCH INDEX idx_json(data_content)", + ] + # Vector indexes are rejected in shared-storage mode on this kernel branch + # (both inline `VECTOR INDEX` in CREATE TABLE and standalone `CREATE VECTOR + # INDEX` fail with OB-1235 "vector index in shared storage mode is not + # supported"). Keep the embedding column but skip the vector index in SS; + # full-text and metadata SEARCH INDEX hybrid_search still work. + if not is_shared_storage: + index_clauses.append(f"VECTOR INDEX idx_vec(embedding) {vector_index_sql}") + index_sql = ",\n ".join(index_clauses) + data_sql = f"""CREATE TABLE `{data_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, ltable_id BIGINT UNSIGNED NOT NULL, @@ -1246,22 +1259,12 @@ def _create_namespace_physical_tables( data_content JSON NOT NULL, created_by VARCHAR(64) DEFAULT '', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - FULLTEXT INDEX idx_fts(document) {fulltext_clause}, - SEARCH INDEX idx_json(data_content) + {index_sql} ) TABLEGROUP=`{tg_name}` COMMENT='逻辑表主数据' DEFAULT CHARSET=utf8mb4 ORGANIZATION HEAP IS_LOGIC_TABLE = TRUE {partition_clause}""" self._execute(data_sql) - # Vector indexes are rejected in shared-storage mode on this kernel branch - # (both inline `VECTOR INDEX` in CREATE TABLE and standalone `CREATE VECTOR - # INDEX` fail with OB-1235 "vector index in shared storage mode is not - # supported"). Keep the embedding column but skip the vector index in SS; - # full-text and metadata SEARCH INDEX hybrid_search still work. - if not is_shared_storage: - ivf_index_sql = f"CREATE VECTOR INDEX idx_vec ON `{data_table}` (embedding) {vector_index_sql}" - self._execute(ivf_index_sql) - if is_shared_storage: hot_table = NamespaceCollectionNames.hot_table_name(collection_id) self._execute(f"""CREATE TABLE `{hot_table}` ( @@ -1312,7 +1315,7 @@ def _resolve_namespace_ltable_id( self, collection_id: str, namespace_id: int ) -> int: """Resolve the default ltable_id for (collection_id, namespace_id) from - sdk_ns_ltables. Cached per (collection_id, namespace_id) on the client + sdk_ltables. Cached per (collection_id, namespace_id) on the client instance to avoid the extra round-trip on every DML/DQL call. The cache is also seeded by `_create_ns_namespace_meta` and @@ -1328,14 +1331,14 @@ def _resolve_namespace_ltable_id( return cached coll_id_escaped = escape_string(str(collection_id)) rows = self._execute( - f"SELECT ltable_id FROM `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"SELECT ltable_id FROM `{NamespaceCollectionNames.sdk_ltables_table()}` " f"WHERE collection_id = '{coll_id_escaped}' AND namespace_id = {int(namespace_id)} " f"AND ltable_name = 'default' LIMIT 1" ) if not rows: raise ValueError( f"No default ltable found for collection_id={collection_id}, " - f"namespace_id={namespace_id} in sdk_ns_ltables" + f"namespace_id={namespace_id} in sdk_ltables" ) row = rows[0] lt_id = int(row[0] if isinstance(row, (list, tuple)) else row["ltable_id"]) @@ -1369,20 +1372,20 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) self._execute( - f"INSERT INTO `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"INSERT INTO `{NamespaceCollectionNames.sdk_namespaces_table()}` " f"(collection_id, namespace_name) VALUES ('{collection_id_escaped}', '{namespace_name_escaped}')" ) rows = self._execute( - f"SELECT namespace_id FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"SELECT namespace_id FROM `{NamespaceCollectionNames.sdk_namespaces_table()}` " f"WHERE collection_id = '{collection_id_escaped}' AND namespace_name = '{namespace_name_escaped}'" ) ns_id = int(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["namespace_id"]) self._execute( - f"INSERT INTO `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"INSERT INTO `{NamespaceCollectionNames.sdk_ltables_table()}` " f"(collection_id, namespace_id, ltable_name) VALUES ('{collection_id_escaped}', {ns_id}, 'default')" ) lt_rows = self._execute( - f"SELECT ltable_id FROM `{NamespaceCollectionNames.sdk_ns_ltables_table()}` " + f"SELECT ltable_id FROM `{NamespaceCollectionNames.sdk_ltables_table()}` " f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {ns_id} AND ltable_name = 'default'" ) lt_id = int(lt_rows[0][0] if isinstance(lt_rows[0], (list, tuple)) else lt_rows[0]["ltable_id"]) @@ -1400,8 +1403,8 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) - ns_table = NamespaceCollectionNames.sdk_ns_namespaces_table() - lt_table = NamespaceCollectionNames.sdk_ns_ltables_table() + ns_table = NamespaceCollectionNames.sdk_namespaces_table() + lt_table = NamespaceCollectionNames.sdk_ltables_table() rows = self._execute( f"SELECT n.namespace_id AS namespace_id, n.namespace_name AS namespace_name, " f"l.ltable_id AS ltable_id " @@ -1456,7 +1459,7 @@ def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> def _list_ns_namespaces(self, collection_id: str) -> list[dict]: collection_id_escaped = escape_string(collection_id) rows = self._execute( - f"SELECT namespace_id, namespace_name FROM `{NamespaceCollectionNames.sdk_ns_namespaces_table()}` " + f"SELECT namespace_id, namespace_name FROM `{NamespaceCollectionNames.sdk_namespaces_table()}` " f"WHERE collection_id = '{collection_id_escaped}' ORDER BY namespace_id" ) results = [] diff --git a/src/pyseekdb/client/configuration.py b/src/pyseekdb/client/configuration.py index 243dcee7..930b905d 100644 --- a/src/pyseekdb/client/configuration.py +++ b/src/pyseekdb/client/configuration.py @@ -336,7 +336,7 @@ class IVFConfiguration: dimension: Vector dimension (number of elements in each vector) distance: Distance metric for similarity calculation (e.g., 'l2', 'cosine', 'inner_product') type: IVF index subtype ('ivf_flat', 'ivf_sq8', 'ivf_pq') - use_spfresh: Whether to enable SPFresh optimization for online index updates + fresh_mode: SPFresh mode for online index updates (e.g., 'spfresh'). Defaults to None (disabled). properties: Optional dictionary of additional IVF index properties """ @@ -344,7 +344,7 @@ class IVFConfiguration: distance: str | DistanceMetric = DistanceMetric.COSINE.value type: str | IVFIndexType = IVFIndexType.IVF_FLAT.value lib: str | IVFIndexLib = IVFIndexLib.OB.value - use_spfresh: bool | None = None + fresh_mode: str | None = None properties: dict[str, PrimitiveValue] | None = None def __post_init__(self): @@ -367,8 +367,8 @@ def __post_init__(self): if self.lib not in valid_libs: raise ValueError(f"lib must be one of {valid_libs}, got {self.lib}") - if self.use_spfresh is not None and not isinstance(self.use_spfresh, bool): - raise TypeError(f"use_spfresh must be a bool, got {type(self.use_spfresh).__name__}") + if self.fresh_mode is not None and not isinstance(self.fresh_mode, str): + raise TypeError(f"fresh_mode must be a str, got {type(self.fresh_mode).__name__}") _ensure_primitive_properties(self.properties) diff --git a/src/pyseekdb/client/meta_info.py b/src/pyseekdb/client/meta_info.py index 53be60e5..4523754b 100644 --- a/src/pyseekdb/client/meta_info.py +++ b/src/pyseekdb/client/meta_info.py @@ -65,12 +65,12 @@ class NamespaceCollectionNames: _TG_SUFFIX = "_tg" @staticmethod - def sdk_ns_namespaces_table() -> str: - return "sdk_ns_namespaces" + def sdk_namespaces_table() -> str: + return "sdk_namespaces" @staticmethod - def sdk_ns_ltables_table() -> str: - return "sdk_ns_ltables" + def sdk_ltables_table() -> str: + return "sdk_ltables" @staticmethod def sdk_namespaces_stats_table() -> str: diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index 74045610..0e709061 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -54,7 +54,7 @@ def _is_ss_mode(client) -> bool: def _fetch_namespace_name(client, collection_id: str, namespace_id: int): rows = _execute( client, - f"SELECT namespace_name FROM sdk_ns_namespaces " + f"SELECT namespace_name FROM sdk_namespaces " f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id}", ) if not rows: @@ -65,7 +65,7 @@ def _fetch_namespace_name(client, collection_id: str, namespace_id: int): def _count_ltables(client, collection_id: str, namespace_id: int) -> int: rows = _execute( client, - f"SELECT COUNT(*) AS c FROM sdk_ns_ltables " + f"SELECT COUNT(*) AS c FROM sdk_ltables " f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id}", ) return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) @@ -92,6 +92,55 @@ def _count_hot_table_rows(client, collection_id: str, namespace_id: int) -> int: return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) +def _count_logic_data_rows(client, collection_id: str, namespace_id: int, ltable_id: int | None = None) -> int: + tbl = NamespaceCollectionNames.data_table_name(collection_id) + if ltable_id is not None: + where = f"namespace_id = {namespace_id} AND ltable_id = {ltable_id}" + else: + where = f"namespace_id = {namespace_id}" + rows = _execute( + client, + f"SELECT COUNT(*) AS c FROM `{tbl}` WHERE {where}", + ) + return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) + + +def _count_kv_data_rows(client, collection_id: str, namespace_id: int) -> int: + tbl = NamespaceCollectionNames.kv_data_table_name(collection_id) + try: + rows = _execute( + client, + f"SELECT COUNT(*) AS c FROM `{tbl}` WHERE namespace_id = {namespace_id}", + ) + except Exception: + return -1 + return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) + + +def _seed_logic_data(client, collection_id: str, namespace_id: int, ltable_id: int, count: int = 3): + tbl = NamespaceCollectionNames.data_table_name(collection_id) + values = [] + for i in range(count): + values.append( + f"({namespace_id}, {ltable_id}, 'doc_{i}', " + f"X'0000803f0000000000000000', " + f"'{{\"id\": \"id_{i}\", \"metadata\": {{\"k\": \"v\"}}}}')" + ) + _execute(client, f"INSERT INTO `{tbl}` (namespace_id, ltable_id, document, embedding, data_content) VALUES {', '.join(values)}") + + +def _fetch_ltable_id(client, collection_id: str, namespace_id: int) -> int | None: + rows = _execute( + client, + f"SELECT ltable_id FROM sdk_ltables " + f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id} " + f"ORDER BY ltable_id LIMIT 1", + ) + if not rows: + return None + return int(rows[0]["ltable_id"] if "ltable_id" in rows[0] else rows[0]["LTABLE_ID"]) + + def _seed_hot_table(client, collection_id: str, namespace_id: int): """In SS mode insert a synthetic hot_table row so the DELETE has something to remove and we can verify the cleanup.""" @@ -167,10 +216,17 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): assert _seed_hot_table(client, coll_id, ns_id) is True assert _count_hot_table_rows(client, coll_id, ns_id) == 1 + # Seed some logic_data rows so we can verify they survive the sync + # drop phase (async background task cleans them up later). + lt_id = _fetch_ltable_id(client, coll_id, ns_id) + assert lt_id is not None, "should have a default ltable" + _seed_logic_data(client, coll_id, ns_id, lt_id, count=3) + assert _count_logic_data_rows(client, coll_id, ns_id, lt_id) == 3 + # Action _drop_namespace_via_pl(client, coll_id, ns_id) - # Post-conditions + # Post-conditions: synchronous phase new_name = _fetch_namespace_name(client, coll_id, ns_id) assert new_name is not None, "namespace row must still exist (renamed, not deleted)" assert new_name.startswith(RECYCLEBIN_PREFIX), ( @@ -180,7 +236,7 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): f"original name should be embedded in recyclebin name, got {new_name!r}" ) assert _count_ltables(client, coll_id, ns_id) == 0, ( - "sdk_ns_ltables rows must be deleted" + "sdk_ltables rows must be deleted" ) assert _count_logic_schema_rows(client, coll_id, ns_id) == 0, ( "logic_schema_table rows must be deleted" @@ -189,6 +245,62 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): assert _count_hot_table_rows(client, coll_id, ns_id) == 0, ( "hot_table row must be deleted in SS mode" ) + + # Logic data rows still exist — async background task handles physical cleanup. + assert _count_logic_data_rows(client, coll_id, ns_id, lt_id) == 3, ( + "logic_data rows must survive the sync drop phase" + ) + finally: + client.delete_collection(name=collection.name) + + # ------------------------------------------------------------- # + # 1b. Async cleanup: wait 32s, verify physical data is gone # + # ------------------------------------------------------------- # + def test_drop_namespace_async_cleanup(self, oceanbase_client): + """After DROP_NAMESPACE renames the namespace, a background task (30s + interval) picks up __recyclebin_ entries and removes physical data. + Wait 32s and verify the renamed namespace and its data are gone.""" + client = oceanbase_client + collection = _make_collection(client) + coll_id = collection.id + try: + ns = collection.create_namespace("ns_async") + ns_id = int(ns.namespace_id) + lt_id = _fetch_ltable_id(client, coll_id, ns_id) + assert lt_id is not None + + # Seed data + _seed_logic_data(client, coll_id, ns_id, lt_id, count=3) + assert _count_logic_data_rows(client, coll_id, ns_id, lt_id) == 3 + + # Drop — sync phase renames namespace, removes catalog entries + _drop_namespace_via_pl(client, coll_id, ns_id) + + recycled_name = _fetch_namespace_name(client, coll_id, ns_id) + assert recycled_name is not None + assert recycled_name.startswith(RECYCLEBIN_PREFIX) + assert "ns_async" in recycled_name + + # Data still present immediately after drop + assert _count_logic_data_rows(client, coll_id, ns_id, lt_id) == 3 + + # Wait for async background task (scans every 30s, 32s should suffice) + time.sleep(32) + + # After background task: renamed namespace entry should be gone + final_name = _fetch_namespace_name(client, coll_id, ns_id) + assert final_name is None, ( + f"async task should have removed the __recyclebin_ namespace row, " + f"got {final_name!r}" + ) + + # Logic data should also be gone + data_rows = _count_logic_data_rows(client, coll_id, ns_id, lt_id) + assert data_rows == 0, ( + f"async task should have cleaned up logic_data rows, " + f"found {data_rows}" + ) + finally: client.delete_collection(name=collection.name) @@ -224,7 +336,7 @@ def test_multiple_ltables_all_deleted(self, oceanbase_client): # Insert two extra ltable rows directly so we can verify bulk delete. _execute( client, - f"INSERT INTO sdk_ns_ltables (collection_id, namespace_id, ltable_name) " + f"INSERT INTO sdk_ltables (collection_id, namespace_id, ltable_name) " f"VALUES ('{collection.id}', {ns_id}, 'extra_lt_a')," f" ('{collection.id}', {ns_id}, 'extra_lt_b')", ) @@ -242,7 +354,7 @@ def test_multiple_ltables_all_deleted(self, oceanbase_client): def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): """If we pre-DROP _logic_schema_table the in-trans DELETE on it fails with OB_TABLE_NOT_EXIST, which must roll back the namespace - rename and the sdk_ns_ltables delete that ran earlier in the same + rename and the sdk_ltables delete that ran earlier in the same transaction.""" client = oceanbase_client collection = _make_collection(client) @@ -269,7 +381,7 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): f"namespace_name must be rolled back to original, got {after_name!r}" ) assert _count_ltables(client, coll_id, ns_id) == ltable_cnt_before, ( - "sdk_ns_ltables delete must be rolled back" + "sdk_ltables delete must be rolled back" ) finally: # Recreate the schema table (empty) so cleanup_namespace_physical_tables diff --git a/tests/unit_tests/test_build_document_query.py b/tests/unit_tests/test_build_document_query.py new file mode 100644 index 00000000..58c93fe2 --- /dev/null +++ b/tests/unit_tests/test_build_document_query.py @@ -0,0 +1,68 @@ +"""Unit tests for BaseClient._build_document_query DSL generation.""" + +from unittest.mock import MagicMock, patch + +import pytest + + +@pytest.fixture() +def client(): + """Create a BaseClient subclass instance with mocked abstract methods.""" + from pyseekdb.client.client_base import BaseClient + + with patch.multiple(BaseClient, __abstractmethods__=set()): + instance = BaseClient.__new__(BaseClient) + return instance + + +class TestBuildDocumentQuery: + + def test_and_contains_has_default_operator_and(self, client): + where_document = {"$and": [{"$contains": "TOKENZPX"}, {"$contains": "TOKENALP"}]} + result = client._build_document_query(where_document) + assert result is not None + qs = result["query_string"] + assert qs["default_operator"] == "and" + assert "TOKENZPX" in qs["query"] + assert "TOKENALP" in qs["query"] + assert qs["fields"] == ["document"] + + def test_or_contains_has_default_operator_or(self, client): + where_document = {"$or": [{"$contains": "TOKENZPX"}, {"$contains": "TOKENALP"}]} + result = client._build_document_query(where_document) + assert result is not None + qs = result["query_string"] + assert qs["default_operator"] == "or" + assert "TOKENZPX" in qs["query"] + assert "TOKENALP" in qs["query"] + assert " OR " not in qs["query"] + + def test_single_contains_no_default_operator(self, client): + where_document = {"$contains": "hello"} + result = client._build_document_query(where_document) + assert result is not None + qs = result["query_string"] + assert "default_operator" not in qs + assert qs["query"] == "hello" + + def test_and_contains_with_boost(self, client): + where_document = {"$and": [{"$contains": "foo"}, {"$contains": "bar"}]} + result = client._build_document_query(where_document, boost=2.0) + assert result is not None + qs = result["query_string"] + assert qs["default_operator"] == "and" + assert qs["boost"] == 2.0 + + def test_and_contains_three_terms(self, client): + where_document = { + "$and": [ + {"$contains": "alpha"}, + {"$contains": "beta"}, + {"$contains": "gamma"}, + ] + } + result = client._build_document_query(where_document) + assert result is not None + qs = result["query_string"] + assert qs["default_operator"] == "and" + assert qs["query"] == "alpha beta gamma" diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index b379b836..df50c1a8 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -32,14 +32,14 @@ def test_valid_defaults(self): config = IVFConfiguration() assert config.dimension == 384 assert config.distance == "cosine" - assert config.use_spfresh is None + assert config.fresh_mode is None assert config.properties is None def test_valid_custom(self): - config = IVFConfiguration(dimension=128, distance="l2", use_spfresh=True) + config = IVFConfiguration(dimension=128, distance="l2", fresh_mode="spfresh") assert config.dimension == 128 assert config.distance == "l2" - assert config.use_spfresh is True + assert config.fresh_mode == "spfresh" def test_valid_inner_product(self): config = IVFConfiguration(dimension=1024, distance="inner_product") @@ -77,9 +77,9 @@ def test_invalid_distance(self): with pytest.raises(ValueError, match="distance must be one of"): IVFConfiguration(distance="invalid") - def test_invalid_use_spfresh_type(self): - with pytest.raises(TypeError, match="use_spfresh must be a bool"): - IVFConfiguration(use_spfresh="yes") + def test_invalid_fresh_mode_type(self): + with pytest.raises(TypeError, match="fresh_mode must be a str"): + IVFConfiguration(fresh_mode=True) def test_properties_valid(self): config = IVFConfiguration(properties={"nlist": 128, "nprobe": 16}) @@ -397,7 +397,7 @@ def _execute(self, sql, params=None): self.executed_sqls.append(sql) return None - # Bypass the sdk_ns_ltables lookup in unit tests: SQL-generation tests don't + # Bypass the sdk_ltables lookup in unit tests: SQL-generation tests don't # have a real database, so return a fixed ltable_id matching test assertions. def _resolve_namespace_ltable_id(self, collection_id, namespace_id): return 1 @@ -906,8 +906,8 @@ def test_tablegroup_name(self): def test_namespace_catalog_table_names(self): from pyseekdb.client.meta_info import NamespaceCollectionNames - assert NamespaceCollectionNames.sdk_ns_namespaces_table() == "sdk_ns_namespaces" - assert NamespaceCollectionNames.sdk_ns_ltables_table() == "sdk_ns_ltables" + assert NamespaceCollectionNames.sdk_namespaces_table() == "sdk_namespaces" + assert NamespaceCollectionNames.sdk_ltables_table() == "sdk_ltables" assert NamespaceCollectionNames.sdk_namespaces_stats_table() == "sdk_namespaces_stats" def test_is_ns_data_table_true(self): @@ -929,11 +929,11 @@ def test_ensure_namespace_catalogs_creates_all_catalog_tables(self): c._ensure_namespace_catalogs() sql = "\n".join(c.executed_sqls) - assert "CREATE TABLE IF NOT EXISTS `sdk_ns_namespaces`" in sql + assert "CREATE TABLE IF NOT EXISTS `sdk_namespaces`" in sql assert "PRIMARY KEY (namespace_id)" in sql assert "UNIQUE KEY uk_sdk_ns_coll_name (collection_id, namespace_name)" in sql - assert "CREATE TABLE IF NOT EXISTS `sdk_ns_ltables`" in sql + assert "CREATE TABLE IF NOT EXISTS `sdk_ltables`" in sql assert "PRIMARY KEY (ltable_id)" in sql assert "UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name)" in sql @@ -950,8 +950,8 @@ def test_ensure_namespace_catalogs_creates_catalog_tables_in_order(self): c._ensure_namespace_catalogs() assert len(c.executed_sqls) == 3 - assert "`sdk_ns_namespaces`" in c.executed_sqls[0] - assert "`sdk_ns_ltables`" in c.executed_sqls[1] + assert "`sdk_namespaces`" in c.executed_sqls[0] + assert "`sdk_ltables`" in c.executed_sqls[1] assert "`sdk_namespaces_stats`" in c.executed_sqls[2] def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): @@ -986,17 +986,17 @@ def test_ob_type_validation(self): c.detect_db_type_and_version = MagicMock(return_value=("mysql", "8.0")) from pyseekdb.client.schema import Schema from pyseekdb.client.configuration import VectorIndexConfig - schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, use_spfresh=True), embedding_function=None)) + schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, fresh_mode="spfresh"), embedding_function=None)) with pytest.raises(ValueError, match="only supported on OceanBase"): c._create_namespace_collection("test", schema) - def test_use_spfresh_false_raises(self): + def test_fresh_mode_false_raises(self): c = FakeClient() c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", "4.3")) from pyseekdb.client.schema import Schema from pyseekdb.client.configuration import VectorIndexConfig - schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, use_spfresh=False), embedding_function=None)) - with pytest.raises(ValueError, match="requires use_spfresh=True"): + schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, fresh_mode="none"), embedding_function=None)) + with pytest.raises(ValueError, match="requires fresh_mode='spfresh'"): c._create_namespace_collection("test", schema) @@ -1038,8 +1038,8 @@ def test_get_ns_namespace_meta_sets_ltable_id(self): meta = c._get_ns_namespace_meta("abc123", "ns1") assert meta == {"namespace_id": "10", "namespace_name": "ns1", "ltable_id": "7"} calls = [str(call) for call in c._execute.call_args_list] - # JOIN against sdk_ns_ltables so we can resolve the default ltable in one round trip. - assert any("sdk_ns_ltables" in s for s in calls) + # JOIN against sdk_ltables so we can resolve the default ltable in one round trip. + assert any("sdk_ltables" in s for s in calls) assert any("SET @namespace_id" in s for s in calls) assert any("SET @ltable_id" in s for s in calls) From 6ea284c142f7466caa37f77b6d014219b02f541c Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 3 Jun 2026 10:29:00 +0800 Subject: [PATCH 22/84] fix(hybrid_search): hoist $not_contains must_not out of scoring bool must OceanBase rejects a bool with only must_not when it is nested under must. Combine $not_contains with metadata filters via outer filter + must_not instead. --- src/pyseekdb/client/client_base.py | 30 +++++-- .../unit_tests/test_build_query_expression.py | 79 +++++++++++++++++++ 2 files changed, 104 insertions(+), 5 deletions(-) create mode 100644 tests/unit_tests/test_build_query_expression.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 77e64ea0..4e729e5c 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -3944,6 +3944,19 @@ def _build_search_parm( # noqa: C901 return search_parm + @staticmethod + def _pure_must_not_clauses(expr: dict[str, Any] | None) -> list[dict[str, Any]] | None: + """Return inner must_not clauses when *expr* is ``{"bool": {"must_not": [...]}}`` only.""" + if not isinstance(expr, dict) or set(expr.keys()) != {"bool"}: + return None + bool_node = expr.get("bool") + if not isinstance(bool_node, dict) or set(bool_node.keys()) != {"must_not"}: + return None + clauses = bool_node.get("must_not") + if not isinstance(clauses, list): + return None + return clauses + def _build_query_expression(self, query: dict[str, Any]) -> dict[str, Any] | None: """ Build query expression from query dict @@ -3996,13 +4009,19 @@ def _build_query_expression(self, query: dict[str, Any]) -> dict[str, Any] | Non if doc_query: # Build filter from where condition filter_conditions = self._build_metadata_filter_for_search_parm(where) + negated = self._pure_must_not_clauses(doc_query) if filter_conditions: - # Full-text search with metadata filtering + # $not_contains becomes a pure must_not bool; nesting that under + # must triggers `bool query ... should have at least one positive clause`. + if negated is not None: + return {"bool": {"filter": filter_conditions, "must_not": negated}} return {"bool": {"must": [doc_query], "filter": filter_conditions}} - else: - # Full-text search only - return doc_query + if negated is not None: + # No metadata filter yet; match_all satisfies the positive-clause rule + # until namespace filters are injected (collection stays standalone). + return {"bool": {"filter": [{"match_all": {}}], "must_not": negated}} + return doc_query return None @@ -4054,13 +4073,14 @@ def _apply_boost(target: Any) -> None: # Handle $not_contains - wrap query_string in must_not bool if "$not_contains" in where_document: + escaped_query = escape_string(where_document["$not_contains"]) return _with_boost({ "bool": { "must_not": [ { "query_string": { "fields": ["document"], - "query": where_document["$not_contains"], + "query": escaped_query, } } ] diff --git a/tests/unit_tests/test_build_query_expression.py b/tests/unit_tests/test_build_query_expression.py new file mode 100644 index 00000000..5df3bfcb --- /dev/null +++ b/tests/unit_tests/test_build_query_expression.py @@ -0,0 +1,79 @@ +"""Unit tests for BaseClient._build_query_expression DSL generation.""" + +from unittest.mock import patch + +import pytest + + +@pytest.fixture() +def client(): + from pyseekdb.client.client_base import BaseClient + + with patch.multiple(BaseClient, __abstractmethods__=set()): + instance = BaseClient.__new__(BaseClient) + return instance + + +class TestBuildQueryExpressionNotContains: + def test_not_contains_with_metadata_filter_hoists_must_not(self, client): + """$not_contains + where must not nest a must_not-only bool inside must.""" + where = {"seq": {"$gte": 43}} + with patch.object( + client, + "_build_metadata_filter_for_search_parm", + return_value=[{"range": {"data_content.metadata.seq": {"gte": 43}}}], + ): + expr = client._build_query_expression({ + "where_document": {"$not_contains": "TOKENZPX"}, + "where": where, + }) + + assert expr == { + "bool": { + "filter": [{"range": {"data_content.metadata.seq": {"gte": 43}}}], + "must_not": [ + { + "query_string": { + "fields": ["document"], + "query": "TOKENZPX", + } + } + ], + } + } + assert "must" not in expr["bool"] + + def test_not_contains_only_uses_match_all_filter(self, client): + with patch.object(client, "_build_metadata_filter_for_search_parm", return_value=[]): + expr = client._build_query_expression({ + "where_document": {"$not_contains": "TOKENZPX"}, + }) + + assert expr == { + "bool": { + "filter": [{"match_all": {}}], + "must_not": [ + { + "query_string": { + "fields": ["document"], + "query": "TOKENZPX", + } + } + ], + } + } + + def test_contains_with_metadata_filter_still_uses_must(self, client): + with patch.object( + client, + "_build_metadata_filter_for_search_parm", + return_value=[{"term": {"data_content.metadata.seq": {"value": 1}}}], + ): + expr = client._build_query_expression({ + "where_document": {"$contains": "TOKENZPX"}, + "where": {"seq": 1}, + }) + + assert "must" in expr["bool"] + assert "query_string" in expr["bool"]["must"][0] + assert "filter" in expr["bool"] From ee117300213d518d2770fa5c9d7f109722ba1ea5 Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Wed, 3 Jun 2026 10:47:45 +0800 Subject: [PATCH 23/84] open spfresh and add case --- .../namespace_dml_helpers.py | 291 +++++- .../namespace_fts_helpers.py | 39 +- .../namespace_hybrid_search_helpers.py | 948 +++++++++++++++++- tests/integration_tests/test_namespace_dml.py | 288 ++++-- .../test_namespace_drop_validation.py | 2 +- .../test_namespace_get_delete_filters.py | 202 ++++ .../test_namespace_hybrid_search_combined.py | 5 +- ...rid_search_fulltext_multi_coll_multi_ns.py | 108 +- ...space_hybrid_search_multi_coll_multi_ns.py | 13 +- ...t_namespace_hybrid_search_triple_branch.py | 97 ++ ...earch_triple_branch_multi_coll_multi_ns.py | 230 +++++ .../test_namespace_hybrid_search_vector.py | 11 +- .../test_namespace_lifecycle.py | 4 +- .../test_namespace_prewarm.py | 2 +- .../integration_tests/test_namespace_query.py | 13 +- tests/unit_tests/test_namespace.py | 12 + 16 files changed, 2127 insertions(+), 138 deletions(-) create mode 100644 tests/integration_tests/test_namespace_get_delete_filters.py create mode 100644 tests/integration_tests/test_namespace_hybrid_search_triple_branch.py create mode 100644 tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index d4293859..7131e4b3 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -5,17 +5,39 @@ from __future__ import annotations import time +from dataclasses import dataclass from typing import Any from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema +LARGE_DML_CORPUS_SIZE = 1000 +DML_BATCH_SIZE = 100 +# Wait for FULLTEXT index after bulk load (same as namespace_fts_helpers). +INDEX_SETTLE_SECONDS = 3 + +# Default sizes for P0 get/delete filter tests (single namespace). +FILTER_KEEP_COUNT = 300 +FILTER_PURGE_COUNT = 200 + +# Multi-namespace orthogonality: 500 keep + 500 purge per namespace. +FILTER_MULTI_NS_KEEP = 500 +FILTER_MULTI_NS_PURGE = 500 + + +@dataclass(frozen=True) +class DmlRecord: + doc_id: str + embedding: list[float] + document: str + metadata: dict[str, Any] + def ns_schema() -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2"), + ivf=IVFConfiguration(dimension=3, distance="l2", use_spfresh=True), embedding_function=None, ), ) @@ -34,6 +56,273 @@ def cleanup(client: Any, *collections: Any) -> None: pass +def build_large_dml_corpus( + size: int = LARGE_DML_CORPUS_SIZE, + *, + id_prefix: str = "dml", + ns_tag: str | None = None, +) -> list[DmlRecord]: + """Deterministic corpus for count/peek scale tests.""" + records: list[DmlRecord] = [] + for i in range(size): + doc_id = f"{id_prefix}_{i:04d}" + embedding = [ + (i % 97) / 97.0, + ((i * 3) % 89) / 89.0, + ((i * 7) % 83) / 83.0, + ] + document = f"namespace dml document {id_prefix} index {i}" + metadata: dict[str, Any] = {"seq": i, "id_prefix": id_prefix} + if ns_tag is not None: + metadata["ns_tag"] = ns_tag + records.append(DmlRecord(doc_id, embedding, document, metadata)) + return records + + +def corpus_id_set(corpus: list[DmlRecord]) -> set[str]: + return {r.doc_id for r in corpus} + + +def _filter_embedding(index: int) -> list[float]: + return [ + (index % 97) / 97.0, + ((index * 3) % 89) / 89.0, + ((index * 7) % 83) / 83.0, + ] + + +def build_filter_corpus( + *, + keep_count: int = FILTER_KEEP_COUNT, + purge_count: int = FILTER_PURGE_COUNT, + ai_in_keep: int = 120, + python_in_keep: int = 50, + id_prefix: str = "flt", + ns_tag: str | None = None, +) -> tuple[list[DmlRecord], dict[str, int]]: + """ + Corpus for conditional get/delete tests. + + - ``tag=keep``: stable rows; first ``ai_in_keep`` have ``category=AI``; + first ``python_in_keep`` documents contain ``python``. + - ``tag=purge``: rows whose documents contain ``obsolete`` (for ``where_document`` delete). + """ + records: list[DmlRecord] = [] + idx = 0 + + for i in range(keep_count): + category = "AI" if i < ai_in_keep else ("Programming" if i % 2 == 0 else "Database") + document = ( + f"python guide {id_prefix} keep index {i}" + if i < python_in_keep + else f"stable document {id_prefix} keep index {i}" + ) + meta: dict[str, Any] = {"tag": "keep", "category": category, "seq": i} + if ns_tag is not None: + meta["ns_tag"] = ns_tag + records.append( + DmlRecord(f"{id_prefix}_keep_{i:04d}", _filter_embedding(idx), document, meta) + ) + idx += 1 + + for i in range(purge_count): + meta = {"tag": "purge", "category": "Database", "seq": i} + if ns_tag is not None: + meta["ns_tag"] = ns_tag + records.append( + DmlRecord( + f"{id_prefix}_purge_{i:04d}", + _filter_embedding(idx), + f"obsolete payload {id_prefix} purge index {i}", + meta, + ) + ) + idx += 1 + + ground_truth = { + "total": keep_count + purge_count, + "keep": keep_count, + "purge": purge_count, + "category_ai": ai_in_keep, + "doc_python": python_in_keep, + "doc_obsolete": purge_count, + } + return records, ground_truth + + +def assert_get_where_count( + ns: Any, + where: dict[str, Any], + expected: int, + *, + limit: int | None = None, + check_meta: dict[str, Any] | None = None, +) -> dict[str, Any]: + """``get(where=...)`` must return exactly ``expected`` rows (up to ``limit``).""" + kwargs: dict[str, Any] = {"where": where, "include": ["metadatas", "documents"]} + if limit is not None: + kwargs["limit"] = limit + result = ns.get(**kwargs) + ids = result.get("ids") or [] + assert len(ids) == expected, f"expected {expected} ids, got {len(ids)}: {ids[:5]!r}..." + if check_meta is not None and result.get("metadatas"): + for meta in result["metadatas"]: + assert meta is not None + for key, value in check_meta.items(): + assert meta.get(key) == value, meta + return result + + +def settle_fts_index() -> None: + """Allow namespace logic-table FULLTEXT index to catch up after bulk insert.""" + time.sleep(INDEX_SETTLE_SECONDS) + + +def assert_get_where_document_count( + ns: Any, + where_document: dict[str, Any], + expected: int, + *, + limit: int | None = None, + substring: str | None = None, + context: str = "", +) -> dict[str, Any]: + kwargs: dict[str, Any] = {"where_document": where_document, "include": ["documents", "metadatas"]} + if limit is not None: + kwargs["limit"] = limit + result = ns.get(**kwargs) + ids = result.get("ids") or [] + prefix = f"{context}: " if context else "" + if len(ids) != expected: + docs = result.get("documents") or [] + raise AssertionError( + f"{prefix}get(where_document={where_document!r}, limit={limit!r}) " + f"expected {expected} row(s), got {len(ids)}. " + f"Uses namespace DML document filter (LIKE for $contains, not hybrid_search FTS). " + f"sample_ids={ids[:5]!r}, sample_documents={docs[:3]!r}" + ) + if substring and result.get("documents"): + for i, doc in enumerate(result["documents"]): + if doc is None: + raise AssertionError( + f"{prefix}document at index {i} is None for id={ids[i]!r}" + ) + if substring not in doc.lower(): + raise AssertionError( + f"{prefix}id={ids[i]!r} document={doc!r} " + f"does not contain substring {substring!r}" + ) + return result + + +def precheck_where_document_hits( + ns: Any, + where_document: dict[str, Any], + expected_hits: int, + *, + context: str = "before delete", +) -> dict[str, Any]: + """Ensure FTS-backed get sees rows before delete(where_document=...).""" + probe_limit = min(expected_hits, 50) if expected_hits > 0 else 10 + return assert_get_where_document_count( + ns, + where_document, + expected_hits, + limit=expected_hits if expected_hits <= 1000 else probe_limit, + context=context, + ) + + +def assert_count_after_document_delete( + ns: Any, + *, + where_document: dict[str, Any], + expected_count: int, + total_before: int, + expected_removed: int, + pre_delete_hit_count: int, +) -> None: + actual = ns.count() + if actual == expected_count: + return + raise AssertionError( + f"after delete(where_document={where_document!r}): " + f"count expected {expected_count}, got {actual} " + f"(total_before={total_before}, expected_removed={expected_removed} rows). " + f"Pre-delete get(where_document=...) returned {pre_delete_hit_count} hit(s). " + "If pre-delete hits matched but count did not drop, suspect logic-table " + "DELETE with MATCH(document) AGAINST on OceanBase." + ) + + +def load_filter_corpus( + namespace: Any, + corpus: list[DmlRecord], + *, + assert_count: bool = True, + settle_fts: bool = True, +) -> None: + """Load filter test corpus; optionally wait for FULLTEXT index (where_document paths).""" + load_dml_corpus(namespace, corpus, assert_count=assert_count) + if settle_fts: + settle_fts_index() + + +def insert_dml_corpus_in_batches( + namespace: Any, + corpus: list[DmlRecord], + batch_size: int = DML_BATCH_SIZE, +) -> None: + for start in range(0, len(corpus), batch_size): + chunk = corpus[start : start + batch_size] + namespace.add( + ids=[r.doc_id for r in chunk], + embeddings=[r.embedding for r in chunk], + documents=[r.document for r in chunk], + metadatas=[r.metadata for r in chunk], + ) + + +def load_dml_corpus( + namespace: Any, + corpus: list[DmlRecord], + *, + assert_count: bool = True, +) -> None: + insert_dml_corpus_in_batches(namespace, corpus) + if assert_count: + expected = len(corpus) + actual = namespace.count() + assert actual == expected, ( + f"namespace {namespace.name!r} expected count {expected}, got {actual}" + ) + + +def assert_peek_result( + result: dict[str, Any], + *, + expected_len: int, + allowed_ids: set[str] | None = None, + ns_tag: str | None = None, +) -> None: + assert "ids" in result + assert "documents" in result + assert "metadatas" in result + assert "embeddings" in result + assert len(result["ids"]) == expected_len + assert len(result["documents"]) == expected_len + assert len(result["metadatas"]) == expected_len + assert len(result["embeddings"]) == expected_len + if allowed_ids is not None: + assert set(result["ids"]).issubset(allowed_ids), ( + f"peek returned foreign ids: {set(result['ids']) - allowed_ids}" + ) + if ns_tag is not None and result["metadatas"]: + for meta in result["metadatas"]: + assert meta is not None + assert meta.get("ns_tag") == ns_tag, meta + + def _default_include( documents: str | None, metadatas: dict | None, diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index 044c640c..4728ce31 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -52,7 +52,19 @@ class FtsQueryCase: def ns_schema() -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2"), + ivf=IVFConfiguration(dimension=3, distance="l2", use_spfresh=True), + embedding_function=None, + ), + ) + + +def flat_hybrid_search_schema() -> Schema: + """Schema for flat collection baseline (``use_namespace=False``, HNSW + FTS).""" + from pyseekdb import HNSWConfiguration + + return Schema( + vector_index=VectorIndexConfig( + hnsw=HNSWConfiguration(dimension=3, distance="l2"), embedding_function=None, ), ) @@ -274,10 +286,11 @@ def expected_best_id( return matched[0][0] -def insert_corpus_in_batches(namespace: Any, corpus: list[CorpusRecord], batch_size: int = BATCH_SIZE) -> None: +def insert_corpus_in_batches(target: Any, corpus: list[CorpusRecord], batch_size: int = BATCH_SIZE) -> None: + """Bulk ``add`` on a namespace or flat collection target.""" for start in range(0, len(corpus), batch_size): chunk = corpus[start : start + batch_size] - namespace.add( + target.add( ids=[r.doc_id for r in chunk], embeddings=[r.embedding for r in chunk], documents=[r.document for r in chunk], @@ -285,6 +298,21 @@ def insert_corpus_in_batches(namespace: Any, corpus: list[CorpusRecord], batch_s ) +def insert_corpus_into_collection(collection: Any, corpus: list[CorpusRecord]) -> None: + """``collection.add(...)`` on a flat collection (``use_namespace=False``).""" + if getattr(collection, "use_namespace", None) is not False: + raise ValueError( + f"insert_corpus_into_collection requires use_namespace=False, " + f"got {getattr(collection, 'use_namespace', None)!r} on {collection!r}" + ) + insert_corpus_in_batches(collection, corpus) + + +def insert_corpus_into_namespace(namespace: Any, corpus: list[CorpusRecord]) -> None: + """``namespace.add(...)`` under a namespace-enabled collection.""" + insert_corpus_in_batches(namespace, corpus) + + def assert_score_order_non_increasing(distances: list[float]) -> None: """Hybrid search exposes BM25-like scores in ``distances`` (higher = more relevant).""" if not distances or len(set(distances)) <= 1: @@ -489,7 +517,7 @@ def setup_fts_namespace_with_corpus( ) -> Any: """Create a namespace under an existing collection and bulk-load the corpus.""" namespace = collection.create_namespace(namespace_name) - insert_corpus_in_batches(namespace, corpus) + insert_corpus_into_namespace(namespace, corpus) expected = len(corpus) actual = namespace.count() assert actual == expected, ( @@ -512,7 +540,8 @@ def setup_large_fts_namespace(db_client: Any, *, namespace_name: str = "ns_hs_ft def teardown_large_fts_collection(db_client: Any, collection: Any) -> None: - db_client.delete_collection(name=collection.name) + with contextlib.suppress(Exception): + db_client.delete_collection(name=collection.name) def teardown_large_fts_namespace(db_client: Any, collection: Any) -> None: diff --git a/tests/integration_tests/namespace_hybrid_search_helpers.py b/tests/integration_tests/namespace_hybrid_search_helpers.py index ba4f50db..aa4332d4 100644 --- a/tests/integration_tests/namespace_hybrid_search_helpers.py +++ b/tests/integration_tests/namespace_hybrid_search_helpers.py @@ -8,8 +8,9 @@ from __future__ import annotations import math +import time from dataclasses import dataclass -from typing import Any +from typing import Any, Literal from namespace_fts_helpers import ( BATCH_SIZE, @@ -26,7 +27,11 @@ doc_matches_where_metadata, get_fts_case, insert_corpus_in_batches, + insert_corpus_into_collection, + insert_corpus_into_namespace, run_hybrid_search_fts_case, + ns_schema, + flat_hybrid_search_schema, setup_fts_namespace_with_corpus, setup_large_fts_collection, setup_multi_coll_multi_ns_fts, @@ -42,6 +47,20 @@ RRF_RANK = {"rrf": {"rank_window_size": 60, "rank_constant": 60}} +# Broad filters so a third branch stays active without narrowing the primary signal under test. +WHERE_DOCUMENT_UNIVERSAL: dict[str, str] = {"$contains": "document"} +WHERE_BROAD_METADATA: dict[str, Any] = {"rel_hint": {"$gte": 0}} +# Filler rows start after the deterministic FTS tier block (~43 rows). +WHERE_FILLER_SEQ: dict[str, Any] = {"seq": {"$gte": 43}} +# ``both_tokens`` row — use numeric fields for baseline-compatible filters. +# Boolean ``has_both`` is not filterable on flat collection (HNSW / GET_SQL) path. +WHERE_BOTH_TOKENS_NUMERIC: dict[str, Any] = { + "$and": [{"zpx_hint": 5}, {"alp_hint": 5}], +} +WHERE_NOT_BOTH_TOKENS_NUMERIC: dict[str, Any] = { + "$not": {"$and": [{"zpx_hint": 5}, {"alp_hint": 5}]}, +} + @dataclass(frozen=True) class SearchIndexQueryCase: @@ -76,16 +95,6 @@ class HybridCombinedCase: rrf_smoke_only: bool = False -def skip_oceanbase_knn(request: Any) -> None: - """OceanBase hybrid_search KNN on namespace logical tables requires HNSW (collections use IVF).""" - import pytest - - if "oceanbase" in request.node.nodeid: - pytest.skip( - "OceanBase hybrid_search KNN requires HNSW; namespace collections use IVF only." - ) - - def corpus_matches_where(record: CorpusRecord, where: dict[str, Any]) -> bool: """Metadata / ``#id`` filter against a corpus row (hybrid_search ``query.where`` / ``knn.where``).""" if "$and" in where: @@ -545,6 +554,877 @@ def get_hybrid_combined_case(name: str) -> HybridCombinedCase: raise KeyError(f"unknown combined hybrid case: {name!r}") +# --- Triple-branch hybrid_search (vector + full-text + search index) --- + + +@dataclass(frozen=True) +class HybridTripleBranchCase: + """ + Namespace hybrid_search with ``query`` (FTS + search index) and ``knn`` all active. + + ``verify`` selects which branch supplies ground-truth assertions; the other branches + use broad aligned filters so every signal participates without masking the operator + under test. + """ + + name: str + n_results: int + verify: Literal["fts", "search_index", "knn", "intersection", "not_contains", "rrf_fusion"] + where_document: dict[str, Any] | str | None = None + where: dict[str, Any] | None = None + knn: dict[str, Any] | None = None + use_rrf: bool = False + check_fts_ranking: bool = True + min_hits: int = 1 + exact_match_count: int | None = None + + +def _default_knn( + n_results: int, + *, + where: dict[str, Any] | None = None, + knn_n_results: int | None = None, +) -> dict[str, Any]: + return { + "query_embeddings": KNN_QUERY_VECTOR, + "n_results": knn_n_results or max(n_results * 2, 20), + "where": where or WHERE_BROAD_METADATA, + } + + +def corpus_matches_triple_intersection( + record: CorpusRecord, + case: HybridTripleBranchCase, +) -> bool: + if case.where_document is not None: + if not doc_matches_where_document(record.document, case.where_document): + return False + if case.where is not None: + if not corpus_matches_where(record, case.where): + return False + if case.knn is not None: + knn_where = case.knn.get("where") + if knn_where is not None and not corpus_matches_where(record, knn_where): + return False + return True + + +def count_triple_intersection_matches( + corpus: list[CorpusRecord], + case: HybridTripleBranchCase, +) -> int: + return sum(1 for rec in corpus if corpus_matches_triple_intersection(rec, case)) + + +def _slice_hybrid_result( + result: dict[str, Any], + corpus: list[CorpusRecord], + predicate, +) -> dict[str, Any]: + corpus_by_id = {rec.doc_id: rec for rec in corpus} + ids = result["ids"][0] + indices = [ + i + for i, doc_id in enumerate(ids) + if doc_id in corpus_by_id and predicate(corpus_by_id[doc_id]) + ] + sliced: dict[str, Any] = {"ids": [[ids[i] for i in indices]]} + for key in ("distances", "documents", "metadatas"): + rows = result.get(key) + if rows and rows[0] is not None: + sliced[key] = [[rows[0][i] for i in indices]] + return sliced + + +def assert_hybrid_triple_branch_result( + corpus: list[CorpusRecord], + result: dict[str, Any], + case: HybridTripleBranchCase, +) -> None: + assert result is not None + assert "ids" in result and result["ids"] + ids = result["ids"][0] + assert len(ids) <= case.n_results + assert len(ids) >= case.min_hits, ( + f"case {case.name!r}: expected at least {case.min_hits} hits, got {len(ids)}" + ) + + corpus_by_id = {rec.doc_id: rec for rec in corpus} + for doc_id in ids: + assert doc_id in corpus_by_id, f"unknown id {doc_id!r}" + + if case.verify == "not_contains": + assert case.where_document is not None + assert_not_contains_no_token_leak(corpus, result, TOKEN_ZPX) + if case.where is not None: + for doc_id in ids: + assert doc_matches_where_metadata( + corpus_by_id[doc_id].metadata, case.where + ), f"id={doc_id!r} violates query.where={case.where!r}" + if case.knn and case.knn.get("where"): + for doc_id in ids: + assert corpus_matches_where(corpus_by_id[doc_id], case.knn["where"]), ( + f"id={doc_id!r} violates knn.where={case.knn['where']!r}" + ) + return + + if case.verify == "fts": + assert case.where_document is not None + + def _fts_row(rec: CorpusRecord) -> bool: + if not doc_matches_where_document(rec.document, case.where_document): + return False + if case.where is not None and not doc_matches_where_metadata(rec.metadata, case.where): + return False + return True + + fts_result = _slice_hybrid_result(result, corpus, _fts_row) + fts_ids = fts_result["ids"][0] + assert len(fts_ids) >= 1, ( + f"case {case.name!r}: expected at least one FTS-matching row in hybrid result" + ) + assert_hybrid_fulltext_result( + corpus, + fts_result, + case.where_document, + min(len(fts_ids), case.n_results), + where=case.where, + check_ranking=case.check_fts_ranking, + ) + if case.knn and case.knn.get("where"): + for doc_id in ids: + assert corpus_matches_where(corpus_by_id[doc_id], case.knn["where"]), ( + f"id={doc_id!r} violates knn.where={case.knn['where']!r}" + ) + return + + if case.verify == "search_index": + assert case.where is not None + + def _si_row(rec: CorpusRecord) -> bool: + if not corpus_matches_where(rec, case.where): + return False + if case.where_document is not None and not doc_matches_where_document( + rec.document, case.where_document + ): + return False + return True + + si_result = _slice_hybrid_result(result, corpus, _si_row) + si_ids = si_result["ids"][0] + assert len(si_ids) >= case.min_hits, ( + f"case {case.name!r}: expected at least {case.min_hits} search-index rows" + ) + assert_hybrid_search_index_result( + corpus, + si_result, + case.where, + min(len(si_ids), case.n_results), + min_hits=case.min_hits, + exact_match_count=case.exact_match_count, + ) + if case.knn and case.knn.get("where"): + for doc_id in ids: + assert corpus_matches_where(corpus_by_id[doc_id], case.knn["where"]), ( + f"id={doc_id!r} violates knn.where={case.knn['where']!r}" + ) + return + + if case.verify == "knn": + assert case.knn is not None + knn_where = case.knn.get("where") + + def _knn_row(rec: CorpusRecord) -> bool: + if knn_where is not None and not corpus_matches_where(rec, knn_where): + return False + if case.where is not None and not doc_matches_where_metadata(rec.metadata, case.where): + return False + if case.where_document is not None and not doc_matches_where_document( + rec.document, case.where_document + ): + return False + return True + + knn_result = _slice_hybrid_result(result, corpus, _knn_row) + knn_ids = knn_result["ids"][0] + assert len(knn_ids) >= 1, ( + f"case {case.name!r}: expected at least one KNN-eligible row in hybrid result" + ) + assert_hybrid_knn_result( + corpus, + knn_result, + case.knn["query_embeddings"], + min(len(knn_ids), case.n_results), + knn_where, + check_top1=True, + ) + return + + if case.verify == "intersection": + total = count_triple_intersection_matches(corpus, case) + if case.exact_match_count is not None: + assert total == case.exact_match_count, ( + f"case {case.name!r}: expected {case.exact_match_count} intersection " + f"matches, corpus has {total}" + ) + for doc_id in ids: + assert corpus_matches_triple_intersection(corpus_by_id[doc_id], case), ( + f"id={doc_id!r} outside triple-branch intersection for case {case.name!r}" + ) + if case.exact_match_count is not None and case.exact_match_count <= case.n_results: + expected_ids = { + rec.doc_id + for rec in corpus + if corpus_matches_triple_intersection(rec, case) + } + assert set(ids) == expected_ids + return + + if case.verify == "rrf_fusion": + if case.where is not None: + for doc_id in ids: + assert doc_matches_where_metadata( + corpus_by_id[doc_id].metadata, case.where + ), f"id={doc_id!r} violates query.where={case.where!r}" + if case.knn and case.knn.get("where"): + for doc_id in ids: + assert corpus_matches_where(corpus_by_id[doc_id], case.knn["where"]), ( + f"id={doc_id!r} violates knn.where={case.knn['where']!r}" + ) + intersection_ids = { + rec.doc_id + for rec in corpus + if corpus_matches_triple_intersection(rec, case) + } + assert intersection_ids, ( + f"case {case.name!r}: intersection must be non-empty for RRF fusion checks" + ) + returned = set(ids) + assert returned & intersection_ids, ( + f"case {case.name!r}: RRF result must include at least one intersection hit, " + f"got {ids[:5]!r}" + ) + if case.where_document is not None: + fts_hits = { + rec.doc_id + for rec in corpus + if doc_matches_where_document(rec.document, case.where_document) + } + assert returned & fts_hits, ( + f"case {case.name!r}: RRF result must include at least one FTS hit" + ) + return + + raise ValueError(f"unsupported verify mode: {case.verify!r}") + + +def _build_triple_branch_hybrid_search_kwargs( + case: HybridTripleBranchCase, +) -> dict[str, Any]: + query: dict[str, Any] | None = None + if case.where_document is not None or case.where is not None: + query = {"n_results": case.n_results} + if case.where_document is not None: + query["where_document"] = case.where_document + if case.where is not None: + query["where"] = case.where + return { + "query": query, + "knn": case.knn, + "rank": RRF_RANK if case.use_rrf else None, + "n_results": case.n_results, + "include": ["documents", "metadatas", "distances"], + } + + +def assert_flat_collection_baseline(collection: Any) -> None: + """Runtime guard: baseline must be a flat collection (HNSW heap table path).""" + if getattr(collection, "use_namespace", None) is not False: + raise TypeError( + f"collection baseline must have use_namespace=False, " + f"got {getattr(collection, 'use_namespace', None)!r} on {collection!r}" + ) + if not hasattr(collection, "hybrid_search"): + raise TypeError(f"collection baseline must expose hybrid_search(), got {type(collection)!r}") + + +def execute_hybrid_triple_branch_on_collection( + collection: Any, + case: HybridTripleBranchCase, +) -> dict[str, Any]: + """ + Triple-branch hybrid_search on flat collection. + + Path: ``collection.add`` (during setup) → ``collection.hybrid_search`` + → ``BaseClient._collection_hybrid_search`` → ``DBMS_HYBRID_SEARCH.GET_SQL``. + Vector index: HNSW (``flat_hybrid_search_schema``). + """ + assert_flat_collection_baseline(collection) + return collection.hybrid_search(**_build_triple_branch_hybrid_search_kwargs(case)) + + +def execute_hybrid_triple_branch_on_namespace( + namespace: Any, + case: HybridTripleBranchCase, +) -> dict[str, Any]: + """ + Triple-branch hybrid_search on namespace. + + Path: ``namespace.add`` (during setup) → ``namespace.hybrid_search`` + → ``BaseClient._namespace_hybrid_search`` → ``hybrid_search(TABLE logic_data_table, ...)``. + Vector index: IVF spfresh (``ns_schema``). + """ + if not hasattr(namespace, "namespace_id"): + raise TypeError(f"expected Namespace, got {type(namespace)!r}") + return namespace.hybrid_search(**_build_triple_branch_hybrid_search_kwargs(case)) + + +def execute_hybrid_triple_branch_search( + target: Any, + case: HybridTripleBranchCase, +) -> dict[str, Any]: + """Dispatch to collection or namespace executor based on target type.""" + if getattr(target, "use_namespace", None) is False: + return execute_hybrid_triple_branch_on_collection(target, case) + if hasattr(target, "namespace_id"): + return execute_hybrid_triple_branch_on_namespace(target, case) + raise TypeError( + f"unsupported hybrid_search target {target!r}: expected flat Collection " + f"(use_namespace=False) or Namespace" + ) + + +def setup_large_fts_flat_collection(db_client: Any) -> tuple[list[CorpusRecord], Any]: + """Create flat collection (HNSW + FTS), load via ``collection.add``.""" + corpus = build_large_fts_corpus(CORPUS_SIZE) + if len(corpus) <= 1000: + raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") + + name = f"test_hs_tb_baseline_{int(time.time() * 1000)}" + collection = db_client.create_collection( + name=name, schema=flat_hybrid_search_schema(), use_namespace=False + ) + assert collection.use_namespace is False + insert_corpus_into_collection(collection, corpus) + expected = len(corpus) + actual = collection.count() + assert actual == expected, ( + f"baseline collection {collection.name!r} expected {expected} rows, got {actual}" + ) + time.sleep(INDEX_SETTLE_SECONDS) + return corpus, collection + + +def _try_assert_hybrid_triple_branch_result( + corpus: list[CorpusRecord], + result: dict[str, Any], + case: HybridTripleBranchCase, +) -> BaseException | None: + try: + assert_hybrid_triple_branch_result(corpus, result, case) + return None + except BaseException as exc: + return exc + + +def assert_hybrid_triple_branch_vs_collection_baseline( + corpus: list[CorpusRecord], + case: HybridTripleBranchCase, + namespace_result: dict[str, Any], + collection_result: dict[str, Any], +) -> None: + """ + Compare namespace (IVF logic table) vs collection baseline (HNSW heap table). + + Collection baseline uses ``use_namespace=False``: data via ``collection.add``, + search via ``collection.hybrid_search`` — the standard collection hybrid path. + + Because HNSW vs IVF may rank differently on KNN/RRF branches, **id list equality + is not required** when both sides satisfy corpus ground truth. Comparison focuses + on pass/fail attribution: + + - ``[test-case]`` — collection baseline also fails; likely expectation / case design issue + - ``[observer/namespace]`` — collection passes ground truth but namespace fails + - ``[test-case/baseline]`` — namespace passes but collection baseline fails + """ + coll_err = _try_assert_hybrid_triple_branch_result(corpus, collection_result, case) + ns_err = _try_assert_hybrid_triple_branch_result(corpus, namespace_result, case) + coll_ids = collection_result["ids"][0] + ns_ids = namespace_result["ids"][0] + + if coll_err is None and ns_err is None: + # Both paths satisfy ground truth; HNSW vs IVF ranking may differ — success. + return + + if coll_err is not None and ns_err is not None: + if coll_ids == ns_ids: + raise AssertionError( + f"[test-case] case {case.name!r}: collection (use_namespace=false, HNSW) and " + f"namespace (IVF) fail the same ground-truth checks with identical ids; " + f"likely test expectation issue.\n" + f" ids={coll_ids[:10]!r}\n" + f" error={coll_err!r}" + ) from coll_err + raise AssertionError( + f"[mixed] case {case.name!r}: collection and namespace both fail ground truth " + f"and ids differ.\n" + f" collection_ids={coll_ids[:10]!r}\n" + f" namespace_ids={ns_ids[:10]!r}\n" + f" collection_error={coll_err!r}\n" + f" namespace_error={ns_err!r}" + ) from ns_err + + if coll_err is None and ns_err is not None: + raise AssertionError( + f"[observer/namespace] case {case.name!r}: collection baseline (HNSW, " + f"use_namespace=false) passes but namespace (IVF) fails.\n" + f" collection_ids={coll_ids!r}\n" + f" namespace_ids={ns_ids!r}\n" + f" namespace_error={ns_err!r}" + ) from ns_err + + raise AssertionError( + f"[test-case/baseline] case {case.name!r}: namespace (IVF) passes but collection " + f"baseline (HNSW, use_namespace=false) fails; revisit case assumptions.\n" + f" collection_ids={coll_ids!r}\n" + f" namespace_ids={ns_ids!r}\n" + f" collection_error={coll_err!r}" + ) from coll_err + + +def run_hybrid_triple_branch_case( + namespace: Any, + corpus: list[CorpusRecord], + case: HybridTripleBranchCase, + *, + collection_baseline: Any | None = None, +) -> dict[str, Any]: + if collection_baseline is not None: + assert_flat_collection_baseline(collection_baseline) + try: + collection_result = execute_hybrid_triple_branch_on_collection( + collection_baseline, case + ) + except BaseException as coll_exec_err: + raise AssertionError( + f"[test-case/baseline] case {case.name!r}: collection.hybrid_search " + f"(use_namespace=false, HNSW/GET_SQL) raised: {coll_exec_err!r}" + ) from coll_exec_err + try: + namespace_result = execute_hybrid_triple_branch_on_namespace(namespace, case) + except BaseException as ns_exec_err: + raise AssertionError( + f"[observer/namespace] case {case.name!r}: collection.hybrid_search " + f"succeeded but namespace.hybrid_search (IVF/logic table) raised: " + f"{ns_exec_err!r}" + ) from ns_exec_err + assert_hybrid_triple_branch_vs_collection_baseline( + corpus, case, namespace_result, collection_result + ) + return namespace_result + + result = execute_hybrid_triple_branch_on_namespace(namespace, case) + assert_hybrid_triple_branch_result(corpus, result, case) + return result + + +TRIPLE_BRANCH_CASES: list[HybridTripleBranchCase] = [ + # --- Full-text operators (FTS is primary; search index + KNN use broad filters) --- + HybridTripleBranchCase( + name="fts_contains_zpx", + verify="fts", + where_document={"$contains": TOKEN_ZPX}, + where=WHERE_BROAD_METADATA, + knn=_default_knn(15), + n_results=15, + ), + HybridTripleBranchCase( + name="fts_contains_alp", + verify="fts", + where_document={"$contains": TOKEN_ALP}, + where=WHERE_BROAD_METADATA, + knn=_default_knn(15), + n_results=15, + ), + HybridTripleBranchCase( + name="fts_string_shorthand", + verify="fts", + where_document=TOKEN_ZPX, + where=WHERE_BROAD_METADATA, + knn=_default_knn(15), + n_results=15, + ), + HybridTripleBranchCase( + name="fts_not_contains", + verify="not_contains", + where_document={"$not_contains": TOKEN_ZPX}, + where=WHERE_FILLER_SEQ, + knn=_default_knn(20, where=WHERE_FILLER_SEQ), + n_results=20, + check_fts_ranking=False, + ), + HybridTripleBranchCase( + name="fts_and_zpx_alp", + verify="fts", + where_document={ + "$and": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + where=WHERE_BROAD_METADATA, + knn=_default_knn(10), + n_results=10, + ), + HybridTripleBranchCase( + name="fts_or_zpx_alp", + verify="fts", + where_document={ + "$or": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + where=WHERE_BROAD_METADATA, + knn=_default_knn(20), + n_results=20, + ), + HybridTripleBranchCase( + name="fts_and_multi_contains", + verify="fts", + where_document={ + "$and": [ + {"$contains": "Primary"}, + {"$contains": TOKEN_ZPX}, + ], + }, + where=WHERE_BROAD_METADATA, + knn=_default_knn(10), + n_results=10, + ), + HybridTripleBranchCase( + name="fts_contains_filter_has_both", + verify="fts", + where_document={"$contains": TOKEN_ZPX}, + where=WHERE_BOTH_TOKENS_NUMERIC, + knn=_default_knn(5, where=WHERE_BOTH_TOKENS_NUMERIC), + n_results=5, + ), + HybridTripleBranchCase( + name="fts_contains_filter_zpx_hint_eq", + verify="fts", + where_document={"$contains": TOKEN_ZPX}, + where={"zpx_hint": 50}, + knn=_default_knn(5, where={"zpx_hint": 50}), + n_results=5, + ), + HybridTripleBranchCase( + name="fts_contains_filter_zpx_hint_gte", + verify="fts", + where_document={"$contains": TOKEN_ZPX}, + where={"zpx_hint": {"$gte": 40}}, + knn=_default_knn(10, where={"zpx_hint": {"$gte": 40}}), + n_results=10, + ), + HybridTripleBranchCase( + name="fts_and_zpx_alp_filter_has_both", + verify="fts", + where_document={ + "$and": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + where=WHERE_BOTH_TOKENS_NUMERIC, + knn=_default_knn(5, where=WHERE_BOTH_TOKENS_NUMERIC), + n_results=5, + ), + HybridTripleBranchCase( + name="fts_contains_filter_zpx_hint_and_alp_zero", + verify="fts", + where_document={"$contains": TOKEN_ZPX}, + where={ + "$and": [ + {"zpx_hint": {"$gte": 40}}, + {"alp_hint": 0}, + ], + }, + knn=_default_knn( + 10, + where={ + "$and": [ + {"zpx_hint": {"$gte": 40}}, + {"alp_hint": 0}, + ], + }, + ), + n_results=10, + ), + # --- Search-index operators (metadata is primary; universal FTS + broad KNN) --- + HybridTripleBranchCase( + name="si_eq_zpx_hint_direct", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"zpx_hint": 50}, + knn=_default_knn(5, where={"zpx_hint": 50}), + n_results=5, + exact_match_count=1, + ), + HybridTripleBranchCase( + name="si_eq_zpx_hint_operator", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"zpx_hint": {"$eq": 50}}, + knn=_default_knn(5, where={"zpx_hint": {"$eq": 50}}), + n_results=5, + exact_match_count=1, + ), + HybridTripleBranchCase( + name="si_ne_zpx_hint_zero", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"zpx_hint": {"$ne": 0}}, + knn=_default_knn(20, where={"zpx_hint": {"$ne": 0}}), + n_results=20, + min_hits=1, + ), + HybridTripleBranchCase( + name="si_gte_zpx_hint_40", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"zpx_hint": {"$gte": 40}}, + knn=_default_knn(15, where={"zpx_hint": {"$gte": 40}}), + n_results=15, + min_hits=2, + ), + HybridTripleBranchCase( + name="si_lt_zpx_hint_10", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"zpx_hint": {"$lt": 10}}, + knn=_default_knn(20, where={"zpx_hint": {"$lt": 10}}), + n_results=20, + min_hits=1, + ), + HybridTripleBranchCase( + name="si_lte_alp_hint_8", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"alp_hint": {"$lte": 8}}, + knn=_default_knn(15, where={"alp_hint": {"$lte": 8}}), + n_results=15, + min_hits=1, + ), + HybridTripleBranchCase( + name="si_gt_rel_hint_zero", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"rel_hint": {"$gt": 0}}, + knn=_default_knn(25, where={"rel_hint": {"$gt": 0}}), + n_results=25, + min_hits=5, + ), + HybridTripleBranchCase( + name="si_in_alp_hint_values", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"alp_hint": {"$in": [32, 24, 16, 8]}}, + knn=_default_knn(10, where={"alp_hint": {"$in": [32, 24, 16, 8]}}), + n_results=10, + min_hits=4, + ), + HybridTripleBranchCase( + name="si_nin_alp_hint_high", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"alp_hint": {"$nin": [32, 24, 16, 8]}}, + knn=_default_knn(15, where={"alp_hint": {"$nin": [32, 24, 16, 8]}}), + n_results=15, + min_hits=1, + ), + HybridTripleBranchCase( + name="si_and_has_both_zpx_gte", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={ + "$and": [ + {"zpx_hint": {"$gte": 5}}, + {"alp_hint": 5}, + ], + }, + knn=_default_knn( + 5, + where={ + "$and": [ + {"zpx_hint": {"$gte": 5}}, + {"alp_hint": 5}, + ], + }, + ), + n_results=5, + exact_match_count=1, + ), + HybridTripleBranchCase( + name="si_or_zpx_alp_hint_high", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={ + "$or": [ + {"zpx_hint": {"$gte": 50}}, + {"alp_hint": {"$gte": 32}}, + ], + }, + knn=_default_knn( + 10, + where={ + "$or": [ + {"zpx_hint": {"$gte": 50}}, + {"alp_hint": {"$gte": 32}}, + ], + }, + ), + n_results=10, + min_hits=2, + ), + HybridTripleBranchCase( + name="si_not_has_both", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where=WHERE_NOT_BOTH_TOKENS_NUMERIC, + knn=_default_knn(20, where=WHERE_NOT_BOTH_TOKENS_NUMERIC), + n_results=20, + min_hits=1, + ), + HybridTripleBranchCase( + name="si_id_eq_both_tokens", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"#id": "both_tokens"}, + knn=_default_knn(5, where={"#id": "both_tokens"}), + n_results=5, + exact_match_count=1, + ), + HybridTripleBranchCase( + name="si_id_in_zpx_tops", + verify="search_index", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"#id": {"$in": ["zpx_top_5", "zpx_top_4", "zpx_top_3"]}}, + knn=_default_knn( + 5, + where={"#id": {"$in": ["zpx_top_5", "zpx_top_4", "zpx_top_3"]}}, + ), + n_results=5, + exact_match_count=3, + ), + # --- Vector / KNN operators (KNN primary; universal FTS + broad metadata) --- + HybridTripleBranchCase( + name="knn_global_top5", + verify="knn", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where=WHERE_BROAD_METADATA, + knn=_default_knn(5, knn_n_results=10), + n_results=5, + ), + HybridTripleBranchCase( + name="knn_filter_has_both", + verify="knn", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where=WHERE_BOTH_TOKENS_NUMERIC, + knn=_default_knn(3, where=WHERE_BOTH_TOKENS_NUMERIC, knn_n_results=10), + n_results=3, + ), + HybridTripleBranchCase( + name="knn_filter_zpx_hint_gte_40", + verify="knn", + where_document=WHERE_DOCUMENT_UNIVERSAL, + where={"zpx_hint": {"$gte": 40}}, + knn=_default_knn(5, where={"zpx_hint": {"$gte": 40}}, knn_n_results=15), + n_results=5, + ), + # --- Aligned intersection (all three branches share the same filter intent) --- + HybridTripleBranchCase( + name="intersection_zpx_gte40", + verify="intersection", + where_document={"$contains": TOKEN_ZPX}, + where={"zpx_hint": {"$gte": 40}}, + knn=_default_knn(10, where={"zpx_hint": {"$gte": 40}}), + n_results=10, + min_hits=2, + ), + HybridTripleBranchCase( + name="intersection_both_tokens", + verify="intersection", + where_document={ + "$and": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + where=WHERE_BOTH_TOKENS_NUMERIC, + knn=_default_knn(5, where=WHERE_BOTH_TOKENS_NUMERIC), + n_results=5, + exact_match_count=1, + ), + HybridTripleBranchCase( + name="intersection_id_zpx_top5", + verify="intersection", + where_document={"$contains": TOKEN_ZPX}, + where={"#id": "zpx_top_5"}, + knn=_default_knn(5, where={"#id": "zpx_top_5"}), + n_results=5, + exact_match_count=1, + ), + # --- RRF fusion across three branches --- + HybridTripleBranchCase( + name="rrf_fts_or_knn_si_gte", + verify="rrf_fusion", + where_document={ + "$or": [ + {"$contains": TOKEN_ZPX}, + {"$contains": TOKEN_ALP}, + ], + }, + where={"rel_hint": {"$gte": 4}}, + knn=_default_knn(10, where={"rel_hint": {"$gte": 4}}, knn_n_results=15), + n_results=10, + use_rrf=True, + min_hits=3, + ), + HybridTripleBranchCase( + name="rrf_all_branches_zpx_gte40", + verify="rrf_fusion", + where_document={"$contains": TOKEN_ZPX}, + where={"zpx_hint": {"$gte": 40}}, + knn=_default_knn(8, where={"zpx_hint": {"$gte": 40}}, knn_n_results=15), + n_results=8, + use_rrf=True, + min_hits=2, + ), +] + + +def get_triple_branch_case(name: str) -> HybridTripleBranchCase: + for case in TRIPLE_BRANCH_CASES: + if case.name == name: + return case + raise KeyError(f"unknown triple-branch case: {name!r}") + + +def run_triple_branch_case_on_quadrants( + ctx: dict[str, Any], + case: HybridTripleBranchCase | str, + quadrant_keys: tuple[str, ...] = MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, + *, + collection_baseline: Any | None = None, +) -> None: + tb_case = case if isinstance(case, HybridTripleBranchCase) else get_triple_branch_case(case) + for key in quadrant_keys: + corpus, namespace = ctx[key] + run_hybrid_triple_branch_case( + namespace, corpus, tb_case, collection_baseline=collection_baseline + ) + + def run_search_index_case_on_quadrants( ctx: dict[str, Any], case: SearchIndexQueryCase | str, @@ -560,11 +1440,7 @@ def run_knn_case_on_quadrants( ctx: dict[str, Any], case: VectorKnnCase | str, quadrant_keys: tuple[str, ...] = MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, - *, - request: Any = None, ) -> None: - if request is not None: - skip_oceanbase_knn(request) knn_case = case if isinstance(case, VectorKnnCase) else get_vector_knn_case(case) for key in quadrant_keys: corpus, namespace = ctx[key] @@ -588,6 +1464,31 @@ def assert_hybrid_search_index_no_hits( ) +def assert_hybrid_triple_branch_no_hits( + namespace: Any, + case: HybridTripleBranchCase, +) -> None: + """Assert triple-branch hybrid_search returns no rows (namespace isolation).""" + query: dict[str, Any] | None = None + if case.where_document is not None or case.where is not None: + query = {"n_results": case.n_results} + if case.where_document is not None: + query["where_document"] = case.where_document + if case.where is not None: + query["where"] = case.where + result = namespace.hybrid_search( + query=query, + knn=case.knn, + rank=RRF_RANK if case.use_rrf else None, + n_results=case.n_results, + include=["documents"], + ) + ids = result.get("ids", [[]])[0] if result.get("ids") else [] + assert len(ids) == 0, ( + f"expected no triple-branch hits in namespace {namespace.name!r}, got {ids[:5]!r}" + ) + + __all__ = [ "BATCH_SIZE", "CORPUS_SIZE", @@ -600,28 +1501,43 @@ def assert_hybrid_search_index_no_hits( "SEARCH_INDEX_CASES", "TOKEN_ALP", "TOKEN_ZPX", + "TRIPLE_BRANCH_CASES", "VECTOR_KNN_CASES", + "WHERE_BROAD_METADATA", + "WHERE_DOCUMENT_UNIVERSAL", + "WHERE_FILLER_SEQ", "HybridCombinedCase", + "HybridTripleBranchCase", "SearchIndexQueryCase", "VectorKnnCase", + "assert_flat_collection_baseline", "assert_hybrid_search_index_no_hits", + "assert_hybrid_triple_branch_no_hits", + "assert_hybrid_triple_branch_result", + "assert_hybrid_triple_branch_vs_collection_baseline", "assert_not_contains_no_token_leak", "build_large_fts_corpus", + "execute_hybrid_triple_branch_on_collection", + "execute_hybrid_triple_branch_on_namespace", + "execute_hybrid_triple_branch_search", "get_fts_case", "get_hybrid_combined_case", "get_search_index_case", + "get_triple_branch_case", "get_vector_knn_case", "run_hybrid_combined_case", "run_hybrid_knn_case", "run_hybrid_search_index_case", "run_hybrid_search_fts_case", + "run_hybrid_triple_branch_case", "run_knn_case_on_quadrants", "run_search_index_case_on_quadrants", + "run_triple_branch_case_on_quadrants", "setup_fts_namespace_with_corpus", "setup_large_fts_collection", + "setup_large_fts_flat_collection", "setup_multi_coll_multi_ns_fts", "setup_multi_coll_multi_ns_fts_single_loaded", - "skip_oceanbase_knn", "teardown_large_fts_collection", "teardown_multi_coll_multi_ns_fts", ] diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py index 612538d7..eaca0277 100644 --- a/tests/integration_tests/test_namespace_dml.py +++ b/tests/integration_tests/test_namespace_dml.py @@ -1,45 +1,45 @@ """ Namespace DML integration tests. -Tests namespace.add, update, upsert, delete, get, count, peek operations. + +Small-scale: add / update / upsert / delete / get. +Large-scale (>=1000 rows): count / peek boundaries and multi-namespace orthogonality. """ -import time +from __future__ import annotations import pytest -from pyseekdb import IVFConfiguration -from pyseekdb.client.configuration import VectorIndexConfig -from pyseekdb.client.schema import Schema - -from namespace_dml_helpers import assert_get_absent, assert_get_present, cleanup, create_ns_collection +from namespace_dml_helpers import ( + LARGE_DML_CORPUS_SIZE, + assert_get_absent, + assert_get_present, + assert_peek_result, + build_large_dml_corpus, + cleanup, + corpus_id_set, + create_ns_collection, + insert_dml_corpus_in_batches, + load_dml_corpus, +) class TestNamespaceDML: - - def _setup(self, client): - name = f"test_ns_dml_{int(time.time() * 1000)}" - schema = Schema( - vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2"), - embedding_function=None, - ), - ) - collection = client.create_collection(name=name, schema=schema, use_namespace=True) - namespace = collection.create_namespace("dml_ns") - return collection, namespace + """Small-scale DML smoke tests.""" def test_add_single(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_add1") + ns = collection.create_namespace("dml_ns") try: ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0], documents="Hello", metadatas={"tag": "a"}) result = ns.get(ids="d1") assert len(result["ids"]) == 1 assert result["ids"][0] == "d1" finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) def test_add_batch(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_addb") + ns = collection.create_namespace("dml_ns") try: ns.add( ids=["d1", "d2", "d3"], @@ -50,10 +50,11 @@ def test_add_batch(self, db_client): result = ns.get(ids=["d1", "d2", "d3"]) assert len(result["ids"]) == 3 finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) def test_get_by_id(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_getid") + ns = collection.create_namespace("dml_ns") try: ns.add( ids=["g1", "g2"], @@ -66,10 +67,11 @@ def test_get_by_id(self, db_client): assert result["documents"][0] == "First" assert result["metadatas"][0]["k"] == 1 finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) def test_get_with_limit(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_getlim") + ns = collection.create_namespace("dml_ns") try: ns.add( ids=["l1", "l2", "l3"], @@ -78,94 +80,255 @@ def test_get_with_limit(self, db_client): result = ns.get(limit=2) assert len(result["ids"]) == 2 finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) def test_update_metadata(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_upd") + ns = collection.create_namespace("dml_ns") try: ns.add(ids="u1", embeddings=[1.0, 2.0, 3.0], metadatas={"score": 10}) ns.update(ids="u1", metadatas={"score": 99}) result = ns.get(ids="u1", include=["metadatas"]) assert result["metadatas"][0]["score"] == 99 finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) def test_update_document_and_embedding(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_upddoc") + ns = collection.create_namespace("dml_ns") try: ns.add(ids="u2", embeddings=[1.0, 2.0, 3.0], documents="Original") ns.update(ids="u2", embeddings=[9.0, 8.0, 7.0], documents="Updated") result = ns.get(ids="u2", include=["documents", "embeddings"]) assert result["documents"][0] == "Updated" finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) def test_upsert_existing(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_upsex") + ns = collection.create_namespace("dml_ns") try: ns.add(ids="up1", embeddings=[1.0, 2.0, 3.0], metadatas={"v": 1}) ns.upsert(ids="up1", embeddings=[4.0, 5.0, 6.0], metadatas={"v": 2}) result = ns.get(ids="up1", include=["metadatas"]) assert result["metadatas"][0]["v"] == 2 finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) def test_upsert_new(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_upsnew") + ns = collection.create_namespace("dml_ns") try: ns.upsert(ids="up_new", embeddings=[1.0, 1.0, 1.0], metadatas={"fresh": True}) result = ns.get(ids="up_new", include=["metadatas"]) assert len(result["ids"]) == 1 assert result["metadatas"][0]["fresh"] is True finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) def test_delete_by_ids(self, db_client): - collection, ns = self._setup(db_client) + collection = create_ns_collection(db_client, suffix="_del") + ns = collection.create_namespace("dml_ns") try: ns.add( ids=["del1", "del2"], embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], ) ns.delete(ids="del1") - result = ns.get(ids="del1") - assert len(result["ids"]) == 0 - - result = ns.get(ids="del2") - assert len(result["ids"]) == 1 + assert len(ns.get(ids="del1")["ids"]) == 0 + assert len(ns.get(ids="del2")["ids"]) == 1 finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) + + +class TestNamespaceDMLCountPeekAtScale: + """count / peek with >=1000 rows in a single namespace.""" - def test_count(self, db_client): - collection, ns = self._setup(db_client) + @pytest.fixture + def large_ns(self, db_client): + collection = create_ns_collection(db_client, suffix="_large_cp") + ns = collection.create_namespace("large_ns") + corpus = build_large_dml_corpus(LARGE_DML_CORPUS_SIZE, id_prefix="large") + load_dml_corpus(ns, corpus) + yield collection, ns, corpus + cleanup(db_client, collection) + + def test_count_empty_namespace(self, db_client): + collection = create_ns_collection(db_client, suffix="_cnt_empty") + ns = collection.create_namespace("empty_ns") try: assert ns.count() == 0 - ns.add( - ids=["c1", "c2", "c3"], - embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], - ) - assert ns.count() == 3 finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) + + def test_count_after_bulk_load_1000(self, large_ns): + _, ns, corpus = large_ns + assert ns.count() == len(corpus) == LARGE_DML_CORPUS_SIZE - def test_peek(self, db_client): - collection, ns = self._setup(db_client) + def test_count_after_partial_delete(self, large_ns): + _, ns, corpus = large_ns + to_delete = [corpus[i].doc_id for i in range(100)] + ns.delete(ids=to_delete) + assert ns.count() == LARGE_DML_CORPUS_SIZE - 100 + assert len(ns.get(ids=to_delete)["ids"]) == 0 + assert len(ns.get(ids=corpus[500].doc_id)["ids"]) == 1 + + def test_peek_empty_namespace(self, db_client): + collection = create_ns_collection(db_client, suffix="_peek_empty") + ns = collection.create_namespace("peek_empty") try: - ns.add( - ids=["p1", "p2"], - embeddings=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], - documents=["Peek A", "Peek B"], - metadatas=[{"x": 1}, {"x": 2}], - ) result = ns.peek(limit=10) - assert len(result["ids"]) == 2 - assert "documents" in result - assert "metadatas" in result + assert_peek_result(result, expected_len=0) finally: - db_client.delete_collection(name=collection.name) + cleanup(db_client, collection) + + @pytest.mark.parametrize( + ("limit", "expected_len"), + [ + (0, 0), + (1, 1), + (10, 10), + (100, 100), + (999, 999), + (1000, 1000), + (1001, 1000), + (5000, 1000), + ], + ids=[ + "limit_0", + "limit_1", + "limit_10", + "limit_100", + "limit_999", + "limit_1000", + "limit_over_total", + "limit_far_over_total", + ], + ) + def test_peek_limit_boundaries(self, large_ns, limit, expected_len): + _, ns, corpus = large_ns + allowed = corpus_id_set(corpus) + result = ns.peek(limit=limit) + assert_peek_result(result, expected_len=expected_len, allowed_ids=allowed) + + def test_peek_default_limit_is_10(self, large_ns): + _, ns, corpus = large_ns + result = ns.peek() + assert_peek_result( + result, + expected_len=10, + allowed_ids=corpus_id_set(corpus), + ) + + +class TestNamespaceDMLCountPeekMultiNs: + """count / peek orthogonality: same collection, different namespaces at scale.""" + + NS_A_SIZE = 1000 + NS_B_SIZE = 800 + + @pytest.fixture + def dual_ns_ctx(self, db_client): + collection = create_ns_collection(db_client, suffix="_dual_cp") + ns_a = collection.create_namespace("ns_alpha") + ns_b = collection.create_namespace("ns_beta") + corpus_a = build_large_dml_corpus( + self.NS_A_SIZE, id_prefix="alpha", ns_tag="alpha" + ) + corpus_b = build_large_dml_corpus( + self.NS_B_SIZE, id_prefix="beta", ns_tag="beta" + ) + load_dml_corpus(ns_a, corpus_a) + load_dml_corpus(ns_b, corpus_b) + ctx = { + "collection": collection, + "ns_a": ns_a, + "ns_b": ns_b, + "corpus_a": corpus_a, + "corpus_b": corpus_b, + "ids_a": corpus_id_set(corpus_a), + "ids_b": corpus_id_set(corpus_b), + } + yield ctx + cleanup(db_client, collection) + + def test_count_isolated_per_namespace(self, dual_ns_ctx): + assert dual_ns_ctx["ns_a"].count() == self.NS_A_SIZE + assert dual_ns_ctx["ns_b"].count() == self.NS_B_SIZE + + @pytest.mark.parametrize( + ("limit", "expected_len"), + [(1, 1), (10, 10), (100, 100), (800, 800), (1000, 1000), (1500, 1000)], + ids=["lim_1", "lim_10", "lim_100", "lim_800", "lim_1000", "lim_over_alpha"], + ) + def test_peek_alpha_boundaries(self, dual_ns_ctx, limit, expected_len): + result = dual_ns_ctx["ns_a"].peek(limit=limit) + assert_peek_result( + result, + expected_len=expected_len, + allowed_ids=dual_ns_ctx["ids_a"], + ns_tag="alpha", + ) + + @pytest.mark.parametrize( + ("limit", "expected_len"), + [(1, 1), (10, 10), (500, 500), (800, 800), (1000, 800), (2000, 800)], + ids=["lim_1", "lim_10", "lim_500", "lim_800", "lim_over_beta", "lim_far_over"], + ) + def test_peek_beta_boundaries(self, dual_ns_ctx, limit, expected_len): + result = dual_ns_ctx["ns_b"].peek(limit=limit) + assert_peek_result( + result, + expected_len=expected_len, + allowed_ids=dual_ns_ctx["ids_b"], + ns_tag="beta", + ) + + def test_peek_does_not_leak_across_namespaces(self, dual_ns_ctx): + peek_a = dual_ns_ctx["ns_a"].peek(limit=50) + peek_b = dual_ns_ctx["ns_b"].peek(limit=50) + assert_peek_result( + peek_a, + expected_len=50, + allowed_ids=dual_ns_ctx["ids_a"], + ns_tag="alpha", + ) + assert_peek_result( + peek_b, + expected_len=50, + allowed_ids=dual_ns_ctx["ids_b"], + ns_tag="beta", + ) + assert set(peek_a["ids"]).isdisjoint(dual_ns_ctx["ids_b"]) + assert set(peek_b["ids"]).isdisjoint(dual_ns_ctx["ids_a"]) + + def test_delete_in_one_namespace_only_affects_its_count(self, dual_ns_ctx): + ns_a = dual_ns_ctx["ns_a"] + ns_b = dual_ns_ctx["ns_b"] + corpus_a = dual_ns_ctx["corpus_a"] + delete_ids = [corpus_a[i].doc_id for i in range(150)] + ns_a.delete(ids=delete_ids) + + assert ns_a.count() == self.NS_A_SIZE - 150 + assert ns_b.count() == self.NS_B_SIZE + + peek_a = ns_a.peek(limit=20) + assert_peek_result( + peek_a, + expected_len=20, + allowed_ids=dual_ns_ctx["ids_a"] - set(delete_ids), + ns_tag="alpha", + ) + assert_peek_result( + ns_b.peek(limit=20), + expected_len=20, + allowed_ids=dual_ns_ctx["ids_b"], + ns_tag="beta", + ) class TestNamespaceDMLFullCycle: + """End-to-end DML cycle (small scale, verified by get).""" def test_full_dml_cycle_verified_by_get(self, db_client): collection = create_ns_collection(db_client, suffix="_cycle") @@ -216,6 +379,7 @@ def test_full_dml_cycle_verified_by_get(self, db_client): ns.delete(ids=doc_id) assert_get_absent(ns, doc_id) + assert ns.count() == 0 finally: cleanup(db_client, collection) diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index 74045610..cb044701 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -24,7 +24,7 @@ def _make_collection(client, suffix: str = ""): name = f"test_ns_drop_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine"), + ivf=IVFConfiguration(dimension=3, distance="cosine", use_spfresh=True), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_get_delete_filters.py b/tests/integration_tests/test_namespace_get_delete_filters.py new file mode 100644 index 00000000..997ab4a4 --- /dev/null +++ b/tests/integration_tests/test_namespace_get_delete_filters.py @@ -0,0 +1,202 @@ +""" +Namespace conditional get/delete integration tests (batch A / P0). + +Covers: + - delete(where=...) / delete(where_document=...) + - get(where=...) / get(where_document=...) + - multi-namespace isolation for metadata filter delete +""" + +from __future__ import annotations + +import pytest + +from namespace_dml_helpers import ( + FILTER_MULTI_NS_KEEP, + FILTER_MULTI_NS_PURGE, + assert_count_after_document_delete, + assert_get_where_count, + assert_get_where_document_count, + build_filter_corpus, + cleanup, + create_ns_collection, + load_dml_corpus, + load_filter_corpus, + precheck_where_document_hits, +) + + +class TestNamespaceDeleteWhere: + """P0 #1–#2: conditional delete on namespace.""" + + def test_delete_where_metadata_tag(self, db_client): + collection = create_ns_collection(db_client, suffix="_del_where") + ns = collection.create_namespace("del_where_ns") + corpus, gt = build_filter_corpus() + try: + load_dml_corpus(ns, corpus) + assert ns.count() == gt["total"] + + ns.delete(where={"tag": "purge"}) + + assert ns.count() == gt["keep"] + assert_get_where_count(ns, {"tag": "purge"}, 0) + assert_get_where_count( + ns, + {"tag": "keep"}, + gt["keep"], + limit=gt["keep"], + check_meta={"tag": "keep"}, + ) + finally: + cleanup(db_client, collection) + + def test_delete_where_document_obsolete(self, db_client): + collection = create_ns_collection(db_client, suffix="_del_wdoc") + ns = collection.create_namespace("del_wdoc_ns") + corpus, gt = build_filter_corpus() + where_document = {"$contains": "obsolete"} + try: + load_filter_corpus(ns, corpus) + pre = precheck_where_document_hits( + ns, + where_document, + gt["doc_obsolete"], + context="pre-delete FTS probe", + ) + ns.delete(where_document=where_document) + + assert_count_after_document_delete( + ns, + where_document=where_document, + expected_count=gt["keep"], + total_before=gt["total"], + expected_removed=gt["doc_obsolete"], + pre_delete_hit_count=len(pre["ids"]), + ) + assert_get_where_document_count( + ns, + where_document, + 0, + limit=10, + substring="obsolete", + context="post-delete", + ) + assert_get_where_count(ns, {"tag": "keep"}, gt["keep"], limit=50) + finally: + # cleanup(db_client, collection) + pass + + +class TestNamespaceGetWhere: + """P0 #3–#4: conditional get on namespace.""" + + def test_get_where_metadata_category_ai(self, db_client): + collection = create_ns_collection(db_client, suffix="_get_where") + ns = collection.create_namespace("get_where_ns") + corpus, gt = build_filter_corpus() + try: + load_dml_corpus(ns, corpus) + + assert_get_where_count( + ns, + {"category": "AI"}, + gt["category_ai"], + limit=gt["category_ai"], + check_meta={"category": "AI"}, + ) + # Capped limit returns fewer rows but all still match. + capped = assert_get_where_count( + ns, + {"category": "AI"}, + 10, + limit=10, + check_meta={"category": "AI"}, + ) + assert len(capped["ids"]) == 10 + finally: + cleanup(db_client, collection) + + def test_get_where_document_contains_python(self, db_client): + collection = create_ns_collection(db_client, suffix="_get_wdoc") + ns = collection.create_namespace("get_wdoc_ns") + corpus, gt = build_filter_corpus() + where_document = {"$contains": "python"} + try: + load_filter_corpus(ns, corpus) + + assert_get_where_document_count( + ns, + where_document, + gt["doc_python"], + limit=gt["doc_python"], + substring="python", + context="full match", + ) + assert_get_where_document_count( + ns, + where_document, + 5, + limit=5, + substring="python", + context="capped limit", + ) + finally: + cleanup(db_client, collection) + + +class TestNamespaceFilterMultiNs: + """P0 #5: filter delete isolated per namespace (same collection).""" + + def test_delete_where_only_affects_target_namespace(self, db_client): + collection = create_ns_collection(db_client, suffix="_filt_mns") + ns_a = collection.create_namespace("ns_alpha") + ns_b = collection.create_namespace("ns_beta") + corpus_a, gt = build_filter_corpus( + keep_count=FILTER_MULTI_NS_KEEP, + purge_count=FILTER_MULTI_NS_PURGE, + ai_in_keep=200, + python_in_keep=80, + id_prefix="alpha", + ns_tag="alpha", + ) + corpus_b, _ = build_filter_corpus( + keep_count=FILTER_MULTI_NS_KEEP, + purge_count=FILTER_MULTI_NS_PURGE, + ai_in_keep=200, + python_in_keep=80, + id_prefix="beta", + ns_tag="beta", + ) + try: + load_dml_corpus(ns_a, corpus_a) + load_dml_corpus(ns_b, corpus_b) + assert ns_a.count() == gt["total"] + assert ns_b.count() == gt["total"] + + ns_a.delete(where={"tag": "purge"}) + + assert ns_a.count() == gt["keep"] + assert ns_b.count() == gt["total"] + + assert_get_where_count(ns_a, {"tag": "purge"}, 0) + assert_get_where_count( + ns_b, + {"tag": "purge"}, + gt["purge"], + limit=10, + check_meta={"tag": "purge", "ns_tag": "beta"}, + ) + assert_get_where_count( + ns_a, + {"tag": "keep"}, + min(20, gt["keep"]), + limit=20, + check_meta={"tag": "keep", "ns_tag": "alpha"}, + ) + finally: + cleanup(db_client, collection) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_combined.py b/tests/integration_tests/test_namespace_hybrid_search_combined.py index 3391da1c..4c2389cf 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_combined.py +++ b/tests/integration_tests/test_namespace_hybrid_search_combined.py @@ -21,7 +21,6 @@ run_hybrid_combined_case, setup_fts_namespace_with_corpus, setup_large_fts_collection, - skip_oceanbase_knn, teardown_large_fts_collection, ) @@ -57,10 +56,8 @@ def _new_namespace(self, case_name: str) -> Any: ) @pytest.mark.parametrize("case_name", [c.name for c in HYBRID_COMBINED_CASES]) - def test_hybrid_search_combined(self, db_client, request, case_name: str): + def test_hybrid_search_combined(self, db_client, case_name: str): case = get_hybrid_combined_case(case_name) - if case.knn is not None: - skip_oceanbase_knn(request) namespace = self._new_namespace(case_name) run_hybrid_combined_case(namespace, self._corpus, case) diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py index 32755516..cc2525c5 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py @@ -1,8 +1,9 @@ """ -Namespace hybrid_search full-text tests: 2 collections x 2 namespaces (use_namespace=True). +Namespace hybrid_search full-text tests: 2 collections x 2 loaded namespaces +plus one empty namespace per collection (use_namespace=True). -Verifies full-text correctness in each quadrant (same assertions as single-namespace tests) -and cross-quadrant isolation (empty namespaces must not leak hits). +Verifies full-text correctness in each loaded namespace, cross-namespace data +isolation, and empty namespace isolation. Run one case in isolation, e.g.:: @@ -12,44 +13,111 @@ from __future__ import annotations +import time + import pytest from namespace_fts_helpers import ( + CORPUS_SIZE, MULTI_COLL_MULTI_NS_QUADRANT_KEYS, TOKEN_ZPX, + CorpusRecord, assert_hybrid_search_no_hits, assert_not_contains_no_token_leak, + build_large_fts_corpus, get_fts_case, + insert_corpus_in_batches, + ns_schema, run_hybrid_search_fts_case, - run_hybrid_search_fts_case_all_quadrants, - setup_multi_coll_multi_ns_fts, - setup_multi_coll_multi_ns_fts_single_loaded, teardown_multi_coll_multi_ns_fts, ) +_EMPTY_NAMESPACE_KEYS: tuple[str, ...] = ("c1_empty", "c2_empty") +_INDEX_SETTLE_SECONDS = 3 + + +def _namespace_corpus(base_corpus: list[CorpusRecord], key: str) -> list[CorpusRecord]: + """Give each namespace a distinct id/document/metadata marker.""" + return [ + CorpusRecord( + doc_id=f"{key}_{record.doc_id}", + document=f"[{key}] {record.document}", + embedding=record.embedding, + metadata={**record.metadata, "namespace_marker": key}, + rel_hint=record.rel_hint, + ) + for record in base_corpus + ] + + +def _setup_multi_coll_multi_ns_isolation_fts(db_client): + """ + 2 collections x 2 loaded namespaces, plus one empty namespace per collection. + + Every loaded namespace receives its own marked corpus, so a leaked row from a + sibling namespace fails the existing corpus-based assertions. + """ + ts = int(time.time() * 1000) + coll_1 = db_client.create_collection( + name=f"test_ns_hs_ft_mcmn_iso_{ts}_c1", + schema=ns_schema(), + use_namespace=True, + ) + coll_2 = db_client.create_collection( + name=f"test_ns_hs_ft_mcmn_iso_{ts}_c2", + schema=ns_schema(), + use_namespace=True, + ) + ctx: dict[str, object] = {"coll_1": coll_1, "coll_2": coll_2} + for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): + for ns_suffix in ("x", "y"): + ctx[f"{coll_tag}_{ns_suffix}"] = collection.create_namespace(f"ns_{ns_suffix}") + ctx[f"{coll_tag}_empty"] = collection.create_namespace("ns_empty") + + base_corpus = build_large_fts_corpus(CORPUS_SIZE) + if len(base_corpus) <= 1000: + raise ValueError(f"corpus must exceed 1000 rows, got {len(base_corpus)}") + + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + corpus = _namespace_corpus(base_corpus, key) + namespace = ctx[key] + insert_corpus_in_batches(namespace, corpus) + ctx[key] = (corpus, namespace) + + time.sleep(_INDEX_SETTLE_SECONDS) + return ctx + class TestNamespaceHybridSearchFulltextMultiCollMultiNs: - """Large-scale hybrid_search FTS across 2 collections x 2 namespaces.""" + """Large-scale hybrid_search FTS across 2 collections x 2 loaded namespaces.""" def _run_fts_case_all_quadrants(self, db_client, case_name: str) -> None: - ctx = setup_multi_coll_multi_ns_fts(db_client) + ctx = _setup_multi_coll_multi_ns_isolation_fts(db_client) try: - run_hybrid_search_fts_case_all_quadrants(ctx, case_name) + case = get_fts_case(case_name) + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + corpus, namespace = ctx[key] + run_hybrid_search_fts_case(namespace, corpus, case) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) + # pass def test_cross_quadrant_fts_isolation(self, db_client): - """Only ``c1_x`` has corpus; other quadrants must return no TOKEN_ZPX hits.""" - ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") + """Loaded namespaces must see only their own rows; empty namespaces return no hits.""" + ctx = _setup_multi_coll_multi_ns_isolation_fts(db_client) try: case = get_fts_case("contains_token_zpx") - corpus, ns_loaded = ctx["c1_x"] - run_hybrid_search_fts_case(ns_loaded, corpus, case) - for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: - if key == "c1_x": - continue - _, namespace = ctx[key] + corpus, namespace = ctx[key] + result = run_hybrid_search_fts_case(namespace, corpus, case) + ids = result["ids"][0] + assert ids + assert all(doc_id.startswith(f"{key}_") for doc_id in ids), ( + f"[{key}] hybrid_search leaked rows from another namespace: {ids[:10]!r}" + ) + + for key in _EMPTY_NAMESPACE_KEYS: + namespace = ctx[key] assert_hybrid_search_no_hits( namespace, case.where_document, @@ -69,7 +137,7 @@ def test_hybrid_search_fulltext_contains_string_shorthand(self, db_client): def test_hybrid_search_fulltext_not_contains(self, db_client): case = get_fts_case("not_contains_token_zpx") - ctx = setup_multi_coll_multi_ns_fts(db_client) + ctx = _setup_multi_coll_multi_ns_isolation_fts(db_client) try: for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: corpus, namespace = ctx[key] @@ -88,7 +156,7 @@ def test_hybrid_search_fulltext_or_zpx_alp(self, db_client): self._run_fts_case_all_quadrants(db_client, "or_zpx_alp") def test_hybrid_search_fulltext_top1_contains_zpx(self, db_client): - ctx = setup_multi_coll_multi_ns_fts(db_client) + ctx = _setup_multi_coll_multi_ns_isolation_fts(db_client) try: for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: _, namespace = ctx[key] @@ -97,7 +165,7 @@ def test_hybrid_search_fulltext_top1_contains_zpx(self, db_client): n_results=1, include=["documents"], ) - assert top_result["ids"][0][0] == "zpx_top_5", ( + assert top_result["ids"][0][0] == f"{key}_zpx_top_5", ( f"[{key}] most relevant TOKEN_ZPX document must rank first, " f"got {top_result['ids'][0]}" ) diff --git a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py index 73e9d16a..fc02a20f 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py @@ -26,7 +26,6 @@ run_search_index_case_on_quadrants, setup_multi_coll_multi_ns_fts, setup_multi_coll_multi_ns_fts_single_loaded, - skip_oceanbase_knn, teardown_multi_coll_multi_ns_fts, ) @@ -102,12 +101,10 @@ def test_hybrid_search_search_index_all_loaded_quadrants(self, db_client, case_n teardown_multi_coll_multi_ns_fts(db_client, ctx) @pytest.mark.parametrize("case_name", ["knn_global_top5", "knn_filter_has_both"]) - def test_hybrid_search_vector_all_loaded_quadrants( - self, db_client, request, case_name: str - ): + def test_hybrid_search_vector_all_loaded_quadrants(self, db_client, case_name: str): ctx = setup_multi_coll_multi_ns_fts(db_client) try: - run_knn_case_on_quadrants(ctx, case_name, request=request) + run_knn_case_on_quadrants(ctx, case_name) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) @@ -122,8 +119,7 @@ def test_hybrid_search_fts_plus_search_index_multi_coll(self, db_client): finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) - def test_hybrid_search_combined_rrf_multi_coll(self, db_client, request): - skip_oceanbase_knn(request) + def test_hybrid_search_combined_rrf_multi_coll(self, db_client): ctx = setup_multi_coll_multi_ns_fts(db_client) try: case = get_hybrid_combined_case("fts_zpx_filter_gte_knn") @@ -133,8 +129,7 @@ def test_hybrid_search_combined_rrf_multi_coll(self, db_client, request): finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) - def test_hybrid_search_vector_plus_search_index_multi_coll(self, db_client, request): - skip_oceanbase_knn(request) + def test_hybrid_search_vector_plus_search_index_multi_coll(self, db_client): ctx = setup_multi_coll_multi_ns_fts(db_client) try: knn_case = get_vector_knn_case("knn_filter_has_both") diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py new file mode 100644 index 00000000..ddde927e --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py @@ -0,0 +1,97 @@ +""" +Namespace hybrid_search triple-branch integration tests. + +Each case activates all three signals in one ``hybrid_search`` call: + - full-text via ``query.where_document`` ($contains / $not_contains / $and / $or) + - search index via ``query.where`` ($eq / $ne / $lt / $lte / $gt / $gte / $in / $nin / $and / $or / $not / #id) + - vector via ``knn`` (KNN distance ordering + optional ``knn.where``) + +Operator coverage is orthogonal: one branch is verified against corpus ground truth while +the other two stay active with broad or aligned filters. + +Run one case:: + + pytest tests/integration_tests/test_namespace_hybrid_search_triple_branch.py \\ + -k "fts_contains_zpx and oceanbase" -v -s +""" + +from __future__ import annotations + +import time +from typing import Any, ClassVar + +import pytest + +from namespace_hybrid_search_helpers import ( + TRIPLE_BRANCH_CASES, + TOKEN_ZPX, + get_triple_branch_case, + run_hybrid_triple_branch_case, + setup_fts_namespace_with_corpus, + setup_large_fts_collection, + teardown_large_fts_collection, +) + + +class TestNamespaceHybridSearchTripleBranch: + """Large-scale namespace vector + full-text + search-index hybrid_search tests.""" + + _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} + + @pytest.fixture(autouse=True) + def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: + mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" + if mode not in self._shared_by_mode: + corpus, collection = setup_large_fts_collection(db_client) + self._shared_by_mode[mode] = { + "db_client": db_client, + "corpus": corpus, + "collection": collection, + } + entry = self._shared_by_mode[mode] + self._corpus = entry["corpus"] + self._collection = entry["collection"] + + @classmethod + def teardown_class(cls) -> None: + for entry in cls._shared_by_mode.values(): + teardown_large_fts_collection(entry["db_client"], entry["collection"]) + cls._shared_by_mode.clear() + + def _new_namespace(self, case_name: str) -> Any: + return setup_fts_namespace_with_corpus( + self._collection, + self._corpus, + namespace_name=f"ns_tb_{case_name}_{int(time.time() * 1000)}", + ) + + @pytest.mark.parametrize("case_name", [c.name for c in TRIPLE_BRANCH_CASES]) + def test_hybrid_search_triple_branch(self, db_client, case_name: str): + case = get_triple_branch_case(case_name) + namespace = self._new_namespace(case_name) + run_hybrid_triple_branch_case(namespace, self._corpus, case) + + def test_hybrid_search_triple_branch_top1_fts_with_knn_active(self, db_client): + """FTS top-1 ranking remains correct when KNN and search-index branches are active.""" + namespace = self._new_namespace("top1_fts_knn_si") + top_result = namespace.hybrid_search( + query={ + "where_document": {"$contains": TOKEN_ZPX}, + "where": {"rel_hint": {"$gte": 0}}, + "n_results": 1, + }, + knn={ + "query_embeddings": [1.0, 1.0, 0.0], + "where": {"rel_hint": {"$gte": 0}}, + "n_results": 20, + }, + n_results=1, + include=["documents"], + ) + assert top_result["ids"][0][0] == "zpx_top_5", ( + f"most relevant TOKEN_ZPX document must rank first, got {top_result['ids'][0]}" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py new file mode 100644 index 00000000..b57e5960 --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py @@ -0,0 +1,230 @@ +""" +Namespace triple-branch hybrid_search: 2 collections x 2 loaded namespaces. + +Extends ``test_namespace_hybrid_search_triple_branch.py`` with multi-collection / +multi-namespace isolation and cross-quadrant correctness checks. + +Run one case:: + + pytest tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py \\ + -k "intersection_both_tokens and oceanbase" -v -s +""" + +from __future__ import annotations + +import time + +import pytest + +from namespace_fts_helpers import ( + CORPUS_SIZE, + MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, + MULTI_COLL_MULTI_NS_QUADRANT_KEYS, + TOKEN_ZPX, + CorpusRecord, + build_large_fts_corpus, + insert_corpus_in_batches, + ns_schema, +) +from namespace_hybrid_search_helpers import ( + TRIPLE_BRANCH_CASES, + assert_hybrid_search_index_no_hits, + assert_hybrid_triple_branch_no_hits, + corpus_matches_triple_intersection, + get_triple_branch_case, + run_hybrid_triple_branch_case, + run_triple_branch_case_on_quadrants, + setup_multi_coll_multi_ns_fts, + setup_multi_coll_multi_ns_fts_single_loaded, + teardown_multi_coll_multi_ns_fts, +) + +_EMPTY_NAMESPACE_KEYS: tuple[str, ...] = ("c1_empty", "c2_empty") +_INDEX_SETTLE_SECONDS = 3 + +_FTS_TRIPLE_CASE_NAMES: tuple[str, ...] = tuple( + c.name for c in TRIPLE_BRANCH_CASES if c.verify in ("fts", "not_contains") +) +_SI_TRIPLE_CASE_NAMES: tuple[str, ...] = tuple( + c.name for c in TRIPLE_BRANCH_CASES if c.verify == "search_index" +) +_KNN_TRIPLE_CASE_NAMES: tuple[str, ...] = tuple( + c.name for c in TRIPLE_BRANCH_CASES if c.verify == "knn" +) +_INTERSECTION_CASE_NAMES: tuple[str, ...] = tuple( + c.name for c in TRIPLE_BRANCH_CASES if c.verify == "intersection" +) +_RRF_CASE_NAMES: tuple[str, ...] = tuple( + c.name for c in TRIPLE_BRANCH_CASES if c.verify == "rrf_fusion" +) + + +def _namespace_corpus(base_corpus: list[CorpusRecord], key: str) -> list[CorpusRecord]: + return [ + CorpusRecord( + doc_id=f"{key}_{record.doc_id}", + document=f"[{key}] {record.document}", + embedding=record.embedding, + metadata={**record.metadata, "namespace_marker": key}, + rel_hint=record.rel_hint, + ) + for record in base_corpus + ] + + +def _setup_multi_coll_multi_ns_isolation_triple(db_client): + """2 collections x 2 loaded namespaces with distinct per-namespace corpus markers.""" + ts = int(time.time() * 1000) + coll_1 = db_client.create_collection( + name=f"test_ns_hs_tb_mcmn_iso_{ts}_c1", + schema=ns_schema(), + use_namespace=True, + ) + coll_2 = db_client.create_collection( + name=f"test_ns_hs_tb_mcmn_iso_{ts}_c2", + schema=ns_schema(), + use_namespace=True, + ) + ctx: dict[str, object] = {"coll_1": coll_1, "coll_2": coll_2} + for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): + for ns_suffix in ("x", "y"): + ctx[f"{coll_tag}_{ns_suffix}"] = collection.create_namespace(f"ns_{ns_suffix}") + ctx[f"{coll_tag}_empty"] = collection.create_namespace("ns_empty") + + base_corpus = build_large_fts_corpus(CORPUS_SIZE) + if len(base_corpus) <= 1000: + raise ValueError(f"corpus must exceed 1000 rows, got {len(base_corpus)}") + + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + corpus = _namespace_corpus(base_corpus, key) + namespace = ctx[key] + insert_corpus_in_batches(namespace, corpus) + ctx[key] = (corpus, namespace) + + time.sleep(_INDEX_SETTLE_SECONDS) + return ctx + + +class TestNamespaceHybridSearchTripleBranchMultiCollMultiNs: + def test_cross_quadrant_triple_branch_isolation(self, db_client): + """Loaded namespaces return only their own rows; empty namespaces stay empty.""" + ctx = _setup_multi_coll_multi_ns_isolation_triple(db_client) + try: + case = get_triple_branch_case("intersection_both_tokens") + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + corpus, namespace = ctx[key] + result = run_hybrid_triple_branch_case(namespace, corpus, case) + ids = result["ids"][0] + assert ids + assert all(doc_id.startswith(f"{key}_") for doc_id in ids), ( + f"[{key}] triple-branch hybrid_search leaked rows: {ids[:10]!r}" + ) + + for key in _EMPTY_NAMESPACE_KEYS: + namespace = ctx[key] + assert_hybrid_triple_branch_no_hits(namespace, case) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_cross_quadrant_search_index_branch_isolation(self, db_client): + ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") + try: + case = get_triple_branch_case("si_and_has_both_zpx_gte") + corpus, ns_loaded = ctx["c1_x"] + run_hybrid_triple_branch_case(ns_loaded, corpus, case) + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + if key == "c1_x": + continue + _, namespace = ctx[key] + assert_hybrid_search_index_no_hits(namespace, case.where, n_results=case.n_results) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_triple_branch_top1_fts_all_quadrants(self, db_client): + ctx = _setup_multi_coll_multi_ns_isolation_triple(db_client) + try: + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + _, namespace = ctx[key] + top_result = namespace.hybrid_search( + query={ + "where_document": {"$contains": TOKEN_ZPX}, + "where": {"rel_hint": {"$gte": 0}}, + "n_results": 1, + }, + knn={ + "query_embeddings": [1.0, 1.0, 0.0], + "where": {"rel_hint": {"$gte": 0}}, + "n_results": 20, + }, + n_results=1, + include=["documents"], + ) + assert top_result["ids"][0][0] == f"{key}_zpx_top_5", ( + f"[{key}] FTS top-1 must be zpx_top_5, got {top_result['ids'][0]}" + ) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + @pytest.mark.parametrize("case_name", _FTS_TRIPLE_CASE_NAMES) + def test_hybrid_search_triple_branch_fts_all_loaded_quadrants(self, db_client, case_name: str): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_triple_branch_case_on_quadrants(ctx, case_name) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + @pytest.mark.parametrize("case_name", _SI_TRIPLE_CASE_NAMES) + def test_hybrid_search_triple_branch_search_index_all_loaded_quadrants( + self, db_client, case_name: str + ): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_triple_branch_case_on_quadrants(ctx, case_name) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + @pytest.mark.parametrize("case_name", _KNN_TRIPLE_CASE_NAMES) + def test_hybrid_search_triple_branch_knn_all_loaded_quadrants(self, db_client, case_name: str): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_triple_branch_case_on_quadrants(ctx, case_name) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + @pytest.mark.parametrize("case_name", _INTERSECTION_CASE_NAMES) + def test_hybrid_search_triple_branch_intersection_all_loaded_quadrants( + self, db_client, case_name: str + ): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_triple_branch_case_on_quadrants(ctx, case_name) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + @pytest.mark.parametrize("case_name", _RRF_CASE_NAMES) + def test_hybrid_search_triple_branch_rrf_all_loaded_quadrants(self, db_client, case_name: str): + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + run_triple_branch_case_on_quadrants(ctx, case_name) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + def test_hybrid_search_triple_branch_intersection_subset_on_quadrants(self, db_client): + """Every hit must lie in the triple-branch intersection on each loaded quadrant.""" + ctx = setup_multi_coll_multi_ns_fts(db_client) + try: + case = get_triple_branch_case("intersection_zpx_gte40") + for key in MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS: + corpus, namespace = ctx[key] + result = run_hybrid_triple_branch_case(namespace, corpus, case) + corpus_by_id = {rec.doc_id: rec for rec in corpus} + for doc_id in result["ids"][0]: + assert corpus_matches_triple_intersection(corpus_by_id[doc_id], case), ( + f"[{key}] id={doc_id!r} outside intersection" + ) + finally: + teardown_multi_coll_multi_ns_fts(db_client, ctx) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_vector.py b/tests/integration_tests/test_namespace_hybrid_search_vector.py index fd45921f..69d54d3b 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_vector.py +++ b/tests/integration_tests/test_namespace_hybrid_search_vector.py @@ -18,7 +18,6 @@ run_hybrid_knn_case, setup_fts_namespace_with_corpus, setup_large_fts_collection, - skip_oceanbase_knn, teardown_large_fts_collection, ) @@ -56,18 +55,16 @@ def _new_namespace(self, case_name: str) -> Any: namespace_name=ns_name, ) - def _run_knn_case(self, request: pytest.FixtureRequest, case_name: str) -> None: - skip_oceanbase_knn(request) + def _run_knn_case(self, case_name: str) -> None: namespace = self._new_namespace(case_name) run_hybrid_knn_case(namespace, self._corpus, get_vector_knn_case(case_name)) @pytest.mark.parametrize("case_name", [c.name for c in VECTOR_KNN_CASES]) - def test_hybrid_search_vector_knn_cases(self, db_client, request, case_name: str): - self._run_knn_case(request, case_name) + def test_hybrid_search_vector_knn_cases(self, db_client, case_name: str): + self._run_knn_case(case_name) - def test_hybrid_search_vector_top1_nearest(self, db_client, request): + def test_hybrid_search_vector_top1_nearest(self, db_client): """Top-1 must be the global nearest neighbor for a fixed query vector.""" - skip_oceanbase_knn(request) namespace = self._new_namespace("top1_nearest") from namespace_hybrid_search_helpers import expected_knn_ids diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 3049e856..9372ee94 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -20,7 +20,7 @@ def _create_ns_collection(self, client, suffix=""): name = f"test_ns_lc_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine"), + ivf=IVFConfiguration(dimension=3, distance="cosine", use_spfresh=True), embedding_function=None, ), ) @@ -172,7 +172,7 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): name = f"test_ns_ss_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine"), + ivf=IVFConfiguration(dimension=3, distance="cosine", use_spfresh=True), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index 4c2f5e3c..636f31e4 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -30,7 +30,7 @@ def _create_ns_collection_and_namespace(self, client): name = f"test_ns_pw_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine"), + ivf=IVFConfiguration(dimension=3, distance="cosine", use_spfresh=True), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index ce942c44..95362d0f 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -18,7 +18,7 @@ def _setup(self, client, suffix=""): name = f"test_ns_q_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2"), + ivf=IVFConfiguration(dimension=3, distance="l2", use_spfresh=True), embedding_function=None, ), ) @@ -316,15 +316,8 @@ def test_hybrid_search_with_data_content_filter(self, db_client): finally: db_client.delete_collection(name=collection.name) - def test_hybrid_search_with_data_content_filter_knn_branch(self, db_client, request): - """ - KNN + metadata / id filters on ``data_content`` (same field mapping as full-text branch). - - OceanBase ``hybrid_search`` KNN currently requires an HNSW vector index; namespace - logical tables are created with IVF only, so this branch is skipped for ``[oceanbase]``. - """ - if "oceanbase" in request.node.nodeid: - pytest.skip("OceanBase hybrid_search KNN requires HNSW; namespace collections use IVF only.") + def test_hybrid_search_with_data_content_filter_knn_branch(self, db_client): + """KNN + metadata / id filters on ``data_content`` (same field mapping as full-text branch).""" collection = self._setup(db_client) ns = collection.create_namespace("qns_hs_dc_knn") try: diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index b379b836..828b74de 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -580,6 +580,18 @@ def test_delete_by_where_sql(self): assert "JSON_EXTRACT" in sql or "JSON_OVERLAPS" in sql assert "metadata.category" in sql + def test_delete_by_where_document_sql(self): + c = self._client() + c._namespace_delete( + **self._common_kwargs(), + where_document={"$contains": "obsolete"}, + ) + sql = c.executed_sqls[-1] + assert "DELETE FROM" in sql + assert "namespace_id = 7" in sql + assert "document LIKE" in sql + assert "MATCH(document)" not in sql + # ---- QUERY ---- def test_query_basic_sql(self): From 65bf56c5cea48f5a9dd16e7c8f4cac1ec8ea4db5 Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Wed, 3 Jun 2026 10:47:57 +0800 Subject: [PATCH 24/84] open spfresh and add case --- tests/integration_tests/test_namespace_session_vars.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py index bc7081ad..bd0af8b9 100644 --- a/tests/integration_tests/test_namespace_session_vars.py +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -20,7 +20,7 @@ def _create_ns_collection(self, client): name = f"test_ns_sessvar_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2"), + ivf=IVFConfiguration(dimension=3, distance="l2", use_spfresh=True), embedding_function=None, ), ) From c18230da14be076d59f47492622bd9fe9cdbe8df Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Wed, 3 Jun 2026 11:40:26 +0800 Subject: [PATCH 25/84] fix case --- .../namespace_dml_helpers.py | 2 +- .../namespace_fts_helpers.py | 2 +- .../namespace_hybrid_search_helpers.py | 4 +-- .../test_namespace_drop_validation.py | 2 +- ...t_namespace_hybrid_search_triple_branch.py | 28 +++++++++++-------- .../test_namespace_lifecycle.py | 4 +-- .../test_namespace_prewarm.py | 2 +- .../integration_tests/test_namespace_query.py | 2 +- .../test_namespace_session_vars.py | 2 +- 9 files changed, 27 insertions(+), 21 deletions(-) diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index 7131e4b3..dc1051e6 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -37,7 +37,7 @@ class DmlRecord: def ns_schema() -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", use_spfresh=True), + ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index 4728ce31..141376cb 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -52,7 +52,7 @@ class FtsQueryCase: def ns_schema() -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", use_spfresh=True), + ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/namespace_hybrid_search_helpers.py b/tests/integration_tests/namespace_hybrid_search_helpers.py index aa4332d4..8430fd1d 100644 --- a/tests/integration_tests/namespace_hybrid_search_helpers.py +++ b/tests/integration_tests/namespace_hybrid_search_helpers.py @@ -574,7 +574,8 @@ class HybridTripleBranchCase: where: dict[str, Any] | None = None knn: dict[str, Any] | None = None use_rrf: bool = False - check_fts_ranking: bool = True + # Triple-branch always runs KNN; fused __score order need not match pure FTS tiers. + check_fts_ranking: bool = False min_hits: int = 1 exact_match_count: int | None = None @@ -1061,7 +1062,6 @@ def run_hybrid_triple_branch_case( where=WHERE_FILLER_SEQ, knn=_default_knn(20, where=WHERE_FILLER_SEQ), n_results=20, - check_fts_ranking=False, ), HybridTripleBranchCase( name="fts_and_zpx_alp", diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index 9aae330a..90821a3c 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -24,7 +24,7 @@ def _make_collection(client, suffix: str = ""): name = f"test_ns_drop_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", use_spfresh=True), + ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py index ddde927e..3a64cd9c 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py @@ -7,7 +7,9 @@ - vector via ``knn`` (KNN distance ordering + optional ``knn.where``) Operator coverage is orthogonal: one branch is verified against corpus ground truth while -the other two stay active with broad or aligned filters. +the other two stay active with broad or aligned filters. For ``verify="fts"``, hits and +filters are checked; fused FTS+KNN ``__score`` order is not asserted (see +``check_fts_ranking`` on ``HybridTripleBranchCase``). Run one case:: @@ -71,26 +73,30 @@ def test_hybrid_search_triple_branch(self, db_client, case_name: str): namespace = self._new_namespace(case_name) run_hybrid_triple_branch_case(namespace, self._corpus, case) - def test_hybrid_search_triple_branch_top1_fts_with_knn_active(self, db_client): - """FTS top-1 ranking remains correct when KNN and search-index branches are active.""" - namespace = self._new_namespace("top1_fts_knn_si") - top_result = namespace.hybrid_search( + def test_hybrid_search_triple_branch_fts_hits_with_knn_active(self, db_client): + """FTS+KNN hybrid returns TOKEN_ZPX hits; fused order is not asserted.""" + namespace = self._new_namespace("fts_hits_knn_si") + result = namespace.hybrid_search( query={ "where_document": {"$contains": TOKEN_ZPX}, "where": {"rel_hint": {"$gte": 0}}, - "n_results": 1, + "n_results": 5, }, knn={ "query_embeddings": [1.0, 1.0, 0.0], "where": {"rel_hint": {"$gte": 0}}, "n_results": 20, }, - n_results=1, - include=["documents"], - ) - assert top_result["ids"][0][0] == "zpx_top_5", ( - f"most relevant TOKEN_ZPX document must rank first, got {top_result['ids'][0]}" + n_results=5, + include=["documents", "metadatas"], ) + ids = result["ids"][0] + assert len(ids) > 0, "expected at least one hybrid FTS+KNN hit" + docs = result.get("documents", [[]])[0] + for doc_id, doc_text in zip(ids, docs): + assert TOKEN_ZPX.lower() in (doc_text or "").lower(), ( + f"id={doc_id!r} must contain {TOKEN_ZPX!r}" + ) if __name__ == "__main__": diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 9372ee94..16fb61d9 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -20,7 +20,7 @@ def _create_ns_collection(self, client, suffix=""): name = f"test_ns_lc_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", use_spfresh=True), + ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), embedding_function=None, ), ) @@ -172,7 +172,7 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): name = f"test_ns_ss_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", use_spfresh=True), + ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index 636f31e4..32f0e052 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -30,7 +30,7 @@ def _create_ns_collection_and_namespace(self, client): name = f"test_ns_pw_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", use_spfresh=True), + ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index 95362d0f..6916c686 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -18,7 +18,7 @@ def _setup(self, client, suffix=""): name = f"test_ns_q_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", use_spfresh=True), + ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py index bd0af8b9..83983a50 100644 --- a/tests/integration_tests/test_namespace_session_vars.py +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -20,7 +20,7 @@ def _create_ns_collection(self, client): name = f"test_ns_sessvar_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", use_spfresh=True), + ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), embedding_function=None, ), ) From 17f89a1ff5dca641f62b172d02e1387317646795 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 3 Jun 2026 11:57:51 +0800 Subject: [PATCH 26/84] fix(hybrid_search): hoist metadata $ne/$nin must_not for FTS and KNN filters OceanBase rejects bool nodes with only must_not when they are nested inside filter arrays. Hoist negation leaves to the outer bool for FTS+where and wrap KNN-only negation with a permissive range/exists positive clause (match_all is unsupported). --- src/pyseekdb/client/client_base.py | 112 +++++++++++++----- tests/unit_tests/test_build_knn_filter.py | 34 ++++++ .../unit_tests/test_build_query_expression.py | 47 +++++++- 3 files changed, 160 insertions(+), 33 deletions(-) create mode 100644 tests/unit_tests/test_build_knn_filter.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 4e729e5c..d9a7bbb0 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -3957,6 +3957,42 @@ def _pure_must_not_clauses(expr: dict[str, Any] | None) -> list[dict[str, Any]] return None return clauses + @staticmethod + def _positive_clause_for_must_not(must_not_clauses: list[dict[str, Any]]) -> dict[str, Any]: + """Build a permissive positive filter leaf for must_not-only bools (OB rejects match_all).""" + for clause in must_not_clauses: + if not isinstance(clause, dict): + continue + term_body = clause.get("term") + if isinstance(term_body, dict) and term_body: + field = next(iter(term_body)) + return {"range": {field: {"gte": -9223372036854775808}}} + terms_body = clause.get("terms") + if isinstance(terms_body, dict) and terms_body: + field = next(iter(terms_body)) + return {"range": {field: {"gte": -9223372036854775808}}} + qs_body = clause.get("query_string") + if isinstance(qs_body, dict): + fields = qs_body.get("fields") or ["document"] + field = fields[0] if fields else "document" + return {"exists": {"field": field}} + return {"exists": {"field": "document"}} + + @staticmethod + def _hoist_must_not_from_filters( + filter_conditions: list[dict[str, Any]], + ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Split a filter clause list into positive filters and hoisted must_not leaves.""" + positive: list[dict[str, Any]] = [] + negative: list[dict[str, Any]] = [] + for cond in filter_conditions: + clauses = BaseClient._pure_must_not_clauses(cond) + if clauses is not None: + negative.extend(clauses) + else: + positive.append(cond) + return positive, negative + def _build_query_expression(self, query: dict[str, Any]) -> dict[str, Any] | None: """ Build query expression from query dict @@ -3983,18 +4019,7 @@ def _build_query_expression(self, query: dict[str, Any]) -> dict[str, Any] | Non # should have at least one positive clause`, so hoist their must_not # clauses onto the outer bool (which gains a positive `filter` once the # namespace filter is injected) instead of nesting them standalone. - positive: list[dict[str, Any]] = [] - negative: list[dict[str, Any]] = [] - for cond in filter_conditions: - if ( - isinstance(cond, dict) - and set(cond.keys()) == {"bool"} - and isinstance(cond["bool"], dict) - and set(cond["bool"].keys()) == {"must_not"} - ): - negative.extend(cond["bool"]["must_not"]) - else: - positive.append(cond) + positive, negative = self._hoist_must_not_from_filters(filter_conditions) bool_q: dict[str, Any] = {} if positive: bool_q["filter"] = positive @@ -4007,21 +4032,28 @@ def _build_query_expression(self, query: dict[str, Any]) -> dict[str, Any] | Non # Build document query using query_string doc_query = self._build_document_query(where_document, boost=boost) if doc_query: - # Build filter from where condition filter_conditions = self._build_metadata_filter_for_search_parm(where) - negated = self._pure_must_not_clauses(doc_query) - - if filter_conditions: - # $not_contains becomes a pure must_not bool; nesting that under - # must triggers `bool query ... should have at least one positive clause`. - if negated is not None: - return {"bool": {"filter": filter_conditions, "must_not": negated}} - return {"bool": {"must": [doc_query], "filter": filter_conditions}} - if negated is not None: - # No metadata filter yet; match_all satisfies the positive-clause rule - # until namespace filters are injected (collection stays standalone). - return {"bool": {"filter": [{"match_all": {}}], "must_not": negated}} - return doc_query + pos_filters, meta_must_not = self._hoist_must_not_from_filters(filter_conditions) + doc_must_not = self._pure_must_not_clauses(doc_query) + must_not_all = list(meta_must_not) + if doc_must_not is not None: + must_not_all.extend(doc_must_not) + + if not filter_conditions and doc_must_not is None: + return doc_query + + bool_q: dict[str, Any] = {} + if doc_must_not is None: + bool_q["must"] = [doc_query] + if pos_filters: + bool_q["filter"] = pos_filters + if must_not_all: + bool_q["must_not"] = must_not_all + # Pure negation ($not_contains, $ne, …) needs a positive clause unless + # a scoring `must` is already present. + if doc_must_not is not None and not pos_filters and "must" not in bool_q: + bool_q["filter"] = [self._positive_clause_for_must_not(must_not_all)] + return {"bool": bool_q} return None @@ -4187,8 +4219,17 @@ def _build_metadata_filter_conditions( # noqa: C901 # Scalar conditions must be ANDed via a (non-scoring) `filter` # clause, not `must`: the kernel rejects scalar term/range queries # inside must/should with `scalar ... query in must/should clause - # not supported`. - result.append({"bool": {"filter": must_conditions}}) + # not supported`. Hoist must_not-only leaves ($ne/$nin/$not) so they + # are not nested as standalone bools inside `filter`. + positive, negative = self._hoist_must_not_from_filters(must_conditions) + bool_q: dict[str, Any] = {} + if positive: + bool_q["filter"] = positive + elif negative: + bool_q["filter"] = [self._positive_clause_for_must_not(negative)] + if negative: + bool_q["must_not"] = negative + result.append({"bool": bool_q}) return result if "$or" in condition: @@ -4335,14 +4376,23 @@ def _normalize_vectors(raw_embeddings: Any) -> list[list[float]]: # Build knn expressions (one per vector) knn_exprs: list[dict[str, Any]] = [] filter_conditions = self._build_metadata_filter_for_search_parm(where) + pos_filters, must_not_clauses = self._hoist_must_not_from_filters(filter_conditions) + knn_filter: list[dict[str, Any]] | None = None + if must_not_clauses: + bool_filter: dict[str, Any] = { + "must_not": must_not_clauses, + "filter": pos_filters or [self._positive_clause_for_must_not(must_not_clauses)], + } + knn_filter = [{"bool": bool_filter}] + elif pos_filters: + knn_filter = pos_filters for vector in vectors: expr = {"field": "embedding", "k": n_results, "query_vector": vector} if boost is not None: expr["boost"] = boost - # Add filter using JSON_EXTRACT format - if filter_conditions: - expr["filter"] = filter_conditions + if knn_filter is not None: + expr["filter"] = knn_filter knn_exprs.append(expr) diff --git a/tests/unit_tests/test_build_knn_filter.py b/tests/unit_tests/test_build_knn_filter.py new file mode 100644 index 00000000..60e0b868 --- /dev/null +++ b/tests/unit_tests/test_build_knn_filter.py @@ -0,0 +1,34 @@ +"""Unit tests for KNN filter DSL generation with hoisted must_not.""" + +from unittest.mock import patch + +import pytest + + +@pytest.fixture() +def client(): + from pyseekdb.client.client_base import BaseClient + + with patch.multiple(BaseClient, __abstractmethods__=set()): + instance = BaseClient.__new__(BaseClient) + return instance + + +class TestBuildKnnFilterNe: + def test_ne_hoists_must_not_in_knn_filter(self, client): + with patch.object( + client, + "_build_metadata_filter_for_search_parm", + return_value=[{"bool": {"must_not": [{"term": {"data_content.metadata.zpx_hint": 0}}]}}], + ): + knn_expr = client._build_knn_expression( + {"query_embeddings": [1.0, 1.0, 0.0], "where": {"zpx_hint": {"$ne": 0}}, "n_results": 40}, + dimension=3, + ) + + assert knn_expr["filter"] == [{ + "bool": { + "filter": [{"range": {"data_content.metadata.zpx_hint": {"gte": -9223372036854775808}}}], + "must_not": [{"term": {"data_content.metadata.zpx_hint": 0}}], + } + }] diff --git a/tests/unit_tests/test_build_query_expression.py b/tests/unit_tests/test_build_query_expression.py index 5df3bfcb..0980988a 100644 --- a/tests/unit_tests/test_build_query_expression.py +++ b/tests/unit_tests/test_build_query_expression.py @@ -43,7 +43,7 @@ def test_not_contains_with_metadata_filter_hoists_must_not(self, client): } assert "must" not in expr["bool"] - def test_not_contains_only_uses_match_all_filter(self, client): + def test_not_contains_only_uses_exists_positive_filter(self, client): with patch.object(client, "_build_metadata_filter_for_search_parm", return_value=[]): expr = client._build_query_expression({ "where_document": {"$not_contains": "TOKENZPX"}, @@ -51,7 +51,7 @@ def test_not_contains_only_uses_match_all_filter(self, client): assert expr == { "bool": { - "filter": [{"match_all": {}}], + "filter": [{"exists": {"field": "document"}}], "must_not": [ { "query_string": { @@ -77,3 +77,46 @@ def test_contains_with_metadata_filter_still_uses_must(self, client): assert "must" in expr["bool"] assert "query_string" in expr["bool"]["must"][0] assert "filter" in expr["bool"] + + +class TestBuildQueryExpressionMetadataNe: + def test_ne_with_fts_hoists_must_not_from_filter(self, client): + """$ne in where must not appear as a must_not-only bool inside filter.""" + with patch.object( + client, + "_build_metadata_filter_for_search_parm", + return_value=[{"bool": {"must_not": [{"term": {"data_content.metadata.zpx_hint": 0}}]}}], + ): + expr = client._build_query_expression({ + "where_document": {"$contains": "document"}, + "where": {"zpx_hint": {"$ne": 0}}, + }) + + assert expr == { + "bool": { + "must": [{"query_string": {"fields": ["document"], "query": "document"}}], + "must_not": [{"term": {"data_content.metadata.zpx_hint": 0}}], + } + } + filters = expr["bool"].get("filter", []) + for item in filters: + nested = item.get("bool", {}) if isinstance(item, dict) else {} + assert set(nested.keys()) != {"must_not"}, nested + + def test_ne_with_not_contains_and_positive_where(self, client): + with patch.object( + client, + "_build_metadata_filter_for_search_parm", + return_value=[ + {"bool": {"must_not": [{"term": {"data_content.metadata.zpx_hint": 0}}]}}, + {"range": {"data_content.metadata.seq": {"gte": 1}}}, + ], + ): + expr = client._build_query_expression({ + "where_document": {"$not_contains": "TOKENZPX"}, + "where": {"zpx_hint": {"$ne": 0}, "seq": {"$gte": 1}}, + }) + + assert expr["bool"]["filter"] == [{"range": {"data_content.metadata.seq": {"gte": 1}}}] + assert len(expr["bool"]["must_not"]) == 2 + assert "must" not in expr["bool"] From a555bdcdf124fa52094af378408fad5422540914 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 3 Jun 2026 14:37:52 +0800 Subject: [PATCH 27/84] =?UTF-8?q?drop=20namespace=E8=81=94=E8=B0=83sdk?= =?UTF-8?q?=E9=80=82=E9=85=8D=E4=BF=AE=E6=94=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 108 ++++-- src/pyseekdb/client/client_seekdb_server.py | 3 + .../test_namespace_drop_validation.py | 337 ++++++++++++++++-- 3 files changed, 389 insertions(+), 59 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index d9a7bbb0..d619f21c 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1028,7 +1028,9 @@ def _get_embedding_function_dimension(self, embedding_function: EmbeddingFunctio def _create_sdk_collections_if_not_exists(self) -> None: try: - create_table_sql = """CREATE TABLE IF NOT EXISTS sdk_collections ( + self._use_catalog_database() + sdk_coll = self._qtable(CollectionNames.sdk_collections_table_name()) + create_table_sql = f"""CREATE TABLE IF NOT EXISTS {sdk_coll} ( collection_id CHAR(32) PRIMARY KEY DEFAULT (replace(uuid(), '-', '')), collection_name STRING, settings JSON COMMENT "Generated by SDK, don't modify", @@ -1094,8 +1096,65 @@ def _create_collection_meta_v1(self, collection_name: str) -> str: # ==================== Namespace Catalog Methods ==================== + def _catalog_database(self) -> str: + """Database that holds sdk_* catalog tables (must match OB bg scheduler scan target).""" + db = getattr(self, "database", None) + if not db: + raise ValueError("client database is not configured") + return db + + def _qtable(self, table: str) -> str: + """Fully-qualified catalog table: `{database}`.`{table}`.""" + return f"`{self._catalog_database()}`.`{table}`" + + def _use_catalog_database(self) -> None: + """Align session with pymysql database= so PL (DROP_NAMESPACE) uses the same DB.""" + self._execute(f"USE `{self._catalog_database()}`") + + def _isolate_sdk_catalog_to_client_database(self) -> None: + """Drop sdk_* catalog tables in other databases so kernel bg task scans self.database. + + ObLTableBGTaskScheduler picks the first tenant database that has sdk_collections. + mysqltest may leave a copy in the `oceanbase` database while pyseekdb uses `test`. + """ + target = self._catalog_database() + try: + rows = self._execute( + "SELECT DISTINCT table_schema AS table_schema " + "FROM information_schema.tables " + "WHERE table_name = 'sdk_collections'" + ) + except Exception: + return + other_schemas: list[str] = [] + for row in rows or []: + schema = row["table_schema"] if isinstance(row, dict) else row[0] + if schema and schema != target: + other_schemas.append(schema) + if not other_schemas: + return + catalog_tables = [ + CollectionNames.sdk_collections_table_name(), + NamespaceCollectionNames.sdk_namespaces_table(), + NamespaceCollectionNames.sdk_ltables_table(), + NamespaceCollectionNames.sdk_namespaces_stats_table(), + ] + for schema in other_schemas: + for table in catalog_tables: + with contextlib.suppress(Exception): + self._execute(f"DROP TABLE IF EXISTS `{schema}`.`{table}`") + logger.info( + "Removed duplicate sdk catalog tables from %s (catalog DB is %s)", + other_schemas, + target, + ) + def _ensure_namespace_catalogs(self) -> None: - ns_namespaces_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_namespaces_table()}` ( + self._use_catalog_database() + self._isolate_sdk_catalog_to_client_database() + ns_namespaces_q = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) + ns_ltables_q = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) + ns_namespaces_sql = f"""CREATE TABLE IF NOT EXISTS {ns_namespaces_q} ( namespace_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, collection_id CHAR(32) NOT NULL, namespace_name VARCHAR(256) NOT NULL, @@ -1106,7 +1165,7 @@ def _ensure_namespace_catalogs(self) -> None: UNIQUE KEY uk_sdk_ns_coll_name (collection_id, namespace_name), KEY idx_sdk_ns_by_collection (collection_id) ) COMMENT='Namespace catalog';""" - ns_ltables_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_ltables_table()}` ( + ns_ltables_sql = f"""CREATE TABLE IF NOT EXISTS {ns_ltables_q} ( ltable_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, collection_id CHAR(32) NOT NULL, namespace_id BIGINT UNSIGNED NOT NULL, @@ -1118,7 +1177,7 @@ def _ensure_namespace_catalogs(self) -> None: UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name), KEY idx_sdk_lt_by_ns (collection_id, namespace_id) ) COMMENT='LTable catalog';""" - namespaces_stats_sql = f"""CREATE TABLE IF NOT EXISTS `{NamespaceCollectionNames.sdk_namespaces_stats_table()}` ( + namespaces_stats_sql = f"""CREATE TABLE IF NOT EXISTS {self._qtable(NamespaceCollectionNames.sdk_namespaces_stats_table())} ( collection_id CHAR(32) NOT NULL COMMENT 'collection id', namespace_id BIGINT UNSIGNED NOT NULL COMMENT 'namespace internal id', ltable_id BIGINT UNSIGNED NOT NULL COMMENT 'logic table internal id, 0 means namespace summary', @@ -1140,14 +1199,15 @@ def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> di self._create_sdk_collections_if_not_exists() settings_str = escape_string(json.dumps(settings)) collection_name_escaped = escape_string(collection_name) + sdk_coll = self._qtable(CollectionNames.sdk_collections_table_name()) insert_sql = ( - f"INSERT INTO `{CollectionNames.sdk_collections_table_name()}` " + f"INSERT INTO {sdk_coll} " f"(collection_name, settings) " f"VALUES ('{collection_name_escaped}', '{settings_str}')" ) self._execute(insert_sql) rows = self._execute( - f"SELECT collection_id FROM `{CollectionNames.sdk_collections_table_name()}` " + f"SELECT collection_id FROM {sdk_coll} " f"WHERE collection_name = '{collection_name_escaped}'" ) collection_id = str(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["collection_id"]) @@ -1159,7 +1219,7 @@ def _get_ns_collection_meta(self, collection_name: str) -> dict | None: try: rows = self._execute( f"SELECT collection_id, collection_name, settings " - f"FROM `{CollectionNames.sdk_collections_table_name()}` " + f"FROM {self._qtable(CollectionNames.sdk_collections_table_name())} " f"WHERE collection_name = '{collection_name_escaped}'" ) except Exception: @@ -1281,7 +1341,7 @@ def _create_namespace_physical_tables( kv_key VARBINARY(1024) NOT NULL, kv_value LONGBLOB NOT NULL, PRIMARY KEY(namespace_id, kv_key) - ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 + ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 LOB_INROW_THRESHOLD = 786432 {partition_clause}""") self._execute(f"""CREATE TABLE `{schema_table}` ( @@ -1331,7 +1391,7 @@ def _resolve_namespace_ltable_id( return cached coll_id_escaped = escape_string(str(collection_id)) rows = self._execute( - f"SELECT ltable_id FROM `{NamespaceCollectionNames.sdk_ltables_table()}` " + f"SELECT ltable_id FROM {self._qtable(NamespaceCollectionNames.sdk_ltables_table())} " f"WHERE collection_id = '{coll_id_escaped}' AND namespace_id = {int(namespace_id)} " f"AND ltable_name = 'default' LIMIT 1" ) @@ -1371,30 +1431,34 @@ def _set_session_ns_context( def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) + ns_table = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) + lt_table = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) self._execute( - f"INSERT INTO `{NamespaceCollectionNames.sdk_namespaces_table()}` " + f"INSERT INTO {ns_table} " f"(collection_id, namespace_name) VALUES ('{collection_id_escaped}', '{namespace_name_escaped}')" ) rows = self._execute( - f"SELECT namespace_id FROM `{NamespaceCollectionNames.sdk_namespaces_table()}` " + f"SELECT namespace_id FROM {ns_table} " f"WHERE collection_id = '{collection_id_escaped}' AND namespace_name = '{namespace_name_escaped}'" ) ns_id = int(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["namespace_id"]) self._execute( - f"INSERT INTO `{NamespaceCollectionNames.sdk_ltables_table()}` " + f"INSERT INTO {lt_table} " f"(collection_id, namespace_id, ltable_name) VALUES ('{collection_id_escaped}', {ns_id}, 'default')" ) lt_rows = self._execute( - f"SELECT ltable_id FROM `{NamespaceCollectionNames.sdk_ltables_table()}` " + f"SELECT ltable_id FROM {lt_table} " f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {ns_id} AND ltable_name = 'default'" ) lt_id = int(lt_rows[0][0] if isinstance(lt_rows[0], (list, tuple)) else lt_rows[0]["ltable_id"]) self._cache_namespace_ltable_id(collection_id, ns_id, lt_id) - schema_table = NamespaceCollectionNames.logic_schema_table_name(collection_id) + schema_table = self._qtable( + NamespaceCollectionNames.logic_schema_table_name(collection_id) + ) schema_content = json.dumps(_build_default_ltable_schema()) with contextlib.suppress(Exception): self._execute( - f"INSERT INTO `{schema_table}` (namespace_id, ltable_id, schema_content) " + f"INSERT INTO {schema_table} (namespace_id, ltable_id, schema_content) " f"VALUES ({ns_id}, {lt_id}, '{escape_string(schema_content)}')" ) self._set_session_ns_context(namespace_id=ns_id, ltable_id=lt_id) @@ -1403,13 +1467,13 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) - ns_table = NamespaceCollectionNames.sdk_namespaces_table() - lt_table = NamespaceCollectionNames.sdk_ltables_table() + ns_table = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) + lt_table = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) rows = self._execute( f"SELECT n.namespace_id AS namespace_id, n.namespace_name AS namespace_name, " f"l.ltable_id AS ltable_id " - f"FROM `{ns_table}` n " - f"LEFT JOIN `{lt_table}` l " + f"FROM {ns_table} n " + f"LEFT JOIN {lt_table} l " f"ON l.collection_id = n.collection_id " f"AND l.namespace_id = n.namespace_id " f"AND l.ltable_name = 'default' " @@ -1447,8 +1511,8 @@ def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> lt_id_raw = meta.get("ltable_id") lt_id = int(lt_id_raw) if lt_id_raw is not None else None collection_id_escaped = escape_string(collection_id) - # Ensure DBMS_LOGIC_TABLE.DROP_NAMESPACE sees the right session context; - # otherwise it may read stale @collection_id / @ltable_id from a previous call. + # PL reads session database_name; must match where catalog tables live. + self._use_catalog_database() self._set_session_ns_context( collection_id=collection_id, namespace_id=int(ns_id), ltable_id=lt_id ) @@ -1459,7 +1523,7 @@ def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> def _list_ns_namespaces(self, collection_id: str) -> list[dict]: collection_id_escaped = escape_string(collection_id) rows = self._execute( - f"SELECT namespace_id, namespace_name FROM `{NamespaceCollectionNames.sdk_namespaces_table()}` " + f"SELECT namespace_id, namespace_name FROM {self._qtable(NamespaceCollectionNames.sdk_namespaces_table())} " f"WHERE collection_id = '{collection_id_escaped}' ORDER BY namespace_id" ) results = [] diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 31e867fc..75cff466 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -3,6 +3,7 @@ Supports both seekdb Server and OceanBase Server """ +import contextlib import logging from collections.abc import Sequence @@ -79,6 +80,8 @@ def _ensure_connection(self) -> pymysql.Connection: **self.kwargs, ) logger.info(f"✅ Connected to remote server: {self.host}:{self.port}/{self.database}") + with contextlib.suppress(Exception): + self._use_catalog_database() return self._connection diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index 0e709061..37a68b2a 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -12,6 +12,10 @@ RECYCLEBIN_PREFIX = "__recyclebin_" +# LTABLE_BG scans every 30s; poll slightly beyond two intervals. +BG_POLL_TIMEOUT_SEC = 65.0 +BG_POLL_INTERVAL_SEC = 1.0 + # --------------------------------------------------------------------------- # # helpers # @@ -35,6 +39,12 @@ def _execute(client, sql: str): return client._server._execute(sql) +def _catalog_table(client, table: str) -> str: + """Qualified catalog table in the client's configured database (default: test).""" + db = client._server.database + return f"`{db}`.`{table}`" + + def _is_ss_mode(client) -> bool: """Best-effort detect whether the connected OB cluster is in SS mode.""" try: @@ -52,9 +62,10 @@ def _is_ss_mode(client) -> bool: def _fetch_namespace_name(client, collection_id: str, namespace_id: int): + ns_table = _catalog_table(client, "sdk_namespaces") rows = _execute( client, - f"SELECT namespace_name FROM sdk_namespaces " + f"SELECT namespace_name FROM {ns_table} " f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id}", ) if not rows: @@ -63,9 +74,10 @@ def _fetch_namespace_name(client, collection_id: str, namespace_id: int): def _count_ltables(client, collection_id: str, namespace_id: int) -> int: + lt_table = _catalog_table(client, "sdk_ltables") rows = _execute( client, - f"SELECT COUNT(*) AS c FROM sdk_ltables " + f"SELECT COUNT(*) AS c FROM {lt_table} " f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id}", ) return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) @@ -73,19 +85,21 @@ def _count_ltables(client, collection_id: str, namespace_id: int) -> int: def _count_logic_schema_rows(client, collection_id: str, namespace_id: int) -> int: tbl = NamespaceCollectionNames.logic_schema_table_name(collection_id) + db = client._server.database rows = _execute( client, - f"SELECT COUNT(*) AS c FROM `{tbl}` WHERE namespace_id = {namespace_id}", + f"SELECT COUNT(*) AS c FROM `{db}`.`{tbl}` WHERE namespace_id = {namespace_id}", ) return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) def _count_hot_table_rows(client, collection_id: str, namespace_id: int) -> int: tbl = NamespaceCollectionNames.hot_table_name(collection_id) + db = client._server.database try: rows = _execute( client, - f"SELECT COUNT(*) AS c FROM `{tbl}` WHERE namespace_id = {namespace_id}", + f"SELECT COUNT(*) AS c FROM `{db}`.`{tbl}` WHERE namespace_id = {namespace_id}", ) except Exception: return -1 # table missing @@ -94,31 +108,66 @@ def _count_hot_table_rows(client, collection_id: str, namespace_id: int) -> int: def _count_logic_data_rows(client, collection_id: str, namespace_id: int, ltable_id: int | None = None) -> int: tbl = NamespaceCollectionNames.data_table_name(collection_id) + db = client._server.database if ltable_id is not None: where = f"namespace_id = {namespace_id} AND ltable_id = {ltable_id}" else: where = f"namespace_id = {namespace_id}" rows = _execute( client, - f"SELECT COUNT(*) AS c FROM `{tbl}` WHERE {where}", + f"SELECT COUNT(*) AS c FROM `{db}`.`{tbl}` WHERE {where}", ) return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) def _count_kv_data_rows(client, collection_id: str, namespace_id: int) -> int: tbl = NamespaceCollectionNames.kv_data_table_name(collection_id) + db = client._server.database try: rows = _execute( client, - f"SELECT COUNT(*) AS c FROM `{tbl}` WHERE namespace_id = {namespace_id}", + f"SELECT COUNT(*) AS c FROM `{db}`.`{tbl}` WHERE namespace_id = {namespace_id}", ) except Exception: return -1 return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) +def _wait_until(predicate, timeout_sec: float = BG_POLL_TIMEOUT_SEC, + interval_sec: float = BG_POLL_INTERVAL_SEC, desc: str = ""): + deadline = time.time() + timeout_sec + last_exc = None + while time.time() < deadline: + try: + if predicate(): + return + except Exception as exc: # noqa: BLE001 + last_exc = exc + time.sleep(interval_sec) + msg = f"timeout after {timeout_sec}s waiting for {desc!r}" + if last_exc is not None: + raise AssertionError(msg) from last_exc + raise AssertionError(msg) + + +def _seed_kv_data(client, collection_id: str, namespace_id: int, count: int = 5): + """Seed rows in _kv_data_table (primary target of LTABLE_BG ns delete).""" + tbl = NamespaceCollectionNames.kv_data_table_name(collection_id) + db = client._server.database + values = [] + for i in range(count): + key_hex = f"{i + 1:064x}" + values.append(f"({namespace_id}, X'{key_hex}', X'00')") + _execute( + client, + f"INSERT INTO `{db}`.`{tbl}` (namespace_id, kv_key, kv_value) " + f"VALUES {', '.join(values)}", + ) + + def _seed_logic_data(client, collection_id: str, namespace_id: int, ltable_id: int, count: int = 3): tbl = NamespaceCollectionNames.data_table_name(collection_id) + db = client._server.database values = [] for i in range(count): values.append( @@ -126,13 +175,18 @@ def _seed_logic_data(client, collection_id: str, namespace_id: int, ltable_id: i f"X'0000803f0000000000000000', " f"'{{\"id\": \"id_{i}\", \"metadata\": {{\"k\": \"v\"}}}}')" ) - _execute(client, f"INSERT INTO `{tbl}` (namespace_id, ltable_id, document, embedding, data_content) VALUES {', '.join(values)}") + _execute( + client, + f"INSERT INTO `{db}`.`{tbl}` (namespace_id, ltable_id, document, embedding, data_content) " + f"VALUES {', '.join(values)}", + ) def _fetch_ltable_id(client, collection_id: str, namespace_id: int) -> int | None: + lt_table = _catalog_table(client, "sdk_ltables") rows = _execute( client, - f"SELECT ltable_id FROM sdk_ltables " + f"SELECT ltable_id FROM {lt_table} " f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id} " f"ORDER BY ltable_id LIMIT 1", ) @@ -141,14 +195,69 @@ def _fetch_ltable_id(client, collection_id: str, namespace_id: int) -> int | Non return int(rows[0]["ltable_id"] if "ltable_id" in rows[0] else rows[0]["LTABLE_ID"]) +def _hot_table_exists(client, collection_id: str) -> bool: + tbl = NamespaceCollectionNames.hot_table_name(collection_id) + db = client._server.database + try: + rows = _execute( + client, + "SELECT 1 AS ok FROM information_schema.tables " + f"WHERE table_schema = '{db}' AND table_name = '{tbl}'", + ) + return bool(rows) + except Exception: + return False + + +def _ensure_hot_table(client, collection_id: str) -> str | None: + """SDK only creates hot_table in SS mode; on SN create it for BG cleanup tests. + + Returns None on success, or an error message string. + """ + if _hot_table_exists(client, collection_id): + return None + tbl = NamespaceCollectionNames.hot_table_name(collection_id) + db = client._server.database + tg = NamespaceCollectionNames.tablegroup_name(collection_id) + ddl_with_tg = ( + f"CREATE TABLE IF NOT EXISTS `{db}`.`{tbl}` (" + f" namespace_id BIGINT UNSIGNED NOT NULL," + f" last_access_time TIMESTAMP(6) NOT NULL," + f" created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP," + f" updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP," + f" PRIMARY KEY(namespace_id)" + f") TABLEGROUP=`{tg}` COMMENT='热点/TTL附属表' DEFAULT CHARSET=utf8mb4 " + f"PARTITION BY KEY(namespace_id) PARTITIONS 1000" + ) + ddl_plain = ( + f"CREATE TABLE IF NOT EXISTS `{db}`.`{tbl}` (" + f" namespace_id BIGINT UNSIGNED NOT NULL," + f" last_access_time TIMESTAMP(6) NOT NULL," + f" PRIMARY KEY(namespace_id)" + f") COMMENT='热点/TTL附属表' DEFAULT CHARSET=utf8mb4 " + f"PARTITION BY KEY(namespace_id) PARTITIONS 1000" + ) + last_err = "" + for ddl in (ddl_with_tg, ddl_plain): + try: + _execute(client, ddl) + if _hot_table_exists(client, collection_id): + return None + except Exception as exc: # noqa: BLE001 + last_err = str(exc) + return last_err or "hot_table not visible after CREATE" + + def _seed_hot_table(client, collection_id: str, namespace_id: int): - """In SS mode insert a synthetic hot_table row so the DELETE has something - to remove and we can verify the cleanup.""" + """Insert a hot_table row (table must exist — use _ensure_hot_table on SN first).""" + if _ensure_hot_table(client, collection_id) is not None: + return False tbl = NamespaceCollectionNames.hot_table_name(collection_id) + db = client._server.database try: _execute( client, - f"INSERT INTO `{tbl}` (namespace_id, last_access_time) " + f"INSERT INTO `{db}`.`{tbl}` (namespace_id, last_access_time) " f"VALUES ({namespace_id}, NOW(6))", ) return True @@ -156,6 +265,48 @@ def _seed_hot_table(client, collection_id: str, namespace_id: int): return False +def _count_active_ltables(client, collection_id: str, namespace_id: int) -> int: + """sdk_ltables rows whose name is not a recyclebin entry (BG skip guard).""" + lt_table = _catalog_table(client, "sdk_ltables") + rows = _execute( + client, + f"SELECT COUNT(*) AS c FROM {lt_table} " + f"WHERE collection_id = '{collection_id}' AND namespace_id = {namespace_id} " + f"AND ltable_name NOT LIKE '{RECYCLEBIN_PREFIX}%'", + ) + return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) + + +def _insert_active_ltable_row( + client, + collection_id: str, + namespace_id: int, + ltable_name: str, + ltable_id: int, +): + lt_table = _catalog_table(client, "sdk_ltables") + _execute( + client, + f"INSERT INTO {lt_table} " + f"(ltable_id, collection_id, namespace_id, ltable_name) " + f"VALUES ({ltable_id}, '{collection_id}', {namespace_id}, '{ltable_name}')", + ) + + +def _fetch_namespace_drop_history(client, collection_id: str, namespace_id: int): + try: + return _execute( + client, + "SELECT operation_type, operation_status, total_deleted_rows " + "FROM oceanbase.__all_virtual_agent_drop_history " + f"WHERE collection_id = '{collection_id}' " + f" AND namespace_id = {namespace_id} AND ltable_id = 0 " + "ORDER BY finish_time DESC LIMIT 5", + ) + except Exception: + return None + + def _drop_namespace_via_pl(client, collection_id: str, namespace_id: int): """Call DBMS_LOGIC_TABLE.DROP_NAMESPACE directly (raw SQL, bypassing SDK's pre-existence guard).""" @@ -188,7 +339,7 @@ def _second_oceanbase_client(): class TestDropNamespaceCatalogValidation: - """Validate everything the DROP_NAMESPACE synchronous transaction must do.""" + """DROP_NAMESPACE sync transaction + LTABLE_BG async namespace deletion.""" # ------------------------------------------------------------- # # 1. Happy path: rename + cleanup all relevant rows # @@ -216,8 +367,8 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): assert _seed_hot_table(client, coll_id, ns_id) is True assert _count_hot_table_rows(client, coll_id, ns_id) == 1 - # Seed some logic_data rows so we can verify they survive the sync - # drop phase (async background task cleans them up later). + # Seed logic_data rows — sync DROP must not touch them; LTABLE_BG ns DAG + # deletes kv_data (see async tests), not logic_data_table. lt_id = _fetch_ltable_id(client, coll_id, ns_id) assert lt_id is not None, "should have a default ltable" _seed_logic_data(client, coll_id, ns_id, lt_id, count=3) @@ -257,9 +408,7 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): # 1b. Async cleanup: wait 32s, verify physical data is gone # # ------------------------------------------------------------- # def test_drop_namespace_async_cleanup(self, oceanbase_client): - """After DROP_NAMESPACE renames the namespace, a background task (30s - interval) picks up __recyclebin_ entries and removes physical data. - Wait 32s and verify the renamed namespace and its data are gone.""" + """LTABLE_BG NAMESPACE_DELETE: batch-delete kv_data, remove sdk_namespaces row.""" client = oceanbase_client collection = _make_collection(client) coll_id = collection.id @@ -269,38 +418,150 @@ def test_drop_namespace_async_cleanup(self, oceanbase_client): lt_id = _fetch_ltable_id(client, coll_id, ns_id) assert lt_id is not None - # Seed data + kv_before = max(_count_kv_data_rows(client, coll_id, ns_id), 0) + _seed_kv_data(client, coll_id, ns_id, count=5) _seed_logic_data(client, coll_id, ns_id, lt_id, count=3) + kv_after_seed = _count_kv_data_rows(client, coll_id, ns_id) + assert kv_after_seed > kv_before, ( + f"kv_data should have rows for BG to delete, before={kv_before} after={kv_after_seed}" + ) assert _count_logic_data_rows(client, coll_id, ns_id, lt_id) == 3 - # Drop — sync phase renames namespace, removes catalog entries _drop_namespace_via_pl(client, coll_id, ns_id) recycled_name = _fetch_namespace_name(client, coll_id, ns_id) assert recycled_name is not None assert recycled_name.startswith(RECYCLEBIN_PREFIX) - assert "ns_async" in recycled_name + assert _count_kv_data_rows(client, coll_id, ns_id) == kv_after_seed + assert _count_ltables(client, coll_id, ns_id) == 0 - # Data still present immediately after drop - assert _count_logic_data_rows(client, coll_id, ns_id, lt_id) == 3 + def _bg_done(): + return ( + _fetch_namespace_name(client, coll_id, ns_id) is None + and _count_kv_data_rows(client, coll_id, ns_id) == 0 + ) + + _wait_until(_bg_done, desc="namespace mapping and kv_data removed by LTABLE_BG") + + finally: + client.delete_collection(name=collection.name) + + def test_async_skipped_while_active_ltable_remains(self, oceanbase_client): + """BG must not delete sdk_namespaces while a non-recyclebin sdk_ltables row exists.""" + client = oceanbase_client + collection = _make_collection(client) + coll_id = collection.id + blocker_lt_id = 9_000_001 + try: + ns = collection.create_namespace("ns_block") + ns_id = int(ns.namespace_id) + _drop_namespace_via_pl(client, coll_id, ns_id) + assert _fetch_namespace_name(client, coll_id, ns_id).startswith(RECYCLEBIN_PREFIX) + assert _count_ltables(client, coll_id, ns_id) == 0 + + # Simulate a stuck active ltable row (sync phase normally clears these). + _insert_active_ltable_row( + client, coll_id, ns_id, "active_ltable_blocker", blocker_lt_id, + ) + assert _count_active_ltables(client, coll_id, ns_id) == 1 + + # Wait one BG scan period; scheduler must skip while active ltable exists. + time.sleep(35) + assert _fetch_namespace_name(client, coll_id, ns_id) is not None, ( + "recyclebin namespace row must remain while active ltable exists" + ) + + _execute( + client, + f"DELETE FROM {_catalog_table(client, 'sdk_ltables')} " + f"WHERE ltable_id = {blocker_lt_id}", + ) + assert _count_active_ltables(client, coll_id, ns_id) == 0 + + _wait_until( + lambda: _fetch_namespace_name(client, coll_id, ns_id) is None, + desc="namespace row removed after blocker ltable deleted", + ) + finally: + client.delete_collection(name=collection.name) + + def test_async_no_kv_data_fast_cleanup(self, oceanbase_client): + """Recyclebin namespace with zero kv rows should still be removed by BG.""" + client = oceanbase_client + collection = _make_collection(client) + coll_id = collection.id + try: + ns = collection.create_namespace("ns_empty_kv") + ns_id = int(ns.namespace_id) + assert _count_kv_data_rows(client, coll_id, ns_id) == 0 - # Wait for async background task (scans every 30s, 32s should suffice) - time.sleep(32) + _drop_namespace_via_pl(client, coll_id, ns_id) + assert _fetch_namespace_name(client, coll_id, ns_id).startswith(RECYCLEBIN_PREFIX) - # After background task: renamed namespace entry should be gone - final_name = _fetch_namespace_name(client, coll_id, ns_id) - assert final_name is None, ( - f"async task should have removed the __recyclebin_ namespace row, " - f"got {final_name!r}" + _wait_until( + lambda: _fetch_namespace_name(client, coll_id, ns_id) is None, + desc="empty-kv recyclebin namespace removed", ) + finally: + client.delete_collection(name=collection.name) + + def test_async_hot_table_cleaned_on_non_ss(self, oceanbase_client): + """When hot_table exists, sync DROP (SN) leaves it; LTABLE_BG deletes it.""" + client = oceanbase_client + if _is_ss_mode(client): + pytest.skip("SS mode deletes hot_table in synchronous DROP_NAMESPACE") + collection = _make_collection(client) + coll_id = collection.id + hot_err = _ensure_hot_table(client, coll_id) + if hot_err is not None: + pytest.fail(f"hot_table setup failed: {hot_err}") + try: + ns = collection.create_namespace("ns_hot_async") + ns_id = int(ns.namespace_id) + assert _seed_hot_table(client, coll_id, ns_id) is True + assert _count_hot_table_rows(client, coll_id, ns_id) == 1 - # Logic data should also be gone - data_rows = _count_logic_data_rows(client, coll_id, ns_id, lt_id) - assert data_rows == 0, ( - f"async task should have cleaned up logic_data rows, " - f"found {data_rows}" + _drop_namespace_via_pl(client, coll_id, ns_id) + assert _fetch_namespace_name(client, coll_id, ns_id).startswith(RECYCLEBIN_PREFIX) + assert _count_hot_table_rows(client, coll_id, ns_id) == 1, ( + "sync DROP must not remove hot_table on non-SS" ) + def _hot_and_ns_gone(): + return ( + _count_hot_table_rows(client, coll_id, ns_id) == 0 + and _fetch_namespace_name(client, coll_id, ns_id) is None + ) + + _wait_until(_hot_and_ns_gone, desc="hot_table and namespace removed by BG") + finally: + client.delete_collection(name=collection.name) + + def test_async_drop_history_recorded(self, oceanbase_client): + """After successful BG delete, audit row appears in agent drop history.""" + client = oceanbase_client + collection = _make_collection(client) + coll_id = collection.id + try: + ns = collection.create_namespace("ns_history") + ns_id = int(ns.namespace_id) + _seed_kv_data(client, coll_id, ns_id, count=2) + _drop_namespace_via_pl(client, coll_id, ns_id) + + def _bg_done(): + return _fetch_namespace_name(client, coll_id, ns_id) is None + + _wait_until(_bg_done, desc="namespace removed for history check") + + history = _fetch_namespace_drop_history(client, coll_id, ns_id) + if history is None: + pytest.skip("__all_virtual_agent_drop_history not available") + assert len(history) >= 1, "expected at least one history row" + row = history[0] + op = row.get("operation_type") or row.get("OPERATION_TYPE") + status = row.get("operation_status") or row.get("OPERATION_STATUS") + assert op == "NAMESPACE_DELETE", f"unexpected operation_type: {op!r}" + assert status == "COMPLETED", f"unexpected operation_status: {status!r}" finally: client.delete_collection(name=collection.name) @@ -334,9 +595,10 @@ def test_multiple_ltables_all_deleted(self, oceanbase_client): ns = collection.create_namespace("ns_multi") ns_id = int(ns.namespace_id) # Insert two extra ltable rows directly so we can verify bulk delete. + lt_table = _catalog_table(client, "sdk_ltables") _execute( client, - f"INSERT INTO sdk_ltables (collection_id, namespace_id, ltable_name) " + f"INSERT INTO {lt_table} (collection_id, namespace_id, ltable_name) " f"VALUES ('{collection.id}', {ns_id}, 'extra_lt_a')," f" ('{collection.id}', {ns_id}, 'extra_lt_b')", ) @@ -360,6 +622,7 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id schema_tbl = NamespaceCollectionNames.logic_schema_table_name(coll_id) + schema_tbl_q = f"`{client._server.database}`.`{schema_tbl}`" try: ns = collection.create_namespace("ns_rollback") ns_id = int(ns.namespace_id) @@ -370,7 +633,7 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): # Sabotage: drop the logic_schema_table so step 6 inside the # DROP_NAMESPACE transaction will fail with OB_TABLE_NOT_EXIST. - _execute(client, f"DROP TABLE `{schema_tbl}`") + _execute(client, f"DROP TABLE {schema_tbl_q}") with pytest.raises(Exception): _drop_namespace_via_pl(client, coll_id, ns_id) @@ -389,7 +652,7 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): try: _execute( client, - f"CREATE TABLE IF NOT EXISTS `{schema_tbl}` (" + f"CREATE TABLE IF NOT EXISTS {schema_tbl_q} (" f" namespace_id BIGINT UNSIGNED NOT NULL," f" ltable_id BIGINT UNSIGNED NOT NULL," f" schema_content JSON NOT NULL," From 6d7ffb60380fd32ae84bebbc28599c3638fcc9d0 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 3 Jun 2026 15:25:24 +0800 Subject: [PATCH 28/84] fix(hybrid_search): replace exists with OB-compatible positive must_not clauses --- src/pyseekdb/client/client_base.py | 38 +++++++++++++------ .../unit_tests/test_build_query_expression.py | 5 ++- 2 files changed, 30 insertions(+), 13 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index d619f21c..7632670c 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -4021,19 +4021,39 @@ def _pure_must_not_clauses(expr: dict[str, Any] | None) -> list[dict[str, Any]] return None return clauses + @staticmethod + def _collect_scalar_fields_from_filter_clause(clause: dict[str, Any]) -> list[str]: + """Collect scalar field names from term/terms leaves (including nested bool clauses).""" + fields: list[str] = [] + if not isinstance(clause, dict): + return fields + term_body = clause.get("term") + if isinstance(term_body, dict): + fields.extend(term_body.keys()) + terms_body = clause.get("terms") + if isinstance(terms_body, dict): + fields.extend(terms_body.keys()) + range_body = clause.get("range") + if isinstance(range_body, dict): + fields.extend(range_body.keys()) + bool_node = clause.get("bool") + if isinstance(bool_node, dict): + for key in ("filter", "must", "should", "must_not"): + sub = bool_node.get(key) + if isinstance(sub, list): + for item in sub: + fields.extend(BaseClient._collect_scalar_fields_from_filter_clause(item)) + elif isinstance(sub, dict): + fields.extend(BaseClient._collect_scalar_fields_from_filter_clause(sub)) + return fields + @staticmethod def _positive_clause_for_must_not(must_not_clauses: list[dict[str, Any]]) -> dict[str, Any]: """Build a permissive positive filter leaf for must_not-only bools (OB rejects match_all).""" for clause in must_not_clauses: if not isinstance(clause, dict): continue - term_body = clause.get("term") - if isinstance(term_body, dict) and term_body: - field = next(iter(term_body)) - return {"range": {field: {"gte": -9223372036854775808}}} - terms_body = clause.get("terms") - if isinstance(terms_body, dict) and terms_body: - field = next(iter(terms_body)) + for field in BaseClient._collect_scalar_fields_from_filter_clause(clause): return {"range": {field: {"gte": -9223372036854775808}}} qs_body = clause.get("query_string") if isinstance(qs_body, dict): @@ -4113,10 +4133,6 @@ def _build_query_expression(self, query: dict[str, Any]) -> dict[str, Any] | Non bool_q["filter"] = pos_filters if must_not_all: bool_q["must_not"] = must_not_all - # Pure negation ($not_contains, $ne, …) needs a positive clause unless - # a scoring `must` is already present. - if doc_must_not is not None and not pos_filters and "must" not in bool_q: - bool_q["filter"] = [self._positive_clause_for_must_not(must_not_all)] return {"bool": bool_q} return None diff --git a/tests/unit_tests/test_build_query_expression.py b/tests/unit_tests/test_build_query_expression.py index 0980988a..4797c56f 100644 --- a/tests/unit_tests/test_build_query_expression.py +++ b/tests/unit_tests/test_build_query_expression.py @@ -43,7 +43,8 @@ def test_not_contains_with_metadata_filter_hoists_must_not(self, client): } assert "must" not in expr["bool"] - def test_not_contains_only_uses_exists_positive_filter(self, client): + def test_not_contains_only_hoists_must_not_without_synthetic_filter(self, client): + """Namespace path injects ns/lt filters later; no exists/match_all leaf here.""" with patch.object(client, "_build_metadata_filter_for_search_parm", return_value=[]): expr = client._build_query_expression({ "where_document": {"$not_contains": "TOKENZPX"}, @@ -51,7 +52,6 @@ def test_not_contains_only_uses_exists_positive_filter(self, client): assert expr == { "bool": { - "filter": [{"exists": {"field": "document"}}], "must_not": [ { "query_string": { @@ -62,6 +62,7 @@ def test_not_contains_only_uses_exists_positive_filter(self, client): ], } } + assert "filter" not in expr["bool"] def test_contains_with_metadata_filter_still_uses_must(self, client): with patch.object( From 78c3d557bc7ba523b3c24ca701f035987dd576a4 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 4 Jun 2026 14:28:40 +0800 Subject: [PATCH 29/84] =?UTF-8?q?prewarm=5Flob=20sdk=E9=80=82=E9=85=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 2 +- .../test_namespace_prewarm_lob.py | 367 ++++++++++++++++++ 2 files changed, 368 insertions(+), 1 deletion(-) create mode 100644 tests/integration_tests/test_namespace_prewarm_lob.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7632670c..ba9ec58b 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1341,7 +1341,7 @@ def _create_namespace_physical_tables( kv_key VARBINARY(1024) NOT NULL, kv_value LONGBLOB NOT NULL, PRIMARY KEY(namespace_id, kv_key) - ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 LOB_INROW_THRESHOLD = 786432 + ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 LOB_INROW_THRESHOLD=786432 {partition_clause}""") self._execute(f"""CREATE TABLE `{schema_table}` ( diff --git a/tests/integration_tests/test_namespace_prewarm_lob.py b/tests/integration_tests/test_namespace_prewarm_lob.py new file mode 100644 index 00000000..8987b07b --- /dev/null +++ b/tests/integration_tests/test_namespace_prewarm_lob.py @@ -0,0 +1,367 @@ +""" +LOB prewarm integration tests (OceanBase shared-storage mode). + +Validates the OB-side change that, after prewarming a namespace's KV data tablet, +the prewarm RPC ALSO prewarms that KV tablet's aux LOB-meta tablet (whole range). + +Observability is layered into three independent tiers so the suite degrades +gracefully depending on what the cluster exposes: + + Tier 1 - structural (always): + Resolve the kv_data_table's aux LOB-meta tablet purely from SQL + (__all_table.data_table_id + table_type=13 -> __all_virtual_tablet_to_ls). + Proves the table the prewarm targets actually owns a LOB-meta tablet. + + Tier 2 - ss-cache (SS mode + real out-of-row LOB): + Force a >0.75MB kv_value so it spills out-of-row into the LOB-meta tablet, + freeze it to a macro block, prewarm, then read + __all_virtual_ss_macro_cache_info filtered by that LOB-meta tablet_id. + Proves blocks were actually pulled into the local cache. + + Tier 3 - log marker (opt-in via OB_OBSERVER_LOG): + Grep observer.log for "prewarm_logic_table lob meta done", which prints the + exact lob_meta_tablet_id. Proves the new code path executed for the right tablet. + +Run: + OB_PORT=11202 OB_TENANT=mysql \ + OB_OBSERVER_LOG=/path/to/observer.log \ + pytest tests/integration_tests/test_namespace_prewarm_lob.py -v -s +""" + +import contextlib +import os +import time + +import pytest + +try: + import pymysql +except ImportError: # pragma: no cover - pymysql ships with the test deps + pymysql = None + +from pyseekdb import ( + IVFConfiguration, + get_namespace_partition_count, + set_namespace_partition_count, +) +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.meta_info import NamespaceCollectionNames +from pyseekdb.client.schema import Schema + +AUX_LOB_META_TABLE_TYPE = 13 +LARGE_BLOB_BYTES = 1_000_000 # > OB_MAX_LOB_INROW_THRESHOLD (786432) -> always out-of-row +LOB_PROBE_ROWS = 3 # 3 x ~1MB incompressible -> ~3MB of real macro blocks +SS_CACHE_VIEW = "oceanbase.__all_virtual_ss_macro_cache_info" +ALL_TABLE_VIEW = "oceanbase.__all_table" +TABLET_TO_LS_VIEW = "oceanbase.__all_tablet_to_ls" + +# Flushing the SS local cache requires a sys-tenant connection with an explicit +# TENANT clause; user-tenant flush returns OB-4016. These are env-overridable so +# the rigorous before/after assertion can run wherever a sys login is available. +OB_HOST = os.environ.get("OB_HOST", "127.0.0.1") +OB_PORT = int(os.environ.get("OB_PORT", "10902")) +OB_TENANT = os.environ.get("OB_TENANT", "test_tenant") +OB_SYS_USER = os.environ.get("OB_SYS_USER", "root@sys") +OB_SYS_PASSWORD = os.environ.get("OB_SYS_PASSWORD", "") + + +# ==================== SQL observability helpers ==================== + +def _exec(client, sql): + return client._server._execute(sql) + + +def _resolve_table_id(client, table_name): + rows = _exec( + client, + f"SELECT table_id FROM {ALL_TABLE_VIEW} WHERE table_name = '{table_name}'", + ) + assert rows, f"table_id not found for {table_name}" + return rows[0]["table_id"] + + +def _resolve_lob_meta_tablets(client, kv_table_id): + """All aux-LOB-meta tablets of a KV data table (one per KV partition).""" + rows = _exec( + client, + f"SELECT table_id FROM {ALL_TABLE_VIEW} " + f"WHERE data_table_id = {kv_table_id} AND table_type = {AUX_LOB_META_TABLE_TYPE}", + ) + assert rows, f"no aux lob meta table for kv table {kv_table_id}" + lob_meta_table_id = rows[0]["table_id"] + + tablet_rows = _exec( + client, + f"SELECT tablet_id FROM {TABLET_TO_LS_VIEW} WHERE table_id = {lob_meta_table_id}", + ) + return {r["tablet_id"] for r in tablet_rows} + + +def _ss_cache_blocks_for_tablets(client, tablet_ids): + """Map tablet_id -> (block_count, total_size) from the SS macro cache view. + + Returns None when the SS cache view is not available (non-SS cluster), so + callers can skip tier-2 assertions cleanly. + """ + if not tablet_ids: + return {} + ids = ",".join(str(t) for t in tablet_ids) + try: + rows = _exec( + client, + f"SELECT tablet_id, COUNT(*) AS blk, COALESCE(SUM(size),0) AS sz " + f"FROM {SS_CACHE_VIEW} WHERE tablet_id IN ({ids}) GROUP BY tablet_id", + ) + except Exception: + return None + return {r["tablet_id"]: (r["blk"], r["sz"]) for r in rows} + + +def _flush_tenant_macro_cache(): + """Evict the tenant's SS local macro cache via a sys-tenant connection. + + Returns True on success, False if a sys login / flush is unavailable so the + caller can skip the rigorous tier-2 assertion instead of failing spuriously. + """ + if pymysql is None: + return False + try: + conn = pymysql.connect( + host=OB_HOST, port=OB_PORT, user=OB_SYS_USER, + password=OB_SYS_PASSWORD, autocommit=True, + ) + except Exception: + return False + try: + with conn.cursor() as cur: + cur.execute( + f"ALTER SYSTEM FLUSH SS_LOCAL_CACHE TENANT = {OB_TENANT} CACHE = macro_cache" + ) + return True + except Exception: + return False + finally: + with contextlib.suppress(Exception): + conn.close() + + +def _is_ss_mode(client, collection_id): + """SS mode is detectable by the per-collection hot_table existing.""" + hot_table = NamespaceCollectionNames.hot_table_name(collection_id) + rows = _exec( + client, + f"SELECT table_id FROM {ALL_TABLE_VIEW} WHERE table_name = '{hot_table}'", + ) + return bool(rows) + + +def _grep_lob_prewarm_log(lob_meta_tablet_ids): + """Tier 3: confirm the new code path logged a lob prewarm for one of the tablets.""" + log_path = os.environ.get("OB_OBSERVER_LOG") + if not log_path or not os.path.exists(log_path): + return None + wanted = {str(t) for t in lob_meta_tablet_ids} + hits = [] + with open(log_path, errors="ignore") as fh: + for line in fh: + if "prewarm_logic_table lob meta done" in line and any(t in line for t in wanted): + hits.append(line.strip()) + return hits + + +# ==================== Fixtures / collection helpers ==================== + +class _BaseLobPrewarm: + + def _make_collection(self, client, name, partitions): + set_namespace_partition_count(partitions) + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine"), + embedding_function=None, + ), + ) + return client.create_collection(name=name, schema=schema, use_namespace=True) + + def _force_out_of_row_lob(self, client, collection_id, namespace_id): + """Insert several >0.75MB incompressible kv_values so they spill out-of-row + into the LOB-meta tablet, then freeze them into macro blocks. + + REPEAT('x', N) compresses to ~nothing and yields a single tiny block, which + is useless for proving prewarm restored real data. os.urandom is + incompressible, so each row produces a genuine ~1MB macro block. + """ + kv_table = NamespaceCollectionNames.kv_data_table_name(collection_id) + # _execute() takes no params, so use the raw DBAPI connection to bind the + # binary blob safely (avoids building multi-MB hex-literal SQL strings). + raw = client._server.get_raw_connection() + with raw.cursor() as cur: + for i in range(LOB_PROBE_ROWS): + cur.execute( + f"INSERT INTO `{kv_table}` (namespace_id, kv_key, kv_value) " + f"VALUES ({namespace_id}, %s, %s)", + (f"__lob_probe_{i}__", os.urandom(LARGE_BLOB_BYTES)), + ) + raw.commit() + # Push the rows into macro blocks so prewarm has remote data to pull. + with contextlib.suppress(Exception): + _exec(client, "ALTER SYSTEM MINOR FREEZE") + + +# ==================== Tier 1 + structural multi-namespace ==================== + +class TestLobPrewarmStructural(_BaseLobPrewarm): + + @pytest.fixture(autouse=True) + def _restore_partitions(self): + original = get_namespace_partition_count() + yield + set_namespace_partition_count(original) + + def test_kv_table_has_lob_meta_tablet(self, oceanbase_client): + """Tier 1: the prewarm target (kv_data_table) owns an aux LOB-meta tablet.""" + name = f"lob_pw_struct_{int(time.time() * 1000)}" + collection = self._make_collection(oceanbase_client, name, partitions=1) + try: + kv_table = NamespaceCollectionNames.kv_data_table_name(collection.id) + kv_table_id = _resolve_table_id(oceanbase_client, kv_table) + lob_tablets = _resolve_lob_meta_tablets(oceanbase_client, kv_table_id) + assert len(lob_tablets) == 1, ( + f"single-partition kv table should have exactly 1 lob meta tablet, " + f"got {lob_tablets}" + ) + finally: + oceanbase_client.delete_collection(name=collection.name) + + def test_multi_namespace_prewarm_is_independent(self, oceanbase_client): + """Multi-namespace: prewarming each namespace succeeds, records its own hot_table + row, and never disturbs another namespace's row. Cached LOB tablets stay within + the collection's own LOB-meta tablet set (no foreign tablets).""" + name = f"lob_pw_multi_{int(time.time() * 1000)}" + collection = self._make_collection(oceanbase_client, name, partitions=8) + try: + kv_table = NamespaceCollectionNames.kv_data_table_name(collection.id) + kv_table_id = _resolve_table_id(oceanbase_client, kv_table) + own_lob_tablets = _resolve_lob_meta_tablets(oceanbase_client, kv_table_id) + + namespaces = [collection.create_namespace(f"ns_{i}") for i in range(3)] + for ns in namespaces: + ns.prewarm() + + hot_table = NamespaceCollectionNames.hot_table_name(collection.id) + rows = _exec( + oceanbase_client, + f"SELECT namespace_id FROM `{hot_table}`", + ) + recorded = {int(r["namespace_id"]) for r in rows} + expected = {int(ns._namespace_id) for ns in namespaces} + assert expected.issubset(recorded), ( + f"every prewarmed namespace must have a hot_table row: " + f"expected {expected}, got {recorded}" + ) + + cache = _ss_cache_blocks_for_tablets(oceanbase_client, own_lob_tablets) + if cache: + assert set(cache).issubset(own_lob_tablets), ( + "cached lob tablets must belong to this collection only" + ) + finally: + oceanbase_client.delete_collection(name=collection.name) + + +# ==================== Tier 2 + Tier 3: real LOB caching ==================== + +class TestLobPrewarmCaching(_BaseLobPrewarm): + + @pytest.fixture(autouse=True) + def _restore_partitions(self): + original = get_namespace_partition_count() + yield + set_namespace_partition_count(original) + + @staticmethod + def _total_bytes(cache): + return sum(v[1] for v in (cache or {}).values()) + + def test_out_of_row_lob_is_prewarmed_into_cache(self, oceanbase_client): + """Tier 2/3: prove prewarm actually *pulls* evicted LOB data back, not that + the write cache merely still holds it. + + Methodology (the only rigorous one): + 1. Insert incompressible out-of-row LOB rows + freeze -> real macro blocks. + 2. Flush the tenant's SS macro cache (sys connection) -> evict those blocks. + 3. Assert cache bytes for the LOB tablet dropped near-zero BEFORE prewarm. + 4. prewarm() -> poll the cache view. + 5. Assert cache bytes were restored (>> the flushed floor) AFTER prewarm. + + Without the flush step, frozen data already sits in the local write cache, so a + plain "blocks > 0 after prewarm" assertion would pass even if prewarm did nothing. + """ + name = f"lob_pw_cache_{int(time.time() * 1000)}" + collection = self._make_collection(oceanbase_client, name, partitions=1) + try: + if not _is_ss_mode(oceanbase_client, collection.id): + pytest.skip("LOB cache prewarm only observable in shared-storage mode") + + namespace = collection.create_namespace("lob_ns") + kv_table = NamespaceCollectionNames.kv_data_table_name(collection.id) + kv_table_id = _resolve_table_id(oceanbase_client, kv_table) + lob_tablets = _resolve_lob_meta_tablets(oceanbase_client, kv_table_id) + + self._force_out_of_row_lob(oceanbase_client, collection.id, namespace._namespace_id) + + if _ss_cache_blocks_for_tablets(oceanbase_client, lob_tablets) is None: + pytest.skip(f"{SS_CACHE_VIEW} not available on this cluster") + + # Give the freeze time to upload macro blocks to object storage. + time.sleep(10) + written = _ss_cache_blocks_for_tablets(oceanbase_client, lob_tablets) or {} + assert self._total_bytes(written) > 0, ( + f"out-of-row LOB should have produced macro blocks in cache, got {written}" + ) + + # Evict so prewarm has something real to restore. Requires sys flush. + if not _flush_tenant_macro_cache(): + pytest.skip( + "sys-tenant SS_LOCAL_CACHE flush unavailable; cannot prove " + "prewarm pulls evicted data (set OB_SYS_USER/OB_SYS_PASSWORD)" + ) + time.sleep(4) + + flushed = _ss_cache_blocks_for_tablets(oceanbase_client, lob_tablets) or {} + flushed_bytes = self._total_bytes(flushed) + written_bytes = self._total_bytes(written) + assert flushed_bytes < written_bytes, ( + f"flush should have evicted LOB macro blocks before prewarm: " + f"written={written_bytes}, flushed={flushed_bytes}" + ) + + namespace.prewarm() + + # Prewarm pulls remote macro blocks asynchronously; poll for restoration. + cached = flushed + deadline = time.time() + 40 + while time.time() < deadline: + cached = _ss_cache_blocks_for_tablets(oceanbase_client, lob_tablets) or {} + if self._total_bytes(cached) > flushed_bytes: + break + time.sleep(2) + + log_hits = _grep_lob_prewarm_log(lob_tablets) + if log_hits is not None: + assert log_hits, ( + "observer.log should contain a 'prewarm_logic_table lob meta done' " + f"line for lob tablets {lob_tablets}" + ) + + restored_bytes = self._total_bytes(cached) + assert restored_bytes > flushed_bytes, ( + f"prewarm should have restored evicted LOB macro blocks: " + f"written={written_bytes}, flushed={flushed_bytes}, restored={restored_bytes}" + ) + finally: + oceanbase_client.delete_collection(name=collection.name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) From 227fdb35d69e2b48ed1aef6b10fdbf89f4f68a0b Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Thu, 4 Jun 2026 14:42:25 +0800 Subject: [PATCH 30/84] add cosine case --- .../namespace_fts_helpers.py | 55 ++-- .../namespace_hybrid_search_helpers.py | 238 +++++++++++++++--- .../test_namespace_get_delete_filters.py | 54 +++- .../test_namespace_hybrid_search_combined.py | 30 ++- ...space_hybrid_search_multi_coll_multi_ns.py | 32 ++- ...t_namespace_hybrid_search_triple_branch.py | 35 ++- ...earch_triple_branch_multi_coll_multi_ns.py | 85 +++++-- .../test_namespace_hybrid_search_vector.py | 42 ++-- 8 files changed, 443 insertions(+), 128 deletions(-) diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index 141376cb..c8acd7e4 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -12,7 +12,10 @@ import math import time from dataclasses import dataclass -from typing import Any +from typing import Any, Literal + +VectorDistanceMetric = Literal["l2", "cosine"] +VECTOR_DISTANCE_METRICS: tuple[VectorDistanceMetric, ...] = ("l2", "cosine") from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig @@ -49,22 +52,22 @@ class FtsQueryCase: where: dict[str, Any] | None = None -def ns_schema() -> Schema: +def ns_schema(distance: VectorDistanceMetric = "l2") -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance=distance, fresh_mode="spfresh"), embedding_function=None, ), ) -def flat_hybrid_search_schema() -> Schema: +def flat_hybrid_search_schema(distance: VectorDistanceMetric = "l2") -> Schema: """Schema for flat collection baseline (``use_namespace=False``, HNSW + FTS).""" from pyseekdb import HNSWConfiguration return Schema( vector_index=VectorIndexConfig( - hnsw=HNSWConfiguration(dimension=3, distance="l2"), + hnsw=HNSWConfiguration(dimension=3, distance=distance), embedding_function=None, ), ) @@ -498,14 +501,20 @@ def get_fts_case(name: str) -> FtsQueryCase: raise KeyError(f"unknown FTS case: {name!r}") -def setup_large_fts_collection(db_client: Any) -> tuple[list[CorpusRecord], Any]: +def setup_large_fts_collection( + db_client: Any, + *, + distance: VectorDistanceMetric = "l2", +) -> tuple[list[CorpusRecord], Any]: """Create one namespace-enabled collection with a deterministic large corpus.""" corpus = build_large_fts_corpus(CORPUS_SIZE) if len(corpus) <= 1000: raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") - name = f"test_ns_hs_ft_{int(time.time() * 1000)}" - collection = db_client.create_collection(name=name, schema=ns_schema(), use_namespace=True) + name = f"test_ns_hs_ft_{distance}_{int(time.time() * 1000)}" + collection = db_client.create_collection( + name=name, schema=ns_schema(distance), use_namespace=True + ) return corpus, collection @@ -555,16 +564,21 @@ def teardown_large_fts_namespace(db_client: Any, collection: Any) -> None: MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS: tuple[str, ...] = ("c1_x", "c2_y") -def _create_multi_coll_multi_ns_layout(db_client: Any, *, name_prefix: str) -> dict[str, Any]: +def _create_multi_coll_multi_ns_layout( + db_client: Any, + *, + name_prefix: str, + distance: VectorDistanceMetric = "l2", +) -> dict[str, Any]: ts = int(time.time() * 1000) coll_1 = db_client.create_collection( - name=f"{name_prefix}_{ts}_c1", - schema=ns_schema(), + name=f"{name_prefix}_{distance}_{ts}_c1", + schema=ns_schema(distance), use_namespace=True, ) coll_2 = db_client.create_collection( - name=f"{name_prefix}_{ts}_c2", - schema=ns_schema(), + name=f"{name_prefix}_{distance}_{ts}_c2", + schema=ns_schema(distance), use_namespace=True, ) ctx: dict[str, Any] = {"coll_1": coll_1, "coll_2": coll_2} @@ -574,7 +588,11 @@ def _create_multi_coll_multi_ns_layout(db_client: Any, *, name_prefix: str) -> d return ctx -def setup_multi_coll_multi_ns_fts(db_client: Any) -> dict[str, Any]: +def setup_multi_coll_multi_ns_fts( + db_client: Any, + *, + distance: VectorDistanceMetric = "l2", +) -> dict[str, Any]: """ 2 collections x 2 namespaces; load the large FTS corpus into ``c1_x`` and ``c2_y`` only. @@ -582,7 +600,9 @@ def setup_multi_coll_multi_ns_fts(db_client: Any) -> dict[str, Any]: only the namespace that received the bulk insert is expected to serve hybrid_search FTS (see ``setup_same_collection_both_ns_fts`` for same-collection multi-ns behavior). """ - ctx = _create_multi_coll_multi_ns_layout(db_client, name_prefix="test_ns_hs_ft_mcmn") + ctx = _create_multi_coll_multi_ns_layout( + db_client, name_prefix="test_ns_hs_ft_mcmn", distance=distance + ) corpus = build_large_fts_corpus(CORPUS_SIZE) if len(corpus) <= 1000: raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") @@ -630,12 +650,15 @@ def setup_multi_coll_multi_ns_fts_single_loaded( db_client: Any, *, loaded_quadrant: str = "c1_x", + distance: VectorDistanceMetric = "l2", ) -> dict[str, Any]: """Same layout as :func:`setup_multi_coll_multi_ns_fts` but corpus only in ``loaded_quadrant``.""" if loaded_quadrant not in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: raise ValueError(f"unknown quadrant {loaded_quadrant!r}") - ctx = _create_multi_coll_multi_ns_layout(db_client, name_prefix="test_ns_hs_ft_mcmn1") + ctx = _create_multi_coll_multi_ns_layout( + db_client, name_prefix="test_ns_hs_ft_mcmn1", distance=distance + ) corpus = build_large_fts_corpus(CORPUS_SIZE) insert_corpus_in_batches(ctx[loaded_quadrant], corpus) diff --git a/tests/integration_tests/namespace_hybrid_search_helpers.py b/tests/integration_tests/namespace_hybrid_search_helpers.py index 8430fd1d..dbbdcfd3 100644 --- a/tests/integration_tests/namespace_hybrid_search_helpers.py +++ b/tests/integration_tests/namespace_hybrid_search_helpers.py @@ -18,6 +18,8 @@ INDEX_SETTLE_SECONDS, TOKEN_ALP, TOKEN_ZPX, + VECTOR_DISTANCE_METRICS, + VectorDistanceMetric, CorpusRecord, assert_hybrid_fulltext_result, assert_not_contains_no_token_leak, @@ -133,19 +135,111 @@ def l2_squared(a: list[float], b: list[float]) -> float: return sum((x - y) ** 2 for x, y in zip(a, b)) -def expected_knn_ids( +def _vector_l2_norm(vec: list[float]) -> float: + return math.sqrt(sum(x * x for x in vec)) + + +def knn_compare_distance( + a: list[float], + b: list[float], + metric: VectorDistanceMetric = "l2", +) -> float: + """Distance value used for KNN ground-truth ordering (smaller = nearer).""" + if metric == "l2": + return l2_squared(a, b) + if metric == "cosine": + norm_a = _vector_l2_norm(a) + norm_b = _vector_l2_norm(b) + if norm_a == 0.0 or norm_b == 0.0: + return 2.0 + cos_sim = sum(x * y for x, y in zip(a, b)) / (norm_a * norm_b) + return 1.0 - cos_sim + raise ValueError(f"unsupported vector distance metric: {metric!r}") + + +def ensure_shared_hybrid_search_collection( + shared_cache: dict[str, dict[str, Any]], + db_client: Any, + request: Any, + vector_distance: VectorDistanceMetric, +) -> dict[str, Any]: + """Return or create a shared (db mode, distance metric) corpus + collection entry.""" + mode = ( + request.node.callspec.params["db_client"] + if request.node.callspec + else "default" + ) + cache_key = f"{mode}:{vector_distance}" + if cache_key not in shared_cache: + corpus, collection = setup_large_fts_collection( + db_client, distance=vector_distance + ) + shared_cache[cache_key] = { + "db_client": db_client, + "corpus": corpus, + "collection": collection, + "vector_distance": vector_distance, + } + return shared_cache[cache_key] + + +def _knn_scored_rows( corpus: list[CorpusRecord], query_vector: list[float], - n_results: int, where: dict[str, Any] | None = None, -) -> list[str]: - scored: list[tuple[str, float]] = [] + *, + distance_metric: VectorDistanceMetric = "l2", +) -> list[tuple[str, float, int]]: + """(doc_id, compare_distance, metadata seq) for rows matching *where*.""" + rows: list[tuple[str, float, int]] = [] for rec in corpus: if where is not None and not corpus_matches_where(rec, where): continue - scored.append((rec.doc_id, l2_squared(rec.embedding, query_vector))) - scored.sort(key=lambda item: (item[1], item[0])) - return [doc_id for doc_id, _ in scored[:n_results]] + seq = rec.metadata.get("seq", 0) + if not isinstance(seq, int): + seq = 0 + rows.append( + ( + rec.doc_id, + knn_compare_distance(rec.embedding, query_vector, distance_metric), + seq, + ) + ) + return rows + + +def nearest_knn_doc_ids( + corpus: list[CorpusRecord], + query_vector: list[float], + where: dict[str, Any] | None = None, + *, + distance_metric: VectorDistanceMetric = "l2", + tol: float = 1e-9, +) -> set[str]: + """Doc ids at minimum compare distance (ties allowed).""" + scored = _knn_scored_rows( + corpus, query_vector, where, distance_metric=distance_metric + ) + if not scored: + return set() + min_d = min(d for _, d, _ in scored) + return {doc_id for doc_id, d, _ in scored if math.isclose(d, min_d, abs_tol=tol)} + + +def expected_knn_ids( + corpus: list[CorpusRecord], + query_vector: list[float], + n_results: int, + where: dict[str, Any] | None = None, + *, + distance_metric: VectorDistanceMetric = "l2", +) -> list[str]: + scored = _knn_scored_rows( + corpus, query_vector, where, distance_metric=distance_metric + ) + # Tie-break: distance, then seq (matches OB hybrid_search), then doc_id. + scored.sort(key=lambda item: (item[1], item[2], item[0])) + return [doc_id for doc_id, _, _ in scored[:n_results]] def assert_hybrid_search_index_result( @@ -194,12 +288,20 @@ def assert_hybrid_knn_result( n_results: int, where: dict[str, Any] | None = None, *, + distance_metric: VectorDistanceMetric = "l2", check_top1: bool = True, + check_score_order: bool = True, ) -> None: + """Assert hybrid_search KNN branch results against corpus ground truth. + + ``result["distances"]`` comes from OB ``__score`` via the client (higher = nearer), + not raw vector distance — only score monotonicity is checked here. Id order uses + corpus ``distance_metric`` ground truth. + """ assert result is not None assert "ids" in result and result["ids"] ids = result["ids"][0] - distances = result.get("distances", [[]])[0] if result.get("distances") else [] + scores = result.get("distances", [[]])[0] if result.get("distances") else [] assert len(ids) <= n_results assert len(ids) > 0, "expected at least one KNN hit" @@ -211,30 +313,49 @@ def assert_hybrid_knn_result( f"id={doc_id!r} metadata does not satisfy knn.where={where!r}" ) - if distances: - assert len(distances) == len(ids) - for dist in distances: - assert dist >= 0 - for i in range(len(distances) - 1): - assert distances[i] <= distances[i + 1] or math.isclose( - distances[i], distances[i + 1] - ), f"KNN distances should be non-decreasing (L2): {distances!r}" - - expected = expected_knn_ids(corpus, query_vector, n_results, where) - if check_top1 and expected: - assert ids[0] == expected[0], ( - f"top-1 KNN must be nearest neighbor, got {ids[0]!r} expected {expected[0]!r}" + if scores: + assert len(scores) == len(ids) + if check_score_order: + for i in range(len(scores) - 1): + assert scores[i] >= scores[i + 1] or math.isclose( + scores[i], scores[i + 1] + ), ( + f"KNN scores should be non-increasing (higher = nearer, " + f"{distance_metric} index): {scores!r}" + ) + + expected = expected_knn_ids( + corpus, query_vector, n_results, where, distance_metric=distance_metric + ) + if len(ids) == len(expected): + assert ids == expected, ( + f"KNN id order mismatch ({distance_metric} ground truth): " + f"got {ids!r} expected {expected!r}" + ) + elif check_top1 and expected: + nearest = nearest_knn_doc_ids( + corpus, query_vector, where, distance_metric=distance_metric + ) + assert ids[0] in nearest, ( + f"top-1 KNN must be among nearest neighbors ({distance_metric} tie set), " + f"got {ids[0]!r} expected one of {sorted(nearest)[:8]!r}" + f"{'…' if len(nearest) > 8 else ''}" ) if len(ids) == n_results: - worst = distances[-1] if distances else float("inf") - for rec in corpus: - if where is not None and not corpus_matches_where(rec, where): - continue - d = l2_squared(rec.embedding, query_vector) - if rec.doc_id not in ids and (not distances or d < worst - 1e-9): + scored = _knn_scored_rows( + corpus, query_vector, where, distance_metric=distance_metric + ) + if scored: + worst_d = max( + d for doc_id, d, _ in scored if doc_id in ids + ) + else: + worst_d = float("inf") + for doc_id, d, _ in scored: + if doc_id not in ids and d < worst_d - 1e-9: raise AssertionError( - f"closer match {rec.doc_id!r} (l2²={d}) missing from top-{n_results}" + f"closer match {doc_id!r} ({distance_metric}={d}) missing from top-{n_results}" ) @@ -263,6 +384,8 @@ def run_hybrid_knn_case( namespace: Any, corpus: list[CorpusRecord], case: VectorKnnCase, + *, + distance_metric: VectorDistanceMetric = "l2", ) -> dict[str, Any]: knn: dict[str, Any] = { "query_embeddings": case.query_vector, @@ -281,6 +404,7 @@ def run_hybrid_knn_case( case.query_vector, case.n_results, case.where, + distance_metric=distance_metric, check_top1=case.check_top1, ) return result @@ -290,6 +414,8 @@ def run_hybrid_combined_case( namespace: Any, corpus: list[CorpusRecord], case: HybridCombinedCase, + *, + distance_metric: VectorDistanceMetric = "l2", ) -> dict[str, Any]: query: dict[str, Any] | None = None if case.where_document is not None or case.where is not None: @@ -326,6 +452,7 @@ def run_hybrid_combined_case( vec, case.n_results, case.knn.get("where"), + distance_metric=distance_metric, check_top1=True, ) return result @@ -641,6 +768,8 @@ def assert_hybrid_triple_branch_result( corpus: list[CorpusRecord], result: dict[str, Any], case: HybridTripleBranchCase, + *, + distance_metric: VectorDistanceMetric = "l2", ) -> None: assert result is not None assert "ids" in result and result["ids"] @@ -757,7 +886,9 @@ def _knn_row(rec: CorpusRecord) -> bool: case.knn["query_embeddings"], min(len(knn_ids), case.n_results), knn_where, + distance_metric=distance_metric, check_top1=True, + check_score_order=False, ) return @@ -895,15 +1026,19 @@ def execute_hybrid_triple_branch_search( ) -def setup_large_fts_flat_collection(db_client: Any) -> tuple[list[CorpusRecord], Any]: +def setup_large_fts_flat_collection( + db_client: Any, + *, + distance: VectorDistanceMetric = "l2", +) -> tuple[list[CorpusRecord], Any]: """Create flat collection (HNSW + FTS), load via ``collection.add``.""" corpus = build_large_fts_corpus(CORPUS_SIZE) if len(corpus) <= 1000: raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") - name = f"test_hs_tb_baseline_{int(time.time() * 1000)}" + name = f"test_hs_tb_baseline_{distance}_{int(time.time() * 1000)}" collection = db_client.create_collection( - name=name, schema=flat_hybrid_search_schema(), use_namespace=False + name=name, schema=flat_hybrid_search_schema(distance), use_namespace=False ) assert collection.use_namespace is False insert_corpus_into_collection(collection, corpus) @@ -920,9 +1055,13 @@ def _try_assert_hybrid_triple_branch_result( corpus: list[CorpusRecord], result: dict[str, Any], case: HybridTripleBranchCase, + *, + distance_metric: VectorDistanceMetric = "l2", ) -> BaseException | None: try: - assert_hybrid_triple_branch_result(corpus, result, case) + assert_hybrid_triple_branch_result( + corpus, result, case, distance_metric=distance_metric + ) return None except BaseException as exc: return exc @@ -933,6 +1072,8 @@ def assert_hybrid_triple_branch_vs_collection_baseline( case: HybridTripleBranchCase, namespace_result: dict[str, Any], collection_result: dict[str, Any], + *, + distance_metric: VectorDistanceMetric = "l2", ) -> None: """ Compare namespace (IVF logic table) vs collection baseline (HNSW heap table). @@ -948,8 +1089,12 @@ def assert_hybrid_triple_branch_vs_collection_baseline( - ``[observer/namespace]`` — collection passes ground truth but namespace fails - ``[test-case/baseline]`` — namespace passes but collection baseline fails """ - coll_err = _try_assert_hybrid_triple_branch_result(corpus, collection_result, case) - ns_err = _try_assert_hybrid_triple_branch_result(corpus, namespace_result, case) + coll_err = _try_assert_hybrid_triple_branch_result( + corpus, collection_result, case, distance_metric=distance_metric + ) + ns_err = _try_assert_hybrid_triple_branch_result( + corpus, namespace_result, case, distance_metric=distance_metric + ) coll_ids = collection_result["ids"][0] ns_ids = namespace_result["ids"][0] @@ -999,6 +1144,7 @@ def run_hybrid_triple_branch_case( case: HybridTripleBranchCase, *, collection_baseline: Any | None = None, + distance_metric: VectorDistanceMetric = "l2", ) -> dict[str, Any]: if collection_baseline is not None: assert_flat_collection_baseline(collection_baseline) @@ -1020,12 +1166,15 @@ def run_hybrid_triple_branch_case( f"{ns_exec_err!r}" ) from ns_exec_err assert_hybrid_triple_branch_vs_collection_baseline( - corpus, case, namespace_result, collection_result + corpus, case, namespace_result, collection_result, + distance_metric=distance_metric, ) return namespace_result result = execute_hybrid_triple_branch_on_namespace(namespace, case) - assert_hybrid_triple_branch_result(corpus, result, case) + assert_hybrid_triple_branch_result( + corpus, result, case, distance_metric=distance_metric + ) return result @@ -1416,12 +1565,17 @@ def run_triple_branch_case_on_quadrants( quadrant_keys: tuple[str, ...] = MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, *, collection_baseline: Any | None = None, + distance_metric: VectorDistanceMetric = "l2", ) -> None: tb_case = case if isinstance(case, HybridTripleBranchCase) else get_triple_branch_case(case) for key in quadrant_keys: corpus, namespace = ctx[key] run_hybrid_triple_branch_case( - namespace, corpus, tb_case, collection_baseline=collection_baseline + namespace, + corpus, + tb_case, + collection_baseline=collection_baseline, + distance_metric=distance_metric, ) @@ -1440,11 +1594,15 @@ def run_knn_case_on_quadrants( ctx: dict[str, Any], case: VectorKnnCase | str, quadrant_keys: tuple[str, ...] = MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, + *, + distance_metric: VectorDistanceMetric = "l2", ) -> None: knn_case = case if isinstance(case, VectorKnnCase) else get_vector_knn_case(case) for key in quadrant_keys: corpus, namespace = ctx[key] - run_hybrid_knn_case(namespace, corpus, knn_case) + run_hybrid_knn_case( + namespace, corpus, knn_case, distance_metric=distance_metric + ) def assert_hybrid_search_index_no_hits( @@ -1502,7 +1660,11 @@ def assert_hybrid_triple_branch_no_hits( "TOKEN_ALP", "TOKEN_ZPX", "TRIPLE_BRANCH_CASES", + "VECTOR_DISTANCE_METRICS", "VECTOR_KNN_CASES", + "VectorDistanceMetric", + "ensure_shared_hybrid_search_collection", + "knn_compare_distance", "WHERE_BROAD_METADATA", "WHERE_DOCUMENT_UNIVERSAL", "WHERE_FILLER_SEQ", diff --git a/tests/integration_tests/test_namespace_get_delete_filters.py b/tests/integration_tests/test_namespace_get_delete_filters.py index 997ab4a4..45c36ca4 100644 --- a/tests/integration_tests/test_namespace_get_delete_filters.py +++ b/tests/integration_tests/test_namespace_get_delete_filters.py @@ -5,6 +5,7 @@ - delete(where=...) / delete(where_document=...) - get(where=...) / get(where_document=...) - multi-namespace isolation for metadata filter delete + - cross-namespace get(where=...) on a shared connection (regression) """ from __future__ import annotations @@ -84,8 +85,7 @@ def test_delete_where_document_obsolete(self, db_client): ) assert_get_where_count(ns, {"tag": "keep"}, gt["keep"], limit=50) finally: - # cleanup(db_client, collection) - pass + cleanup(db_client, collection) class TestNamespaceGetWhere: @@ -183,7 +183,7 @@ def test_delete_where_only_affects_target_namespace(self, db_client): assert_get_where_count( ns_b, {"tag": "purge"}, - gt["purge"], + 10, limit=10, check_meta={"tag": "purge", "ns_tag": "beta"}, ) @@ -198,5 +198,53 @@ def test_delete_where_only_affects_target_namespace(self, db_client): cleanup(db_client, collection) +class TestNamespaceCrossNsGetWhere: + """Regression: get(where=...) after switching namespace on one connection. + + Client sets @collection_id, @namespace_id, and @ltable_id before each DQL call; + OB logic-table kernel must honor the session context when metadata filters are used. + """ + + def test_get_where_alternates_namespaces_after_delete(self, db_client): + collection = create_ns_collection(db_client, suffix="_cross_ns_get") + ns_a = collection.create_namespace("cross_alpha") + ns_b = collection.create_namespace("cross_beta") + corpus_a, gt = build_filter_corpus( + keep_count=50, + purge_count=50, + id_prefix="alpha", + ns_tag="alpha", + ) + corpus_b, _ = build_filter_corpus( + keep_count=50, + purge_count=50, + id_prefix="beta", + ns_tag="beta", + ) + try: + load_dml_corpus(ns_a, corpus_a) + load_dml_corpus(ns_b, corpus_b) + ns_a.delete(where={"tag": "purge"}) + + assert_get_where_count(ns_a, {"tag": "purge"}, 0) + assert_get_where_count( + ns_b, + {"tag": "purge"}, + 10, + limit=10, + check_meta={"tag": "purge", "ns_tag": "beta"}, + ) + assert_get_where_count( + ns_a, + {"tag": "keep"}, + min(20, gt["keep"]), + limit=20, + check_meta={"tag": "keep", "ns_tag": "alpha"}, + ) + finally: + # cleanup(db_client, collection) + pass + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/integration_tests/test_namespace_hybrid_search_combined.py b/tests/integration_tests/test_namespace_hybrid_search_combined.py index 4c2389cf..7c5c74fe 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_combined.py +++ b/tests/integration_tests/test_namespace_hybrid_search_combined.py @@ -6,6 +6,8 @@ - vector + search index (``knn.where``) - vector + full-text (+ RRF smoke) - vector + full-text + search index (+ RRF smoke) + +Vector branches run under both L2 and cosine IVF index metrics. """ from __future__ import annotations @@ -17,30 +19,32 @@ from namespace_hybrid_search_helpers import ( HYBRID_COMBINED_CASES, + VectorDistanceMetric, + ensure_shared_hybrid_search_collection, get_hybrid_combined_case, run_hybrid_combined_case, setup_fts_namespace_with_corpus, - setup_large_fts_collection, teardown_large_fts_collection, ) +@pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) class TestNamespaceHybridSearchCombined: _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} @pytest.fixture(autouse=True) - def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: - mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" - if mode not in self._shared_by_mode: - corpus, collection = setup_large_fts_collection(db_client) - self._shared_by_mode[mode] = { - "corpus": corpus, - "collection": collection, - "db_client": db_client, - } - entry = self._shared_by_mode[mode] + def _bind_shared_collection( + self, + db_client: Any, + vector_distance: VectorDistanceMetric, + request: pytest.FixtureRequest, + ) -> None: + entry = ensure_shared_hybrid_search_collection( + self._shared_by_mode, db_client, request, vector_distance + ) self._corpus = entry["corpus"] self._collection = entry["collection"] + self._vector_distance = vector_distance @classmethod def teardown_class(cls) -> None: @@ -59,7 +63,9 @@ def _new_namespace(self, case_name: str) -> Any: def test_hybrid_search_combined(self, db_client, case_name: str): case = get_hybrid_combined_case(case_name) namespace = self._new_namespace(case_name) - run_hybrid_combined_case(namespace, self._corpus, case) + run_hybrid_combined_case( + namespace, self._corpus, case, distance_metric=self._vector_distance + ) if __name__ == "__main__": diff --git a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py index fc02a20f..80097b18 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py @@ -10,6 +10,7 @@ from namespace_fts_helpers import ( MULTI_COLL_MULTI_NS_QUADRANT_KEYS, + VectorDistanceMetric, assert_hybrid_search_no_hits, get_fts_case, run_hybrid_search_fts_case, @@ -100,11 +101,14 @@ def test_hybrid_search_search_index_all_loaded_quadrants(self, db_client, case_n finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) @pytest.mark.parametrize("case_name", ["knn_global_top5", "knn_filter_has_both"]) - def test_hybrid_search_vector_all_loaded_quadrants(self, db_client, case_name: str): - ctx = setup_multi_coll_multi_ns_fts(db_client) + def test_hybrid_search_vector_all_loaded_quadrants( + self, db_client, vector_distance: VectorDistanceMetric, case_name: str + ): + ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: - run_knn_case_on_quadrants(ctx, case_name) + run_knn_case_on_quadrants(ctx, case_name, distance_metric=vector_distance) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) @@ -119,23 +123,33 @@ def test_hybrid_search_fts_plus_search_index_multi_coll(self, db_client): finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) - def test_hybrid_search_combined_rrf_multi_coll(self, db_client): - ctx = setup_multi_coll_multi_ns_fts(db_client) + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) + def test_hybrid_search_combined_rrf_multi_coll( + self, db_client, vector_distance: VectorDistanceMetric + ): + ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: case = get_hybrid_combined_case("fts_zpx_filter_gte_knn") for key in ("c1_x", "c2_y"): corpus, namespace = ctx[key] - run_hybrid_combined_case(namespace, corpus, case) + run_hybrid_combined_case( + namespace, corpus, case, distance_metric=vector_distance + ) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) - def test_hybrid_search_vector_plus_search_index_multi_coll(self, db_client): - ctx = setup_multi_coll_multi_ns_fts(db_client) + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) + def test_hybrid_search_vector_plus_search_index_multi_coll( + self, db_client, vector_distance: VectorDistanceMetric + ): + ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: knn_case = get_vector_knn_case("knn_filter_has_both") for key in ("c1_x", "c2_y"): corpus, namespace = ctx[key] - run_hybrid_knn_case(namespace, corpus, knn_case) + run_hybrid_knn_case( + namespace, corpus, knn_case, distance_metric=vector_distance + ) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py index 3a64cd9c..d60ba04b 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py @@ -11,10 +11,12 @@ filters are checked; fused FTS+KNN ``__score`` order is not asserted (see ``check_fts_ranking`` on ``HybridTripleBranchCase``). +Vector index distance: L2 and cosine (IVF spfresh). + Run one case:: pytest tests/integration_tests/test_namespace_hybrid_search_triple_branch.py \\ - -k "fts_contains_zpx and oceanbase" -v -s + -k "fts_contains_zpx and oceanbase and l2" -v -s """ from __future__ import annotations @@ -27,32 +29,34 @@ from namespace_hybrid_search_helpers import ( TRIPLE_BRANCH_CASES, TOKEN_ZPX, + VectorDistanceMetric, + ensure_shared_hybrid_search_collection, get_triple_branch_case, run_hybrid_triple_branch_case, setup_fts_namespace_with_corpus, - setup_large_fts_collection, teardown_large_fts_collection, ) +@pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) class TestNamespaceHybridSearchTripleBranch: """Large-scale namespace vector + full-text + search-index hybrid_search tests.""" _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} @pytest.fixture(autouse=True) - def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: - mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" - if mode not in self._shared_by_mode: - corpus, collection = setup_large_fts_collection(db_client) - self._shared_by_mode[mode] = { - "db_client": db_client, - "corpus": corpus, - "collection": collection, - } - entry = self._shared_by_mode[mode] + def _bind_shared_collection( + self, + db_client: Any, + vector_distance: VectorDistanceMetric, + request: pytest.FixtureRequest, + ) -> None: + entry = ensure_shared_hybrid_search_collection( + self._shared_by_mode, db_client, request, vector_distance + ) self._corpus = entry["corpus"] self._collection = entry["collection"] + self._vector_distance = vector_distance @classmethod def teardown_class(cls) -> None: @@ -71,7 +75,12 @@ def _new_namespace(self, case_name: str) -> Any: def test_hybrid_search_triple_branch(self, db_client, case_name: str): case = get_triple_branch_case(case_name) namespace = self._new_namespace(case_name) - run_hybrid_triple_branch_case(namespace, self._corpus, case) + run_hybrid_triple_branch_case( + namespace, + self._corpus, + case, + distance_metric=self._vector_distance, + ) def test_hybrid_search_triple_branch_fts_hits_with_knn_active(self, db_client): """FTS+KNN hybrid returns TOKEN_ZPX hits; fused order is not asserted.""" diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py index b57e5960..355eef53 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py @@ -22,6 +22,7 @@ MULTI_COLL_MULTI_NS_QUADRANT_KEYS, TOKEN_ZPX, CorpusRecord, + VectorDistanceMetric, build_large_fts_corpus, insert_corpus_in_batches, ns_schema, @@ -72,17 +73,21 @@ def _namespace_corpus(base_corpus: list[CorpusRecord], key: str) -> list[CorpusR ] -def _setup_multi_coll_multi_ns_isolation_triple(db_client): +def _setup_multi_coll_multi_ns_isolation_triple( + db_client, + *, + distance: VectorDistanceMetric = "l2", +): """2 collections x 2 loaded namespaces with distinct per-namespace corpus markers.""" ts = int(time.time() * 1000) coll_1 = db_client.create_collection( - name=f"test_ns_hs_tb_mcmn_iso_{ts}_c1", - schema=ns_schema(), + name=f"test_ns_hs_tb_mcmn_iso_{distance}_{ts}_c1", + schema=ns_schema(distance), use_namespace=True, ) coll_2 = db_client.create_collection( - name=f"test_ns_hs_tb_mcmn_iso_{ts}_c2", - schema=ns_schema(), + name=f"test_ns_hs_tb_mcmn_iso_{distance}_{ts}_c2", + schema=ns_schema(distance), use_namespace=True, ) ctx: dict[str, object] = {"coll_1": coll_1, "coll_2": coll_2} @@ -106,14 +111,24 @@ def _setup_multi_coll_multi_ns_isolation_triple(db_client): class TestNamespaceHybridSearchTripleBranchMultiCollMultiNs: - def test_cross_quadrant_triple_branch_isolation(self, db_client): + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) + def test_cross_quadrant_triple_branch_isolation( + self, db_client, vector_distance: VectorDistanceMetric + ): """Loaded namespaces return only their own rows; empty namespaces stay empty.""" - ctx = _setup_multi_coll_multi_ns_isolation_triple(db_client) + ctx = _setup_multi_coll_multi_ns_isolation_triple( + db_client, distance=vector_distance + ) try: case = get_triple_branch_case("intersection_both_tokens") for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: corpus, namespace = ctx[key] - result = run_hybrid_triple_branch_case(namespace, corpus, case) + result = run_hybrid_triple_branch_case( + namespace, + corpus, + case, + distance_metric=vector_distance, + ) ids = result["ids"][0] assert ids assert all(doc_id.startswith(f"{key}_") for doc_id in ids), ( @@ -140,8 +155,13 @@ def test_cross_quadrant_search_index_branch_isolation(self, db_client): finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) - def test_triple_branch_top1_fts_all_quadrants(self, db_client): - ctx = _setup_multi_coll_multi_ns_isolation_triple(db_client) + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) + def test_triple_branch_top1_fts_all_quadrants( + self, db_client, vector_distance: VectorDistanceMetric + ): + ctx = _setup_multi_coll_multi_ns_isolation_triple( + db_client, distance=vector_distance + ) try: for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: _, namespace = ctx[key] @@ -183,40 +203,61 @@ def test_hybrid_search_triple_branch_search_index_all_loaded_quadrants( finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) @pytest.mark.parametrize("case_name", _KNN_TRIPLE_CASE_NAMES) - def test_hybrid_search_triple_branch_knn_all_loaded_quadrants(self, db_client, case_name: str): - ctx = setup_multi_coll_multi_ns_fts(db_client) + def test_hybrid_search_triple_branch_knn_all_loaded_quadrants( + self, db_client, vector_distance: VectorDistanceMetric, case_name: str + ): + ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: - run_triple_branch_case_on_quadrants(ctx, case_name) + run_triple_branch_case_on_quadrants( + ctx, case_name, distance_metric=vector_distance + ) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) @pytest.mark.parametrize("case_name", _INTERSECTION_CASE_NAMES) def test_hybrid_search_triple_branch_intersection_all_loaded_quadrants( - self, db_client, case_name: str + self, db_client, vector_distance: VectorDistanceMetric, case_name: str ): - ctx = setup_multi_coll_multi_ns_fts(db_client) + ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: - run_triple_branch_case_on_quadrants(ctx, case_name) + run_triple_branch_case_on_quadrants( + ctx, case_name, distance_metric=vector_distance + ) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) @pytest.mark.parametrize("case_name", _RRF_CASE_NAMES) - def test_hybrid_search_triple_branch_rrf_all_loaded_quadrants(self, db_client, case_name: str): - ctx = setup_multi_coll_multi_ns_fts(db_client) + def test_hybrid_search_triple_branch_rrf_all_loaded_quadrants( + self, db_client, vector_distance: VectorDistanceMetric, case_name: str + ): + ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: - run_triple_branch_case_on_quadrants(ctx, case_name) + run_triple_branch_case_on_quadrants( + ctx, case_name, distance_metric=vector_distance + ) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) - def test_hybrid_search_triple_branch_intersection_subset_on_quadrants(self, db_client): + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) + def test_hybrid_search_triple_branch_intersection_subset_on_quadrants( + self, db_client, vector_distance: VectorDistanceMetric + ): """Every hit must lie in the triple-branch intersection on each loaded quadrant.""" - ctx = setup_multi_coll_multi_ns_fts(db_client) + ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: case = get_triple_branch_case("intersection_zpx_gte40") for key in MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS: corpus, namespace = ctx[key] - result = run_hybrid_triple_branch_case(namespace, corpus, case) + result = run_hybrid_triple_branch_case( + namespace, + corpus, + case, + distance_metric=vector_distance, + ) corpus_by_id = {rec.doc_id: rec for rec in corpus} for doc_id in result["ids"][0]: assert corpus_matches_triple_intersection(corpus_by_id[doc_id], case), ( diff --git a/tests/integration_tests/test_namespace_hybrid_search_vector.py b/tests/integration_tests/test_namespace_hybrid_search_vector.py index 69d54d3b..cafef9e0 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_vector.py +++ b/tests/integration_tests/test_namespace_hybrid_search_vector.py @@ -2,6 +2,7 @@ Namespace hybrid_search pure vector (KNN) integration tests. Validates KNN distance ordering and metadata filters on the ``knn`` branch. +Covers both L2 and cosine IVF index distance metrics. """ from __future__ import annotations @@ -14,32 +15,35 @@ from namespace_hybrid_search_helpers import ( KNN_QUERY_VECTOR, VECTOR_KNN_CASES, + VectorDistanceMetric, + ensure_shared_hybrid_search_collection, + expected_knn_ids, get_vector_knn_case, run_hybrid_knn_case, setup_fts_namespace_with_corpus, - setup_large_fts_collection, teardown_large_fts_collection, ) +@pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) class TestNamespaceHybridSearchVector: """Large-scale namespace pure KNN hybrid_search tests.""" _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} @pytest.fixture(autouse=True) - def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: - mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" - if mode not in self._shared_by_mode: - corpus, collection = setup_large_fts_collection(db_client) - self._shared_by_mode[mode] = { - "db_client": db_client, - "corpus": corpus, - "collection": collection, - } - entry = self._shared_by_mode[mode] + def _bind_shared_collection( + self, + db_client: Any, + vector_distance: VectorDistanceMetric, + request: pytest.FixtureRequest, + ) -> None: + entry = ensure_shared_hybrid_search_collection( + self._shared_by_mode, db_client, request, vector_distance + ) self._corpus = entry["corpus"] self._collection = entry["collection"] + self._vector_distance = vector_distance @classmethod def teardown_class(cls) -> None: @@ -57,7 +61,12 @@ def _new_namespace(self, case_name: str) -> Any: def _run_knn_case(self, case_name: str) -> None: namespace = self._new_namespace(case_name) - run_hybrid_knn_case(namespace, self._corpus, get_vector_knn_case(case_name)) + run_hybrid_knn_case( + namespace, + self._corpus, + get_vector_knn_case(case_name), + distance_metric=self._vector_distance, + ) @pytest.mark.parametrize("case_name", [c.name for c in VECTOR_KNN_CASES]) def test_hybrid_search_vector_knn_cases(self, db_client, case_name: str): @@ -66,14 +75,17 @@ def test_hybrid_search_vector_knn_cases(self, db_client, case_name: str): def test_hybrid_search_vector_top1_nearest(self, db_client): """Top-1 must be the global nearest neighbor for a fixed query vector.""" namespace = self._new_namespace("top1_nearest") - from namespace_hybrid_search_helpers import expected_knn_ids - result = namespace.hybrid_search( knn={"query_embeddings": KNN_QUERY_VECTOR, "n_results": 1}, n_results=1, include=["distances"], ) - assert result["ids"][0][0] == expected_knn_ids(self._corpus, KNN_QUERY_VECTOR, 1)[0] + assert result["ids"][0][0] == expected_knn_ids( + self._corpus, + KNN_QUERY_VECTOR, + 1, + distance_metric=self._vector_distance, + )[0] if __name__ == "__main__": From d67cbfb44921022bc112f712dacb6517dafb9273 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 4 Jun 2026 14:28:40 +0800 Subject: [PATCH 31/84] =?UTF-8?q?prewarm=5Flob=20sdk=E9=80=82=E9=85=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 2 +- .../test_namespace_prewarm_lob.py | 387 ++++++++++++++++++ 2 files changed, 388 insertions(+), 1 deletion(-) create mode 100644 tests/integration_tests/test_namespace_prewarm_lob.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7632670c..ba9ec58b 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1341,7 +1341,7 @@ def _create_namespace_physical_tables( kv_key VARBINARY(1024) NOT NULL, kv_value LONGBLOB NOT NULL, PRIMARY KEY(namespace_id, kv_key) - ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 LOB_INROW_THRESHOLD = 786432 + ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 LOB_INROW_THRESHOLD=786432 {partition_clause}""") self._execute(f"""CREATE TABLE `{schema_table}` ( diff --git a/tests/integration_tests/test_namespace_prewarm_lob.py b/tests/integration_tests/test_namespace_prewarm_lob.py new file mode 100644 index 00000000..72c8d520 --- /dev/null +++ b/tests/integration_tests/test_namespace_prewarm_lob.py @@ -0,0 +1,387 @@ +""" +LOB prewarm integration tests (OceanBase shared-storage mode). + +Validates the OB-side change that, after prewarming a namespace's KV data tablet, +the prewarm RPC ALSO prewarms that KV tablet's aux LOB-meta tablet (whole range). + +Observability is layered into three independent tiers so the suite degrades +gracefully depending on what the cluster exposes: + + Tier 1 - structural (always): + Resolve the kv_data_table's aux LOB-meta tablet purely from SQL + (__all_table.data_table_id + table_type=13 -> __all_virtual_tablet_to_ls). + Proves the table the prewarm targets actually owns a LOB-meta tablet. + + Tier 2 - ss-cache (SS mode + real out-of-row LOB): + Force a >0.75MB kv_value so it spills out-of-row into the LOB-meta tablet, + freeze it to a macro block, prewarm, then read + __all_virtual_ss_macro_cache_info filtered by that LOB-meta tablet_id. + Proves blocks were actually pulled into the local cache. + + Tier 3 - log marker (opt-in via OB_OBSERVER_LOG): + Grep observer.log for "prewarm_logic_table lob meta done", which prints the + exact lob_meta_tablet_id. Proves the new code path executed for the right tablet. + +Run: + OB_PORT=11202 OB_TENANT=mysql \ + OB_OBSERVER_LOG=/path/to/observer.log \ + pytest tests/integration_tests/test_namespace_prewarm_lob.py -v -s +""" + +import contextlib +import os +import time + +import pytest + +try: + import pymysql +except ImportError: # pragma: no cover - pymysql ships with the test deps + pymysql = None + +from pyseekdb import ( + IVFConfiguration, + get_namespace_partition_count, + set_namespace_partition_count, +) +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.meta_info import NamespaceCollectionNames +from pyseekdb.client.schema import Schema + +AUX_LOB_META_TABLE_TYPE = 13 +LARGE_BLOB_BYTES = 1_000_000 # > OB_MAX_LOB_INROW_THRESHOLD (786432) -> always out-of-row +LOB_PROBE_ROWS = 3 # 3 x ~1MB incompressible -> ~3MB of real macro blocks +SS_CACHE_VIEW = "oceanbase.__all_virtual_ss_macro_cache_info" +ALL_TABLE_VIEW = "oceanbase.__all_table" +TABLET_TO_LS_VIEW = "oceanbase.__all_tablet_to_ls" + +# Flushing the SS local cache requires a sys-tenant connection with an explicit +# TENANT clause; user-tenant flush returns OB-4016. These are env-overridable so +# the rigorous before/after assertion can run wherever a sys login is available. +OB_HOST = os.environ.get("OB_HOST", "127.0.0.1") +OB_PORT = int(os.environ.get("OB_PORT", "10902")) +OB_TENANT = os.environ.get("OB_TENANT", "test_tenant") +OB_SYS_USER = os.environ.get("OB_SYS_USER", "root@sys") +OB_SYS_PASSWORD = os.environ.get("OB_SYS_PASSWORD", "") + + +# ==================== SQL observability helpers ==================== + +def _exec(client, sql): + return client._server._execute(sql) + + +def _resolve_table_id(client, table_name): + rows = _exec( + client, + f"SELECT table_id FROM {ALL_TABLE_VIEW} WHERE table_name = '{table_name}'", + ) + assert rows, f"table_id not found for {table_name}" + return rows[0]["table_id"] + + +def _resolve_lob_meta_tablets(client, kv_table_id): + """All aux-LOB-meta tablets of a KV data table (one per KV partition).""" + rows = _exec( + client, + f"SELECT table_id FROM {ALL_TABLE_VIEW} " + f"WHERE data_table_id = {kv_table_id} AND table_type = {AUX_LOB_META_TABLE_TYPE}", + ) + assert rows, f"no aux lob meta table for kv table {kv_table_id}" + lob_meta_table_id = rows[0]["table_id"] + + tablet_rows = _exec( + client, + f"SELECT tablet_id FROM {TABLET_TO_LS_VIEW} WHERE table_id = {lob_meta_table_id}", + ) + return {r["tablet_id"] for r in tablet_rows} + + +def _ss_cache_blocks_for_tablets(client, tablet_ids): + """Map tablet_id -> (block_count, total_size) from the SS macro cache view. + + Returns None when the SS cache view is not available (non-SS cluster), so + callers can skip tier-2 assertions cleanly. + """ + if not tablet_ids: + return {} + ids = ",".join(str(t) for t in tablet_ids) + try: + rows = _exec( + client, + f"SELECT tablet_id, COUNT(*) AS blk, COALESCE(SUM(size),0) AS sz " + f"FROM {SS_CACHE_VIEW} WHERE tablet_id IN ({ids}) GROUP BY tablet_id", + ) + except Exception: + return None + return {r["tablet_id"]: (r["blk"], r["sz"]) for r in rows} + + +def _flush_tenant_macro_cache(): + """Evict the tenant's SS local macro cache via a sys-tenant connection. + + Returns True on success, False if a sys login / flush is unavailable so the + caller can skip the rigorous tier-2 assertion instead of failing spuriously. + """ + if pymysql is None: + return False + try: + conn = pymysql.connect( + host=OB_HOST, port=OB_PORT, user=OB_SYS_USER, + password=OB_SYS_PASSWORD, autocommit=True, + ) + except Exception: + return False + try: + with conn.cursor() as cur: + cur.execute( + f"ALTER SYSTEM FLUSH SS_LOCAL_CACHE TENANT = {OB_TENANT} CACHE = macro_cache" + ) + return True + except Exception: + return False + finally: + with contextlib.suppress(Exception): + conn.close() + + +def _is_ss_mode(client, collection_id): + """SS mode is detectable by the per-collection hot_table existing.""" + hot_table = NamespaceCollectionNames.hot_table_name(collection_id) + rows = _exec( + client, + f"SELECT table_id FROM {ALL_TABLE_VIEW} WHERE table_name = '{hot_table}'", + ) + return bool(rows) + + +def _observer_log_files(): + """The configured observer.log plus any rotated siblings in the same dir. + + OB rotates observer.log aggressively (observer.log.), so a marker emitted + during the test may already have moved to a rotated file by the time we grep. + """ + log_path = os.environ.get("OB_OBSERVER_LOG") + if not log_path or not os.path.exists(log_path): + return None + directory = os.path.dirname(log_path) or "." + base = os.path.basename(log_path) + files = [] + for fn in os.listdir(directory): + # observer.log and observer.log., but not observer.log.wf* + if (fn == base or fn.startswith(base + ".")) and ".wf" not in fn: + files.append(os.path.join(directory, fn)) + return files + + +def _grep_lob_prewarm_log(lob_meta_tablet_ids): + """Tier 3: confirm the new code path logged a lob prewarm for one of the tablets.""" + files = _observer_log_files() + if files is None: + return None + wanted = {str(t) for t in lob_meta_tablet_ids} + hits = [] + for path in files: + with open(path, errors="ignore") as fh: + for line in fh: + if "prewarm_logic_table lob meta done" in line and any(t in line for t in wanted): + hits.append(line.strip()) + return hits + + +# ==================== Fixtures / collection helpers ==================== + +class _BaseLobPrewarm: + + def _make_collection(self, client, name, partitions): + set_namespace_partition_count(partitions) + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine"), + embedding_function=None, + ), + ) + return client.create_collection(name=name, schema=schema, use_namespace=True) + + def _force_out_of_row_lob(self, client, collection_id, namespace_id): + """Insert several >0.75MB incompressible kv_values so they spill out-of-row + into the LOB-meta tablet, then freeze them into macro blocks. + + REPEAT('x', N) compresses to ~nothing and yields a single tiny block, which + is useless for proving prewarm restored real data. os.urandom is + incompressible, so each row produces a genuine ~1MB macro block. + """ + kv_table = NamespaceCollectionNames.kv_data_table_name(collection_id) + # _execute() takes no params, so use the raw DBAPI connection to bind the + # binary blob safely (avoids building multi-MB hex-literal SQL strings). + raw = client._server.get_raw_connection() + with raw.cursor() as cur: + for i in range(LOB_PROBE_ROWS): + cur.execute( + f"INSERT INTO `{kv_table}` (namespace_id, kv_key, kv_value) " + f"VALUES ({namespace_id}, %s, %s)", + (f"__lob_probe_{i}__", os.urandom(LARGE_BLOB_BYTES)), + ) + raw.commit() + # Push the rows into macro blocks so prewarm has remote data to pull. + with contextlib.suppress(Exception): + _exec(client, "ALTER SYSTEM MINOR FREEZE") + + +# ==================== Tier 1 + structural multi-namespace ==================== + +class TestLobPrewarmStructural(_BaseLobPrewarm): + + @pytest.fixture(autouse=True) + def _restore_partitions(self): + original = get_namespace_partition_count() + yield + set_namespace_partition_count(original) + + def test_kv_table_has_lob_meta_tablet(self, oceanbase_client): + """Tier 1: the prewarm target (kv_data_table) owns an aux LOB-meta tablet.""" + name = f"lob_pw_struct_{int(time.time() * 1000)}" + collection = self._make_collection(oceanbase_client, name, partitions=1) + try: + kv_table = NamespaceCollectionNames.kv_data_table_name(collection.id) + kv_table_id = _resolve_table_id(oceanbase_client, kv_table) + lob_tablets = _resolve_lob_meta_tablets(oceanbase_client, kv_table_id) + assert len(lob_tablets) == 1, ( + f"single-partition kv table should have exactly 1 lob meta tablet, " + f"got {lob_tablets}" + ) + finally: + oceanbase_client.delete_collection(name=collection.name) + + def test_multi_namespace_prewarm_is_independent(self, oceanbase_client): + """Multi-namespace: prewarming each namespace succeeds, records its own hot_table + row, and never disturbs another namespace's row. Cached LOB tablets stay within + the collection's own LOB-meta tablet set (no foreign tablets).""" + name = f"lob_pw_multi_{int(time.time() * 1000)}" + collection = self._make_collection(oceanbase_client, name, partitions=8) + try: + kv_table = NamespaceCollectionNames.kv_data_table_name(collection.id) + kv_table_id = _resolve_table_id(oceanbase_client, kv_table) + own_lob_tablets = _resolve_lob_meta_tablets(oceanbase_client, kv_table_id) + + namespaces = [collection.create_namespace(f"ns_{i}") for i in range(3)] + for ns in namespaces: + ns.prewarm() + + hot_table = NamespaceCollectionNames.hot_table_name(collection.id) + rows = _exec( + oceanbase_client, + f"SELECT namespace_id FROM `{hot_table}`", + ) + recorded = {int(r["namespace_id"]) for r in rows} + expected = {int(ns._namespace_id) for ns in namespaces} + assert expected.issubset(recorded), ( + f"every prewarmed namespace must have a hot_table row: " + f"expected {expected}, got {recorded}" + ) + + cache = _ss_cache_blocks_for_tablets(oceanbase_client, own_lob_tablets) + if cache: + assert set(cache).issubset(own_lob_tablets), ( + "cached lob tablets must belong to this collection only" + ) + finally: + oceanbase_client.delete_collection(name=collection.name) + + +# ==================== Tier 2 + Tier 3: real LOB caching ==================== + +class TestLobPrewarmCaching(_BaseLobPrewarm): + + @pytest.fixture(autouse=True) + def _restore_partitions(self): + original = get_namespace_partition_count() + yield + set_namespace_partition_count(original) + + @staticmethod + def _total_bytes(cache): + return sum(v[1] for v in (cache or {}).values()) + + def test_out_of_row_lob_is_prewarmed_into_cache(self, oceanbase_client): + """Tier 2/3: prove prewarm actually *pulls* evicted LOB data back, not that + the write cache merely still holds it. + + Methodology (the only rigorous one): + 1. Insert incompressible out-of-row LOB rows + freeze -> real macro blocks. + 2. Flush the tenant's SS macro cache (sys connection) -> evict those blocks. + 3. Assert cache bytes for the LOB tablet dropped near-zero BEFORE prewarm. + 4. prewarm() -> poll the cache view. + 5. Assert cache bytes were restored (>> the flushed floor) AFTER prewarm. + + Without the flush step, frozen data already sits in the local write cache, so a + plain "blocks > 0 after prewarm" assertion would pass even if prewarm did nothing. + """ + name = f"lob_pw_cache_{int(time.time() * 1000)}" + collection = self._make_collection(oceanbase_client, name, partitions=1) + try: + if not _is_ss_mode(oceanbase_client, collection.id): + pytest.skip("LOB cache prewarm only observable in shared-storage mode") + + namespace = collection.create_namespace("lob_ns") + kv_table = NamespaceCollectionNames.kv_data_table_name(collection.id) + kv_table_id = _resolve_table_id(oceanbase_client, kv_table) + lob_tablets = _resolve_lob_meta_tablets(oceanbase_client, kv_table_id) + + self._force_out_of_row_lob(oceanbase_client, collection.id, namespace._namespace_id) + + if _ss_cache_blocks_for_tablets(oceanbase_client, lob_tablets) is None: + pytest.skip(f"{SS_CACHE_VIEW} not available on this cluster") + + # Give the freeze time to upload macro blocks to object storage. + time.sleep(10) + written = _ss_cache_blocks_for_tablets(oceanbase_client, lob_tablets) or {} + assert self._total_bytes(written) > 0, ( + f"out-of-row LOB should have produced macro blocks in cache, got {written}" + ) + + # Evict so prewarm has something real to restore. Requires sys flush. + if not _flush_tenant_macro_cache(): + pytest.skip( + "sys-tenant SS_LOCAL_CACHE flush unavailable; cannot prove " + "prewarm pulls evicted data (set OB_SYS_USER/OB_SYS_PASSWORD)" + ) + time.sleep(4) + + flushed = _ss_cache_blocks_for_tablets(oceanbase_client, lob_tablets) or {} + flushed_bytes = self._total_bytes(flushed) + written_bytes = self._total_bytes(written) + assert flushed_bytes < written_bytes, ( + f"flush should have evicted LOB macro blocks before prewarm: " + f"written={written_bytes}, flushed={flushed_bytes}" + ) + + namespace.prewarm() + + # Prewarm pulls remote macro blocks asynchronously; poll for restoration. + cached = flushed + deadline = time.time() + 40 + while time.time() < deadline: + cached = _ss_cache_blocks_for_tablets(oceanbase_client, lob_tablets) or {} + if self._total_bytes(cached) > flushed_bytes: + break + time.sleep(2) + + log_hits = _grep_lob_prewarm_log(lob_tablets) + if log_hits is not None: + assert log_hits, ( + "observer.log should contain a 'prewarm_logic_table lob meta done' " + f"line for lob tablets {lob_tablets}" + ) + + restored_bytes = self._total_bytes(cached) + assert restored_bytes > flushed_bytes, ( + f"prewarm should have restored evicted LOB macro blocks: " + f"written={written_bytes}, flushed={flushed_bytes}, restored={restored_bytes}" + ) + finally: + oceanbase_client.delete_collection(name=collection.name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) From c564ff99e55aa62e1e1682de400a6c6d4e8e34d7 Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Sat, 6 Jun 2026 16:36:17 +0800 Subject: [PATCH 32/84] fix case --- tests/integration_tests/test_namespace_get_delete_filters.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/integration_tests/test_namespace_get_delete_filters.py b/tests/integration_tests/test_namespace_get_delete_filters.py index 45c36ca4..61130084 100644 --- a/tests/integration_tests/test_namespace_get_delete_filters.py +++ b/tests/integration_tests/test_namespace_get_delete_filters.py @@ -83,7 +83,7 @@ def test_delete_where_document_obsolete(self, db_client): substring="obsolete", context="post-delete", ) - assert_get_where_count(ns, {"tag": "keep"}, gt["keep"], limit=50) + assert_get_where_count(ns, {"tag": "keep"}, gt["keep"], limit=gt["keep"]) finally: cleanup(db_client, collection) @@ -242,8 +242,7 @@ def test_get_where_alternates_namespaces_after_delete(self, db_client): check_meta={"tag": "keep", "ns_tag": "alpha"}, ) finally: - # cleanup(db_client, collection) - pass + cleanup(db_client, collection) if __name__ == "__main__": From c77f23b4d68bbf1c634f46d3e3f3884a17a5b45e Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 8 Jun 2026 11:51:15 +0800 Subject: [PATCH 33/84] =?UTF-8?q?=E4=BF=AEquery=E9=87=8C=E9=9D=A2=E6=9C=89?= =?UTF-8?q?=E5=90=91=E9=87=8F=E6=9F=A5=E8=AF=A2SDK=E8=A6=81=E8=BD=AC?= =?UTF-8?q?=E6=88=90DSL?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 206 ++++++++++++----------------- tests/unit_tests/test_namespace.py | 66 +++++++-- 2 files changed, 136 insertions(+), 136 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index ba9ec58b..9c83d9ca 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -125,6 +125,9 @@ def _validate_collection_name(name: str) -> None: from .validators import _MAX_NAMESPACE_BATCH_SIZE, _validate_namespace_name, _validate_record_ids # noqa: F401 _NS_PARTITION_COUNT = 1000 +# Unquoted id for WHERE/CASE; plain JSON_EXTRACT returns a quoted JSON string and +# can route through SEARCH INDEX on SS logic tables, breaking cross-namespace id lookups. +_NS_DATA_CONTENT_ID_EXPR = "JSON_UNQUOTE(JSON_EXTRACT(data_content, '$.id'))" def set_namespace_partition_count(n: int) -> None: @@ -1298,18 +1301,19 @@ def _create_namespace_physical_tables( try: self._execute(f"CREATE TABLEGROUP `{tg_name}` SHARDING='ADAPTIVE'") - index_clauses = [ - f"FULLTEXT INDEX idx_fts(document) {fulltext_clause}", - "SEARCH INDEX idx_json(data_content)", - ] - # Vector indexes are rejected in shared-storage mode on this kernel branch - # (both inline `VECTOR INDEX` in CREATE TABLE and standalone `CREATE VECTOR - # INDEX` fail with OB-1235 "vector index in shared storage mode is not - # supported"). Keep the embedding column but skip the vector index in SS; - # full-text and metadata SEARCH INDEX hybrid_search still work. - if not is_shared_storage: - index_clauses.append(f"VECTOR INDEX idx_vec(embedding) {vector_index_sql}") - index_sql = ",\n ".join(index_clauses) + # SN: post-create ``CREATE VECTOR INDEX`` on logic tables (documented path). + # SS: logic tables reject post-create vector indexes (OB-5703); inline only. + if is_shared_storage: + index_sql = ( + f"FULLTEXT INDEX idx_fts(document) {fulltext_clause},\n" + f" SEARCH INDEX idx_json(data_content),\n" + f" VECTOR INDEX idx_vec(embedding) {vector_index_sql}" + ) + else: + index_sql = ( + f"FULLTEXT INDEX idx_fts(document) {fulltext_clause},\n" + f" SEARCH INDEX idx_json(data_content)" + ) data_sql = f"""CREATE TABLE `{data_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, @@ -1325,6 +1329,12 @@ def _create_namespace_physical_tables( self._execute(data_sql) + if not is_shared_storage: + ivf_index_sql = ( + f"CREATE VECTOR INDEX idx_vec ON `{data_table}` (embedding) {vector_index_sql}" + ) + self._execute(ivf_index_sql) + if is_shared_storage: hot_table = NamespaceCollectionNames.hot_table_name(collection_id) self._execute(f"""CREATE TABLE `{hot_table}` ( @@ -4844,7 +4854,7 @@ def _namespace_update( table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) - id_expr = "JSON_EXTRACT(data_content, '$.id')" + id_expr = _NS_DATA_CONTENT_ID_EXPR active_ids = [] for i, record_id in enumerate(ids): has_update = ( @@ -4975,10 +4985,10 @@ def _namespace_upsert( ns_id = int(namespace_id) existing_ids = set() - id_expr = "JSON_EXTRACT(data_content, '$.id')" + id_expr = _NS_DATA_CONTENT_ID_EXPR id_placeholders = ", ".join(["%s"] * len(ids)) check_sql = ( - f"SELECT {id_expr} AS rid FROM `{table_name}` " + f"SELECT JSON_EXTRACT(data_content, '$.id') AS rid FROM `{table_name}` " f"WHERE namespace_id = {ns_id} AND ltable_id = {ltable_id} " f"AND {id_expr} IN ({id_placeholders})" ) @@ -5051,7 +5061,7 @@ def _namespace_delete( if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) - id_placeholders = " OR ".join(["JSON_EXTRACT(data_content, '$.id') = %s"] * len(ids)) + id_placeholders = " OR ".join([f"{_NS_DATA_CONTENT_ID_EXPR} = %s"] * len(ids)) conditions.append(f"({id_placeholders})") params.extend(ids) @@ -5120,123 +5130,63 @@ def _namespace_query( # noqa: C901 query_embeddings = self._normalize_query_embeddings(query_embeddings) include_fields = self._normalize_include_fields(include) - table_name = NamespaceCollectionNames.data_table_name(collection_id) - ns_id = int(namespace_id) - - select_parts = ["JSON_EXTRACT(data_content, '$.id') AS record_id"] - if include_fields.get("documents") or include_fields.get("document") or include is None: - select_parts.append("document") - if include_fields.get("metadatas") or include_fields.get("metadata") or include is None: - select_parts.append("JSON_EXTRACT(data_content, '$.metadata') AS metadata") - if include_fields.get("embeddings") or include_fields.get("embedding"): - select_parts.append("embedding") - - user_conditions = [] - filter_params = [] - - if where is not None: - rewritten = self._rewrite_where_for_ns(where) - meta_clause, meta_params = FilterBuilder.build_metadata_filter(rewritten, "data_content") - if meta_clause: - user_conditions.append(meta_clause) - filter_params.extend(meta_params) - - if where_document is not None: - doc_clause, doc_params = FilterBuilder.build_document_filter(where_document, "document") - if doc_clause: - user_conditions.append(doc_clause) - filter_params.extend(doc_params) - - user_where = f"WHERE {' AND '.join(user_conditions)}" if user_conditions else "" - where_clause, filter_params = self._append_namespace_filter(user_where, filter_params, ns_id, ltable_id) - - distance_function_map = { - "l2": "l2_distance", - "cosine": "cosine_distance", - "inner_product": "inner_product", + # Logical-table vector search must use hybrid_search DSL; direct SQL vector + # index scans on namespace partitions are not supported on OceanBase. + all_ids: list[list[Any]] = [] + all_documents: list[list[Any]] = [] + all_metadatas: list[list[Any]] = [] + all_embeddings: list[list[Any]] = [] + all_distances: list[list[float]] = [] + + hybrid_kwargs = { + k: v for k, v in kwargs.items() + if k not in ("embedding_function", "distance", "dimension") } - distance_func = distance_function_map.get(distance, "l2_distance") - conn = self._ensure_connection() - use_context_manager = self._use_context_manager_for_cursor() + for query_vector in query_embeddings: + knn_cfg: dict[str, Any] = { + "query_embeddings": query_vector, + "n_results": n_results, + } + if where is not None: + knn_cfg["where"] = where - all_ids = [] - all_documents = [] - all_metadatas = [] - all_embeddings = [] - all_distances = [] + query_cfg: dict[str, Any] | None = None + if where_document is not None: + query_cfg = {"where_document": where_document, "n_results": n_results} + if where is not None: + query_cfg["where"] = where - for query_vector in query_embeddings: - vector_str = _embedding_to_hexstring(query_vector) - select_clause = ", ".join(select_parts) - sql = ( - f"SELECT {select_clause}, " - f"{distance_func}(embedding, {vector_str}) AS distance " - f"FROM `{table_name}` " - f"{where_clause} " - f"ORDER BY {distance_func}(embedding, {vector_str}) " - f"APPROXIMATE LIMIT %s" + batch = self._namespace_hybrid_search( + collection_id=collection_id, + collection_name=collection_name, + namespace_id=namespace_id, + namespace_name=namespace_name, + query=query_cfg, + knn=knn_cfg, + n_results=n_results, + include=include, + embedding_function=embedding_function, + distance=distance, + dimension=kwargs.get("dimension"), + **hybrid_kwargs, ) - query_params = [*filter_params, n_results] - rows = self._execute_query_with_cursor(conn, sql, query_params, use_context_manager) - q_ids = [] - q_documents = [] - q_metadatas = [] - q_embeddings = [] - q_distances = [] + batch_ids = batch.get("ids") or [[]] + all_ids.append(batch_ids[0] if batch_ids else []) + all_distances.append((batch.get("distances") or [[]])[0]) - for row in rows: - if isinstance(row, dict): - rid_raw = row.get("record_id") - rid = json.loads(rid_raw) if isinstance(rid_raw, str) else rid_raw - q_ids.append(rid) - if "documents" in include_fields or "document" in include_fields or include is None: - q_documents.append(row.get("document")) - if "metadatas" in include_fields or "metadata" in include_fields or include is None: - meta_raw = row.get("metadata") - if isinstance(meta_raw, str): - meta_raw = json.loads(meta_raw) - q_metadatas.append(meta_raw or {}) - if "embeddings" in include_fields or "embedding" in include_fields: - emb = row.get("embedding") - if isinstance(emb, bytes): - emb = self._parse_embedding_from_bytes(emb) - elif isinstance(emb, str): - emb = json.loads(emb) - q_embeddings.append(emb) - q_distances.append(row.get("distance")) - elif isinstance(row, (list, tuple)): - idx = 0 - rid_raw = row[idx]; idx += 1 - rid = json.loads(rid_raw) if isinstance(rid_raw, str) else rid_raw - q_ids.append(rid) - if "documents" in include_fields or "document" in include_fields or include is None: - q_documents.append(row[idx]); idx += 1 - if "metadatas" in include_fields or "metadata" in include_fields or include is None: - meta_raw = row[idx]; idx += 1 - if isinstance(meta_raw, str): - meta_raw = json.loads(meta_raw) - q_metadatas.append(meta_raw or {}) - if "embeddings" in include_fields or "embedding" in include_fields: - emb = row[idx]; idx += 1 - if isinstance(emb, bytes): - emb = self._parse_embedding_from_bytes(emb) - elif isinstance(emb, str): - emb = json.loads(emb) - q_embeddings.append(emb) - q_distances.append(row[idx]) - - all_ids.append(q_ids) if "documents" in include_fields or "document" in include_fields or include is None: - all_documents.append(q_documents) + batch_docs = batch.get("documents") or [[]] + all_documents.append(batch_docs[0] if batch_docs else []) if "metadatas" in include_fields or "metadata" in include_fields or include is None: - all_metadatas.append(q_metadatas) + batch_meta = batch.get("metadatas") or [[]] + all_metadatas.append(batch_meta[0] if batch_meta else []) if "embeddings" in include_fields or "embedding" in include_fields: - all_embeddings.append(q_embeddings) - all_distances.append(q_distances) + batch_emb = batch.get("embeddings") or [[]] + all_embeddings.append(batch_emb[0] if batch_emb else []) - result = {"ids": all_ids, "distances": all_distances} + result: dict[str, Any] = {"ids": all_ids, "distances": all_distances} if "documents" in include_fields or "document" in include_fields or include is None: result["documents"] = all_documents if "metadatas" in include_fields or "metadata" in include_fields or include is None: @@ -5284,7 +5234,7 @@ def _namespace_get( # noqa: C901 id_conds = [] for rid in ids: id_escaped = escape_string(rid) - id_conds.append(f"JSON_EXTRACT(data_content, '$.id') = '{id_escaped}'") + id_conds.append(f"{_NS_DATA_CONTENT_ID_EXPR} = '{id_escaped}'") user_conditions.append(f"({' OR '.join(id_conds)})") if where is not None: @@ -5541,6 +5491,16 @@ def _namespace_hybrid_search( search_parm.pop("_source", None) + if "knn" in search_parm: + if query_hint is None: + query_hint = QueryHint(vector_index=True) + elif query_hint.vector_index is None: + query_hint = QueryHint( + parallel=query_hint.parallel, + query_timeout=query_hint.query_timeout, + vector_index=True, + ) + search_parm_json = json.dumps(search_parm, ensure_ascii=False) use_context_manager = self._use_context_manager_for_cursor() diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index df50c1a8..ede8cd02 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -565,7 +565,7 @@ def test_delete_by_ids_sql(self): assert sql.startswith("DELETE FROM") assert f"`{self.TABLE}`" in sql assert "namespace_id = 7" in sql - assert "JSON_EXTRACT(data_content, '$.id') = " in sql + assert "JSON_UNQUOTE(JSON_EXTRACT(data_content, '$.id')) = " in sql assert "'d1'" in sql def test_delete_by_where_sql(self): @@ -582,7 +582,7 @@ def test_delete_by_where_sql(self): # ---- QUERY ---- - def test_query_basic_sql(self): + def test_query_basic_dsl(self): c = self._client() c.query_return_value = [] c._namespace_query( @@ -592,16 +592,14 @@ def test_query_basic_sql(self): distance="l2", ) sql = c.query_sqls[-1] - assert "SELECT" in sql - assert "JSON_EXTRACT(data_content, '$.id') AS record_id" in sql - assert "l2_distance(embedding," in sql - assert "AS distance" in sql + assert "hybrid_search(TABLE" in sql assert f"`{self.TABLE}`" in sql - assert "namespace_id = 7" in sql - assert "ORDER BY l2_distance(embedding," in sql - assert "APPROXIMATE LIMIT" in sql + assert '"knn"' in sql + assert '"query_vector"' in sql + assert "l2_distance(embedding," not in sql + assert "APPROXIMATE LIMIT" not in sql - def test_query_with_where_sql(self): + def test_query_with_where_dsl(self): c = self._client() c.query_return_value = [] c._namespace_query( @@ -612,8 +610,9 @@ def test_query_with_where_sql(self): distance="cosine", ) sql = c.query_sqls[-1] - assert "cosine_distance(embedding," in sql - assert "metadata.category" in sql + assert "hybrid_search(TABLE" in sql + assert "data_content.metadata.category" in sql + assert "cosine_distance(embedding," not in sql # ---- GET ---- @@ -628,7 +627,7 @@ def test_get_by_ids_sql(self): assert "SELECT" in sql assert f"`{self.TABLE}`" in sql assert "namespace_id = 7" in sql - assert """JSON_EXTRACT(data_content, '$.id') = 'g1'""" in sql + assert """JSON_UNQUOTE(JSON_EXTRACT(data_content, '$.id')) = 'g1'""" in sql def test_get_with_limit_sql(self): c = self._client() @@ -698,6 +697,47 @@ def test_ivf_lib_ob_in_sql(self): sql = _get_ivf_vector_index_sql(config) assert "LIB=OB" in sql + def test_create_namespace_physical_tables_sn_post_create_vector_index(self): + """SN logic_data_table: FTS + SEARCH in CREATE TABLE, vector index post-create.""" + c = self._client() + ivf_config = IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh") + c._create_namespace_physical_tables( + collection_id=self.COLLECTION_ID, + dimension=3, + ivf_config=ivf_config, + is_shared_storage=False, + ) + data_create = next( + s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s + ) + assert "VECTOR INDEX" not in data_create + assert "FULLTEXT INDEX idx_fts(document) WITH PARSER ik" in data_create + assert "SEARCH INDEX idx_json(data_content)" in data_create + vec_create = next(s for s in c.executed_sqls if s.startswith("CREATE VECTOR INDEX")) + assert f"`{self.TABLE}`" in vec_create + assert "fresh_mode=spfresh" in vec_create + + def test_create_namespace_physical_tables_ss_inline_vector_index(self): + """SS logic_data_table: inline VECTOR INDEX (post-create fails on logic tables).""" + c = self._client() + ivf_config = IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh") + c._create_namespace_physical_tables( + collection_id=self.COLLECTION_ID, + dimension=3, + ivf_config=ivf_config, + is_shared_storage=True, + ) + data_create = next( + s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s + ) + assert "VECTOR INDEX idx_vec(embedding)" in data_create + assert "fresh_mode=spfresh" in data_create + assert not any(s.startswith("CREATE VECTOR INDEX") for s in c.executed_sqls) + hot_create = next( + s for s in c.executed_sqls if "CREATE TABLE" in s and "_hot_table" in s + ) + assert hot_create + # ==================== Namespace Name Validation Tests ==================== From 9df48f6041fa16e7632399e422bd9b8b60de5432 Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Mon, 8 Jun 2026 20:38:12 +0800 Subject: [PATCH 34/84] fix case --- ...rid_search_fulltext_multi_coll_multi_ns.py | 1 - ...st_namespace_hybrid_search_search_index.py | 21 +++++++++---------- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py index cc2525c5..32fa76e0 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py @@ -100,7 +100,6 @@ def _run_fts_case_all_quadrants(self, db_client, case_name: str) -> None: run_hybrid_search_fts_case(namespace, corpus, case) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) - # pass def test_cross_quadrant_fts_isolation(self, db_client): """Loaded namespaces must see only their own rows; empty namespaces return no hits.""" diff --git a/tests/integration_tests/test_namespace_hybrid_search_search_index.py b/tests/integration_tests/test_namespace_hybrid_search_search_index.py index d899ec4e..448e4168 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_search_index.py +++ b/tests/integration_tests/test_namespace_hybrid_search_search_index.py @@ -7,7 +7,6 @@ from __future__ import annotations -import time from typing import Any, ClassVar import pytest @@ -32,14 +31,21 @@ def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" if mode not in self._shared_by_mode: corpus, collection = setup_large_fts_collection(db_client) + namespace = setup_fts_namespace_with_corpus( + collection, + corpus, + namespace_name=f"ns_si_shared_{mode}", + ) self._shared_by_mode[mode] = { "db_client": db_client, "corpus": corpus, "collection": collection, + "namespace": namespace, } entry = self._shared_by_mode[mode] self._corpus = entry["corpus"] self._collection = entry["collection"] + self._namespace = entry["namespace"] @classmethod def teardown_class(cls) -> None: @@ -47,17 +53,10 @@ def teardown_class(cls) -> None: teardown_large_fts_collection(entry["db_client"], entry["collection"]) cls._shared_by_mode.clear() - def _new_namespace(self, case_name: str) -> Any: - ns_name = f"ns_si_{case_name}_{int(time.time() * 1000)}" - return setup_fts_namespace_with_corpus( - self._collection, - self._corpus, - namespace_name=ns_name, - ) - def _run_case(self, case_name: str) -> None: - namespace = self._new_namespace(case_name) - run_hybrid_search_index_case(namespace, self._corpus, get_search_index_case(case_name)) + run_hybrid_search_index_case( + self._namespace, self._corpus, get_search_index_case(case_name) + ) @pytest.mark.parametrize("case_name", [c.name for c in SEARCH_INDEX_CASES]) def test_hybrid_search_search_index_operators(self, db_client, case_name: str): From 1ee5a99df7f99cd83cb958f0b099079abd68d63f Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Tue, 9 Jun 2026 11:37:18 +0800 Subject: [PATCH 35/84] del print --- src/pyseekdb/client/client_base.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 20a76e9e..e0e5a354 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -124,7 +124,7 @@ def _validate_collection_name(name: str) -> None: from .validators import _MAX_NAMESPACE_BATCH_SIZE, _validate_namespace_name, _validate_record_ids # noqa: F401 -_NS_PARTITION_COUNT = 8 +_NS_PARTITION_COUNT = 4 def set_namespace_partition_count(n: int) -> None: @@ -5558,7 +5558,6 @@ def _namespace_hybrid_search( f"SELECT {hint_sql + ' ' if hint_sql else ''}* " f"FROM hybrid_search(TABLE `{table_name}`, '{escaped_params}')" ) - print(f"hybrid_sql: {hybrid_sql}") result_rows = self._execute_query_with_cursor(conn, hybrid_sql, [], use_context_manager) if not result_rows: return { From 551e3516a0a6c64043f8fbb3f3d6b9021d61c90e Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 9 Jun 2026 16:22:55 +0800 Subject: [PATCH 36/84] =?UTF-8?q?ss=E6=A8=A1=E5=BC=8F=E5=92=8Csn=E6=A8=A1?= =?UTF-8?q?=E5=BC=8F=E9=80=BB=E8=BE=91=E8=A1=A8=E5=90=91=E9=87=8F=E7=B4=A2?= =?UTF-8?q?=E5=BC=95=E9=9A=8F=E8=A1=A8=E5=88=9B=E5=BB=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 25 ++++++------------------- tests/unit_tests/test_namespace.py | 14 +++++++------- 2 files changed, 13 insertions(+), 26 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index f9830f35..2bbbe4e6 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1301,19 +1301,12 @@ def _create_namespace_physical_tables( try: self._execute(f"CREATE TABLEGROUP `{tg_name}` SHARDING='ADAPTIVE'") - # SN: post-create ``CREATE VECTOR INDEX`` on logic tables (documented path). - # SS: logic tables reject post-create vector indexes (OB-5703); inline only. - if is_shared_storage: - index_sql = ( - f"FULLTEXT INDEX idx_fts(document) {fulltext_clause},\n" - f" SEARCH INDEX idx_json(data_content),\n" - f" VECTOR INDEX idx_vec(embedding) {vector_index_sql}" - ) - else: - index_sql = ( - f"FULLTEXT INDEX idx_fts(document) {fulltext_clause},\n" - f" SEARCH INDEX idx_json(data_content)" - ) + # Inline VECTOR INDEX in CREATE TABLE for both SN and SS logic tables. + index_sql = ( + f"FULLTEXT INDEX idx_fts(document) {fulltext_clause},\n" + f" SEARCH INDEX idx_json(data_content),\n" + f" VECTOR INDEX idx_vec(embedding) {vector_index_sql}" + ) data_sql = f"""CREATE TABLE `{data_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, @@ -1329,12 +1322,6 @@ def _create_namespace_physical_tables( self._execute(data_sql) - if not is_shared_storage: - ivf_index_sql = ( - f"CREATE VECTOR INDEX idx_vec ON `{data_table}` (embedding) {vector_index_sql}" - ) - self._execute(ivf_index_sql) - if is_shared_storage: hot_table = NamespaceCollectionNames.hot_table_name(collection_id) self._execute(f"""CREATE TABLE `{hot_table}` ( diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 1f79ed70..0ec5e306 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -709,8 +709,8 @@ def test_ivf_lib_ob_in_sql(self): sql = _get_ivf_vector_index_sql(config) assert "LIB=OB" in sql - def test_create_namespace_physical_tables_sn_post_create_vector_index(self): - """SN logic_data_table: FTS + SEARCH in CREATE TABLE, vector index post-create.""" + def test_create_namespace_physical_tables_sn_inline_vector_index(self): + """SN logic_data_table: inline VECTOR INDEX in CREATE TABLE (same as SS).""" c = self._client() ivf_config = IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh") c._create_namespace_physical_tables( @@ -722,15 +722,15 @@ def test_create_namespace_physical_tables_sn_post_create_vector_index(self): data_create = next( s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s ) - assert "VECTOR INDEX" not in data_create + assert "VECTOR INDEX idx_vec(embedding)" in data_create assert "FULLTEXT INDEX idx_fts(document) WITH PARSER ik" in data_create assert "SEARCH INDEX idx_json(data_content)" in data_create - vec_create = next(s for s in c.executed_sqls if s.startswith("CREATE VECTOR INDEX")) - assert f"`{self.TABLE}`" in vec_create - assert "fresh_mode=spfresh" in vec_create + assert "fresh_mode=spfresh" in data_create + assert not any(s.startswith("CREATE VECTOR INDEX") for s in c.executed_sqls) + assert not any("_hot_table" in s for s in c.executed_sqls if s.startswith("CREATE TABLE")) def test_create_namespace_physical_tables_ss_inline_vector_index(self): - """SS logic_data_table: inline VECTOR INDEX (post-create fails on logic tables).""" + """SS logic_data_table: inline VECTOR INDEX in CREATE TABLE.""" c = self._client() ivf_config = IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh") c._create_namespace_physical_tables( From 0a722a7c9d7396a94c0111e8db70c20342e8d90b Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 10 Jun 2026 12:00:00 +0800 Subject: [PATCH 37/84] =?UTF-8?q?=E7=94=A8=E6=88=B7=E8=87=AA=E5=AE=9A?= =?UTF-8?q?=E4=B9=89=E5=88=9B=E5=BB=BA=E5=90=91=E9=87=8F,=E5=85=A8?= =?UTF-8?q?=E6=96=87=E7=B4=A2=E5=BC=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 29 ++- .../namespace_dml_helpers.py | 3 +- .../namespace_fts_helpers.py | 3 +- .../test_namespace_optional_indexes.py | 189 ++++++++++++++++++ .../integration_tests/test_namespace_query.py | 3 +- .../test_namespace_reconnect.py | 146 ++++++++++++++ tests/unit_tests/test_namespace.py | 105 ++++++++-- 7 files changed, 444 insertions(+), 34 deletions(-) create mode 100644 tests/integration_tests/test_namespace_optional_indexes.py create mode 100644 tests/integration_tests/test_namespace_reconnect.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 2bbbe4e6..e4229886 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -944,16 +944,12 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " if ivf_config is not None: dimension = ivf_config.dimension distance = ivf_config.distance - if ivf_config.fresh_mode is not None and ivf_config.fresh_mode != "spfresh": - raise ValueError("use_namespace=True requires fresh_mode='spfresh'") else: if dense_embedding_function is not None: dimension = self._get_embedding_function_dimension(dense_embedding_function) else: dimension = DEFAULT_VECTOR_DIMENSION distance = DEFAULT_DISTANCE_METRIC - from .configuration import IVFConfiguration - ivf_config = IVFConfiguration(dimension=dimension, distance=distance, fresh_mode="spfresh") if dimension < 1 or dimension > 4096: raise ValueError(f"Dimension must be between 1 and 4096, got {dimension}") @@ -962,12 +958,14 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " settings = { "version": 2, "use_namespace": True, - "dense_index_type": "ivf", - "fresh_mode": "spfresh", "storage_mode": "ss" if is_ss else "sn", "dimension": dimension, "distance": distance, } + if ivf_config is not None: + settings["dense_index_type"] = "ivf" + if ivf_config.fresh_mode is not None: + settings["fresh_mode"] = ivf_config.fresh_mode if dense_embedding_function is not None and EmbeddingFunction.support_persistence(dense_embedding_function): settings["embedding_function"] = { "name": dense_embedding_function.name(), @@ -1285,7 +1283,7 @@ def _create_namespace_physical_tables( self, collection_id: str, dimension: int, - ivf_config, + ivf_config=None, fulltext_config=None, is_shared_storage: bool = False, ) -> None: @@ -1294,20 +1292,19 @@ def _create_namespace_physical_tables( kv_table = NamespaceCollectionNames.kv_data_table_name(collection_id) schema_table = NamespaceCollectionNames.logic_schema_table_name(collection_id) - fulltext_clause = _get_fulltext_index_sql(fulltext_config) - vector_index_sql = _get_ivf_vector_index_sql(ivf_config) + index_parts = ["SEARCH INDEX idx_json(data_content)"] + if fulltext_config is not None: + fulltext_clause = _get_fulltext_index_sql(fulltext_config) + index_parts.insert(0, f"FULLTEXT INDEX idx_fts(document) {fulltext_clause}") + if ivf_config is not None: + vector_index_sql = _get_ivf_vector_index_sql(ivf_config) + index_parts.append(f"VECTOR INDEX idx_vec(embedding) {vector_index_sql}") + index_sql = ",\n ".join(index_parts) partition_clause = f"PARTITION BY KEY(namespace_id) PARTITIONS {_NS_PARTITION_COUNT}" try: self._execute(f"CREATE TABLEGROUP `{tg_name}` SHARDING='ADAPTIVE'") - # Inline VECTOR INDEX in CREATE TABLE for both SN and SS logic tables. - index_sql = ( - f"FULLTEXT INDEX idx_fts(document) {fulltext_clause},\n" - f" SEARCH INDEX idx_json(data_content),\n" - f" VECTOR INDEX idx_vec(embedding) {vector_index_sql}" - ) - data_sql = f"""CREATE TABLE `{data_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, ltable_id BIGINT UNSIGNED NOT NULL, diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index dc1051e6..87adb309 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -9,7 +9,7 @@ from typing import Any from pyseekdb import IVFConfiguration -from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema LARGE_DML_CORPUS_SIZE = 1000 @@ -40,6 +40,7 @@ def ns_schema() -> Schema: ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), embedding_function=None, ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), ) diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index c8acd7e4..df47d110 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -18,7 +18,7 @@ VECTOR_DISTANCE_METRICS: tuple[VectorDistanceMetric, ...] = ("l2", "cosine") from pyseekdb import IVFConfiguration -from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema # Rare ASCII tokens to reduce IK segmentation surprises in assertions. @@ -58,6 +58,7 @@ def ns_schema(distance: VectorDistanceMetric = "l2") -> Schema: ivf=IVFConfiguration(dimension=3, distance=distance, fresh_mode="spfresh"), embedding_function=None, ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), ) diff --git a/tests/integration_tests/test_namespace_optional_indexes.py b/tests/integration_tests/test_namespace_optional_indexes.py new file mode 100644 index 00000000..24e9a712 --- /dev/null +++ b/tests/integration_tests/test_namespace_optional_indexes.py @@ -0,0 +1,189 @@ +""" +Integration tests for optional namespace collection indexes. + +Product behavior (Agent Database): +- ``schema.vector_index.ivf`` unset → no VECTOR INDEX +- ``schema.fulltext_index`` unset → no FULLTEXT INDEX +- SEARCH INDEX on ``data_content`` is always created + +These tests verify DDL, settings metadata, and query behavior for: +1. fulltext only +2. vector (IVF) only +3. search index only (no fulltext, no vector) +""" + +from __future__ import annotations + +import json +import time +from typing import Any + +import pymysql.err +import pytest + +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig +from pyseekdb.client.meta_info import NamespaceCollectionNames +from pyseekdb.client.schema import Schema + + +def _schema_fts_only() -> Schema: + return Schema( + vector_index=VectorIndexConfig(embedding_function=None), + fulltext_index=FulltextIndexConfig(analyzer="ik"), + ) + + +def _schema_vector_only() -> Schema: + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), + embedding_function=None, + ), + ) + + +def _schema_search_only() -> Schema: + return Schema(vector_index=VectorIndexConfig(embedding_function=None)) + + +def _index_names(client: Any, collection_id: str) -> set[str]: + table = NamespaceCollectionNames.data_table_name(collection_id) + rows = client._server._execute(f"SHOW INDEX FROM `{table}`") + return {(r.get("Key_name") if isinstance(r, dict) else r[2]) for r in (rows or [])} + + +def _collection_settings(client: Any, collection_id: str) -> dict: + rows = client._server._execute( + f"SELECT settings FROM sdk_collections WHERE collection_id = '{collection_id}'" + ) + raw = rows[0]["settings"] if isinstance(rows[0], dict) else rows[0][0] + return json.loads(raw) if raw else {} + + +def _create_collection(client: Any, label: str, schema: Schema) -> Any: + name = f"test_ns_opt_idx_{label}_{int(time.time() * 1000)}" + return client.create_collection(name=name, schema=schema, use_namespace=True) + + +def _unit_vector(dimension: int, axis: int = 0) -> list[float]: + vec = [0.0] * dimension + vec[axis % dimension] = 1.0 + return vec + + +def _seed_namespace(ns: Any, dimension: int) -> None: + ns.add( + ids=["d1", "d2"], + embeddings=[_unit_vector(dimension, 0), _unit_vector(dimension, 1)], + documents=["machine learning guide", "python database tutorial"], + metadatas=[{"category": "AI", "score": 90}, {"category": "Programming", "score": 80}], + ) + + +class TestNamespaceOptionalIndexes: + """DDL + query matrix for optional VECTOR / FULLTEXT indexes.""" + + def test_fts_only_indexes_and_queries(self, db_client): + collection = _create_collection(db_client, "fts_only", _schema_fts_only()) + try: + keys = _index_names(db_client, collection.id) + settings = _collection_settings(db_client, collection.id) + assert keys == {"idx_fts", "idx_json"} + assert "dense_index_type" not in settings + + ns = collection.create_namespace("ns") + _seed_namespace(ns, collection.dimension) + time.sleep(1) + + result = ns.hybrid_search( + query={"where_document": {"$contains": "machine"}}, + n_results=5, + include=["documents"], + ) + assert result["ids"] and result["ids"][0] + + with pytest.raises(pymysql.err.Error, match="knn search without vector index not supported"): + ns.query(query_embeddings=_unit_vector(collection.dimension), n_results=3) + + with pytest.raises(pymysql.err.Error, match="knn search without vector index not supported"): + ns.hybrid_search( + knn={"query_embeddings": _unit_vector(collection.dimension), "n_results": 3}, + n_results=3, + ) + finally: + db_client.delete_collection(name=collection.name) + + def test_vector_only_indexes_and_queries(self, db_client): + collection = _create_collection(db_client, "vec_only", _schema_vector_only()) + try: + keys = _index_names(db_client, collection.id) + settings = _collection_settings(db_client, collection.id) + assert keys == {"idx_json", "idx_vec"} + assert settings.get("dense_index_type") == "ivf" + assert settings.get("fresh_mode") == "spfresh" + + ns = collection.create_namespace("ns") + _seed_namespace(ns, collection.dimension) + time.sleep(1) + + result = ns.query(query_embeddings=_unit_vector(3), n_results=3) + assert result["ids"] and result["ids"][0] + + knn = ns.hybrid_search( + knn={"query_embeddings": _unit_vector(3), "n_results": 3}, + n_results=3, + ) + assert knn["ids"] and knn["ids"][0] + + with pytest.raises(pymysql.err.OperationalError, match="FULLTEXT"): + ns.hybrid_search( + query={"where_document": {"$contains": "machine"}}, + n_results=5, + ) + finally: + db_client.delete_collection(name=collection.name) + + def test_search_only_indexes_and_queries(self, db_client): + collection = _create_collection(db_client, "search_only", _schema_search_only()) + try: + keys = _index_names(db_client, collection.id) + settings = _collection_settings(db_client, collection.id) + assert keys == {"idx_json"} + assert "dense_index_type" not in settings + + ns = collection.create_namespace("ns") + _seed_namespace(ns, collection.dimension) + time.sleep(1) + + scalar = ns.hybrid_search( + query={"where": {"category": "AI"}}, + n_results=10, + include=["metadatas"], + ) + ids = scalar["ids"][0] if scalar["ids"] else [] + assert ids == ["d1"] + + got = ns.get(where={"category": "Programming"}) + assert sorted(got["ids"]) == ["d2"] + + with pytest.raises(pymysql.err.OperationalError, match="FULLTEXT"): + ns.hybrid_search( + query={"where_document": {"$contains": "machine"}}, + n_results=5, + ) + + with pytest.raises(pymysql.err.Error, match="knn search without vector index not supported"): + ns.query(query_embeddings=_unit_vector(collection.dimension), n_results=3) + + with pytest.raises(pymysql.err.Error, match="knn search without vector index not supported"): + ns.hybrid_search( + knn={"query_embeddings": _unit_vector(collection.dimension), "n_results": 3}, + n_results=3, + ) + finally: + db_client.delete_collection(name=collection.name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s", "-k", "oceanbase"]) diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index 6916c686..51269f49 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -8,7 +8,7 @@ import pytest from pyseekdb import IVFConfiguration -from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema @@ -21,6 +21,7 @@ def _setup(self, client, suffix=""): ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), embedding_function=None, ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), ) collection = client.create_collection(name=name, schema=schema, use_namespace=True) return collection diff --git a/tests/integration_tests/test_namespace_reconnect.py b/tests/integration_tests/test_namespace_reconnect.py new file mode 100644 index 00000000..3d4cd23a --- /dev/null +++ b/tests/integration_tests/test_namespace_reconnect.py @@ -0,0 +1,146 @@ +""" +Namespace integration tests for connection lifecycle and multi-client access. + +Verifies that namespace data persists across: +1. Disconnect + lazy reconnect on the same Client instance +2. A separate Client connecting to the same collection/namespace +3. Cross-client read/write visibility on shared namespace data +""" + +from __future__ import annotations + +import contextlib +import os +import time +from typing import Any + +import pyseekdb + +from namespace_dml_helpers import cleanup, create_ns_collection + + +def _new_oceanbase_client() -> Any: + """Second pymysql-backed client (separate TCP connection).""" + client = pyseekdb.Client( + host=os.environ.get("OB_HOST", "localhost"), + port=int(os.environ.get("OB_PORT", "11202")), + tenant=os.environ.get("OB_TENANT", "mysql"), + database=os.environ.get("OB_DATABASE", "test"), + user=os.environ.get("OB_USER", "root"), + password=os.environ.get("OB_PASSWORD", ""), + ) + result = client._server._execute("SELECT 1 AS test") + assert result and result[0].get("test") == 1 + return client + + +class TestNamespaceReconnectAndMultiClient: + + def test_same_client_reconnect_reads_namespace_data(self, oceanbase_client): + """After _cleanup(), the next SDK call should reconnect and read prior rows.""" + client = oceanbase_client + collection = create_ns_collection(client, suffix="_reconn") + coll_name = collection.name + try: + ns = collection.create_namespace("persist_ns") + ns.add( + ids="doc1", + embeddings=[[1.0, 0.0, 0.0]], + documents=["persist after reconnect"], + metadatas={"tag": "before_disconnect"}, + ) + + assert client._server.is_connected() + client._server._cleanup() + assert not client._server.is_connected() + + collection2 = client.get_collection(coll_name) + assert client._server.is_connected() + + ns2 = collection2.get_namespace("persist_ns") + res = ns2.get(ids="doc1", include=["documents", "metadatas"]) + assert res["ids"] == ["doc1"] + assert res["documents"][0] == "persist after reconnect" + assert res["metadatas"][0]["tag"] == "before_disconnect" + + query = ns2.query(query_embeddings=[1.0, 0.0, 0.0], n_results=3, include=["metadatas"]) + assert "doc1" in query["ids"][0] + finally: + cleanup(client, collection) + + def test_new_client_reads_existing_namespace(self, oceanbase_client): + """A fresh Client should see data written by another connection.""" + client_a = oceanbase_client + client_b = _new_oceanbase_client() + collection = create_ns_collection(client_a, suffix="_newcli") + coll_name = collection.name + try: + ns_a = collection.create_namespace("shared_read") + ns_a.add( + ids="shared1", + embeddings=[[0.0, 1.0, 0.0]], + documents=["written by client A"], + metadatas={"writer": "A"}, + ) + + coll_b = client_b.get_collection(coll_name) + assert coll_b.use_namespace is True + assert coll_b.id == collection.id + + ns_b = coll_b.get_namespace("shared_read") + res = ns_b.get(ids="shared1", include=["documents", "metadatas"]) + assert res["ids"] == ["shared1"] + assert res["documents"][0] == "written by client A" + assert res["metadatas"][0]["writer"] == "A" + finally: + cleanup(client_a, collection) + with contextlib.suppress(Exception): + client_b._server._cleanup() + + def test_two_clients_read_write_same_namespace(self, oceanbase_client): + """Writes from client B should be visible to client A on the same namespace.""" + client_a = oceanbase_client + client_b = _new_oceanbase_client() + collection = create_ns_collection(client_a, suffix="_multi") + coll_name = collection.name + try: + ns_a = collection.create_namespace("shared_rw") + ns_a.add( + ids="from_a", + embeddings=[[1.0, 0.0, 0.0]], + documents=["seed from A"], + metadatas={"src": "A"}, + ) + + coll_b = client_b.get_collection(coll_name) + ns_b = coll_b.get_namespace("shared_rw") + + got_b = ns_b.get(ids="from_a", include=["metadatas"]) + assert got_b["ids"] == ["from_a"] + assert got_b["metadatas"][0]["src"] == "A" + + ns_b.add( + ids="from_b", + embeddings=[[0.0, 1.0, 0.0]], + documents=["written by B"], + metadatas={"src": "B"}, + ) + + # Small pause so count/get on A observes B's insert on a separate session. + time.sleep(0.2) + got_a = ns_a.get(where={"src": "B"}, include=["documents"]) + assert got_a["ids"] == ["from_b"] + assert got_a["documents"][0] == "written by B" + + assert ns_a.count() == 2 + assert ns_b.count() == 2 + finally: + cleanup(client_a, collection) + with contextlib.suppress(Exception): + client_b._server._cleanup() + + +if __name__ == "__main__": + import pytest + + pytest.main([__file__, "-v", "-s", "-k", "oceanbase"]) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 0ec5e306..d36fe02e 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -355,6 +355,7 @@ class FakeClient(BaseClient): """Concrete BaseClient subclass that captures SQL without executing.""" def __init__(self): + self.database = "test" self.executed_sqls = [] self.query_sqls = [] self.query_return_value = [] @@ -589,8 +590,8 @@ def test_delete_by_where_document_sql(self): sql = c.executed_sqls[-1] assert "DELETE FROM" in sql assert "namespace_id = 7" in sql - assert "document LIKE" in sql - assert "MATCH(document)" not in sql + assert "MATCH(document) AGAINST" in sql + assert "obsolete" in sql # ---- QUERY ---- @@ -711,12 +712,15 @@ def test_ivf_lib_ob_in_sql(self): def test_create_namespace_physical_tables_sn_inline_vector_index(self): """SN logic_data_table: inline VECTOR INDEX in CREATE TABLE (same as SS).""" + from pyseekdb.client.configuration import FulltextIndexConfig + c = self._client() ivf_config = IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh") c._create_namespace_physical_tables( collection_id=self.COLLECTION_ID, dimension=3, ivf_config=ivf_config, + fulltext_config=FulltextIndexConfig(analyzer="ik"), is_shared_storage=False, ) data_create = next( @@ -743,6 +747,8 @@ def test_create_namespace_physical_tables_ss_inline_vector_index(self): s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s ) assert "VECTOR INDEX idx_vec(embedding)" in data_create + assert "FULLTEXT INDEX" not in data_create + assert "SEARCH INDEX idx_json(data_content)" in data_create assert "fresh_mode=spfresh" in data_create assert not any(s.startswith("CREATE VECTOR INDEX") for s in c.executed_sqls) hot_create = next( @@ -750,6 +756,37 @@ def test_create_namespace_physical_tables_ss_inline_vector_index(self): ) assert hot_create + def test_create_namespace_physical_tables_search_index_only(self): + """Without ivf/fulltext config, only SEARCH INDEX is created.""" + c = self._client() + c._create_namespace_physical_tables( + collection_id=self.COLLECTION_ID, + dimension=3, + is_shared_storage=False, + ) + data_create = next( + s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s + ) + assert "SEARCH INDEX idx_json(data_content)" in data_create + assert "VECTOR INDEX" not in data_create + assert "FULLTEXT INDEX" not in data_create + + def test_create_namespace_physical_tables_ivf_without_fresh_mode(self): + """IVF without fresh_mode omits fresh_mode from VECTOR INDEX DDL.""" + c = self._client() + ivf_config = IVFConfiguration(dimension=3, distance="l2") + c._create_namespace_physical_tables( + collection_id=self.COLLECTION_ID, + dimension=3, + ivf_config=ivf_config, + is_shared_storage=False, + ) + data_create = next( + s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s + ) + assert "VECTOR INDEX idx_vec(embedding)" in data_create + assert "fresh_mode" not in data_create + # ==================== Namespace Name Validation Tests ==================== @@ -981,15 +1018,15 @@ def test_ensure_namespace_catalogs_creates_all_catalog_tables(self): c._ensure_namespace_catalogs() sql = "\n".join(c.executed_sqls) - assert "CREATE TABLE IF NOT EXISTS `sdk_namespaces`" in sql + assert "CREATE TABLE IF NOT EXISTS `test`.`sdk_namespaces`" in sql assert "PRIMARY KEY (namespace_id)" in sql assert "UNIQUE KEY uk_sdk_ns_coll_name (collection_id, namespace_name)" in sql - assert "CREATE TABLE IF NOT EXISTS `sdk_ltables`" in sql + assert "CREATE TABLE IF NOT EXISTS `test`.`sdk_ltables`" in sql assert "PRIMARY KEY (ltable_id)" in sql assert "UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name)" in sql - assert "CREATE TABLE IF NOT EXISTS `sdk_namespaces_stats`" in sql + assert "CREATE TABLE IF NOT EXISTS `test`.`sdk_namespaces_stats`" in sql assert "collection_id CHAR(32) NOT NULL" in sql assert "namespace_id BIGINT UNSIGNED NOT NULL" in sql assert "ltable_id BIGINT UNSIGNED NOT NULL" in sql @@ -1001,10 +1038,11 @@ def test_ensure_namespace_catalogs_creates_catalog_tables_in_order(self): c = FakeClient() c._ensure_namespace_catalogs() - assert len(c.executed_sqls) == 3 - assert "`sdk_namespaces`" in c.executed_sqls[0] - assert "`sdk_ltables`" in c.executed_sqls[1] - assert "`sdk_namespaces_stats`" in c.executed_sqls[2] + assert len(c.executed_sqls) == 5 + assert c.executed_sqls[0] == "USE `test`" + assert "`test`.`sdk_namespaces`" in c.executed_sqls[2] + assert "`test`.`sdk_ltables`" in c.executed_sqls[3] + assert "`test`.`sdk_namespaces_stats`" in c.executed_sqls[4] def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): c = FakeClient() @@ -1042,14 +1080,50 @@ def test_ob_type_validation(self): with pytest.raises(ValueError, match="only supported on OceanBase"): c._create_namespace_collection("test", schema) - def test_fresh_mode_false_raises(self): + def test_create_namespace_collection_without_ivf_skips_vector_index(self): c = FakeClient() c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", "4.3")) + c._is_shared_storage_mode = MagicMock(return_value=False) + c._create_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) + c._ensure_namespace_catalogs = MagicMock() from pyseekdb.client.schema import Schema from pyseekdb.client.configuration import VectorIndexConfig - schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, fresh_mode="none"), embedding_function=None)) - with pytest.raises(ValueError, match="requires fresh_mode='spfresh'"): - c._create_namespace_collection("test", schema) + + schema = Schema(vector_index=VectorIndexConfig(embedding_function=None)) + c._create_namespace_collection("test", schema) + data_create = next( + s for s in c.executed_sqls if "CREATE TABLE" in s and "logic_data_table" in s + ) + assert "VECTOR INDEX" not in data_create + assert "SEARCH INDEX idx_json(data_content)" in data_create + settings = c._create_ns_collection_meta.call_args[0][1] + assert "dense_index_type" not in settings + assert "fresh_mode" not in settings + + def test_create_namespace_collection_ivf_without_fresh_mode(self): + c = FakeClient() + c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", "4.3")) + c._is_shared_storage_mode = MagicMock(return_value=False) + c._create_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) + c._ensure_namespace_catalogs = MagicMock() + from pyseekdb.client.schema import Schema + from pyseekdb.client.configuration import VectorIndexConfig + + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2"), + embedding_function=None, + ) + ) + c._create_namespace_collection("test", schema) + data_create = next( + s for s in c.executed_sqls if "CREATE TABLE" in s and "logic_data_table" in s + ) + assert "VECTOR INDEX idx_vec(embedding)" in data_create + assert "fresh_mode" not in data_create + settings = c._create_ns_collection_meta.call_args[0][1] + assert settings["dense_index_type"] == "ivf" + assert "fresh_mode" not in settings # ==================== Delete Namespace Uses Kernel ==================== @@ -1061,8 +1135,9 @@ def test_delete_namespace_calls_dbms_logic_table(self): c = FakeClient() c._execute = MagicMock(side_effect=[ [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], - None, # _get_ns_namespace_meta -> _set_session_ns_context SET @namespace_id - None, # _get_ns_namespace_meta -> _set_session_ns_context SET @ltable_id + None, # _get_ns_namespace_meta -> SET @namespace_id + None, # _get_ns_namespace_meta -> SET @ltable_id + None, # _delete_ns_namespace_meta -> USE catalog database None, # _delete_ns_namespace_meta -> SET @collection_id None, # _delete_ns_namespace_meta -> SET @namespace_id None, # _delete_ns_namespace_meta -> SET @ltable_id From 94dcf5fea181c255f5644b06c304457585c4889f Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 11 Jun 2026 16:25:45 +0800 Subject: [PATCH 38/84] =?UTF-8?q?set=5Fnamespace=5Fpartition=5Fcount?= =?UTF-8?q?=E6=8E=A5=E5=8F=A3=E6=94=B9=E5=90=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/__init__.py | 8 ++++---- src/pyseekdb/client/__init__.py | 6 +++--- src/pyseekdb/client/client_base.py | 6 +++--- tests/integration_tests/conftest.py | 8 ++++---- .../integration_tests/test_namespace_lifecycle.py | 12 ++++++------ tests/integration_tests/test_namespace_prewarm.py | 8 ++++---- .../test_namespace_prewarm_lob.py | 14 +++++++------- 7 files changed, 31 insertions(+), 31 deletions(-) diff --git a/src/pyseekdb/__init__.py b/src/pyseekdb/__init__.py index da587421..1995ce48 100644 --- a/src/pyseekdb/__init__.py +++ b/src/pyseekdb/__init__.py @@ -102,10 +102,10 @@ VectorIndexConfig, Version, get_default_embedding_function, - get_namespace_partition_count, + get_collection_partition_count, register_embedding_function, register_sparse_embedding_function, - set_namespace_partition_count, + set_collection_partition_count, ) from .client.collection import Collection from .client.namespace import Namespace @@ -151,8 +151,8 @@ "VectorIndexConfig", "Version", "get_default_embedding_function", - "get_namespace_partition_count", + "get_collection_partition_count", "register_embedding_function", "register_sparse_embedding_function", - "set_namespace_partition_count", + "set_collection_partition_count", ] diff --git a/src/pyseekdb/client/__init__.py b/src/pyseekdb/client/__init__.py index 0fc90630..7d92d918 100644 --- a/src/pyseekdb/client/__init__.py +++ b/src/pyseekdb/client/__init__.py @@ -22,7 +22,7 @@ from .admin_client import AdminAPI, _AdminClientProxy, _ClientProxy from .base_connection import BaseConnection -from .client_base import BaseClient, ClientAPI, get_namespace_partition_count, set_namespace_partition_count +from .client_base import BaseClient, ClientAPI, get_collection_partition_count, set_collection_partition_count from .client_seekdb_server import RemoteServerClient from .configuration import ( BengProperties, @@ -188,10 +188,10 @@ def __getattr__(name: str) -> Any: "VectorIndexConfig", "Version", "get_default_embedding_function", - "get_namespace_partition_count", + "get_collection_partition_count", "register_embedding_function", "register_sparse_embedding_function", - "set_namespace_partition_count", + "set_collection_partition_count", ] diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index e4229886..f7e34353 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -130,14 +130,14 @@ def _validate_collection_name(name: str) -> None: _NS_DATA_CONTENT_ID_EXPR = "JSON_UNQUOTE(JSON_EXTRACT(data_content, '$.id'))" -def set_namespace_partition_count(n: int) -> None: +def set_collection_partition_count(n: int) -> None: global _NS_PARTITION_COUNT if n < 1: - raise ValueError("namespace partition count must be >= 1") + raise ValueError("collection partition count must be >= 1") _NS_PARTITION_COUNT = n -def get_namespace_partition_count() -> int: +def get_collection_partition_count() -> int: return _NS_PARTITION_COUNT diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index b8e805db..ef05111c 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -275,8 +275,8 @@ def oceanbase_admin_client(): @pytest.fixture(autouse=True) def _reduce_namespace_partitions(): """Use a small partition count in integration tests to speed up DDL.""" - from pyseekdb import get_namespace_partition_count, set_namespace_partition_count - original = get_namespace_partition_count() - set_namespace_partition_count(8) + from pyseekdb import get_collection_partition_count, set_collection_partition_count + original = get_collection_partition_count() + set_collection_partition_count(8) yield - set_namespace_partition_count(original) + set_collection_partition_count(original) diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 16fb61d9..d86ca86a 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -202,14 +202,14 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): client.delete_collection(name=name) - def test_custom_namespace_partition_count(self, oceanbase_client): - """Verify set_namespace_partition_count controls the PARTITIONS clause.""" - from pyseekdb import get_namespace_partition_count, set_namespace_partition_count + def test_custom_collection_partition_count(self, oceanbase_client): + """Verify set_collection_partition_count controls the PARTITIONS clause.""" + from pyseekdb import get_collection_partition_count, set_collection_partition_count from pyseekdb.client.meta_info import NamespaceCollectionNames - original = get_namespace_partition_count() + original = get_collection_partition_count() try: - set_namespace_partition_count(4) + set_collection_partition_count(4) collection = self._create_ns_collection(oceanbase_client, suffix="_pc") try: data_table = NamespaceCollectionNames.data_table_name(collection.id) @@ -223,7 +223,7 @@ def test_custom_namespace_partition_count(self, oceanbase_client): finally: oceanbase_client.delete_collection(name=collection.name) finally: - set_namespace_partition_count(original) + set_collection_partition_count(original) if __name__ == "__main__": diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index 32f0e052..4617483d 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -20,11 +20,11 @@ class TestNamespacePrewarm: @pytest.fixture(autouse=True) def _reduce_namespace_partitions(self): """Override conftest: use 8 partitions for prewarm until DDL is faster.""" - from pyseekdb import get_namespace_partition_count, set_namespace_partition_count - original = get_namespace_partition_count() - set_namespace_partition_count(8) + from pyseekdb import get_collection_partition_count, set_collection_partition_count + original = get_collection_partition_count() + set_collection_partition_count(8) yield - set_namespace_partition_count(original) + set_collection_partition_count(original) def _create_ns_collection_and_namespace(self, client): name = f"test_ns_pw_{int(time.time() * 1000)}" diff --git a/tests/integration_tests/test_namespace_prewarm_lob.py b/tests/integration_tests/test_namespace_prewarm_lob.py index 72c8d520..48eaf424 100644 --- a/tests/integration_tests/test_namespace_prewarm_lob.py +++ b/tests/integration_tests/test_namespace_prewarm_lob.py @@ -41,8 +41,8 @@ from pyseekdb import ( IVFConfiguration, - get_namespace_partition_count, - set_namespace_partition_count, + get_collection_partition_count, + set_collection_partition_count, ) from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.meta_info import NamespaceCollectionNames @@ -194,7 +194,7 @@ def _grep_lob_prewarm_log(lob_meta_tablet_ids): class _BaseLobPrewarm: def _make_collection(self, client, name, partitions): - set_namespace_partition_count(partitions) + set_collection_partition_count(partitions) schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine"), @@ -234,9 +234,9 @@ class TestLobPrewarmStructural(_BaseLobPrewarm): @pytest.fixture(autouse=True) def _restore_partitions(self): - original = get_namespace_partition_count() + original = get_collection_partition_count() yield - set_namespace_partition_count(original) + set_collection_partition_count(original) def test_kv_table_has_lob_meta_tablet(self, oceanbase_client): """Tier 1: the prewarm target (kv_data_table) owns an aux LOB-meta tablet.""" @@ -295,9 +295,9 @@ class TestLobPrewarmCaching(_BaseLobPrewarm): @pytest.fixture(autouse=True) def _restore_partitions(self): - original = get_namespace_partition_count() + original = get_collection_partition_count() yield - set_namespace_partition_count(original) + set_collection_partition_count(original) @staticmethod def _total_bytes(cache): From 52dcbaddb05d16290eae346420e5355b495353c6 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 12 Jun 2026 14:15:33 +0800 Subject: [PATCH 39/84] =?UTF-8?q?ob=E5=86=85=E6=A0=B8=E7=89=88=E6=9C=AC?= =?UTF-8?q?=E6=8B=A6=E6=88=AA+=E5=8F=AA=E6=94=AF=E6=8C=81ivf=5Fflat?= =?UTF-8?q?=E6=8B=A6=E6=88=AA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 16 +++- .../test_namespace_constraints.py | 95 +++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 tests/integration_tests/test_namespace_constraints.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index f7e34353..6d889e70 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -27,6 +27,7 @@ FulltextIndexConfig, HNSWConfiguration, IVFConfiguration, + IVFIndexType, VectorIndexConfig, ) from .database import Database @@ -60,6 +61,9 @@ # Maximum allowed length for user-facing collection names. _MAX_COLLECTION_NAME_LENGTH = 512 +# Minimum OceanBase version that supports namespace-enabled collections. +NAMESPACE_MIN_OB_VERSION = Version("4.6.1.0") + logger = logging.getLogger(__name__) from .types import _NOT_PROVIDED, _NotProvided # noqa: E402, F401 @@ -412,9 +416,14 @@ class BaseClient(BaseConnection, AdminAPI): # ==================== Database Type Detection ==================== def _validate_ob_database_type(self) -> None: - db_type, _version = self.detect_db_type_and_version() + db_type, version = self.detect_db_type_and_version() if db_type.lower() != "oceanbase": raise ValueError("use_namespace=True is only supported on OceanBase") + if version < NAMESPACE_MIN_OB_VERSION: + raise ValueError( + f"use_namespace=True requires OceanBase version >= {NAMESPACE_MIN_OB_VERSION}, " + f"current version is {version}" + ) def _is_shared_storage_mode(self) -> bool: try: @@ -939,6 +948,11 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " if hnsw_config is not None: raise ValueError("use_namespace=True only supports IVF index type, HNSW is not allowed") + if ivf_config is not None and ivf_config.type != IVFIndexType.IVF_FLAT.value: + raise ValueError( + f"use_namespace=True currently only supports IVF index type '{IVFIndexType.IVF_FLAT.value}', " + f"got '{ivf_config.type}'" + ) self._validate_ob_database_type() if ivf_config is not None: diff --git a/tests/integration_tests/test_namespace_constraints.py b/tests/integration_tests/test_namespace_constraints.py new file mode 100644 index 00000000..76253eb8 --- /dev/null +++ b/tests/integration_tests/test_namespace_constraints.py @@ -0,0 +1,95 @@ +""" +Namespace constraint integration tests. + +Covers two SDK-level constraints for use_namespace=True collections: +1. Minimum OceanBase version: namespace-enabled collections require OB >= 4.6.1. +2. Vector index type: only IVF_FLAT is currently supported; ivf_sq8 / ivf_pq are + rejected at the SDK layer before any DDL is issued. +""" + +import time + +import pytest + +from pyseekdb import IVFConfiguration +from pyseekdb.client.client_base import NAMESPACE_MIN_OB_VERSION +from pyseekdb.client.configuration import VectorIndexConfig +from pyseekdb.client.schema import Schema +from pyseekdb.client.version import Version + + +def _make_schema(ivf_type: str) -> Schema: + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine", type=ivf_type, fresh_mode="spfresh"), + embedding_function=None, + ), + ) + + +def _unique_name(suffix: str) -> str: + return f"test_ns_constraint_{int(time.time() * 1000)}{suffix}" + + +class TestNamespaceIvfTypeConstraint: + @pytest.mark.parametrize("ivf_type", ["ivf_pq", "ivf_sq8"]) + def test_non_ivf_flat_rejected(self, oceanbase_client, ivf_type): + """Non-ivf_flat IVF types must be rejected at the SDK layer.""" + name = _unique_name(f"_{ivf_type}") + with pytest.raises(ValueError, match="only supports IVF index type 'ivf_flat'"): + oceanbase_client.create_collection( + name=name, schema=_make_schema(ivf_type), use_namespace=True + ) + # Rejection happens before any DDL: no collection metadata is left behind. + assert not oceanbase_client.has_collection(name) + + def test_ivf_flat_accepted(self, oceanbase_client): + """Positive control: ivf_flat namespace collection creation succeeds.""" + name = _unique_name("_ivf_flat") + collection = oceanbase_client.create_collection( + name=name, schema=_make_schema("ivf_flat"), use_namespace=True + ) + try: + assert collection.use_namespace is True + assert collection.dimension == 3 + finally: + oceanbase_client.delete_collection(name=name) + + +class TestNamespaceMinVersionConstraint: + def test_connected_ob_meets_min_version(self, oceanbase_client): + """The kernel under test must already be >= 4.6.1, and creation succeeds.""" + db_type, version = oceanbase_client._server.detect_db_type_and_version() + assert db_type.lower() == "oceanbase" + assert version >= NAMESPACE_MIN_OB_VERSION, ( + f"connected OB version {version} < required {NAMESPACE_MIN_OB_VERSION}" + ) + + name = _unique_name("_ver_ok") + collection = oceanbase_client.create_collection( + name=name, schema=_make_schema("ivf_flat"), use_namespace=True + ) + try: + assert collection.use_namespace is True + finally: + oceanbase_client.delete_collection(name=name) + + def test_old_ob_version_rejected(self, oceanbase_client, monkeypatch): + """Simulate an older OB kernel: namespace creation must fail with a clear error.""" + server = oceanbase_client._server + monkeypatch.setattr( + type(server), + "detect_db_type_and_version", + lambda self: ("oceanbase", Version("4.6.0.0")), + ) + name = _unique_name("_ver_old") + with pytest.raises(ValueError, match=r"requires OceanBase version >= 4\.6\.1"): + oceanbase_client.create_collection( + name=name, schema=_make_schema("ivf_flat"), use_namespace=True + ) + # restore before checking leftovers + monkeypatch.undo() + assert not oceanbase_client.has_collection(name) + + def test_min_version_constant(self): + assert NAMESPACE_MIN_OB_VERSION == Version("4.6.1.0") From 9e284ee067727bcb8eb14bf10d645cab60ade786 Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Mon, 15 Jun 2026 14:11:41 +0800 Subject: [PATCH 40/84] add prewarm for namespace test case --- src/pyseekdb/client/client_base.py | 4 +- .../namespace_dml_helpers.py | 9 +++++ .../namespace_fts_helpers.py | 11 +++++- tests/integration_tests/test_namespace_dml.py | 38 +++++++++++-------- .../test_namespace_dml_multi_coll.py | 4 ++ .../test_namespace_dml_multi_coll_multi_ns.py | 5 ++- .../test_namespace_dml_multi_ns.py | 10 +++++ .../test_namespace_drop_validation.py | 33 +++++++++------- .../test_namespace_get_delete_filters.py | 8 ++++ ...rid_search_fulltext_multi_coll_multi_ns.py | 10 ++++- ...earch_triple_branch_multi_coll_multi_ns.py | 10 ++++- .../test_namespace_lifecycle.py | 3 ++ .../test_namespace_prewarm.py | 12 ++---- .../integration_tests/test_namespace_query.py | 2 + .../test_namespace_session_vars.py | 2 + 15 files changed, 114 insertions(+), 47 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 2bbbe4e6..8d9bf953 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1358,10 +1358,10 @@ def _create_namespace_physical_tables( def _cleanup_namespace_physical_tables(self, collection_id: str) -> None: for suffix_fn in [ + NamespaceCollectionNames.data_table_name, NamespaceCollectionNames.logic_schema_table_name, NamespaceCollectionNames.kv_data_table_name, - NamespaceCollectionNames.hot_table_name, - NamespaceCollectionNames.data_table_name, + NamespaceCollectionNames.hot_table_name ]: with contextlib.suppress(Exception): self._execute(f"DROP TABLE IF EXISTS `{suffix_fn(collection_id)}`") diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index dc1051e6..74d0559b 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -12,6 +12,8 @@ from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema +NAMESPACE_TEST_PARTITION_COUNT = 4 + LARGE_DML_CORPUS_SIZE = 1000 DML_BATCH_SIZE = 100 # Wait for FULLTEXT index after bulk load (same as namespace_fts_helpers). @@ -43,7 +45,14 @@ def ns_schema() -> Schema: ) +def use_namespace_test_partitions() -> None: + from pyseekdb import set_namespace_partition_count + + set_namespace_partition_count(NAMESPACE_TEST_PARTITION_COUNT) + + def create_ns_collection(client: Any, suffix: str = "") -> Any: + use_namespace_test_partitions() name = f"test_ns_dml_{int(time.time() * 1000)}{suffix}" return client.create_collection(name=name, schema=ns_schema(), use_namespace=True) diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index c8acd7e4..ed7fb374 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -17,6 +17,7 @@ VectorDistanceMetric = Literal["l2", "cosine"] VECTOR_DISTANCE_METRICS: tuple[VectorDistanceMetric, ...] = ("l2", "cosine") +from namespace_dml_helpers import use_namespace_test_partitions from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema @@ -512,6 +513,7 @@ def setup_large_fts_collection( raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") name = f"test_ns_hs_ft_{distance}_{int(time.time() * 1000)}" + use_namespace_test_partitions() collection = db_client.create_collection( name=name, schema=ns_schema(distance), use_namespace=True ) @@ -526,6 +528,7 @@ def setup_fts_namespace_with_corpus( ) -> Any: """Create a namespace under an existing collection and bulk-load the corpus.""" namespace = collection.create_namespace(namespace_name) + namespace.prewarm() insert_corpus_into_namespace(namespace, corpus) expected = len(corpus) actual = namespace.count() @@ -571,6 +574,7 @@ def _create_multi_coll_multi_ns_layout( distance: VectorDistanceMetric = "l2", ) -> dict[str, Any]: ts = int(time.time() * 1000) + use_namespace_test_partitions() coll_1 = db_client.create_collection( name=f"{name_prefix}_{distance}_{ts}_c1", schema=ns_schema(distance), @@ -584,7 +588,9 @@ def _create_multi_coll_multi_ns_layout( ctx: dict[str, Any] = {"coll_1": coll_1, "coll_2": coll_2} for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): for ns_suffix in ("x", "y"): - ctx[f"{coll_tag}_{ns_suffix}"] = collection.create_namespace(f"ns_{ns_suffix}") + ns = collection.create_namespace(f"ns_{ns_suffix}") + ns.prewarm() + ctx[f"{coll_tag}_{ns_suffix}"] = ns return ctx @@ -627,6 +633,7 @@ def setup_same_collection_both_ns_fts(db_client: Any) -> dict[str, Any]: remain visible via get. """ ts = int(time.time() * 1000) + use_namespace_test_partitions() collection = db_client.create_collection( name=f"test_ns_hs_ft_sc2ns_{ts}", schema=ns_schema(), @@ -635,6 +642,8 @@ def setup_same_collection_both_ns_fts(db_client: Any) -> dict[str, Any]: corpus = build_large_fts_corpus(CORPUS_SIZE) ns_x = collection.create_namespace("ns_x") ns_y = collection.create_namespace("ns_y") + ns_x.prewarm() + ns_y.prewarm() insert_corpus_in_batches(ns_x, corpus) insert_corpus_in_batches(ns_y, corpus) time.sleep(INDEX_SETTLE_SECONDS) diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py index eaca0277..b3225a71 100644 --- a/tests/integration_tests/test_namespace_dml.py +++ b/tests/integration_tests/test_namespace_dml.py @@ -23,12 +23,18 @@ ) +def _create_namespace(collection, name: str): + ns = collection.create_namespace(name) + ns.prewarm() + return ns + + class TestNamespaceDML: """Small-scale DML smoke tests.""" def test_add_single(self, db_client): collection = create_ns_collection(db_client, suffix="_add1") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0], documents="Hello", metadatas={"tag": "a"}) result = ns.get(ids="d1") @@ -39,7 +45,7 @@ def test_add_single(self, db_client): def test_add_batch(self, db_client): collection = create_ns_collection(db_client, suffix="_addb") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.add( ids=["d1", "d2", "d3"], @@ -54,7 +60,7 @@ def test_add_batch(self, db_client): def test_get_by_id(self, db_client): collection = create_ns_collection(db_client, suffix="_getid") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.add( ids=["g1", "g2"], @@ -71,7 +77,7 @@ def test_get_by_id(self, db_client): def test_get_with_limit(self, db_client): collection = create_ns_collection(db_client, suffix="_getlim") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.add( ids=["l1", "l2", "l3"], @@ -84,7 +90,7 @@ def test_get_with_limit(self, db_client): def test_update_metadata(self, db_client): collection = create_ns_collection(db_client, suffix="_upd") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.add(ids="u1", embeddings=[1.0, 2.0, 3.0], metadatas={"score": 10}) ns.update(ids="u1", metadatas={"score": 99}) @@ -95,7 +101,7 @@ def test_update_metadata(self, db_client): def test_update_document_and_embedding(self, db_client): collection = create_ns_collection(db_client, suffix="_upddoc") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.add(ids="u2", embeddings=[1.0, 2.0, 3.0], documents="Original") ns.update(ids="u2", embeddings=[9.0, 8.0, 7.0], documents="Updated") @@ -106,7 +112,7 @@ def test_update_document_and_embedding(self, db_client): def test_upsert_existing(self, db_client): collection = create_ns_collection(db_client, suffix="_upsex") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.add(ids="up1", embeddings=[1.0, 2.0, 3.0], metadatas={"v": 1}) ns.upsert(ids="up1", embeddings=[4.0, 5.0, 6.0], metadatas={"v": 2}) @@ -117,7 +123,7 @@ def test_upsert_existing(self, db_client): def test_upsert_new(self, db_client): collection = create_ns_collection(db_client, suffix="_upsnew") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.upsert(ids="up_new", embeddings=[1.0, 1.0, 1.0], metadatas={"fresh": True}) result = ns.get(ids="up_new", include=["metadatas"]) @@ -128,7 +134,7 @@ def test_upsert_new(self, db_client): def test_delete_by_ids(self, db_client): collection = create_ns_collection(db_client, suffix="_del") - ns = collection.create_namespace("dml_ns") + ns = _create_namespace(collection,"dml_ns") try: ns.add( ids=["del1", "del2"], @@ -147,7 +153,7 @@ class TestNamespaceDMLCountPeekAtScale: @pytest.fixture def large_ns(self, db_client): collection = create_ns_collection(db_client, suffix="_large_cp") - ns = collection.create_namespace("large_ns") + ns = _create_namespace(collection,"large_ns") corpus = build_large_dml_corpus(LARGE_DML_CORPUS_SIZE, id_prefix="large") load_dml_corpus(ns, corpus) yield collection, ns, corpus @@ -155,7 +161,7 @@ def large_ns(self, db_client): def test_count_empty_namespace(self, db_client): collection = create_ns_collection(db_client, suffix="_cnt_empty") - ns = collection.create_namespace("empty_ns") + ns = _create_namespace(collection,"empty_ns") try: assert ns.count() == 0 finally: @@ -175,7 +181,7 @@ def test_count_after_partial_delete(self, large_ns): def test_peek_empty_namespace(self, db_client): collection = create_ns_collection(db_client, suffix="_peek_empty") - ns = collection.create_namespace("peek_empty") + ns = _create_namespace(collection,"peek_empty") try: result = ns.peek(limit=10) assert_peek_result(result, expected_len=0) @@ -230,8 +236,8 @@ class TestNamespaceDMLCountPeekMultiNs: @pytest.fixture def dual_ns_ctx(self, db_client): collection = create_ns_collection(db_client, suffix="_dual_cp") - ns_a = collection.create_namespace("ns_alpha") - ns_b = collection.create_namespace("ns_beta") + ns_a = _create_namespace(collection,"ns_alpha") + ns_b = _create_namespace(collection,"ns_beta") corpus_a = build_large_dml_corpus( self.NS_A_SIZE, id_prefix="alpha", ns_tag="alpha" ) @@ -332,7 +338,7 @@ class TestNamespaceDMLFullCycle: def test_full_dml_cycle_verified_by_get(self, db_client): collection = create_ns_collection(db_client, suffix="_cycle") - ns = collection.create_namespace("cycle_ns") + ns = _create_namespace(collection,"cycle_ns") doc_id = "cycle_doc" try: ns.add( @@ -385,7 +391,7 @@ def test_full_dml_cycle_verified_by_get(self, db_client): def test_batch_add_then_get_each(self, db_client): collection = create_ns_collection(db_client, suffix="_batch") - ns = collection.create_namespace("batch_ns") + ns = _create_namespace(collection,"batch_ns") try: ns.add( ids=["b1", "b2", "b3"], diff --git a/tests/integration_tests/test_namespace_dml_multi_coll.py b/tests/integration_tests/test_namespace_dml_multi_coll.py index 72d064a4..061d951c 100644 --- a/tests/integration_tests/test_namespace_dml_multi_coll.py +++ b/tests/integration_tests/test_namespace_dml_multi_coll.py @@ -20,6 +20,8 @@ def test_same_namespace_name_across_collections(self, db_client): coll_b = create_ns_collection(db_client, suffix="_mc_b") ns_a = coll_a.create_namespace("tenant") ns_b = coll_b.create_namespace("tenant") + ns_a.prewarm() + ns_b.prewarm() try: ns_a.add( ids="doc1", @@ -48,6 +50,8 @@ def test_update_delete_isolated_by_collection(self, db_client): coll_b = create_ns_collection(db_client, suffix="_mc_upd_b") ns_a = coll_a.create_namespace("tenant") ns_b = coll_b.create_namespace("tenant") + ns_a.prewarm() + ns_b.prewarm() try: ns_a.add(ids="doc1", embeddings=[1.0, 0.0, 0.0], documents="Original A") ns_b.add(ids="doc1", embeddings=[0.0, 1.0, 0.0], documents="Original B") diff --git a/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py index 8e0dc09f..668fde8c 100644 --- a/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py @@ -18,7 +18,7 @@ class TestNamespaceDMLMultiCollMultiNs: def _setup_quadrants(self, db_client): coll_1 = create_ns_collection(db_client, suffix="_mcmn_c1") coll_2 = create_ns_collection(db_client, suffix="_mcmn_c2") - return { + quadrants = { "coll_1": coll_1, "coll_2": coll_2, "c1_x": coll_1.create_namespace("ns_x"), @@ -26,6 +26,9 @@ def _setup_quadrants(self, db_client): "c2_x": coll_2.create_namespace("ns_x"), "c2_y": coll_2.create_namespace("ns_y"), } + for key in ("c1_x", "c1_y", "c2_x", "c2_y"): + quadrants[key].prewarm() + return quadrants def test_cross_quadrant_isolation(self, db_client): q = self._setup_quadrants(db_client) diff --git a/tests/integration_tests/test_namespace_dml_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_ns.py index 5a98feee..f77094da 100644 --- a/tests/integration_tests/test_namespace_dml_multi_ns.py +++ b/tests/integration_tests/test_namespace_dml_multi_ns.py @@ -19,6 +19,8 @@ def test_add_isolation(self, db_client): collection = create_ns_collection(db_client, suffix="_mns_add") ns_a = collection.create_namespace("ns_a") ns_b = collection.create_namespace("ns_b") + ns_a.prewarm() + ns_b.prewarm() try: ns_a.add(ids="a1", embeddings=[1.0, 0.0, 0.0], metadatas={"owner": "A"}) ns_b.add(ids="b1", embeddings=[0.0, 1.0, 0.0], metadatas={"owner": "B"}) @@ -34,6 +36,8 @@ def test_same_id_different_namespaces(self, db_client): collection = create_ns_collection(db_client, suffix="_mns_sameid") ns_x = collection.create_namespace("ns_x") ns_y = collection.create_namespace("ns_y") + ns_x.prewarm() + ns_y.prewarm() try: ns_x.add(ids="shared", embeddings=[1.0, 0.0, 0.0], metadatas={"src": "X"}) ns_y.add(ids="shared", embeddings=[0.0, 1.0, 0.0], metadatas={"src": "Y"}) @@ -47,6 +51,8 @@ def test_update_does_not_affect_other_ns(self, db_client): collection = create_ns_collection(db_client, suffix="_mns_upd") ns_a = collection.create_namespace("ns_a") ns_b = collection.create_namespace("ns_b") + ns_a.prewarm() + ns_b.prewarm() doc_id = "shared_upd" try: ns_a.add(ids=doc_id, embeddings=[1.0, 0.0, 0.0], documents="A original") @@ -65,6 +71,8 @@ def test_delete_only_target_ns(self, db_client): collection = create_ns_collection(db_client, suffix="_mns_del") ns_a = collection.create_namespace("ns_a") ns_b = collection.create_namespace("ns_b") + ns_a.prewarm() + ns_b.prewarm() doc_id = "shared_del" try: ns_a.add(ids=doc_id, embeddings=[1.0, 0.0, 0.0]) @@ -83,6 +91,8 @@ def test_upsert_isolation(self, db_client): collection = create_ns_collection(db_client, suffix="_mns_ups") ns_a = collection.create_namespace("ns_a") ns_b = collection.create_namespace("ns_b") + ns_a.prewarm() + ns_b.prewarm() try: ns_a.add(ids="a_only", embeddings=[1.0, 0.0, 0.0], metadatas={"ns": "a"}) ns_b.upsert(ids="b_only", embeddings=[0.0, 1.0, 0.0], metadatas={"ns": "b"}) diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index b60f5b10..3ad46205 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -4,6 +4,7 @@ import pytest import pyseekdb +from namespace_dml_helpers import use_namespace_test_partitions from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.meta_info import NamespaceCollectionNames @@ -22,10 +23,17 @@ # --------------------------------------------------------------------------- # +def _create_namespace(collection, name: str): + ns = collection.create_namespace(name) + ns.prewarm() + return ns + + def _make_collection(client, suffix: str = ""): """Create a namespace-mode collection with an IVF index (so we get a full set of physical tables: logic_data / kv_data / logic_schema / hot_table).""" name = f"test_ns_drop_{int(time.time() * 1000)}{suffix}" + use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), @@ -350,7 +358,7 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = collection.create_namespace("ns_validate") + ns = _create_namespace(collection,"ns_validate") ns_id = int(ns.namespace_id) # Pre-conditions @@ -363,8 +371,6 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): "create_namespace should have inserted a logic_schema row" ) if is_ss: - # Seed a hot_table row so we can verify the SS-only delete. - assert _seed_hot_table(client, coll_id, ns_id) is True assert _count_hot_table_rows(client, coll_id, ns_id) == 1 # Seed logic_data rows — sync DROP must not touch them; LTABLE_BG ns DAG @@ -413,7 +419,7 @@ def test_drop_namespace_async_cleanup(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = collection.create_namespace("ns_async") + ns = _create_namespace(collection,"ns_async") ns_id = int(ns.namespace_id) lt_id = _fetch_ltable_id(client, coll_id, ns_id) assert lt_id is not None @@ -453,7 +459,7 @@ def test_async_skipped_while_active_ltable_remains(self, oceanbase_client): coll_id = collection.id blocker_lt_id = 9_000_001 try: - ns = collection.create_namespace("ns_block") + ns = _create_namespace(collection,"ns_block") ns_id = int(ns.namespace_id) _drop_namespace_via_pl(client, coll_id, ns_id) assert _fetch_namespace_name(client, coll_id, ns_id).startswith(RECYCLEBIN_PREFIX) @@ -491,7 +497,7 @@ def test_async_no_kv_data_fast_cleanup(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = collection.create_namespace("ns_empty_kv") + ns = _create_namespace(collection,"ns_empty_kv") ns_id = int(ns.namespace_id) assert _count_kv_data_rows(client, coll_id, ns_id) == 0 @@ -516,9 +522,8 @@ def test_async_hot_table_cleaned_on_non_ss(self, oceanbase_client): if hot_err is not None: pytest.fail(f"hot_table setup failed: {hot_err}") try: - ns = collection.create_namespace("ns_hot_async") + ns = _create_namespace(collection, "ns_hot_async") ns_id = int(ns.namespace_id) - assert _seed_hot_table(client, coll_id, ns_id) is True assert _count_hot_table_rows(client, coll_id, ns_id) == 1 _drop_namespace_via_pl(client, coll_id, ns_id) @@ -543,7 +548,7 @@ def test_async_drop_history_recorded(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = collection.create_namespace("ns_history") + ns = _create_namespace(collection,"ns_history") ns_id = int(ns.namespace_id) _seed_kv_data(client, coll_id, ns_id, count=2) _drop_namespace_via_pl(client, coll_id, ns_id) @@ -572,7 +577,7 @@ def test_recyclebin_name_format(self, oceanbase_client): client = oceanbase_client collection = _make_collection(client) try: - ns = collection.create_namespace("fmt_ns") + ns = _create_namespace(collection,"fmt_ns") ns_id = int(ns.namespace_id) _drop_namespace_via_pl(client, collection.id, ns_id) new_name = _fetch_namespace_name(client, collection.id, ns_id) @@ -592,7 +597,7 @@ def test_multiple_ltables_all_deleted(self, oceanbase_client): client = oceanbase_client collection = _make_collection(client) try: - ns = collection.create_namespace("ns_multi") + ns = _create_namespace(collection,"ns_multi") ns_id = int(ns.namespace_id) # Insert two extra ltable rows directly so we can verify bulk delete. lt_table = _catalog_table(client, "sdk_ltables") @@ -624,7 +629,7 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): schema_tbl = NamespaceCollectionNames.logic_schema_table_name(coll_id) schema_tbl_q = f"`{client._server.database}`.`{schema_tbl}`" try: - ns = collection.create_namespace("ns_rollback") + ns = _create_namespace(collection,"ns_rollback") ns_id = int(ns.namespace_id) orig_name = _fetch_namespace_name(client, coll_id, ns_id) assert orig_name == "ns_rollback" @@ -671,7 +676,7 @@ def test_drop_namespace_is_idempotent(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = collection.create_namespace("ns_idem") + ns = _create_namespace(collection,"ns_idem") ns_id = int(ns.namespace_id) _drop_namespace_via_pl(client, coll_id, ns_id) first_name = _fetch_namespace_name(client, coll_id, ns_id) @@ -706,7 +711,7 @@ def test_concurrent_drop_from_two_sdks(self, oceanbase_client): collection = _make_collection(client_a, suffix="_conc") coll_id = collection.id try: - ns = collection.create_namespace("ns_concurrent") + ns = _create_namespace(collection,"ns_concurrent") ns_id = int(ns.namespace_id) results = {} diff --git a/tests/integration_tests/test_namespace_get_delete_filters.py b/tests/integration_tests/test_namespace_get_delete_filters.py index 61130084..dcac545d 100644 --- a/tests/integration_tests/test_namespace_get_delete_filters.py +++ b/tests/integration_tests/test_namespace_get_delete_filters.py @@ -33,6 +33,7 @@ class TestNamespaceDeleteWhere: def test_delete_where_metadata_tag(self, db_client): collection = create_ns_collection(db_client, suffix="_del_where") ns = collection.create_namespace("del_where_ns") + ns.prewarm() corpus, gt = build_filter_corpus() try: load_dml_corpus(ns, corpus) @@ -55,6 +56,7 @@ def test_delete_where_metadata_tag(self, db_client): def test_delete_where_document_obsolete(self, db_client): collection = create_ns_collection(db_client, suffix="_del_wdoc") ns = collection.create_namespace("del_wdoc_ns") + ns.prewarm() corpus, gt = build_filter_corpus() where_document = {"$contains": "obsolete"} try: @@ -94,6 +96,7 @@ class TestNamespaceGetWhere: def test_get_where_metadata_category_ai(self, db_client): collection = create_ns_collection(db_client, suffix="_get_where") ns = collection.create_namespace("get_where_ns") + ns.prewarm() corpus, gt = build_filter_corpus() try: load_dml_corpus(ns, corpus) @@ -120,6 +123,7 @@ def test_get_where_metadata_category_ai(self, db_client): def test_get_where_document_contains_python(self, db_client): collection = create_ns_collection(db_client, suffix="_get_wdoc") ns = collection.create_namespace("get_wdoc_ns") + ns.prewarm() corpus, gt = build_filter_corpus() where_document = {"$contains": "python"} try: @@ -152,6 +156,8 @@ def test_delete_where_only_affects_target_namespace(self, db_client): collection = create_ns_collection(db_client, suffix="_filt_mns") ns_a = collection.create_namespace("ns_alpha") ns_b = collection.create_namespace("ns_beta") + ns_a.prewarm() + ns_b.prewarm() corpus_a, gt = build_filter_corpus( keep_count=FILTER_MULTI_NS_KEEP, purge_count=FILTER_MULTI_NS_PURGE, @@ -209,6 +215,8 @@ def test_get_where_alternates_namespaces_after_delete(self, db_client): collection = create_ns_collection(db_client, suffix="_cross_ns_get") ns_a = collection.create_namespace("cross_alpha") ns_b = collection.create_namespace("cross_beta") + ns_a.prewarm() + ns_b.prewarm() corpus_a, gt = build_filter_corpus( keep_count=50, purge_count=50, diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py index 32fa76e0..e906ccb4 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py @@ -17,6 +17,7 @@ import pytest +from namespace_dml_helpers import use_namespace_test_partitions from namespace_fts_helpers import ( CORPUS_SIZE, MULTI_COLL_MULTI_NS_QUADRANT_KEYS, @@ -58,6 +59,7 @@ def _setup_multi_coll_multi_ns_isolation_fts(db_client): sibling namespace fails the existing corpus-based assertions. """ ts = int(time.time() * 1000) + use_namespace_test_partitions() coll_1 = db_client.create_collection( name=f"test_ns_hs_ft_mcmn_iso_{ts}_c1", schema=ns_schema(), @@ -71,8 +73,12 @@ def _setup_multi_coll_multi_ns_isolation_fts(db_client): ctx: dict[str, object] = {"coll_1": coll_1, "coll_2": coll_2} for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): for ns_suffix in ("x", "y"): - ctx[f"{coll_tag}_{ns_suffix}"] = collection.create_namespace(f"ns_{ns_suffix}") - ctx[f"{coll_tag}_empty"] = collection.create_namespace("ns_empty") + ns = collection.create_namespace(f"ns_{ns_suffix}") + ns.prewarm() + ctx[f"{coll_tag}_{ns_suffix}"] = ns + empty_ns = collection.create_namespace("ns_empty") + empty_ns.prewarm() + ctx[f"{coll_tag}_empty"] = empty_ns base_corpus = build_large_fts_corpus(CORPUS_SIZE) if len(base_corpus) <= 1000: diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py index 355eef53..c46c1936 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py @@ -16,6 +16,7 @@ import pytest +from namespace_dml_helpers import use_namespace_test_partitions from namespace_fts_helpers import ( CORPUS_SIZE, MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, @@ -80,6 +81,7 @@ def _setup_multi_coll_multi_ns_isolation_triple( ): """2 collections x 2 loaded namespaces with distinct per-namespace corpus markers.""" ts = int(time.time() * 1000) + use_namespace_test_partitions() coll_1 = db_client.create_collection( name=f"test_ns_hs_tb_mcmn_iso_{distance}_{ts}_c1", schema=ns_schema(distance), @@ -93,8 +95,12 @@ def _setup_multi_coll_multi_ns_isolation_triple( ctx: dict[str, object] = {"coll_1": coll_1, "coll_2": coll_2} for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): for ns_suffix in ("x", "y"): - ctx[f"{coll_tag}_{ns_suffix}"] = collection.create_namespace(f"ns_{ns_suffix}") - ctx[f"{coll_tag}_empty"] = collection.create_namespace("ns_empty") + ns = collection.create_namespace(f"ns_{ns_suffix}") + ns.prewarm() + ctx[f"{coll_tag}_{ns_suffix}"] = ns + empty_ns = collection.create_namespace("ns_empty") + empty_ns.prewarm() + ctx[f"{coll_tag}_empty"] = empty_ns base_corpus = build_large_fts_corpus(CORPUS_SIZE) if len(base_corpus) <= 1000: diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 16fb61d9..dcf6ef5b 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -8,6 +8,7 @@ import pytest import pyseekdb +from namespace_dml_helpers import use_namespace_test_partitions from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema @@ -18,6 +19,7 @@ class TestNamespaceLifecycle: def _create_ns_collection(self, client, suffix=""): name = f"test_ns_lc_{int(time.time() * 1000)}{suffix}" + use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), @@ -170,6 +172,7 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): type(client._server), "_is_shared_storage_mode", return_value=True ): name = f"test_ns_ss_{int(time.time() * 1000)}" + use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index 32f0e052..2c7c86f3 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -17,17 +17,11 @@ class TestNamespacePrewarm: - @pytest.fixture(autouse=True) - def _reduce_namespace_partitions(self): - """Override conftest: use 8 partitions for prewarm until DDL is faster.""" - from pyseekdb import get_namespace_partition_count, set_namespace_partition_count - original = get_namespace_partition_count() - set_namespace_partition_count(8) - yield - set_namespace_partition_count(original) - def _create_ns_collection_and_namespace(self, client): + from namespace_dml_helpers import use_namespace_test_partitions + name = f"test_ns_pw_{int(time.time() * 1000)}" + use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index 6916c686..0ecd0795 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -7,6 +7,7 @@ import pytest +from namespace_dml_helpers import use_namespace_test_partitions from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema @@ -16,6 +17,7 @@ class TestNamespaceQuery: def _setup(self, client, suffix=""): name = f"test_ns_q_{int(time.time() * 1000)}{suffix}" + use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py index 83983a50..ddf7d379 100644 --- a/tests/integration_tests/test_namespace_session_vars.py +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -9,6 +9,7 @@ import pytest +from namespace_dml_helpers import use_namespace_test_partitions from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema @@ -18,6 +19,7 @@ class TestNamespaceSessionVars: def _create_ns_collection(self, client): name = f"test_ns_sessvar_{int(time.time() * 1000)}" + use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), From da44ebea20bb9de0108a380e8f72a381df3232b5 Mon Sep 17 00:00:00 2001 From: "luohongdi.lhd" Date: Mon, 15 Jun 2026 14:24:52 +0800 Subject: [PATCH 41/84] fix run --- tests/integration_tests/namespace_dml_helpers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index 6791e007..ad38c595 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -47,9 +47,9 @@ def ns_schema() -> Schema: def use_namespace_test_partitions() -> None: - from pyseekdb import set_namespace_partition_count + from pyseekdb import set_collection_partition_count - set_namespace_partition_count(NAMESPACE_TEST_PARTITION_COUNT) + set_collection_partition_count(NAMESPACE_TEST_PARTITION_COUNT) def create_ns_collection(client: Any, suffix: str = "") -> Any: From ad3697ac16f784caa1cedf06cafe6305e4d88f53 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 16 Jun 2026 10:20:10 +0800 Subject: [PATCH 42/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91collection?= =?UTF-8?q?=E8=A2=AB=E5=88=A0=E9=99=A4=E4=B9=8B=E5=90=8E=EF=BC=8C=E4=BB=8D?= =?UTF-8?q?=E8=83=BD=E6=89=A7=E8=A1=8Cns=20=3D=20coll.create=5Fnamespace?= =?UTF-8?q?=20https://project.alipay.com/project/W24001004732/P26001026476?= =?UTF-8?q?/workitemDetail=3FopenWorkItemId=3D2026061500116762982?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 13 +++++++++ src/pyseekdb/client/collection.py | 27 ++++++++++--------- .../test_namespace_lifecycle.py | 13 +++++++++ 3 files changed, 41 insertions(+), 12 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 18d104a7..46d232d8 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1266,6 +1266,19 @@ def _has_ns_collection(self, collection_name: str) -> bool: except Exception: return False + def _ns_collection_exists_by_id(self, collection_id: str) -> bool: + """Whether a namespace-enabled collection with this id still exists. + + Used to reject namespace operations on a stale Collection handle whose + underlying collection was deleted (the in-memory handle keeps its old id). + """ + collection_id_escaped = escape_string(str(collection_id)) + rows = self._execute( + f"SELECT collection_id FROM {self._qtable(CollectionNames.sdk_collections_table_name())} " + f"WHERE collection_id = '{collection_id_escaped}'" + ) + return bool(rows) + def _delete_ns_collection_meta(self, collection_name: str) -> None: meta = self._get_ns_collection_meta(collection_name) if meta is None: diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 18e2ff19..7aec5ff7 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -115,6 +115,15 @@ def _guard_collection_data_api(self) -> None: "Get a namespace via collection.create_namespace() or collection.get_namespace()." ) + def _guard_namespace_enabled(self) -> None: + if not self._use_namespace: + raise ValueError("Namespace is not enabled for this collection. Use use_namespace=True when creating the collection.") + if not self._client._ns_collection_exists_by_id(self._id): + raise ValueError( + f"Collection '{self._name}' no longer exists (it may have been deleted). " + "Namespace operations are not allowed on a deleted collection." + ) + def __repr__(self) -> str: ns_str = ", use_namespace=True" if self._use_namespace else "" return f"Collection(name='{self._name}', dimension={self._dimension}{ns_str}, client={self._client.mode})" @@ -122,16 +131,14 @@ def __repr__(self) -> str: # ==================== Namespace Management ==================== def create_namespace(self, name: str) -> "Namespace": - if not self._use_namespace: - raise ValueError("Namespace is not enabled for this collection. Use use_namespace=True when creating the collection.") + self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace meta = self._client._create_ns_namespace_meta(self._id, name) return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) def get_namespace(self, name: str) -> "Namespace": - if not self._use_namespace: - raise ValueError("Namespace is not enabled for this collection.") + self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace meta = self._client._get_ns_namespace_meta(self._id, name) @@ -140,8 +147,7 @@ def get_namespace(self, name: str) -> "Namespace": return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) def get_or_create_namespace(self, name: str) -> "Namespace": - if not self._use_namespace: - raise ValueError("Namespace is not enabled for this collection.") + self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace meta = self._client._get_ns_namespace_meta(self._id, name) @@ -150,14 +156,12 @@ def get_or_create_namespace(self, name: str) -> "Namespace": return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) def delete_namespace(self, name: str) -> None: - if not self._use_namespace: - raise ValueError("Namespace is not enabled for this collection.") + self._guard_namespace_enabled() _validate_namespace_name(name) self._client._delete_ns_namespace_meta(self._id, name) def list_namespaces(self) -> list["Namespace"]: - if not self._use_namespace: - raise ValueError("Namespace is not enabled for this collection.") + self._guard_namespace_enabled() from .namespace import Namespace metas = self._client._list_ns_namespaces(self._id) return [ @@ -166,8 +170,7 @@ def list_namespaces(self) -> list["Namespace"]: ] def has_namespace(self, name: str) -> bool: - if not self._use_namespace: - raise ValueError("Namespace is not enabled for this collection.") + self._guard_namespace_enabled() _validate_namespace_name(name) return self._client._has_ns_namespace(self._id, name) diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index b3080220..b4244ee4 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -120,6 +120,19 @@ def test_delete_collection_cleans_namespaces(self, db_client): assert db_client.has_collection(coll_name) is False + def test_namespace_ops_blocked_after_collection_deleted(self, db_client): + collection = self._create_ns_collection(db_client) + db_client.delete_collection(name=collection.name) + + with pytest.raises(ValueError, match="no longer exists"): + collection.create_namespace("demo") + with pytest.raises(ValueError, match="no longer exists"): + collection.get_or_create_namespace("demo") + with pytest.raises(ValueError, match="no longer exists"): + collection.list_namespaces() + with pytest.raises(ValueError, match="no longer exists"): + collection.has_namespace("demo") + def test_collection_data_api_blocked_when_namespace_enabled(self, db_client): collection = self._create_ns_collection(db_client) try: From 81b36cfde63d8a31a315d73649e83a070ebc4bcf Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 16 Jun 2026 10:37:26 +0800 Subject: [PATCH 43/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91namespace?= =?UTF-8?q?=E5=88=A0=E9=99=A4=E5=90=8E=EF=BC=8C=E6=89=A7=E8=A1=8Cns.add?= =?UTF-8?q?=EF=BC=8C=E4=BB=8D=E6=97=A7=E6=89=A7=E8=A1=8C=E6=88=90=E5=8A=9F?= =?UTF-8?q?=20https://project.alipay.com/project/W24001004732/P26001026476?= =?UTF-8?q?/workitemDetail=3FopenWorkItemId=3D2026061500116761230?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 18 +++++++++++++ src/pyseekdb/client/namespace.py | 17 ++++++++++++ tests/integration_tests/test_namespace_dml.py | 27 +++++++++++++++++++ 3 files changed, 62 insertions(+) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 46d232d8..57daeb42 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1524,6 +1524,24 @@ def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dic def _has_ns_namespace(self, collection_id: str, namespace_name: str) -> bool: return self._get_ns_namespace_meta(collection_id, namespace_name) is not None + def _ns_namespace_exists_by_id(self, collection_id: str, namespace_id: str) -> bool: + """Whether a namespace with this id is still live in the catalog. + + Used to reject DML/DQL on a stale Namespace handle whose namespace (or + whole collection) was deleted. delete_namespace soft-deletes by renaming + the row to '__recyclebin__' (kernel async cleanup follows), so + a recyclebin-prefixed row counts as gone; a deleted collection removes + the row outright. Underlying data may linger after either, so we trust + the catalog, not the data table. + """ + collection_id_escaped = escape_string(str(collection_id)) + rows = self._execute( + f"SELECT namespace_id FROM {self._qtable(NamespaceCollectionNames.sdk_namespaces_table())} " + f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {int(namespace_id)} " + f"AND LEFT(namespace_name, 13) <> '__recyclebin_'" + ) + return bool(rows) + def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> None: meta = self._get_ns_namespace_meta(collection_id, namespace_name) if meta is None: diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py index f0179fc2..0c7ad6ee 100644 --- a/src/pyseekdb/client/namespace.py +++ b/src/pyseekdb/client/namespace.py @@ -49,6 +49,13 @@ def __repr__(self) -> str: f"collection='{self._collection.name}')" ) + def _guard_exists(self) -> None: + if not self._client._ns_namespace_exists_by_id(self._collection.id, self._namespace_id): + raise ValueError( + f"Namespace '{self._name}' no longer exists (it or its collection may have been deleted). " + "Operations are not allowed on a deleted namespace." + ) + # ==================== DML Operations ==================== def add( @@ -59,6 +66,7 @@ def add( documents: str | list[str] | None = None, **kwargs, ) -> None: + self._guard_exists() return self._client._namespace_add( collection_id=self._collection.id, collection_name=self._collection.name, @@ -80,6 +88,7 @@ def update( documents: str | list[str] | None = None, **kwargs, ) -> None: + self._guard_exists() return self._client._namespace_update( collection_id=self._collection.id, collection_name=self._collection.name, @@ -101,6 +110,7 @@ def upsert( documents: str | list[str] | None = None, **kwargs, ) -> None: + self._guard_exists() return self._client._namespace_upsert( collection_id=self._collection.id, collection_name=self._collection.name, @@ -121,6 +131,7 @@ def delete( where_document: dict[str, Any] | None = None, **kwargs, ) -> None: + self._guard_exists() return self._client._namespace_delete( collection_id=self._collection.id, collection_name=self._collection.name, @@ -144,6 +155,7 @@ def query( include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + self._guard_exists() return self._client._namespace_query( collection_id=self._collection.id, collection_name=self._collection.name, @@ -169,6 +181,7 @@ def hybrid_search( include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + self._guard_exists() if include is None and not query and not knn: include = [] return self._client._namespace_hybrid_search( @@ -196,6 +209,7 @@ def get( include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + self._guard_exists() return self._client._namespace_get( collection_id=self._collection.id, collection_name=self._collection.name, @@ -211,6 +225,7 @@ def get( ) def count(self) -> int: + self._guard_exists() return self._client._namespace_count( collection_id=self._collection.id, collection_name=self._collection.name, @@ -219,6 +234,7 @@ def count(self) -> int: ) def peek(self, limit: int = 10) -> dict[str, Any]: + self._guard_exists() return self._client._namespace_peek( collection_id=self._collection.id, collection_name=self._collection.name, @@ -228,6 +244,7 @@ def peek(self, limit: int = 10) -> dict[str, Any]: ) def prewarm(self) -> None: + self._guard_exists() return self._client._namespace_prewarm( collection_id=self._collection.id, collection_name=self._collection.name, diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py index b3225a71..f3adb5c4 100644 --- a/tests/integration_tests/test_namespace_dml.py +++ b/tests/integration_tests/test_namespace_dml.py @@ -43,6 +43,33 @@ def test_add_single(self, db_client): finally: cleanup(db_client, collection) + def test_ops_blocked_after_namespace_deleted(self, db_client): + collection = create_ns_collection(db_client, suffix="_delns") + ns = _create_namespace(collection, "dml_ns") + try: + ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0]) + collection.delete_namespace("dml_ns") + with pytest.raises(ValueError, match="no longer exists"): + ns.add(ids="d2", embeddings=[4.0, 5.0, 6.0]) + with pytest.raises(ValueError, match="no longer exists"): + ns.get(ids="d1") + with pytest.raises(ValueError, match="no longer exists"): + ns.query(query_embeddings=[1.0, 2.0, 3.0], n_results=1) + with pytest.raises(ValueError, match="no longer exists"): + ns.count() + finally: + cleanup(db_client, collection) + + def test_ops_blocked_after_collection_deleted(self, db_client): + collection = create_ns_collection(db_client, suffix="_delcoll") + ns = _create_namespace(collection, "dml_ns") + ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0]) + db_client.delete_collection(name=collection.name) + with pytest.raises(ValueError, match="no longer exists"): + ns.add(ids="d2", embeddings=[4.0, 5.0, 6.0]) + with pytest.raises(ValueError, match="no longer exists"): + ns.query(query_embeddings=[1.0, 2.0, 3.0], n_results=1) + def test_add_batch(self, db_client): collection = create_ns_collection(db_client, suffix="_addb") ns = _create_namespace(collection,"dml_ns") From 28b7d92c3ea4371680daf5520a7af99cc4451362 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 16 Jun 2026 10:54:47 +0800 Subject: [PATCH 44/84] =?UTF-8?q?fix:[pyseekdb][namespace]=20get=5Fcollect?= =?UTF-8?q?ion=20=E6=89=93=E5=BC=80=E5=B8=A6=20embedding=5Ffunction=20sett?= =?UTF-8?q?ings=20=E7=9A=84=20namespace=20collection=20=E6=8A=A5=20Attribu?= =?UTF-8?q?teError=20https://project.alipay.com/project/W24001004732/P2600?= =?UTF-8?q?1026476/workitemDetail=3FopenWorkItemId=3D2026061500116759681?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 7 ++++--- .../test_namespace_lifecycle.py | 21 +++++++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 57daeb42..e0cb89ba 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1605,9 +1605,10 @@ def _build_ns_collection_from_meta(self, meta: dict, embedding_function=_NOT_PRO ef = embedding_function elif "embedding_function" in settings: ef_info = settings["embedding_function"] - ef = EmbeddingFunctionRegistry.get(ef_info["name"]) - if ef and hasattr(ef, "set_config"): - ef.set_config(ef_info.get("properties", {})) + ef_class = EmbeddingFunctionRegistry.get_class(ef_info["name"]) + if ef_class is None: + raise ValueError(f"Embedding function class '{ef_info['name']}' not found") + ef = ef_class.build_from_config(ef_info.get("properties", {})) return Collection( client=self, name=meta["collection_name"], diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index b4244ee4..43beb759 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -47,6 +47,27 @@ def test_get_collection_preserves_namespace_flag(self, db_client): finally: db_client.delete_collection(name=collection.name) + def test_get_collection_restores_embedding_function(self, db_client): + from pyseekdb import DefaultEmbeddingFunction + + ef = DefaultEmbeddingFunction() + name = f"test_ns_ef_{int(time.time() * 1000)}" + use_namespace_test_partitions() + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=ef.dimension, distance="cosine", fresh_mode="spfresh"), + embedding_function=ef, + ), + ) + collection = db_client.create_collection(name=name, schema=schema, use_namespace=True) + try: + reopened = db_client.get_collection(name) + assert reopened.use_namespace is True + assert reopened.embedding_function is not None + assert reopened.embedding_function.name() == "default" + finally: + db_client.delete_collection(name=name) + def test_create_and_get_namespace(self, db_client): collection = self._create_ns_collection(db_client) try: From 210c2d15e9a66b8938a7b8e4e9391645a5286fa4 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 16 Jun 2026 14:54:47 +0800 Subject: [PATCH 45/84] =?UTF-8?q?=E8=AE=BE=E7=BD=AE=E5=88=86=E5=8C=BA?= =?UTF-8?q?=E6=95=B0=E6=94=B9=E4=B8=BAcollection=E7=B2=92=E5=BA=A6?= =?UTF-8?q?=EF=BC=8C=E6=94=B9=E6=88=90create=20collection=E7=9A=84?= =?UTF-8?q?=E9=85=8D=E7=BD=AE=E9=A1=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/__init__.py | 4 -- src/pyseekdb/client/__init__.py | 4 +- src/pyseekdb/client/client_base.py | 38 +++++----- src/pyseekdb/client/collection.py | 12 ++++ tests/integration_tests/conftest.py | 8 --- .../namespace_dml_helpers.py | 12 ++-- .../namespace_fts_helpers.py | 11 +-- .../test_namespace_constraints.py | 7 +- .../test_namespace_drop_validation.py | 8 ++- ...rid_search_fulltext_multi_coll_multi_ns.py | 5 +- ...earch_triple_branch_multi_coll_multi_ns.py | 5 +- .../test_namespace_lifecycle.py | 71 ++++++++++++------- .../test_namespace_optional_indexes.py | 6 +- .../test_namespace_prewarm.py | 17 ++--- .../test_namespace_prewarm_lob.py | 23 ++---- .../integration_tests/test_namespace_query.py | 8 ++- .../test_namespace_session_vars.py | 8 ++- tests/unit_tests/test_namespace.py | 11 +++ 18 files changed, 143 insertions(+), 115 deletions(-) diff --git a/src/pyseekdb/__init__.py b/src/pyseekdb/__init__.py index 1995ce48..1e2a33af 100644 --- a/src/pyseekdb/__init__.py +++ b/src/pyseekdb/__init__.py @@ -102,10 +102,8 @@ VectorIndexConfig, Version, get_default_embedding_function, - get_collection_partition_count, register_embedding_function, register_sparse_embedding_function, - set_collection_partition_count, ) from .client.collection import Collection from .client.namespace import Namespace @@ -151,8 +149,6 @@ "VectorIndexConfig", "Version", "get_default_embedding_function", - "get_collection_partition_count", "register_embedding_function", "register_sparse_embedding_function", - "set_collection_partition_count", ] diff --git a/src/pyseekdb/client/__init__.py b/src/pyseekdb/client/__init__.py index 7d92d918..1ce0aaef 100644 --- a/src/pyseekdb/client/__init__.py +++ b/src/pyseekdb/client/__init__.py @@ -22,7 +22,7 @@ from .admin_client import AdminAPI, _AdminClientProxy, _ClientProxy from .base_connection import BaseConnection -from .client_base import BaseClient, ClientAPI, get_collection_partition_count, set_collection_partition_count +from .client_base import BaseClient, ClientAPI from .client_seekdb_server import RemoteServerClient from .configuration import ( BengProperties, @@ -188,10 +188,8 @@ def __getattr__(name: str) -> Any: "VectorIndexConfig", "Version", "get_default_embedding_function", - "get_collection_partition_count", "register_embedding_function", "register_sparse_embedding_function", - "set_collection_partition_count", ] diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index e0cb89ba..04b910e6 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -128,23 +128,12 @@ def _validate_collection_name(name: str) -> None: from .validators import _MAX_NAMESPACE_BATCH_SIZE, _validate_namespace_name, _validate_record_ids # noqa: F401 -_NS_PARTITION_COUNT = 1000 +_DEFAULT_PARTITION_COUNT = 1000 # Unquoted id for WHERE/CASE; plain JSON_EXTRACT returns a quoted JSON string and # can route through SEARCH INDEX on SS logic tables, breaking cross-namespace id lookups. _NS_DATA_CONTENT_ID_EXPR = "JSON_UNQUOTE(JSON_EXTRACT(data_content, '$.id'))" -def set_collection_partition_count(n: int) -> None: - global _NS_PARTITION_COUNT - if n < 1: - raise ValueError("collection partition count must be >= 1") - _NS_PARTITION_COUNT = n - - -def get_collection_partition_count() -> int: - return _NS_PARTITION_COUNT - - def _build_default_ltable_schema() -> dict: return { "col_info": [ @@ -801,6 +790,7 @@ def create_collection( configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, use_namespace: bool = False, + partition_count: int | None = None, **kwargs, ) -> "Collection": """Create a new collection. @@ -817,6 +807,9 @@ def create_collection( Defaults to ``DefaultEmbeddingFunction`` (all-MiniLM-L6-v2). If set to None, no embedding function will be used (embeddings must be provided manually). use_namespace: If True, create a namespace-enabled collection. Defaults to False. + partition_count: Number of partitions for the namespace physical tables. + Only valid when ``use_namespace=True``. Defaults to 1000 when not provided. + Passing it for a non-namespace collection raises ``ValueError``. **kwargs: Additional parameters for collection creation. Returns: @@ -849,6 +842,10 @@ def create_collection( ... ) """ _validate_collection_name(name) + if partition_count is not None and not use_namespace: + raise ValueError( + "partition_count is only supported for namespace-enabled collections (use_namespace=True)." + ) if self.has_collection(name): raise ValueError(f"Collection '{name}' already exists") @@ -866,7 +863,7 @@ def create_collection( logger.debug(f"schema: {schema}") if use_namespace: - return self._create_namespace_collection(name, schema, **kwargs) + return self._create_namespace_collection(name, schema, partition_count=partition_count, **kwargs) # Resolve HNSW configuration dimension if not set hnsw_config = schema.vector_index.hnsw @@ -941,7 +938,10 @@ def create_collection( **kwargs, ) - def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> "Collection": + def _create_namespace_collection(self, name: str, schema: Schema, partition_count: int | None = None, **kwargs) -> "Collection": + pc = _DEFAULT_PARTITION_COUNT if partition_count is None else partition_count + if pc < 1: + raise ValueError("partition_count must be >= 1") dense_embedding_function = schema.vector_index.embedding_function ivf_config = schema.vector_index.ivf hnsw_config = schema.vector_index.hnsw @@ -975,6 +975,7 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " "storage_mode": "ss" if is_ss else "sn", "dimension": dimension, "distance": distance, + "partition_count": pc, } if ivf_config is not None: settings["dense_index_type"] = "ivf" @@ -998,6 +999,7 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " ivf_config=ivf_config, fulltext_config=schema.fulltext_index, is_shared_storage=is_ss, + partition_count=pc, ) except Exception: with contextlib.suppress(Exception): @@ -1016,6 +1018,7 @@ def _create_namespace_collection(self, name: str, schema: Schema, **kwargs) -> " embedding_function=dense_embedding_function, distance=distance, use_namespace=True, + partition_count=pc, ) def _get_embedding_function_dimension(self, embedding_function: EmbeddingFunction) -> int: @@ -1205,7 +1208,7 @@ def _ensure_namespace_catalogs(self) -> None: PRIMARY KEY (namespace_id, ltable_id, included_index), KEY idx_sdk_ns_stat_by_collection (collection_id) ) COMMENT='Logic table row count and storage size statistics' DEFAULT CHARSET=utf8mb4 - PARTITION BY KEY(namespace_id) PARTITIONS {_NS_PARTITION_COUNT};""" + PARTITION BY KEY(namespace_id) PARTITIONS 8;""" self._execute(ns_namespaces_sql) self._execute(ns_ltables_sql) self._execute(namespaces_stats_sql) @@ -1313,6 +1316,7 @@ def _create_namespace_physical_tables( ivf_config=None, fulltext_config=None, is_shared_storage: bool = False, + partition_count: int = _DEFAULT_PARTITION_COUNT, ) -> None: tg_name = NamespaceCollectionNames.tablegroup_name(collection_id) data_table = NamespaceCollectionNames.data_table_name(collection_id) @@ -1327,7 +1331,7 @@ def _create_namespace_physical_tables( vector_index_sql = _get_ivf_vector_index_sql(ivf_config) index_parts.append(f"VECTOR INDEX idx_vec(embedding) {vector_index_sql}") index_sql = ",\n ".join(index_parts) - partition_clause = f"PARTITION BY KEY(namespace_id) PARTITIONS {_NS_PARTITION_COUNT}" + partition_clause = f"PARTITION BY KEY(namespace_id) PARTITIONS {partition_count}" try: self._execute(f"CREATE TABLEGROUP `{tg_name}` SHARDING='ADAPTIVE'") @@ -1600,6 +1604,7 @@ def _build_ns_collection_from_meta(self, meta: dict, embedding_function=_NOT_PRO settings = meta.get("settings", {}) dimension = settings.get("dimension") distance = settings.get("distance", DEFAULT_DISTANCE_METRIC) + partition_count = settings.get("partition_count") ef = None if embedding_function is not _NOT_PROVIDED: ef = embedding_function @@ -1617,6 +1622,7 @@ def _build_ns_collection_from_meta(self, meta: dict, embedding_function=_NOT_PRO embedding_function=ef, distance=distance, use_namespace=True, + partition_count=partition_count, ) def _resolve_collection_metadata_from_sdk_collections(self, collection_name: str) -> _CollectionMeta | None: diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 7aec5ff7..253041ec 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -42,6 +42,7 @@ def __init__( distance: str | None = None, sparse_vector_index_config: Optional["SparseVectorIndexConfig"] = None, use_namespace: bool = False, + partition_count: int | None = None, **metadata, ): self._client = client @@ -52,6 +53,7 @@ def __init__( self._distance = distance self._sparse_vector_index_config = sparse_vector_index_config self._use_namespace = use_namespace + self._partition_count = partition_count self._metadata = metadata # ==================== Properties ==================== @@ -107,6 +109,16 @@ def sparse_embedding_function(self) -> Optional["SparseEmbeddingFunction"]: def use_namespace(self) -> bool: return self._use_namespace + @property + def partition_count(self) -> int | None: + """Number of partitions for the namespace physical tables. + + Returns the configured partition count for namespace-enabled collections + (defaults to 1000 when not set at creation), or None for non-namespace + collections, which are not partitioned. + """ + return self._partition_count + def _guard_collection_data_api(self) -> None: if self._use_namespace: raise ValueError( diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index ef05111c..ad87f57e 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -272,11 +272,3 @@ def oceanbase_admin_client(): client.close() -@pytest.fixture(autouse=True) -def _reduce_namespace_partitions(): - """Use a small partition count in integration tests to speed up DDL.""" - from pyseekdb import get_collection_partition_count, set_collection_partition_count - original = get_collection_partition_count() - set_collection_partition_count(8) - yield - set_collection_partition_count(original) diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index ad38c595..835644b9 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -46,16 +46,12 @@ def ns_schema() -> Schema: ) -def use_namespace_test_partitions() -> None: - from pyseekdb import set_collection_partition_count - - set_collection_partition_count(NAMESPACE_TEST_PARTITION_COUNT) - - def create_ns_collection(client: Any, suffix: str = "") -> Any: - use_namespace_test_partitions() name = f"test_ns_dml_{int(time.time() * 1000)}{suffix}" - return client.create_collection(name=name, schema=ns_schema(), use_namespace=True) + return client.create_collection( + name=name, schema=ns_schema(), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) def cleanup(client: Any, *collections: Any) -> None: diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index 33e25067..9dcf6deb 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -17,7 +17,7 @@ VectorDistanceMetric = Literal["l2", "cosine"] VECTOR_DISTANCE_METRICS: tuple[VectorDistanceMetric, ...] = ("l2", "cosine") -from namespace_dml_helpers import use_namespace_test_partitions +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema @@ -514,9 +514,9 @@ def setup_large_fts_collection( raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") name = f"test_ns_hs_ft_{distance}_{int(time.time() * 1000)}" - use_namespace_test_partitions() collection = db_client.create_collection( - name=name, schema=ns_schema(distance), use_namespace=True + name=name, schema=ns_schema(distance), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) return corpus, collection @@ -575,16 +575,17 @@ def _create_multi_coll_multi_ns_layout( distance: VectorDistanceMetric = "l2", ) -> dict[str, Any]: ts = int(time.time() * 1000) - use_namespace_test_partitions() coll_1 = db_client.create_collection( name=f"{name_prefix}_{distance}_{ts}_c1", schema=ns_schema(distance), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) coll_2 = db_client.create_collection( name=f"{name_prefix}_{distance}_{ts}_c2", schema=ns_schema(distance), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) ctx: dict[str, Any] = {"coll_1": coll_1, "coll_2": coll_2} for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): @@ -634,11 +635,11 @@ def setup_same_collection_both_ns_fts(db_client: Any) -> dict[str, Any]: remain visible via get. """ ts = int(time.time() * 1000) - use_namespace_test_partitions() collection = db_client.create_collection( name=f"test_ns_hs_ft_sc2ns_{ts}", schema=ns_schema(), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) corpus = build_large_fts_corpus(CORPUS_SIZE) ns_x = collection.create_namespace("ns_x") diff --git a/tests/integration_tests/test_namespace_constraints.py b/tests/integration_tests/test_namespace_constraints.py index 76253eb8..11a3f113 100644 --- a/tests/integration_tests/test_namespace_constraints.py +++ b/tests/integration_tests/test_namespace_constraints.py @@ -11,6 +11,7 @@ import pytest +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.client_base import NAMESPACE_MIN_OB_VERSION from pyseekdb.client.configuration import VectorIndexConfig @@ -47,7 +48,8 @@ def test_ivf_flat_accepted(self, oceanbase_client): """Positive control: ivf_flat namespace collection creation succeeds.""" name = _unique_name("_ivf_flat") collection = oceanbase_client.create_collection( - name=name, schema=_make_schema("ivf_flat"), use_namespace=True + name=name, schema=_make_schema("ivf_flat"), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) try: assert collection.use_namespace is True @@ -67,7 +69,8 @@ def test_connected_ob_meets_min_version(self, oceanbase_client): name = _unique_name("_ver_ok") collection = oceanbase_client.create_collection( - name=name, schema=_make_schema("ivf_flat"), use_namespace=True + name=name, schema=_make_schema("ivf_flat"), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) try: assert collection.use_namespace is True diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index 3ad46205..f314b6d0 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -4,7 +4,7 @@ import pytest import pyseekdb -from namespace_dml_helpers import use_namespace_test_partitions +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.meta_info import NamespaceCollectionNames @@ -33,14 +33,16 @@ def _make_collection(client, suffix: str = ""): """Create a namespace-mode collection with an IVF index (so we get a full set of physical tables: logic_data / kv_data / logic_schema / hot_table).""" name = f"test_ns_drop_{int(time.time() * 1000)}{suffix}" - use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), embedding_function=None, ), ) - return client.create_collection(name=name, schema=schema, use_namespace=True) + return client.create_collection( + name=name, schema=schema, use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) def _execute(client, sql: str): diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py index e906ccb4..26fde4c0 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py @@ -17,7 +17,7 @@ import pytest -from namespace_dml_helpers import use_namespace_test_partitions +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from namespace_fts_helpers import ( CORPUS_SIZE, MULTI_COLL_MULTI_NS_QUADRANT_KEYS, @@ -59,16 +59,17 @@ def _setup_multi_coll_multi_ns_isolation_fts(db_client): sibling namespace fails the existing corpus-based assertions. """ ts = int(time.time() * 1000) - use_namespace_test_partitions() coll_1 = db_client.create_collection( name=f"test_ns_hs_ft_mcmn_iso_{ts}_c1", schema=ns_schema(), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) coll_2 = db_client.create_collection( name=f"test_ns_hs_ft_mcmn_iso_{ts}_c2", schema=ns_schema(), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) ctx: dict[str, object] = {"coll_1": coll_1, "coll_2": coll_2} for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py index c46c1936..1be13535 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py @@ -16,7 +16,7 @@ import pytest -from namespace_dml_helpers import use_namespace_test_partitions +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from namespace_fts_helpers import ( CORPUS_SIZE, MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, @@ -81,16 +81,17 @@ def _setup_multi_coll_multi_ns_isolation_triple( ): """2 collections x 2 loaded namespaces with distinct per-namespace corpus markers.""" ts = int(time.time() * 1000) - use_namespace_test_partitions() coll_1 = db_client.create_collection( name=f"test_ns_hs_tb_mcmn_iso_{distance}_{ts}_c1", schema=ns_schema(distance), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) coll_2 = db_client.create_collection( name=f"test_ns_hs_tb_mcmn_iso_{distance}_{ts}_c2", schema=ns_schema(distance), use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) ctx: dict[str, object] = {"coll_1": coll_1, "coll_2": coll_2} for coll_tag, collection in (("c1", coll_1), ("c2", coll_2)): diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 43beb759..2b587476 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -8,7 +8,7 @@ import pytest import pyseekdb -from namespace_dml_helpers import use_namespace_test_partitions +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema @@ -19,14 +19,16 @@ class TestNamespaceLifecycle: def _create_ns_collection(self, client, suffix=""): name = f"test_ns_lc_{int(time.time() * 1000)}{suffix}" - use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), embedding_function=None, ), ) - collection = client.create_collection(name=name, schema=schema, use_namespace=True) + collection = client.create_collection( + name=name, schema=schema, use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) return collection def test_create_namespace_collection(self, db_client): @@ -52,14 +54,16 @@ def test_get_collection_restores_embedding_function(self, db_client): ef = DefaultEmbeddingFunction() name = f"test_ns_ef_{int(time.time() * 1000)}" - use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=ef.dimension, distance="cosine", fresh_mode="spfresh"), embedding_function=ef, ), ) - collection = db_client.create_collection(name=name, schema=schema, use_namespace=True) + collection = db_client.create_collection( + name=name, schema=schema, use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) try: reopened = db_client.get_collection(name) assert reopened.use_namespace is True @@ -206,7 +210,6 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): type(client._server), "_is_shared_storage_mode", return_value=True ): name = f"test_ns_ss_{int(time.time() * 1000)}" - use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), @@ -214,7 +217,8 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): ), ) collection = client.create_collection( - name=name, schema=schema, use_namespace=True + name=name, schema=schema, use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) try: @@ -240,27 +244,46 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): def test_custom_collection_partition_count(self, oceanbase_client): - """Verify set_collection_partition_count controls the PARTITIONS clause.""" - from pyseekdb import get_collection_partition_count, set_collection_partition_count + """Verify create_collection(partition_count=...) controls the PARTITIONS clause + and is exposed via collection.partition_count.""" + import re from pyseekdb.client.meta_info import NamespaceCollectionNames - original = get_collection_partition_count() + name = f"test_ns_lc_pc_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), + embedding_function=None, + ), + ) + collection = oceanbase_client.create_collection( + name=name, schema=schema, use_namespace=True, partition_count=4 + ) try: - set_collection_partition_count(4) - collection = self._create_ns_collection(oceanbase_client, suffix="_pc") - try: - data_table = NamespaceCollectionNames.data_table_name(collection.id) - rows = oceanbase_client._server._execute(f"SHOW CREATE TABLE `{data_table}`") - create_sql = rows[0].get("Create Table", "") if rows else "" - import re - partitions = re.findall(r"partition `p\d+`", create_sql) - assert len(partitions) == 4, ( - f"Expected 4 partitions, found {len(partitions)}: {create_sql[-300:]}" - ) - finally: - oceanbase_client.delete_collection(name=collection.name) + assert collection.partition_count == 4 + # Reopened handle restores partition_count from settings. + assert oceanbase_client.get_collection(name).partition_count == 4 + + data_table = NamespaceCollectionNames.data_table_name(collection.id) + rows = oceanbase_client._server._execute(f"SHOW CREATE TABLE `{data_table}`") + create_sql = rows[0].get("Create Table", "") if rows else "" + partitions = re.findall(r"partition `p\d+`", create_sql) + assert len(partitions) == 4, ( + f"Expected 4 partitions, found {len(partitions)}: {create_sql[-300:]}" + ) finally: - set_collection_partition_count(original) + oceanbase_client.delete_collection(name=collection.name) + + def test_partition_count_rejected_for_non_namespace(self, db_client): + name = f"test_nons_pc_{int(time.time() * 1000)}" + with pytest.raises(ValueError, match="partition_count is only supported"): + db_client.create_collection( + name=name, + configuration=pyseekdb.HNSWConfiguration(dimension=3), + embedding_function=None, + partition_count=4, + ) + assert not db_client.has_collection(name) if __name__ == "__main__": diff --git a/tests/integration_tests/test_namespace_optional_indexes.py b/tests/integration_tests/test_namespace_optional_indexes.py index 24e9a712..dd01980b 100644 --- a/tests/integration_tests/test_namespace_optional_indexes.py +++ b/tests/integration_tests/test_namespace_optional_indexes.py @@ -21,6 +21,7 @@ import pymysql.err import pytest +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.meta_info import NamespaceCollectionNames @@ -63,7 +64,10 @@ def _collection_settings(client: Any, collection_id: str) -> dict: def _create_collection(client: Any, label: str, schema: Schema) -> Any: name = f"test_ns_opt_idx_{label}_{int(time.time() * 1000)}" - return client.create_collection(name=name, schema=schema, use_namespace=True) + return client.create_collection( + name=name, schema=schema, use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) def _unit_vector(dimension: int, axis: int = 0) -> list[float]: diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index 9c0e9742..f5458c14 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -17,27 +17,20 @@ class TestNamespacePrewarm: - @pytest.fixture(autouse=True) - def _reduce_namespace_partitions(self): - """Override conftest: use 8 partitions for prewarm until DDL is faster.""" - from pyseekdb import get_collection_partition_count, set_collection_partition_count - original = get_collection_partition_count() - set_collection_partition_count(8) - yield - set_collection_partition_count(original) - def _create_ns_collection_and_namespace(self, client): - from namespace_dml_helpers import use_namespace_test_partitions + from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT name = f"test_ns_pw_{int(time.time() * 1000)}" - use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), embedding_function=None, ), ) - collection = client.create_collection(name=name, schema=schema, use_namespace=True) + collection = client.create_collection( + name=name, schema=schema, use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) namespace = collection.create_namespace("pw_ns") return collection, namespace diff --git a/tests/integration_tests/test_namespace_prewarm_lob.py b/tests/integration_tests/test_namespace_prewarm_lob.py index 48eaf424..a2f758cc 100644 --- a/tests/integration_tests/test_namespace_prewarm_lob.py +++ b/tests/integration_tests/test_namespace_prewarm_lob.py @@ -39,11 +39,7 @@ except ImportError: # pragma: no cover - pymysql ships with the test deps pymysql = None -from pyseekdb import ( - IVFConfiguration, - get_collection_partition_count, - set_collection_partition_count, -) +from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.meta_info import NamespaceCollectionNames from pyseekdb.client.schema import Schema @@ -194,14 +190,15 @@ def _grep_lob_prewarm_log(lob_meta_tablet_ids): class _BaseLobPrewarm: def _make_collection(self, client, name, partitions): - set_collection_partition_count(partitions) schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine"), embedding_function=None, ), ) - return client.create_collection(name=name, schema=schema, use_namespace=True) + return client.create_collection( + name=name, schema=schema, use_namespace=True, partition_count=partitions + ) def _force_out_of_row_lob(self, client, collection_id, namespace_id): """Insert several >0.75MB incompressible kv_values so they spill out-of-row @@ -232,12 +229,6 @@ def _force_out_of_row_lob(self, client, collection_id, namespace_id): class TestLobPrewarmStructural(_BaseLobPrewarm): - @pytest.fixture(autouse=True) - def _restore_partitions(self): - original = get_collection_partition_count() - yield - set_collection_partition_count(original) - def test_kv_table_has_lob_meta_tablet(self, oceanbase_client): """Tier 1: the prewarm target (kv_data_table) owns an aux LOB-meta tablet.""" name = f"lob_pw_struct_{int(time.time() * 1000)}" @@ -293,12 +284,6 @@ def test_multi_namespace_prewarm_is_independent(self, oceanbase_client): class TestLobPrewarmCaching(_BaseLobPrewarm): - @pytest.fixture(autouse=True) - def _restore_partitions(self): - original = get_collection_partition_count() - yield - set_collection_partition_count(original) - @staticmethod def _total_bytes(cache): return sum(v[1] for v in (cache or {}).values()) diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index 3d978948..cbd82228 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -7,7 +7,7 @@ import pytest -from namespace_dml_helpers import use_namespace_test_partitions +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema @@ -17,7 +17,6 @@ class TestNamespaceQuery: def _setup(self, client, suffix=""): name = f"test_ns_q_{int(time.time() * 1000)}{suffix}" - use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), @@ -25,7 +24,10 @@ def _setup(self, client, suffix=""): ), fulltext_index=FulltextIndexConfig(analyzer="ik"), ) - collection = client.create_collection(name=name, schema=schema, use_namespace=True) + collection = client.create_collection( + name=name, schema=schema, use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) return collection def _insert_data(self, ns): diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py index ddf7d379..178339f7 100644 --- a/tests/integration_tests/test_namespace_session_vars.py +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -9,7 +9,7 @@ import pytest -from namespace_dml_helpers import use_namespace_test_partitions +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema @@ -19,14 +19,16 @@ class TestNamespaceSessionVars: def _create_ns_collection(self, client): name = f"test_ns_sessvar_{int(time.time() * 1000)}" - use_namespace_test_partitions() schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), embedding_function=None, ), ) - return client.create_collection(name=name, schema=schema, use_namespace=True) + return client.create_collection( + name=name, schema=schema, use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) def _query_session_vars(self, client): rows = client._server._execute( diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index d36fe02e..348edaeb 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -862,6 +862,17 @@ def test_collection_has_namespace_validates_name(self): with pytest.raises(TypeError, match="must be a string"): coll.has_namespace(42) + def test_partition_count_property(self): + mock_client = MagicMock() + ns_coll = Collection( + client=mock_client, name="c", collection_id="1", dimension=3, + use_namespace=True, partition_count=4, + ) + assert ns_coll.partition_count == 4 + # Non-namespace collections are not partitioned. + plain = Collection(client=mock_client, name="p", collection_id="2", dimension=3) + assert plain.partition_count is None + # ==================== Record ID Validation Tests ==================== From 553c86199b58ba82c89616e3099d2b5b49dc6013 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 16 Jun 2026 16:36:55 +0800 Subject: [PATCH 46/84] =?UTF-8?q?fix:n=5Fresults=E5=80=BC=E6=9C=89?= =?UTF-8?q?=E8=AF=AF=E5=9C=A8sdk=E7=9B=B4=E6=8E=A5=E6=8A=A5=E9=94=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 9 ++++++++- src/pyseekdb/client/collection.py | 13 ++++++++++++- src/pyseekdb/client/namespace.py | 14 ++++++++++++++ src/pyseekdb/client/validators.py | 1 + 4 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 04b910e6..be96b9b1 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -126,7 +126,7 @@ def _validate_collection_name(name: str) -> None: ) -from .validators import _MAX_NAMESPACE_BATCH_SIZE, _validate_namespace_name, _validate_record_ids # noqa: F401 +from .validators import _MAX_N_RESULTS, _MAX_NAMESPACE_BATCH_SIZE, _validate_namespace_name, _validate_record_ids # noqa: F401 _DEFAULT_PARTITION_COUNT = 1000 # Unquoted id for WHERE/CASE; plain JSON_EXTRACT returns a quoted JSON string and @@ -4450,6 +4450,13 @@ def _build_knn_expression( # noqa: C901 query_embeddings = knn.get("query_embeddings") where = knn.get("where") n_results = knn.get("n_results", 10) + if not isinstance(n_results, int) or n_results < 1: + raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") + if n_results > _MAX_N_RESULTS: + raise ValueError( + f"n_results must be <= {_MAX_N_RESULTS}, got {n_results}. " + "Use a smaller value or paginate with offset/limit." + ) boost = knn.get("boost") embedding_function = kwargs.get("embedding_function") diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 253041ec..3b7f4e0a 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Optional -from .validators import _validate_namespace_name +from .validators import _MAX_N_RESULTS, _validate_namespace_name if TYPE_CHECKING: from .embedding_function import Documents as EmbeddingDocuments @@ -119,6 +119,16 @@ def partition_count(self) -> int | None: """ return self._partition_count + @staticmethod + def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> None: + if not isinstance(n_results, int) or n_results < 1: + raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") + if n_results > max_results: + raise ValueError( + f"n_results must be <= {max_results}, got {n_results}. " + "Use a smaller value or paginate with offset/limit." + ) + def _guard_collection_data_api(self) -> None: if self._use_namespace: raise ValueError( @@ -486,6 +496,7 @@ def query( ) """ self._guard_collection_data_api() + self._validate_n_results(n_results) return self._client._collection_query( collection_id=self._id, collection_name=self._name, diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py index 0c7ad6ee..c95be242 100644 --- a/src/pyseekdb/client/namespace.py +++ b/src/pyseekdb/client/namespace.py @@ -9,6 +9,8 @@ from typing import TYPE_CHECKING, Any +from .validators import _MAX_N_RESULTS + if TYPE_CHECKING: from .collection import Collection from .embedding_function import EmbeddingFunction @@ -56,6 +58,16 @@ def _guard_exists(self) -> None: "Operations are not allowed on a deleted namespace." ) + @staticmethod + def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> None: + if not isinstance(n_results, int) or n_results < 1: + raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") + if n_results > max_results: + raise ValueError( + f"n_results must be <= {max_results}, got {n_results}. " + "Use a smaller value or paginate with offset/limit." + ) + # ==================== DML Operations ==================== def add( @@ -156,6 +168,7 @@ def query( **kwargs, ) -> dict[str, Any]: self._guard_exists() + self._validate_n_results(n_results) return self._client._namespace_query( collection_id=self._collection.id, collection_name=self._collection.name, @@ -182,6 +195,7 @@ def hybrid_search( **kwargs, ) -> dict[str, Any]: self._guard_exists() + self._validate_n_results(n_results) if include is None and not query and not knn: include = [] return self._client._namespace_hybrid_search( diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py index 71fc25fa..7a6eb253 100644 --- a/src/pyseekdb/client/validators.py +++ b/src/pyseekdb/client/validators.py @@ -4,6 +4,7 @@ _MAX_NAME_LENGTH = 512 _MAX_NAMESPACE_NAME_LENGTH = 256 _MAX_NAMESPACE_BATCH_SIZE = 100 +_MAX_N_RESULTS = 16384 # OceanBase vector-search k upper bound def _validate_namespace_name(name: str) -> None: From 1454126e0cc4add8e9c4b921c40f18fe391f0800 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 16 Jun 2026 19:46:43 +0800 Subject: [PATCH 47/84] =?UTF-8?q?update=E4=B8=8Ecollection=E5=9C=BA?= =?UTF-8?q?=E6=99=AF=E4=BF=9D=E6=8C=81=E4=B8=80=E4=B8=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 57 ++++++++++++++++++++++++------ 1 file changed, 47 insertions(+), 10 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index be96b9b1..5d3233fd 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -415,16 +415,31 @@ def _validate_ob_database_type(self) -> None: ) def _is_shared_storage_mode(self) -> bool: + cached = getattr(self, "_shared_storage", None) + if cached is not None: + return cached + result = False try: rows = self._execute( "SELECT VALUE FROM oceanbase.GV$OB_PARAMETERS WHERE name = 'ob_startup_mode'" ) if rows: val = rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["VALUE"] - return str(val).upper() == "SHARED_STORAGE" + result = str(val).upper() == "SHARED_STORAGE" except Exception: - pass - return False + result = False + self._shared_storage = result + return result + + def _stg_cache_policy_clause(self) -> str: + """SS mode: make catalog tables global-hot so metadata is locally cached. + + STORAGE_CACHE_POLICY is only supported in shared-storage mode; SN mode + returns an empty string. + """ + if self._is_shared_storage_mode(): + return 'STORAGE_CACHE_POLICY = (GLOBAL = "hot")' + return "" def detect_db_type_and_version(self) -> tuple[str, "Version"]: # noqa: C901 """ @@ -1048,6 +1063,7 @@ def _create_sdk_collections_if_not_exists(self) -> None: try: self._use_catalog_database() sdk_coll = self._qtable(CollectionNames.sdk_collections_table_name()) + scp = self._stg_cache_policy_clause() create_table_sql = f"""CREATE TABLE IF NOT EXISTS {sdk_coll} ( collection_id CHAR(32) PRIMARY KEY DEFAULT (replace(uuid(), '-', '')), collection_name STRING, @@ -1055,7 +1071,7 @@ def _create_sdk_collections_if_not_exists(self) -> None: created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, INDEX idx_name(collection_name) - ) COMMENT='Settings of collections created by SDK';""" + ) COMMENT='Settings of collections created by SDK' {scp};""" self._execute(create_table_sql) except Exception as e: raise ValueError(f"Failed to create sdk_collections table: {e}") from e @@ -1172,6 +1188,7 @@ def _ensure_namespace_catalogs(self) -> None: self._isolate_sdk_catalog_to_client_database() ns_namespaces_q = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) ns_ltables_q = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) + scp = self._stg_cache_policy_clause() ns_namespaces_sql = f"""CREATE TABLE IF NOT EXISTS {ns_namespaces_q} ( namespace_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, collection_id CHAR(32) NOT NULL, @@ -1182,7 +1199,7 @@ def _ensure_namespace_catalogs(self) -> None: PRIMARY KEY (namespace_id), UNIQUE KEY uk_sdk_ns_coll_name (collection_id, namespace_name), KEY idx_sdk_ns_by_collection (collection_id) - ) COMMENT='Namespace catalog';""" + ) COMMENT='Namespace catalog' {scp};""" ns_ltables_sql = f"""CREATE TABLE IF NOT EXISTS {ns_ltables_q} ( ltable_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, collection_id CHAR(32) NOT NULL, @@ -1194,7 +1211,7 @@ def _ensure_namespace_catalogs(self) -> None: PRIMARY KEY (ltable_id), UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name), KEY idx_sdk_lt_by_ns (collection_id, namespace_id) - ) COMMENT='LTable catalog';""" + ) COMMENT='LTable catalog' {scp};""" namespaces_stats_sql = f"""CREATE TABLE IF NOT EXISTS {self._qtable(NamespaceCollectionNames.sdk_namespaces_stats_table())} ( collection_id CHAR(32) NOT NULL COMMENT 'collection id', namespace_id BIGINT UNSIGNED NOT NULL COMMENT 'namespace internal id', @@ -4822,10 +4839,18 @@ def _namespace_add( # noqa: C901 embeddings = embedding_function(documents) else: raise ValueError( - "Documents provided but no embeddings and no embedding function." + "Documents provided but no embeddings and no embedding function. " + "Either:\n" + " 1. Provide embeddings directly when calling add(), or\n" + " 2. Provide embedding_function to auto-generate embeddings from documents." ) else: - raise ValueError("Neither embeddings nor documents provided.") + raise ValueError( + "Neither embeddings nor documents provided. " + "Please provide either:\n" + " 1. embeddings directly, or\n" + " 2. documents with embedding_function to generate embeddings." + ) num_items = len(ids) if documents and len(documents) != num_items: @@ -4898,8 +4923,20 @@ def _namespace_update( ): embeddings = [embeddings] - if embeddings is None and documents is not None and embedding_function is not None: - embeddings = embedding_function(documents) + if embeddings: + # embeddings provided, use them directly without embedding + pass + elif documents: + # embeddings not provided but documents are provided, check for embedding_function + if embedding_function is not None: + embeddings = embedding_function(documents) + else: + raise ValueError( + "Documents provided but no embeddings and no embedding function. " + "Either:\n" + " 1. Provide embeddings directly when calling update(), or\n" + " 2. Provide embedding_function to auto-generate embeddings from documents." + ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) From d3c3879a147e4d732e49c706143c9569cf571e50 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 16 Jun 2026 20:50:24 +0800 Subject: [PATCH 48/84] =?UTF-8?q?=E6=B5=8B=E8=AF=95=E7=94=A8=E4=BE=8Bembed?= =?UTF-8?q?ding=E7=94=A8=E4=BE=8B=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/integration_tests/test_namespace_dml_multi_coll.py | 2 +- tests/integration_tests/test_namespace_dml_multi_ns.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration_tests/test_namespace_dml_multi_coll.py b/tests/integration_tests/test_namespace_dml_multi_coll.py index 061d951c..212fa2c4 100644 --- a/tests/integration_tests/test_namespace_dml_multi_coll.py +++ b/tests/integration_tests/test_namespace_dml_multi_coll.py @@ -56,7 +56,7 @@ def test_update_delete_isolated_by_collection(self, db_client): ns_a.add(ids="doc1", embeddings=[1.0, 0.0, 0.0], documents="Original A") ns_b.add(ids="doc1", embeddings=[0.0, 1.0, 0.0], documents="Original B") - ns_a.update(ids="doc1", documents="Updated A", metadatas={"updated": True}) + ns_a.update(ids="doc1", embeddings=[0.5, 0.5, 0.5], documents="Updated A", metadatas={"updated": True}) assert_get_present(ns_a, "doc1", documents="Updated A", metadatas={"updated": True}) assert_get_present(ns_b, "doc1", documents="Original B") diff --git a/tests/integration_tests/test_namespace_dml_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_ns.py index f77094da..8e854857 100644 --- a/tests/integration_tests/test_namespace_dml_multi_ns.py +++ b/tests/integration_tests/test_namespace_dml_multi_ns.py @@ -60,7 +60,7 @@ def test_update_does_not_affect_other_ns(self, db_client): assert_get_present(ns_a, doc_id, documents="A original") assert_get_present(ns_b, doc_id, documents="B original") - ns_a.update(ids=doc_id, documents="A updated", metadatas={"updated": True}) + ns_a.update(ids=doc_id, embeddings=[0.5, 0.5, 0.5], documents="A updated", metadatas={"updated": True}) assert_get_present(ns_a, doc_id, documents="A updated", metadatas={"updated": True}) assert_get_present(ns_b, doc_id, documents="B original") From c1cdef5a7e7f0171369db3a264be4feaa39602d7 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 17 Jun 2026 16:07:19 +0800 Subject: [PATCH 49/84] =?UTF-8?q?fix:[pyseekdb]=E5=88=9B=E5=BB=BAcollectio?= =?UTF-8?q?n=E8=BF=87=E7=A8=8B=E4=B8=AD=EF=BC=8Cob=E9=87=8D=E5=90=AF?= =?UTF-8?q?=E6=83=85=E5=86=B5=E4=B8=8B=EF=BC=8C=E5=86=8D=E6=AC=A1=E5=88=9B?= =?UTF-8?q?=E5=BB=BAcollection=E7=AC=AC=E4=B8=80=E6=AC=A1=E5=88=9B?= =?UTF-8?q?=E5=BB=BA=E4=BC=9A=E5=A4=B1=E8=B4=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 162 +++++++++++++----- .../test_namespace_lifecycle.py | 45 +++++ 2 files changed, 161 insertions(+), 46 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 5d3233fd..5316623e 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -862,7 +862,12 @@ def create_collection( "partition_count is only supported for namespace-enabled collections (use_namespace=True)." ) if self.has_collection(name): - raise ValueError(f"Collection '{name}' already exists") + # A namespace collection whose catalog row exists but whose physical + # tables are incomplete (e.g. a crash interrupted creation) is allowed + # through so the create can be resumed/finished idempotently. A complete + # collection still errors as "already exists". + if not (use_namespace and self._is_incomplete_ns_collection(name)): + raise ValueError(f"Collection '{name}' already exists") # Resolve schema: either use the provided schema or build one from legacy params if schema is not None: @@ -954,9 +959,6 @@ def create_collection( ) def _create_namespace_collection(self, name: str, schema: Schema, partition_count: int | None = None, **kwargs) -> "Collection": - pc = _DEFAULT_PARTITION_COUNT if partition_count is None else partition_count - if pc < 1: - raise ValueError("partition_count must be >= 1") dense_embedding_function = schema.vector_index.embedding_function ivf_config = schema.vector_index.ivf hnsw_config = schema.vector_index.hnsw @@ -970,40 +972,60 @@ def _create_namespace_collection(self, name: str, schema: Schema, partition_coun ) self._validate_ob_database_type() - if ivf_config is not None: - dimension = ivf_config.dimension - distance = ivf_config.distance + # Resume an incomplete collection (meta row exists but physical tables are + # missing, e.g. a crash interrupted creation): reuse its id/settings and + # idempotently (re)build only the missing physical tables. Otherwise create + # from scratch. Both paths are safe to call repeatedly. + existing = self._get_ns_collection_meta(name) + if existing is not None: + collection_id = existing["collection_id"] + settings = existing.get("settings", {}) + dimension = settings.get("dimension") + distance = settings.get("distance", DEFAULT_DISTANCE_METRIC) + pc = int(settings.get("partition_count", _DEFAULT_PARTITION_COUNT)) + is_ss = settings.get("storage_mode") == "ss" + logger.info( + f"Namespace collection '{name}' already has a catalog entry; resuming creation " + f"(reusing collection_id={collection_id}, rebuilding any missing physical tables)." + ) else: - if dense_embedding_function is not None: - dimension = self._get_embedding_function_dimension(dense_embedding_function) + pc = _DEFAULT_PARTITION_COUNT if partition_count is None else partition_count + if pc < 1: + raise ValueError("partition_count must be >= 1") + if ivf_config is not None: + dimension = ivf_config.dimension + distance = ivf_config.distance else: - dimension = DEFAULT_VECTOR_DIMENSION - distance = DEFAULT_DISTANCE_METRIC - - if dimension < 1 or dimension > 4096: - raise ValueError(f"Dimension must be between 1 and 4096, got {dimension}") - - is_ss = self._is_shared_storage_mode() - settings = { - "version": 2, - "use_namespace": True, - "storage_mode": "ss" if is_ss else "sn", - "dimension": dimension, - "distance": distance, - "partition_count": pc, - } - if ivf_config is not None: - settings["dense_index_type"] = "ivf" - if ivf_config.fresh_mode is not None: - settings["fresh_mode"] = ivf_config.fresh_mode - if dense_embedding_function is not None and EmbeddingFunction.support_persistence(dense_embedding_function): - settings["embedding_function"] = { - "name": dense_embedding_function.name(), - "properties": dense_embedding_function.get_config(), + if dense_embedding_function is not None: + dimension = self._get_embedding_function_dimension(dense_embedding_function) + else: + dimension = DEFAULT_VECTOR_DIMENSION + distance = DEFAULT_DISTANCE_METRIC + + if dimension < 1 or dimension > 4096: + raise ValueError(f"Dimension must be between 1 and 4096, got {dimension}") + + is_ss = self._is_shared_storage_mode() + settings = { + "version": 2, + "use_namespace": True, + "storage_mode": "ss" if is_ss else "sn", + "dimension": dimension, + "distance": distance, + "partition_count": pc, } + if ivf_config is not None: + settings["dense_index_type"] = "ivf" + if ivf_config.fresh_mode is not None: + settings["fresh_mode"] = ivf_config.fresh_mode + if dense_embedding_function is not None and EmbeddingFunction.support_persistence(dense_embedding_function): + settings["embedding_function"] = { + "name": dense_embedding_function.name(), + "properties": dense_embedding_function.get_config(), + } - collection_meta = self._create_ns_collection_meta(name, settings) - collection_id = collection_meta["collection_id"] + collection_meta = self._create_ns_collection_meta(name, settings) + collection_id = collection_meta["collection_id"] self._ensure_namespace_catalogs() @@ -1015,14 +1037,19 @@ def _create_namespace_collection(self, name: str, schema: Schema, partition_coun fulltext_config=schema.fulltext_index, is_shared_storage=is_ss, partition_count=pc, + # On a fresh create, roll back partial tables so a failure leaves no + # trace. On resume, keep whatever already exists so a later retry can + # finish the job. + cleanup_on_error=existing is None, ) except Exception: - with contextlib.suppress(Exception): - collection_id_escaped = escape_string(collection_id) - self._execute( - f"DELETE FROM `{CollectionNames.sdk_collections_table_name()}` " - f"WHERE collection_id = '{collection_id_escaped}'" - ) + if existing is None: + with contextlib.suppress(Exception): + collection_id_escaped = escape_string(collection_id) + self._execute( + f"DELETE FROM `{CollectionNames.sdk_collections_table_name()}` " + f"WHERE collection_id = '{collection_id_escaped}'" + ) raise return Collection( @@ -1334,6 +1361,7 @@ def _create_namespace_physical_tables( fulltext_config=None, is_shared_storage: bool = False, partition_count: int = _DEFAULT_PARTITION_COUNT, + cleanup_on_error: bool = True, ) -> None: tg_name = NamespaceCollectionNames.tablegroup_name(collection_id) data_table = NamespaceCollectionNames.data_table_name(collection_id) @@ -1350,10 +1378,13 @@ def _create_namespace_physical_tables( index_sql = ",\n ".join(index_parts) partition_clause = f"PARTITION BY KEY(namespace_id) PARTITIONS {partition_count}" + # All CREATEs use IF NOT EXISTS so this is idempotent: a fresh create builds + # everything, while a resume (after a crash left some tables behind) skips the + # existing ones and only fills in the gaps. try: - self._execute(f"CREATE TABLEGROUP `{tg_name}` SHARDING='ADAPTIVE'") + self._execute(f"CREATE TABLEGROUP IF NOT EXISTS `{tg_name}` SHARDING='ADAPTIVE'") - data_sql = f"""CREATE TABLE `{data_table}` ( + data_sql = f"""CREATE TABLE IF NOT EXISTS `{data_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, ltable_id BIGINT UNSIGNED NOT NULL, document LONGTEXT, @@ -1369,7 +1400,7 @@ def _create_namespace_physical_tables( if is_shared_storage: hot_table = NamespaceCollectionNames.hot_table_name(collection_id) - self._execute(f"""CREATE TABLE `{hot_table}` ( + self._execute(f"""CREATE TABLE IF NOT EXISTS `{hot_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, last_access_time TIMESTAMP(6) NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, @@ -1378,7 +1409,7 @@ def _create_namespace_physical_tables( ) TABLEGROUP=`{tg_name}` COMMENT='热点/TTL附属表' DEFAULT CHARSET=utf8mb4 {partition_clause}""") - self._execute(f"""CREATE TABLE `{kv_table}` ( + self._execute(f"""CREATE TABLE IF NOT EXISTS `{kv_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, kv_key VARBINARY(1024) NOT NULL, kv_value LONGBLOB NOT NULL, @@ -1386,7 +1417,7 @@ def _create_namespace_physical_tables( ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 LOB_INROW_THRESHOLD=786432 {partition_clause}""") - self._execute(f"""CREATE TABLE `{schema_table}` ( + self._execute(f"""CREATE TABLE IF NOT EXISTS `{schema_table}` ( namespace_id BIGINT UNSIGNED NOT NULL, ltable_id BIGINT UNSIGNED NOT NULL, schema_content JSON NOT NULL, @@ -1398,7 +1429,8 @@ def _create_namespace_physical_tables( {partition_clause}""") except Exception: - self._cleanup_namespace_physical_tables(collection_id) + if cleanup_on_error: + self._cleanup_namespace_physical_tables(collection_id) raise def _cleanup_namespace_physical_tables(self, collection_id: str) -> None: @@ -1413,6 +1445,44 @@ def _cleanup_namespace_physical_tables(self, collection_id: str) -> None: with contextlib.suppress(Exception): self._execute(f"DROP TABLEGROUP IF EXISTS `{NamespaceCollectionNames.tablegroup_name(collection_id)}`") + def _table_exists(self, table_name: str) -> bool: + """Whether `table_name` exists in the current (catalog) database.""" + name_escaped = escape_string(table_name) + try: + rows = self._execute( + "SELECT 1 FROM information_schema.TABLES " + f"WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = '{name_escaped}'" + ) + return bool(rows) + except Exception: + return False + + def _ns_missing_physical_tables(self, collection_id: str, is_shared_storage: bool) -> list[str]: + """Return the namespace physical tables that are expected but absent. + + An empty list means the collection's physical tables are complete. + """ + expected = [ + NamespaceCollectionNames.data_table_name(collection_id), + NamespaceCollectionNames.kv_data_table_name(collection_id), + NamespaceCollectionNames.logic_schema_table_name(collection_id), + ] + if is_shared_storage: + expected.append(NamespaceCollectionNames.hot_table_name(collection_id)) + return [t for t in expected if not self._table_exists(t)] + + def _is_incomplete_ns_collection(self, name: str) -> bool: + """Whether `name` is a namespace collection whose catalog row exists but + whose physical tables are not all present (e.g. creation was interrupted + by a crash). Such a collection can be finished by re-running create. + """ + meta = self._get_ns_collection_meta(name) + if meta is None: + return False + is_ss = meta.get("settings", {}).get("storage_mode") == "ss" + self._use_catalog_database() + return len(self._ns_missing_physical_tables(meta["collection_id"], is_ss)) > 0 + def _resolve_namespace_ltable_id( self, collection_id: str, namespace_id: int ) -> int: diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 2b587476..b7b23665 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -145,6 +145,51 @@ def test_delete_collection_cleans_namespaces(self, db_client): assert db_client.has_collection(coll_name) is False + def test_create_resumes_incomplete_collection(self, oceanbase_client): + """A namespace collection whose physical tables were partially lost (e.g. a + crash mid-create) is finished by re-running create_collection, instead of + being blocked by 'already exists'. A complete collection still errors.""" + from pyseekdb.client.meta_info import NamespaceCollectionNames + + client = oceanbase_client + srv = client._server + name = f"test_ns_resume_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), + embedding_function=None, + ), + ) + collection = client.create_collection( + name=name, schema=schema, use_namespace=True, partition_count=4 + ) + try: + cid = collection.id + # Simulate an interrupted create: drop one physical table. + srv._use_catalog_database() + srv._execute(f"DROP TABLE IF EXISTS `{NamespaceCollectionNames.kv_data_table_name(cid)}`") + assert srv._is_incomplete_ns_collection(name) is True + + # Re-create resumes (same id, rebuilds the missing table) instead of raising. + resumed = client.create_collection( + name=name, schema=schema, use_namespace=True, partition_count=4 + ) + assert resumed.id == cid + assert srv._is_incomplete_ns_collection(name) is False + + # Data path works after resume. + ns = resumed.create_namespace("ns_x") + ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0]) + assert ns.count() == 1 + + # A complete collection still rejects a duplicate create. + with pytest.raises(ValueError, match="already exists"): + client.create_collection( + name=name, schema=schema, use_namespace=True, partition_count=4 + ) + finally: + client.delete_collection(name=name) + def test_namespace_ops_blocked_after_collection_deleted(self, db_client): collection = self._create_ns_collection(db_client) db_client.delete_collection(name=collection.name) From fb708d2bf8004c41bb862e544afb5327bdd840c6 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 17 Jun 2026 16:17:14 +0800 Subject: [PATCH 50/84] =?UTF-8?q?fix:=E5=8E=BB=E6=8E=89drop=5Fnamespace?= =?UTF-8?q?=E7=9A=84mock=20https://project.alipay.com/project/W24001004732?= =?UTF-8?q?/P26001026476/workitemDetail=3FopenWorkItemId=3D202606170011681?= =?UTF-8?q?0861?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 39 ------------------------------ 1 file changed, 39 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 5316623e..d05a195c 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1172,47 +1172,8 @@ def _use_catalog_database(self) -> None: """Align session with pymysql database= so PL (DROP_NAMESPACE) uses the same DB.""" self._execute(f"USE `{self._catalog_database()}`") - def _isolate_sdk_catalog_to_client_database(self) -> None: - """Drop sdk_* catalog tables in other databases so kernel bg task scans self.database. - - ObLTableBGTaskScheduler picks the first tenant database that has sdk_collections. - mysqltest may leave a copy in the `oceanbase` database while pyseekdb uses `test`. - """ - target = self._catalog_database() - try: - rows = self._execute( - "SELECT DISTINCT table_schema AS table_schema " - "FROM information_schema.tables " - "WHERE table_name = 'sdk_collections'" - ) - except Exception: - return - other_schemas: list[str] = [] - for row in rows or []: - schema = row["table_schema"] if isinstance(row, dict) else row[0] - if schema and schema != target: - other_schemas.append(schema) - if not other_schemas: - return - catalog_tables = [ - CollectionNames.sdk_collections_table_name(), - NamespaceCollectionNames.sdk_namespaces_table(), - NamespaceCollectionNames.sdk_ltables_table(), - NamespaceCollectionNames.sdk_namespaces_stats_table(), - ] - for schema in other_schemas: - for table in catalog_tables: - with contextlib.suppress(Exception): - self._execute(f"DROP TABLE IF EXISTS `{schema}`.`{table}`") - logger.info( - "Removed duplicate sdk catalog tables from %s (catalog DB is %s)", - other_schemas, - target, - ) - def _ensure_namespace_catalogs(self) -> None: self._use_catalog_database() - self._isolate_sdk_catalog_to_client_database() ns_namespaces_q = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) ns_ltables_q = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) scp = self._stg_cache_policy_clause() From c9b44b6dcf353ec685fa2ec23adc6aef5558453c Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 18 Jun 2026 13:47:44 +0800 Subject: [PATCH 51/84] =?UTF-8?q?=E6=9B=B4=E6=96=B0centroids=5Ffresh=5Fmod?= =?UTF-8?q?e=E5=90=8D=E9=80=82=E9=85=8D=E5=86=85=E6=A0=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 43 +++++++++++++++++++++++++++++- tests/unit_tests/test_namespace.py | 4 +-- 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index d05a195c..19c9d223 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -228,7 +228,7 @@ def _get_ivf_vector_index_sql(ivf_config: "IVFConfiguration") -> str: else: property_parts.append(f"{k}={v}") if ivf_config.fresh_mode: - property_parts.append(f"fresh_mode={ivf_config.fresh_mode}") + property_parts.append(f"centroids_fresh_mode={ivf_config.fresh_mode}") property_str = ", ".join(property_parts) properties_str = f", {property_str}" if property_str else "" return f"WITH (DISTANCE={ivf_config.distance}, TYPE={ivf_config.type.upper()}, LIB={ivf_config.lib.upper()}{properties_str})" @@ -1172,8 +1172,49 @@ def _use_catalog_database(self) -> None: """Align session with pymysql database= so PL (DROP_NAMESPACE) uses the same DB.""" self._execute(f"USE `{self._catalog_database()}`") + # NOTE: temporarily restored (reverting commit fb708d2) to reproduce the + # multi-database primary-key conflict. Remove again after repro. + def _isolate_sdk_catalog_to_client_database(self) -> None: + """Drop sdk_* catalog tables in other databases so kernel bg task scans self.database. + + ObLTableBGTaskScheduler picks the first tenant database that has sdk_collections. + mysqltest may leave a copy in the `oceanbase` database while pyseekdb uses `test`. + """ + target = self._catalog_database() + try: + rows = self._execute( + "SELECT DISTINCT table_schema AS table_schema " + "FROM information_schema.tables " + "WHERE table_name = 'sdk_collections'" + ) + except Exception: + return + other_schemas: list[str] = [] + for row in rows or []: + schema = row["table_schema"] if isinstance(row, dict) else row[0] + if schema and schema != target: + other_schemas.append(schema) + if not other_schemas: + return + catalog_tables = [ + CollectionNames.sdk_collections_table_name(), + NamespaceCollectionNames.sdk_namespaces_table(), + NamespaceCollectionNames.sdk_ltables_table(), + NamespaceCollectionNames.sdk_namespaces_stats_table(), + ] + for schema in other_schemas: + for table in catalog_tables: + with contextlib.suppress(Exception): + self._execute(f"DROP TABLE IF EXISTS `{schema}`.`{table}`") + logger.info( + "Removed duplicate sdk catalog tables from %s (catalog DB is %s)", + other_schemas, + target, + ) + def _ensure_namespace_catalogs(self) -> None: self._use_catalog_database() + self._isolate_sdk_catalog_to_client_database() ns_namespaces_q = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) ns_ltables_q = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) scp = self._stg_cache_policy_clause() diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 348edaeb..051c77fb 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -729,7 +729,7 @@ def test_create_namespace_physical_tables_sn_inline_vector_index(self): assert "VECTOR INDEX idx_vec(embedding)" in data_create assert "FULLTEXT INDEX idx_fts(document) WITH PARSER ik" in data_create assert "SEARCH INDEX idx_json(data_content)" in data_create - assert "fresh_mode=spfresh" in data_create + assert "centroids_fresh_mode=spfresh" in data_create assert not any(s.startswith("CREATE VECTOR INDEX") for s in c.executed_sqls) assert not any("_hot_table" in s for s in c.executed_sqls if s.startswith("CREATE TABLE")) @@ -749,7 +749,7 @@ def test_create_namespace_physical_tables_ss_inline_vector_index(self): assert "VECTOR INDEX idx_vec(embedding)" in data_create assert "FULLTEXT INDEX" not in data_create assert "SEARCH INDEX idx_json(data_content)" in data_create - assert "fresh_mode=spfresh" in data_create + assert "centroids_fresh_mode=spfresh" in data_create assert not any(s.startswith("CREATE VECTOR INDEX") for s in c.executed_sqls) hot_create = next( s for s in c.executed_sqls if "CREATE TABLE" in s and "_hot_table" in s From c77839e690aaa6396adb8785ca5cd077ee0939cb Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 18 Jun 2026 15:44:41 +0800 Subject: [PATCH 52/84] fix_drop_table --- src/pyseekdb/client/client_base.py | 41 ------------------------------ 1 file changed, 41 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 19c9d223..aa3ea71e 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1172,49 +1172,8 @@ def _use_catalog_database(self) -> None: """Align session with pymysql database= so PL (DROP_NAMESPACE) uses the same DB.""" self._execute(f"USE `{self._catalog_database()}`") - # NOTE: temporarily restored (reverting commit fb708d2) to reproduce the - # multi-database primary-key conflict. Remove again after repro. - def _isolate_sdk_catalog_to_client_database(self) -> None: - """Drop sdk_* catalog tables in other databases so kernel bg task scans self.database. - - ObLTableBGTaskScheduler picks the first tenant database that has sdk_collections. - mysqltest may leave a copy in the `oceanbase` database while pyseekdb uses `test`. - """ - target = self._catalog_database() - try: - rows = self._execute( - "SELECT DISTINCT table_schema AS table_schema " - "FROM information_schema.tables " - "WHERE table_name = 'sdk_collections'" - ) - except Exception: - return - other_schemas: list[str] = [] - for row in rows or []: - schema = row["table_schema"] if isinstance(row, dict) else row[0] - if schema and schema != target: - other_schemas.append(schema) - if not other_schemas: - return - catalog_tables = [ - CollectionNames.sdk_collections_table_name(), - NamespaceCollectionNames.sdk_namespaces_table(), - NamespaceCollectionNames.sdk_ltables_table(), - NamespaceCollectionNames.sdk_namespaces_stats_table(), - ] - for schema in other_schemas: - for table in catalog_tables: - with contextlib.suppress(Exception): - self._execute(f"DROP TABLE IF EXISTS `{schema}`.`{table}`") - logger.info( - "Removed duplicate sdk catalog tables from %s (catalog DB is %s)", - other_schemas, - target, - ) - def _ensure_namespace_catalogs(self) -> None: self._use_catalog_database() - self._isolate_sdk_catalog_to_client_database() ns_namespaces_q = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) ns_ltables_q = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) scp = self._stg_cache_policy_clause() From c27ea4a98eb3dbe49df5f98d6ab2531a209800ff Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 22 Jun 2026 11:44:35 +0800 Subject: [PATCH 53/84] =?UTF-8?q?fix:=E9=99=A4data=5Ftable=E6=84=8F?= =?UTF-8?q?=E5=A4=96=E9=BB=98=E8=AE=A4=E4=BD=BF=E7=94=A8index=E8=A1=A8=20h?= =?UTF-8?q?ttps://project.alipay.com/project/W24001004732/P26001026476/wor?= =?UTF-8?q?kitemDetail=3FopenWorkItemId=3D2026061700116808472?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index aa3ea71e..cc041b15 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1098,7 +1098,7 @@ def _create_sdk_collections_if_not_exists(self) -> None: created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, INDEX idx_name(collection_name) - ) COMMENT='Settings of collections created by SDK' {scp};""" + ) COMMENT='Settings of collections created by SDK' ORGANIZATION INDEX {scp};""" self._execute(create_table_sql) except Exception as e: raise ValueError(f"Failed to create sdk_collections table: {e}") from e @@ -1187,7 +1187,7 @@ def _ensure_namespace_catalogs(self) -> None: PRIMARY KEY (namespace_id), UNIQUE KEY uk_sdk_ns_coll_name (collection_id, namespace_name), KEY idx_sdk_ns_by_collection (collection_id) - ) COMMENT='Namespace catalog' {scp};""" + ) COMMENT='Namespace catalog' ORGANIZATION INDEX {scp};""" ns_ltables_sql = f"""CREATE TABLE IF NOT EXISTS {ns_ltables_q} ( ltable_id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, collection_id CHAR(32) NOT NULL, @@ -1199,7 +1199,7 @@ def _ensure_namespace_catalogs(self) -> None: PRIMARY KEY (ltable_id), UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name), KEY idx_sdk_lt_by_ns (collection_id, namespace_id) - ) COMMENT='LTable catalog' {scp};""" + ) COMMENT='LTable catalog' ORGANIZATION INDEX {scp};""" namespaces_stats_sql = f"""CREATE TABLE IF NOT EXISTS {self._qtable(NamespaceCollectionNames.sdk_namespaces_stats_table())} ( collection_id CHAR(32) NOT NULL COMMENT 'collection id', namespace_id BIGINT UNSIGNED NOT NULL COMMENT 'namespace internal id', @@ -1212,7 +1212,7 @@ def _ensure_namespace_catalogs(self) -> None: included_index BOOL NOT NULL DEFAULT FALSE COMMENT 'whether stats include index data', PRIMARY KEY (namespace_id, ltable_id, included_index), KEY idx_sdk_ns_stat_by_collection (collection_id) - ) COMMENT='Logic table row count and storage size statistics' DEFAULT CHARSET=utf8mb4 + ) COMMENT='Logic table row count and storage size statistics' DEFAULT CHARSET=utf8mb4 ORGANIZATION INDEX PARTITION BY KEY(namespace_id) PARTITIONS 8;""" self._execute(ns_namespaces_sql) self._execute(ns_ltables_sql) @@ -1367,7 +1367,7 @@ def _create_namespace_physical_tables( created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, PRIMARY KEY(namespace_id) - ) TABLEGROUP=`{tg_name}` COMMENT='热点/TTL附属表' DEFAULT CHARSET=utf8mb4 + ) TABLEGROUP=`{tg_name}` COMMENT='热点/TTL附属表' DEFAULT CHARSET=utf8mb4 ORGANIZATION INDEX {partition_clause}""") self._execute(f"""CREATE TABLE IF NOT EXISTS `{kv_table}` ( @@ -1375,7 +1375,7 @@ def _create_namespace_physical_tables( kv_key VARBINARY(1024) NOT NULL, kv_value LONGBLOB NOT NULL, PRIMARY KEY(namespace_id, kv_key) - ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 LOB_INROW_THRESHOLD=786432 + ) TABLEGROUP=`{tg_name}` COMMENT='索引与映射KV表' DEFAULT CHARSET=utf8mb4 ORGANIZATION INDEX LOB_INROW_THRESHOLD=786432 {partition_clause}""") self._execute(f"""CREATE TABLE IF NOT EXISTS `{schema_table}` ( @@ -1386,7 +1386,7 @@ def _create_namespace_physical_tables( created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, PRIMARY KEY(namespace_id, ltable_id) - ) TABLEGROUP=`{tg_name}` COMMENT='LTable schema定义' DEFAULT CHARSET=utf8mb4 + ) TABLEGROUP=`{tg_name}` COMMENT='LTable schema定义' DEFAULT CHARSET=utf8mb4 ORGANIZATION INDEX {partition_clause}""") except Exception: From 626cacaea149ea4c243dc27a5c5d38f6473a0f99 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 22 Jun 2026 17:01:46 +0800 Subject: [PATCH 54/84] =?UTF-8?q?pyseekdb=E5=AF=B9=E5=A4=96ivf=E5=8F=82?= =?UTF-8?q?=E6=95=B0fresh=5Fmode=E6=94=B9=E4=B8=BAcentroids=5Ffresh=5Fmode?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 8 ++--- src/pyseekdb/client/configuration.py | 10 +++--- .../namespace_dml_helpers.py | 2 +- .../namespace_fts_helpers.py | 2 +- .../test_namespace_constraints.py | 2 +- .../test_namespace_drop_validation.py | 2 +- .../test_namespace_lifecycle.py | 10 +++--- .../test_namespace_optional_indexes.py | 4 +-- .../test_namespace_prewarm.py | 2 +- .../integration_tests/test_namespace_query.py | 2 +- .../test_namespace_session_vars.py | 2 +- tests/unit_tests/test_namespace.py | 32 +++++++++---------- 12 files changed, 40 insertions(+), 38 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index ea36ad2a..35a3208d 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -249,8 +249,8 @@ def _get_ivf_vector_index_sql(ivf_config: "IVFConfiguration") -> str: property_parts.append(f"{k}='{v}'") else: property_parts.append(f"{k}={v}") - if ivf_config.fresh_mode: - property_parts.append(f"centroids_fresh_mode={ivf_config.fresh_mode}") + if ivf_config.centroids_fresh_mode: + property_parts.append(f"centroids_fresh_mode={ivf_config.centroids_fresh_mode}") property_str = ", ".join(property_parts) properties_str = f", {property_str}" if property_str else "" return f"WITH (DISTANCE={ivf_config.distance}, TYPE={ivf_config.type.upper()}, LIB={ivf_config.lib.upper()}{properties_str})" @@ -1040,8 +1040,8 @@ def _create_namespace_collection(self, name: str, schema: Schema, partition_coun } if ivf_config is not None: settings["dense_index_type"] = "ivf" - if ivf_config.fresh_mode is not None: - settings["fresh_mode"] = ivf_config.fresh_mode + if ivf_config.centroids_fresh_mode is not None: + settings["centroids_fresh_mode"] = ivf_config.centroids_fresh_mode if dense_embedding_function is not None and EmbeddingFunction.support_persistence(dense_embedding_function): settings["embedding_function"] = { "name": dense_embedding_function.name(), diff --git a/src/pyseekdb/client/configuration.py b/src/pyseekdb/client/configuration.py index 930b905d..3c683c60 100644 --- a/src/pyseekdb/client/configuration.py +++ b/src/pyseekdb/client/configuration.py @@ -336,7 +336,7 @@ class IVFConfiguration: dimension: Vector dimension (number of elements in each vector) distance: Distance metric for similarity calculation (e.g., 'l2', 'cosine', 'inner_product') type: IVF index subtype ('ivf_flat', 'ivf_sq8', 'ivf_pq') - fresh_mode: SPFresh mode for online index updates (e.g., 'spfresh'). Defaults to None (disabled). + centroids_fresh_mode: SPFresh mode for online index updates (e.g., 'spfresh'). Defaults to None (disabled). properties: Optional dictionary of additional IVF index properties """ @@ -344,7 +344,7 @@ class IVFConfiguration: distance: str | DistanceMetric = DistanceMetric.COSINE.value type: str | IVFIndexType = IVFIndexType.IVF_FLAT.value lib: str | IVFIndexLib = IVFIndexLib.OB.value - fresh_mode: str | None = None + centroids_fresh_mode: str | None = None properties: dict[str, PrimitiveValue] | None = None def __post_init__(self): @@ -367,8 +367,10 @@ def __post_init__(self): if self.lib not in valid_libs: raise ValueError(f"lib must be one of {valid_libs}, got {self.lib}") - if self.fresh_mode is not None and not isinstance(self.fresh_mode, str): - raise TypeError(f"fresh_mode must be a str, got {type(self.fresh_mode).__name__}") + if self.centroids_fresh_mode is not None and not isinstance(self.centroids_fresh_mode, str): + raise TypeError( + f"centroids_fresh_mode must be a str, got {type(self.centroids_fresh_mode).__name__}" + ) _ensure_primitive_properties(self.properties) diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index 835644b9..3f130b95 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -39,7 +39,7 @@ class DmlRecord: def ns_schema() -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), embedding_function=None, ), fulltext_index=FulltextIndexConfig(analyzer="ik"), diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index 9dcf6deb..98ca5de0 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -56,7 +56,7 @@ class FtsQueryCase: def ns_schema(distance: VectorDistanceMetric = "l2") -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance=distance, fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance=distance, centroids_fresh_mode="spfresh"), embedding_function=None, ), fulltext_index=FulltextIndexConfig(analyzer="ik"), diff --git a/tests/integration_tests/test_namespace_constraints.py b/tests/integration_tests/test_namespace_constraints.py index 11a3f113..9bf5e5d1 100644 --- a/tests/integration_tests/test_namespace_constraints.py +++ b/tests/integration_tests/test_namespace_constraints.py @@ -22,7 +22,7 @@ def _make_schema(ivf_type: str) -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", type=ivf_type, fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="cosine", type=ivf_type, centroids_fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index f314b6d0..8dcd4195 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -35,7 +35,7 @@ def _make_collection(client, suffix: str = ""): name = f"test_ns_drop_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index b7b23665..33b57abd 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -21,7 +21,7 @@ def _create_ns_collection(self, client, suffix=""): name = f"test_ns_lc_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), embedding_function=None, ), ) @@ -56,7 +56,7 @@ def test_get_collection_restores_embedding_function(self, db_client): name = f"test_ns_ef_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=ef.dimension, distance="cosine", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=ef.dimension, distance="cosine", centroids_fresh_mode="spfresh"), embedding_function=ef, ), ) @@ -156,7 +156,7 @@ def test_create_resumes_incomplete_collection(self, oceanbase_client): name = f"test_ns_resume_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), embedding_function=None, ), ) @@ -257,7 +257,7 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): name = f"test_ns_ss_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), embedding_function=None, ), ) @@ -297,7 +297,7 @@ def test_custom_collection_partition_count(self, oceanbase_client): name = f"test_ns_lc_pc_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_optional_indexes.py b/tests/integration_tests/test_namespace_optional_indexes.py index dd01980b..ba957da2 100644 --- a/tests/integration_tests/test_namespace_optional_indexes.py +++ b/tests/integration_tests/test_namespace_optional_indexes.py @@ -38,7 +38,7 @@ def _schema_fts_only() -> Schema: def _schema_vector_only() -> Schema: return Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), embedding_function=None, ), ) @@ -125,7 +125,7 @@ def test_vector_only_indexes_and_queries(self, db_client): settings = _collection_settings(db_client, collection.id) assert keys == {"idx_json", "idx_vec"} assert settings.get("dense_index_type") == "ivf" - assert settings.get("fresh_mode") == "spfresh" + assert settings.get("centroids_fresh_mode") == "spfresh" ns = collection.create_namespace("ns") _seed_namespace(ns, collection.dimension) diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index f5458c14..b63bd12f 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -23,7 +23,7 @@ def _create_ns_collection_and_namespace(self, client): name = f"test_ns_pw_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="cosine", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index cbd82228..a0613300 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -19,7 +19,7 @@ def _setup(self, client, suffix=""): name = f"test_ns_q_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), embedding_function=None, ), fulltext_index=FulltextIndexConfig(analyzer="ik"), diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py index 178339f7..166b145b 100644 --- a/tests/integration_tests/test_namespace_session_vars.py +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -21,7 +21,7 @@ def _create_ns_collection(self, client): name = f"test_ns_sessvar_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh"), + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), embedding_function=None, ), ) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 051c77fb..391fa234 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -32,14 +32,14 @@ def test_valid_defaults(self): config = IVFConfiguration() assert config.dimension == 384 assert config.distance == "cosine" - assert config.fresh_mode is None + assert config.centroids_fresh_mode is None assert config.properties is None def test_valid_custom(self): - config = IVFConfiguration(dimension=128, distance="l2", fresh_mode="spfresh") + config = IVFConfiguration(dimension=128, distance="l2", centroids_fresh_mode="spfresh") assert config.dimension == 128 assert config.distance == "l2" - assert config.fresh_mode == "spfresh" + assert config.centroids_fresh_mode == "spfresh" def test_valid_inner_product(self): config = IVFConfiguration(dimension=1024, distance="inner_product") @@ -77,9 +77,9 @@ def test_invalid_distance(self): with pytest.raises(ValueError, match="distance must be one of"): IVFConfiguration(distance="invalid") - def test_invalid_fresh_mode_type(self): - with pytest.raises(TypeError, match="fresh_mode must be a str"): - IVFConfiguration(fresh_mode=True) + def test_invalid_centroids_fresh_mode_type(self): + with pytest.raises(TypeError, match="centroids_fresh_mode must be a str"): + IVFConfiguration(centroids_fresh_mode=True) def test_properties_valid(self): config = IVFConfiguration(properties={"nlist": 128, "nprobe": 16}) @@ -715,7 +715,7 @@ def test_create_namespace_physical_tables_sn_inline_vector_index(self): from pyseekdb.client.configuration import FulltextIndexConfig c = self._client() - ivf_config = IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh") + ivf_config = IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh") c._create_namespace_physical_tables( collection_id=self.COLLECTION_ID, dimension=3, @@ -736,7 +736,7 @@ def test_create_namespace_physical_tables_sn_inline_vector_index(self): def test_create_namespace_physical_tables_ss_inline_vector_index(self): """SS logic_data_table: inline VECTOR INDEX in CREATE TABLE.""" c = self._client() - ivf_config = IVFConfiguration(dimension=3, distance="l2", fresh_mode="spfresh") + ivf_config = IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh") c._create_namespace_physical_tables( collection_id=self.COLLECTION_ID, dimension=3, @@ -771,8 +771,8 @@ def test_create_namespace_physical_tables_search_index_only(self): assert "VECTOR INDEX" not in data_create assert "FULLTEXT INDEX" not in data_create - def test_create_namespace_physical_tables_ivf_without_fresh_mode(self): - """IVF without fresh_mode omits fresh_mode from VECTOR INDEX DDL.""" + def test_create_namespace_physical_tables_ivf_without_centroids_fresh_mode(self): + """IVF without centroids_fresh_mode omits centroids_fresh_mode from VECTOR INDEX DDL.""" c = self._client() ivf_config = IVFConfiguration(dimension=3, distance="l2") c._create_namespace_physical_tables( @@ -785,7 +785,7 @@ def test_create_namespace_physical_tables_ivf_without_fresh_mode(self): s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s ) assert "VECTOR INDEX idx_vec(embedding)" in data_create - assert "fresh_mode" not in data_create + assert "centroids_fresh_mode" not in data_create # ==================== Namespace Name Validation Tests ==================== @@ -1087,7 +1087,7 @@ def test_ob_type_validation(self): c.detect_db_type_and_version = MagicMock(return_value=("mysql", "8.0")) from pyseekdb.client.schema import Schema from pyseekdb.client.configuration import VectorIndexConfig - schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, fresh_mode="spfresh"), embedding_function=None)) + schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, centroids_fresh_mode="spfresh"), embedding_function=None)) with pytest.raises(ValueError, match="only supported on OceanBase"): c._create_namespace_collection("test", schema) @@ -1109,9 +1109,9 @@ def test_create_namespace_collection_without_ivf_skips_vector_index(self): assert "SEARCH INDEX idx_json(data_content)" in data_create settings = c._create_ns_collection_meta.call_args[0][1] assert "dense_index_type" not in settings - assert "fresh_mode" not in settings + assert "centroids_fresh_mode" not in settings - def test_create_namespace_collection_ivf_without_fresh_mode(self): + def test_create_namespace_collection_ivf_without_centroids_fresh_mode(self): c = FakeClient() c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", "4.3")) c._is_shared_storage_mode = MagicMock(return_value=False) @@ -1131,10 +1131,10 @@ def test_create_namespace_collection_ivf_without_fresh_mode(self): s for s in c.executed_sqls if "CREATE TABLE" in s and "logic_data_table" in s ) assert "VECTOR INDEX idx_vec(embedding)" in data_create - assert "fresh_mode" not in data_create + assert "centroids_fresh_mode" not in data_create settings = c._create_ns_collection_meta.call_args[0][1] assert settings["dense_index_type"] == "ivf" - assert "fresh_mode" not in settings + assert "centroids_fresh_mode" not in settings # ==================== Delete Namespace Uses Kernel ==================== From dff351fce64b5d9b161ec5a7208651909727daf3 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 22 Jun 2026 17:51:19 +0800 Subject: [PATCH 55/84] =?UTF-8?q?fix:=E3=80=90AgentDB/=E9=80=BB=E8=BE=91?= =?UTF-8?q?=E8=A1=A8=E3=80=91=E5=B9=B6=E5=8F=91=20get=5For=5Fcreate=5Fname?= =?UTF-8?q?space=20=E5=90=8C=E5=90=8D=20namespace=20=E4=B8=8D=E5=B9=82?= =?UTF-8?q?=E7=AD=89=EF=BC=8C=E5=94=AF=E4=B8=80=E9=94=AE=E5=86=B2=E7=AA=81?= =?UTF-8?q?=201062=20https://project.alipay.com/project/W24001004732/P2600?= =?UTF-8?q?1026476/workitemDetail=3FopenWorkItemId=3D2026062200116890620&s?= =?UTF-8?q?tatus=3Dstatus:?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 146 ++++++++++++++++-- src/pyseekdb/client/collection.py | 4 +- .../test_logic_table_monitoring_p1.py | 99 ++++++++++++ ...est_get_or_create_namespace_concurrency.py | 89 +++++++++++ 4 files changed, 318 insertions(+), 20 deletions(-) create mode 100644 tests/integration_tests/test_logic_table_monitoring_p1.py create mode 100644 tests/unit_tests/test_get_or_create_namespace_concurrency.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 35a3208d..d311b5bd 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -91,6 +91,25 @@ def _is_collection_conflict_error(exc: BaseException) -> bool: return False +def _is_namespace_catalog_conflict_error(exc: BaseException) -> bool: + current: BaseException | None = exc + while current is not None: + message = str(current).lower() + if ( + "duplicate entry" in message + and ( + "uk_sdk_ns_coll_name" in message + or "uk_sdk_lt_coll_ns_name" in message + or "code=1062" in message + ) + ): + return True + if type(current).__name__ == "IntegrityError" and "1062" in message: + return True + current = current.__cause__ + return False + + def _extract_hnsw_config(config: ConfigurationParam) -> HNSWConfiguration | None: if config is None: return None @@ -1530,30 +1549,95 @@ def _set_session_ns_context( if ltable_id is not None: self._execute(f"SET @ltable_id = {int(ltable_id)}") - def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: + def _fetch_ns_namespace_id(self, collection_id: str, namespace_name: str) -> int: namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) ns_table = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) - lt_table = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) - self._execute( - f"INSERT INTO {ns_table} " - f"(collection_id, namespace_name) VALUES ('{collection_id_escaped}', '{namespace_name_escaped}')" - ) rows = self._execute( f"SELECT namespace_id FROM {ns_table} " f"WHERE collection_id = '{collection_id_escaped}' AND namespace_name = '{namespace_name_escaped}'" ) - ns_id = int(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["namespace_id"]) - self._execute( - f"INSERT INTO {lt_table} " - f"(collection_id, namespace_id, ltable_name) VALUES ('{collection_id_escaped}', {ns_id}, 'default')" - ) + if not rows: + raise ValueError( + f"Namespace '{namespace_name}' not found for collection_id={collection_id} in sdk_namespaces" + ) + return int(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["namespace_id"]) + + def _fetch_ns_ltable_id( + self, + collection_id: str, + namespace_id: int, + ltable_name: str = "default", + ) -> int: + collection_id_escaped = escape_string(collection_id) + ltable_name_escaped = escape_string(ltable_name) + lt_table = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) lt_rows = self._execute( f"SELECT ltable_id FROM {lt_table} " - f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {ns_id} AND ltable_name = 'default'" + f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {int(namespace_id)} " + f"AND ltable_name = '{ltable_name_escaped}'" ) - lt_id = int(lt_rows[0][0] if isinstance(lt_rows[0], (list, tuple)) else lt_rows[0]["ltable_id"]) - self._cache_namespace_ltable_id(collection_id, ns_id, lt_id) + if not lt_rows: + raise ValueError( + f"LTable '{ltable_name}' not found for collection_id={collection_id}, " + f"namespace_id={namespace_id} in sdk_ltables" + ) + return int(lt_rows[0][0] if isinstance(lt_rows[0], (list, tuple)) else lt_rows[0]["ltable_id"]) + + def _insert_ns_namespace_catalog_row( + self, + collection_id: str, + namespace_name: str, + *, + idempotent: bool, + ) -> int: + namespace_name_escaped = escape_string(namespace_name) + collection_id_escaped = escape_string(collection_id) + ns_table = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) + try: + self._execute( + f"INSERT INTO {ns_table} " + f"(collection_id, namespace_name) VALUES ('{collection_id_escaped}', '{namespace_name_escaped}')" + ) + except Exception as exc: + if idempotent and _is_namespace_catalog_conflict_error(exc): + pass + else: + raise + return self._fetch_ns_namespace_id(collection_id, namespace_name) + + def _insert_ns_ltable_catalog_row( + self, + collection_id: str, + namespace_id: int, + *, + idempotent: bool, + ltable_name: str = "default", + ) -> int: + collection_id_escaped = escape_string(collection_id) + ltable_name_escaped = escape_string(ltable_name) + lt_table = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) + try: + self._execute( + f"INSERT INTO {lt_table} " + f"(collection_id, namespace_id, ltable_name) " + f"VALUES ('{collection_id_escaped}', {int(namespace_id)}, '{ltable_name_escaped}')" + ) + except Exception as exc: + if idempotent and _is_namespace_catalog_conflict_error(exc): + pass + else: + raise + return self._fetch_ns_ltable_id(collection_id, namespace_id, ltable_name) + + def _finalize_ns_namespace_meta( + self, + collection_id: str, + namespace_name: str, + namespace_id: int, + ltable_id: int, + ) -> dict: + self._cache_namespace_ltable_id(collection_id, namespace_id, ltable_id) schema_table = self._qtable( NamespaceCollectionNames.logic_schema_table_name(collection_id) ) @@ -1561,10 +1645,38 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> with contextlib.suppress(Exception): self._execute( f"INSERT INTO {schema_table} (namespace_id, ltable_id, schema_content) " - f"VALUES ({ns_id}, {lt_id}, '{escape_string(schema_content)}')" + f"VALUES ({namespace_id}, {ltable_id}, '{escape_string(schema_content)}')" + ) + self._set_session_ns_context(namespace_id=namespace_id, ltable_id=ltable_id) + return {"namespace_id": str(namespace_id), "namespace_name": namespace_name, "ltable_id": str(ltable_id)} + + def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: + ns_id = self._insert_ns_namespace_catalog_row( + collection_id, namespace_name, idempotent=False + ) + lt_id = self._insert_ns_ltable_catalog_row( + collection_id, ns_id, idempotent=False + ) + return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) + + def _get_or_create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: + meta = self._get_ns_namespace_meta(collection_id, namespace_name) + if meta is not None: + if meta.get("ltable_id") is not None: + return meta + lt_id = self._insert_ns_ltable_catalog_row( + collection_id, int(meta["namespace_id"]), idempotent=True + ) + return self._finalize_ns_namespace_meta( + collection_id, namespace_name, int(meta["namespace_id"]), lt_id ) - self._set_session_ns_context(namespace_id=ns_id, ltable_id=lt_id) - return {"namespace_id": str(ns_id), "namespace_name": namespace_name} + ns_id = self._insert_ns_namespace_catalog_row( + collection_id, namespace_name, idempotent=True + ) + lt_id = self._insert_ns_ltable_catalog_row( + collection_id, ns_id, idempotent=True + ) + return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: namespace_name_escaped = escape_string(namespace_name) diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 87fa9ff9..1ac5e8a7 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -172,9 +172,7 @@ def get_or_create_namespace(self, name: str) -> "Namespace": self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace - meta = self._client._get_ns_namespace_meta(self._id, name) - if meta is None: - meta = self._client._create_ns_namespace_meta(self._id, name) + meta = self._client._get_or_create_ns_namespace_meta(self._id, name) return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) def delete_namespace(self, name: str) -> None: diff --git a/tests/integration_tests/test_logic_table_monitoring_p1.py b/tests/integration_tests/test_logic_table_monitoring_p1.py new file mode 100644 index 00000000..9294bf50 --- /dev/null +++ b/tests/integration_tests/test_logic_table_monitoring_p1.py @@ -0,0 +1,99 @@ +""" +Logic table monitoring P1 integration tests. + +Includes concurrent get_or_create_namespace idempotency coverage. +""" + +from __future__ import annotations + +import threading +import time +import uuid +from concurrent.futures import ThreadPoolExecutor, as_completed + +import pytest + +import pyseekdb +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, ns_schema + + +def _make_oceanbase_client(): + import os + + return pyseekdb.Client( + host=os.environ.get("OB_HOST", "127.0.0.1"), + port=int(os.environ.get("OB_PORT", "10902")), + tenant=os.environ.get("OB_TENANT", "mysql"), + database=os.environ.get("OB_DATABASE", "test"), + user=os.environ.get("OB_USER", "root"), + password=os.environ.get("OB_PASSWORD", ""), + ) + + +def _count_sdk_namespaces(client, collection_id: str, namespace_name: str) -> int: + rows = client._server._execute( + "SELECT COUNT(*) AS cnt FROM sdk_namespaces " + f"WHERE collection_id = '{collection_id}' AND namespace_name = '{namespace_name}'" + ) + row = rows[0] + return int(row["cnt"] if isinstance(row, dict) else row[0]) + + +def _count_sdk_ltables(client, collection_id: str, namespace_id: int) -> int: + rows = client._server._execute( + "SELECT COUNT(*) AS cnt FROM sdk_ltables " + f"WHERE collection_id = '{collection_id}' AND namespace_id = {int(namespace_id)} " + "AND ltable_name = 'default'" + ) + row = rows[0] + return int(row["cnt"] if isinstance(row, dict) else row[0]) + + +class TestLogicTableMonitoringP1: + def test_concurrent_get_or_create_same_namespace_is_idempotent(self, oceanbase_client): + """8 threads with independent clients race on the same namespace name.""" + owner = oceanbase_client + collection_name = f"ltmon_p1_{uuid.uuid4().hex[:12]}" + collection = owner.create_collection( + name=collection_name, + schema=ns_schema(), + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + coll_id = collection.id + namespace_name = "ns_shared" + barrier = threading.Barrier(8) + results: dict[int, dict] = {} + errors: list[str] = [] + + def _worker(thread_id: int) -> None: + client = _make_oceanbase_client() + try: + coll = client.get_collection(collection_name) + barrier.wait(timeout=30) + ns = coll.get_or_create_namespace(namespace_name) + results[thread_id] = { + "namespace_id": str(ns.namespace_id), + "name": ns.name, + } + except Exception as exc: # noqa: BLE001 + errors.append(f"thread {thread_id}: {exc!r}") + finally: + if hasattr(client, "close"): + client.close() + + with ThreadPoolExecutor(max_workers=8) as pool: + futures = [pool.submit(_worker, i) for i in range(8)] + for fut in as_completed(futures, timeout=60): + fut.result() + + try: + assert not errors, f"concurrent get_or_create_namespace failed: {errors}" + assert len(results) == 8, results + ns_ids = {item["namespace_id"] for item in results.values()} + assert len(ns_ids) == 1, f"expected one namespace_id, got {ns_ids}" + ns_id = next(iter(ns_ids)) + assert _count_sdk_namespaces(owner, coll_id, namespace_name) == 1 + assert _count_sdk_ltables(owner, coll_id, int(ns_id)) == 1 + finally: + owner.delete_collection(name=collection_name) diff --git a/tests/unit_tests/test_get_or_create_namespace_concurrency.py b/tests/unit_tests/test_get_or_create_namespace_concurrency.py new file mode 100644 index 00000000..52f23fe6 --- /dev/null +++ b/tests/unit_tests/test_get_or_create_namespace_concurrency.py @@ -0,0 +1,89 @@ +""" +Unit tests for concurrent-safe get_or_create_namespace helpers. +""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock, call + +import pytest + +project_root = Path(__file__).parent.parent.parent +src_root = project_root / "src" +sys.path.insert(0, str(src_root)) + +from pyseekdb.client.client_base import ( # noqa: E402 + BaseClient, + _is_namespace_catalog_conflict_error, +) + + +class TestNamespaceCatalogConflictDetection: + def test_detects_integrity_error_on_sdk_namespaces(self): + class IntegrityError(Exception): + pass + + exc = IntegrityError( + '(1062, "Duplicate entry \'cid-ns_shared\' for key \'uk_sdk_ns_coll_name\'")' + ) + assert _is_namespace_catalog_conflict_error(exc) + + def test_detects_integrity_error_on_sdk_ltables(self): + class IntegrityError(Exception): + pass + + exc = IntegrityError( + '(1062, "Duplicate entry \'cid-1-default\' for key \'uk_sdk_lt_coll_ns_name\'")' + ) + assert _is_namespace_catalog_conflict_error(exc) + + def test_ignores_unrelated_errors(self): + assert not _is_namespace_catalog_conflict_error(ValueError("invalid namespace name")) + + +class TestGetOrCreateNamespaceMetaRecovery: + def test_idempotent_insert_reuses_existing_namespace(self): + client = MagicMock(spec=BaseClient) + client._get_ns_namespace_meta.return_value = None + client._insert_ns_namespace_catalog_row.return_value = 42 + client._insert_ns_ltable_catalog_row.return_value = 7 + client._finalize_ns_namespace_meta.return_value = { + "namespace_id": "42", + "namespace_name": "ns_shared", + "ltable_id": "7", + } + + result = BaseClient._get_or_create_ns_namespace_meta(client, "cid", "ns_shared") + + assert result["namespace_id"] == "42" + client._insert_ns_namespace_catalog_row.assert_called_once_with( + "cid", "ns_shared", idempotent=True + ) + client._insert_ns_ltable_catalog_row.assert_called_once_with( + "cid", 42, idempotent=True + ) + + def test_existing_namespace_without_ltable_creates_default_ltable(self): + client = MagicMock(spec=BaseClient) + client._get_ns_namespace_meta.return_value = { + "namespace_id": "42", + "namespace_name": "ns_shared", + } + client._insert_ns_ltable_catalog_row.return_value = 7 + client._finalize_ns_namespace_meta.return_value = { + "namespace_id": "42", + "namespace_name": "ns_shared", + "ltable_id": "7", + } + + result = BaseClient._get_or_create_ns_namespace_meta(client, "cid", "ns_shared") + + assert result["ltable_id"] == "7" + client._insert_ns_namespace_catalog_row.assert_not_called() + client._insert_ns_ltable_catalog_row.assert_called_once_with( + "cid", 42, idempotent=True + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From d16c80eeff76aed7f97584f37e574b5dbc2d2c40 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 22 Jun 2026 19:39:06 +0800 Subject: [PATCH 56/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91=E5=B9=B6?= =?UTF-8?q?=E5=8F=91=E6=89=A7=E8=A1=8Cns.upsert,=E4=BF=9D=E5=AD=98?= =?UTF-8?q?=E5=A4=9A=E8=A1=8C=E7=9B=B8=E5=90=8Cids=20https://project.alipa?= =?UTF-8?q?y.com/project/W24001004732/P26001026476/workitemDetail=3FopenWo?= =?UTF-8?q?rkItemId=3D2026062200116892785&status=3Dstatus?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 163 ++++++++++++++++++ .../test_namespace_upsert_concurrency.py | 80 +++++++++ .../test_namespace_upsert_reconcile.py | 79 +++++++++ 3 files changed, 322 insertions(+) create mode 100644 tests/integration_tests/test_namespace_upsert_concurrency.py create mode 100644 tests/unit_tests/test_namespace_upsert_reconcile.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index d311b5bd..eb075ecc 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -3,11 +3,13 @@ """ import contextlib +import hashlib import json import logging import os import re import struct +import time import warnings from abc import ABC, abstractmethod from collections.abc import Sequence @@ -4981,6 +4983,152 @@ def _append_namespace_filter( return f"WHERE {ns_cond} AND ({inner})", params return f"WHERE {ns_cond} AND ({where_clause})", params + @staticmethod + def _namespace_record_lock_name( + collection_id: str, + namespace_id: str | int, + ltable_id: int, + record_id: str, + ) -> str: + key = f"{collection_id}:{namespace_id}:{ltable_id}:{record_id}" + digest = hashlib.sha256(key.encode()).hexdigest()[:40] + return f"pyseekdb:nsu:{digest}" + + @contextlib.contextmanager + def _namespace_record_lock( + self, + collection_id: str, + namespace_id: str | int, + ltable_id: int, + record_id: str, + *, + total_wait_seconds: float = 30.0, + ): + lock_name = self._namespace_record_lock_name( + collection_id, namespace_id, ltable_id, record_id + ) + lock_name_escaped = escape_string(lock_name) + acquired = False + deadline = time.monotonic() + total_wait_seconds + try: + while time.monotonic() < deadline: + rows = self._execute( + f"SELECT GET_LOCK('{lock_name_escaped}', 1) AS got" + ) + row = rows[0] + got = row.get("got") if isinstance(row, dict) else row[0] + if got == 1: + acquired = True + yield True + return + time.sleep(0.05) + yield False + finally: + if acquired: + with contextlib.suppress(Exception): + self._execute(f"SELECT RELEASE_LOCK('{lock_name_escaped}')") + + def _count_namespace_records_by_id( + self, + table_name: str, + namespace_id: int, + ltable_id: int, + record_id: str, + ) -> int: + id_expr = _NS_DATA_CONTENT_ID_EXPR + sql = ( + f"SELECT COUNT(*) AS cnt FROM `{table_name}` " + f"WHERE namespace_id = {int(namespace_id)} AND ltable_id = {int(ltable_id)} " + f"AND {id_expr} = %s" + ) + conn = self._ensure_connection() + use_ctx = self._use_context_manager_for_cursor() + rows = self._execute_query_with_cursor(conn, sql, [record_id], use_ctx) + if not rows: + return 0 + row = rows[0] + if isinstance(row, dict): + return int(row.get("cnt", 0)) + if isinstance(row, (tuple, list)): + return int(row[0]) + return int(row) + + def _delete_namespace_records_by_id( + self, + table_name: str, + namespace_id: int, + ltable_id: int, + record_id: str, + ) -> None: + id_expr = _NS_DATA_CONTENT_ID_EXPR + sql = ( + f"DELETE FROM `{table_name}` " + f"WHERE namespace_id = {int(namespace_id)} AND ltable_id = {int(ltable_id)} " + f"AND {id_expr} = %s" + ) + conn = self._ensure_connection() + use_ctx = self._use_context_manager_for_cursor() + if use_ctx: + with conn.cursor() as cursor: + cursor.execute(sql, [record_id]) + else: + cursor = conn.cursor() + try: + cursor.execute(sql, [record_id]) + finally: + cursor.close() + + def _reconcile_namespace_duplicate_records( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + ltable_id: int, + table_name: str, + ids: list[str], + documents: list[str] | None, + metadatas: list[dict] | None, + embeddings: list[list[float]] | None, + embedding_function: EmbeddingFunction[EmbeddingDocuments] | None, + **kwargs: Any, + ) -> None: + """Collapse concurrent upsert races to a single row per business id.""" + if not ids or collection_id is None: + return + ns_id = int(namespace_id) + for i, record_id in enumerate(ids): + doc_val = documents[i] if documents and i < len(documents) else None + meta_val = metadatas[i] if metadatas and i < len(metadatas) else None + emb_val = embeddings[i] if embeddings and i < len(embeddings) else None + for _ in range(60): + with self._namespace_record_lock( + collection_id, namespace_id, ltable_id, record_id + ) as acquired: + if not acquired: + continue + duplicate_count = self._count_namespace_records_by_id( + table_name, ns_id, ltable_id, record_id + ) + if duplicate_count <= 1: + break + self._delete_namespace_records_by_id( + table_name, ns_id, ltable_id, record_id + ) + self._namespace_add( + collection_id=collection_id, + collection_name=collection_name, + namespace_id=namespace_id, + namespace_name=namespace_name, + ids=[record_id], + embeddings=[emb_val] if emb_val is not None else None, + metadatas=[meta_val] if meta_val is not None else None, + documents=[doc_val] if doc_val is not None else None, + embedding_function=embedding_function, + **kwargs, + ) + break + def _namespace_add( # noqa: C901 self, collection_id: str | None, @@ -5305,6 +5453,21 @@ def _namespace_upsert( documents=upd_docs, embedding_function=embedding_function, **kwargs, ) + self._reconcile_namespace_duplicate_records( + collection_id=collection_id, + collection_name=collection_name, + namespace_id=namespace_id, + namespace_name=namespace_name, + ltable_id=ltable_id, + table_name=table_name, + ids=ids, + documents=documents, + metadatas=metadatas, + embeddings=embeddings, + embedding_function=embedding_function, + **kwargs, + ) + def _namespace_delete( self, collection_id: str | None, diff --git a/tests/integration_tests/test_namespace_upsert_concurrency.py b/tests/integration_tests/test_namespace_upsert_concurrency.py new file mode 100644 index 00000000..46e64269 --- /dev/null +++ b/tests/integration_tests/test_namespace_upsert_concurrency.py @@ -0,0 +1,80 @@ +"""Concurrent namespace.upsert() repro for duplicate business ids.""" + +from __future__ import annotations + +import contextlib +import os +import threading +import time +import uuid + +import pytest + +import pyseekdb +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, cleanup, ns_schema + + +def _unique_name(prefix: str) -> str: + return f"{prefix}_{time.time_ns()}" + + +def _new_oceanbase_client(): + return pyseekdb.Client( + host=os.environ.get("OB_HOST", "127.0.0.1"), + port=int(os.environ.get("OB_PORT", "10902")), + tenant=os.environ.get("OB_TENANT", "mysql"), + database=os.environ.get("OB_DATABASE", "test"), + user=os.environ.get("OB_USER", "root"), + password=os.environ.get("OB_PASSWORD", ""), + ) + + +def test_multi_client_concurrent_upsert_same_new_id_should_keep_single_record( + oceanbase_client, +): + collection = oceanbase_client.create_collection( + name=_unique_name("test_ns_upsert_race_repro"), + schema=ns_schema(), + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + clients = [_new_oceanbase_client() for _ in range(4)] + try: + ns = collection.create_namespace("race_ns") + namespaces = [ + client.get_collection(collection.name).get_namespace("race_ns") + for client in clients + ] + barrier = threading.Barrier(len(namespaces)) + errors: list[Exception] = [] + + def _upsert_once(target_ns, idx: int) -> None: + try: + barrier.wait() + target_ns.upsert( + ids="same_new_id", + embeddings=[1.0, 2.0, 3.0], + documents=f"written by client {idx}", + metadatas={"client": idx}, + ) + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + threads = [ + threading.Thread(target=_upsert_once, args=(target_ns, idx)) + for idx, target_ns in enumerate(namespaces) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert errors == [] + result = ns.get(ids="same_new_id", include=["documents", "metadatas"]) + assert result["ids"] == ["same_new_id"] + assert ns.count() == 1 + finally: + for client in clients: + with contextlib.suppress(Exception): + client.close() + cleanup(oceanbase_client, collection) diff --git a/tests/unit_tests/test_namespace_upsert_reconcile.py b/tests/unit_tests/test_namespace_upsert_reconcile.py new file mode 100644 index 00000000..1a79c9eb --- /dev/null +++ b/tests/unit_tests/test_namespace_upsert_reconcile.py @@ -0,0 +1,79 @@ +""" +Unit tests for namespace upsert duplicate-record reconciliation. +""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock, call + +import pytest + +project_root = Path(__file__).parent.parent.parent +src_root = project_root / "src" +sys.path.insert(0, str(src_root)) + +from pyseekdb.client.client_base import BaseClient # noqa: E402 + + +class TestNamespaceUpsertReconcile: + def test_lock_name_is_stable_and_bounded(self): + name = BaseClient._namespace_record_lock_name("c" * 32, 7, 9, "same_new_id") + assert name.startswith("pyseekdb:nsu:") + assert len(name) <= 64 + + def test_reconcile_skips_when_single_row_exists(self): + client = MagicMock(spec=BaseClient) + client._namespace_record_lock.return_value.__enter__.return_value = True + client._count_namespace_records_by_id.return_value = 1 + + BaseClient._reconcile_namespace_duplicate_records( + client, + collection_id="c" * 32, + collection_name="items", + namespace_id="7", + namespace_name="race_ns", + ltable_id=9, + table_name="logic_data_table", + ids=["same_new_id"], + documents=["doc"], + metadatas=[{"client": 0}], + embeddings=[[1.0, 2.0, 3.0]], + embedding_function=None, + ) + + client._delete_namespace_records_by_id.assert_not_called() + client._namespace_add.assert_not_called() + + def test_reconcile_collapses_duplicate_rows(self): + client = MagicMock(spec=BaseClient) + client._namespace_record_lock.return_value.__enter__.return_value = True + client._count_namespace_records_by_id.return_value = 4 + + BaseClient._reconcile_namespace_duplicate_records( + client, + collection_id="c" * 32, + collection_name="items", + namespace_id="7", + namespace_name="race_ns", + ltable_id=9, + table_name="logic_data_table", + ids=["same_new_id"], + documents=["winner"], + metadatas=[{"client": 2}], + embeddings=[[1.0, 2.0, 3.0]], + embedding_function=None, + ) + + client._delete_namespace_records_by_id.assert_called_once_with( + "logic_data_table", 7, 9, "same_new_id" + ) + client._namespace_add.assert_called_once() + add_kwargs = client._namespace_add.call_args.kwargs + assert add_kwargs["ids"] == ["same_new_id"] + assert add_kwargs["documents"] == ["winner"] + assert add_kwargs["metadatas"] == [{"client": 2}] + assert add_kwargs["embeddings"] == [[1.0, 2.0, 3.0]] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 0be3947eb0bf19f73d18429aa9fdb1f31de520d3 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 22 Jun 2026 21:10:21 +0800 Subject: [PATCH 57/84] =?UTF-8?q?fix:=E5=A4=9A=E4=B8=AAclient=E5=B9=B6?= =?UTF-8?q?=E5=8F=91=E6=89=A7=E8=A1=8Cdelete=5Fcollection=E4=B8=8Ehas=5Fco?= =?UTF-8?q?llection,=E5=87=BA=E7=8E=B0=E4=B8=8D=E4=B8=80=E8=87=B4=20https:?= =?UTF-8?q?//project.alipay.com/project/W24001004732/P26001026476/workitem?= =?UTF-8?q?Detail=3FopenWorkItemId=3D2026062200116895034&status=3Dstatus?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 134 +++++++++++----- .../test_namespace_lifecycle_concurrency.py | 151 ++++++++++++++++++ tests/unit_tests/test_namespace.py | 4 + .../test_namespace_lifecycle_lock.py | 65 ++++++++ 4 files changed, 316 insertions(+), 38 deletions(-) create mode 100644 tests/integration_tests/test_namespace_lifecycle_concurrency.py create mode 100644 tests/unit_tests/test_namespace_lifecycle_lock.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index eb075ecc..cae671fa 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1653,32 +1653,34 @@ def _finalize_ns_namespace_meta( return {"namespace_id": str(namespace_id), "namespace_name": namespace_name, "ltable_id": str(ltable_id)} def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: - ns_id = self._insert_ns_namespace_catalog_row( - collection_id, namespace_name, idempotent=False - ) - lt_id = self._insert_ns_ltable_catalog_row( - collection_id, ns_id, idempotent=False - ) - return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) + with self._namespace_lifecycle_guard(collection_id, namespace_name): + ns_id = self._insert_ns_namespace_catalog_row( + collection_id, namespace_name, idempotent=False + ) + lt_id = self._insert_ns_ltable_catalog_row( + collection_id, ns_id, idempotent=False + ) + return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) def _get_or_create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: - meta = self._get_ns_namespace_meta(collection_id, namespace_name) - if meta is not None: - if meta.get("ltable_id") is not None: - return meta - lt_id = self._insert_ns_ltable_catalog_row( - collection_id, int(meta["namespace_id"]), idempotent=True + with self._namespace_lifecycle_guard(collection_id, namespace_name): + meta = self._get_ns_namespace_meta(collection_id, namespace_name) + if meta is not None: + if meta.get("ltable_id") is not None: + return meta + lt_id = self._insert_ns_ltable_catalog_row( + collection_id, int(meta["namespace_id"]), idempotent=True + ) + return self._finalize_ns_namespace_meta( + collection_id, namespace_name, int(meta["namespace_id"]), lt_id + ) + ns_id = self._insert_ns_namespace_catalog_row( + collection_id, namespace_name, idempotent=True ) - return self._finalize_ns_namespace_meta( - collection_id, namespace_name, int(meta["namespace_id"]), lt_id + lt_id = self._insert_ns_ltable_catalog_row( + collection_id, ns_id, idempotent=True ) - ns_id = self._insert_ns_namespace_catalog_row( - collection_id, namespace_name, idempotent=True - ) - lt_id = self._insert_ns_ltable_catalog_row( - collection_id, ns_id, idempotent=True - ) - return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) + return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: namespace_name_escaped = escape_string(namespace_name) @@ -1717,7 +1719,8 @@ def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dic return meta def _has_ns_namespace(self, collection_id: str, namespace_name: str) -> bool: - return self._get_ns_namespace_meta(collection_id, namespace_name) is not None + with self._namespace_lifecycle_guard(collection_id, namespace_name): + return self._get_ns_namespace_meta(collection_id, namespace_name) is not None def _ns_namespace_exists_by_id(self, collection_id: str, namespace_id: str) -> bool: """Whether a namespace with this id is still live in the catalog. @@ -1738,21 +1741,22 @@ def _ns_namespace_exists_by_id(self, collection_id: str, namespace_id: str) -> b return bool(rows) def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> None: - meta = self._get_ns_namespace_meta(collection_id, namespace_name) - if meta is None: - raise ValueError(f"Namespace '{namespace_name}' not found") - ns_id = meta["namespace_id"] - lt_id_raw = meta.get("ltable_id") - lt_id = int(lt_id_raw) if lt_id_raw is not None else None - collection_id_escaped = escape_string(collection_id) - # PL reads session database_name; must match where catalog tables live. - self._use_catalog_database() - self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(ns_id), ltable_id=lt_id - ) - self._execute( - f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})" - ) + with self._namespace_lifecycle_guard(collection_id, namespace_name): + meta = self._get_ns_namespace_meta(collection_id, namespace_name) + if meta is None: + raise ValueError(f"Namespace '{namespace_name}' not found") + ns_id = meta["namespace_id"] + lt_id_raw = meta.get("ltable_id") + lt_id = int(lt_id_raw) if lt_id_raw is not None else None + collection_id_escaped = escape_string(collection_id) + # PL reads session database_name; must match where catalog tables live. + self._use_catalog_database() + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(ns_id), ltable_id=lt_id + ) + self._execute( + f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})" + ) def _list_ns_namespaces(self, collection_id: str) -> list[dict]: collection_id_escaped = escape_string(collection_id) @@ -4983,6 +4987,60 @@ def _append_namespace_filter( return f"WHERE {ns_cond} AND ({inner})", params return f"WHERE {ns_cond} AND ({where_clause})", params + @staticmethod + def _namespace_lifecycle_lock_name(collection_id: str, namespace_name: str) -> str: + key = f"{collection_id}:{namespace_name}" + digest = hashlib.sha256(key.encode()).hexdigest()[:40] + return f"pyseekdb:nslc:{digest}" + + @contextlib.contextmanager + def _namespace_lifecycle_lock( + self, + collection_id: str, + namespace_name: str, + *, + total_wait_seconds: float = 30.0, + ): + lock_name = self._namespace_lifecycle_lock_name(collection_id, namespace_name) + lock_name_escaped = escape_string(lock_name) + acquired = False + deadline = time.monotonic() + total_wait_seconds + try: + while time.monotonic() < deadline: + rows = self._execute( + f"SELECT GET_LOCK('{lock_name_escaped}', 1) AS got" + ) + row = rows[0] + got = row.get("got") if isinstance(row, dict) else row[0] + if got == 1: + acquired = True + yield True + return + time.sleep(0.05) + yield False + finally: + if acquired: + with contextlib.suppress(Exception): + self._execute(f"SELECT RELEASE_LOCK('{lock_name_escaped}')") + + @contextlib.contextmanager + def _namespace_lifecycle_guard( + self, + collection_id: str, + namespace_name: str, + *, + total_wait_seconds: float = 30.0, + ): + with self._namespace_lifecycle_lock( + collection_id, namespace_name, total_wait_seconds=total_wait_seconds + ) as acquired: + if not acquired: + raise TimeoutError( + f"timed out waiting for namespace lifecycle lock on " + f"{namespace_name!r} in collection {collection_id!r}" + ) + yield + @staticmethod def _namespace_record_lock_name( collection_id: str, diff --git a/tests/integration_tests/test_namespace_lifecycle_concurrency.py b/tests/integration_tests/test_namespace_lifecycle_concurrency.py new file mode 100644 index 00000000..cceffcff --- /dev/null +++ b/tests/integration_tests/test_namespace_lifecycle_concurrency.py @@ -0,0 +1,151 @@ +"""Concurrent namespace lifecycle tests (drop vs has across clients).""" + +from __future__ import annotations + +import contextlib +import os +import threading +import time +from unittest.mock import patch + +import pytest + +import pyseekdb +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, cleanup, ns_schema + + +def _unique_name(prefix: str) -> str: + return f"{prefix}_{time.time_ns()}" + + +def _new_oceanbase_client(): + return pyseekdb.Client( + host=os.environ.get("OB_HOST", "127.0.0.1"), + port=int(os.environ.get("OB_PORT", "10902")), + tenant=os.environ.get("OB_TENANT", "mysql"), + database=os.environ.get("OB_DATABASE", "test"), + user=os.environ.get("OB_USER", "root"), + password=os.environ.get("OB_PASSWORD", ""), + ) + + +def test_multi_client_has_false_after_drop_returns(oceanbase_client): + """After delete_namespace returns, has_namespace on another client is always False.""" + collection = oceanbase_client.create_collection( + name=_unique_name("test_ns_lc_drop_has"), + schema=ns_schema(), + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + clients = [_new_oceanbase_client() for _ in range(2)] + ns_name = "drop_has_ns" + drop_done = threading.Event() + errors: list[Exception] = [] + post_drop_has: list[bool] = [] + + try: + collection.create_namespace(ns_name) + + def _drop_worker() -> None: + try: + clients[0].get_collection(collection.name).delete_namespace(ns_name) + except Exception as exc: # noqa: BLE001 + errors.append(exc) + finally: + drop_done.set() + + def _has_worker() -> None: + try: + coll = clients[1].get_collection(collection.name) + drop_done.wait(timeout=60) + for _ in range(30): + post_drop_has.append(coll.has_namespace(ns_name)) + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + drop_thread = threading.Thread(target=_drop_worker) + has_thread = threading.Thread(target=_has_worker) + drop_thread.start() + has_thread.start() + drop_thread.join(timeout=120) + has_thread.join(timeout=120) + + assert errors == [] + assert drop_thread.is_alive() is False + assert has_thread.is_alive() is False + assert post_drop_has + assert all(value is False for value in post_drop_has) + finally: + for client in clients: + with contextlib.suppress(Exception): + client.close() + cleanup(oceanbase_client, collection) + + +def test_multi_client_has_blocks_until_drop_finishes(oceanbase_client): + """has_namespace on client B waits for client A drop to release lifecycle lock.""" + collection = oceanbase_client.create_collection( + name=_unique_name("test_ns_lc_drop_block"), + schema=ns_schema(), + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + clients = [_new_oceanbase_client() for _ in range(2)] + ns_name = "block_ns" + drop_started = threading.Event() + drop_can_finish = threading.Event() + has_started = threading.Event() + errors: list[Exception] = [] + has_result: list[bool] = [] + + try: + collection.create_namespace(ns_name) + dropper = clients[0].get_collection(collection.name) + checker = clients[1].get_collection(collection.name) + real_execute = clients[0]._server._execute + + def _slow_drop_execute(sql, *args, **kwargs): + if "DROP_NAMESPACE" in str(sql): + drop_started.set() + assert drop_can_finish.wait(timeout=30), "drop blocked too long" + return real_execute(sql, *args, **kwargs) + + def _drop_worker() -> None: + try: + with patch.object(clients[0]._server, "_execute", side_effect=_slow_drop_execute): + dropper.delete_namespace(ns_name) + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + def _has_worker() -> None: + try: + assert drop_started.wait(timeout=30), "drop did not start" + has_started.set() + has_result.append(checker.has_namespace(ns_name)) + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + drop_thread = threading.Thread(target=_drop_worker) + has_thread = threading.Thread(target=_has_worker) + drop_thread.start() + has_thread.start() + + assert has_started.wait(timeout=30), "has did not start while drop in progress" + time.sleep(0.2) + assert drop_thread.is_alive(), "drop finished before lifecycle lock blocked has" + drop_can_finish.set() + + drop_thread.join(timeout=120) + has_thread.join(timeout=120) + + assert errors == [] + assert has_result == [False] + finally: + for client in clients: + with contextlib.suppress(Exception): + client.close() + cleanup(oceanbase_client, collection) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 391fa234..b2a5d41a 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -1145,6 +1145,7 @@ class TestDeleteNamespaceUsesKernel: def test_delete_namespace_calls_dbms_logic_table(self): c = FakeClient() c._execute = MagicMock(side_effect=[ + [{"got": 1}], # GET_LOCK [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], None, # _get_ns_namespace_meta -> SET @namespace_id None, # _get_ns_namespace_meta -> SET @ltable_id @@ -1153,9 +1154,12 @@ def test_delete_namespace_calls_dbms_logic_table(self): None, # _delete_ns_namespace_meta -> SET @namespace_id None, # _delete_ns_namespace_meta -> SET @ltable_id None, # CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE + [{"got": 1}], # RELEASE_LOCK ]) c._delete_ns_namespace_meta("abc123", "ns1") calls = [str(call) for call in c._execute.call_args_list] + assert any("GET_LOCK" in s and "pyseekdb:nslc:" in s for s in calls) + assert any("RELEASE_LOCK" in s for s in calls) assert any("DBMS_LOGIC_TABLE.DROP_NAMESPACE" in s for s in calls) # session context must be set BEFORE DROP_NAMESPACE so the kernel sees the # right @collection_id / @namespace_id / @ltable_id for this call. diff --git a/tests/unit_tests/test_namespace_lifecycle_lock.py b/tests/unit_tests/test_namespace_lifecycle_lock.py new file mode 100644 index 00000000..7824f03d --- /dev/null +++ b/tests/unit_tests/test_namespace_lifecycle_lock.py @@ -0,0 +1,65 @@ +"""Unit tests for namespace lifecycle GET_LOCK serialization.""" + +import contextlib +import sys +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +project_root = Path(__file__).parent.parent.parent +src_root = project_root / "src" +sys.path.insert(0, str(src_root)) + +from pyseekdb.client.client_base import BaseClient # noqa: E402 + + +class TestNamespaceLifecycleLockName: + def test_lock_name_is_stable_and_bounded(self): + name_a = BaseClient._namespace_lifecycle_lock_name("cid", "ns1") + name_b = BaseClient._namespace_lifecycle_lock_name("cid", "ns1") + name_c = BaseClient._namespace_lifecycle_lock_name("cid", "ns2") + assert name_a == name_b + assert name_a != name_c + assert name_a.startswith("pyseekdb:nslc:") + assert len(name_a) <= 64 + + +class TestNamespaceLifecycleGuard: + def test_has_namespace_acquires_and_releases_lock(self): + client = MagicMock(spec=BaseClient) + client._get_ns_namespace_meta.return_value = None + + @contextlib.contextmanager + def _lock_cm(*_args, **_kwargs): + yield True + + client._namespace_lifecycle_lock.side_effect = _lock_cm + client._namespace_lifecycle_guard = ( + BaseClient._namespace_lifecycle_guard.__get__(client, BaseClient) + ) + + assert BaseClient._has_ns_namespace(client, "cid", "ns1") is False + client._namespace_lifecycle_lock.assert_called_once_with( + "cid", "ns1", total_wait_seconds=30.0 + ) + client._get_ns_namespace_meta.assert_called_once_with("cid", "ns1") + + def test_guard_raises_when_lock_times_out(self): + client = MagicMock(spec=BaseClient) + + @contextlib.contextmanager + def _no_lock(): + yield False + + client._namespace_lifecycle_lock.return_value = _no_lock() + client._namespace_lifecycle_guard = ( + BaseClient._namespace_lifecycle_guard.__get__(client, BaseClient) + ) + + with pytest.raises(TimeoutError, match="namespace lifecycle lock"): + BaseClient._has_ns_namespace(client, "cid", "ns1") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 8b50723e81c1877a4506913614e0b9e8420c1bde Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 23 Jun 2026 14:00:51 +0800 Subject: [PATCH 58/84] fix:issue_fix --- src/pyseekdb/client/client_base.py | 21 +++++++-- src/pyseekdb/client/collection.py | 12 +++++ src/pyseekdb/client/configuration.py | 47 +++++++++++++++++++ src/pyseekdb/client/meta_info.py | 20 +++++++- src/pyseekdb/client/namespace.py | 20 ++++++++ src/pyseekdb/client/validators.py | 8 ++++ tests/integration_tests/Untitled | 1 - ...st_get_or_create_collection_concurrency.py | 37 +++++++++++++++ tests/unit_tests/test_namespace.py | 5 ++ 9 files changed, 166 insertions(+), 5 deletions(-) delete mode 100644 tests/integration_tests/Untitled diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index cae671fa..c5cd6a82 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1762,7 +1762,9 @@ def _list_ns_namespaces(self, collection_id: str) -> list[dict]: collection_id_escaped = escape_string(collection_id) rows = self._execute( f"SELECT namespace_id, namespace_name FROM {self._qtable(NamespaceCollectionNames.sdk_namespaces_table())} " - f"WHERE collection_id = '{collection_id_escaped}' ORDER BY namespace_id" + f"WHERE collection_id = '{collection_id_escaped}' " + f"AND LEFT(namespace_name, 13) <> '__recyclebin_' " + f"ORDER BY namespace_id" ) results = [] for row in rows: @@ -2318,7 +2320,11 @@ def _collection_table_exists(self, table_name: str) -> bool: def _has_collection_v2(self, name: str) -> bool: try: - query_sql = f"SELECT COLLECTION_ID FROM {CollectionNames.sdk_collections_table_name()} WHERE COLLECTION_NAME = '{name}'" + name_escaped = escape_string(name) + query_sql = ( + f"SELECT COLLECTION_ID FROM {CollectionNames.sdk_collections_table_name()} " + f"WHERE COLLECTION_NAME = '{name_escaped}'" + ) rows = self._execute(query_sql) if not rows or len(rows) == 0: return False @@ -2392,6 +2398,15 @@ def get_or_create_collection( _validate_collection_name(name) if self.has_collection(name): + if use_namespace and self._is_incomplete_ns_collection(name): + return self.create_collection( + name=name, + schema=schema, + configuration=configuration, + embedding_function=embedding_function, + use_namespace=use_namespace, + **kwargs, + ) return self.get_collection(name, embedding_function=embedding_function) try: @@ -2404,7 +2419,7 @@ def get_or_create_collection( **kwargs, ) except Exception as exc: - if _is_collection_conflict_error(exc) or self.has_collection(name): + if _is_collection_conflict_error(exc): return self.get_collection(name, embedding_function=embedding_function) raise diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 1ac5e8a7..5c38157b 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -45,6 +45,7 @@ def __init__( partition_count: int | None = None, **metadata, ): + """Initialize a lightweight collection handle bound to a client implementation.""" self._client = client self._name = name self._id = collection_id @@ -107,6 +108,7 @@ def sparse_embedding_function(self) -> Optional["SparseEmbeddingFunction"]: @property def use_namespace(self) -> bool: + """Whether this collection routes data operations through namespaces.""" return self._use_namespace @property @@ -121,6 +123,7 @@ def partition_count(self) -> int | None: @staticmethod def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> None: + """Validate ``n_results`` is a positive integer within the engine limit.""" if not isinstance(n_results, int) or n_results < 1: raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") if n_results > max_results: @@ -130,6 +133,7 @@ def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> ) def _guard_collection_data_api(self) -> None: + """Raise if collection-level DML/DQL is used on a namespace-enabled collection.""" if self._use_namespace: raise ValueError( "This collection has namespace enabled. " @@ -138,6 +142,7 @@ def _guard_collection_data_api(self) -> None: ) def _guard_namespace_enabled(self) -> None: + """Raise if namespace APIs are used on a non-namespace collection.""" if not self._use_namespace: raise ValueError("Namespace is not enabled for this collection. Use use_namespace=True when creating the collection.") if not self._client._ns_collection_exists_by_id(self._id): @@ -147,12 +152,14 @@ def _guard_namespace_enabled(self) -> None: ) def __repr__(self) -> str: + """Return a debug-friendly representation of this collection.""" ns_str = ", use_namespace=True" if self._use_namespace else "" return f"Collection(name='{self._name}', dimension={self._dimension}{ns_str}, client={self._client.mode})" # ==================== Namespace Management ==================== def create_namespace(self, name: str) -> "Namespace": + """Create a new namespace and return a handle for scoped operations.""" self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace @@ -160,6 +167,7 @@ def create_namespace(self, name: str) -> "Namespace": return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) def get_namespace(self, name: str) -> "Namespace": + """Return an existing namespace handle or raise if it does not exist.""" self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace @@ -169,6 +177,7 @@ def get_namespace(self, name: str) -> "Namespace": return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) def get_or_create_namespace(self, name: str) -> "Namespace": + """Return an existing namespace or create it if missing.""" self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace @@ -176,11 +185,13 @@ def get_or_create_namespace(self, name: str) -> "Namespace": return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) def delete_namespace(self, name: str) -> None: + """Delete a namespace and its records from this collection.""" self._guard_namespace_enabled() _validate_namespace_name(name) self._client._delete_ns_namespace_meta(self._id, name) def list_namespaces(self) -> list["Namespace"]: + """List all active namespaces in this collection.""" self._guard_namespace_enabled() from .namespace import Namespace metas = self._client._list_ns_namespaces(self._id) @@ -190,6 +201,7 @@ def list_namespaces(self) -> list["Namespace"]: ] def has_namespace(self, name: str) -> bool: + """Return whether a namespace with the given name exists.""" self._guard_namespace_enabled() _validate_namespace_name(name) return self._client._has_ns_namespace(self._id, name) diff --git a/src/pyseekdb/client/configuration.py b/src/pyseekdb/client/configuration.py index 3c683c60..8dec0c28 100644 --- a/src/pyseekdb/client/configuration.py +++ b/src/pyseekdb/client/configuration.py @@ -1,3 +1,5 @@ +"""Index and analyzer configuration types for collection and schema creation.""" + import warnings from dataclasses import dataclass from enum import Enum @@ -16,6 +18,7 @@ def _ensure_primitive_properties(properties: dict[str, Any] | None, *, field_name: str = "properties") -> None: + """Ensure index property values are primitive JSON-compatible scalars.""" if properties is None: return if not isinstance(properties, dict): @@ -26,6 +29,7 @@ def _ensure_primitive_properties(properties: dict[str, Any] | None, *, field_nam def _normalize_str_enum(value: str | Enum, *, field_name: str) -> str: + """Normalize enum or string values to lowercase strings.""" if isinstance(value, Enum): value = value.value if not isinstance(value, str): @@ -34,6 +38,7 @@ def _normalize_str_enum(value: str | Enum, *, field_name: str) -> str: def _validate_int_range(value: Any, *, key: str, min_value: int, max_value: int) -> None: + """Validate an integer option is within the allowed inclusive range.""" if isinstance(value, bool) or not isinstance(value, int): raise TypeError(f"{key} must be an integer, got {type(value).__name__}") if value < min_value or value > max_value: @@ -41,6 +46,7 @@ def _validate_int_range(value: Any, *, key: str, min_value: int, max_value: int) def _validate_float_range(value: Any, *, key: str, min_value: float, max_value: float) -> None: + """Validate a numeric option is within the allowed inclusive range.""" if isinstance(value, bool) or not isinstance(value, (int, float)): raise TypeError(f"{key} must be a number, got {type(value).__name__}") numeric_value = float(value) @@ -49,6 +55,7 @@ def _validate_float_range(value: Any, *, key: str, min_value: float, max_value: def _validate_space_or_beng_properties(properties: dict[str, PrimitiveValue]) -> None: + """Validate tokenizer size bounds for space and beng analyzers.""" if not properties: return min_token_size = properties.get("min_token_size") @@ -65,6 +72,7 @@ def _validate_space_or_beng_properties(properties: dict[str, PrimitiveValue]) -> def _validate_ngram_properties(properties: dict[str, PrimitiveValue]) -> None: + """Validate ngram analyzer token size bounds.""" if not properties: return if "ngram_token_size" in properties: @@ -72,6 +80,7 @@ def _validate_ngram_properties(properties: dict[str, PrimitiveValue]) -> None: def _validate_ngram2_properties(properties: dict[str, PrimitiveValue]) -> None: + """Validate ngram2 analyzer min/max ngram size bounds.""" if not properties: return min_ngram_size = properties.get("min_ngram_size") @@ -88,6 +97,7 @@ def _validate_ngram2_properties(properties: dict[str, PrimitiveValue]) -> None: def _normalize_ik_mode(properties: dict[str, PrimitiveValue]) -> None: + """Normalize and validate IK analyzer mode values.""" if not properties: return if "ik_mode" not in properties: @@ -104,6 +114,7 @@ def _normalize_ik_mode(properties: dict[str, PrimitiveValue]) -> None: def _validate_fulltext_properties_by_analyzer(analyzer: str, properties: dict[str, PrimitiveValue] | None) -> None: + """Validate fulltext analyzer properties for the selected analyzer type.""" if not properties: return if analyzer in {FulltextAnalyzer.SPACE.value, FulltextAnalyzer.BENG.value}: @@ -117,6 +128,7 @@ def _validate_fulltext_properties_by_analyzer(analyzer: str, properties: dict[st def _validate_hnsw_base_fields(config: "HNSWConfiguration") -> None: + """Validate core HNSW index fields shared across index subtypes.""" if isinstance(config.dimension, bool) or not isinstance(config.dimension, int): raise TypeError(f"dimension must be an integer, got {type(config.dimension).__name__}") _validate_int_range(config.dimension, key="dimension", min_value=1, max_value=4096) @@ -138,6 +150,7 @@ def _validate_hnsw_base_fields(config: "HNSWConfiguration") -> None: def _normalize_hnsw_properties(properties: dict[str, PrimitiveValue]) -> None: + """Strip reserved HNSW keys from user-supplied property bags.""" reserved_keys = [ key for key in properties @@ -162,6 +175,7 @@ def _normalize_hnsw_properties(properties: dict[str, PrimitiveValue]) -> None: def _validate_hnsw_configuration(config: "HNSWConfiguration") -> None: + """Validate optional HNSW tuning parameters.""" if config.M is not None: _validate_int_range(config.M, key="M", min_value=5, max_value=128) if config.ef_construction is not None: @@ -196,27 +210,37 @@ class DistanceMetric(str, Enum): class HNSWIndexType(str, Enum): + """Supported HNSW index subtypes.""" + HNSW = "hnsw" HNSW_SQ = "hnsw_sq" HNSW_BQ = "hnsw_bq" class HNSWIndexLib(str, Enum): + """Supported HNSW index libraries.""" + VSAG = "vsag" class IVFIndexType(str, Enum): + """Supported IVF index subtypes for namespace-enabled collections.""" + IVF_FLAT = "ivf_flat" IVF_SQ8 = "ivf_sq8" IVF_PQ = "ivf_pq" class IVFIndexLib(str, Enum): + """Supported IVF index libraries.""" + OB = "ob" VSAG = "vsag" class FulltextAnalyzer(str, Enum): + """Supported fulltext analyzers.""" + SPACE = "space" NGRAM = "ngram" BENG = "beng" @@ -225,11 +249,15 @@ class FulltextAnalyzer(str, Enum): class IKMode(str, Enum): + """Supported IK analyzer segmentation modes.""" + SMART = "smart" MAX_WORD = "max_word" class BQRefineType(str, Enum): + """Supported binary-quantization refine types for HNSW BQ indexes.""" + SQ8 = "sq8" FP32 = "fp32" @@ -248,6 +276,7 @@ class FulltextIndexConfig: properties: dict[str, PrimitiveValue] | None = None def __post_init__(self): + """Validate analyzer name and analyzer-specific properties.""" self.analyzer = _normalize_str_enum(self.analyzer, field_name="analyzer") _ensure_primitive_properties(self.properties) @@ -294,6 +323,7 @@ class HNSWConfiguration: properties: dict[str, PrimitiveValue] | None = None def __post_init__(self): + """Validate HNSW configuration fields and normalize properties.""" _validate_hnsw_base_fields(self) _validate_hnsw_configuration(self) @@ -305,24 +335,34 @@ def __post_init__(self): class IKProperties(TypedDict, total=False): + """Typed fulltext properties for the IK analyzer.""" + ik_mode: str | IKMode class SpaceProperties(TypedDict, total=False): + """Typed fulltext properties for the space analyzer.""" + min_token_size: int max_token_size: int class NgramProperties(TypedDict, total=False): + """Typed fulltext properties for the ngram analyzer.""" + ngram_token_size: int class Ngram2Properties(TypedDict, total=False): + """Typed fulltext properties for the ngram2 analyzer.""" + min_ngram_size: int max_ngram_size: int class BengProperties(TypedDict, total=False): + """Typed fulltext properties for the beng analyzer.""" + min_token_size: int max_token_size: int @@ -348,6 +388,7 @@ class IVFConfiguration: properties: dict[str, PrimitiveValue] | None = None def __post_init__(self): + """Validate IVF configuration fields and normalize properties.""" if isinstance(self.dimension, bool) or not isinstance(self.dimension, int): raise TypeError(f"dimension must be an integer, got {type(self.dimension).__name__}") _validate_int_range(self.dimension, key="dimension", min_value=1, max_value=4096) @@ -377,11 +418,14 @@ def __post_init__(self): @dataclass class VectorIndexConfig: + """Dense vector index configuration wrapper for HNSW or IVF indexes.""" + ivf: IVFConfiguration | None = None hnsw: HNSWConfiguration | None = None embedding_function: EmbeddingFunction | None = _NOT_PROVIDED def __post_init__(self): + """Resolve defaults and validate the selected dense vector index.""" if self.ivf is not None and self.hnsw is not None: raise ValueError("Only one of ivf or hnsw can be configured") if self.embedding_function is _NOT_PROVIDED: @@ -449,6 +493,7 @@ class SparseVectorIndexConfig: properties: dict[str, PrimitiveValue] | None = None def __post_init__(self): # noqa: C901 + """Validate sparse vector index options and embedding function persistence.""" if self.distance != DistanceMetric.INNER_PRODUCT.value: raise ValueError( f"Sparse vector index only supports {DistanceMetric.INNER_PRODUCT.value} distance, got '{self.distance}'" @@ -481,6 +526,7 @@ def __post_init__(self): # noqa: C901 ) def _validate_source_key(self) -> None: + """Normalize and validate the sparse vector source field.""" if self.source_key is None: self.source_key = K.DOCUMENT return @@ -528,6 +574,7 @@ def __init__( hnsw: HNSWConfiguration | None = None, fulltext_config: FulltextIndexConfig | None = None, ): + """Initialize deprecated collection configuration wrapper.""" self.hnsw = hnsw self.fulltext_config = fulltext_config warnings.warn("Configuration is deprecated. Please use Schema instead.", DeprecationWarning, stacklevel=2) diff --git a/src/pyseekdb/client/meta_info.py b/src/pyseekdb/client/meta_info.py index 4523754b..78df4ce2 100644 --- a/src/pyseekdb/client/meta_info.py +++ b/src/pyseekdb/client/meta_info.py @@ -1,11 +1,13 @@ """ -Metadata information for collection fields. +Metadata information for collection fields and namespace catalog naming. """ from typing import ClassVar class CollectionFieldNames: + """Standard field names used in collection record payloads.""" + ID = "_id" DOCUMENT = "document" EMBEDDING = "embedding" @@ -16,6 +18,8 @@ class CollectionFieldNames: class CollectionNames: + """Helpers for mapping between collection names and physical table names.""" + # Version prefix for collection tables _PREFIX = "c$v1$" _PREFIX_V2 = "c$v2$" @@ -54,10 +58,13 @@ def prefix() -> str: @staticmethod def sdk_collections_table_name() -> str: + """Return the SDK catalog table that stores collection metadata.""" return "sdk_collections" class NamespaceCollectionNames: + """Naming helpers for namespace-enabled collection physical tables.""" + _LOGIC_DATA_SUFFIX = "_logic_data_table" _HOT_SUFFIX = "_hot_table" _KV_DATA_SUFFIX = "_kv_data_table" @@ -66,42 +73,53 @@ class NamespaceCollectionNames: @staticmethod def sdk_namespaces_table() -> str: + """Return the SDK catalog table that stores namespace metadata.""" return "sdk_namespaces" @staticmethod def sdk_ltables_table() -> str: + """Return the SDK catalog table that stores logic-table metadata.""" return "sdk_ltables" @staticmethod def sdk_namespaces_stats_table() -> str: + """Return the SDK catalog table that stores namespace statistics.""" return "sdk_namespaces_stats" @staticmethod def data_table_name(collection_id: str) -> str: + """Build the logic data table name for a namespace-enabled collection.""" return f"{collection_id}{NamespaceCollectionNames._LOGIC_DATA_SUFFIX}" @staticmethod def hot_table_name(collection_id: str) -> str: + """Build the hot data table name for a namespace-enabled collection.""" return f"{collection_id}{NamespaceCollectionNames._HOT_SUFFIX}" @staticmethod def kv_data_table_name(collection_id: str) -> str: + """Build the KV data table name for a namespace-enabled collection.""" return f"{collection_id}{NamespaceCollectionNames._KV_DATA_SUFFIX}" @staticmethod def logic_schema_table_name(collection_id: str) -> str: + """Build the logic schema table name for a namespace-enabled collection.""" return f"{collection_id}{NamespaceCollectionNames._LOGIC_SCHEMA_SUFFIX}" @staticmethod def tablegroup_name(collection_id: str) -> str: + """Build the table group name for a namespace-enabled collection.""" return f"{collection_id}{NamespaceCollectionNames._TG_SUFFIX}" @staticmethod def is_ns_data_table(table_name: str) -> bool: + """Return True if ``table_name`` is a namespace logic data table.""" return table_name.endswith(NamespaceCollectionNames._LOGIC_DATA_SUFFIX) class NamespaceFieldNames: + """Standard field names used in namespace record payloads.""" + NAMESPACE_ID = "namespace_id" LTABLE_ID = "ltable_id" DOCUMENT = "document" diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py index c95be242..3e726bef 100644 --- a/src/pyseekdb/client/namespace.py +++ b/src/pyseekdb/client/namespace.py @@ -17,6 +17,8 @@ class Namespace: + """Scoped view of a single namespace within a namespace-enabled collection.""" + def __init__( self, client: Any, @@ -24,6 +26,7 @@ def __init__( name: str, namespace_id: str, ): + """Bind a namespace handle to its parent collection and catalog identifiers.""" self._client = client self._collection = collection self._name = name @@ -31,27 +34,33 @@ def __init__( @property def name(self) -> str: + """Human-readable namespace name.""" return self._name @property def namespace_id(self) -> str: + """Stable namespace identifier assigned by the catalog.""" return self._namespace_id @property def collection(self) -> "Collection": + """Parent collection that owns this namespace.""" return self._collection @property def embedding_function(self) -> "EmbeddingFunction | None": + """Embedding function inherited from the parent collection.""" return self._collection.embedding_function def __repr__(self) -> str: + """Return a debug-friendly representation of this namespace.""" return ( f"Namespace(name='{self._name}', namespace_id={self._namespace_id}, " f"collection='{self._collection.name}')" ) def _guard_exists(self) -> None: + """Raise if this namespace or its collection was deleted.""" if not self._client._ns_namespace_exists_by_id(self._collection.id, self._namespace_id): raise ValueError( f"Namespace '{self._name}' no longer exists (it or its collection may have been deleted). " @@ -60,6 +69,7 @@ def _guard_exists(self) -> None: @staticmethod def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> None: + """Validate ``n_results`` is a positive integer within the engine limit.""" if not isinstance(n_results, int) or n_results < 1: raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") if n_results > max_results: @@ -78,6 +88,7 @@ def add( documents: str | list[str] | None = None, **kwargs, ) -> None: + """Insert records into this namespace.""" self._guard_exists() return self._client._namespace_add( collection_id=self._collection.id, @@ -100,6 +111,7 @@ def update( documents: str | list[str] | None = None, **kwargs, ) -> None: + """Update existing records in this namespace.""" self._guard_exists() return self._client._namespace_update( collection_id=self._collection.id, @@ -122,6 +134,7 @@ def upsert( documents: str | list[str] | None = None, **kwargs, ) -> None: + """Insert or update records in this namespace.""" self._guard_exists() return self._client._namespace_upsert( collection_id=self._collection.id, @@ -143,6 +156,7 @@ def delete( where_document: dict[str, Any] | None = None, **kwargs, ) -> None: + """Delete records from this namespace by ids or filters.""" self._guard_exists() return self._client._namespace_delete( collection_id=self._collection.id, @@ -167,6 +181,7 @@ def query( include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + """Run vector similarity search within this namespace.""" self._guard_exists() self._validate_n_results(n_results) return self._client._namespace_query( @@ -194,6 +209,7 @@ def hybrid_search( include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + """Run hybrid fulltext + vector search within this namespace.""" self._guard_exists() self._validate_n_results(n_results) if include is None and not query and not knn: @@ -223,6 +239,7 @@ def get( include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + """Fetch records from this namespace by ids or filters.""" self._guard_exists() return self._client._namespace_get( collection_id=self._collection.id, @@ -239,6 +256,7 @@ def get( ) def count(self) -> int: + """Return the number of records in this namespace.""" self._guard_exists() return self._client._namespace_count( collection_id=self._collection.id, @@ -248,6 +266,7 @@ def count(self) -> int: ) def peek(self, limit: int = 10) -> dict[str, Any]: + """Return up to ``limit`` records from this namespace.""" self._guard_exists() return self._client._namespace_peek( collection_id=self._collection.id, @@ -258,6 +277,7 @@ def peek(self, limit: int = 10) -> dict[str, Any]: ) def prewarm(self) -> None: + """Prewarm namespace physical tables to reduce cold-start latency.""" self._guard_exists() return self._client._namespace_prewarm( collection_id=self._collection.id, diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py index 7a6eb253..8c7515af 100644 --- a/src/pyseekdb/client/validators.py +++ b/src/pyseekdb/client/validators.py @@ -1,3 +1,5 @@ +"""Input validators for namespace names and record identifiers.""" + import re _NAME_PATTERN = re.compile(r"^[A-Za-z0-9_]+$") @@ -8,6 +10,7 @@ def _validate_namespace_name(name: str) -> None: + """Validate a namespace name for type, length, and allowed characters.""" if not isinstance(name, str): raise TypeError( f"Invalid namespace name: '{name}'. Namespace name must be a string, got {type(name).__name__}" @@ -26,6 +29,11 @@ def _validate_namespace_name(name: str) -> None: def _validate_record_ids(ids: list[str]) -> None: + """Validate namespace/collection record id list shape and per-id constraints.""" + if not isinstance(ids, list): + raise TypeError( + f"Invalid record ids: expected list[str], got {type(ids).__name__}" + ) for rid in ids: if not isinstance(rid, str): raise TypeError( diff --git a/tests/integration_tests/Untitled b/tests/integration_tests/Untitled deleted file mode 100644 index 4a0ec5c6..00000000 --- a/tests/integration_tests/Untitled +++ /dev/null @@ -1 +0,0 @@ -test_collection_data_api_blocked_when_namespace_enabled[ \ No newline at end of file diff --git a/tests/unit_tests/test_get_or_create_collection_concurrency.py b/tests/unit_tests/test_get_or_create_collection_concurrency.py index 370456da..00a963d7 100644 --- a/tests/unit_tests/test_get_or_create_collection_concurrency.py +++ b/tests/unit_tests/test_get_or_create_collection_concurrency.py @@ -73,6 +73,43 @@ def test_retries_get_after_wrapped_table_conflict(self): assert result is existing client.get_collection.assert_called_once_with("items", embedding_function=_NOT_PROVIDED) + def test_resumes_incomplete_namespace_collection_when_present(self): + client = MagicMock(spec=BaseClient) + resumed = object() + client.has_collection.return_value = True + client._is_incomplete_ns_collection.return_value = True + client.create_collection.return_value = resumed + + result = BaseClient.get_or_create_collection(client, "items", use_namespace=True) + + assert result is resumed + client.create_collection.assert_called_once() + client.get_collection.assert_not_called() + + def test_does_not_mask_unrelated_create_errors(self): + client = MagicMock(spec=BaseClient) + client.has_collection.return_value = False + client.create_collection.side_effect = ValueError("invalid dimension") + + with pytest.raises(ValueError, match="invalid dimension"): + BaseClient.get_or_create_collection(client, "items") + + +class TestListNsNamespacesRecyclebinFilter: + def test_sql_excludes_recyclebin_rows(self): + client = MagicMock(spec=BaseClient) + client._qtable.return_value = "`sdk_namespaces`" + client._execute.return_value = [ + ("1", "active_ns"), + ] + + result = BaseClient._list_ns_namespaces(client, "coll_1") + + assert result == [{"namespace_id": "1", "namespace_name": "active_ns"}] + sql = client._execute.call_args[0][0] + assert "__recyclebin_" in sql + assert "LEFT(namespace_name, 13) <> '__recyclebin_'" in sql + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index b2a5d41a..46bf9ed1 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -921,6 +921,11 @@ def test_mixed_valid_and_invalid_raises_on_first_bad(self): with pytest.raises(ValueError, match="invalid characters"): _validate_record_ids(["good_id", "bad-id"]) + def test_non_list_ids_raises(self): + from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(TypeError, match="expected list\\[str\\]"): + _validate_record_ids("doc_1") + # ==================== Namespace Batch Limit Tests ==================== From 444d4ba37df63f509ce1922f0ea2c8b312098117 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 23 Jun 2026 16:34:31 +0800 Subject: [PATCH 59/84] =?UTF-8?q?fix:=20[pyseekdb]=20=E5=A4=9A=E4=B8=AAcli?= =?UTF-8?q?ent=E5=B9=B6=E5=8F=91=E8=B0=83=E7=94=A8get=5For=5Fcreate=5Fcoll?= =?UTF-8?q?ection=E6=8E=A5=E5=8F=A3=EF=BC=8C=E5=9C=A8sdk=5Fcollection?= =?UTF-8?q?=E5=AD=98=E5=9C=A8=E5=A4=9A=E4=B8=AA=E5=90=8C=E5=90=8Dcollectio?= =?UTF-8?q?n?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit https://project.alipay.com/project/W24001004732/P26001026476/workitemDetail?openWorkItemId=2026062300116920158&status=status --- src/pyseekdb/client/client_base.py | 60 ++++++++++-- ...st_get_or_create_collection_concurrency.py | 92 +++++++++++++++++++ ...t_get_or_create_collection_multiprocess.py | 10 ++ .../test_logic_table_monitoring_p1.py | 3 - ...st_get_or_create_collection_concurrency.py | 53 +++++++++++ 5 files changed, 208 insertions(+), 10 deletions(-) create mode 100644 tests/integration_tests/test_get_or_create_collection_concurrency.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index c5cd6a82..7232a38d 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -73,7 +73,7 @@ def _extract_collection_id_from_sdk_row(row: Any) -> str: if isinstance(row, dict): - collection_id = row.get("COLLECTION_ID", "") + collection_id = row.get("COLLECTION_ID") or row.get("collection_id") or "" elif isinstance(row, (tuple, list)): collection_id = row[0] if len(row) > 0 else "" else: @@ -112,6 +112,26 @@ def _is_namespace_catalog_conflict_error(exc: BaseException) -> bool: return False +def _is_sdk_collection_catalog_conflict_error(exc: BaseException) -> bool: + current: BaseException | None = exc + while current is not None: + message = str(current).lower() + if ( + "duplicate entry" in message + and ( + "uk_sdk_coll_name" in message + or "idx_name" in message + or "collection_name" in message + or "code=1062" in message + ) + ): + return True + if type(current).__name__ == "IntegrityError" and "1062" in message: + return True + current = current.__cause__ + return False + + def _extract_hnsw_config(config: ConfigurationParam) -> HNSWConfiguration | None: if config is None: return None @@ -1142,9 +1162,13 @@ def _create_sdk_collections_if_not_exists(self) -> None: settings JSON COMMENT "Generated by SDK, don't modify", created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, - INDEX idx_name(collection_name) + UNIQUE KEY uk_sdk_coll_name (collection_name) ) COMMENT='Settings of collections created by SDK' ORGANIZATION INDEX {scp};""" self._execute(create_table_sql) + with contextlib.suppress(Exception): + self._execute( + f"CREATE UNIQUE INDEX uk_sdk_coll_name ON {sdk_coll} (collection_name)" + ) except Exception as e: raise ValueError(f"Failed to create sdk_collections table: {e}") from e @@ -1189,7 +1213,15 @@ def _create_collection_meta_v2( f"INSERT INTO `{CollectionNames.sdk_collections_table_name()}` " f"(COLLECTION_NAME, SETTINGS) VALUES ('{collection_name_in_table}', '{settings_str}')" ) - self._execute(insert_sql) + try: + self._execute(insert_sql) + except Exception as exc: + if not _is_sdk_collection_catalog_conflict_error(exc): + raise + conn_getter = getattr(self, "_ensure_connection", None) + if conn_getter is not None: + with contextlib.suppress(Exception): + conn_getter().rollback() collection_id = self._get_collection_id(collection_name) results["collection_id"] = collection_id @@ -1278,7 +1310,15 @@ def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> di f"(collection_name, settings) " f"VALUES ('{collection_name_escaped}', '{settings_str}')" ) - self._execute(insert_sql) + try: + self._execute(insert_sql) + except Exception as exc: + if not _is_sdk_collection_catalog_conflict_error(exc): + raise + conn_getter = getattr(self, "_ensure_connection", None) + if conn_getter is not None: + with contextlib.suppress(Exception): + conn_getter().rollback() rows = self._execute( f"SELECT collection_id FROM {sdk_coll} " f"WHERE collection_name = '{collection_name_escaped}'" @@ -2322,8 +2362,9 @@ def _has_collection_v2(self, name: str) -> bool: try: name_escaped = escape_string(name) query_sql = ( - f"SELECT COLLECTION_ID FROM {CollectionNames.sdk_collections_table_name()} " - f"WHERE COLLECTION_NAME = '{name_escaped}'" + f"SELECT collection_id FROM {self._qtable(CollectionNames.sdk_collections_table_name())} " + f"WHERE collection_name = '{name_escaped}' " + f"ORDER BY created_at, collection_id LIMIT 1" ) rows = self._execute(query_sql) if not rows or len(rows) == 0: @@ -2462,7 +2503,12 @@ def refresh_index(self) -> None: self._execute("CALL dbms_index_manager.refresh();") def _get_collection_id(self, collection_name: str) -> str: - collection_id_query_sql = f"SELECT COLLECTION_ID FROM `{CollectionNames.sdk_collections_table_name()}` WHERE COLLECTION_NAME = '{collection_name}'" + collection_name_escaped = escape_string(collection_name) + collection_id_query_sql = ( + f"SELECT collection_id FROM {self._qtable(CollectionNames.sdk_collections_table_name())} " + f"WHERE collection_name = '{collection_name_escaped}' " + f"ORDER BY created_at, collection_id LIMIT 1" + ) collection_id_query_result = self._execute(collection_id_query_sql) if not collection_id_query_result or len(collection_id_query_result) == 0: raise ValueError(f"Collection not found: '{collection_name}'") diff --git a/tests/integration_tests/test_get_or_create_collection_concurrency.py b/tests/integration_tests/test_get_or_create_collection_concurrency.py new file mode 100644 index 00000000..a24a5bcf --- /dev/null +++ b/tests/integration_tests/test_get_or_create_collection_concurrency.py @@ -0,0 +1,92 @@ +""" +Integration tests for concurrent-safe get_or_create_collection. +""" + +from __future__ import annotations + +import contextlib +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor, as_completed + +import pytest + +import pyseekdb +from pymysql.converters import escape_string + + +def _make_oceanbase_client(): + import os + + return pyseekdb.Client( + host=os.environ.get("OB_HOST", "127.0.0.1"), + port=int(os.environ.get("OB_PORT", "10902")), + tenant=os.environ.get("OB_TENANT", "mysql"), + database=os.environ.get("OB_DATABASE", "test"), + user=os.environ.get("OB_USER", "root"), + password=os.environ.get("OB_PASSWORD", ""), + ) + + +def _count_sdk_collections(client, collection_name: str) -> int: + name_escaped = escape_string(collection_name) + rows = client._server._execute( + "SELECT COUNT(*) AS cnt FROM sdk_collections " + f"WHERE collection_name = '{name_escaped}'" + ) + row = rows[0] + return int(row["cnt"] if isinstance(row, dict) else row[0]) + + +def _has_unique_name_index(client) -> bool: + rows = client._server._execute("SHOW INDEX FROM sdk_collections") + for row in rows: + key_name = row["Key_name"] if isinstance(row, dict) else row[2] + if key_name == "uk_sdk_coll_name": + return True + return False + + +class TestGetOrCreateCollectionConcurrencyOceanBase: + def test_concurrent_get_or_create_collection_is_idempotent(self, oceanbase_client): + """8 independent clients race on get_or_create_collection with the same name.""" + owner = oceanbase_client + collection_name = f"goc_conc_{uuid.uuid4().hex[:12]}" + barrier = threading.Barrier(8) + results: dict[int, dict] = {} + errors: list[str] = [] + + def _worker(thread_id: int) -> None: + client = _make_oceanbase_client() + try: + barrier.wait(timeout=30) + coll = client.get_or_create_collection( + collection_name, + configuration=pyseekdb.HNSWConfiguration(dimension=3, distance="cosine"), + embedding_function=None, + ) + results[thread_id] = { + "collection_id": str(coll.id), + "name": coll.name, + } + except Exception as exc: # noqa: BLE001 + errors.append(f"thread {thread_id}: {exc!r}") + finally: + if hasattr(client, "close"): + client.close() + + with ThreadPoolExecutor(max_workers=8) as pool: + futures = [pool.submit(_worker, i) for i in range(8)] + for fut in as_completed(futures, timeout=120): + fut.result() + + try: + assert not errors, f"concurrent get_or_create_collection failed: {errors}" + assert len(results) == 8, results + coll_ids = {item["collection_id"] for item in results.values()} + assert len(coll_ids) == 1, f"expected one collection_id, got {coll_ids}" + assert _count_sdk_collections(owner, collection_name) == 1 + assert _has_unique_name_index(owner) + finally: + with contextlib.suppress(Exception): + owner.delete_collection(name=collection_name) diff --git a/tests/integration_tests/test_get_or_create_collection_multiprocess.py b/tests/integration_tests/test_get_or_create_collection_multiprocess.py index 7d16fc01..d53fba0b 100644 --- a/tests/integration_tests/test_get_or_create_collection_multiprocess.py +++ b/tests/integration_tests/test_get_or_create_collection_multiprocess.py @@ -560,6 +560,16 @@ def test_concurrent_get_or_create_collection(self, multiprocess_db, _mode): assert client.has_collection(collection_name) collection = _get_collection(client, collection_name) assert collection.name == collection_name + rows = client._server._execute( + "SELECT collection_id FROM sdk_collections " + f"WHERE collection_name = '{collection_name}'" + ) + assert len(rows) == 1, ( + f"expected exactly one sdk_collections row for {collection_name!r}, got {len(rows)}" + ) + row = rows[0] + catalog_id = row["collection_id"] if isinstance(row, dict) else row[0] + assert str(catalog_id) == str(collection.id) finally: with contextlib.suppress(Exception): client.delete_collection(collection_name) diff --git a/tests/integration_tests/test_logic_table_monitoring_p1.py b/tests/integration_tests/test_logic_table_monitoring_p1.py index 9294bf50..ef9d5261 100644 --- a/tests/integration_tests/test_logic_table_monitoring_p1.py +++ b/tests/integration_tests/test_logic_table_monitoring_p1.py @@ -7,12 +7,9 @@ from __future__ import annotations import threading -import time import uuid from concurrent.futures import ThreadPoolExecutor, as_completed -import pytest - import pyseekdb from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, ns_schema diff --git a/tests/unit_tests/test_get_or_create_collection_concurrency.py b/tests/unit_tests/test_get_or_create_collection_concurrency.py index 00a963d7..6defdbd7 100644 --- a/tests/unit_tests/test_get_or_create_collection_concurrency.py +++ b/tests/unit_tests/test_get_or_create_collection_concurrency.py @@ -15,10 +15,63 @@ from pyseekdb.client.client_base import ( # noqa: E402 BaseClient, _is_collection_conflict_error, + _is_sdk_collection_catalog_conflict_error, ) from pyseekdb.client.types import _NOT_PROVIDED # noqa: E402 +class TestCollectionCatalogConflictDetection: + def test_detects_integrity_error_on_sdk_collections(self): + class IntegrityError(Exception): + pass + + exc = IntegrityError( + '(1062, "Duplicate entry \'my_coll\' for key \'uk_sdk_coll_name\'")' + ) + assert _is_sdk_collection_catalog_conflict_error(exc) + + def test_ignores_unrelated_errors(self): + assert not _is_sdk_collection_catalog_conflict_error(ValueError("invalid dimension")) + + +class TestCollectionCatalogInsertRecovery: + def test_insert_conflict_reuses_existing_collection_id(self): + client = MagicMock(spec=BaseClient) + client._get_collection_id.side_effect = [ValueError("not found"), "existing_id"] + conn = MagicMock() + client._ensure_connection.return_value = conn + + class IntegrityError(Exception): + pass + + def execute_side_effect(sql): + if "INSERT INTO" in sql: + raise IntegrityError( + '(1062, "Duplicate entry \'items\' for key \'uk_sdk_coll_name\'")' + ) + return [] + + client._execute.side_effect = execute_side_effect + + result = BaseClient._create_collection_meta_v2(client, "items", None) + + assert result["collection_id"] == "existing_id" + conn.rollback.assert_called_once() + assert client._get_collection_id.call_count == 2 + + def test_existing_catalog_row_is_reused_without_insert(self): + client = MagicMock(spec=BaseClient) + client._get_collection_id.return_value = "existing_id" + + result = BaseClient._create_collection_meta_v2(client, "items", None) + + assert result["collection_id"] == "existing_id" + insert_calls = [ + call for call in client._execute.call_args_list if "INSERT INTO" in str(call) + ] + assert not insert_calls + + class TestCollectionConflictDetection: def test_detects_value_error_for_existing_collection(self): assert _is_collection_conflict_error(ValueError("Collection 'items' already exists")) From f87406642a13805889962783cc978e0be7793b6e Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 23 Jun 2026 17:35:22 +0800 Subject: [PATCH 60/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91=20?= =?UTF-8?q?=E5=88=9B=E5=BB=BAnamespcae=E6=8A=A5=E9=94=991235=EF=BC=8C?= =?UTF-8?q?=E5=BA=94=E8=AF=A5=E6=98=AFobproxy=E4=B8=8D=E6=94=AF=E6=8C=81ge?= =?UTF-8?q?t=5Flock=E5=87=BD=E6=95=B0=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 228 ++++++------------ .../test_namespace_lifecycle_concurrency.py | 19 +- tests/unit_tests/test_namespace.py | 5 +- .../test_namespace_lifecycle_lock.py | 50 +--- .../test_namespace_upsert_reconcile.py | 11 +- 5 files changed, 98 insertions(+), 215 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7232a38d..7ede6737 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -3,7 +3,6 @@ """ import contextlib -import hashlib import json import logging import os @@ -1299,6 +1298,22 @@ def _ensure_namespace_catalogs(self) -> None: self._execute(ns_namespaces_sql) self._execute(ns_ltables_sql) self._execute(namespaces_stats_sql) + with contextlib.suppress(Exception): + self._execute( + f"CREATE UNIQUE INDEX uk_sdk_ns_coll_name ON {ns_namespaces_q} " + f"(collection_id, namespace_name)" + ) + with contextlib.suppress(Exception): + self._execute( + f"CREATE UNIQUE INDEX uk_sdk_lt_coll_ns_name ON {ns_ltables_q} " + f"(collection_id, namespace_id, ltable_name)" + ) + + def _rollback_connection_if_supported(self) -> None: + conn_getter = getattr(self, "_ensure_connection", None) + if conn_getter is not None: + with contextlib.suppress(Exception): + conn_getter().rollback() def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> dict: self._create_sdk_collections_if_not_exists() @@ -1643,7 +1658,7 @@ def _insert_ns_namespace_catalog_row( ) except Exception as exc: if idempotent and _is_namespace_catalog_conflict_error(exc): - pass + self._rollback_connection_if_supported() else: raise return self._fetch_ns_namespace_id(collection_id, namespace_name) @@ -1667,7 +1682,7 @@ def _insert_ns_ltable_catalog_row( ) except Exception as exc: if idempotent and _is_namespace_catalog_conflict_error(exc): - pass + self._rollback_connection_if_supported() else: raise return self._fetch_ns_ltable_id(collection_id, namespace_id, ltable_name) @@ -1693,34 +1708,43 @@ def _finalize_ns_namespace_meta( return {"namespace_id": str(namespace_id), "namespace_name": namespace_name, "ltable_id": str(ltable_id)} def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: - with self._namespace_lifecycle_guard(collection_id, namespace_name): + if self._get_ns_namespace_meta(collection_id, namespace_name) is not None: + raise ValueError(f"Namespace '{namespace_name}' already exists") + try: ns_id = self._insert_ns_namespace_catalog_row( collection_id, namespace_name, idempotent=False ) lt_id = self._insert_ns_ltable_catalog_row( collection_id, ns_id, idempotent=False ) - return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) + except Exception as exc: + if _is_namespace_catalog_conflict_error(exc): + self._rollback_connection_if_supported() + if self._get_ns_namespace_meta(collection_id, namespace_name) is not None: + raise ValueError( + f"Namespace '{namespace_name}' already exists" + ) from exc + raise + return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) def _get_or_create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: - with self._namespace_lifecycle_guard(collection_id, namespace_name): - meta = self._get_ns_namespace_meta(collection_id, namespace_name) - if meta is not None: - if meta.get("ltable_id") is not None: - return meta - lt_id = self._insert_ns_ltable_catalog_row( - collection_id, int(meta["namespace_id"]), idempotent=True - ) - return self._finalize_ns_namespace_meta( - collection_id, namespace_name, int(meta["namespace_id"]), lt_id - ) - ns_id = self._insert_ns_namespace_catalog_row( - collection_id, namespace_name, idempotent=True - ) + meta = self._get_ns_namespace_meta(collection_id, namespace_name) + if meta is not None: + if meta.get("ltable_id") is not None: + return meta lt_id = self._insert_ns_ltable_catalog_row( - collection_id, ns_id, idempotent=True + collection_id, int(meta["namespace_id"]), idempotent=True ) - return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) + return self._finalize_ns_namespace_meta( + collection_id, namespace_name, int(meta["namespace_id"]), lt_id + ) + ns_id = self._insert_ns_namespace_catalog_row( + collection_id, namespace_name, idempotent=True + ) + lt_id = self._insert_ns_ltable_catalog_row( + collection_id, ns_id, idempotent=True + ) + return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: namespace_name_escaped = escape_string(namespace_name) @@ -1759,8 +1783,7 @@ def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dic return meta def _has_ns_namespace(self, collection_id: str, namespace_name: str) -> bool: - with self._namespace_lifecycle_guard(collection_id, namespace_name): - return self._get_ns_namespace_meta(collection_id, namespace_name) is not None + return self._get_ns_namespace_meta(collection_id, namespace_name) is not None def _ns_namespace_exists_by_id(self, collection_id: str, namespace_id: str) -> bool: """Whether a namespace with this id is still live in the catalog. @@ -1781,22 +1804,21 @@ def _ns_namespace_exists_by_id(self, collection_id: str, namespace_id: str) -> b return bool(rows) def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> None: - with self._namespace_lifecycle_guard(collection_id, namespace_name): - meta = self._get_ns_namespace_meta(collection_id, namespace_name) - if meta is None: - raise ValueError(f"Namespace '{namespace_name}' not found") - ns_id = meta["namespace_id"] - lt_id_raw = meta.get("ltable_id") - lt_id = int(lt_id_raw) if lt_id_raw is not None else None - collection_id_escaped = escape_string(collection_id) - # PL reads session database_name; must match where catalog tables live. - self._use_catalog_database() - self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(ns_id), ltable_id=lt_id - ) - self._execute( - f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})" - ) + meta = self._get_ns_namespace_meta(collection_id, namespace_name) + if meta is None: + raise ValueError(f"Namespace '{namespace_name}' not found") + ns_id = meta["namespace_id"] + lt_id_raw = meta.get("ltable_id") + lt_id = int(lt_id_raw) if lt_id_raw is not None else None + collection_id_escaped = escape_string(collection_id) + # PL reads session database_name; must match where catalog tables live. + self._use_catalog_database() + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(ns_id), ltable_id=lt_id + ) + self._execute( + f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})" + ) def _list_ns_namespaces(self, collection_id: str) -> list[dict]: collection_id_escaped = escape_string(collection_id) @@ -5048,105 +5070,6 @@ def _append_namespace_filter( return f"WHERE {ns_cond} AND ({inner})", params return f"WHERE {ns_cond} AND ({where_clause})", params - @staticmethod - def _namespace_lifecycle_lock_name(collection_id: str, namespace_name: str) -> str: - key = f"{collection_id}:{namespace_name}" - digest = hashlib.sha256(key.encode()).hexdigest()[:40] - return f"pyseekdb:nslc:{digest}" - - @contextlib.contextmanager - def _namespace_lifecycle_lock( - self, - collection_id: str, - namespace_name: str, - *, - total_wait_seconds: float = 30.0, - ): - lock_name = self._namespace_lifecycle_lock_name(collection_id, namespace_name) - lock_name_escaped = escape_string(lock_name) - acquired = False - deadline = time.monotonic() + total_wait_seconds - try: - while time.monotonic() < deadline: - rows = self._execute( - f"SELECT GET_LOCK('{lock_name_escaped}', 1) AS got" - ) - row = rows[0] - got = row.get("got") if isinstance(row, dict) else row[0] - if got == 1: - acquired = True - yield True - return - time.sleep(0.05) - yield False - finally: - if acquired: - with contextlib.suppress(Exception): - self._execute(f"SELECT RELEASE_LOCK('{lock_name_escaped}')") - - @contextlib.contextmanager - def _namespace_lifecycle_guard( - self, - collection_id: str, - namespace_name: str, - *, - total_wait_seconds: float = 30.0, - ): - with self._namespace_lifecycle_lock( - collection_id, namespace_name, total_wait_seconds=total_wait_seconds - ) as acquired: - if not acquired: - raise TimeoutError( - f"timed out waiting for namespace lifecycle lock on " - f"{namespace_name!r} in collection {collection_id!r}" - ) - yield - - @staticmethod - def _namespace_record_lock_name( - collection_id: str, - namespace_id: str | int, - ltable_id: int, - record_id: str, - ) -> str: - key = f"{collection_id}:{namespace_id}:{ltable_id}:{record_id}" - digest = hashlib.sha256(key.encode()).hexdigest()[:40] - return f"pyseekdb:nsu:{digest}" - - @contextlib.contextmanager - def _namespace_record_lock( - self, - collection_id: str, - namespace_id: str | int, - ltable_id: int, - record_id: str, - *, - total_wait_seconds: float = 30.0, - ): - lock_name = self._namespace_record_lock_name( - collection_id, namespace_id, ltable_id, record_id - ) - lock_name_escaped = escape_string(lock_name) - acquired = False - deadline = time.monotonic() + total_wait_seconds - try: - while time.monotonic() < deadline: - rows = self._execute( - f"SELECT GET_LOCK('{lock_name_escaped}', 1) AS got" - ) - row = rows[0] - got = row.get("got") if isinstance(row, dict) else row[0] - if got == 1: - acquired = True - yield True - return - time.sleep(0.05) - yield False - finally: - if acquired: - with contextlib.suppress(Exception): - self._execute(f"SELECT RELEASE_LOCK('{lock_name_escaped}')") - def _count_namespace_records_by_id( self, table_name: str, @@ -5220,20 +5143,18 @@ def _reconcile_namespace_duplicate_records( doc_val = documents[i] if documents and i < len(documents) else None meta_val = metadatas[i] if metadatas and i < len(metadatas) else None emb_val = embeddings[i] if embeddings and i < len(embeddings) else None - for _ in range(60): - with self._namespace_record_lock( - collection_id, namespace_id, ltable_id, record_id - ) as acquired: - if not acquired: - continue - duplicate_count = self._count_namespace_records_by_id( - table_name, ns_id, ltable_id, record_id - ) - if duplicate_count <= 1: - break - self._delete_namespace_records_by_id( - table_name, ns_id, ltable_id, record_id - ) + for attempt in range(120): + duplicate_count = self._count_namespace_records_by_id( + table_name, ns_id, ltable_id, record_id + ) + if duplicate_count <= 1: + break + self._delete_namespace_records_by_id( + table_name, ns_id, ltable_id, record_id + ) + if self._count_namespace_records_by_id( + table_name, ns_id, ltable_id, record_id + ) == 0: self._namespace_add( collection_id=collection_id, collection_name=collection_name, @@ -5246,7 +5167,8 @@ def _reconcile_namespace_duplicate_records( embedding_function=embedding_function, **kwargs, ) - break + if attempt < 119: + time.sleep(0.05 * min(attempt + 1, 10)) def _namespace_add( # noqa: C901 self, diff --git a/tests/integration_tests/test_namespace_lifecycle_concurrency.py b/tests/integration_tests/test_namespace_lifecycle_concurrency.py index cceffcff..581c30b0 100644 --- a/tests/integration_tests/test_namespace_lifecycle_concurrency.py +++ b/tests/integration_tests/test_namespace_lifecycle_concurrency.py @@ -82,8 +82,8 @@ def _has_worker() -> None: cleanup(oceanbase_client, collection) -def test_multi_client_has_blocks_until_drop_finishes(oceanbase_client): - """has_namespace on client B waits for client A drop to release lifecycle lock.""" +def test_multi_client_has_does_not_error_during_drop(oceanbase_client): + """has_namespace stays callable while another client drops the namespace.""" collection = oceanbase_client.create_collection( name=_unique_name("test_ns_lc_drop_block"), schema=ns_schema(), @@ -94,9 +94,8 @@ def test_multi_client_has_blocks_until_drop_finishes(oceanbase_client): ns_name = "block_ns" drop_started = threading.Event() drop_can_finish = threading.Event() - has_started = threading.Event() errors: list[Exception] = [] - has_result: list[bool] = [] + has_results: list[bool] = [] try: collection.create_namespace(ns_name) @@ -120,8 +119,9 @@ def _drop_worker() -> None: def _has_worker() -> None: try: assert drop_started.wait(timeout=30), "drop did not start" - has_started.set() - has_result.append(checker.has_namespace(ns_name)) + for _ in range(10): + has_results.append(checker.has_namespace(ns_name)) + time.sleep(0.05) except Exception as exc: # noqa: BLE001 errors.append(exc) @@ -130,16 +130,13 @@ def _has_worker() -> None: drop_thread.start() has_thread.start() - assert has_started.wait(timeout=30), "has did not start while drop in progress" - time.sleep(0.2) - assert drop_thread.is_alive(), "drop finished before lifecycle lock blocked has" drop_can_finish.set() - drop_thread.join(timeout=120) has_thread.join(timeout=120) assert errors == [] - assert has_result == [False] + assert has_results + assert checker.has_namespace(ns_name) is False finally: for client in clients: with contextlib.suppress(Exception): diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 46bf9ed1..2c5c7bfa 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -1150,7 +1150,6 @@ class TestDeleteNamespaceUsesKernel: def test_delete_namespace_calls_dbms_logic_table(self): c = FakeClient() c._execute = MagicMock(side_effect=[ - [{"got": 1}], # GET_LOCK [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], None, # _get_ns_namespace_meta -> SET @namespace_id None, # _get_ns_namespace_meta -> SET @ltable_id @@ -1159,12 +1158,10 @@ def test_delete_namespace_calls_dbms_logic_table(self): None, # _delete_ns_namespace_meta -> SET @namespace_id None, # _delete_ns_namespace_meta -> SET @ltable_id None, # CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE - [{"got": 1}], # RELEASE_LOCK ]) c._delete_ns_namespace_meta("abc123", "ns1") calls = [str(call) for call in c._execute.call_args_list] - assert any("GET_LOCK" in s and "pyseekdb:nslc:" in s for s in calls) - assert any("RELEASE_LOCK" in s for s in calls) + assert not any("GET_LOCK" in s for s in calls) assert any("DBMS_LOGIC_TABLE.DROP_NAMESPACE" in s for s in calls) # session context must be set BEFORE DROP_NAMESPACE so the kernel sees the # right @collection_id / @namespace_id / @ltable_id for this call. diff --git a/tests/unit_tests/test_namespace_lifecycle_lock.py b/tests/unit_tests/test_namespace_lifecycle_lock.py index 7824f03d..929a47dd 100644 --- a/tests/unit_tests/test_namespace_lifecycle_lock.py +++ b/tests/unit_tests/test_namespace_lifecycle_lock.py @@ -1,6 +1,5 @@ -"""Unit tests for namespace lifecycle GET_LOCK serialization.""" +"""Unit tests for namespace catalog concurrency without GET_LOCK.""" -import contextlib import sys from pathlib import Path from unittest.mock import MagicMock @@ -14,51 +13,26 @@ from pyseekdb.client.client_base import BaseClient # noqa: E402 -class TestNamespaceLifecycleLockName: - def test_lock_name_is_stable_and_bounded(self): - name_a = BaseClient._namespace_lifecycle_lock_name("cid", "ns1") - name_b = BaseClient._namespace_lifecycle_lock_name("cid", "ns1") - name_c = BaseClient._namespace_lifecycle_lock_name("cid", "ns2") - assert name_a == name_b - assert name_a != name_c - assert name_a.startswith("pyseekdb:nslc:") - assert len(name_a) <= 64 - - -class TestNamespaceLifecycleGuard: - def test_has_namespace_acquires_and_releases_lock(self): +class TestNamespaceLifecycleWithoutGetLock: + def test_has_namespace_queries_catalog_directly(self): client = MagicMock(spec=BaseClient) client._get_ns_namespace_meta.return_value = None - @contextlib.contextmanager - def _lock_cm(*_args, **_kwargs): - yield True - - client._namespace_lifecycle_lock.side_effect = _lock_cm - client._namespace_lifecycle_guard = ( - BaseClient._namespace_lifecycle_guard.__get__(client, BaseClient) - ) - assert BaseClient._has_ns_namespace(client, "cid", "ns1") is False - client._namespace_lifecycle_lock.assert_called_once_with( - "cid", "ns1", total_wait_seconds=30.0 - ) client._get_ns_namespace_meta.assert_called_once_with("cid", "ns1") - def test_guard_raises_when_lock_times_out(self): + def test_create_namespace_raises_when_catalog_row_exists(self): client = MagicMock(spec=BaseClient) + client._get_ns_namespace_meta.return_value = { + "namespace_id": "42", + "namespace_name": "ns1", + "ltable_id": "7", + } - @contextlib.contextmanager - def _no_lock(): - yield False - - client._namespace_lifecycle_lock.return_value = _no_lock() - client._namespace_lifecycle_guard = ( - BaseClient._namespace_lifecycle_guard.__get__(client, BaseClient) - ) + with pytest.raises(ValueError, match="already exists"): + BaseClient._create_ns_namespace_meta(client, "cid", "ns1") - with pytest.raises(TimeoutError, match="namespace lifecycle lock"): - BaseClient._has_ns_namespace(client, "cid", "ns1") + client._insert_ns_namespace_catalog_row.assert_not_called() if __name__ == "__main__": diff --git a/tests/unit_tests/test_namespace_upsert_reconcile.py b/tests/unit_tests/test_namespace_upsert_reconcile.py index 1a79c9eb..6ddf7018 100644 --- a/tests/unit_tests/test_namespace_upsert_reconcile.py +++ b/tests/unit_tests/test_namespace_upsert_reconcile.py @@ -4,7 +4,7 @@ import sys from pathlib import Path -from unittest.mock import MagicMock, call +from unittest.mock import MagicMock import pytest @@ -16,14 +16,8 @@ class TestNamespaceUpsertReconcile: - def test_lock_name_is_stable_and_bounded(self): - name = BaseClient._namespace_record_lock_name("c" * 32, 7, 9, "same_new_id") - assert name.startswith("pyseekdb:nsu:") - assert len(name) <= 64 - def test_reconcile_skips_when_single_row_exists(self): client = MagicMock(spec=BaseClient) - client._namespace_record_lock.return_value.__enter__.return_value = True client._count_namespace_records_by_id.return_value = 1 BaseClient._reconcile_namespace_duplicate_records( @@ -46,8 +40,7 @@ def test_reconcile_skips_when_single_row_exists(self): def test_reconcile_collapses_duplicate_rows(self): client = MagicMock(spec=BaseClient) - client._namespace_record_lock.return_value.__enter__.return_value = True - client._count_namespace_records_by_id.return_value = 4 + client._count_namespace_records_by_id.side_effect = [4, 0, 1] BaseClient._reconcile_namespace_duplicate_records( client, From c688423004e571fe05c901dd299c3e9882890cd6 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 24 Jun 2026 11:38:10 +0800 Subject: [PATCH 61/84] =?UTF-8?q?Docstring=20=E8=A6=86=E7=9B=96=E7=8E=87?= =?UTF-8?q?=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/__init__.py | 1 + src/pyseekdb/client/admin_client.py | 2 + src/pyseekdb/client/client_base.py | 144 ++++++++++++++++-- src/pyseekdb/client/client_seekdb_embedded.py | 3 + src/pyseekdb/client/client_seekdb_server.py | 9 +- src/pyseekdb/client/collection.py | 15 +- src/pyseekdb/client/namespace.py | 17 +-- src/pyseekdb/client/validators.py | 31 ++++ .../test_namespace_drop_validation.py | 6 +- ...space_hybrid_search_multi_coll_multi_ns.py | 3 +- ...t_namespace_hybrid_search_triple_branch.py | 2 +- .../test_namespace_lifecycle.py | 2 +- .../test_namespace_upsert_concurrency.py | 5 +- ...st_get_or_create_collection_concurrency.py | 44 +++++- tests/unit_tests/test_namespace.py | 14 +- 15 files changed, 247 insertions(+), 51 deletions(-) diff --git a/src/pyseekdb/client/__init__.py b/src/pyseekdb/client/__init__.py index 1ce0aaef..03827b28 100644 --- a/src/pyseekdb/client/__init__.py +++ b/src/pyseekdb/client/__init__.py @@ -71,6 +71,7 @@ def _resolve_password(password: str) -> str: def _default_seekdb_path() -> str: # Keep existing behavior: default to "seekdb.db" under current working directory. + """Return the default on-disk path for the embedded SeekDB store.""" return os.path.abspath("seekdb.db") diff --git a/src/pyseekdb/client/admin_client.py b/src/pyseekdb/client/admin_client.py index 4cdb46ae..aef416be 100644 --- a/src/pyseekdb/client/admin_client.py +++ b/src/pyseekdb/client/admin_client.py @@ -157,6 +157,7 @@ def fork_database(self, source_name: str, destination_name: str, tenant: str = D return self._server.fork_database(source_name=source_name, destination_name=destination_name, tenant=tenant) def __repr__(self): + """Return the developer-readable representation.""" return f"" def __enter__(self): @@ -247,6 +248,7 @@ def count_collection(self) -> int: return self._server.count_collection() def __repr__(self): + """Return the developer-readable representation.""" return f"" def __enter__(self): diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 7ede6737..01867da2 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -71,6 +71,7 @@ def _extract_collection_id_from_sdk_row(row: Any) -> str: + """Extract the collection_id from an sdk_collections row (dict/tuple/scalar).""" if isinstance(row, dict): collection_id = row.get("COLLECTION_ID") or row.get("collection_id") or "" elif isinstance(row, (tuple, list)): @@ -81,6 +82,7 @@ def _extract_collection_id_from_sdk_row(row: Any) -> str: def _is_collection_conflict_error(exc: BaseException) -> bool: + """Whether the exception (or its cause chain) indicates a collection/table already exists.""" current: BaseException | None = exc while current is not None: message = str(current).lower() @@ -93,6 +95,7 @@ def _is_collection_conflict_error(exc: BaseException) -> bool: def _is_namespace_catalog_conflict_error(exc: BaseException) -> bool: + """Whether the exception indicates a duplicate sdk_namespaces/sdk_ltables unique-key conflict (1062).""" current: BaseException | None = exc while current is not None: message = str(current).lower() @@ -112,6 +115,7 @@ def _is_namespace_catalog_conflict_error(exc: BaseException) -> bool: def _is_sdk_collection_catalog_conflict_error(exc: BaseException) -> bool: + """Whether the exception indicates a duplicate sdk_collections collection_name conflict (1062).""" current: BaseException | None = exc while current is not None: message = str(current).lower() @@ -132,6 +136,7 @@ def _is_sdk_collection_catalog_conflict_error(exc: BaseException) -> bool: def _extract_hnsw_config(config: ConfigurationParam) -> HNSWConfiguration | None: + """Return the HNSW config from a Configuration/HNSWConfiguration, or None.""" if config is None: return None elif isinstance(config, HNSWConfiguration): @@ -145,6 +150,7 @@ def _extract_hnsw_config(config: ConfigurationParam) -> HNSWConfiguration | None def _extract_fulltext_config( config: ConfigurationParam, ) -> FulltextIndexConfig | None: + """Return the fulltext index config from a Configuration, or None.""" if config is None: return None elif isinstance(config, HNSWConfiguration): @@ -188,7 +194,14 @@ def _validate_collection_name(name: str) -> None: ) -from .validators import _MAX_N_RESULTS, _MAX_NAMESPACE_BATCH_SIZE, _validate_namespace_name, _validate_record_ids # noqa: F401 +from .validators import ( # noqa: F401 + _MAX_N_RESULTS, + _MAX_NAMESPACE_BATCH_SIZE, + _quote_sql_identifier, + _validate_database_name, + _validate_namespace_name, + _validate_record_ids, +) _DEFAULT_PARTITION_COUNT = 1000 # Unquoted id for WHERE/CASE; plain JSON_EXTRACT returns a quoted JSON string and @@ -197,6 +210,7 @@ def _validate_collection_name(name: str) -> None: def _build_default_ltable_schema() -> dict: + """Return the default logical-table schema (metadata/content/embedding columns and indexes).""" return { "col_info": [ {"col_idx": 1, "col_name": "metadata", "col_type": "JSON"}, @@ -282,6 +296,7 @@ def _get_vector_index_sql(hnsw_config: HNSWConfiguration) -> str: def _get_ivf_vector_index_sql(ivf_config: "IVFConfiguration") -> str: + """Build the IVF vector index DDL fragment from an IVFConfiguration.""" property_parts = [] if ivf_config.properties: for k, v in ivf_config.properties.items(): @@ -289,7 +304,7 @@ def _get_ivf_vector_index_sql(ivf_config: "IVFConfiguration") -> str: property_parts.append(f"{k}='{v}'") else: property_parts.append(f"{k}={v}") - if ivf_config.centroids_fresh_mode: + if ivf_config.centroids_fresh_mode is not None: property_parts.append(f"centroids_fresh_mode={ivf_config.centroids_fresh_mode}") property_str = ", ".join(property_parts) properties_str = f", {property_str}" if property_str else "" @@ -432,6 +447,7 @@ class _CollectionMeta: @staticmethod def from_row(row: Any) -> "_CollectionMeta": + """Construct a _CollectionMeta from a catalog row (dict for server, tuple for embedded).""" if isinstance(row, dict): # Server client returns dict, get the first value collection_id = row["COLLECTION_ID"] @@ -467,6 +483,7 @@ class BaseClient(BaseConnection, AdminAPI): # ==================== Database Type Detection ==================== def _validate_ob_database_type(self) -> None: + """Validate that the backend is OceanBase and meets the minimum version for namespaces.""" db_type, version = self.detect_db_type_and_version() if db_type.lower() != "oceanbase": raise ValueError("use_namespace=True is only supported on OceanBase") @@ -477,6 +494,7 @@ def _validate_ob_database_type(self) -> None: ) def _is_shared_storage_mode(self) -> bool: + """Return whether the OceanBase deployment runs in shared-storage mode.""" cached = getattr(self, "_shared_storage", None) if cached is not None: return cached @@ -588,6 +606,7 @@ def _extract_seekdb_version(version_str: str) -> str | None: # Truncate potentially verbose or sensitive database responses in error message def _truncate(val, length=20): + """Truncate a value to a short string for logging.""" if val is None: return "None" val_str = str(val) @@ -605,9 +624,11 @@ def _database_tenant(self, tenant: str) -> str | None: return None def _database_context(self, tenant: str | None) -> str: + """Yield a context with the active database selected for the connection.""" return f" in tenant: {tenant}" if tenant else "" def _parse_schema_row(self, row: Any) -> tuple[str | None, str | None, str | None]: + """Parse a raw catalog row into a schema descriptor.""" if isinstance(row, dict): return ( row.get("SCHEMA_NAME"), @@ -764,6 +785,7 @@ def _prepare_schema_parameters( # noqa: C901 ) -> Schema: # Handle embedding function first # If not provided (sentinel), use default embedding function + """Normalize and validate schema parameters before creating a collection.""" if embedding_function is _NOT_PROVIDED: embedding_function = get_default_embedding_function() @@ -1023,6 +1045,7 @@ def create_collection( ) def _create_namespace_collection(self, name: str, schema: Schema, partition_count: int | None = None, **kwargs) -> "Collection": + """Create a namespace-enabled collection and its catalog/physical tables.""" dense_embedding_function = schema.vector_index.embedding_function ivf_config = schema.vector_index.ivf hnsw_config = schema.vector_index.hnsw @@ -1151,6 +1174,7 @@ def _get_embedding_function_dimension(self, embedding_function: EmbeddingFunctio ) from e def _create_sdk_collections_if_not_exists(self) -> None: + """Create the sdk_collections catalog table if it does not already exist.""" try: self._use_catalog_database() sdk_coll = self._qtable(CollectionNames.sdk_collections_table_name()) @@ -1177,6 +1201,7 @@ def _create_collection_meta_v2( embedding_function, sparse_vector_index_config: SparseVectorIndexConfig | None = None, ) -> dict[str, str]: + """Insert collection metadata into sdk_collections, resolving unique-key conflicts idempotently.""" try: results = {} settings = {"version": 2} @@ -1230,6 +1255,7 @@ def _create_collection_meta_v2( raise ValueError(f"Failed to create collection metadata: {e}") from e def _create_collection_meta_v1(self, collection_name: str) -> str: + """Insert collection metadata using the legacy v1 catalog layout.""" try: table_name = CollectionNames.table_name(collection_name) return table_name # noqa: TRY300 @@ -1243,17 +1269,21 @@ def _catalog_database(self) -> str: db = getattr(self, "database", None) if not db: raise ValueError("client database is not configured") + _validate_database_name(db) return db def _qtable(self, table: str) -> str: """Fully-qualified catalog table: `{database}`.`{table}`.""" - return f"`{self._catalog_database()}`.`{table}`" + db = _quote_sql_identifier(self._catalog_database()) + table_quoted = _quote_sql_identifier(table) + return f"{db}.{table_quoted}" def _use_catalog_database(self) -> None: """Align session with pymysql database= so PL (DROP_NAMESPACE) uses the same DB.""" - self._execute(f"USE `{self._catalog_database()}`") + self._execute(f"USE {_quote_sql_identifier(self._catalog_database())}") def _ensure_namespace_catalogs(self) -> None: + """Create the sdk_namespaces and sdk_ltables catalog tables and their unique indexes.""" self._use_catalog_database() ns_namespaces_q = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) ns_ltables_q = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) @@ -1310,12 +1340,14 @@ def _ensure_namespace_catalogs(self) -> None: ) def _rollback_connection_if_supported(self) -> None: + """Roll back the current connection transaction if the backend supports it.""" conn_getter = getattr(self, "_ensure_connection", None) if conn_getter is not None: with contextlib.suppress(Exception): conn_getter().rollback() def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> dict: + """Insert namespace collection metadata, resolving unique-key conflicts idempotently.""" self._create_sdk_collections_if_not_exists() settings_str = escape_string(json.dumps(settings)) collection_name_escaped = escape_string(collection_name) @@ -1343,6 +1375,7 @@ def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> di return {"collection_id": collection_id, "collection_name": collection_name} def _get_ns_collection_meta(self, collection_name: str) -> dict | None: + """Fetch namespace collection metadata from the catalog.""" collection_name_escaped = escape_string(collection_name) try: rows = self._execute( @@ -1374,6 +1407,7 @@ def _get_ns_collection_meta(self, collection_name: str) -> dict | None: } def _has_ns_collection(self, collection_name: str) -> bool: + """Return whether a namespace collection with the given name exists.""" try: return self._get_ns_collection_meta(collection_name) is not None except Exception: @@ -1393,6 +1427,7 @@ def _ns_collection_exists_by_id(self, collection_id: str) -> bool: return bool(rows) def _delete_ns_collection_meta(self, collection_name: str) -> None: + """Delete namespace collection metadata from the catalog.""" meta = self._get_ns_collection_meta(collection_name) if meta is None: raise ValueError(f"Namespace collection '{collection_name}' not found") @@ -1429,6 +1464,7 @@ def _create_namespace_physical_tables( partition_count: int = _DEFAULT_PARTITION_COUNT, cleanup_on_error: bool = True, ) -> None: + """Create the physical tables backing a namespace collection.""" tg_name = NamespaceCollectionNames.tablegroup_name(collection_id) data_table = NamespaceCollectionNames.data_table_name(collection_id) kv_table = NamespaceCollectionNames.kv_data_table_name(collection_id) @@ -1500,6 +1536,7 @@ def _create_namespace_physical_tables( raise def _cleanup_namespace_physical_tables(self, collection_id: str) -> None: + """Drop the physical tables backing a namespace collection.""" for suffix_fn in [ NamespaceCollectionNames.data_table_name, NamespaceCollectionNames.logic_schema_table_name, @@ -1586,6 +1623,7 @@ def _resolve_namespace_ltable_id( def _cache_namespace_ltable_id( self, collection_id: str, namespace_id: int, ltable_id: int ) -> None: + """Cache the resolved logical-table id for a namespace.""" key = (str(collection_id), int(namespace_id)) cache = getattr(self, "_ns_ltable_id_cache", None) if cache is None: @@ -1599,6 +1637,7 @@ def _set_session_ns_context( namespace_id: int | None = None, ltable_id: int | None = None, ) -> None: + """Set session variables identifying the active namespace context.""" if collection_id is not None: self._execute(f"SET @collection_id = '{escape_string(str(collection_id))}'") if namespace_id is not None: @@ -1607,6 +1646,7 @@ def _set_session_ns_context( self._execute(f"SET @ltable_id = {int(ltable_id)}") def _fetch_ns_namespace_id(self, collection_id: str, namespace_name: str) -> int: + """Fetch the namespace id for a collection/namespace pair from the catalog.""" namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) ns_table = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) @@ -1626,6 +1666,7 @@ def _fetch_ns_ltable_id( namespace_id: int, ltable_name: str = "default", ) -> int: + """Fetch the logical-table id for a namespace from the catalog.""" collection_id_escaped = escape_string(collection_id) ltable_name_escaped = escape_string(ltable_name) lt_table = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) @@ -1648,6 +1689,7 @@ def _insert_ns_namespace_catalog_row( *, idempotent: bool, ) -> int: + """Insert a row into the sdk_namespaces catalog table.""" namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) ns_table = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) @@ -1671,6 +1713,7 @@ def _insert_ns_ltable_catalog_row( idempotent: bool, ltable_name: str = "default", ) -> int: + """Insert a row into the sdk_ltables catalog table.""" collection_id_escaped = escape_string(collection_id) ltable_name_escaped = escape_string(ltable_name) lt_table = self._qtable(NamespaceCollectionNames.sdk_ltables_table()) @@ -1694,6 +1737,7 @@ def _finalize_ns_namespace_meta( namespace_id: int, ltable_id: int, ) -> dict: + """Finalize namespace metadata after catalog rows and physical tables are created.""" self._cache_namespace_ltable_id(collection_id, namespace_id, ltable_id) schema_table = self._qtable( NamespaceCollectionNames.logic_schema_table_name(collection_id) @@ -1708,6 +1752,7 @@ def _finalize_ns_namespace_meta( return {"namespace_id": str(namespace_id), "namespace_name": namespace_name, "ltable_id": str(ltable_id)} def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: + """Create namespace metadata, resolving unique-key conflicts idempotently.""" if self._get_ns_namespace_meta(collection_id, namespace_name) is not None: raise ValueError(f"Namespace '{namespace_name}' already exists") try: @@ -1728,6 +1773,7 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) def _get_or_create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict: + """Get existing namespace metadata or create it if absent.""" meta = self._get_ns_namespace_meta(collection_id, namespace_name) if meta is not None: if meta.get("ltable_id") is not None: @@ -1747,6 +1793,7 @@ def _get_or_create_ns_namespace_meta(self, collection_id: str, namespace_name: s return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: + """Fetch namespace metadata from the catalog.""" namespace_name_escaped = escape_string(namespace_name) collection_id_escaped = escape_string(collection_id) ns_table = self._qtable(NamespaceCollectionNames.sdk_namespaces_table()) @@ -1783,6 +1830,7 @@ def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dic return meta def _has_ns_namespace(self, collection_id: str, namespace_name: str) -> bool: + """Return whether a namespace with the given name exists.""" return self._get_ns_namespace_meta(collection_id, namespace_name) is not None def _ns_namespace_exists_by_id(self, collection_id: str, namespace_id: str) -> bool: @@ -1804,6 +1852,7 @@ def _ns_namespace_exists_by_id(self, collection_id: str, namespace_id: str) -> b return bool(rows) def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> None: + """Delete namespace metadata from the catalog.""" meta = self._get_ns_namespace_meta(collection_id, namespace_name) if meta is None: raise ValueError(f"Namespace '{namespace_name}' not found") @@ -1821,6 +1870,7 @@ def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> ) def _list_ns_namespaces(self, collection_id: str) -> list[dict]: + """List namespaces registered for a collection.""" collection_id_escaped = escape_string(collection_id) rows = self._execute( f"SELECT namespace_id, namespace_name FROM {self._qtable(NamespaceCollectionNames.sdk_namespaces_table())} " @@ -1837,6 +1887,7 @@ def _list_ns_namespaces(self, collection_id: str) -> list[dict]: return results def _get_ns_namespace_id(self, collection_id: str, namespace_name: str) -> str: + """Resolve the namespace id for a collection/namespace pair.""" meta = self._get_ns_namespace_meta(collection_id, namespace_name) if meta is None: raise ValueError(f"Namespace '{namespace_name}' not found in collection {collection_id}") @@ -1845,6 +1896,7 @@ def _get_ns_namespace_id(self, collection_id: str, namespace_name: str) -> str: # ==================== End Namespace Catalog Methods ==================== def get_collection(self, name: str, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED) -> "Collection": + """Get an existing collection by name.""" ns_meta = None try: ns_meta = self._get_ns_collection_meta(name) @@ -1860,6 +1912,7 @@ def get_collection(self, name: str, embedding_function: EmbeddingFunctionParam = return collection def _build_ns_collection_from_meta(self, meta: dict, embedding_function=_NOT_PROVIDED) -> "Collection": + """Build a namespace Collection facade from catalog metadata.""" settings = meta.get("settings", {}) dimension = settings.get("dimension") distance = settings.get("distance", DEFAULT_DISTANCE_METRIC) @@ -1987,6 +2040,7 @@ def _resolve_collection_metadata_from_table(self, table_name: str, collection_na return metadata def _resolve_embedding_function(self, settings: str | None) -> EmbeddingFunction[EmbeddingDocuments]: + """Resolve the embedding function to use for a collection.""" if not settings: return None settings_json = json.loads(settings) @@ -2070,6 +2124,7 @@ def _validate_embedding_function( return embedding_function def _get_collection_v2(self, name: str, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED) -> "Collection": + """Fetch a collection using the v2 catalog layout.""" collection_meta = self._resolve_collection_metadata_from_sdk_collections(name) if not collection_meta or not collection_meta.collection_id: raise ValueError(f"Collection '{name}' does not exist") @@ -2201,6 +2256,7 @@ def list_collections(self) -> list["Collection"]: return collections def _list_ns_collections(self) -> list["Collection"]: + """List namespace-enabled collections.""" result = [] try: sdk_table = CollectionNames.sdk_collections_table_name() @@ -2240,6 +2296,7 @@ def _list_ns_collections(self) -> list["Collection"]: return result def _list_collections_v2(self) -> list["Collection"]: + """List collections using the v2 catalog layout.""" collections = [] try: # Detect if the sdk_collections table exists before querying it @@ -2374,6 +2431,7 @@ def has_collection(self, name: str) -> bool: return self._has_ns_collection(name) or self._has_collection_v2(name) or self._has_collection_v1(name) def _collection_table_exists(self, table_name: str) -> bool: + """Return whether the physical table for a collection exists.""" try: table_info = self._execute(f"DESCRIBE `{table_name}`") return table_info is not None and len(table_info) > 0 @@ -2381,6 +2439,7 @@ def _collection_table_exists(self, table_name: str) -> bool: return False def _has_collection_v2(self, name: str) -> bool: + """Return whether a collection exists using the v2 catalog layout.""" try: name_escaped = escape_string(name) query_sql = ( @@ -2483,9 +2542,43 @@ def get_or_create_collection( ) except Exception as exc: if _is_collection_conflict_error(exc): - return self.get_collection(name, embedding_function=embedding_function) + return self._get_or_resume_existing_collection( + name, + schema=schema, + configuration=configuration, + embedding_function=embedding_function, + use_namespace=use_namespace, + **kwargs, + ) raise + def _get_or_resume_existing_collection( + self, + name: str, + *, + schema: Schema | None, + configuration: ConfigurationParam, + embedding_function: EmbeddingFunctionParam, + use_namespace: bool, + **kwargs, + ) -> "Collection": + """Return an existing collection or resume an incomplete namespace-enabled one.""" + if use_namespace: + if self._get_ns_collection_meta(name) is None: + raise ValueError( + f"Collection '{name}' conflicted during create but namespace metadata is missing" + ) + if self._is_incomplete_ns_collection(name): + return self.create_collection( + name=name, + schema=schema, + configuration=configuration, + embedding_function=embedding_function, + use_namespace=use_namespace, + **kwargs, + ) + return self.get_collection(name, embedding_function=embedding_function) + def _get_collection_table_name(self, collection_id: str | None, collection_name: str) -> str: """ Get collection table name @@ -2495,18 +2588,21 @@ def _get_collection_table_name(self, collection_id: str | None, collection_name: return CollectionNames.table_name(collection_name) def _fork_table_enabled(self) -> bool: + """Return whether table fork is enabled on the backend.""" db_type, version = self.detect_db_type_and_version() version_110 = Version("1.1.0.0") logger.debug(f"db_type: {db_type}, version: {version}") return db_type.lower() == "seekdb" and version >= version_110 def _fork_database_enabled(self) -> bool: + """Return whether database fork is enabled on the backend.""" db_type, version = self.detect_db_type_and_version() version_120 = Version("1.2.0.0") logger.debug(f"db_type: {db_type}, version: {version}") return db_type.lower() == "seekdb" and version >= version_120 def _refresh_enabled(self) -> bool: + """Return whether index refresh is enabled on the backend.""" db_type, version = self.detect_db_type_and_version() version_130 = Version("1.3.0.0") logger.debug(f"db_type: {db_type}, version: {version}") @@ -2525,6 +2621,7 @@ def refresh_index(self) -> None: self._execute("CALL dbms_index_manager.refresh();") def _get_collection_id(self, collection_name: str) -> str: + """Resolve the collection id for a collection name.""" collection_name_escaped = escape_string(collection_name) collection_id_query_sql = ( f"SELECT collection_id FROM {self._qtable(CollectionNames.sdk_collections_table_name())} " @@ -3622,12 +3719,14 @@ def _use_context_manager_for_cursor(self) -> bool: return True def _should_fetch_results(self, cursor: Any, sql: str) -> bool: + """Return whether the given SQL statement is expected to yield result rows.""" description = getattr(cursor, "description", None) if description is not None: return True return is_query_sql(sql) def _execute(self, sql: str) -> Any: + """Execute a SQL statement against the connection and return any result rows.""" if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): print(f"[pyseekdb SQL] {sql}", flush=True) conn = self._ensure_connection() @@ -4499,10 +4598,12 @@ def _build_document_query( # noqa: C901 return None def _with_boost(expr: dict[str, Any] | None) -> dict[str, Any] | None: + """Apply field boosting to a document query expression.""" if boost is None or not expr: return expr def _apply_boost(target: Any) -> None: + """Apply a boost factor to a single field expression.""" if not isinstance(target, dict): return if "query_string" in target and isinstance(target["query_string"], dict): @@ -4761,6 +4862,7 @@ def _build_knn_expression( # noqa: C901 embedding_function = kwargs.get("embedding_function") def _normalize_vectors(raw_embeddings: Any) -> list[list[float]]: + """Normalize input vectors to a consistent list-of-floats form.""" if raw_embeddings is None: return [] if isinstance(raw_embeddings, list) and raw_embeddings and isinstance(raw_embeddings[0], list): @@ -5043,6 +5145,7 @@ def _collection_count(self, collection_id: str | None, collection_name: str) -> @staticmethod def _rewrite_where_for_ns(where: dict[str, Any] | None) -> dict[str, Any] | None: + """Rewrite a WHERE clause to target namespace-scoped columns.""" if where is None: return None rewritten = {} @@ -5062,6 +5165,7 @@ def _append_namespace_filter( namespace_id: int, ltable_id: int, ) -> tuple[str, list[Any]]: + """Append the namespace id filter to a WHERE clause.""" ns_cond = f"namespace_id = {namespace_id} AND ltable_id = {ltable_id}" if not where_clause: return f"WHERE {ns_cond}", params @@ -5077,6 +5181,7 @@ def _count_namespace_records_by_id( ltable_id: int, record_id: str, ) -> int: + """Count namespace records matching the given id.""" id_expr = _NS_DATA_CONTENT_ID_EXPR sql = ( f"SELECT COUNT(*) AS cnt FROM `{table_name}` " @@ -5102,6 +5207,7 @@ def _delete_namespace_records_by_id( ltable_id: int, record_id: str, ) -> None: + """Delete namespace records matching the given id.""" id_expr = _NS_DATA_CONTENT_ID_EXPR sql = ( f"DELETE FROM `{table_name}` " @@ -5183,6 +5289,7 @@ def _namespace_add( # noqa: C901 embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + """Add records to a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5218,12 +5325,15 @@ def _namespace_add( # noqa: C901 " 1. Provide embeddings directly when calling add(), or\n" " 2. Provide embedding_function to auto-generate embeddings from documents." ) + elif metadatas: + pass else: raise ValueError( - "Neither embeddings nor documents provided. " - "Please provide either:\n" - " 1. embeddings directly, or\n" - " 2. documents with embedding_function to generate embeddings." + "Neither embeddings, documents, nor metadatas provided. " + "Please provide at least one of:\n" + " 1. embeddings directly,\n" + " 2. documents with embedding_function to generate embeddings, or\n" + " 3. metadatas for metadata-only add." ) num_items = len(ids) @@ -5274,6 +5384,7 @@ def _namespace_update( embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + """Update existing records in a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5419,6 +5530,7 @@ def _namespace_upsert( embedding_function: EmbeddingFunction[EmbeddingDocuments] | None = None, **kwargs, ) -> None: + """Insert or update records in a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5520,6 +5632,7 @@ def _namespace_delete( where_document: dict[str, Any] | None = None, **kwargs, ) -> None: + """Delete records from a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5586,6 +5699,7 @@ def _namespace_query( # noqa: C901 include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + """Run a vector query against a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5685,6 +5799,7 @@ def _namespace_get( # noqa: C901 include: list[str] | None = None, **kwargs, ) -> dict[str, Any]: + """Fetch records from a namespace collection by id/where/pagination.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5730,6 +5845,8 @@ def _namespace_get( # noqa: C901 where_str, params = self._append_namespace_filter(user_where, params, ns_id, ltable_id) select_clause = ", ".join(select_parts) sql = f"SELECT {select_clause} FROM `{table_name}` {where_str}" + if limit is None and offset is not None: + limit = 100 if limit is not None: sql += f" LIMIT {int(limit)}" if offset is not None: @@ -5800,6 +5917,7 @@ def _namespace_count( namespace_name: str, **kwargs, ) -> int: + """Count records in a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5829,6 +5947,7 @@ def _namespace_peek( limit: int = 10, **kwargs, ) -> dict[str, Any]: + """Return a small sample of records from a namespace collection.""" return self._namespace_get( collection_id=collection_id, collection_name=collection_name, @@ -5840,12 +5959,14 @@ def _namespace_peek( ) def _adapt_search_parm_for_ns(self, search_parm: dict[str, Any], ns_id: int, lt_id: int) -> dict[str, Any]: + """Adapt search parameters to the namespace-scoped table layout.""" ns_filter = [ {"term": {"namespace_id": ns_id}}, {"term": {"ltable_id": int(lt_id)}}, ] def _rewrite_field_refs(obj): + """Rewrite field references to their namespace-scoped equivalents.""" if isinstance(obj, dict): new_dict = {} for k, v in obj.items(): @@ -5879,6 +6000,7 @@ def _rewrite_field_refs(obj): search_parm["_source"] = new_source def _inject_filter_into_knn(node): + """Inject the namespace filter into a KNN search expression.""" if isinstance(node, dict): if "filter" in node: existing = node["filter"] @@ -5898,6 +6020,7 @@ def _inject_filter_into_knn(node): _scoring_leaf_keys = {"query_string", "match", "multi_match", "match_phrase"} def _inject_filter_into_query(node): + """Inject the namespace filter into a query expression.""" if not isinstance(node, dict): return node if "bool" in node: @@ -5949,6 +6072,7 @@ def _namespace_hybrid_search( query_hint: QueryHint | None = None, **kwargs, ) -> dict[str, Any]: + """Run a hybrid (vector + fulltext) search against a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5998,6 +6122,7 @@ def _namespace_hybrid_search( def _transform_ns_hybrid_result( self, result_rows: list[dict[str, Any]], include: list[str] | None ) -> dict[str, Any]: + """Transform raw hybrid-search rows into the public result shape.""" if not result_rows: return { "ids": [[]], "distances": [[]], "metadatas": [[]], @@ -6074,4 +6199,5 @@ def _namespace_prewarm( namespace_name: str, **kwargs, ) -> None: + """Prewarm the namespace logical table to reduce first-query latency.""" raise NotImplementedError("prewarm is not supported in this client mode") diff --git a/src/pyseekdb/client/client_seekdb_embedded.py b/src/pyseekdb/client/client_seekdb_embedded.py index 293d4846..fb714e40 100644 --- a/src/pyseekdb/client/client_seekdb_embedded.py +++ b/src/pyseekdb/client/client_seekdb_embedded.py @@ -105,6 +105,7 @@ def get_raw_connection(self) -> Any: # seekdb.Connection @property def mode(self) -> str: + """Return the client mode identifier.""" return "SeekdbEmbeddedClient" def _use_context_manager_for_cursor(self) -> bool: @@ -278,8 +279,10 @@ def _namespace_prewarm( namespace_name: str, **kwargs, ) -> None: + """Prewarm the namespace logical table to reduce first-query latency.""" raise ValueError("prewarm is only supported in shared-storage remote deployment") def __repr__(self): + """Return the developer-readable representation.""" status = "connected" if self.is_connected() else "disconnected" return f"" diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 75cff466..5fb19205 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -80,8 +80,10 @@ def _ensure_connection(self) -> pymysql.Connection: **self.kwargs, ) logger.info(f"✅ Connected to remote server: {self.host}:{self.port}/{self.database}") - with contextlib.suppress(Exception): + try: self._use_catalog_database() + except Exception as exc: + logger.warning("Failed to initialize catalog database on connect: %s", exc) return self._connection @@ -102,6 +104,7 @@ def get_raw_connection(self) -> pymysql.Connection: @property def mode(self) -> str: + """Return the client mode identifier.""" return "RemoteServerClient" # ==================== Collection Management (framework) ==================== @@ -197,6 +200,7 @@ def list_databases( return super().list_databases(limit=limit, offset=offset, tenant=tenant) def _database_tenant(self, tenant: str) -> str | None: + """Return the tenant associated with the active database.""" if tenant != self.tenant and tenant != DEFAULT_TENANT: logger.warning( f"Specified tenant '{tenant}' differs from client tenant '{self.tenant}', using client tenant" @@ -211,7 +215,9 @@ def _namespace_prewarm( namespace_name: str, **kwargs, ) -> None: + """Prewarm the namespace logical table to reduce first-query latency.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) + self._use_catalog_database() self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, ) @@ -219,5 +225,6 @@ def _namespace_prewarm( self._execute(sql) def __repr__(self): + """Return the developer-readable representation.""" status = "connected" if self.is_connected() else "disconnected" return f"" diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 5c38157b..80063f83 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Optional -from .validators import _MAX_N_RESULTS, _validate_namespace_name +from .validators import _MAX_N_RESULTS, _validate_namespace_name, _validate_n_results if TYPE_CHECKING: from .embedding_function import Documents as EmbeddingDocuments @@ -121,17 +121,6 @@ def partition_count(self) -> int | None: """ return self._partition_count - @staticmethod - def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> None: - """Validate ``n_results`` is a positive integer within the engine limit.""" - if not isinstance(n_results, int) or n_results < 1: - raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") - if n_results > max_results: - raise ValueError( - f"n_results must be <= {max_results}, got {n_results}. " - "Use a smaller value or paginate with offset/limit." - ) - def _guard_collection_data_api(self) -> None: """Raise if collection-level DML/DQL is used on a namespace-enabled collection.""" if self._use_namespace: @@ -506,7 +495,7 @@ def query( ) """ self._guard_collection_data_api() - self._validate_n_results(n_results) + _validate_n_results(n_results) return self._client._collection_query( collection_id=self._id, collection_name=self._name, diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py index 3e726bef..d4739dde 100644 --- a/src/pyseekdb/client/namespace.py +++ b/src/pyseekdb/client/namespace.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any -from .validators import _MAX_N_RESULTS +from .validators import _validate_n_results if TYPE_CHECKING: from .collection import Collection @@ -67,17 +67,6 @@ def _guard_exists(self) -> None: "Operations are not allowed on a deleted namespace." ) - @staticmethod - def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> None: - """Validate ``n_results`` is a positive integer within the engine limit.""" - if not isinstance(n_results, int) or n_results < 1: - raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") - if n_results > max_results: - raise ValueError( - f"n_results must be <= {max_results}, got {n_results}. " - "Use a smaller value or paginate with offset/limit." - ) - # ==================== DML Operations ==================== def add( @@ -183,7 +172,7 @@ def query( ) -> dict[str, Any]: """Run vector similarity search within this namespace.""" self._guard_exists() - self._validate_n_results(n_results) + _validate_n_results(n_results) return self._client._namespace_query( collection_id=self._collection.id, collection_name=self._collection.name, @@ -211,7 +200,7 @@ def hybrid_search( ) -> dict[str, Any]: """Run hybrid fulltext + vector search within this namespace.""" self._guard_exists() - self._validate_n_results(n_results) + _validate_n_results(n_results) if include is None and not query and not knn: include = [] return self._client._namespace_hybrid_search( diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py index 8c7515af..d0d6c787 100644 --- a/src/pyseekdb/client/validators.py +++ b/src/pyseekdb/client/validators.py @@ -28,6 +28,37 @@ def _validate_namespace_name(name: str) -> None: ) +def _validate_database_name(name: str) -> None: + """Validate a SQL database identifier for catalog table qualification.""" + if not isinstance(name, str): + raise TypeError( + f"Invalid database name: '{name}'. Database name must be a string, got {type(name).__name__}" + ) + if not name: + raise ValueError(f"Invalid database name: '{name}'. Database name must not be empty") + if _NAME_PATTERN.match(name) is None: + raise ValueError( + f"Invalid database name: '{name}'. Database name contains invalid characters. " + "Only letters, digits, and underscore are allowed: [a-zA-Z0-9_]" + ) + + +def _quote_sql_identifier(identifier: str) -> str: + """Quote a SQL identifier and escape embedded backticks.""" + return "`" + identifier.replace("`", "``") + "`" + + +def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> None: + """Validate ``n_results`` is a positive integer within the engine limit.""" + if not isinstance(n_results, int) or n_results < 1: + raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") + if n_results > max_results: + raise ValueError( + f"n_results must be <= {max_results}, got {n_results}. " + "Use a smaller value or paginate with offset/limit." + ) + + def _validate_record_ids(ids: list[str]) -> None: """Validate namespace/collection record id list shape and per-id constraints.""" if not isinstance(ids, list): diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index 8dcd4195..62245c36 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -729,8 +729,10 @@ def _worker(tag, c): ta = threading.Thread(target=_worker, args=("A", client_a)) tb = threading.Thread(target=_worker, args=("B", client_b)) - ta.start(); tb.start() - ta.join(timeout=30); tb.join(timeout=30) + ta.start() + tb.start() + ta.join(timeout=30) + tb.join(timeout=30) # Both calls must succeed (one does the work, the other no-ops). assert results.get("A") == "ok", f"A failed: {results.get('A')!r}" diff --git a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py index 80097b18..3fbc0243 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py @@ -17,6 +17,7 @@ run_hybrid_search_fts_case_all_quadrants, ) from namespace_hybrid_search_helpers import ( + assert_hybrid_search_index_no_hits, get_hybrid_combined_case, get_search_index_case, get_vector_knn_case, @@ -43,8 +44,6 @@ def test_cross_quadrant_search_index_isolation(self, db_client): if key == "c1_x": continue _, namespace = ctx[key] - from namespace_hybrid_search_helpers import assert_hybrid_search_index_no_hits - assert_hybrid_search_index_no_hits(namespace, case.where, n_results=case.n_results) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py index d60ba04b..ddfaa52c 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py @@ -102,7 +102,7 @@ def test_hybrid_search_triple_branch_fts_hits_with_knn_active(self, db_client): ids = result["ids"][0] assert len(ids) > 0, "expected at least one hybrid FTS+KNN hit" docs = result.get("documents", [[]])[0] - for doc_id, doc_text in zip(ids, docs): + for doc_id, doc_text in zip(ids, docs, strict=True): assert TOKEN_ZPX.lower() in (doc_text or "").lower(), ( f"id={doc_id!r} must contain {TOKEN_ZPX!r}" ) diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 33b57abd..ef7d9e21 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -224,7 +224,7 @@ def test_namespace_on_non_namespace_collection_raises(self, db_client): finally: try: db_client.delete_collection(name=name) - except Exception: + except (ValueError, RuntimeError): db_client._server._execute(f"DROP TABLE IF EXISTS `{name}`") db_client._server._execute( f"DELETE FROM `sdk_collections` WHERE COLLECTION_NAME = '{name}'" diff --git a/tests/integration_tests/test_namespace_upsert_concurrency.py b/tests/integration_tests/test_namespace_upsert_concurrency.py index 46e64269..c0f27fa6 100644 --- a/tests/integration_tests/test_namespace_upsert_concurrency.py +++ b/tests/integration_tests/test_namespace_upsert_concurrency.py @@ -50,7 +50,7 @@ def test_multi_client_concurrent_upsert_same_new_id_should_keep_single_record( def _upsert_once(target_ns, idx: int) -> None: try: - barrier.wait() + barrier.wait(timeout=30) target_ns.upsert( ids="same_new_id", embeddings=[1.0, 2.0, 3.0], @@ -67,7 +67,8 @@ def _upsert_once(target_ns, idx: int) -> None: for thread in threads: thread.start() for thread in threads: - thread.join() + thread.join(timeout=120) + assert all(not thread.is_alive() for thread in threads), "upsert worker thread timed out" assert errors == [] result = ns.get(ids="same_new_id", include=["documents", "metadatas"]) diff --git a/tests/unit_tests/test_get_or_create_collection_concurrency.py b/tests/unit_tests/test_get_or_create_collection_concurrency.py index 6defdbd7..89d259f1 100644 --- a/tests/unit_tests/test_get_or_create_collection_concurrency.py +++ b/tests/unit_tests/test_get_or_create_collection_concurrency.py @@ -99,9 +99,16 @@ def test_detects_conflict_in_cause_chain(self): class TestGetOrCreateCollectionRecovery: + @staticmethod + def _bind_resume_helper(client): + client._get_or_resume_existing_collection = ( + BaseClient._get_or_resume_existing_collection.__get__(client, BaseClient) + ) + def test_returns_existing_collection_after_create_conflict(self): client = MagicMock(spec=BaseClient) - client.has_collection.side_effect = [False, True] + self._bind_resume_helper(client) + client.has_collection.return_value = False existing = object() client.get_collection.return_value = existing client.create_collection.side_effect = ValueError("Collection 'items' already exists") @@ -113,7 +120,8 @@ def test_returns_existing_collection_after_create_conflict(self): def test_retries_get_after_wrapped_table_conflict(self): client = MagicMock(spec=BaseClient) - client.has_collection.side_effect = [False, True] + self._bind_resume_helper(client) + client.has_collection.return_value = False existing = object() client.get_collection.return_value = existing inner = Exception("Table 'c$v2$abc' already exists failed: code=1050") @@ -126,6 +134,38 @@ def test_retries_get_after_wrapped_table_conflict(self): assert result is existing client.get_collection.assert_called_once_with("items", embedding_function=_NOT_PROVIDED) + def test_conflict_on_namespace_collection_resumes_incomplete_handle(self): + client = MagicMock(spec=BaseClient) + self._bind_resume_helper(client) + resumed = object() + client.has_collection.return_value = False + client._get_ns_collection_meta.return_value = { + "collection_id": "abc", + "collection_name": "items", + "settings": {"use_namespace": True}, + } + client._is_incomplete_ns_collection.return_value = True + client.create_collection.side_effect = [ + ValueError("Collection 'items' already exists"), + resumed, + ] + + result = BaseClient.get_or_create_collection(client, "items", use_namespace=True) + + assert result is resumed + assert client.create_collection.call_count == 2 + client.get_collection.assert_not_called() + + def test_conflict_on_namespace_collection_without_metadata_reraises(self): + client = MagicMock(spec=BaseClient) + self._bind_resume_helper(client) + client.has_collection.return_value = False + client._get_ns_collection_meta.return_value = None + client.create_collection.side_effect = ValueError("Collection 'items' already exists") + + with pytest.raises(ValueError, match="namespace metadata is missing"): + BaseClient.get_or_create_collection(client, "items", use_namespace=True) + def test_resumes_incomplete_namespace_collection_when_present(self): client = MagicMock(spec=BaseClient) resumed = object() diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 2c5c7bfa..aa917ade 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -1048,17 +1048,19 @@ def test_ensure_namespace_catalogs_creates_all_catalog_tables(self): assert "ltable_id BIGINT UNSIGNED NOT NULL" in sql assert "included_index BOOL" in sql assert "PRIMARY KEY (namespace_id, ltable_id, included_index)" in sql - assert "PARTITION BY KEY(namespace_id) PARTITIONS 1000" in sql + assert "PARTITION BY KEY(namespace_id) PARTITIONS 8" in sql def test_ensure_namespace_catalogs_creates_catalog_tables_in_order(self): c = FakeClient() c._ensure_namespace_catalogs() - assert len(c.executed_sqls) == 5 + assert len(c.executed_sqls) == 7 assert c.executed_sqls[0] == "USE `test`" assert "`test`.`sdk_namespaces`" in c.executed_sqls[2] assert "`test`.`sdk_ltables`" in c.executed_sqls[3] assert "`test`.`sdk_namespaces_stats`" in c.executed_sqls[4] + assert "CREATE UNIQUE INDEX uk_sdk_ns_coll_name" in c.executed_sqls[5] + assert "CREATE UNIQUE INDEX uk_sdk_lt_coll_ns_name" in c.executed_sqls[6] def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): c = FakeClient() @@ -1097,8 +1099,10 @@ def test_ob_type_validation(self): c._create_namespace_collection("test", schema) def test_create_namespace_collection_without_ivf_skips_vector_index(self): + from pyseekdb.client.version import Version + c = FakeClient() - c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", "4.3")) + c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", Version("4.6.1.0"))) c._is_shared_storage_mode = MagicMock(return_value=False) c._create_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) c._ensure_namespace_catalogs = MagicMock() @@ -1117,8 +1121,10 @@ def test_create_namespace_collection_without_ivf_skips_vector_index(self): assert "centroids_fresh_mode" not in settings def test_create_namespace_collection_ivf_without_centroids_fresh_mode(self): + from pyseekdb.client.version import Version + c = FakeClient() - c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", "4.3")) + c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", Version("4.6.1.0"))) c._is_shared_storage_mode = MagicMock(return_value=False) c._create_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) c._ensure_namespace_catalogs = MagicMock() From 4d4d0dcda984cac4149734036de6f9c7549f04c9 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Wed, 24 Jun 2026 17:10:43 +0800 Subject: [PATCH 62/84] =?UTF-8?q?fix=EF=BC=9A=E3=80=90pyseekdb=E3=80=91?= =?UTF-8?q?=E5=AF=B9=E4=BA=8E=E6=B2=A1=E6=9C=89=E7=89=A9=E7=90=86=E8=A1=A8?= =?UTF-8?q?=E7=9A=84collection=EF=BC=8C=E4=BB=8D=E6=97=A7=E5=8F=AF?= =?UTF-8?q?=E4=BB=A5get=E5=88=B0=EF=BC=8C=E5=B9=B6=E4=B8=94=E8=83=BD?= =?UTF-8?q?=E5=A4=9F=E5=88=9B=E5=BB=BAnamespace=20https://project.alipay.c?= =?UTF-8?q?om/project/W24001004732/P26001026476/workitemDetail=3FopenWorkI?= =?UTF-8?q?temId=3D2026062400116954267&status=3Dstatus?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 73 ++++++++++++++++--- .../test_namespace_lifecycle.py | 37 ++++++++++ tests/unit_tests/test_namespace.py | 48 ++++++++++++ 3 files changed, 146 insertions(+), 12 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 01867da2..6304fd84 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1560,19 +1560,39 @@ def _table_exists(self, table_name: str) -> bool: except Exception: return False - def _ns_missing_physical_tables(self, collection_id: str, is_shared_storage: bool) -> list[str]: - """Return the namespace physical tables that are expected but absent. + def _tablegroup_exists(self, tablegroup_name: str) -> bool: + """Whether `tablegroup_name` exists in the current OceanBase tenant.""" + name_escaped = escape_string(tablegroup_name) + try: + rows = self._execute( + "SELECT 1 FROM oceanbase.DBA_OB_TABLEGROUPS " + f"WHERE TABLEGROUP_NAME = '{name_escaped}'" + ) + return bool(rows) + except Exception: + return False - An empty list means the collection's physical tables are complete. - """ - expected = [ - NamespaceCollectionNames.data_table_name(collection_id), - NamespaceCollectionNames.kv_data_table_name(collection_id), - NamespaceCollectionNames.logic_schema_table_name(collection_id), + + def _ns_missing_physical_resources(self, collection_id: str, is_shared_storage: bool) -> list[str]: + """Return expected tablegroup/tables that are absent for this namespace collection.""" + resources: list[tuple[str, bool]] = [ + (NamespaceCollectionNames.tablegroup_name(collection_id), True), + (NamespaceCollectionNames.data_table_name(collection_id), False), + (NamespaceCollectionNames.kv_data_table_name(collection_id), False), + (NamespaceCollectionNames.logic_schema_table_name(collection_id), False), ] if is_shared_storage: - expected.append(NamespaceCollectionNames.hot_table_name(collection_id)) - return [t for t in expected if not self._table_exists(t)] + resources.append((NamespaceCollectionNames.hot_table_name(collection_id), False)) + missing: list[str] = [] + for resource_name, is_tablegroup in resources: + exists = ( + self._tablegroup_exists(resource_name) + if is_tablegroup + else self._table_exists(resource_name) + ) + if not exists: + missing.append(resource_name) + return missing def _is_incomplete_ns_collection(self, name: str) -> bool: """Whether `name` is a namespace collection whose catalog row exists but @@ -1584,7 +1604,33 @@ def _is_incomplete_ns_collection(self, name: str) -> bool: return False is_ss = meta.get("settings", {}).get("storage_mode") == "ss" self._use_catalog_database() - return len(self._ns_missing_physical_tables(meta["collection_id"], is_ss)) > 0 + return len(self._ns_missing_physical_resources(meta["collection_id"], is_ss)) > 0 + + def _purge_broken_ns_collection_if_incomplete( + self, collection_name: str, meta: dict | None = None + ) -> bool: + """Purge a namespace collection whose catalog row exists but physical resources are incomplete. + + Returns True when the collection was purged and should be treated as non-existent. + """ + if meta is None: + meta = self._get_ns_collection_meta(collection_name) + if meta is None: + return False + self._use_catalog_database() + is_ss = meta.get("settings", {}).get("storage_mode") == "ss" + missing = self._ns_missing_physical_resources(meta["collection_id"], is_ss) + if not missing: + return False + logger.warning( + "Namespace collection '%s' (collection_id=%s) is missing physical resources %s; " + "purging catalog metadata and cleaning up leftovers.", + collection_name, + meta["collection_id"], + missing, + ) + self._delete_ns_collection_meta(collection_name) + return True def _resolve_namespace_ltable_id( self, collection_id: str, namespace_id: int @@ -1903,7 +1949,10 @@ def get_collection(self, name: str, embedding_function: EmbeddingFunctionParam = except Exception: pass if ns_meta is not None: - return self._build_ns_collection_from_meta(ns_meta, embedding_function) + if self._purge_broken_ns_collection_if_incomplete(collection_name=name, meta=ns_meta): + ns_meta = None + else: + return self._build_ns_collection_from_meta(ns_meta, embedding_function) try: collection = self._get_collection_v1(name, embedding_function) except ValueError as e: diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index ef7d9e21..07097976 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -3,6 +3,7 @@ Tests collection creation with use_namespace=True, namespace CRUD, and collection deletion. """ +import contextlib import time import pytest @@ -190,6 +191,42 @@ def test_create_resumes_incomplete_collection(self, oceanbase_client): finally: client.delete_collection(name=name) + def test_get_collection_purges_broken_ns_collection(self, oceanbase_client): + """get_collection on a namespace collection with missing physical tables + should treat it as non-existent and purge catalog leftovers.""" + from pyseekdb.client.meta_info import NamespaceCollectionNames + + client = oceanbase_client + srv = client._server + name = f"test_ns_broken_get_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + collection = client.create_collection( + name=name, schema=schema, use_namespace=True, partition_count=4 + ) + cid = collection.id + try: + srv._use_catalog_database() + srv._execute(f"DROP TABLE IF EXISTS `{NamespaceCollectionNames.kv_data_table_name(cid)}`") + assert srv._is_incomplete_ns_collection(name) is True + + with pytest.raises(ValueError, match="does not exist"): + client.get_collection(name) + + assert client.has_collection(name) is False + assert srv._get_ns_collection_meta(name) is None + assert not srv._table_exists(NamespaceCollectionNames.data_table_name(cid)) + assert not srv._table_exists(NamespaceCollectionNames.kv_data_table_name(cid)) + assert not srv._table_exists(NamespaceCollectionNames.logic_schema_table_name(cid)) + assert not srv._tablegroup_exists(NamespaceCollectionNames.tablegroup_name(cid)) + finally: + with contextlib.suppress(Exception): + client.delete_collection(name=name) + def test_namespace_ops_blocked_after_collection_deleted(self, db_client): collection = self._create_ns_collection(db_client) db_client.delete_collection(name=collection.name) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index aa917ade..e2931698 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -1075,6 +1075,54 @@ def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): assert any("WHERE collection_id = 'abc123'" in s for s in calls) +class TestBrokenNsCollectionPurge: + + def test_purge_broken_ns_collection_if_incomplete_calls_delete(self): + c = FakeClient() + meta = { + "collection_id": "cid1", + "collection_name": "coll", + "settings": {"use_namespace": True, "storage_mode": "sn"}, + } + c._use_catalog_database = MagicMock() + c._ns_missing_physical_resources = MagicMock(return_value=["cid1_logic_data_table"]) + c._delete_ns_collection_meta = MagicMock() + + assert c._purge_broken_ns_collection_if_incomplete("coll", meta=meta) is True + c._delete_ns_collection_meta.assert_called_once_with("coll") + + def test_purge_skips_complete_collection(self): + c = FakeClient() + meta = {"collection_id": "cid1", "collection_name": "coll", "settings": {"storage_mode": "sn"}} + c._use_catalog_database = MagicMock() + c._ns_missing_physical_resources = MagicMock(return_value=[]) + c._delete_ns_collection_meta = MagicMock() + + assert c._purge_broken_ns_collection_if_incomplete("coll", meta=meta) is False + c._delete_ns_collection_meta.assert_not_called() + + def test_get_collection_purges_incomplete_namespace_collection(self): + class GetClient(FakeClient): + get_collection = BaseClient.get_collection + + c = GetClient() + meta = { + "collection_id": "cid1", + "collection_name": "coll", + "settings": {"use_namespace": True, "storage_mode": "sn", "dimension": 3}, + } + c._get_ns_collection_meta = MagicMock(return_value=meta) + c._purge_broken_ns_collection_if_incomplete = MagicMock(return_value=True) + c._get_collection_v1 = MagicMock(side_effect=ValueError("not v1")) + c._get_collection_v2 = MagicMock(side_effect=ValueError("Collection 'coll' does not exist")) + + with pytest.raises(ValueError, match="does not exist"): + c.get_collection("coll") + c._purge_broken_ns_collection_if_incomplete.assert_called_once_with( + collection_name="coll", meta=meta + ) + + # ==================== UseNamespace Validation Tests ==================== From 9139b9ecda7f197602bc1d2a9364677d20a875a8 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 25 Jun 2026 10:21:32 +0800 Subject: [PATCH 63/84] fix:review --- src/pyseekdb/client/client_base.py | 29 +++++++++++++++++-- src/pyseekdb/client/validators.py | 2 +- tests/unit_tests/test_namespace.py | 10 +++++++ .../test_namespace_upsert_reconcile.py | 23 ++++++++++++++- 4 files changed, 59 insertions(+), 5 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 6304fd84..fedf20e7 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -135,6 +135,19 @@ def _is_sdk_collection_catalog_conflict_error(exc: BaseException) -> bool: return False +def _reraise_unless_unique_index_exists(exc: BaseException) -> None: + """Re-raise unless the exception indicates the unique index is already present.""" + message = str(exc).lower() + if ( + "already exists" in message + or "duplicate key name" in message + or "code=1061" in message + or "1061" in message and "duplicate" in message + ): + return + raise exc + + def _extract_hnsw_config(config: ConfigurationParam) -> HNSWConfiguration | None: """Return the HNSW config from a Configuration/HNSWConfiguration, or None.""" if config is None: @@ -1188,10 +1201,12 @@ def _create_sdk_collections_if_not_exists(self) -> None: UNIQUE KEY uk_sdk_coll_name (collection_name) ) COMMENT='Settings of collections created by SDK' ORGANIZATION INDEX {scp};""" self._execute(create_table_sql) - with contextlib.suppress(Exception): + try: self._execute( f"CREATE UNIQUE INDEX uk_sdk_coll_name ON {sdk_coll} (collection_name)" ) + except Exception as exc: + _reraise_unless_unique_index_exists(exc) except Exception as e: raise ValueError(f"Failed to create sdk_collections table: {e}") from e @@ -1328,16 +1343,20 @@ def _ensure_namespace_catalogs(self) -> None: self._execute(ns_namespaces_sql) self._execute(ns_ltables_sql) self._execute(namespaces_stats_sql) - with contextlib.suppress(Exception): + try: self._execute( f"CREATE UNIQUE INDEX uk_sdk_ns_coll_name ON {ns_namespaces_q} " f"(collection_id, namespace_name)" ) - with contextlib.suppress(Exception): + except Exception as exc: + _reraise_unless_unique_index_exists(exc) + try: self._execute( f"CREATE UNIQUE INDEX uk_sdk_lt_coll_ns_name ON {ns_ltables_q} " f"(collection_id, namespace_id, ltable_name)" ) + except Exception as exc: + _reraise_unless_unique_index_exists(exc) def _rollback_connection_if_supported(self) -> None: """Roll back the current connection transaction if the backend supports it.""" @@ -5324,6 +5343,10 @@ def _reconcile_namespace_duplicate_records( ) if attempt < 119: time.sleep(0.05 * min(attempt + 1, 10)) + else: + raise ValueError( + f"Failed to reconcile duplicate namespace rows for record_id={record_id!r}" + ) def _namespace_add( # noqa: C901 self, diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py index d0d6c787..78434711 100644 --- a/src/pyseekdb/client/validators.py +++ b/src/pyseekdb/client/validators.py @@ -50,7 +50,7 @@ def _quote_sql_identifier(identifier: str) -> str: def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> None: """Validate ``n_results`` is a positive integer within the engine limit.""" - if not isinstance(n_results, int) or n_results < 1: + if isinstance(n_results, bool) or not isinstance(n_results, int) or n_results < 1: raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") if n_results > max_results: raise ValueError( diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index e2931698..251c29dc 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -21,6 +21,7 @@ from pyseekdb.client.collection import Collection # noqa: E402 from pyseekdb.client.configuration import HNSWConfiguration, IVFIndexType, VectorIndexConfig # noqa: E402 from pyseekdb.client.namespace import Namespace # noqa: E402 +from pyseekdb.client.validators import _validate_n_results # noqa: E402 # ==================== IVFConfiguration Tests ==================== @@ -1123,6 +1124,15 @@ class GetClient(FakeClient): ) +class TestValidateNResults: + + def test_rejects_boolean_values(self): + with pytest.raises(ValueError, match="n_results must be an integer"): + _validate_n_results(True) + with pytest.raises(ValueError, match="n_results must be an integer"): + _validate_n_results(False) + + # ==================== UseNamespace Validation Tests ==================== diff --git a/tests/unit_tests/test_namespace_upsert_reconcile.py b/tests/unit_tests/test_namespace_upsert_reconcile.py index 6ddf7018..b03e2b99 100644 --- a/tests/unit_tests/test_namespace_upsert_reconcile.py +++ b/tests/unit_tests/test_namespace_upsert_reconcile.py @@ -4,7 +4,7 @@ import sys from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -67,6 +67,27 @@ def test_reconcile_collapses_duplicate_rows(self): assert add_kwargs["metadatas"] == [{"client": 2}] assert add_kwargs["embeddings"] == [[1.0, 2.0, 3.0]] + def test_reconcile_raises_when_retries_exhausted(self): + client = MagicMock(spec=BaseClient) + client._count_namespace_records_by_id.return_value = 4 + + with patch("pyseekdb.client.client_base.time.sleep"): + with pytest.raises(ValueError, match="Failed to reconcile duplicate namespace rows"): + BaseClient._reconcile_namespace_duplicate_records( + client, + collection_id="c" * 32, + collection_name="items", + namespace_id="7", + namespace_name="race_ns", + ltable_id=9, + table_name="logic_data_table", + ids=["same_new_id"], + documents=["doc"], + metadatas=None, + embeddings=None, + embedding_function=None, + ) + if __name__ == "__main__": pytest.main([__file__, "-v"]) From dde32f12c90c3f2a350fb3f0770659bba20d9e4e Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 25 Jun 2026 11:16:06 +0800 Subject: [PATCH 64/84] docstring_add --- src/pyseekdb/client/schema.py | 2 + .../namespace_dml_helpers.py | 14 ++ .../namespace_fts_helpers.py | 16 ++ .../namespace_hybrid_search_helpers.py | 29 +++ ...st_get_or_create_collection_concurrency.py | 5 + ...t_get_or_create_collection_multiprocess.py | 36 ++++ .../test_logic_table_monitoring_p1.py | 5 + .../test_namespace_constraints.py | 5 + tests/integration_tests/test_namespace_dml.py | 27 +++ .../test_namespace_dml_multi_coll.py | 3 + .../test_namespace_dml_multi_coll_multi_ns.py | 4 + .../test_namespace_dml_multi_ns.py | 6 + .../test_namespace_drop_validation.py | 24 +++ .../test_namespace_get_delete_filters.py | 6 + .../test_namespace_hybrid_search_combined.py | 5 + .../test_namespace_hybrid_search_fulltext.py | 4 + ...rid_search_fulltext_multi_coll_multi_ns.py | 9 + ...space_hybrid_search_multi_coll_multi_ns.py | 8 + ...st_namespace_hybrid_search_search_index.py | 3 + ...t_namespace_hybrid_search_triple_branch.py | 4 + ...earch_triple_branch_multi_coll_multi_ns.py | 9 + .../test_namespace_hybrid_search_vector.py | 5 + .../test_namespace_lifecycle.py | 17 ++ .../test_namespace_lifecycle_concurrency.py | 7 + .../test_namespace_optional_indexes.py | 11 ++ .../test_namespace_prewarm.py | 3 + .../test_namespace_prewarm_lob.py | 7 + .../integration_tests/test_namespace_query.py | 14 ++ .../test_namespace_reconnect.py | 1 + .../test_namespace_session_vars.py | 6 + .../test_namespace_upsert_concurrency.py | 4 + tests/unit_tests/test_build_document_query.py | 6 + tests/unit_tests/test_build_knn_filter.py | 3 + .../unit_tests/test_build_query_expression.py | 5 + ...st_get_or_create_collection_concurrency.py | 26 +++ ...est_get_or_create_namespace_concurrency.py | 9 + tests/unit_tests/test_namespace.py | 167 ++++++++++++++++++ .../test_namespace_lifecycle_lock.py | 3 + .../test_namespace_upsert_reconcile.py | 4 + 39 files changed, 522 insertions(+) diff --git a/src/pyseekdb/client/schema.py b/src/pyseekdb/client/schema.py index ce99adff..41baa0fd 100644 --- a/src/pyseekdb/client/schema.py +++ b/src/pyseekdb/client/schema.py @@ -88,6 +88,7 @@ def __init__( sparse_vector_index: SparseVectorIndexConfig | None = None, fulltext_index: FulltextIndexConfig | None = None, ): + """Init.""" if isinstance(vector_index, VectorIndexConfig): self.vector_index = vector_index elif isinstance(vector_index, HNSWConfiguration): @@ -149,6 +150,7 @@ def create_index(self, config: Any) -> Schema: return self def __repr__(self) -> str: + """Repr.""" parts = [] if self.vector_index is not None: parts.append(f"vector_index={self.vector_index}") diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index 3f130b95..95ceaeec 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -30,6 +30,7 @@ @dataclass(frozen=True) class DmlRecord: + """DmlRecord class.""" doc_id: str embedding: list[float] document: str @@ -37,6 +38,7 @@ class DmlRecord: def ns_schema() -> Schema: + """Ns schema.""" return Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), @@ -47,6 +49,7 @@ def ns_schema() -> Schema: def create_ns_collection(client: Any, suffix: str = "") -> Any: + """Create ns collection.""" name = f"test_ns_dml_{int(time.time() * 1000)}{suffix}" return client.create_collection( name=name, schema=ns_schema(), use_namespace=True, @@ -55,6 +58,7 @@ def create_ns_collection(client: Any, suffix: str = "") -> Any: def cleanup(client: Any, *collections: Any) -> None: + """Cleanup.""" for collection in collections: try: client.delete_collection(name=collection.name) @@ -86,10 +90,12 @@ def build_large_dml_corpus( def corpus_id_set(corpus: list[DmlRecord]) -> set[str]: + """Corpus id set.""" return {r.doc_id for r in corpus} def _filter_embedding(index: int) -> list[float]: + """Filter embedding.""" return [ (index % 97) / 97.0, ((index * 3) % 89) / 89.0, @@ -193,6 +199,7 @@ def assert_get_where_document_count( substring: str | None = None, context: str = "", ) -> dict[str, Any]: + """Assert get where document count.""" kwargs: dict[str, Any] = {"where_document": where_document, "include": ["documents", "metadatas"]} if limit is not None: kwargs["limit"] = limit @@ -248,6 +255,7 @@ def assert_count_after_document_delete( expected_removed: int, pre_delete_hit_count: int, ) -> None: + """Assert count after document delete.""" actual = ns.count() if actual == expected_count: return @@ -279,6 +287,7 @@ def insert_dml_corpus_in_batches( corpus: list[DmlRecord], batch_size: int = DML_BATCH_SIZE, ) -> None: + """Insert dml corpus in batches.""" for start in range(0, len(corpus), batch_size): chunk = corpus[start : start + batch_size] namespace.add( @@ -295,6 +304,7 @@ def load_dml_corpus( *, assert_count: bool = True, ) -> None: + """Load dml corpus.""" insert_dml_corpus_in_batches(namespace, corpus) if assert_count: expected = len(corpus) @@ -311,6 +321,7 @@ def assert_peek_result( allowed_ids: set[str] | None = None, ns_tag: str | None = None, ) -> None: + """Assert peek result.""" assert "ids" in result assert "documents" in result assert "metadatas" in result @@ -335,6 +346,7 @@ def _default_include( embeddings: list[float] | None, include: list[str] | None, ) -> list[str] | None: + """Default include.""" if include is not None: return include fields: list[str] = [] @@ -356,6 +368,7 @@ def assert_get_present( embeddings: list[float] | None = None, include: list[str] | None = None, ) -> dict[str, Any]: + """Assert get present.""" result = ns.get(ids=doc_id, include=_default_include(documents, metadatas, embeddings, include)) indices = [i for i, rid in enumerate(result["ids"]) if rid == doc_id] assert indices, f"expected doc {doc_id!r}, got {result['ids']}" @@ -370,5 +383,6 @@ def assert_get_present( def assert_get_absent(ns: Any, doc_id: str) -> None: + """Assert get absent.""" result = ns.get(ids=doc_id) assert len(result["ids"]) == 0, f"expected doc {doc_id!r} absent, got {result['ids']}" diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index 98ca5de0..8325f7f1 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -37,6 +37,7 @@ @dataclass(frozen=True) class CorpusRecord: + """CorpusRecord class.""" doc_id: str document: str embedding: list[float] @@ -46,6 +47,7 @@ class CorpusRecord: @dataclass(frozen=True) class FtsQueryCase: + """FtsQueryCase class.""" name: str where_document: dict[str, Any] | str n_results: int @@ -54,6 +56,7 @@ class FtsQueryCase: def ns_schema(distance: VectorDistanceMetric = "l2") -> Schema: + """Ns schema.""" return Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance=distance, centroids_fresh_mode="spfresh"), @@ -86,6 +89,7 @@ def _add( embedding: list[float] | None = None, extra_meta: dict[str, Any] | None = None, ) -> None: + """Add.""" idx = len(records) meta = {"rel_hint": rel_hint, "seq": idx} if extra_meta: @@ -169,6 +173,7 @@ def _add( def _embedding_for_index(index: int) -> list[float]: + """Embedding for index.""" return [ float((index + 1) % 7) / 7.0, float((index + 2) % 5) / 5.0, @@ -177,6 +182,7 @@ def _embedding_for_index(index: int) -> list[float]: def doc_matches_where_document(document: str, where_document: dict[str, Any] | str) -> bool: + """Doc matches where document.""" text = document.lower() if isinstance(where_document, str): return where_document.lower() in text @@ -237,6 +243,7 @@ def corpus_matches_fts( where_document: dict[str, Any] | str, where: dict[str, Any] | None = None, ) -> bool: + """Corpus matches fts.""" if not doc_matches_where_document(record.document, where_document): return False if where is not None and not doc_matches_where_metadata(record.metadata, where): @@ -282,6 +289,7 @@ def expected_best_id( where_document: dict[str, Any] | str, where: dict[str, Any] | None = None, ) -> str: + """Expected best id.""" matched = [ (rec.doc_id, effective_rel_hint(rec, where_document)) for rec in corpus @@ -338,6 +346,7 @@ def assert_hybrid_fulltext_result( where: dict[str, Any] | None = None, check_ranking: bool = True, ) -> None: + """Assert hybrid fulltext result.""" assert result is not None assert "ids" in result and result["ids"] ids = result["ids"][0] @@ -497,6 +506,7 @@ def assert_hybrid_fulltext_result( def get_fts_case(name: str) -> FtsQueryCase: + """Get fts case.""" for case in HYBRID_SEARCH_FTS_CASES: if case.name == name: return case @@ -553,11 +563,13 @@ def setup_large_fts_namespace(db_client: Any, *, namespace_name: str = "ns_hs_ft def teardown_large_fts_collection(db_client: Any, collection: Any) -> None: + """Teardown large fts collection.""" with contextlib.suppress(Exception): db_client.delete_collection(name=collection.name) def teardown_large_fts_namespace(db_client: Any, collection: Any) -> None: + """Teardown large fts namespace.""" teardown_large_fts_collection(db_client, collection) @@ -574,6 +586,7 @@ def _create_multi_coll_multi_ns_layout( name_prefix: str, distance: VectorDistanceMetric = "l2", ) -> dict[str, Any]: + """Create multi coll multi ns layout.""" ts = int(time.time() * 1000) coll_1 = db_client.create_collection( name=f"{name_prefix}_{distance}_{ts}_c1", @@ -683,6 +696,7 @@ def setup_multi_coll_multi_ns_fts_single_loaded( def teardown_multi_coll_multi_ns_fts(db_client: Any, ctx: dict[str, Any]) -> None: + """Teardown multi coll multi ns fts.""" for coll_key in ("coll_1", "coll_2", "collection"): collection = ctx.get(coll_key) if collection is not None: @@ -734,6 +748,7 @@ def run_hybrid_search_fts_case( corpus: list[CorpusRecord], case: FtsQueryCase, ) -> dict[str, Any]: + """Run hybrid search fts case.""" query: dict[str, Any] = { "where_document": case.where_document, "n_results": case.n_results, @@ -761,6 +776,7 @@ def assert_not_contains_no_token_leak( result: dict[str, Any], forbidden_token: str, ) -> None: + """Assert not contains no token leak.""" result_ids = set(result["ids"][0]) forbidden_ids = { rec.doc_id diff --git a/tests/integration_tests/namespace_hybrid_search_helpers.py b/tests/integration_tests/namespace_hybrid_search_helpers.py index dbbdcfd3..28a94f05 100644 --- a/tests/integration_tests/namespace_hybrid_search_helpers.py +++ b/tests/integration_tests/namespace_hybrid_search_helpers.py @@ -66,6 +66,7 @@ @dataclass(frozen=True) class SearchIndexQueryCase: + """SearchIndexQueryCase class.""" name: str where: dict[str, Any] n_results: int @@ -75,6 +76,7 @@ class SearchIndexQueryCase: @dataclass(frozen=True) class VectorKnnCase: + """VectorKnnCase class.""" name: str query_vector: list[float] n_results: int @@ -128,14 +130,17 @@ def corpus_matches_where(record: CorpusRecord, where: dict[str, Any]) -> bool: def count_corpus_matches(corpus: list[CorpusRecord], where: dict[str, Any]) -> int: + """Count corpus matches.""" return sum(1 for rec in corpus if corpus_matches_where(rec, where)) def l2_squared(a: list[float], b: list[float]) -> float: + """L2 squared.""" return sum((x - y) ** 2 for x, y in zip(a, b)) def _vector_l2_norm(vec: list[float]) -> float: + """Vector l2 norm.""" return math.sqrt(sum(x * x for x in vec)) @@ -234,6 +239,7 @@ def expected_knn_ids( *, distance_metric: VectorDistanceMetric = "l2", ) -> list[str]: + """Expected knn ids.""" scored = _knn_scored_rows( corpus, query_vector, where, distance_metric=distance_metric ) @@ -251,6 +257,7 @@ def assert_hybrid_search_index_result( min_hits: int = 1, exact_match_count: int | None = None, ) -> None: + """Assert hybrid search index result.""" assert result is not None assert "ids" in result and result["ids"] ids = result["ids"][0] @@ -364,6 +371,7 @@ def run_hybrid_search_index_case( corpus: list[CorpusRecord], case: SearchIndexQueryCase, ) -> dict[str, Any]: + """Run hybrid search index case.""" result = namespace.hybrid_search( query={"where": case.where, "n_results": case.n_results}, n_results=case.n_results, @@ -387,6 +395,7 @@ def run_hybrid_knn_case( *, distance_metric: VectorDistanceMetric = "l2", ) -> dict[str, Any]: + """Run hybrid knn case.""" knn: dict[str, Any] = { "query_embeddings": case.query_vector, "n_results": case.n_results, @@ -417,6 +426,7 @@ def run_hybrid_combined_case( *, distance_metric: VectorDistanceMetric = "l2", ) -> dict[str, Any]: + """Run hybrid combined case.""" query: dict[str, Any] | None = None if case.where_document is not None or case.where is not None: query = {"n_results": case.n_results} @@ -574,6 +584,7 @@ def run_hybrid_combined_case( def get_search_index_case(name: str) -> SearchIndexQueryCase: + """Get search index case.""" for case in SEARCH_INDEX_CASES: if case.name == name: return case @@ -605,6 +616,7 @@ def get_search_index_case(name: str) -> SearchIndexQueryCase: def get_vector_knn_case(name: str) -> VectorKnnCase: + """Get vector knn case.""" for case in VECTOR_KNN_CASES: if case.name == name: return case @@ -675,6 +687,7 @@ def get_vector_knn_case(name: str) -> VectorKnnCase: def get_hybrid_combined_case(name: str) -> HybridCombinedCase: + """Get hybrid combined case.""" for case in HYBRID_COMBINED_CASES: if case.name == name: return case @@ -713,6 +726,7 @@ def _default_knn( where: dict[str, Any] | None = None, knn_n_results: int | None = None, ) -> dict[str, Any]: + """Default knn.""" return { "query_embeddings": KNN_QUERY_VECTOR, "n_results": knn_n_results or max(n_results * 2, 20), @@ -724,6 +738,7 @@ def corpus_matches_triple_intersection( record: CorpusRecord, case: HybridTripleBranchCase, ) -> bool: + """Corpus matches triple intersection.""" if case.where_document is not None: if not doc_matches_where_document(record.document, case.where_document): return False @@ -741,6 +756,7 @@ def count_triple_intersection_matches( corpus: list[CorpusRecord], case: HybridTripleBranchCase, ) -> int: + """Count triple intersection matches.""" return sum(1 for rec in corpus if corpus_matches_triple_intersection(rec, case)) @@ -749,6 +765,7 @@ def _slice_hybrid_result( corpus: list[CorpusRecord], predicate, ) -> dict[str, Any]: + """Slice hybrid result.""" corpus_by_id = {rec.doc_id: rec for rec in corpus} ids = result["ids"][0] indices = [ @@ -771,6 +788,7 @@ def assert_hybrid_triple_branch_result( *, distance_metric: VectorDistanceMetric = "l2", ) -> None: + """Assert hybrid triple branch result.""" assert result is not None assert "ids" in result and result["ids"] ids = result["ids"][0] @@ -802,6 +820,7 @@ def assert_hybrid_triple_branch_result( assert case.where_document is not None def _fts_row(rec: CorpusRecord) -> bool: + """Fts row.""" if not doc_matches_where_document(rec.document, case.where_document): return False if case.where is not None and not doc_matches_where_metadata(rec.metadata, case.where): @@ -832,6 +851,7 @@ def _fts_row(rec: CorpusRecord) -> bool: assert case.where is not None def _si_row(rec: CorpusRecord) -> bool: + """Si row.""" if not corpus_matches_where(rec, case.where): return False if case.where_document is not None and not doc_matches_where_document( @@ -865,6 +885,7 @@ def _si_row(rec: CorpusRecord) -> bool: knn_where = case.knn.get("where") def _knn_row(rec: CorpusRecord) -> bool: + """Knn row.""" if knn_where is not None and not corpus_matches_where(rec, knn_where): return False if case.where is not None and not doc_matches_where_metadata(rec.metadata, case.where): @@ -953,6 +974,7 @@ def _knn_row(rec: CorpusRecord) -> bool: def _build_triple_branch_hybrid_search_kwargs( case: HybridTripleBranchCase, ) -> dict[str, Any]: + """Build triple branch hybrid search kwargs.""" query: dict[str, Any] | None = None if case.where_document is not None or case.where is not None: query = {"n_results": case.n_results} @@ -1058,6 +1080,7 @@ def _try_assert_hybrid_triple_branch_result( *, distance_metric: VectorDistanceMetric = "l2", ) -> BaseException | None: + """Try assert hybrid triple branch result.""" try: assert_hybrid_triple_branch_result( corpus, result, case, distance_metric=distance_metric @@ -1146,6 +1169,7 @@ def run_hybrid_triple_branch_case( collection_baseline: Any | None = None, distance_metric: VectorDistanceMetric = "l2", ) -> dict[str, Any]: + """Run hybrid triple branch case.""" if collection_baseline is not None: assert_flat_collection_baseline(collection_baseline) try: @@ -1553,6 +1577,7 @@ def run_hybrid_triple_branch_case( def get_triple_branch_case(name: str) -> HybridTripleBranchCase: + """Get triple branch case.""" for case in TRIPLE_BRANCH_CASES: if case.name == name: return case @@ -1567,6 +1592,7 @@ def run_triple_branch_case_on_quadrants( collection_baseline: Any | None = None, distance_metric: VectorDistanceMetric = "l2", ) -> None: + """Run triple branch case on quadrants.""" tb_case = case if isinstance(case, HybridTripleBranchCase) else get_triple_branch_case(case) for key in quadrant_keys: corpus, namespace = ctx[key] @@ -1584,6 +1610,7 @@ def run_search_index_case_on_quadrants( case: SearchIndexQueryCase | str, quadrant_keys: tuple[str, ...] = MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, ) -> None: + """Run search index case on quadrants.""" si_case = case if isinstance(case, SearchIndexQueryCase) else get_search_index_case(case) for key in quadrant_keys: corpus, namespace = ctx[key] @@ -1597,6 +1624,7 @@ def run_knn_case_on_quadrants( *, distance_metric: VectorDistanceMetric = "l2", ) -> None: + """Run knn case on quadrants.""" knn_case = case if isinstance(case, VectorKnnCase) else get_vector_knn_case(case) for key in quadrant_keys: corpus, namespace = ctx[key] @@ -1611,6 +1639,7 @@ def assert_hybrid_search_index_no_hits( *, n_results: int = 10, ) -> None: + """Assert hybrid search index no hits.""" result = namespace.hybrid_search( query={"where": where, "n_results": n_results}, n_results=n_results, diff --git a/tests/integration_tests/test_get_or_create_collection_concurrency.py b/tests/integration_tests/test_get_or_create_collection_concurrency.py index a24a5bcf..3f9498e6 100644 --- a/tests/integration_tests/test_get_or_create_collection_concurrency.py +++ b/tests/integration_tests/test_get_or_create_collection_concurrency.py @@ -16,6 +16,7 @@ def _make_oceanbase_client(): + """Make oceanbase client.""" import os return pyseekdb.Client( @@ -29,6 +30,7 @@ def _make_oceanbase_client(): def _count_sdk_collections(client, collection_name: str) -> int: + """Count sdk collections.""" name_escaped = escape_string(collection_name) rows = client._server._execute( "SELECT COUNT(*) AS cnt FROM sdk_collections " @@ -39,6 +41,7 @@ def _count_sdk_collections(client, collection_name: str) -> int: def _has_unique_name_index(client) -> bool: + """Has unique name index.""" rows = client._server._execute("SHOW INDEX FROM sdk_collections") for row in rows: key_name = row["Key_name"] if isinstance(row, dict) else row[2] @@ -48,6 +51,7 @@ def _has_unique_name_index(client) -> bool: class TestGetOrCreateCollectionConcurrencyOceanBase: + """TestGetOrCreateCollectionConcurrencyOceanBase class.""" def test_concurrent_get_or_create_collection_is_idempotent(self, oceanbase_client): """8 independent clients race on get_or_create_collection with the same name.""" owner = oceanbase_client @@ -57,6 +61,7 @@ def test_concurrent_get_or_create_collection_is_idempotent(self, oceanbase_clien errors: list[str] = [] def _worker(thread_id: int) -> None: + """Worker.""" client = _make_oceanbase_client() try: barrier.wait(timeout=30) diff --git a/tests/integration_tests/test_get_or_create_collection_multiprocess.py b/tests/integration_tests/test_get_or_create_collection_multiprocess.py index d53fba0b..8208bb96 100644 --- a/tests/integration_tests/test_get_or_create_collection_multiprocess.py +++ b/tests/integration_tests/test_get_or_create_collection_multiprocess.py @@ -60,25 +60,30 @@ def _purge_pyseekdb_modules() -> None: + """Purge pyseekdb modules.""" for module_name in list(sys.modules): if module_name == "pyseekdb" or module_name.startswith("pyseekdb."): del sys.modules[module_name] def _build_embedding(seed: int, embed_dim: int = EMBED_DIM) -> list[float]: + """Build embedding.""" base = float(seed + 1) return [((base + i) % 13) / 13.0 for i in range(embed_dim)] def _record_id(process_id: int, thread_id: int, seq: int) -> str: + """Record id.""" return f"p{process_id:02d}_t{thread_id:02d}_{seq:04d}" def _import_pyseekdb(): + """Import pyseekdb.""" return importlib.import_module("pyseekdb") def _make_client(client_config: dict[str, Any]): + """Make client.""" pyseekdb = _import_pyseekdb() mode = client_config["mode"] if mode == "embedded": @@ -96,6 +101,7 @@ def _make_client(client_config: dict[str, Any]): def _make_admin_client(client_config: dict[str, Any]): + """Make admin client.""" pyseekdb = _import_pyseekdb() mode = client_config["mode"] if mode == "embedded": @@ -112,16 +118,19 @@ def _make_admin_client(client_config: dict[str, Any]): def _get_collection(client, collection_name: str): + """Get collection.""" return client.get_collection(collection_name, embedding_function=None) def _refresh_collection(client_config: dict[str, Any], collection_name: str) -> None: + """Refresh collection.""" client = _make_client(client_config) collection = _get_collection(client, collection_name) collection.refresh_index() def _require_embedded_pylibseekdb() -> None: + """Require embedded pylibseekdb.""" try: import pylibseekdb # noqa: F401 except ImportError: @@ -144,6 +153,7 @@ def _run_processes( # noqa: C901 expected_count: int, timeout: float = WORKER_TIMEOUT_SECONDS, ) -> list[dict[str, Any]]: + """Run processes.""" result_queue: mp.Queue[dict[str, Any]] = _MP_CONTEXT.Queue() processes = [_MP_CONTEXT.Process(target=target, args=(*args, result_queue)) for args in args_list] @@ -184,6 +194,7 @@ def _get_or_create_worker( delay: float, output: mp.Queue[dict[str, Any]], ) -> None: + """Get or create worker.""" _purge_pyseekdb_modules() time.sleep(delay) try: @@ -207,9 +218,11 @@ def _add_worker( items_per_thread: int, output: mp.Queue[dict[str, Any]], ) -> None: + """Add worker.""" _purge_pyseekdb_modules() def add_in_thread(thread_id: int) -> int: + """Add in thread.""" client = _make_client(client_config) collection = _get_collection(client, collection_name) ids = [_record_id(process_id, thread_id, seq) for seq in range(items_per_thread)] @@ -238,9 +251,11 @@ def _get_worker( items_per_thread: int, output: mp.Queue[dict[str, Any]], ) -> None: + """Get worker.""" _purge_pyseekdb_modules() def get_in_thread(thread_id: int) -> int: + """Get in thread.""" client = _make_client(client_config) collection = _get_collection(client, collection_name) ids = [_record_id(process_id, thread_id, seq) for seq in range(items_per_thread)] @@ -268,9 +283,11 @@ def _query_worker( queries_per_thread: int, output: mp.Queue[dict[str, Any]], ) -> None: + """Query worker.""" _purge_pyseekdb_modules() def query_in_thread(thread_id: int) -> int: + """Query in thread.""" client = _make_client(client_config) collection = _get_collection(client, collection_name) completed = 0 @@ -305,9 +322,11 @@ def _update_worker( items_per_thread: int, output: mp.Queue[dict[str, Any]], ) -> None: + """Update worker.""" _purge_pyseekdb_modules() def update_in_thread(thread_id: int) -> int: + """Update in thread.""" client = _make_client(client_config) collection = _get_collection(client, collection_name) ids = [_record_id(process_id, thread_id, seq) for seq in range(items_per_thread)] @@ -343,9 +362,11 @@ def _delete_worker( items_per_thread: int, output: mp.Queue[dict[str, Any]], ) -> None: + """Delete worker.""" _purge_pyseekdb_modules() def delete_in_thread(thread_id: int) -> int: + """Delete in thread.""" client = _make_client(client_config) collection = _get_collection(client, collection_name) ids = [_record_id(process_id, thread_id, seq) for seq in range(items_per_thread)] @@ -374,9 +395,11 @@ def _mixed_crud_worker( items_per_thread: int, output: mp.Queue[dict[str, Any]], ) -> None: + """Mixed crud worker.""" _purge_pyseekdb_modules() def mixed_in_thread(thread_id: int) -> dict[str, int]: + """Mixed in thread.""" client = _make_client(client_config) collection = _get_collection(client, collection_name) ids = [_record_id(process_id, thread_id, seq) for seq in range(items_per_thread)] @@ -443,6 +466,7 @@ def mixed_in_thread(thread_id: int) -> dict[str, int]: def _build_client_config(mode: str) -> tuple[dict[str, Any], Path | None]: + """Build client config.""" database = f"test_mp_{uuid.uuid4().hex[:8]}" temp_db_path: Path | None = None @@ -482,6 +506,7 @@ def _build_client_config(mode: str) -> tuple[dict[str, Any], Path | None]: @pytest.fixture def multiprocess_db(_mode): + """Multiprocess db.""" client_config, temp_db_path = _build_client_config(_mode) yield client_config if temp_db_path is not None: @@ -491,6 +516,7 @@ def multiprocess_db(_mode): @pytest.fixture def crud_collection(multiprocess_db): + """Crud collection.""" client_config = multiprocess_db collection_name = f"mp_crud_{uuid.uuid4().hex}" @@ -515,6 +541,7 @@ def _seed_collection_rows( threads_per_process: int, items_per_thread: int, ) -> int: + """Seed collection rows.""" client = _make_client(client_config) collection = _get_collection(client, collection_name) @@ -538,7 +565,9 @@ def _seed_collection_rows( class TestGetOrCreateCollectionMultiprocess: + """TestGetOrCreateCollectionMultiprocess class.""" def test_concurrent_get_or_create_collection(self, multiprocess_db, _mode): + """Test concurrent get or create collection.""" client_config = multiprocess_db collection_name = f"mp_collection_{uuid.uuid4().hex}" @@ -576,7 +605,9 @@ def test_concurrent_get_or_create_collection(self, multiprocess_db, _mode): class TestMultiprocessMultithreadCrud: + """TestMultiprocessMultithreadCrud class.""" def test_concurrent_add(self, crud_collection, _mode): + """Test concurrent add.""" client_config, collection_name = crud_collection results = _run_processes( @@ -600,6 +631,7 @@ def test_concurrent_add(self, crud_collection, _mode): assert collection.count() == expected_rows def test_concurrent_get(self, crud_collection, _mode): + """Test concurrent get.""" client_config, collection_name = crud_collection expected_rows = _seed_collection_rows( client_config, collection_name, NUM_PROCESSES, THREADS_PER_PROCESS, ITEMS_PER_THREAD @@ -619,6 +651,7 @@ def test_concurrent_get(self, crud_collection, _mode): assert sum(result["fetched"] for result in results) == expected_rows def test_concurrent_query(self, crud_collection, _mode): + """Test concurrent query.""" client_config, collection_name = crud_collection _seed_collection_rows(client_config, collection_name, NUM_PROCESSES, THREADS_PER_PROCESS, ITEMS_PER_THREAD) @@ -637,6 +670,7 @@ def test_concurrent_query(self, crud_collection, _mode): assert sum(result["queried"] for result in results) == expected_queries def test_concurrent_update(self, crud_collection, _mode): + """Test concurrent update.""" client_config, collection_name = crud_collection expected_rows = _seed_collection_rows( client_config, collection_name, NUM_PROCESSES, THREADS_PER_PROCESS, ITEMS_PER_THREAD @@ -656,6 +690,7 @@ def test_concurrent_update(self, crud_collection, _mode): assert sum(result["updated"] for result in results) == expected_rows def test_concurrent_delete(self, crud_collection, _mode): + """Test concurrent delete.""" client_config, collection_name = crud_collection expected_rows = _seed_collection_rows( client_config, collection_name, NUM_PROCESSES, THREADS_PER_PROCESS, ITEMS_PER_THREAD @@ -679,6 +714,7 @@ def test_concurrent_delete(self, crud_collection, _mode): assert collection.count() == 0 def test_concurrent_mixed_crud(self, crud_collection, _mode): + """Test concurrent mixed crud.""" client_config, collection_name = crud_collection results = _run_processes( diff --git a/tests/integration_tests/test_logic_table_monitoring_p1.py b/tests/integration_tests/test_logic_table_monitoring_p1.py index ef9d5261..070c08b7 100644 --- a/tests/integration_tests/test_logic_table_monitoring_p1.py +++ b/tests/integration_tests/test_logic_table_monitoring_p1.py @@ -15,6 +15,7 @@ def _make_oceanbase_client(): + """Make oceanbase client.""" import os return pyseekdb.Client( @@ -28,6 +29,7 @@ def _make_oceanbase_client(): def _count_sdk_namespaces(client, collection_id: str, namespace_name: str) -> int: + """Count sdk namespaces.""" rows = client._server._execute( "SELECT COUNT(*) AS cnt FROM sdk_namespaces " f"WHERE collection_id = '{collection_id}' AND namespace_name = '{namespace_name}'" @@ -37,6 +39,7 @@ def _count_sdk_namespaces(client, collection_id: str, namespace_name: str) -> in def _count_sdk_ltables(client, collection_id: str, namespace_id: int) -> int: + """Count sdk ltables.""" rows = client._server._execute( "SELECT COUNT(*) AS cnt FROM sdk_ltables " f"WHERE collection_id = '{collection_id}' AND namespace_id = {int(namespace_id)} " @@ -47,6 +50,7 @@ def _count_sdk_ltables(client, collection_id: str, namespace_id: int) -> int: class TestLogicTableMonitoringP1: + """TestLogicTableMonitoringP1 class.""" def test_concurrent_get_or_create_same_namespace_is_idempotent(self, oceanbase_client): """8 threads with independent clients race on the same namespace name.""" owner = oceanbase_client @@ -64,6 +68,7 @@ def test_concurrent_get_or_create_same_namespace_is_idempotent(self, oceanbase_c errors: list[str] = [] def _worker(thread_id: int) -> None: + """Worker.""" client = _make_oceanbase_client() try: coll = client.get_collection(collection_name) diff --git a/tests/integration_tests/test_namespace_constraints.py b/tests/integration_tests/test_namespace_constraints.py index 9bf5e5d1..5a0c1af4 100644 --- a/tests/integration_tests/test_namespace_constraints.py +++ b/tests/integration_tests/test_namespace_constraints.py @@ -20,6 +20,7 @@ def _make_schema(ivf_type: str) -> Schema: + """Make schema.""" return Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine", type=ivf_type, centroids_fresh_mode="spfresh"), @@ -29,10 +30,12 @@ def _make_schema(ivf_type: str) -> Schema: def _unique_name(suffix: str) -> str: + """Unique name.""" return f"test_ns_constraint_{int(time.time() * 1000)}{suffix}" class TestNamespaceIvfTypeConstraint: + """TestNamespaceIvfTypeConstraint class.""" @pytest.mark.parametrize("ivf_type", ["ivf_pq", "ivf_sq8"]) def test_non_ivf_flat_rejected(self, oceanbase_client, ivf_type): """Non-ivf_flat IVF types must be rejected at the SDK layer.""" @@ -59,6 +62,7 @@ def test_ivf_flat_accepted(self, oceanbase_client): class TestNamespaceMinVersionConstraint: + """TestNamespaceMinVersionConstraint class.""" def test_connected_ob_meets_min_version(self, oceanbase_client): """The kernel under test must already be >= 4.6.1, and creation succeeds.""" db_type, version = oceanbase_client._server.detect_db_type_and_version() @@ -95,4 +99,5 @@ def test_old_ob_version_rejected(self, oceanbase_client, monkeypatch): assert not oceanbase_client.has_collection(name) def test_min_version_constant(self): + """Test min version constant.""" assert NAMESPACE_MIN_OB_VERSION == Version("4.6.1.0") diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py index f3adb5c4..cb157dd8 100644 --- a/tests/integration_tests/test_namespace_dml.py +++ b/tests/integration_tests/test_namespace_dml.py @@ -24,6 +24,7 @@ def _create_namespace(collection, name: str): + """Create namespace.""" ns = collection.create_namespace(name) ns.prewarm() return ns @@ -33,6 +34,7 @@ class TestNamespaceDML: """Small-scale DML smoke tests.""" def test_add_single(self, db_client): + """Test add single.""" collection = create_ns_collection(db_client, suffix="_add1") ns = _create_namespace(collection,"dml_ns") try: @@ -44,6 +46,7 @@ def test_add_single(self, db_client): cleanup(db_client, collection) def test_ops_blocked_after_namespace_deleted(self, db_client): + """Test ops blocked after namespace deleted.""" collection = create_ns_collection(db_client, suffix="_delns") ns = _create_namespace(collection, "dml_ns") try: @@ -61,6 +64,7 @@ def test_ops_blocked_after_namespace_deleted(self, db_client): cleanup(db_client, collection) def test_ops_blocked_after_collection_deleted(self, db_client): + """Test ops blocked after collection deleted.""" collection = create_ns_collection(db_client, suffix="_delcoll") ns = _create_namespace(collection, "dml_ns") ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0]) @@ -71,6 +75,7 @@ def test_ops_blocked_after_collection_deleted(self, db_client): ns.query(query_embeddings=[1.0, 2.0, 3.0], n_results=1) def test_add_batch(self, db_client): + """Test add batch.""" collection = create_ns_collection(db_client, suffix="_addb") ns = _create_namespace(collection,"dml_ns") try: @@ -86,6 +91,7 @@ def test_add_batch(self, db_client): cleanup(db_client, collection) def test_get_by_id(self, db_client): + """Test get by id.""" collection = create_ns_collection(db_client, suffix="_getid") ns = _create_namespace(collection,"dml_ns") try: @@ -103,6 +109,7 @@ def test_get_by_id(self, db_client): cleanup(db_client, collection) def test_get_with_limit(self, db_client): + """Test get with limit.""" collection = create_ns_collection(db_client, suffix="_getlim") ns = _create_namespace(collection,"dml_ns") try: @@ -116,6 +123,7 @@ def test_get_with_limit(self, db_client): cleanup(db_client, collection) def test_update_metadata(self, db_client): + """Test update metadata.""" collection = create_ns_collection(db_client, suffix="_upd") ns = _create_namespace(collection,"dml_ns") try: @@ -127,6 +135,7 @@ def test_update_metadata(self, db_client): cleanup(db_client, collection) def test_update_document_and_embedding(self, db_client): + """Test update document and embedding.""" collection = create_ns_collection(db_client, suffix="_upddoc") ns = _create_namespace(collection,"dml_ns") try: @@ -138,6 +147,7 @@ def test_update_document_and_embedding(self, db_client): cleanup(db_client, collection) def test_upsert_existing(self, db_client): + """Test upsert existing.""" collection = create_ns_collection(db_client, suffix="_upsex") ns = _create_namespace(collection,"dml_ns") try: @@ -149,6 +159,7 @@ def test_upsert_existing(self, db_client): cleanup(db_client, collection) def test_upsert_new(self, db_client): + """Test upsert new.""" collection = create_ns_collection(db_client, suffix="_upsnew") ns = _create_namespace(collection,"dml_ns") try: @@ -160,6 +171,7 @@ def test_upsert_new(self, db_client): cleanup(db_client, collection) def test_delete_by_ids(self, db_client): + """Test delete by ids.""" collection = create_ns_collection(db_client, suffix="_del") ns = _create_namespace(collection,"dml_ns") try: @@ -179,6 +191,7 @@ class TestNamespaceDMLCountPeekAtScale: @pytest.fixture def large_ns(self, db_client): + """Large ns.""" collection = create_ns_collection(db_client, suffix="_large_cp") ns = _create_namespace(collection,"large_ns") corpus = build_large_dml_corpus(LARGE_DML_CORPUS_SIZE, id_prefix="large") @@ -187,6 +200,7 @@ def large_ns(self, db_client): cleanup(db_client, collection) def test_count_empty_namespace(self, db_client): + """Test count empty namespace.""" collection = create_ns_collection(db_client, suffix="_cnt_empty") ns = _create_namespace(collection,"empty_ns") try: @@ -195,10 +209,12 @@ def test_count_empty_namespace(self, db_client): cleanup(db_client, collection) def test_count_after_bulk_load_1000(self, large_ns): + """Test count after bulk load 1000.""" _, ns, corpus = large_ns assert ns.count() == len(corpus) == LARGE_DML_CORPUS_SIZE def test_count_after_partial_delete(self, large_ns): + """Test count after partial delete.""" _, ns, corpus = large_ns to_delete = [corpus[i].doc_id for i in range(100)] ns.delete(ids=to_delete) @@ -207,6 +223,7 @@ def test_count_after_partial_delete(self, large_ns): assert len(ns.get(ids=corpus[500].doc_id)["ids"]) == 1 def test_peek_empty_namespace(self, db_client): + """Test peek empty namespace.""" collection = create_ns_collection(db_client, suffix="_peek_empty") ns = _create_namespace(collection,"peek_empty") try: @@ -239,12 +256,14 @@ def test_peek_empty_namespace(self, db_client): ], ) def test_peek_limit_boundaries(self, large_ns, limit, expected_len): + """Test peek limit boundaries.""" _, ns, corpus = large_ns allowed = corpus_id_set(corpus) result = ns.peek(limit=limit) assert_peek_result(result, expected_len=expected_len, allowed_ids=allowed) def test_peek_default_limit_is_10(self, large_ns): + """Test peek default limit is 10.""" _, ns, corpus = large_ns result = ns.peek() assert_peek_result( @@ -262,6 +281,7 @@ class TestNamespaceDMLCountPeekMultiNs: @pytest.fixture def dual_ns_ctx(self, db_client): + """Dual ns ctx.""" collection = create_ns_collection(db_client, suffix="_dual_cp") ns_a = _create_namespace(collection,"ns_alpha") ns_b = _create_namespace(collection,"ns_beta") @@ -286,6 +306,7 @@ def dual_ns_ctx(self, db_client): cleanup(db_client, collection) def test_count_isolated_per_namespace(self, dual_ns_ctx): + """Test count isolated per namespace.""" assert dual_ns_ctx["ns_a"].count() == self.NS_A_SIZE assert dual_ns_ctx["ns_b"].count() == self.NS_B_SIZE @@ -295,6 +316,7 @@ def test_count_isolated_per_namespace(self, dual_ns_ctx): ids=["lim_1", "lim_10", "lim_100", "lim_800", "lim_1000", "lim_over_alpha"], ) def test_peek_alpha_boundaries(self, dual_ns_ctx, limit, expected_len): + """Test peek alpha boundaries.""" result = dual_ns_ctx["ns_a"].peek(limit=limit) assert_peek_result( result, @@ -309,6 +331,7 @@ def test_peek_alpha_boundaries(self, dual_ns_ctx, limit, expected_len): ids=["lim_1", "lim_10", "lim_500", "lim_800", "lim_over_beta", "lim_far_over"], ) def test_peek_beta_boundaries(self, dual_ns_ctx, limit, expected_len): + """Test peek beta boundaries.""" result = dual_ns_ctx["ns_b"].peek(limit=limit) assert_peek_result( result, @@ -318,6 +341,7 @@ def test_peek_beta_boundaries(self, dual_ns_ctx, limit, expected_len): ) def test_peek_does_not_leak_across_namespaces(self, dual_ns_ctx): + """Test peek does not leak across namespaces.""" peek_a = dual_ns_ctx["ns_a"].peek(limit=50) peek_b = dual_ns_ctx["ns_b"].peek(limit=50) assert_peek_result( @@ -336,6 +360,7 @@ def test_peek_does_not_leak_across_namespaces(self, dual_ns_ctx): assert set(peek_b["ids"]).isdisjoint(dual_ns_ctx["ids_a"]) def test_delete_in_one_namespace_only_affects_its_count(self, dual_ns_ctx): + """Test delete in one namespace only affects its count.""" ns_a = dual_ns_ctx["ns_a"] ns_b = dual_ns_ctx["ns_b"] corpus_a = dual_ns_ctx["corpus_a"] @@ -364,6 +389,7 @@ class TestNamespaceDMLFullCycle: """End-to-end DML cycle (small scale, verified by get).""" def test_full_dml_cycle_verified_by_get(self, db_client): + """Test full dml cycle verified by get.""" collection = create_ns_collection(db_client, suffix="_cycle") ns = _create_namespace(collection,"cycle_ns") doc_id = "cycle_doc" @@ -417,6 +443,7 @@ def test_full_dml_cycle_verified_by_get(self, db_client): cleanup(db_client, collection) def test_batch_add_then_get_each(self, db_client): + """Test batch add then get each.""" collection = create_ns_collection(db_client, suffix="_batch") ns = _create_namespace(collection,"batch_ns") try: diff --git a/tests/integration_tests/test_namespace_dml_multi_coll.py b/tests/integration_tests/test_namespace_dml_multi_coll.py index 212fa2c4..0b18413c 100644 --- a/tests/integration_tests/test_namespace_dml_multi_coll.py +++ b/tests/integration_tests/test_namespace_dml_multi_coll.py @@ -15,7 +15,9 @@ class TestNamespaceDMLMultiColl: + """TestNamespaceDMLMultiColl class.""" def test_same_namespace_name_across_collections(self, db_client): + """Test same namespace name across collections.""" coll_a = create_ns_collection(db_client, suffix="_mc_a") coll_b = create_ns_collection(db_client, suffix="_mc_b") ns_a = coll_a.create_namespace("tenant") @@ -46,6 +48,7 @@ def test_same_namespace_name_across_collections(self, db_client): cleanup(db_client, coll_a, coll_b) def test_update_delete_isolated_by_collection(self, db_client): + """Test update delete isolated by collection.""" coll_a = create_ns_collection(db_client, suffix="_mc_upd_a") coll_b = create_ns_collection(db_client, suffix="_mc_upd_b") ns_a = coll_a.create_namespace("tenant") diff --git a/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py index 668fde8c..89d3ce91 100644 --- a/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py @@ -15,7 +15,9 @@ class TestNamespaceDMLMultiCollMultiNs: + """TestNamespaceDMLMultiCollMultiNs class.""" def _setup_quadrants(self, db_client): + """Setup quadrants.""" coll_1 = create_ns_collection(db_client, suffix="_mcmn_c1") coll_2 = create_ns_collection(db_client, suffix="_mcmn_c2") quadrants = { @@ -31,6 +33,7 @@ def _setup_quadrants(self, db_client): return quadrants def test_cross_quadrant_isolation(self, db_client): + """Test cross quadrant isolation.""" q = self._setup_quadrants(db_client) marker_id = "marker_id" try: @@ -57,6 +60,7 @@ def test_cross_quadrant_isolation(self, db_client): cleanup(db_client, q["coll_1"], q["coll_2"]) def test_delete_in_one_quadrant(self, db_client): + """Test delete in one quadrant.""" q = self._setup_quadrants(db_client) doc_id = "shared_quad_del" try: diff --git a/tests/integration_tests/test_namespace_dml_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_ns.py index 8e854857..6380326f 100644 --- a/tests/integration_tests/test_namespace_dml_multi_ns.py +++ b/tests/integration_tests/test_namespace_dml_multi_ns.py @@ -15,7 +15,9 @@ class TestNamespaceDMLMultiNs: + """TestNamespaceDMLMultiNs class.""" def test_add_isolation(self, db_client): + """Test add isolation.""" collection = create_ns_collection(db_client, suffix="_mns_add") ns_a = collection.create_namespace("ns_a") ns_b = collection.create_namespace("ns_b") @@ -33,6 +35,7 @@ def test_add_isolation(self, db_client): cleanup(db_client, collection) def test_same_id_different_namespaces(self, db_client): + """Test same id different namespaces.""" collection = create_ns_collection(db_client, suffix="_mns_sameid") ns_x = collection.create_namespace("ns_x") ns_y = collection.create_namespace("ns_y") @@ -48,6 +51,7 @@ def test_same_id_different_namespaces(self, db_client): cleanup(db_client, collection) def test_update_does_not_affect_other_ns(self, db_client): + """Test update does not affect other ns.""" collection = create_ns_collection(db_client, suffix="_mns_upd") ns_a = collection.create_namespace("ns_a") ns_b = collection.create_namespace("ns_b") @@ -68,6 +72,7 @@ def test_update_does_not_affect_other_ns(self, db_client): cleanup(db_client, collection) def test_delete_only_target_ns(self, db_client): + """Test delete only target ns.""" collection = create_ns_collection(db_client, suffix="_mns_del") ns_a = collection.create_namespace("ns_a") ns_b = collection.create_namespace("ns_b") @@ -88,6 +93,7 @@ def test_delete_only_target_ns(self, db_client): cleanup(db_client, collection) def test_upsert_isolation(self, db_client): + """Test upsert isolation.""" collection = create_ns_collection(db_client, suffix="_mns_ups") ns_a = collection.create_namespace("ns_a") ns_b = collection.create_namespace("ns_b") diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index 62245c36..a6ed5c37 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -1,3 +1,5 @@ +"""test namespace drop validation module.""" + import threading import time @@ -24,6 +26,7 @@ def _create_namespace(collection, name: str): + """Create namespace.""" ns = collection.create_namespace(name) ns.prewarm() return ns @@ -46,6 +49,7 @@ def _make_collection(client, suffix: str = ""): def _execute(client, sql: str): + """Execute.""" return client._server._execute(sql) @@ -72,6 +76,7 @@ def _is_ss_mode(client) -> bool: def _fetch_namespace_name(client, collection_id: str, namespace_id: int): + """Fetch namespace name.""" ns_table = _catalog_table(client, "sdk_namespaces") rows = _execute( client, @@ -84,6 +89,7 @@ def _fetch_namespace_name(client, collection_id: str, namespace_id: int): def _count_ltables(client, collection_id: str, namespace_id: int) -> int: + """Count ltables.""" lt_table = _catalog_table(client, "sdk_ltables") rows = _execute( client, @@ -94,6 +100,7 @@ def _count_ltables(client, collection_id: str, namespace_id: int) -> int: def _count_logic_schema_rows(client, collection_id: str, namespace_id: int) -> int: + """Count logic schema rows.""" tbl = NamespaceCollectionNames.logic_schema_table_name(collection_id) db = client._server.database rows = _execute( @@ -104,6 +111,7 @@ def _count_logic_schema_rows(client, collection_id: str, namespace_id: int) -> i def _count_hot_table_rows(client, collection_id: str, namespace_id: int) -> int: + """Count hot table rows.""" tbl = NamespaceCollectionNames.hot_table_name(collection_id) db = client._server.database try: @@ -117,6 +125,7 @@ def _count_hot_table_rows(client, collection_id: str, namespace_id: int) -> int: def _count_logic_data_rows(client, collection_id: str, namespace_id: int, ltable_id: int | None = None) -> int: + """Count logic data rows.""" tbl = NamespaceCollectionNames.data_table_name(collection_id) db = client._server.database if ltable_id is not None: @@ -131,6 +140,7 @@ def _count_logic_data_rows(client, collection_id: str, namespace_id: int, ltable def _count_kv_data_rows(client, collection_id: str, namespace_id: int) -> int: + """Count kv data rows.""" tbl = NamespaceCollectionNames.kv_data_table_name(collection_id) db = client._server.database try: @@ -145,6 +155,7 @@ def _count_kv_data_rows(client, collection_id: str, namespace_id: int) -> int: def _wait_until(predicate, timeout_sec: float = BG_POLL_TIMEOUT_SEC, interval_sec: float = BG_POLL_INTERVAL_SEC, desc: str = ""): + """Wait until.""" deadline = time.time() + timeout_sec last_exc = None while time.time() < deadline: @@ -176,6 +187,7 @@ def _seed_kv_data(client, collection_id: str, namespace_id: int, count: int = 5) def _seed_logic_data(client, collection_id: str, namespace_id: int, ltable_id: int, count: int = 3): + """Seed logic data.""" tbl = NamespaceCollectionNames.data_table_name(collection_id) db = client._server.database values = [] @@ -193,6 +205,7 @@ def _seed_logic_data(client, collection_id: str, namespace_id: int, ltable_id: i def _fetch_ltable_id(client, collection_id: str, namespace_id: int) -> int | None: + """Fetch ltable id.""" lt_table = _catalog_table(client, "sdk_ltables") rows = _execute( client, @@ -206,6 +219,7 @@ def _fetch_ltable_id(client, collection_id: str, namespace_id: int) -> int | Non def _hot_table_exists(client, collection_id: str) -> bool: + """Hot table exists.""" tbl = NamespaceCollectionNames.hot_table_name(collection_id) db = client._server.database try: @@ -294,6 +308,7 @@ def _insert_active_ltable_row( ltable_name: str, ltable_id: int, ): + """Insert active ltable row.""" lt_table = _catalog_table(client, "sdk_ltables") _execute( client, @@ -304,6 +319,7 @@ def _insert_active_ltable_row( def _fetch_namespace_drop_history(client, collection_id: str, namespace_id: int): + """Fetch namespace drop history.""" try: return _execute( client, @@ -355,6 +371,7 @@ class TestDropNamespaceCatalogValidation: # 1. Happy path: rename + cleanup all relevant rows # # ------------------------------------------------------------- # def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): + """Test drop namespace renames and cleans catalog.""" client = oceanbase_client is_ss = _is_ss_mode(client) collection = _make_collection(client) @@ -444,6 +461,7 @@ def test_drop_namespace_async_cleanup(self, oceanbase_client): assert _count_ltables(client, coll_id, ns_id) == 0 def _bg_done(): + """Bg done.""" return ( _fetch_namespace_name(client, coll_id, ns_id) is None and _count_kv_data_rows(client, coll_id, ns_id) == 0 @@ -535,6 +553,7 @@ def test_async_hot_table_cleaned_on_non_ss(self, oceanbase_client): ) def _hot_and_ns_gone(): + """Hot and ns gone.""" return ( _count_hot_table_rows(client, coll_id, ns_id) == 0 and _fetch_namespace_name(client, coll_id, ns_id) is None @@ -556,6 +575,7 @@ def test_async_drop_history_recorded(self, oceanbase_client): _drop_namespace_via_pl(client, coll_id, ns_id) def _bg_done(): + """Bg done.""" return _fetch_namespace_name(client, coll_id, ns_id) is None _wait_until(_bg_done, desc="namespace removed for history check") @@ -576,6 +596,7 @@ def _bg_done(): # 2. Namespace-name rename format # # ------------------------------------------------------------- # def test_recyclebin_name_format(self, oceanbase_client): + """Test recyclebin name format.""" client = oceanbase_client collection = _make_collection(client) try: @@ -596,6 +617,7 @@ def test_recyclebin_name_format(self, oceanbase_client): # 3. Multiple ltables under one namespace all get deleted # # ------------------------------------------------------------- # def test_multiple_ltables_all_deleted(self, oceanbase_client): + """Test multiple ltables all deleted.""" client = oceanbase_client collection = _make_collection(client) try: @@ -674,6 +696,7 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): # 5. Idempotent: second drop on the same ns is a no-op success # # ------------------------------------------------------------- # def test_drop_namespace_is_idempotent(self, oceanbase_client): + """Test drop namespace is idempotent.""" client = oceanbase_client collection = _make_collection(client) coll_id = collection.id @@ -720,6 +743,7 @@ def test_concurrent_drop_from_two_sdks(self, oceanbase_client): barrier = threading.Barrier(2) def _worker(tag, c): + """Worker.""" try: barrier.wait(timeout=10) _drop_namespace_via_pl(c, coll_id, ns_id) diff --git a/tests/integration_tests/test_namespace_get_delete_filters.py b/tests/integration_tests/test_namespace_get_delete_filters.py index dcac545d..5a076c16 100644 --- a/tests/integration_tests/test_namespace_get_delete_filters.py +++ b/tests/integration_tests/test_namespace_get_delete_filters.py @@ -31,6 +31,7 @@ class TestNamespaceDeleteWhere: """P0 #1–#2: conditional delete on namespace.""" def test_delete_where_metadata_tag(self, db_client): + """Test delete where metadata tag.""" collection = create_ns_collection(db_client, suffix="_del_where") ns = collection.create_namespace("del_where_ns") ns.prewarm() @@ -54,6 +55,7 @@ def test_delete_where_metadata_tag(self, db_client): cleanup(db_client, collection) def test_delete_where_document_obsolete(self, db_client): + """Test delete where document obsolete.""" collection = create_ns_collection(db_client, suffix="_del_wdoc") ns = collection.create_namespace("del_wdoc_ns") ns.prewarm() @@ -94,6 +96,7 @@ class TestNamespaceGetWhere: """P0 #3–#4: conditional get on namespace.""" def test_get_where_metadata_category_ai(self, db_client): + """Test get where metadata category ai.""" collection = create_ns_collection(db_client, suffix="_get_where") ns = collection.create_namespace("get_where_ns") ns.prewarm() @@ -121,6 +124,7 @@ def test_get_where_metadata_category_ai(self, db_client): cleanup(db_client, collection) def test_get_where_document_contains_python(self, db_client): + """Test get where document contains python.""" collection = create_ns_collection(db_client, suffix="_get_wdoc") ns = collection.create_namespace("get_wdoc_ns") ns.prewarm() @@ -153,6 +157,7 @@ class TestNamespaceFilterMultiNs: """P0 #5: filter delete isolated per namespace (same collection).""" def test_delete_where_only_affects_target_namespace(self, db_client): + """Test delete where only affects target namespace.""" collection = create_ns_collection(db_client, suffix="_filt_mns") ns_a = collection.create_namespace("ns_alpha") ns_b = collection.create_namespace("ns_beta") @@ -212,6 +217,7 @@ class TestNamespaceCrossNsGetWhere: """ def test_get_where_alternates_namespaces_after_delete(self, db_client): + """Test get where alternates namespaces after delete.""" collection = create_ns_collection(db_client, suffix="_cross_ns_get") ns_a = collection.create_namespace("cross_alpha") ns_b = collection.create_namespace("cross_beta") diff --git a/tests/integration_tests/test_namespace_hybrid_search_combined.py b/tests/integration_tests/test_namespace_hybrid_search_combined.py index 7c5c74fe..b80fdeae 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_combined.py +++ b/tests/integration_tests/test_namespace_hybrid_search_combined.py @@ -30,6 +30,7 @@ @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) class TestNamespaceHybridSearchCombined: + """TestNamespaceHybridSearchCombined class.""" _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} @pytest.fixture(autouse=True) @@ -39,6 +40,7 @@ def _bind_shared_collection( vector_distance: VectorDistanceMetric, request: pytest.FixtureRequest, ) -> None: + """Bind shared collection.""" entry = ensure_shared_hybrid_search_collection( self._shared_by_mode, db_client, request, vector_distance ) @@ -48,11 +50,13 @@ def _bind_shared_collection( @classmethod def teardown_class(cls) -> None: + """Teardown class.""" for entry in cls._shared_by_mode.values(): teardown_large_fts_collection(entry["db_client"], entry["collection"]) cls._shared_by_mode.clear() def _new_namespace(self, case_name: str) -> Any: + """New namespace.""" return setup_fts_namespace_with_corpus( self._collection, self._corpus, @@ -61,6 +65,7 @@ def _new_namespace(self, case_name: str) -> Any: @pytest.mark.parametrize("case_name", [c.name for c in HYBRID_COMBINED_CASES]) def test_hybrid_search_combined(self, db_client, case_name: str): + """Test hybrid search combined.""" case = get_hybrid_combined_case(case_name) namespace = self._new_namespace(case_name) run_hybrid_combined_case( diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext.py index 0c30cddb..19375cc1 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_fulltext.py +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext.py @@ -35,6 +35,7 @@ class TestNamespaceHybridSearchFulltext: @pytest.fixture(autouse=True) def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: + """Bind shared collection.""" mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" if mode not in self._shared_by_mode: corpus, collection = setup_large_fts_collection(db_client) @@ -50,11 +51,13 @@ def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest @classmethod def teardown_class(cls) -> None: + """Teardown class.""" for entry in cls._shared_by_mode.values(): teardown_large_fts_collection(entry["db_client"], entry["collection"]) cls._shared_by_mode.clear() def _new_namespace(self, case_name: str) -> Any: + """New namespace.""" ns_name = f"ns_{case_name}_{int(time.time() * 1000)}" return setup_fts_namespace_with_corpus( self._collection, @@ -63,6 +66,7 @@ def _new_namespace(self, case_name: str) -> Any: ) def _run_fts_case(self, case_name: str) -> None: + """Run fts case.""" namespace = self._new_namespace(case_name) run_hybrid_search_fts_case(namespace, self._corpus, get_fts_case(case_name)) diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py index 26fde4c0..8826c4f1 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py @@ -99,6 +99,7 @@ class TestNamespaceHybridSearchFulltextMultiCollMultiNs: """Large-scale hybrid_search FTS across 2 collections x 2 loaded namespaces.""" def _run_fts_case_all_quadrants(self, db_client, case_name: str) -> None: + """Run fts case all quadrants.""" ctx = _setup_multi_coll_multi_ns_isolation_fts(db_client) try: case = get_fts_case(case_name) @@ -133,15 +134,19 @@ def test_cross_quadrant_fts_isolation(self, db_client): teardown_multi_coll_multi_ns_fts(db_client, ctx) def test_hybrid_search_fulltext_contains_token_zpx(self, db_client): + """Test hybrid search fulltext contains token zpx.""" self._run_fts_case_all_quadrants(db_client, "contains_token_zpx") def test_hybrid_search_fulltext_contains_token_alp(self, db_client): + """Test hybrid search fulltext contains token alp.""" self._run_fts_case_all_quadrants(db_client, "contains_token_alp") def test_hybrid_search_fulltext_contains_string_shorthand(self, db_client): + """Test hybrid search fulltext contains string shorthand.""" self._run_fts_case_all_quadrants(db_client, "string_shorthand_token_zpx") def test_hybrid_search_fulltext_not_contains(self, db_client): + """Test hybrid search fulltext not contains.""" case = get_fts_case("not_contains_token_zpx") ctx = _setup_multi_coll_multi_ns_isolation_fts(db_client) try: @@ -153,15 +158,19 @@ def test_hybrid_search_fulltext_not_contains(self, db_client): teardown_multi_coll_multi_ns_fts(db_client, ctx) def test_hybrid_search_fulltext_and_zpx_alp(self, db_client): + """Test hybrid search fulltext and zpx alp.""" self._run_fts_case_all_quadrants(db_client, "and_zpx_alp") def test_hybrid_search_fulltext_and_multi_contains(self, db_client): + """Test hybrid search fulltext and multi contains.""" self._run_fts_case_all_quadrants(db_client, "and_multi_contains_phrase") def test_hybrid_search_fulltext_or_zpx_alp(self, db_client): + """Test hybrid search fulltext or zpx alp.""" self._run_fts_case_all_quadrants(db_client, "or_zpx_alp") def test_hybrid_search_fulltext_top1_contains_zpx(self, db_client): + """Test hybrid search fulltext top1 contains zpx.""" ctx = _setup_multi_coll_multi_ns_isolation_fts(db_client) try: for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: diff --git a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py index 3fbc0243..2c3fe958 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py @@ -33,6 +33,7 @@ class TestNamespaceHybridSearchMultiCollMultiNs: + """TestNamespaceHybridSearchMultiCollMultiNs class.""" def test_cross_quadrant_search_index_isolation(self, db_client): """Only the loaded quadrant returns ``has_both`` rows.""" ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") @@ -49,6 +50,7 @@ def test_cross_quadrant_search_index_isolation(self, db_client): teardown_multi_coll_multi_ns_fts(db_client, ctx) def test_cross_quadrant_fts_isolation(self, db_client): + """Test cross quadrant fts isolation.""" ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") try: case = get_fts_case("contains_token_zpx") @@ -76,6 +78,7 @@ def test_cross_quadrant_fts_isolation(self, db_client): ], ) def test_hybrid_search_fulltext_all_loaded_quadrants(self, db_client, case_name: str): + """Test hybrid search fulltext all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client) try: run_hybrid_search_fts_case_all_quadrants(ctx, case_name) @@ -94,6 +97,7 @@ def test_hybrid_search_fulltext_all_loaded_quadrants(self, db_client, case_name: ], ) def test_hybrid_search_search_index_all_loaded_quadrants(self, db_client, case_name: str): + """Test hybrid search search index all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client) try: run_search_index_case_on_quadrants(ctx, case_name) @@ -105,6 +109,7 @@ def test_hybrid_search_search_index_all_loaded_quadrants(self, db_client, case_n def test_hybrid_search_vector_all_loaded_quadrants( self, db_client, vector_distance: VectorDistanceMetric, case_name: str ): + """Test hybrid search vector all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: run_knn_case_on_quadrants(ctx, case_name, distance_metric=vector_distance) @@ -112,6 +117,7 @@ def test_hybrid_search_vector_all_loaded_quadrants( teardown_multi_coll_multi_ns_fts(db_client, ctx) def test_hybrid_search_fts_plus_search_index_multi_coll(self, db_client): + """Test hybrid search fts plus search index multi coll.""" ctx = setup_multi_coll_multi_ns_fts(db_client) try: for key in ("c1_x", "c2_y"): @@ -126,6 +132,7 @@ def test_hybrid_search_fts_plus_search_index_multi_coll(self, db_client): def test_hybrid_search_combined_rrf_multi_coll( self, db_client, vector_distance: VectorDistanceMetric ): + """Test hybrid search combined rrf multi coll.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: case = get_hybrid_combined_case("fts_zpx_filter_gte_knn") @@ -141,6 +148,7 @@ def test_hybrid_search_combined_rrf_multi_coll( def test_hybrid_search_vector_plus_search_index_multi_coll( self, db_client, vector_distance: VectorDistanceMetric ): + """Test hybrid search vector plus search index multi coll.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: knn_case = get_vector_knn_case("knn_filter_has_both") diff --git a/tests/integration_tests/test_namespace_hybrid_search_search_index.py b/tests/integration_tests/test_namespace_hybrid_search_search_index.py index 448e4168..b716ac63 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_search_index.py +++ b/tests/integration_tests/test_namespace_hybrid_search_search_index.py @@ -28,6 +28,7 @@ class TestNamespaceHybridSearchSearchIndex: @pytest.fixture(autouse=True) def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest) -> None: + """Bind shared collection.""" mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" if mode not in self._shared_by_mode: corpus, collection = setup_large_fts_collection(db_client) @@ -49,11 +50,13 @@ def _bind_shared_collection(self, db_client: Any, request: pytest.FixtureRequest @classmethod def teardown_class(cls) -> None: + """Teardown class.""" for entry in cls._shared_by_mode.values(): teardown_large_fts_collection(entry["db_client"], entry["collection"]) cls._shared_by_mode.clear() def _run_case(self, case_name: str) -> None: + """Run case.""" run_hybrid_search_index_case( self._namespace, self._corpus, get_search_index_case(case_name) ) diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py index ddfaa52c..55a76796 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py @@ -51,6 +51,7 @@ def _bind_shared_collection( vector_distance: VectorDistanceMetric, request: pytest.FixtureRequest, ) -> None: + """Bind shared collection.""" entry = ensure_shared_hybrid_search_collection( self._shared_by_mode, db_client, request, vector_distance ) @@ -60,11 +61,13 @@ def _bind_shared_collection( @classmethod def teardown_class(cls) -> None: + """Teardown class.""" for entry in cls._shared_by_mode.values(): teardown_large_fts_collection(entry["db_client"], entry["collection"]) cls._shared_by_mode.clear() def _new_namespace(self, case_name: str) -> Any: + """New namespace.""" return setup_fts_namespace_with_corpus( self._collection, self._corpus, @@ -73,6 +76,7 @@ def _new_namespace(self, case_name: str) -> Any: @pytest.mark.parametrize("case_name", [c.name for c in TRIPLE_BRANCH_CASES]) def test_hybrid_search_triple_branch(self, db_client, case_name: str): + """Test hybrid search triple branch.""" case = get_triple_branch_case(case_name) namespace = self._new_namespace(case_name) run_hybrid_triple_branch_case( diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py index 1be13535..12200003 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py @@ -62,6 +62,7 @@ def _namespace_corpus(base_corpus: list[CorpusRecord], key: str) -> list[CorpusRecord]: + """Namespace corpus.""" return [ CorpusRecord( doc_id=f"{key}_{record.doc_id}", @@ -118,6 +119,7 @@ def _setup_multi_coll_multi_ns_isolation_triple( class TestNamespaceHybridSearchTripleBranchMultiCollMultiNs: + """TestNamespaceHybridSearchTripleBranchMultiCollMultiNs class.""" @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) def test_cross_quadrant_triple_branch_isolation( self, db_client, vector_distance: VectorDistanceMetric @@ -149,6 +151,7 @@ def test_cross_quadrant_triple_branch_isolation( teardown_multi_coll_multi_ns_fts(db_client, ctx) def test_cross_quadrant_search_index_branch_isolation(self, db_client): + """Test cross quadrant search index branch isolation.""" ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") try: case = get_triple_branch_case("si_and_has_both_zpx_gte") @@ -166,6 +169,7 @@ def test_cross_quadrant_search_index_branch_isolation(self, db_client): def test_triple_branch_top1_fts_all_quadrants( self, db_client, vector_distance: VectorDistanceMetric ): + """Test triple branch top1 fts all quadrants.""" ctx = _setup_multi_coll_multi_ns_isolation_triple( db_client, distance=vector_distance ) @@ -194,6 +198,7 @@ def test_triple_branch_top1_fts_all_quadrants( @pytest.mark.parametrize("case_name", _FTS_TRIPLE_CASE_NAMES) def test_hybrid_search_triple_branch_fts_all_loaded_quadrants(self, db_client, case_name: str): + """Test hybrid search triple branch fts all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client) try: run_triple_branch_case_on_quadrants(ctx, case_name) @@ -204,6 +209,7 @@ def test_hybrid_search_triple_branch_fts_all_loaded_quadrants(self, db_client, c def test_hybrid_search_triple_branch_search_index_all_loaded_quadrants( self, db_client, case_name: str ): + """Test hybrid search triple branch search index all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client) try: run_triple_branch_case_on_quadrants(ctx, case_name) @@ -215,6 +221,7 @@ def test_hybrid_search_triple_branch_search_index_all_loaded_quadrants( def test_hybrid_search_triple_branch_knn_all_loaded_quadrants( self, db_client, vector_distance: VectorDistanceMetric, case_name: str ): + """Test hybrid search triple branch knn all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: run_triple_branch_case_on_quadrants( @@ -228,6 +235,7 @@ def test_hybrid_search_triple_branch_knn_all_loaded_quadrants( def test_hybrid_search_triple_branch_intersection_all_loaded_quadrants( self, db_client, vector_distance: VectorDistanceMetric, case_name: str ): + """Test hybrid search triple branch intersection all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: run_triple_branch_case_on_quadrants( @@ -241,6 +249,7 @@ def test_hybrid_search_triple_branch_intersection_all_loaded_quadrants( def test_hybrid_search_triple_branch_rrf_all_loaded_quadrants( self, db_client, vector_distance: VectorDistanceMetric, case_name: str ): + """Test hybrid search triple branch rrf all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: run_triple_branch_case_on_quadrants( diff --git a/tests/integration_tests/test_namespace_hybrid_search_vector.py b/tests/integration_tests/test_namespace_hybrid_search_vector.py index cafef9e0..086ce2b5 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_vector.py +++ b/tests/integration_tests/test_namespace_hybrid_search_vector.py @@ -38,6 +38,7 @@ def _bind_shared_collection( vector_distance: VectorDistanceMetric, request: pytest.FixtureRequest, ) -> None: + """Bind shared collection.""" entry = ensure_shared_hybrid_search_collection( self._shared_by_mode, db_client, request, vector_distance ) @@ -47,11 +48,13 @@ def _bind_shared_collection( @classmethod def teardown_class(cls) -> None: + """Teardown class.""" for entry in cls._shared_by_mode.values(): teardown_large_fts_collection(entry["db_client"], entry["collection"]) cls._shared_by_mode.clear() def _new_namespace(self, case_name: str) -> Any: + """New namespace.""" ns_name = f"ns_knn_{case_name}_{int(time.time() * 1000)}" return setup_fts_namespace_with_corpus( self._collection, @@ -60,6 +63,7 @@ def _new_namespace(self, case_name: str) -> Any: ) def _run_knn_case(self, case_name: str) -> None: + """Run knn case.""" namespace = self._new_namespace(case_name) run_hybrid_knn_case( namespace, @@ -70,6 +74,7 @@ def _run_knn_case(self, case_name: str) -> None: @pytest.mark.parametrize("case_name", [c.name for c in VECTOR_KNN_CASES]) def test_hybrid_search_vector_knn_cases(self, db_client, case_name: str): + """Test hybrid search vector knn cases.""" self._run_knn_case(case_name) def test_hybrid_search_vector_top1_nearest(self, db_client): diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 07097976..7053a485 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -18,7 +18,9 @@ class TestNamespaceLifecycle: + """TestNamespaceLifecycle class.""" def _create_ns_collection(self, client, suffix=""): + """Create ns collection.""" name = f"test_ns_lc_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( @@ -33,6 +35,7 @@ def _create_ns_collection(self, client, suffix=""): return collection def test_create_namespace_collection(self, db_client): + """Test create namespace collection.""" collection = self._create_ns_collection(db_client) try: assert collection.use_namespace is True @@ -42,6 +45,7 @@ def test_create_namespace_collection(self, db_client): db_client.delete_collection(name=collection.name) def test_get_collection_preserves_namespace_flag(self, db_client): + """Test get collection preserves namespace flag.""" collection = self._create_ns_collection(db_client) try: retrieved = db_client.get_collection(collection.name) @@ -51,6 +55,7 @@ def test_get_collection_preserves_namespace_flag(self, db_client): db_client.delete_collection(name=collection.name) def test_get_collection_restores_embedding_function(self, db_client): + """Test get collection restores embedding function.""" from pyseekdb import DefaultEmbeddingFunction ef = DefaultEmbeddingFunction() @@ -74,6 +79,7 @@ def test_get_collection_restores_embedding_function(self, db_client): db_client.delete_collection(name=name) def test_create_and_get_namespace(self, db_client): + """Test create and get namespace.""" collection = self._create_ns_collection(db_client) try: ns = collection.create_namespace("ns_a") @@ -87,6 +93,7 @@ def test_create_and_get_namespace(self, db_client): db_client.delete_collection(name=collection.name) def test_get_or_create_namespace(self, db_client): + """Test get or create namespace.""" collection = self._create_ns_collection(db_client) try: ns1 = collection.get_or_create_namespace("ns_goc") @@ -98,6 +105,7 @@ def test_get_or_create_namespace(self, db_client): db_client.delete_collection(name=collection.name) def test_has_namespace(self, db_client): + """Test has namespace.""" collection = self._create_ns_collection(db_client) try: assert collection.has_namespace("nonexistent") is False @@ -107,6 +115,7 @@ def test_has_namespace(self, db_client): db_client.delete_collection(name=collection.name) def test_list_namespaces(self, db_client): + """Test list namespaces.""" collection = self._create_ns_collection(db_client) try: collection.create_namespace("ns_x") @@ -120,6 +129,7 @@ def test_list_namespaces(self, db_client): db_client.delete_collection(name=collection.name) def test_delete_namespace(self, db_client): + """Test delete namespace.""" collection = self._create_ns_collection(db_client) try: collection.create_namespace("ns_del") @@ -131,6 +141,7 @@ def test_delete_namespace(self, db_client): db_client.delete_collection(name=collection.name) def test_get_nonexistent_namespace_raises(self, db_client): + """Test get nonexistent namespace raises.""" collection = self._create_ns_collection(db_client) try: with pytest.raises(ValueError): @@ -139,6 +150,7 @@ def test_get_nonexistent_namespace_raises(self, db_client): db_client.delete_collection(name=collection.name) def test_delete_collection_cleans_namespaces(self, db_client): + """Test delete collection cleans namespaces.""" collection = self._create_ns_collection(db_client) coll_name = collection.name collection.create_namespace("ns_cleanup") @@ -228,6 +240,7 @@ def test_get_collection_purges_broken_ns_collection(self, oceanbase_client): client.delete_collection(name=name) def test_namespace_ops_blocked_after_collection_deleted(self, db_client): + """Test namespace ops blocked after collection deleted.""" collection = self._create_ns_collection(db_client) db_client.delete_collection(name=collection.name) @@ -241,6 +254,7 @@ def test_namespace_ops_blocked_after_collection_deleted(self, db_client): collection.has_namespace("demo") def test_collection_data_api_blocked_when_namespace_enabled(self, db_client): + """Test collection data api blocked when namespace enabled.""" collection = self._create_ns_collection(db_client) try: with pytest.raises(ValueError, match="namespace enabled"): @@ -249,6 +263,7 @@ def test_collection_data_api_blocked_when_namespace_enabled(self, db_client): db_client.delete_collection(name=collection.name) def test_namespace_on_non_namespace_collection_raises(self, db_client): + """Test namespace on non namespace collection raises.""" name = f"test_nons_{int(time.time() * 1000)}" collection = db_client.create_collection( name=name, @@ -269,6 +284,7 @@ def test_namespace_on_non_namespace_collection_raises(self, db_client): def test_create_namespace_collection_with_hnsw_raises(self, db_client): + """Test create namespace collection with hnsw raises.""" from pyseekdb.client.configuration import HNSWConfiguration name = f"test_ns_hnsw_{int(time.time() * 1000)}" schema = Schema( @@ -357,6 +373,7 @@ def test_custom_collection_partition_count(self, oceanbase_client): oceanbase_client.delete_collection(name=collection.name) def test_partition_count_rejected_for_non_namespace(self, db_client): + """Test partition count rejected for non namespace.""" name = f"test_nons_pc_{int(time.time() * 1000)}" with pytest.raises(ValueError, match="partition_count is only supported"): db_client.create_collection( diff --git a/tests/integration_tests/test_namespace_lifecycle_concurrency.py b/tests/integration_tests/test_namespace_lifecycle_concurrency.py index 581c30b0..e50df851 100644 --- a/tests/integration_tests/test_namespace_lifecycle_concurrency.py +++ b/tests/integration_tests/test_namespace_lifecycle_concurrency.py @@ -15,10 +15,12 @@ def _unique_name(prefix: str) -> str: + """Unique name.""" return f"{prefix}_{time.time_ns()}" def _new_oceanbase_client(): + """New oceanbase client.""" return pyseekdb.Client( host=os.environ.get("OB_HOST", "127.0.0.1"), port=int(os.environ.get("OB_PORT", "10902")), @@ -47,6 +49,7 @@ def test_multi_client_has_false_after_drop_returns(oceanbase_client): collection.create_namespace(ns_name) def _drop_worker() -> None: + """Drop worker.""" try: clients[0].get_collection(collection.name).delete_namespace(ns_name) except Exception as exc: # noqa: BLE001 @@ -55,6 +58,7 @@ def _drop_worker() -> None: drop_done.set() def _has_worker() -> None: + """Has worker.""" try: coll = clients[1].get_collection(collection.name) drop_done.wait(timeout=60) @@ -104,12 +108,14 @@ def test_multi_client_has_does_not_error_during_drop(oceanbase_client): real_execute = clients[0]._server._execute def _slow_drop_execute(sql, *args, **kwargs): + """Slow drop execute.""" if "DROP_NAMESPACE" in str(sql): drop_started.set() assert drop_can_finish.wait(timeout=30), "drop blocked too long" return real_execute(sql, *args, **kwargs) def _drop_worker() -> None: + """Drop worker.""" try: with patch.object(clients[0]._server, "_execute", side_effect=_slow_drop_execute): dropper.delete_namespace(ns_name) @@ -117,6 +123,7 @@ def _drop_worker() -> None: errors.append(exc) def _has_worker() -> None: + """Has worker.""" try: assert drop_started.wait(timeout=30), "drop did not start" for _ in range(10): diff --git a/tests/integration_tests/test_namespace_optional_indexes.py b/tests/integration_tests/test_namespace_optional_indexes.py index ba957da2..cdcebeb0 100644 --- a/tests/integration_tests/test_namespace_optional_indexes.py +++ b/tests/integration_tests/test_namespace_optional_indexes.py @@ -29,6 +29,7 @@ def _schema_fts_only() -> Schema: + """Schema fts only.""" return Schema( vector_index=VectorIndexConfig(embedding_function=None), fulltext_index=FulltextIndexConfig(analyzer="ik"), @@ -36,6 +37,7 @@ def _schema_fts_only() -> Schema: def _schema_vector_only() -> Schema: + """Schema vector only.""" return Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), @@ -45,16 +47,19 @@ def _schema_vector_only() -> Schema: def _schema_search_only() -> Schema: + """Schema search only.""" return Schema(vector_index=VectorIndexConfig(embedding_function=None)) def _index_names(client: Any, collection_id: str) -> set[str]: + """Index names.""" table = NamespaceCollectionNames.data_table_name(collection_id) rows = client._server._execute(f"SHOW INDEX FROM `{table}`") return {(r.get("Key_name") if isinstance(r, dict) else r[2]) for r in (rows or [])} def _collection_settings(client: Any, collection_id: str) -> dict: + """Collection settings.""" rows = client._server._execute( f"SELECT settings FROM sdk_collections WHERE collection_id = '{collection_id}'" ) @@ -63,6 +68,7 @@ def _collection_settings(client: Any, collection_id: str) -> dict: def _create_collection(client: Any, label: str, schema: Schema) -> Any: + """Create collection.""" name = f"test_ns_opt_idx_{label}_{int(time.time() * 1000)}" return client.create_collection( name=name, schema=schema, use_namespace=True, @@ -71,12 +77,14 @@ def _create_collection(client: Any, label: str, schema: Schema) -> Any: def _unit_vector(dimension: int, axis: int = 0) -> list[float]: + """Unit vector.""" vec = [0.0] * dimension vec[axis % dimension] = 1.0 return vec def _seed_namespace(ns: Any, dimension: int) -> None: + """Seed namespace.""" ns.add( ids=["d1", "d2"], embeddings=[_unit_vector(dimension, 0), _unit_vector(dimension, 1)], @@ -89,6 +97,7 @@ class TestNamespaceOptionalIndexes: """DDL + query matrix for optional VECTOR / FULLTEXT indexes.""" def test_fts_only_indexes_and_queries(self, db_client): + """Test fts only indexes and queries.""" collection = _create_collection(db_client, "fts_only", _schema_fts_only()) try: keys = _index_names(db_client, collection.id) @@ -119,6 +128,7 @@ def test_fts_only_indexes_and_queries(self, db_client): db_client.delete_collection(name=collection.name) def test_vector_only_indexes_and_queries(self, db_client): + """Test vector only indexes and queries.""" collection = _create_collection(db_client, "vec_only", _schema_vector_only()) try: keys = _index_names(db_client, collection.id) @@ -149,6 +159,7 @@ def test_vector_only_indexes_and_queries(self, db_client): db_client.delete_collection(name=collection.name) def test_search_only_indexes_and_queries(self, db_client): + """Test search only indexes and queries.""" collection = _create_collection(db_client, "search_only", _schema_search_only()) try: keys = _index_names(db_client, collection.id) diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index b63bd12f..6744eb29 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -17,7 +17,9 @@ class TestNamespacePrewarm: + """TestNamespacePrewarm class.""" def _create_ns_collection_and_namespace(self, client): + """Create ns collection and namespace.""" from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT name = f"test_ns_pw_{int(time.time() * 1000)}" @@ -35,6 +37,7 @@ def _create_ns_collection_and_namespace(self, client): return collection, namespace def _query_hot_table(self, client, collection, namespace_id): + """Query hot table.""" coll_id = collection.id hot_table = NamespaceCollectionNames.hot_table_name(coll_id) rows = client._server._execute( diff --git a/tests/integration_tests/test_namespace_prewarm_lob.py b/tests/integration_tests/test_namespace_prewarm_lob.py index a2f758cc..a46c03d3 100644 --- a/tests/integration_tests/test_namespace_prewarm_lob.py +++ b/tests/integration_tests/test_namespace_prewarm_lob.py @@ -64,10 +64,12 @@ # ==================== SQL observability helpers ==================== def _exec(client, sql): + """Exec.""" return client._server._execute(sql) def _resolve_table_id(client, table_name): + """Resolve table id.""" rows = _exec( client, f"SELECT table_id FROM {ALL_TABLE_VIEW} WHERE table_name = '{table_name}'", @@ -189,7 +191,9 @@ def _grep_lob_prewarm_log(lob_meta_tablet_ids): class _BaseLobPrewarm: + """BaseLobPrewarm class.""" def _make_collection(self, client, name, partitions): + """Make collection.""" schema = Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration(dimension=3, distance="cosine"), @@ -229,6 +233,7 @@ def _force_out_of_row_lob(self, client, collection_id, namespace_id): class TestLobPrewarmStructural(_BaseLobPrewarm): + """TestLobPrewarmStructural class.""" def test_kv_table_has_lob_meta_tablet(self, oceanbase_client): """Tier 1: the prewarm target (kv_data_table) owns an aux LOB-meta tablet.""" name = f"lob_pw_struct_{int(time.time() * 1000)}" @@ -283,9 +288,11 @@ def test_multi_namespace_prewarm_is_independent(self, oceanbase_client): # ==================== Tier 2 + Tier 3: real LOB caching ==================== class TestLobPrewarmCaching(_BaseLobPrewarm): + """TestLobPrewarmCaching class.""" @staticmethod def _total_bytes(cache): + """Total bytes.""" return sum(v[1] for v in (cache or {}).values()) def test_out_of_row_lob_is_prewarmed_into_cache(self, oceanbase_client): diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index a0613300..2f98d6dd 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -15,7 +15,9 @@ class TestNamespaceQuery: + """TestNamespaceQuery class.""" def _setup(self, client, suffix=""): + """Setup.""" name = f"test_ns_q_{int(time.time() * 1000)}{suffix}" schema = Schema( vector_index=VectorIndexConfig( @@ -31,6 +33,7 @@ def _setup(self, client, suffix=""): return collection def _insert_data(self, ns): + """Insert data.""" ns.add( ids=["q1", "q2", "q3", "q4", "q5"], embeddings=[ @@ -57,6 +60,7 @@ def _insert_data(self, ns): ) def test_basic_vector_query(self, db_client): + """Test basic vector query.""" collection = self._setup(db_client) ns = collection.create_namespace("qns") try: @@ -70,6 +74,7 @@ def test_basic_vector_query(self, db_client): db_client.delete_collection(name=collection.name) def test_query_with_metadata_filter(self, db_client): + """Test query with metadata filter.""" collection = self._setup(db_client) ns = collection.create_namespace("qns_meta") try: @@ -88,6 +93,7 @@ def test_query_with_metadata_filter(self, db_client): db_client.delete_collection(name=collection.name) def test_query_with_score_filter(self, db_client): + """Test query with score filter.""" collection = self._setup(db_client) ns = collection.create_namespace("qns_score") try: @@ -103,6 +109,7 @@ def test_query_with_score_filter(self, db_client): db_client.delete_collection(name=collection.name) def test_query_with_include(self, db_client): + """Test query with include.""" collection = self._setup(db_client) ns = collection.create_namespace("qns_inc") try: @@ -120,6 +127,7 @@ def test_query_with_include(self, db_client): db_client.delete_collection(name=collection.name) def test_query_include_embeddings(self, db_client): + """Test query include embeddings.""" collection = self._setup(db_client) ns = collection.create_namespace("qns_emb") try: @@ -136,6 +144,7 @@ def test_query_include_embeddings(self, db_client): db_client.delete_collection(name=collection.name) def test_multi_namespace_isolation(self, db_client): + """Test multi namespace isolation.""" collection = self._setup(db_client, suffix="_iso") ns_a = collection.create_namespace("tenant_a") ns_b = collection.create_namespace("tenant_b") @@ -167,6 +176,7 @@ def test_multi_namespace_isolation(self, db_client): db_client.delete_collection(name=collection.name) def test_same_id_different_namespaces(self, db_client): + """Test same id different namespaces.""" collection = self._setup(db_client, suffix="_sameid") ns_x = collection.create_namespace("ns_x") ns_y = collection.create_namespace("ns_y") @@ -188,6 +198,7 @@ def test_same_id_different_namespaces(self, db_client): # ==================== hybrid_search tests ==================== def test_hybrid_search_fulltext_only(self, db_client): + """Test hybrid search fulltext only.""" collection = self._setup(db_client) ns = collection.create_namespace("qns_hs_ft") try: @@ -210,6 +221,7 @@ def test_hybrid_search_fulltext_only(self, db_client): db_client.delete_collection(name=collection.name) def test_hybrid_search_vector_only(self, db_client): + """Test hybrid search vector only.""" collection = self._setup(db_client) ns = collection.create_namespace("qns_hs_vec") try: @@ -231,6 +243,7 @@ def test_hybrid_search_vector_only(self, db_client): db_client.delete_collection(name=collection.name) def test_hybrid_search_combined(self, db_client): + """Test hybrid search combined.""" collection = self._setup(db_client) ns = collection.create_namespace("qns_hs_comb") try: @@ -251,6 +264,7 @@ def test_hybrid_search_combined(self, db_client): db_client.delete_collection(name=collection.name) def test_hybrid_search_namespace_isolation(self, db_client): + """Test hybrid search namespace isolation.""" collection = self._setup(db_client, suffix="_hs_iso") ns_a = collection.create_namespace("hs_tenant_a") ns_b = collection.create_namespace("hs_tenant_b") diff --git a/tests/integration_tests/test_namespace_reconnect.py b/tests/integration_tests/test_namespace_reconnect.py index 3d4cd23a..487dc02d 100644 --- a/tests/integration_tests/test_namespace_reconnect.py +++ b/tests/integration_tests/test_namespace_reconnect.py @@ -36,6 +36,7 @@ def _new_oceanbase_client() -> Any: class TestNamespaceReconnectAndMultiClient: + """TestNamespaceReconnectAndMultiClient class.""" def test_same_client_reconnect_reads_namespace_data(self, oceanbase_client): """After _cleanup(), the next SDK call should reconnect and read prior rows.""" client = oceanbase_client diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py index 166b145b..9c350047 100644 --- a/tests/integration_tests/test_namespace_session_vars.py +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -17,7 +17,9 @@ class TestNamespaceSessionVars: + """TestNamespaceSessionVars class.""" def _create_ns_collection(self, client): + """Create ns collection.""" name = f"test_ns_sessvar_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( @@ -31,6 +33,7 @@ def _create_ns_collection(self, client): ) def _query_session_vars(self, client): + """Query session vars.""" rows = client._server._execute( "SELECT @collection_id AS cid, @namespace_id AS nsid, @ltable_id AS ltid" ) @@ -40,6 +43,7 @@ def _query_session_vars(self, client): return {"cid": row.get("cid"), "nsid": row.get("nsid"), "ltid": row.get("ltid")} def test_session_collection_id_after_create(self, db_client): + """Test session collection id after create.""" collection = self._create_ns_collection(db_client) try: vars_ = self._query_session_vars(db_client) @@ -49,6 +53,7 @@ def test_session_collection_id_after_create(self, db_client): db_client.delete_collection(name=collection.name) def test_session_ns_ids_after_create_namespace(self, db_client): + """Test session ns ids after create namespace.""" collection = self._create_ns_collection(db_client) try: ns = collection.create_namespace("sess_ns") @@ -61,6 +66,7 @@ def test_session_ns_ids_after_create_namespace(self, db_client): db_client.delete_collection(name=collection.name) def test_session_ns_id_after_get_namespace(self, db_client): + """Test session ns id after get namespace.""" collection = self._create_ns_collection(db_client) try: ns1 = collection.create_namespace("sess_get_ns") diff --git a/tests/integration_tests/test_namespace_upsert_concurrency.py b/tests/integration_tests/test_namespace_upsert_concurrency.py index c0f27fa6..8e4b6975 100644 --- a/tests/integration_tests/test_namespace_upsert_concurrency.py +++ b/tests/integration_tests/test_namespace_upsert_concurrency.py @@ -15,10 +15,12 @@ def _unique_name(prefix: str) -> str: + """Unique name.""" return f"{prefix}_{time.time_ns()}" def _new_oceanbase_client(): + """New oceanbase client.""" return pyseekdb.Client( host=os.environ.get("OB_HOST", "127.0.0.1"), port=int(os.environ.get("OB_PORT", "10902")), @@ -32,6 +34,7 @@ def _new_oceanbase_client(): def test_multi_client_concurrent_upsert_same_new_id_should_keep_single_record( oceanbase_client, ): + """Test multi client concurrent upsert same new id should keep single record.""" collection = oceanbase_client.create_collection( name=_unique_name("test_ns_upsert_race_repro"), schema=ns_schema(), @@ -49,6 +52,7 @@ def test_multi_client_concurrent_upsert_same_new_id_should_keep_single_record( errors: list[Exception] = [] def _upsert_once(target_ns, idx: int) -> None: + """Upsert once.""" try: barrier.wait(timeout=30) target_ns.upsert( diff --git a/tests/unit_tests/test_build_document_query.py b/tests/unit_tests/test_build_document_query.py index 58c93fe2..347f04ce 100644 --- a/tests/unit_tests/test_build_document_query.py +++ b/tests/unit_tests/test_build_document_query.py @@ -17,7 +17,9 @@ def client(): class TestBuildDocumentQuery: + """TestBuildDocumentQuery class.""" def test_and_contains_has_default_operator_and(self, client): + """Test and contains has default operator and.""" where_document = {"$and": [{"$contains": "TOKENZPX"}, {"$contains": "TOKENALP"}]} result = client._build_document_query(where_document) assert result is not None @@ -28,6 +30,7 @@ def test_and_contains_has_default_operator_and(self, client): assert qs["fields"] == ["document"] def test_or_contains_has_default_operator_or(self, client): + """Test or contains has default operator or.""" where_document = {"$or": [{"$contains": "TOKENZPX"}, {"$contains": "TOKENALP"}]} result = client._build_document_query(where_document) assert result is not None @@ -38,6 +41,7 @@ def test_or_contains_has_default_operator_or(self, client): assert " OR " not in qs["query"] def test_single_contains_no_default_operator(self, client): + """Test single contains no default operator.""" where_document = {"$contains": "hello"} result = client._build_document_query(where_document) assert result is not None @@ -46,6 +50,7 @@ def test_single_contains_no_default_operator(self, client): assert qs["query"] == "hello" def test_and_contains_with_boost(self, client): + """Test and contains with boost.""" where_document = {"$and": [{"$contains": "foo"}, {"$contains": "bar"}]} result = client._build_document_query(where_document, boost=2.0) assert result is not None @@ -54,6 +59,7 @@ def test_and_contains_with_boost(self, client): assert qs["boost"] == 2.0 def test_and_contains_three_terms(self, client): + """Test and contains three terms.""" where_document = { "$and": [ {"$contains": "alpha"}, diff --git a/tests/unit_tests/test_build_knn_filter.py b/tests/unit_tests/test_build_knn_filter.py index 60e0b868..43959129 100644 --- a/tests/unit_tests/test_build_knn_filter.py +++ b/tests/unit_tests/test_build_knn_filter.py @@ -7,6 +7,7 @@ @pytest.fixture() def client(): + """Client.""" from pyseekdb.client.client_base import BaseClient with patch.multiple(BaseClient, __abstractmethods__=set()): @@ -15,7 +16,9 @@ def client(): class TestBuildKnnFilterNe: + """TestBuildKnnFilterNe class.""" def test_ne_hoists_must_not_in_knn_filter(self, client): + """Test ne hoists must not in knn filter.""" with patch.object( client, "_build_metadata_filter_for_search_parm", diff --git a/tests/unit_tests/test_build_query_expression.py b/tests/unit_tests/test_build_query_expression.py index 4797c56f..476e59bc 100644 --- a/tests/unit_tests/test_build_query_expression.py +++ b/tests/unit_tests/test_build_query_expression.py @@ -7,6 +7,7 @@ @pytest.fixture() def client(): + """Client.""" from pyseekdb.client.client_base import BaseClient with patch.multiple(BaseClient, __abstractmethods__=set()): @@ -15,6 +16,7 @@ def client(): class TestBuildQueryExpressionNotContains: + """TestBuildQueryExpressionNotContains class.""" def test_not_contains_with_metadata_filter_hoists_must_not(self, client): """$not_contains + where must not nest a must_not-only bool inside must.""" where = {"seq": {"$gte": 43}} @@ -65,6 +67,7 @@ def test_not_contains_only_hoists_must_not_without_synthetic_filter(self, client assert "filter" not in expr["bool"] def test_contains_with_metadata_filter_still_uses_must(self, client): + """Test contains with metadata filter still uses must.""" with patch.object( client, "_build_metadata_filter_for_search_parm", @@ -81,6 +84,7 @@ def test_contains_with_metadata_filter_still_uses_must(self, client): class TestBuildQueryExpressionMetadataNe: + """TestBuildQueryExpressionMetadataNe class.""" def test_ne_with_fts_hoists_must_not_from_filter(self, client): """$ne in where must not appear as a must_not-only bool inside filter.""" with patch.object( @@ -105,6 +109,7 @@ def test_ne_with_fts_hoists_must_not_from_filter(self, client): assert set(nested.keys()) != {"must_not"}, nested def test_ne_with_not_contains_and_positive_where(self, client): + """Test ne with not contains and positive where.""" with patch.object( client, "_build_metadata_filter_for_search_parm", diff --git a/tests/unit_tests/test_get_or_create_collection_concurrency.py b/tests/unit_tests/test_get_or_create_collection_concurrency.py index 89d259f1..53a863fc 100644 --- a/tests/unit_tests/test_get_or_create_collection_concurrency.py +++ b/tests/unit_tests/test_get_or_create_collection_concurrency.py @@ -21,8 +21,11 @@ class TestCollectionCatalogConflictDetection: + """TestCollectionCatalogConflictDetection class.""" def test_detects_integrity_error_on_sdk_collections(self): + """Test detects integrity error on sdk collections.""" class IntegrityError(Exception): + """IntegrityError class.""" pass exc = IntegrityError( @@ -31,20 +34,25 @@ class IntegrityError(Exception): assert _is_sdk_collection_catalog_conflict_error(exc) def test_ignores_unrelated_errors(self): + """Test ignores unrelated errors.""" assert not _is_sdk_collection_catalog_conflict_error(ValueError("invalid dimension")) class TestCollectionCatalogInsertRecovery: + """TestCollectionCatalogInsertRecovery class.""" def test_insert_conflict_reuses_existing_collection_id(self): + """Test insert conflict reuses existing collection id.""" client = MagicMock(spec=BaseClient) client._get_collection_id.side_effect = [ValueError("not found"), "existing_id"] conn = MagicMock() client._ensure_connection.return_value = conn class IntegrityError(Exception): + """IntegrityError class.""" pass def execute_side_effect(sql): + """Execute side effect.""" if "INSERT INTO" in sql: raise IntegrityError( '(1062, "Duplicate entry \'items\' for key \'uk_sdk_coll_name\'")' @@ -60,6 +68,7 @@ def execute_side_effect(sql): assert client._get_collection_id.call_count == 2 def test_existing_catalog_row_is_reused_without_insert(self): + """Test existing catalog row is reused without insert.""" client = MagicMock(spec=BaseClient) client._get_collection_id.return_value = "existing_id" @@ -73,25 +82,32 @@ def test_existing_catalog_row_is_reused_without_insert(self): class TestCollectionConflictDetection: + """TestCollectionConflictDetection class.""" def test_detects_value_error_for_existing_collection(self): + """Test detects value error for existing collection.""" assert _is_collection_conflict_error(ValueError("Collection 'items' already exists")) def test_detects_seekdb_table_exists_error(self): + """Test detects seekdb table exists error.""" class SeekdbError(Exception): + """SeekdbError class.""" pass exc = SeekdbError("Table 'c$v2$abc' already exists failed: code=1050") assert _is_collection_conflict_error(exc) def test_ignores_unrelated_errors(self): + """Test ignores unrelated errors.""" assert not _is_collection_conflict_error(ValueError("invalid dimension")) def test_ignores_metadata_failure_without_conflict_cause(self): + """Test ignores metadata failure without conflict cause.""" assert not _is_collection_conflict_error( ValueError("Failed to create collection metadata: Collection not found: 'items'") ) def test_detects_conflict_in_cause_chain(self): + """Test detects conflict in cause chain.""" inner = Exception("Table 'c$v2$abc' already exists failed: code=1050") outer = ValueError("Failed to create collection metadata: duplicate entry") outer.__cause__ = inner @@ -99,13 +115,16 @@ def test_detects_conflict_in_cause_chain(self): class TestGetOrCreateCollectionRecovery: + """TestGetOrCreateCollectionRecovery class.""" @staticmethod def _bind_resume_helper(client): + """Bind resume helper.""" client._get_or_resume_existing_collection = ( BaseClient._get_or_resume_existing_collection.__get__(client, BaseClient) ) def test_returns_existing_collection_after_create_conflict(self): + """Test returns existing collection after create conflict.""" client = MagicMock(spec=BaseClient) self._bind_resume_helper(client) client.has_collection.return_value = False @@ -119,6 +138,7 @@ def test_returns_existing_collection_after_create_conflict(self): client.get_collection.assert_called_once_with("items", embedding_function=_NOT_PROVIDED) def test_retries_get_after_wrapped_table_conflict(self): + """Test retries get after wrapped table conflict.""" client = MagicMock(spec=BaseClient) self._bind_resume_helper(client) client.has_collection.return_value = False @@ -135,6 +155,7 @@ def test_retries_get_after_wrapped_table_conflict(self): client.get_collection.assert_called_once_with("items", embedding_function=_NOT_PROVIDED) def test_conflict_on_namespace_collection_resumes_incomplete_handle(self): + """Test conflict on namespace collection resumes incomplete handle.""" client = MagicMock(spec=BaseClient) self._bind_resume_helper(client) resumed = object() @@ -157,6 +178,7 @@ def test_conflict_on_namespace_collection_resumes_incomplete_handle(self): client.get_collection.assert_not_called() def test_conflict_on_namespace_collection_without_metadata_reraises(self): + """Test conflict on namespace collection without metadata reraises.""" client = MagicMock(spec=BaseClient) self._bind_resume_helper(client) client.has_collection.return_value = False @@ -167,6 +189,7 @@ def test_conflict_on_namespace_collection_without_metadata_reraises(self): BaseClient.get_or_create_collection(client, "items", use_namespace=True) def test_resumes_incomplete_namespace_collection_when_present(self): + """Test resumes incomplete namespace collection when present.""" client = MagicMock(spec=BaseClient) resumed = object() client.has_collection.return_value = True @@ -180,6 +203,7 @@ def test_resumes_incomplete_namespace_collection_when_present(self): client.get_collection.assert_not_called() def test_does_not_mask_unrelated_create_errors(self): + """Test does not mask unrelated create errors.""" client = MagicMock(spec=BaseClient) client.has_collection.return_value = False client.create_collection.side_effect = ValueError("invalid dimension") @@ -189,7 +213,9 @@ def test_does_not_mask_unrelated_create_errors(self): class TestListNsNamespacesRecyclebinFilter: + """TestListNsNamespacesRecyclebinFilter class.""" def test_sql_excludes_recyclebin_rows(self): + """Test sql excludes recyclebin rows.""" client = MagicMock(spec=BaseClient) client._qtable.return_value = "`sdk_namespaces`" client._execute.return_value = [ diff --git a/tests/unit_tests/test_get_or_create_namespace_concurrency.py b/tests/unit_tests/test_get_or_create_namespace_concurrency.py index 52f23fe6..5f2935e7 100644 --- a/tests/unit_tests/test_get_or_create_namespace_concurrency.py +++ b/tests/unit_tests/test_get_or_create_namespace_concurrency.py @@ -19,8 +19,11 @@ class TestNamespaceCatalogConflictDetection: + """TestNamespaceCatalogConflictDetection class.""" def test_detects_integrity_error_on_sdk_namespaces(self): + """Test detects integrity error on sdk namespaces.""" class IntegrityError(Exception): + """IntegrityError class.""" pass exc = IntegrityError( @@ -29,7 +32,9 @@ class IntegrityError(Exception): assert _is_namespace_catalog_conflict_error(exc) def test_detects_integrity_error_on_sdk_ltables(self): + """Test detects integrity error on sdk ltables.""" class IntegrityError(Exception): + """IntegrityError class.""" pass exc = IntegrityError( @@ -38,11 +43,14 @@ class IntegrityError(Exception): assert _is_namespace_catalog_conflict_error(exc) def test_ignores_unrelated_errors(self): + """Test ignores unrelated errors.""" assert not _is_namespace_catalog_conflict_error(ValueError("invalid namespace name")) class TestGetOrCreateNamespaceMetaRecovery: + """TestGetOrCreateNamespaceMetaRecovery class.""" def test_idempotent_insert_reuses_existing_namespace(self): + """Test idempotent insert reuses existing namespace.""" client = MagicMock(spec=BaseClient) client._get_ns_namespace_meta.return_value = None client._insert_ns_namespace_catalog_row.return_value = 42 @@ -64,6 +72,7 @@ def test_idempotent_insert_reuses_existing_namespace(self): ) def test_existing_namespace_without_ltable_creates_default_ltable(self): + """Test existing namespace without ltable creates default ltable.""" client = MagicMock(spec=BaseClient) client._get_ns_namespace_meta.return_value = { "namespace_id": "42", diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 251c29dc..01d32193 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -29,7 +29,9 @@ class TestIVFConfiguration: + """TestIVFConfiguration class.""" def test_valid_defaults(self): + """Test valid defaults.""" config = IVFConfiguration() assert config.dimension == 384 assert config.distance == "cosine" @@ -37,93 +39,115 @@ def test_valid_defaults(self): assert config.properties is None def test_valid_custom(self): + """Test valid custom.""" config = IVFConfiguration(dimension=128, distance="l2", centroids_fresh_mode="spfresh") assert config.dimension == 128 assert config.distance == "l2" assert config.centroids_fresh_mode == "spfresh" def test_valid_inner_product(self): + """Test valid inner product.""" config = IVFConfiguration(dimension=1024, distance="inner_product") assert config.distance == "inner_product" def test_dimension_boundary_min(self): + """Test dimension boundary min.""" config = IVFConfiguration(dimension=1) assert config.dimension == 1 def test_dimension_boundary_max(self): + """Test dimension boundary max.""" config = IVFConfiguration(dimension=4096) assert config.dimension == 4096 def test_invalid_dimension_zero(self): + """Test invalid dimension zero.""" with pytest.raises(ValueError, match="must be between"): IVFConfiguration(dimension=0) def test_invalid_dimension_negative(self): + """Test invalid dimension negative.""" with pytest.raises(ValueError, match="must be between"): IVFConfiguration(dimension=-1) def test_invalid_dimension_too_large(self): + """Test invalid dimension too large.""" with pytest.raises(ValueError, match="must be between"): IVFConfiguration(dimension=4097) def test_invalid_dimension_type(self): + """Test invalid dimension type.""" with pytest.raises(TypeError, match="dimension must be an integer"): IVFConfiguration(dimension="128") def test_invalid_dimension_bool(self): + """Test invalid dimension bool.""" with pytest.raises(TypeError, match="dimension must be an integer"): IVFConfiguration(dimension=True) def test_invalid_distance(self): + """Test invalid distance.""" with pytest.raises(ValueError, match="distance must be one of"): IVFConfiguration(distance="invalid") def test_invalid_centroids_fresh_mode_type(self): + """Test invalid centroids fresh mode type.""" with pytest.raises(TypeError, match="centroids_fresh_mode must be a str"): IVFConfiguration(centroids_fresh_mode=True) def test_properties_valid(self): + """Test properties valid.""" config = IVFConfiguration(properties={"nlist": 128, "nprobe": 16}) assert config.properties["nlist"] == 128 def test_valid_type_default(self): + """Test valid type default.""" config = IVFConfiguration() assert config.type == "ivf_flat" def test_valid_type_ivf_sq8(self): + """Test valid type ivf sq8.""" config = IVFConfiguration(type="ivf_sq8") assert config.type == "ivf_sq8" def test_valid_type_ivf_pq(self): + """Test valid type ivf pq.""" config = IVFConfiguration(type="ivf_pq") assert config.type == "ivf_pq" def test_valid_type_enum(self): + """Test valid type enum.""" config = IVFConfiguration(type=IVFIndexType.IVF_SQ8) assert config.type == "ivf_sq8" def test_invalid_type(self): + """Test invalid type.""" with pytest.raises(ValueError, match="type must be one of"): IVFConfiguration(type="invalid") def test_valid_lib_default(self): + """Test valid lib default.""" config = IVFConfiguration() assert config.lib == "ob" def test_valid_lib_vsag(self): + """Test valid lib vsag.""" config = IVFConfiguration(lib="vsag") assert config.lib == "vsag" def test_valid_lib_enum(self): + """Test valid lib enum.""" from pyseekdb.client.configuration import IVFIndexLib config = IVFConfiguration(lib=IVFIndexLib.VSAG) assert config.lib == "vsag" def test_invalid_lib(self): + """Test invalid lib.""" with pytest.raises(ValueError, match="lib must be one of"): IVFConfiguration(lib="invalid") def test_properties_invalid_nested(self): + """Test properties invalid nested.""" with pytest.raises(TypeError): IVFConfiguration(properties={"bad": {"nested": True}}) @@ -133,25 +157,30 @@ def test_properties_invalid_nested(self): class TestVectorIndexConfig: + """TestVectorIndexConfig class.""" def test_ivf_only(self): + """Test ivf only.""" ivf = IVFConfiguration(dimension=128) config = VectorIndexConfig(ivf=ivf, embedding_function=None) assert config.ivf is ivf assert config.hnsw is None def test_hnsw_only(self): + """Test hnsw only.""" hnsw = HNSWConfiguration(dimension=128) config = VectorIndexConfig(hnsw=hnsw, embedding_function=None) assert config.hnsw is hnsw assert config.ivf is None def test_ivf_and_hnsw_mutually_exclusive(self): + """Test ivf and hnsw mutually exclusive.""" ivf = IVFConfiguration(dimension=128) hnsw = HNSWConfiguration(dimension=128) with pytest.raises(ValueError, match="Only one of ivf or hnsw"): VectorIndexConfig(ivf=ivf, hnsw=hnsw, embedding_function=None) def test_neither_ivf_nor_hnsw(self): + """Test neither ivf nor hnsw.""" config = VectorIndexConfig(embedding_function=None) assert config.ivf is None assert config.hnsw is None @@ -162,7 +191,9 @@ def test_neither_ivf_nor_hnsw(self): class TestCollectionNamespaceGuard: + """TestCollectionNamespaceGuard class.""" def _make_collection(self, use_namespace: bool) -> Collection: + """Make collection.""" mock_client = MagicMock() return Collection( client=mock_client, @@ -173,41 +204,49 @@ def _make_collection(self, use_namespace: bool) -> Collection: ) def test_guard_blocks_add_when_namespace_enabled(self): + """Test guard blocks add when namespace enabled.""" coll = self._make_collection(use_namespace=True) with pytest.raises(ValueError, match="namespace enabled"): coll.add(ids="1", embeddings=[1.0, 2.0, 3.0]) def test_guard_blocks_update_when_namespace_enabled(self): + """Test guard blocks update when namespace enabled.""" coll = self._make_collection(use_namespace=True) with pytest.raises(ValueError, match="namespace enabled"): coll.update(ids="1", metadatas={"k": "v"}) def test_guard_blocks_upsert_when_namespace_enabled(self): + """Test guard blocks upsert when namespace enabled.""" coll = self._make_collection(use_namespace=True) with pytest.raises(ValueError, match="namespace enabled"): coll.upsert(ids="1", embeddings=[1.0, 2.0, 3.0]) def test_guard_blocks_delete_when_namespace_enabled(self): + """Test guard blocks delete when namespace enabled.""" coll = self._make_collection(use_namespace=True) with pytest.raises(ValueError, match="namespace enabled"): coll.delete(ids="1") def test_guard_blocks_query_when_namespace_enabled(self): + """Test guard blocks query when namespace enabled.""" coll = self._make_collection(use_namespace=True) with pytest.raises(ValueError, match="namespace enabled"): coll.query(query_embeddings=[1.0, 2.0, 3.0]) def test_guard_blocks_get_when_namespace_enabled(self): + """Test guard blocks get when namespace enabled.""" coll = self._make_collection(use_namespace=True) with pytest.raises(ValueError, match="namespace enabled"): coll.get(ids="1") def test_guard_blocks_count_when_namespace_enabled(self): + """Test guard blocks count when namespace enabled.""" coll = self._make_collection(use_namespace=True) with pytest.raises(ValueError, match="namespace enabled"): coll.count() def test_guard_blocks_peek_when_namespace_enabled(self): + """Test guard blocks peek when namespace enabled.""" coll = self._make_collection(use_namespace=True) with pytest.raises(ValueError, match="namespace enabled"): coll.peek() @@ -218,7 +257,9 @@ def test_guard_blocks_peek_when_namespace_enabled(self): class TestCollectionNamespaceManagementGuard: + """TestCollectionNamespaceManagementGuard class.""" def _make_collection(self, use_namespace: bool) -> Collection: + """Make collection.""" mock_client = MagicMock() return Collection( client=mock_client, @@ -229,31 +270,37 @@ def _make_collection(self, use_namespace: bool) -> Collection: ) def test_create_namespace_requires_namespace_enabled(self): + """Test create namespace requires namespace enabled.""" coll = self._make_collection(use_namespace=False) with pytest.raises(ValueError, match="not enabled"): coll.create_namespace("ns1") def test_get_namespace_requires_namespace_enabled(self): + """Test get namespace requires namespace enabled.""" coll = self._make_collection(use_namespace=False) with pytest.raises(ValueError, match="not enabled"): coll.get_namespace("ns1") def test_get_or_create_namespace_requires_namespace_enabled(self): + """Test get or create namespace requires namespace enabled.""" coll = self._make_collection(use_namespace=False) with pytest.raises(ValueError, match="not enabled"): coll.get_or_create_namespace("ns1") def test_delete_namespace_requires_namespace_enabled(self): + """Test delete namespace requires namespace enabled.""" coll = self._make_collection(use_namespace=False) with pytest.raises(ValueError, match="not enabled"): coll.delete_namespace("ns1") def test_list_namespaces_requires_namespace_enabled(self): + """Test list namespaces requires namespace enabled.""" coll = self._make_collection(use_namespace=False) with pytest.raises(ValueError, match="not enabled"): coll.list_namespaces() def test_has_namespace_requires_namespace_enabled(self): + """Test has namespace requires namespace enabled.""" coll = self._make_collection(use_namespace=False) with pytest.raises(ValueError, match="not enabled"): coll.has_namespace("ns1") @@ -264,7 +311,9 @@ def test_has_namespace_requires_namespace_enabled(self): class TestNamespaceObject: + """TestNamespaceObject class.""" def _make_namespace(self) -> Namespace: + """Make namespace.""" mock_client = MagicMock() mock_ef = MagicMock() coll = Collection( @@ -278,22 +327,27 @@ def _make_namespace(self) -> Namespace: return Namespace(client=mock_client, collection=coll, name="ns1", namespace_id="100") def test_name_property(self): + """Test name property.""" ns = self._make_namespace() assert ns.name == "ns1" def test_namespace_id_property(self): + """Test namespace id property.""" ns = self._make_namespace() assert ns.namespace_id == "100" def test_collection_property(self): + """Test collection property.""" ns = self._make_namespace() assert ns.collection.name == "test_coll" def test_embedding_function_inherited(self): + """Test embedding function inherited.""" ns = self._make_namespace() assert ns.embedding_function is ns.collection.embedding_function def test_repr(self): + """Test repr.""" ns = self._make_namespace() r = repr(ns) assert "ns1" in r @@ -301,49 +355,58 @@ def test_repr(self): assert "test_coll" in r def test_add_delegates_to_client(self): + """Test add delegates to client.""" ns = self._make_namespace() ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0]) ns._client._namespace_add.assert_called_once() def test_update_delegates_to_client(self): + """Test update delegates to client.""" ns = self._make_namespace() ns.update(ids="d1", metadatas={"k": "v"}) ns._client._namespace_update.assert_called_once() def test_upsert_delegates_to_client(self): + """Test upsert delegates to client.""" ns = self._make_namespace() ns.upsert(ids="d1", embeddings=[1.0, 2.0, 3.0]) ns._client._namespace_upsert.assert_called_once() def test_delete_delegates_to_client(self): + """Test delete delegates to client.""" ns = self._make_namespace() ns.delete(ids="d1") ns._client._namespace_delete.assert_called_once() def test_query_delegates_to_client(self): + """Test query delegates to client.""" ns = self._make_namespace() ns._client._namespace_query.return_value = {"ids": [["d1"]]} ns.query(query_embeddings=[1.0, 2.0, 3.0]) ns._client._namespace_query.assert_called_once() def test_get_delegates_to_client(self): + """Test get delegates to client.""" ns = self._make_namespace() ns._client._namespace_get.return_value = {"ids": ["d1"]} ns.get(ids="d1") ns._client._namespace_get.assert_called_once() def test_count_delegates_to_client(self): + """Test count delegates to client.""" ns = self._make_namespace() ns._client._namespace_count.return_value = 42 assert ns.count() == 42 def test_peek_delegates_to_client(self): + """Test peek delegates to client.""" ns = self._make_namespace() ns._client._namespace_peek.return_value = {"ids": ["d1"]} ns.peek(limit=5) ns._client._namespace_peek.assert_called_once() def test_prewarm_delegates_to_client(self): + """Test prewarm delegates to client.""" ns = self._make_namespace() ns.prewarm() ns._client._namespace_prewarm.assert_called_once() @@ -356,17 +419,21 @@ class FakeClient(BaseClient): """Concrete BaseClient subclass that captures SQL without executing.""" def __init__(self): + """Init.""" self.database = "test" self.executed_sqls = [] self.query_sqls = [] self.query_return_value = [] def _ensure_connection(self): + """Ensure connection.""" mock_conn = MagicMock() client = self class CaptureCursor: + """CaptureCursor class.""" def execute(self, sql, params=None): + """Execute.""" resolved = sql if params: for p in params: @@ -376,21 +443,26 @@ def execute(self, sql, params=None): client.executed_sqls.append(resolved) def __enter__(self): + """Enter.""" return self def __exit__(self, *args): + """Exit.""" pass def close(self): + """Close.""" pass mock_conn.cursor.return_value = CaptureCursor() return mock_conn def _use_context_manager_for_cursor(self): + """Use context manager for cursor.""" return True def _execute(self, sql, params=None): + """Execute.""" if params: for p in params: sql = sql.replace("%s", repr(p), 1) @@ -402,9 +474,11 @@ def _execute(self, sql, params=None): # Bypass the sdk_ltables lookup in unit tests: SQL-generation tests don't # have a real database, so return a fixed ltable_id matching test assertions. def _resolve_namespace_ltable_id(self, collection_id, namespace_id): + """Resolve namespace ltable id.""" return 1 def _execute_query_with_cursor(self, conn, sql, params, use_context_manager=True): + """Execute query with cursor.""" resolved = sql for p in params: resolved = resolved.replace("%s", repr(p), 1) @@ -414,46 +488,60 @@ def _execute_query_with_cursor(self, conn, sql, params, use_context_manager=True return self.query_return_value def is_connected(self): + """Is connected.""" return True def _cleanup(self): + """Cleanup.""" pass def get_raw_connection(self): + """Get raw connection.""" return None @property def mode(self): + """Mode.""" return "FakeClient" def create_collection(self, *a, **kw): + """Create collection.""" pass def get_collection(self, *a, **kw): + """Get collection.""" pass def delete_collection(self, *a, **kw): + """Delete collection.""" pass def list_collections(self): + """List collections.""" return [] def has_collection(self, name): + """Has collection.""" return False def create_database(self, *a, **kw): + """Create database.""" pass def get_database(self, *a, **kw): + """Get database.""" pass def delete_database(self, *a, **kw): + """Delete database.""" pass def list_databases(self, *a, **kw): + """List databases.""" return [] def fork_database(self, *a, **kw): + """Fork database.""" pass @@ -465,9 +553,11 @@ class TestNamespaceSQLGeneration: TABLE = "42_logic_data_table" def _client(self): + """Client.""" return FakeClient() def _common_kwargs(self): + """Common kwargs.""" return dict( collection_id=self.COLLECTION_ID, collection_name="test_coll", @@ -478,6 +568,7 @@ def _common_kwargs(self): # ---- ADD ---- def test_add_single_sql(self): + """Test add single sql.""" c = self._client() c._namespace_add( **self._common_kwargs(), @@ -495,6 +586,7 @@ def test_add_single_sql(self): assert '\\"tag\\": \\"a\\"' in sql or '"tag": "a"' in sql def test_add_batch_sql(self): + """Test add batch sql.""" c = self._client() c._namespace_add( **self._common_kwargs(), @@ -510,6 +602,7 @@ def test_add_batch_sql(self): assert '\\"id\\": \\"d3\\"' in sql or '"id": "d3"' in sql def test_add_without_metadata_sql(self): + """Test add without metadata sql.""" c = self._client() c._namespace_add( **self._common_kwargs(), @@ -523,6 +616,7 @@ def test_add_without_metadata_sql(self): # ---- UPDATE ---- def test_update_metadata_sql(self): + """Test update metadata sql.""" c = self._client() c._namespace_update( **self._common_kwargs(), @@ -540,6 +634,7 @@ def test_update_metadata_sql(self): assert "'d1'" in sql def test_update_embedding_and_document_sql(self): + """Test update embedding and document sql.""" c = self._client() c._namespace_update( **self._common_kwargs(), @@ -558,6 +653,7 @@ def test_update_embedding_and_document_sql(self): # ---- DELETE ---- def test_delete_by_ids_sql(self): + """Test delete by ids sql.""" c = self._client() c._namespace_delete( **self._common_kwargs(), @@ -571,6 +667,7 @@ def test_delete_by_ids_sql(self): assert "'d1'" in sql def test_delete_by_where_sql(self): + """Test delete by where sql.""" c = self._client() c._namespace_delete( **self._common_kwargs(), @@ -583,6 +680,7 @@ def test_delete_by_where_sql(self): assert "metadata.category" in sql def test_delete_by_where_document_sql(self): + """Test delete by where document sql.""" c = self._client() c._namespace_delete( **self._common_kwargs(), @@ -597,6 +695,7 @@ def test_delete_by_where_document_sql(self): # ---- QUERY ---- def test_query_basic_dsl(self): + """Test query basic dsl.""" c = self._client() c.query_return_value = [] c._namespace_query( @@ -614,6 +713,7 @@ def test_query_basic_dsl(self): assert "APPROXIMATE LIMIT" not in sql def test_query_with_where_dsl(self): + """Test query with where dsl.""" c = self._client() c.query_return_value = [] c._namespace_query( @@ -631,6 +731,7 @@ def test_query_with_where_dsl(self): # ---- GET ---- def test_get_by_ids_sql(self): + """Test get by ids sql.""" c = self._client() c.query_return_value = [] c._namespace_get( @@ -644,6 +745,7 @@ def test_get_by_ids_sql(self): assert """JSON_UNQUOTE(JSON_EXTRACT(data_content, '$.id')) = 'g1'""" in sql def test_get_with_limit_sql(self): + """Test get with limit sql.""" c = self._client() c.query_return_value = [] c._namespace_get( @@ -656,6 +758,7 @@ def test_get_with_limit_sql(self): # ---- COUNT ---- def test_count_sql(self): + """Test count sql.""" c = self._client() c.query_return_value = [{"cnt": 0}] result = c._namespace_count(**self._common_kwargs()) @@ -794,76 +897,92 @@ def test_create_namespace_physical_tables_ivf_without_centroids_fresh_mode(self) class TestValidateNamespaceName: + """TestValidateNamespaceName class.""" def test_valid_simple_name(self): + """Test valid simple name.""" from pyseekdb.client.client_base import _validate_namespace_name _validate_namespace_name("my_namespace") def test_valid_with_digits_and_underscore(self): + """Test valid with digits and underscore.""" from pyseekdb.client.client_base import _validate_namespace_name _validate_namespace_name("tenant_123_abc") def test_valid_single_char(self): + """Test valid single char.""" from pyseekdb.client.client_base import _validate_namespace_name _validate_namespace_name("a") def test_valid_boundary_256_chars(self): + """Test valid boundary 256 chars.""" from pyseekdb.client.client_base import _validate_namespace_name _validate_namespace_name("a" * 256) def test_empty_name_raises(self): + """Test empty name raises.""" from pyseekdb.client.client_base import _validate_namespace_name with pytest.raises(ValueError, match="must not be empty"): _validate_namespace_name("") def test_non_string_raises(self): + """Test non string raises.""" from pyseekdb.client.client_base import _validate_namespace_name with pytest.raises(TypeError, match="must be a string"): _validate_namespace_name(123) def test_too_long_raises(self): + """Test too long raises.""" from pyseekdb.client.client_base import _validate_namespace_name with pytest.raises(ValueError, match="too long"): _validate_namespace_name("a" * 257) def test_hyphen_raises(self): + """Test hyphen raises.""" from pyseekdb.client.client_base import _validate_namespace_name with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my-namespace") def test_space_raises(self): + """Test space raises.""" from pyseekdb.client.client_base import _validate_namespace_name with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my namespace") def test_dot_raises(self): + """Test dot raises.""" from pyseekdb.client.client_base import _validate_namespace_name with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my.namespace") def test_chinese_raises(self): + """Test chinese raises.""" from pyseekdb.client.client_base import _validate_namespace_name with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("命名空间") def test_collection_create_namespace_validates_name(self): + """Test collection create namespace validates name.""" mock_client = MagicMock() coll = Collection(client=mock_client, name="c", collection_id="1", dimension=3, use_namespace=True) with pytest.raises(ValueError, match="invalid characters"): coll.create_namespace("bad-name") def test_collection_get_namespace_validates_name(self): + """Test collection get namespace validates name.""" mock_client = MagicMock() coll = Collection(client=mock_client, name="c", collection_id="1", dimension=3, use_namespace=True) with pytest.raises(ValueError, match="must not be empty"): coll.get_namespace("") def test_collection_has_namespace_validates_name(self): + """Test collection has namespace validates name.""" mock_client = MagicMock() coll = Collection(client=mock_client, name="c", collection_id="1", dimension=3, use_namespace=True) with pytest.raises(TypeError, match="must be a string"): coll.has_namespace(42) def test_partition_count_property(self): + """Test partition count property.""" mock_client = MagicMock() ns_coll = Collection( client=mock_client, name="c", collection_id="1", dimension=3, @@ -880,49 +999,60 @@ def test_partition_count_property(self): class TestValidateRecordIds: + """TestValidateRecordIds class.""" def test_valid_single_id(self): + """Test valid single id.""" from pyseekdb.client.client_base import _validate_record_ids _validate_record_ids(["doc_1"]) def test_valid_multiple_ids(self): + """Test valid multiple ids.""" from pyseekdb.client.client_base import _validate_record_ids _validate_record_ids(["id1", "id2", "id_3"]) def test_valid_boundary_512_chars(self): + """Test valid boundary 512 chars.""" from pyseekdb.client.client_base import _validate_record_ids _validate_record_ids(["a" * 512]) def test_empty_id_raises(self): + """Test empty id raises.""" from pyseekdb.client.client_base import _validate_record_ids with pytest.raises(ValueError, match="must not be empty"): _validate_record_ids([""]) def test_non_string_id_raises(self): + """Test non string id raises.""" from pyseekdb.client.client_base import _validate_record_ids with pytest.raises(TypeError, match="must be a string"): _validate_record_ids([123]) def test_too_long_id_raises(self): + """Test too long id raises.""" from pyseekdb.client.client_base import _validate_record_ids with pytest.raises(ValueError, match="too long"): _validate_record_ids(["a" * 513]) def test_invalid_chars_raises(self): + """Test invalid chars raises.""" from pyseekdb.client.client_base import _validate_record_ids with pytest.raises(ValueError, match="invalid characters"): _validate_record_ids(["doc-1"]) def test_space_in_id_raises(self): + """Test space in id raises.""" from pyseekdb.client.client_base import _validate_record_ids with pytest.raises(ValueError, match="invalid characters"): _validate_record_ids(["doc 1"]) def test_mixed_valid_and_invalid_raises_on_first_bad(self): + """Test mixed valid and invalid raises on first bad.""" from pyseekdb.client.client_base import _validate_record_ids with pytest.raises(ValueError, match="invalid characters"): _validate_record_ids(["good_id", "bad-id"]) def test_non_list_ids_raises(self): + """Test non list ids raises.""" from pyseekdb.client.client_base import _validate_record_ids with pytest.raises(TypeError, match="expected list\\[str\\]"): _validate_record_ids("doc_1") @@ -933,10 +1063,13 @@ def test_non_list_ids_raises(self): class TestNamespaceBatchLimit: + """TestNamespaceBatchLimit class.""" def _client(self): + """Client.""" return FakeClient() def _common_kwargs(self): + """Common kwargs.""" return dict( collection_id="42", collection_name="test_coll", @@ -945,16 +1078,19 @@ def _common_kwargs(self): ) def test_add_single_record_ok(self): + """Test add single record ok.""" c = self._client() c._namespace_add(**self._common_kwargs(), ids="d1", embeddings=[1.0, 2.0, 3.0]) def test_add_100_records_ok(self): + """Test add 100 records ok.""" c = self._client() ids = [f"d{i}" for i in range(100)] embeddings = [[float(i)] * 3 for i in range(100)] c._namespace_add(**self._common_kwargs(), ids=ids, embeddings=embeddings) def test_add_101_records_raises(self): + """Test add 101 records raises.""" c = self._client() ids = [f"d{i}" for i in range(101)] embeddings = [[float(i)] * 3 for i in range(101)] @@ -962,12 +1098,14 @@ def test_add_101_records_raises(self): c._namespace_add(**self._common_kwargs(), ids=ids, embeddings=embeddings) def test_update_101_records_raises(self): + """Test update 101 records raises.""" c = self._client() ids = [f"d{i}" for i in range(101)] with pytest.raises(ValueError, match="exceeds maximum allowed 100"): c._namespace_update(**self._common_kwargs(), ids=ids, metadatas=[{"k": "v"}] * 101) def test_upsert_101_records_raises(self): + """Test upsert 101 records raises.""" c = self._client() ids = [f"d{i}" for i in range(101)] embeddings = [[float(i)] * 3 for i in range(101)] @@ -975,11 +1113,13 @@ def test_upsert_101_records_raises(self): c._namespace_upsert(**self._common_kwargs(), ids=ids, embeddings=embeddings) def test_add_invalid_id_raises(self): + """Test add invalid id raises.""" c = self._client() with pytest.raises(ValueError, match="invalid characters"): c._namespace_add(**self._common_kwargs(), ids="bad-id", embeddings=[1.0, 2.0, 3.0]) def test_delete_invalid_id_raises(self): + """Test delete invalid id raises.""" c = self._client() with pytest.raises(ValueError, match="invalid characters"): c._namespace_delete(**self._common_kwargs(), ids="bad-id") @@ -990,37 +1130,46 @@ def test_delete_invalid_id_raises(self): class TestPhysicalTableNames: + """TestPhysicalTableNames class.""" def test_data_table_name(self): + """Test data table name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.data_table_name("my_coll") == "my_coll_logic_data_table" def test_hot_table_name(self): + """Test hot table name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.hot_table_name("my_coll") == "my_coll_hot_table" def test_kv_data_table_name(self): + """Test kv data table name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.kv_data_table_name("my_coll") == "my_coll_kv_data_table" def test_logic_schema_table_name(self): + """Test logic schema table name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.logic_schema_table_name("my_coll") == "my_coll_logic_schema_table" def test_tablegroup_name(self): + """Test tablegroup name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.tablegroup_name("my_coll") == "my_coll_tg" def test_namespace_catalog_table_names(self): + """Test namespace catalog table names.""" from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.sdk_namespaces_table() == "sdk_namespaces" assert NamespaceCollectionNames.sdk_ltables_table() == "sdk_ltables" assert NamespaceCollectionNames.sdk_namespaces_stats_table() == "sdk_namespaces_stats" def test_is_ns_data_table_true(self): + """Test is ns data table true.""" from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.is_ns_data_table("my_coll_logic_data_table") is True def test_is_ns_data_table_false(self): + """Test is ns data table false.""" from pyseekdb.client.meta_info import NamespaceCollectionNames assert NamespaceCollectionNames.is_ns_data_table("my_coll_hot_table") is False @@ -1030,7 +1179,9 @@ def test_is_ns_data_table_false(self): class TestNamespaceCatalogs: + """TestNamespaceCatalogs class.""" def test_ensure_namespace_catalogs_creates_all_catalog_tables(self): + """Test ensure namespace catalogs creates all catalog tables.""" c = FakeClient() c._ensure_namespace_catalogs() @@ -1052,6 +1203,7 @@ def test_ensure_namespace_catalogs_creates_all_catalog_tables(self): assert "PARTITION BY KEY(namespace_id) PARTITIONS 8" in sql def test_ensure_namespace_catalogs_creates_catalog_tables_in_order(self): + """Test ensure namespace catalogs creates catalog tables in order.""" c = FakeClient() c._ensure_namespace_catalogs() @@ -1064,6 +1216,7 @@ def test_ensure_namespace_catalogs_creates_catalog_tables_in_order(self): assert "CREATE UNIQUE INDEX uk_sdk_lt_coll_ns_name" in c.executed_sqls[6] def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): + """Test delete ns collection meta cleans namespaces stats table.""" c = FakeClient() c._get_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) c._cleanup_namespace_physical_tables = MagicMock() @@ -1078,7 +1231,9 @@ def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): class TestBrokenNsCollectionPurge: + """TestBrokenNsCollectionPurge class.""" def test_purge_broken_ns_collection_if_incomplete_calls_delete(self): + """Test purge broken ns collection if incomplete calls delete.""" c = FakeClient() meta = { "collection_id": "cid1", @@ -1093,6 +1248,7 @@ def test_purge_broken_ns_collection_if_incomplete_calls_delete(self): c._delete_ns_collection_meta.assert_called_once_with("coll") def test_purge_skips_complete_collection(self): + """Test purge skips complete collection.""" c = FakeClient() meta = {"collection_id": "cid1", "collection_name": "coll", "settings": {"storage_mode": "sn"}} c._use_catalog_database = MagicMock() @@ -1103,7 +1259,9 @@ def test_purge_skips_complete_collection(self): c._delete_ns_collection_meta.assert_not_called() def test_get_collection_purges_incomplete_namespace_collection(self): + """Test get collection purges incomplete namespace collection.""" class GetClient(FakeClient): + """GetClient class.""" get_collection = BaseClient.get_collection c = GetClient() @@ -1126,7 +1284,9 @@ class GetClient(FakeClient): class TestValidateNResults: + """TestValidateNResults class.""" def test_rejects_boolean_values(self): + """Test rejects boolean values.""" with pytest.raises(ValueError, match="n_results must be an integer"): _validate_n_results(True) with pytest.raises(ValueError, match="n_results must be an integer"): @@ -1138,7 +1298,9 @@ def test_rejects_boolean_values(self): class TestUseNamespaceValidation: + """TestUseNamespaceValidation class.""" def test_hnsw_raises(self): + """Test hnsw raises.""" c = FakeClient() from pyseekdb.client.schema import Schema from pyseekdb.client.configuration import VectorIndexConfig, HNSWConfiguration @@ -1148,6 +1310,7 @@ def test_hnsw_raises(self): c._create_namespace_collection("test", schema) def test_ob_type_validation(self): + """Test ob type validation.""" c = FakeClient() c.detect_db_type_and_version = MagicMock(return_value=("mysql", "8.0")) from pyseekdb.client.schema import Schema @@ -1157,6 +1320,7 @@ def test_ob_type_validation(self): c._create_namespace_collection("test", schema) def test_create_namespace_collection_without_ivf_skips_vector_index(self): + """Test create namespace collection without ivf skips vector index.""" from pyseekdb.client.version import Version c = FakeClient() @@ -1179,6 +1343,7 @@ def test_create_namespace_collection_without_ivf_skips_vector_index(self): assert "centroids_fresh_mode" not in settings def test_create_namespace_collection_ivf_without_centroids_fresh_mode(self): + """Test create namespace collection ivf without centroids fresh mode.""" from pyseekdb.client.version import Version c = FakeClient() @@ -1211,7 +1376,9 @@ def test_create_namespace_collection_ivf_without_centroids_fresh_mode(self): class TestDeleteNamespaceUsesKernel: + """TestDeleteNamespaceUsesKernel class.""" def test_delete_namespace_calls_dbms_logic_table(self): + """Test delete namespace calls dbms logic table.""" c = FakeClient() c._execute = MagicMock(side_effect=[ [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], diff --git a/tests/unit_tests/test_namespace_lifecycle_lock.py b/tests/unit_tests/test_namespace_lifecycle_lock.py index 929a47dd..38ea7dcf 100644 --- a/tests/unit_tests/test_namespace_lifecycle_lock.py +++ b/tests/unit_tests/test_namespace_lifecycle_lock.py @@ -14,7 +14,9 @@ class TestNamespaceLifecycleWithoutGetLock: + """TestNamespaceLifecycleWithoutGetLock class.""" def test_has_namespace_queries_catalog_directly(self): + """Test has namespace queries catalog directly.""" client = MagicMock(spec=BaseClient) client._get_ns_namespace_meta.return_value = None @@ -22,6 +24,7 @@ def test_has_namespace_queries_catalog_directly(self): client._get_ns_namespace_meta.assert_called_once_with("cid", "ns1") def test_create_namespace_raises_when_catalog_row_exists(self): + """Test create namespace raises when catalog row exists.""" client = MagicMock(spec=BaseClient) client._get_ns_namespace_meta.return_value = { "namespace_id": "42", diff --git a/tests/unit_tests/test_namespace_upsert_reconcile.py b/tests/unit_tests/test_namespace_upsert_reconcile.py index b03e2b99..72552139 100644 --- a/tests/unit_tests/test_namespace_upsert_reconcile.py +++ b/tests/unit_tests/test_namespace_upsert_reconcile.py @@ -16,7 +16,9 @@ class TestNamespaceUpsertReconcile: + """TestNamespaceUpsertReconcile class.""" def test_reconcile_skips_when_single_row_exists(self): + """Test reconcile skips when single row exists.""" client = MagicMock(spec=BaseClient) client._count_namespace_records_by_id.return_value = 1 @@ -39,6 +41,7 @@ def test_reconcile_skips_when_single_row_exists(self): client._namespace_add.assert_not_called() def test_reconcile_collapses_duplicate_rows(self): + """Test reconcile collapses duplicate rows.""" client = MagicMock(spec=BaseClient) client._count_namespace_records_by_id.side_effect = [4, 0, 1] @@ -68,6 +71,7 @@ def test_reconcile_collapses_duplicate_rows(self): assert add_kwargs["embeddings"] == [[1.0, 2.0, 3.0]] def test_reconcile_raises_when_retries_exhausted(self): + """Test reconcile raises when retries exhausted.""" client = MagicMock(spec=BaseClient) client._count_namespace_records_by_id.return_value = 4 From a9d0f59fc7f226cbba771ac9cc86c455bb7c813c Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 25 Jun 2026 20:49:05 +0800 Subject: [PATCH 65/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91=E9=9D=9E?= =?UTF-8?q?=E6=B3=95=20include=20=E5=AD=97=E6=AE=B5=E6=B2=A1=E6=9C=89?= =?UTF-8?q?=E6=8A=A5=E9=94=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/collection.py | 11 ++++++- src/pyseekdb/client/namespace.py | 5 +++- src/pyseekdb/client/validators.py | 30 +++++++++++++++++++ .../integration_tests/test_namespace_query.py | 15 ++++++++++ tests/unit_tests/test_namespace.py | 26 +++++++++++++++- 5 files changed, 84 insertions(+), 3 deletions(-) diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 80063f83..517fc007 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -10,7 +10,12 @@ from typing import TYPE_CHECKING, Any, Optional -from .validators import _MAX_N_RESULTS, _validate_namespace_name, _validate_n_results +from .validators import ( + _MAX_N_RESULTS, + _validate_include, + _validate_namespace_name, + _validate_n_results, +) if TYPE_CHECKING: from .embedding_function import Documents as EmbeddingDocuments @@ -496,6 +501,7 @@ def query( """ self._guard_collection_data_api() _validate_n_results(n_results) + _validate_include(include) return self._client._collection_query( collection_id=self._id, collection_name=self._name, @@ -578,6 +584,7 @@ def get( ) """ self._guard_collection_data_api() + _validate_include(include) return self._client._collection_get( collection_id=self._id, collection_name=self._name, @@ -666,6 +673,8 @@ def hybrid_search( ) """ self._guard_collection_data_api() + _validate_n_results(n_results) + _validate_include(include) # When no query/knn provided, return only ids/distances by default if include is None and not query and not knn: include = [] diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py index d4739dde..f20f6906 100644 --- a/src/pyseekdb/client/namespace.py +++ b/src/pyseekdb/client/namespace.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any -from .validators import _validate_n_results +from .validators import _validate_include, _validate_n_results if TYPE_CHECKING: from .collection import Collection @@ -173,6 +173,7 @@ def query( """Run vector similarity search within this namespace.""" self._guard_exists() _validate_n_results(n_results) + _validate_include(include) return self._client._namespace_query( collection_id=self._collection.id, collection_name=self._collection.name, @@ -201,6 +202,7 @@ def hybrid_search( """Run hybrid fulltext + vector search within this namespace.""" self._guard_exists() _validate_n_results(n_results) + _validate_include(include) if include is None and not query and not knn: include = [] return self._client._namespace_hybrid_search( @@ -230,6 +232,7 @@ def get( ) -> dict[str, Any]: """Fetch records from this namespace by ids or filters.""" self._guard_exists() + _validate_include(include) return self._client._namespace_get( collection_id=self._collection.id, collection_name=self._collection.name, diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py index 78434711..f9b4c1a8 100644 --- a/src/pyseekdb/client/validators.py +++ b/src/pyseekdb/client/validators.py @@ -7,6 +7,36 @@ _MAX_NAMESPACE_NAME_LENGTH = 256 _MAX_NAMESPACE_BATCH_SIZE = 100 _MAX_N_RESULTS = 16384 # OceanBase vector-search k upper bound +_VALID_INCLUDE_FIELDS = frozenset( + { + "documents", + "document", + "metadatas", + "metadata", + "embeddings", + "embedding", + "distances", + "distance", + } +) + + +def _validate_include(include: list[str] | None) -> None: + """Validate ``include`` is a list of supported result field names.""" + if include is None: + return + if not isinstance(include, list): + raise TypeError( + f"include must be a list[str] or None, got {type(include).__name__}" + ) + invalid = [field for field in include if field not in _VALID_INCLUDE_FIELDS] + if invalid: + allowed = "documents, metadatas, embeddings, distances" + raise ValueError( + f"Invalid include field(s): {invalid!r}. " + f"Allowed values: {allowed} " + f"(singular forms document, metadata, embedding, distance are also accepted)." + ) def _validate_namespace_name(name: str) -> None: diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index 2f98d6dd..00832584 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -143,6 +143,21 @@ def test_query_include_embeddings(self, db_client): finally: db_client.delete_collection(name=collection.name) + def test_query_rejects_invalid_include(self, db_client): + """Invalid include fields must fail before executing the query.""" + collection = self._setup(db_client) + ns = collection.create_namespace("qns_bad_inc") + try: + self._insert_data(ns) + with pytest.raises(ValueError, match="Invalid include field"): + ns.query( + query_embeddings=[1.0, 0.0, 0.0], + n_results=3, + include=["invalid_field_name"], + ) + finally: + db_client.delete_collection(name=collection.name) + def test_multi_namespace_isolation(self, db_client): """Test multi namespace isolation.""" collection = self._setup(db_client, suffix="_iso") diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 01d32193..31a05010 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -21,7 +21,7 @@ from pyseekdb.client.collection import Collection # noqa: E402 from pyseekdb.client.configuration import HNSWConfiguration, IVFIndexType, VectorIndexConfig # noqa: E402 from pyseekdb.client.namespace import Namespace # noqa: E402 -from pyseekdb.client.validators import _validate_n_results # noqa: E402 +from pyseekdb.client.validators import _validate_include, _validate_n_results # noqa: E402 # ==================== IVFConfiguration Tests ==================== @@ -1293,6 +1293,30 @@ def test_rejects_boolean_values(self): _validate_n_results(False) +class TestValidateInclude: + + """TestValidateInclude class.""" + + def test_accepts_none_and_valid_fields(self): + """Test accepts none and valid fields.""" + _validate_include(None) + _validate_include([]) + _validate_include(["documents", "metadatas", "embeddings", "distances"]) + _validate_include(["document", "metadata", "embedding", "distance"]) + + def test_rejects_invalid_field_names(self): + """Test rejects invalid field names.""" + with pytest.raises(ValueError, match="Invalid include field"): + _validate_include(["invalid_field_name"]) + with pytest.raises(ValueError, match="Invalid include field"): + _validate_include(["documents", "bad_field"]) + + def test_rejects_non_list(self): + """Test rejects non list.""" + with pytest.raises(TypeError, match="include must be a list"): + _validate_include("documents") # type: ignore[arg-type] + + # ==================== UseNamespace Validation Tests ==================== From 081d24fcb3a541774a802a4e87b9a2c6cddb0b2f Mon Sep 17 00:00:00 2001 From: lc533134 Date: Thu, 25 Jun 2026 21:06:19 +0800 Subject: [PATCH 66/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91sdk?= =?UTF-8?q?=E6=B2=A1=E6=9C=89=E5=AF=B9=E5=86=85=E6=A0=B8=E6=8A=A5=E9=94=99?= =?UTF-8?q?=E8=BF=9B=E8=A1=8C=E5=8C=85=E8=A3=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 70 ++++--- src/pyseekdb/client/client_seekdb_server.py | 2 + src/pyseekdb/client/kernel_errors.py | 183 ++++++++++++++++++ src/pyseekdb/client/namespace.py | 5 + tests/integration_tests/test_namespace_dml.py | 6 +- .../test_namespace_lifecycle.py | 21 ++ .../test_namespace_kernel_errors.py | 103 ++++++++++ 7 files changed, 361 insertions(+), 29 deletions(-) create mode 100644 src/pyseekdb/client/kernel_errors.py create mode 100644 tests/unit_tests/test_namespace_kernel_errors.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index fedf20e7..fccc3165 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -17,6 +17,7 @@ from pymysql.converters import escape_string +from .kernel_errors import maybe_reraise_friendly_kernel_error, namespace_kernel_error_guard from .admin_client import DEFAULT_TENANT, AdminAPI from .base_connection import BaseConnection from .collection import Collection @@ -1830,10 +1831,9 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> except Exception as exc: if _is_namespace_catalog_conflict_error(exc): self._rollback_connection_if_supported() - if self._get_ns_namespace_meta(collection_id, namespace_name) is not None: - raise ValueError( - f"Namespace '{namespace_name}' already exists" - ) from exc + raise ValueError( + f"Namespace '{namespace_name}' already exists" + ) from exc raise return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) @@ -3540,18 +3540,18 @@ def _execute_query_with_cursor( """ if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): print(f"[pyseekdb SQL] {sql} -- params={params}", flush=True) - if use_context_manager: - with conn.cursor() as cursor: - cursor.execute(sql, params) - if not self._should_fetch_results(cursor, sql): - return [] - rows = cursor.fetchall() - # Normalize rows - normalized_rows = [] - for row in rows: - normalized_rows.append(self._normalize_row(row, cursor.description)) - return normalized_rows - else: + try: + if use_context_manager: + with conn.cursor() as cursor: + cursor.execute(sql, params) + if not self._should_fetch_results(cursor, sql): + return [] + rows = cursor.fetchall() + # Normalize rows + normalized_rows = [] + for row in rows: + normalized_rows.append(self._normalize_row(row, cursor.description)) + return normalized_rows cursor = conn.cursor() try: cursor.execute(sql, params) @@ -3565,6 +3565,9 @@ def _execute_query_with_cursor( return normalized_rows finally: cursor.close() + except Exception as exc: + maybe_reraise_friendly_kernel_error(exc) + raise def _build_select_clause(self, include_fields: dict[str, bool]) -> str: """ @@ -3800,21 +3803,25 @@ def _execute(self, sql: str) -> Any: conn = self._ensure_connection() use_context_manager = self._use_context_manager_for_cursor() - if use_context_manager: - with conn.cursor() as cursor: + try: + if use_context_manager: + with conn.cursor() as cursor: + cursor.execute(sql) + if self._should_fetch_results(cursor, sql): + return cursor.fetchall() + return None + + cursor = conn.cursor() + try: cursor.execute(sql) if self._should_fetch_results(cursor, sql): return cursor.fetchall() return None - - cursor = conn.cursor() - try: - cursor.execute(sql) - if self._should_fetch_results(cursor, sql): - return cursor.fetchall() - return None - finally: - cursor.close() + finally: + cursor.close() + except Exception as exc: + maybe_reraise_friendly_kernel_error(exc) + raise # -------------------- DQL Operations (Common Implementation) -------------------- @@ -5348,6 +5355,7 @@ def _reconcile_namespace_duplicate_records( f"Failed to reconcile duplicate namespace rows for record_id={record_id!r}" ) + @namespace_kernel_error_guard def _namespace_add( # noqa: C901 self, collection_id: str | None, @@ -5443,6 +5451,7 @@ def _namespace_add( # noqa: C901 sql = f"INSERT INTO `{table_name}` ({columns}) VALUES {','.join(values_list)}" self._execute(sql) + @namespace_kernel_error_guard def _namespace_update( self, collection_id: str | None, @@ -5589,6 +5598,7 @@ def _namespace_update( finally: cursor.close() + @namespace_kernel_error_guard def _namespace_upsert( self, collection_id: str | None, @@ -5693,6 +5703,7 @@ def _namespace_upsert( **kwargs, ) + @namespace_kernel_error_guard def _namespace_delete( self, collection_id: str | None, @@ -5757,6 +5768,7 @@ def _namespace_delete( else: self._execute(sql) + @namespace_kernel_error_guard def _namespace_query( # noqa: C901 self, collection_id: str | None, @@ -5857,6 +5869,7 @@ def _namespace_query( # noqa: C901 result["embeddings"] = all_embeddings return result + @namespace_kernel_error_guard def _namespace_get( # noqa: C901 self, collection_id: str | None, @@ -5981,6 +5994,7 @@ def _namespace_get( # noqa: C901 result["embeddings"] = result_embeddings return result + @namespace_kernel_error_guard def _namespace_count( self, collection_id: str | None, @@ -6010,6 +6024,7 @@ def _namespace_count( return row[0] if len(row) > 0 else 0 return int(row) if row else 0 + @namespace_kernel_error_guard def _namespace_peek( self, collection_id: str | None, @@ -6130,6 +6145,7 @@ def _inject_filter_into_query(node): return search_parm + @namespace_kernel_error_guard def _namespace_hybrid_search( self, collection_id: str | None, diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 5fb19205..90d4d033 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -13,6 +13,7 @@ from .admin_client import DEFAULT_TENANT from .client_base import BaseClient from .database import Database +from .kernel_errors import namespace_kernel_error_guard logger = logging.getLogger(__name__) @@ -207,6 +208,7 @@ def _database_tenant(self, tenant: str) -> str | None: ) return self.tenant + @namespace_kernel_error_guard def _namespace_prewarm( self, collection_id: str | None, diff --git a/src/pyseekdb/client/kernel_errors.py b/src/pyseekdb/client/kernel_errors.py new file mode 100644 index 00000000..78c686df --- /dev/null +++ b/src/pyseekdb/client/kernel_errors.py @@ -0,0 +1,183 @@ +"""Translate OceanBase kernel errors into SDK-friendly ValueError messages.""" + +from __future__ import annotations + +import contextlib +import contextvars +import functools +import re +from typing import Any, Callable, Iterator, TypeVar + +_NS_KERNEL_ERROR_CTX: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( + "ns_kernel_error_ctx", default=None +) + +_NS_PHYSICAL_TABLE_SUFFIXES = ( + "_logic_data_table", + "_kv_data_table", + "_logic_schema_table", + "_hot_table", +) + + +def _iter_exception_chain(exc: BaseException) -> Iterator[BaseException]: + """Yield the exception and its cause/context chain.""" + current: BaseException | None = exc + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + yield current + cause = current.__cause__ + context = current.__context__ + current = cause if cause is not None else context + + +def _exception_text(exc: BaseException) -> str: + """Flatten the exception chain into one searchable string.""" + return " ".join(str(item) for item in _iter_exception_chain(exc)) + + +def _has_ob_error_code(exc: BaseException, *codes: int) -> bool: + """Return whether any chained exception carries one of the OceanBase error codes.""" + text = _exception_text(exc) + lowered = text.lower() + for code in codes: + code_str = str(code) + if ( + f"code={code_str}" in lowered + or f"code=-{code_str}" in lowered + or f"errcode=-{code_str}" in lowered + or re.search(rf"\({code_str},", text) is not None + or re.search(rf"(?:^|[\s,])(?:-{code_str}|{code_str})(?:[\s,)]|$)", text) is not None + ): + return True + return False + + +def _is_namespace_dropping_error(exc: BaseException) -> bool: + """Whether the kernel reports a namespace is still being dropped.""" + text = _exception_text(exc).lower() + return ( + _has_ob_error_code(exc, 4109) + or "namespace is dropping" in text + or "namespace is in drop" in text + or "being dropped" in text and "namespace" in text + ) + + +def _is_namespace_missing_kernel_error(exc: BaseException) -> bool: + """Whether the kernel reports the namespace no longer exists.""" + text = _exception_text(exc).lower() + return ( + _has_ob_error_code(exc, 4018) + or "namespace not exist" in text + or "namespace does not exist" in text + or "namespace doesn't exist" in text + or "no such namespace" in text + ) + + +def _is_missing_collection_physical_error(exc: BaseException, collection_id: str | None = None) -> bool: + """Whether the error indicates namespace collection physical tables are gone.""" + text = _exception_text(exc).lower() + if collection_id: + cid = collection_id.lower() + for suffix in _NS_PHYSICAL_TABLE_SUFFIXES: + if f"{cid}{suffix}" in text: + return True + if any(suffix in text for suffix in _NS_PHYSICAL_TABLE_SUFFIXES): + return True + if "table doesn't exist" in text or "table does not exist" in text or "table not exist" in text: + return True + return _has_ob_error_code(exc, 1146, 5019, 942) + + +def _friendly_kernel_error_message( + exc: BaseException, + *, + namespace_name: str | None, + collection_name: str | None, + collection_id: str | None, +) -> str | None: + """Map a kernel exception to a user-facing message, or None to keep the original.""" + if _is_namespace_dropping_error(exc): + ns_label = f"'{namespace_name}'" if namespace_name else "The namespace" + return ( + f"Namespace {ns_label} is being dropped and is not available. " + "Wait for the drop to finish or use a different namespace." + ) + if _is_namespace_missing_kernel_error(exc): + ns_label = f"'{namespace_name}'" if namespace_name else "The namespace" + return ( + f"Namespace {ns_label} no longer exists (it may have been deleted). " + "Operations are not allowed on a deleted namespace." + ) + if _is_missing_collection_physical_error(exc, collection_id): + coll_label = f"'{collection_name}'" if collection_name else "The collection" + return ( + f"Collection {coll_label} does not exist (it may have been deleted). " + "Namespace operations are not allowed on a deleted collection." + ) + return None + + +def maybe_reraise_friendly_kernel_error(exc: BaseException) -> None: + """Re-raise a wrapped ValueError when a namespace-scoped kernel error is recognized.""" + ctx = _NS_KERNEL_ERROR_CTX.get() + if ctx is None: + return + friendly = _friendly_kernel_error_message( + exc, + namespace_name=ctx.get("namespace_name"), + collection_name=ctx.get("collection_name"), + collection_id=ctx.get("collection_id"), + ) + if friendly is None: + return + raise ValueError(friendly) from exc + + +@contextlib.contextmanager +def namespace_kernel_error_scope( + *, + namespace_name: str | None = None, + collection_name: str | None = None, + collection_id: str | None = None, +): + """Attach namespace/collection context used to translate kernel SQL errors.""" + token = _NS_KERNEL_ERROR_CTX.set( + { + "namespace_name": namespace_name or "", + "collection_name": collection_name or "", + "collection_id": collection_id or "", + } + ) + try: + yield + finally: + _NS_KERNEL_ERROR_CTX.reset(token) + + +_F = TypeVar("_F", bound=Callable[..., Any]) + + +def namespace_kernel_error_guard(method: _F) -> _F: + """Decorator for ``BaseClient._namespace_*`` methods with standard leading args.""" + @functools.wraps(method) + def wrapper( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + namespace_name: str, + *args: Any, + **kwargs: Any, + ): + with namespace_kernel_error_scope( + namespace_name=namespace_name, + collection_name=collection_name, + collection_id=str(collection_id) if collection_id is not None else "", + ): + return method(self, collection_id, collection_name, namespace_id, namespace_name, *args, **kwargs) + + return wrapper # type: ignore[return-value] diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py index f20f6906..4213f191 100644 --- a/src/pyseekdb/client/namespace.py +++ b/src/pyseekdb/client/namespace.py @@ -61,6 +61,11 @@ def __repr__(self) -> str: def _guard_exists(self) -> None: """Raise if this namespace or its collection was deleted.""" + if not self._client._ns_collection_exists_by_id(self._collection.id): + raise ValueError( + f"Collection '{self._collection.name}' no longer exists (it may have been deleted). " + "Namespace operations are not allowed on a deleted collection." + ) if not self._client._ns_namespace_exists_by_id(self._collection.id, self._namespace_id): raise ValueError( f"Namespace '{self._name}' no longer exists (it or its collection may have been deleted). " diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py index cb157dd8..f76f56e8 100644 --- a/tests/integration_tests/test_namespace_dml.py +++ b/tests/integration_tests/test_namespace_dml.py @@ -60,6 +60,8 @@ def test_ops_blocked_after_namespace_deleted(self, db_client): ns.query(query_embeddings=[1.0, 2.0, 3.0], n_results=1) with pytest.raises(ValueError, match="no longer exists"): ns.count() + with pytest.raises(ValueError, match="no longer exists"): + ns.prewarm() finally: cleanup(db_client, collection) @@ -69,9 +71,9 @@ def test_ops_blocked_after_collection_deleted(self, db_client): ns = _create_namespace(collection, "dml_ns") ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0]) db_client.delete_collection(name=collection.name) - with pytest.raises(ValueError, match="no longer exists"): + with pytest.raises(ValueError, match="no longer exists|does not exist"): ns.add(ids="d2", embeddings=[4.0, 5.0, 6.0]) - with pytest.raises(ValueError, match="no longer exists"): + with pytest.raises(ValueError, match="no longer exists|does not exist"): ns.query(query_embeddings=[1.0, 2.0, 3.0], n_results=1) def test_add_batch(self, db_client): diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 7053a485..2d9bc391 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -149,6 +149,27 @@ def test_get_nonexistent_namespace_raises(self, db_client): finally: db_client.delete_collection(name=collection.name) + def test_create_duplicate_namespace_raises_friendly_error(self, db_client): + """Duplicate namespace creation should surface a clear SDK error.""" + collection = self._create_ns_collection(db_client) + try: + collection.create_namespace("dup_ns") + with pytest.raises(ValueError, match="already exists"): + collection.create_namespace("dup_ns") + finally: + db_client.delete_collection(name=collection.name) + + def test_prewarm_after_namespace_deleted_raises_friendly_error(self, db_client): + """Prewarm on a deleted namespace should not leak raw kernel error codes.""" + collection = self._create_ns_collection(db_client) + ns = collection.create_namespace("prewarm_del") + try: + collection.delete_namespace("prewarm_del") + with pytest.raises(ValueError, match="no longer exists|being dropped"): + ns.prewarm() + finally: + db_client.delete_collection(name=collection.name) + def test_delete_collection_cleans_namespaces(self, db_client): """Test delete collection cleans namespaces.""" collection = self._create_ns_collection(db_client) diff --git a/tests/unit_tests/test_namespace_kernel_errors.py b/tests/unit_tests/test_namespace_kernel_errors.py new file mode 100644 index 00000000..7d7a7bb4 --- /dev/null +++ b/tests/unit_tests/test_namespace_kernel_errors.py @@ -0,0 +1,103 @@ +""" +Unit tests for OceanBase kernel error translation in namespace operations. +""" + +import sys +from pathlib import Path + +import pytest + +project_root = Path(__file__).parent.parent.parent +src_root = project_root / "src" +sys.path.insert(0, str(src_root)) + +from pyseekdb.client.kernel_errors import ( # noqa: E402 + _friendly_kernel_error_message, + _is_missing_collection_physical_error, + _is_namespace_dropping_error, + _is_namespace_missing_kernel_error, + maybe_reraise_friendly_kernel_error, + namespace_kernel_error_scope, +) + + +class IntegrityError(Exception): + """Stand-in for pymysql IntegrityError in unit tests.""" + + +class TestNamespaceKernelErrorDetection: + """TestNamespaceKernelErrorDetection class.""" + + def test_detects_namespace_dropping_4109(self): + """Test detects namespace dropping 4109.""" + exc = Exception("DBMS_LOGIC_TABLE.PREWARM failed: errcode=-4109") + assert _is_namespace_dropping_error(exc) is True + + def test_detects_namespace_missing_4018(self): + """Test detects namespace missing 4018.""" + exc = Exception("namespace not exist, code=4018") + assert _is_namespace_missing_kernel_error(exc) is True + + def test_detects_missing_logic_data_table(self): + """Test detects missing logic data table.""" + exc = Exception("Table 'abc123_logic_data_table' doesn't exist") + assert _is_missing_collection_physical_error(exc, "abc123") is True + + def test_maps_duplicate_namespace_to_friendly_message(self): + """Duplicate sdk_namespaces insert is handled in create path, not here.""" + exc = IntegrityError( + '(1062, "Duplicate entry \'cid-demo_ns\' for key \'uk_sdk_ns_coll_name\'")' + ) + assert _friendly_kernel_error_message( + exc, + namespace_name="demo_ns", + collection_name="coll", + collection_id="cid", + ) is None + + +class TestNamespaceKernelErrorTranslation: + """TestNamespaceKernelErrorTranslation class.""" + + def test_wraps_namespace_dropping_error(self): + """Test wraps namespace dropping error.""" + exc = Exception("prewarm failed, errcode=-4109") + with namespace_kernel_error_scope( + namespace_name="demo_ns", + collection_name="coll", + collection_id="cid", + ): + with pytest.raises(ValueError, match="being dropped"): + maybe_reraise_friendly_kernel_error(exc) + + def test_wraps_namespace_missing_error(self): + """Test wraps namespace missing error.""" + exc = Exception("namespace does not exist, code=4018") + with namespace_kernel_error_scope( + namespace_name="demo_ns", + collection_name="coll", + collection_id="cid", + ): + with pytest.raises(ValueError, match="no longer exists"): + maybe_reraise_friendly_kernel_error(exc) + + def test_wraps_missing_collection_physical_error(self): + """Test wraps missing collection physical error.""" + exc = Exception("Table 'abc123_logic_data_table' doesn't exist") + with namespace_kernel_error_scope( + namespace_name="demo_ns", + collection_name="coll", + collection_id="abc123", + ): + with pytest.raises(ValueError, match="Collection 'coll' does not exist"): + maybe_reraise_friendly_kernel_error(exc) + + def test_leaves_unrelated_errors_untouched(self): + """Test leaves unrelated errors untouched.""" + exc = ValueError("syntax error") + with namespace_kernel_error_scope( + namespace_name="demo_ns", + collection_name="coll", + collection_id="cid", + ): + maybe_reraise_friendly_kernel_error(exc) From 908186ad71f068891ea3d52aaeb1d2035aa83546 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 26 Jun 2026 10:34:19 +0800 Subject: [PATCH 67/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91ns.query?= =?UTF-8?q?=E6=94=AF=E6=8C=81=E7=9A=84where=5Fdocument=E6=93=8D=E4=BD=9C?= =?UTF-8?q?=E7=AC=A6=E4=B8=8D=E5=AE=8C=E5=85=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 162 +++++----- src/pyseekdb/client/document_query_builder.py | 275 +++++++++++++++++ src/pyseekdb/client/filters.py | 8 +- .../test_namespace_query_where_document.py | 284 ++++++++++++++++++ .../unit_tests/test_document_query_builder.py | 61 ++++ 5 files changed, 716 insertions(+), 74 deletions(-) create mode 100644 src/pyseekdb/client/document_query_builder.py create mode 100644 tests/integration_tests/test_namespace_query_where_document.py create mode 100644 tests/unit_tests/test_document_query_builder.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index fccc3165..d22b57d7 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -17,6 +17,13 @@ from pymysql.converters import escape_string +from .document_query_builder import ( + build_document_hybrid_expression, + doc_matches_where_document, + document_expr_as_knn_filter, + merge_into_knn_filter, + where_document_knn_prefilterable, +) from .kernel_errors import maybe_reraise_friendly_kernel_error, namespace_kernel_error_guard from .admin_client import DEFAULT_TENANT, AdminAPI from .base_connection import BaseConnection @@ -4697,72 +4704,9 @@ def _apply_boost(target: Any) -> None: _apply_boost(expr) return expr - # Handle $contains - use query_string - if "$contains" in where_document: - # Use pymysql's escape_string for safe escaping of query content - query_content = where_document["$contains"] - escaped_query = escape_string(query_content) - return _with_boost({"query_string": {"fields": ["document"], "query": escaped_query}}) - - # Handle $not_contains - wrap query_string in must_not bool - if "$not_contains" in where_document: - escaped_query = escape_string(where_document["$not_contains"]) - return _with_boost({ - "bool": { - "must_not": [ - { - "query_string": { - "fields": ["document"], - "query": escaped_query, - } - } - ] - } - }) - - # Handle $and with $contains - if "$and" in where_document: - and_conditions = where_document["$and"] - contains_queries = [] - for condition in and_conditions: - if isinstance(condition, dict) and "$contains" in condition: - contains_queries.append(condition["$contains"]) - - if contains_queries: - # Combine multiple $contains with AND (escape each query) - escaped_queries = [escape_string(q) for q in contains_queries] - return _with_boost({ - "query_string": { - "fields": ["document"], - "query": " ".join(escaped_queries), - "default_operator": "and", - } - }) - - # Handle $or with $contains - if "$or" in where_document: - or_conditions = where_document["$or"] - contains_queries = [] - for condition in or_conditions: - if isinstance(condition, dict) and "$contains" in condition: - contains_queries.append(condition["$contains"]) - - if contains_queries: - # Combine multiple $contains with OR (escape each query) - escaped_queries = [escape_string(q) for q in contains_queries] - return _with_boost({ - "query_string": { - "fields": ["document"], - "query": " ".join(escaped_queries), - "default_operator": "or", - } - }) - - # Default: if it's a string, treat as $contains - if isinstance(where_document, str): - return _with_boost({"query_string": {"fields": ["document"], "query": where_document}}) - - return None + return _with_boost( + build_document_hybrid_expression(where_document, boost=boost) + ) def _build_metadata_filter_for_search_parm(self, where: dict[str, Any] | None) -> list[dict[str, Any]]: """ @@ -4924,6 +4868,7 @@ def _build_knn_expression( # noqa: C901 query_texts = knn.get("query_texts") query_embeddings = knn.get("query_embeddings") where = knn.get("where") + where_document = knn.get("where_document") n_results = knn.get("n_results", 10) if not isinstance(n_results, int) or n_results < 1: raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") @@ -4995,6 +4940,17 @@ def _normalize_vectors(raw_embeddings: Any) -> list[list[float]]: knn_filter = [{"bool": bool_filter}] elif pos_filters: knn_filter = pos_filters + + if where_document is not None and where_document_knn_prefilterable(where_document): + doc_filter = document_expr_as_knn_filter( + build_document_hybrid_expression(where_document) + ) + if doc_filter is not None: + if knn_filter is None: + knn_filter = [doc_filter] + else: + knn_filter = [*knn_filter, doc_filter] + for vector in vectors: expr = {"field": "embedding", "k": n_results, "query_vector": vector} if boost is not None: @@ -5007,6 +4963,54 @@ def _normalize_vectors(raw_embeddings: Any) -> list[list[float]]: return knn_exprs if len(knn_exprs) > 1 else knn_exprs[0] + def _post_filter_namespace_query_result( + self, + result: dict[str, Any], + where_document: dict[str, Any] | str, + *, + n_results: int, + ) -> dict[str, Any]: + """Drop hybrid_search rows that violate a where_document predicate.""" + ids_groups = result.get("ids") or [] + if not ids_groups: + return result + + filtered: dict[str, Any] = {"ids": []} + if result.get("distances") is not None: + filtered["distances"] = [] + for optional_key in ("documents", "metadatas", "embeddings"): + if optional_key in result: + filtered[optional_key] = [] + + for qi, ids in enumerate(ids_groups): + kept_indices: list[int] = [] + docs_group = (result.get("documents") or [[]])[qi] if result.get("documents") else None + for idx, _doc_id in enumerate(ids): + if not docs_group or idx >= len(docs_group) or docs_group[idx] is None: + continue + doc_text = str(docs_group[idx]) + if doc_matches_where_document(doc_text, where_document): + kept_indices.append(idx) + if len(kept_indices) >= n_results: + break + + def _pick(key: str) -> list[Any]: + """Pick kept elements from one result group.""" + groups = result.get(key) + if not groups or qi >= len(groups): + return [] + group = groups[qi] + return [group[i] for i in kept_indices if i < len(group)] + + filtered["ids"].append([ids[i] for i in kept_indices]) + if "distances" in result and result["distances"]: + filtered["distances"].append(_pick("distances")) + for optional_key in ("documents", "metadatas", "embeddings"): + if optional_key in filtered: + filtered[optional_key].append(_pick(optional_key)) + + return filtered + def _build_source_fields(self, include: list[str] | None) -> list[str]: """ Infer OceanBase GET_SQL `_source` allowlist from include. @@ -5825,27 +5829,39 @@ def _namespace_query( # noqa: C901 if where is not None: knn_cfg["where"] = where - query_cfg: dict[str, Any] | None = None + post_filter_wd: dict[str, Any] | str | None = None + hybrid_include = include if where_document is not None: - query_cfg = {"where_document": where_document, "n_results": n_results} - if where is not None: - query_cfg["where"] = where + if where_document_knn_prefilterable(where_document): + knn_cfg["where_document"] = where_document + else: + post_filter_wd = where_document + knn_cfg["n_results"] = min(max(n_results * 5, n_results), _MAX_N_RESULTS) + if include is None: + hybrid_include = ["documents", "metadatas"] + elif not {"documents", "document"} & {*(include or [])}: + hybrid_include = ["documents", *include] batch = self._namespace_hybrid_search( collection_id=collection_id, collection_name=collection_name, namespace_id=namespace_id, namespace_name=namespace_name, - query=query_cfg, + query=None, knn=knn_cfg, - n_results=n_results, - include=include, + n_results=knn_cfg["n_results"], + include=hybrid_include, embedding_function=embedding_function, distance=distance, dimension=kwargs.get("dimension"), **hybrid_kwargs, ) + if post_filter_wd is not None: + batch = self._post_filter_namespace_query_result( + batch, post_filter_wd, n_results=n_results, + ) + batch_ids = batch.get("ids") or [[]] all_ids.append(batch_ids[0] if batch_ids else []) all_distances.append((batch.get("distances") or [[]])[0]) diff --git a/src/pyseekdb/client/document_query_builder.py b/src/pyseekdb/client/document_query_builder.py new file mode 100644 index 00000000..501222c4 --- /dev/null +++ b/src/pyseekdb/client/document_query_builder.py @@ -0,0 +1,275 @@ +"""Build hybrid-search and SQL document filter expressions from ``where_document``.""" + +from __future__ import annotations + +import re +from typing import Any + +from pymysql.converters import escape_string + +_DOCUMENT_FIELD = "document" +_SCORING_LEAF_KEYS = frozenset({"query_string", "match", "multi_match", "match_phrase"}) + + +def _query_string_contains(query: str, *, boost: float | None = None) -> dict[str, Any]: + """Build a ``query_string`` leaf for document full-text match.""" + body: dict[str, Any] = { + "fields": [_DOCUMENT_FIELD], + "query": escape_string(query), + } + if boost is not None: + body["boost"] = boost + return {"query_string": body} + + +def _regexp_document(pattern: str, *, boost: float | None = None) -> dict[str, Any]: + """Build a regexp leaf on the document column.""" + body: dict[str, Any] = {"value": pattern} + if boost is not None: + body["boost"] = boost + return {"regexp": {_DOCUMENT_FIELD: body}} + + +def build_document_hybrid_expression( + where_document: dict[str, Any] | str, + *, + boost: float | None = None, +) -> dict[str, Any] | None: + """ + Recursively translate ``where_document`` into a hybrid_search document expression. + + Supports ``$contains``, ``$not_contains``, ``$regex``, ``$and``, ``$or`` and nesting. + """ + if isinstance(where_document, str): + return _query_string_contains(where_document, boost=boost) + + if not isinstance(where_document, dict) or not where_document: + return None + + if len(where_document) == 1: + op, value = next(iter(where_document.items())) + if op == "$contains" and isinstance(value, str): + return _query_string_contains(value, boost=boost) + if op == "$not_contains" and isinstance(value, str): + return { + "bool": { + "must_not": [_query_string_contains(value, boost=boost)], + } + } + if op == "$regex" and isinstance(value, str): + return _regexp_document(value, boost=boost) + if op == "$and" and isinstance(value, list): + return _combine_document_bool(value, combiner="and", boost=boost) + if op == "$or" and isinstance(value, list): + return _combine_document_bool(value, combiner="or", boost=boost) + + # Multiple top-level keys: treat as implicit AND. + parts = [ + build_document_hybrid_expression({key: val}, boost=boost) + for key, val in where_document.items() + ] + parts = [part for part in parts if part is not None] + if not parts: + return None + if len(parts) == 1: + return parts[0] + return {"bool": {"must": parts}} + + +def _contains_operator(where_document: dict[str, Any] | str, operator: str) -> bool: + """Return whether *operator* appears anywhere in a where_document tree.""" + if isinstance(where_document, str): + return False + if not isinstance(where_document, dict): + return False + if operator in where_document: + return True + for key in ("$and", "$or"): + children = where_document.get(key) + if isinstance(children, list): + for child in children: + if _contains_operator(child, operator): + return True + return False + + +def where_document_knn_prefilterable(where_document: dict[str, Any] | str | None) -> bool: + """ + Whether *where_document* can be enforced via ``knn.filter`` on OceanBase. + + Negative predicates (``$not_contains``) and ``$regex`` are not supported inside + knn filters and must be applied as SDK post-filters instead. + """ + if where_document is None: + return False + if _contains_operator(where_document, "$not_contains"): + return False + if _contains_operator(where_document, "$regex"): + return False + return True + + +def doc_matches_where_document(document: str, where_document: dict[str, Any] | str) -> bool: + """Evaluate a where_document predicate against a single document string.""" + text = document.lower() + if isinstance(where_document, str): + return where_document.lower() in text + if "$contains" in where_document: + return str(where_document["$contains"]).lower() in text + if "$not_contains" in where_document: + return str(where_document["$not_contains"]).lower() not in text + if "$regex" in where_document: + return re.search(str(where_document["$regex"]), document) is not None + if "$and" in where_document: + return all( + doc_matches_where_document(document, sub) + for sub in where_document["$and"] + ) + if "$or" in where_document: + return any( + doc_matches_where_document(document, sub) + for sub in where_document["$or"] + ) + raise ValueError(f"Unsupported where_document: {where_document!r}") + + +def _combine_document_bool( + conditions: list[Any], + *, + combiner: str, + boost: float | None, +) -> dict[str, Any] | None: + """Combine child ``where_document`` clauses with AND or OR.""" + if combiner == "and": + if all(isinstance(c, dict) and "$contains" in c and isinstance(c["$contains"], str) for c in conditions): + queries = [c["$contains"] for c in conditions if isinstance(c, dict)] + body: dict[str, Any] = { + "fields": [_DOCUMENT_FIELD], + "query": " ".join(escape_string(q) for q in queries), + "default_operator": "and", + } + if boost is not None: + body["boost"] = boost + return {"query_string": body} + + if all(isinstance(c, dict) and "$not_contains" in c and isinstance(c["$not_contains"], str) for c in conditions): + return { + "bool": { + "must_not": [ + _query_string_contains(c["$not_contains"], boost=boost) + for c in conditions + if isinstance(c, dict) + ], + "filter": [{"exists": {"field": _DOCUMENT_FIELD}}], + } + } + + if combiner == "or": + if all(isinstance(c, dict) and "$contains" in c and isinstance(c["$contains"], str) for c in conditions): + queries = [c["$contains"] for c in conditions if isinstance(c, dict)] + body: dict[str, Any] = { + "fields": [_DOCUMENT_FIELD], + "query": " ".join(escape_string(q) for q in queries), + "default_operator": "or", + } + if boost is not None: + body["boost"] = boost + return {"query_string": body} + + child_exprs: list[dict[str, Any]] = [] + for condition in conditions: + if not isinstance(condition, dict): + continue + expr = build_document_hybrid_expression(condition, boost=boost) + if expr is not None: + child_exprs.append(expr) + + if not child_exprs: + return None + + if combiner == "and": + must: list[dict[str, Any]] = [] + must_not: list[dict[str, Any]] = [] + for expr in child_exprs: + neg = _pure_must_not_clauses(expr) + if neg is not None: + must_not.extend(neg) + else: + must.append(expr) + bool_q: dict[str, Any] = {} + if must: + bool_q["must"] = must + if must_not: + bool_q["must_not"] = must_not + if not must: + bool_q["filter"] = [{"exists": {"field": _DOCUMENT_FIELD}}] + if not bool_q: + return None + return {"bool": bool_q} + + return { + "bool": { + "should": child_exprs, + "minimum_should_match": 1, + } + } + + +def _pure_must_not_clauses(expr: dict[str, Any] | None) -> list[dict[str, Any]] | None: + """Return inner ``must_not`` leaves when *expr* is a pure negation bool.""" + if not isinstance(expr, dict) or set(expr.keys()) != {"bool"}: + return None + bool_node = expr.get("bool") + if not isinstance(bool_node, dict) or set(bool_node.keys()) != {"must_not"}: + return None + clauses = bool_node.get("must_not") + if not isinstance(clauses, list): + return None + return clauses + + +def document_expr_as_knn_filter(doc_expr: dict[str, Any] | None) -> dict[str, Any] | None: + """ + Wrap a document expression for use inside ``knn.filter``. + + Scoring leaves (``query_string`` / ``match``) are wrapped in a non-scoring bool. + Pure ``must_not`` bools gain a permissive positive ``filter`` leaf (OB requirement). + """ + if doc_expr is None: + return None + + if "bool" in doc_expr: + bool_node = dict(doc_expr["bool"]) + must_not = bool_node.get("must_not") + if must_not and not bool_node.get("must") and not bool_node.get("filter"): + bool_node.setdefault("filter", [{"exists": {"field": _DOCUMENT_FIELD}}]) + return {"bool": bool_node} + + if _SCORING_LEAF_KEYS & doc_expr.keys(): + return {"bool": {"must": [doc_expr]}} + + return doc_expr + + +def merge_into_knn_filter( + knn_expr: dict[str, Any] | list[dict[str, Any]], + extra_filter: dict[str, Any] | None, +) -> dict[str, Any] | list[dict[str, Any]]: + """Append *extra_filter* to every knn expression's ``filter`` list.""" + if extra_filter is None: + return knn_expr + + def _merge_one(expr: dict[str, Any]) -> dict[str, Any]: + """Merge one knn expression.""" + existing = expr.get("filter") + if existing is None: + expr["filter"] = [extra_filter] + elif isinstance(existing, list): + expr["filter"] = [*existing, extra_filter] + else: + expr["filter"] = [existing, extra_filter] + return expr + + if isinstance(knn_expr, list): + return [_merge_one(item) for item in knn_expr] + return _merge_one(knn_expr) diff --git a/src/pyseekdb/client/filters.py b/src/pyseekdb/client/filters.py index 4e73926b..501d6d1f 100644 --- a/src/pyseekdb/client/filters.py +++ b/src/pyseekdb/client/filters.py @@ -28,7 +28,7 @@ class FilterBuilder: LOGICAL_OPS: ClassVar[list[str]] = ["$and", "$or", "$not"] # Document operators - DOCUMENT_OPS: ClassVar[list[str]] = ["$contains", "$regex"] + DOCUMENT_OPS: ClassVar[list[str]] = ["$contains", "$not_contains", "$regex"] @staticmethod def build_metadata_filter(where: dict[str, Any], metadata_column: str = "metadata") -> tuple[str, list[Any]]: @@ -188,6 +188,12 @@ def _build_document_condition( clauses.append(f"MATCH({document_column}) AGAINST (%s IN NATURAL LANGUAGE MODE)") params.append(value) + elif key == "$not_contains": + clauses.append( + f"NOT (MATCH({document_column}) AGAINST (%s IN NATURAL LANGUAGE MODE))" + ) + params.append(value) + elif key == "$regex": # Regular expression matching clauses.append(f"{document_column} REGEXP %s") diff --git a/tests/integration_tests/test_namespace_query_where_document.py b/tests/integration_tests/test_namespace_query_where_document.py new file mode 100644 index 00000000..8561f08b --- /dev/null +++ b/tests/integration_tests/test_namespace_query_where_document.py @@ -0,0 +1,284 @@ +""" +Integration tests: ns.query where_document operators (1.6.10). + +Vector query with document pre-filters applied via knn.filter. +""" + +from __future__ import annotations + +import time + +import pytest + +from namespace_dml_helpers import cleanup +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig +from pyseekdb.client.schema import Schema + + +def _schema() -> Schema: + """Namespace collection schema for where_document query tests.""" + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), + ) + + +def _create_ns(client, suffix: str): + """Create a namespace-enabled test collection.""" + name = f"test_ns_qb_{int(time.time() * 1000)}{suffix}" + return client.create_collection( + name=name, schema=_schema(), use_namespace=True, partition_count=4, + ) + + +def _seed(ns) -> None: + """Load deterministic corpus for where_document operator checks.""" + ns.add( + ids=["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k"], + embeddings=[ + [1.0, 0.0, 0.0], [1.0, 0.2, 0.0], [1.0, 0.4, 0.0], + [1.0, 0.6, 0.0], [1.0, 0.8, 0.0], [0.8, 1.0, 0.0], + [0.6, 1.0, 0.0], [0.4, 1.0, 0.0], [0.2, 1.0, 0.0], + [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], + ], + documents=[ + "machine learning AI research", + "Python programming language", + "OceanBase distributed database", + "deep learning neural networks", + "data science machine learning", + "web development Python Django", + "cloud computing AWS", + "database SQL optimization", + "container Docker Kubernetes", + "microservices architecture", + "blockchain distributed ledger", + ], + metadatas=[ + {"cat": "AI", "score": 95}, {"cat": "Prog", "score": 88}, + {"cat": "DB", "score": 92}, {"cat": "AI", "score": 90}, + {"cat": "DS", "score": 85}, {"cat": "Prog", "score": 75}, + {"cat": "Cloud", "score": 70}, {"cat": "DB", "score": 80}, + {"cat": "Cloud", "score": 65}, {"cat": "Arch", "score": 60}, + {"cat": "DB", "score": 55}, + ], + ) + + +class TestQueryWhereDocumentOperators: + """1.6.10: where_document $not_contains / $and / $or / $regex / nested combos.""" + + def test_where_document_contains_baseline(self, db_client): + """$contains baseline: single operator works under ns.query.""" + collection = _create_ns(db_client, "_wd_contains") + try: + ns = collection.create_namespace("ns_wd_contains") + _seed(ns) + time.sleep(1) + + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + where_document={"$contains": "machine"}, + n_results=10, + include=["documents"], + ) + docs = result.get("documents", [[]])[0] or [] + assert docs, "expected at least one hit for $contains 'machine'" + for doc in docs: + assert "machine" in doc.lower(), ( + f"$contains 'machine' returned doc without 'machine': {doc}" + ) + finally: + cleanup(db_client, collection) + + def test_where_document_not_contains(self, db_client): + """$not_contains excludes documents containing the keyword.""" + collection = _create_ns(db_client, "_wd_not_ctn") + try: + ns = collection.create_namespace("ns_wd_not_ctn") + _seed(ns) + time.sleep(1) + + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + where_document={"$not_contains": "machine"}, + n_results=10, + include=["documents"], + ) + docs = result.get("documents", [[]])[0] or [] + violations = [d for d in docs if "machine" in d.lower()] + assert not violations, ( + f"$not_contains 'machine' should exclude docs with 'machine', " + f"got violations: {violations}" + ) + finally: + cleanup(db_client, collection) + + def test_where_document_and(self, db_client): + """$and: multiple $contains must all match.""" + collection = _create_ns(db_client, "_wd_and") + try: + ns = collection.create_namespace("ns_wd_and") + _seed(ns) + time.sleep(1) + + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + where_document={ + "$and": [{"$contains": "machine"}, {"$contains": "learning"}] + }, + n_results=10, + include=["documents"], + ) + docs = result.get("documents", [[]])[0] or [] + violations = [ + d for d in docs + if not ("machine" in d.lower() and "learning" in d.lower()) + ] + assert not violations, ( + f"$and should require both 'machine' and 'learning', " + f"got violations: {violations}" + ) + finally: + cleanup(db_client, collection) + + def test_where_document_or(self, db_client): + """$or: any $contains child may match.""" + collection = _create_ns(db_client, "_wd_or") + try: + ns = collection.create_namespace("ns_wd_or") + _seed(ns) + time.sleep(1) + + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + where_document={ + "$or": [{"$contains": "blockchain"}, {"$contains": "Kubernetes"}] + }, + n_results=10, + include=["documents"], + ) + docs = result.get("documents", [[]])[0] or [] + violations = [ + d for d in docs + if not ("blockchain" in d.lower() or "kubernetes" in d.lower()) + ] + assert not violations, ( + f"$or should accept 'blockchain' or 'Kubernetes', " + f"got violations: {violations}" + ) + finally: + cleanup(db_client, collection) + + def test_where_document_regex(self, db_client): + """$regex filters documents by regexp on the document column.""" + collection = _create_ns(db_client, "_wd_regex") + try: + ns = collection.create_namespace("ns_wd_regex") + _seed(ns) + time.sleep(1) + + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + where_document={"$regex": r"machine.*learning"}, + n_results=10, + include=["documents"], + ) + docs = result.get("documents", [[]])[0] or [] + violations = [ + d for d in docs + if "machine" not in d.lower() or "learning" not in d.lower() + ] + assert not violations, ( + f"$regex should match docs with machine+learning, " + f"got violations: {violations}" + ) + finally: + cleanup(db_client, collection) + + def test_where_document_and_or_nested(self, db_client): + """Nested $and($or): (machine|database) AND learning.""" + collection = _create_ns(db_client, "_wd_and_or") + try: + ns = collection.create_namespace("ns_wd_and_or") + _seed(ns) + time.sleep(1) + + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + where_document={ + "$and": [ + { + "$or": [ + {"$contains": "machine"}, + {"$contains": "database"}, + ] + }, + {"$contains": "learning"}, + ] + }, + n_results=10, + include=["documents"], + ) + docs = result.get("documents", [[]])[0] or [] + violations = [ + d for d in docs + if not ( + ("machine" in d.lower() or "database" in d.lower()) + and "learning" in d.lower() + ) + ] + assert not violations, ( + f"Nested $and($or) should filter (machine|database) + learning, " + f"got violations: {violations}" + ) + finally: + cleanup(db_client, collection) + + def test_where_document_or_and_nested(self, db_client): + """Nested $or($and): (machine+learning) OR (distributed+database).""" + collection = _create_ns(db_client, "_wd_or_and") + try: + ns = collection.create_namespace("ns_wd_or_and") + _seed(ns) + time.sleep(1) + + result = ns.query( + query_embeddings=[1.0, 0.0, 0.0], + where_document={ + "$or": [ + { + "$and": [ + {"$contains": "machine"}, + {"$contains": "learning"}, + ] + }, + { + "$and": [ + {"$contains": "distributed"}, + {"$contains": "database"}, + ] + }, + ] + }, + n_results=10, + include=["documents"], + ) + docs = result.get("documents", [[]])[0] or [] + violations = [ + d for d in docs + if not ( + ("machine" in d.lower() and "learning" in d.lower()) + or ("distributed" in d.lower() and "database" in d.lower()) + ) + ] + assert not violations, ( + f"Nested $or($and) should filter (machine+learning) OR (distributed+database), " + f"got violations: {violations}" + ) + finally: + cleanup(db_client, collection) diff --git a/tests/unit_tests/test_document_query_builder.py b/tests/unit_tests/test_document_query_builder.py new file mode 100644 index 00000000..314d92b7 --- /dev/null +++ b/tests/unit_tests/test_document_query_builder.py @@ -0,0 +1,61 @@ +"""Unit tests for document_query_builder.""" + +import sys +from pathlib import Path + +import pytest + +project_root = Path(__file__).parent.parent.parent +sys.path.insert(0, str(project_root / "src")) + +from pyseekdb.client.document_query_builder import ( # noqa: E402 + build_document_hybrid_expression, + document_expr_as_knn_filter, +) + + +class TestBuildDocumentHybridExpression: + """TestBuildDocumentHybridExpression class.""" + + def test_not_contains_builds_must_not(self): + """Test not contains builds must not.""" + expr = build_document_hybrid_expression({"$not_contains": "machine"}) + assert expr == { + "bool": { + "must_not": [{"query_string": {"fields": ["document"], "query": "machine"}}], + } + } + + def test_and_or_nested(self): + """Test and or nested.""" + expr = build_document_hybrid_expression({ + "$and": [ + {"$or": [{"$contains": "a"}, {"$contains": "b"}]}, + {"$contains": "c"}, + ] + }) + assert expr is not None + assert "bool" in expr + assert "must" in expr["bool"] + + def test_regex_leaf(self): + """Test regex leaf.""" + expr = build_document_hybrid_expression({"$regex": r"machine.*learning"}) + assert expr == {"regexp": {"document": {"value": r"machine.*learning"}}} + + +class TestDocumentExprAsKnnFilter: + """TestDocumentExprAsKnnFilter class.""" + + def test_pure_must_not_adds_exists_filter(self): + """Test pure must not adds exists filter.""" + expr = build_document_hybrid_expression({"$not_contains": "x"}) + wrapped = document_expr_as_knn_filter(expr) + assert wrapped is not None + assert wrapped["bool"]["filter"] == [{"exists": {"field": "document"}}] + + def test_contains_wrapped_in_bool_must(self): + """Test contains wrapped in bool must.""" + expr = build_document_hybrid_expression({"$contains": "hello"}) + wrapped = document_expr_as_knn_filter(expr) + assert wrapped == {"bool": {"must": [expr]}} From a0c9599165ad25cd0e6732d0cdf7b920f5f044ab Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 26 Jun 2026 11:09:30 +0800 Subject: [PATCH 68/84] fix:issue --- docs/guide/dql.md | 5 ++- src/pyseekdb/client/client_base.py | 3 +- src/pyseekdb/client/document_query_builder.py | 40 +++++-------------- src/pyseekdb/client/kernel_errors.py | 2 +- .../namespace_hybrid_search_helpers.py | 2 +- .../test_namespace_query_where_document.py | 6 +++ .../unit_tests/test_document_query_builder.py | 35 ++++++++++++++-- 7 files changed, 55 insertions(+), 38 deletions(-) diff --git a/docs/guide/dql.md b/docs/guide/dql.md index 1585719a..b19b1fd7 100644 --- a/docs/guide/dql.md +++ b/docs/guide/dql.md @@ -101,7 +101,7 @@ results = collection.query( - `n_results` (int, required): Number of similar results to return (default: 10) - `where` (dict, optional): Metadata filter conditions (see Filter Operators section) - `where_document` (dict, optional): Document content filter -- `include` (List[str], optional): List of fields to include: `["documents", "metadatas", "embeddings"]` +- `include` (List[str], optional): List of fields to include: `["documents", "metadatas", "embeddings", "distances"]` (singular aliases accepted). Invalid field names raise `ValueError` before the query runs. **Returns:** Dict with keys (chromadb-compatible format): @@ -288,7 +288,8 @@ results = collection.hybrid_search( ### Document Filters (`where_document` parameter) - `$contains`: full-text match - `$not_contains`: exclude matches -- `$or` / `$and` combining multiple `$contains` clauses +- `$regex`: regular expression match (namespace `query` applies via SDK post-filter when OB knn filter cannot enforce it) +- `$or` / `$and` combining multiple document clauses (including nested combinations) ## 5.5 Collection Information Methods diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index d22b57d7..a103bdc8 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -21,7 +21,6 @@ build_document_hybrid_expression, doc_matches_where_document, document_expr_as_knn_filter, - merge_into_knn_filter, where_document_knn_prefilterable, ) from .kernel_errors import maybe_reraise_friendly_kernel_error, namespace_kernel_error_guard @@ -150,7 +149,7 @@ def _reraise_unless_unique_index_exists(exc: BaseException) -> None: "already exists" in message or "duplicate key name" in message or "code=1061" in message - or "1061" in message and "duplicate" in message + or ("1061" in message and "duplicate" in message) ): return raise exc diff --git a/src/pyseekdb/client/document_query_builder.py b/src/pyseekdb/client/document_query_builder.py index 501222c4..712f76d5 100644 --- a/src/pyseekdb/client/document_query_builder.py +++ b/src/pyseekdb/client/document_query_builder.py @@ -5,17 +5,23 @@ import re from typing import Any -from pymysql.converters import escape_string - _DOCUMENT_FIELD = "document" _SCORING_LEAF_KEYS = frozenset({"query_string", "match", "multi_match", "match_phrase"}) +# Lucene query_string reserved characters (Elasticsearch-style). +_QUERY_STRING_RESERVED_RE = re.compile(r'([+\-=&|> str: + """Escape Lucene ``query_string`` syntax characters in a user term.""" + escaped = text.replace("\\", "\\\\") + return _QUERY_STRING_RESERVED_RE.sub(r"\\\1", escaped) def _query_string_contains(query: str, *, boost: float | None = None) -> dict[str, Any]: """Build a ``query_string`` leaf for document full-text match.""" body: dict[str, Any] = { "fields": [_DOCUMENT_FIELD], - "query": escape_string(query), + "query": _escape_query_string_term(query), } if boost is not None: body["boost"] = boost @@ -145,7 +151,7 @@ def _combine_document_bool( queries = [c["$contains"] for c in conditions if isinstance(c, dict)] body: dict[str, Any] = { "fields": [_DOCUMENT_FIELD], - "query": " ".join(escape_string(q) for q in queries), + "query": " ".join(_escape_query_string_term(q) for q in queries), "default_operator": "and", } if boost is not None: @@ -169,7 +175,7 @@ def _combine_document_bool( queries = [c["$contains"] for c in conditions if isinstance(c, dict)] body: dict[str, Any] = { "fields": [_DOCUMENT_FIELD], - "query": " ".join(escape_string(q) for q in queries), + "query": " ".join(_escape_query_string_term(q) for q in queries), "default_operator": "or", } if boost is not None: @@ -249,27 +255,3 @@ def document_expr_as_knn_filter(doc_expr: dict[str, Any] | None) -> dict[str, An return {"bool": {"must": [doc_expr]}} return doc_expr - - -def merge_into_knn_filter( - knn_expr: dict[str, Any] | list[dict[str, Any]], - extra_filter: dict[str, Any] | None, -) -> dict[str, Any] | list[dict[str, Any]]: - """Append *extra_filter* to every knn expression's ``filter`` list.""" - if extra_filter is None: - return knn_expr - - def _merge_one(expr: dict[str, Any]) -> dict[str, Any]: - """Merge one knn expression.""" - existing = expr.get("filter") - if existing is None: - expr["filter"] = [extra_filter] - elif isinstance(existing, list): - expr["filter"] = [*existing, extra_filter] - else: - expr["filter"] = [existing, extra_filter] - return expr - - if isinstance(knn_expr, list): - return [_merge_one(item) for item in knn_expr] - return _merge_one(knn_expr) diff --git a/src/pyseekdb/client/kernel_errors.py b/src/pyseekdb/client/kernel_errors.py index 78c686df..6bb0445e 100644 --- a/src/pyseekdb/client/kernel_errors.py +++ b/src/pyseekdb/client/kernel_errors.py @@ -61,7 +61,7 @@ def _is_namespace_dropping_error(exc: BaseException) -> bool: _has_ob_error_code(exc, 4109) or "namespace is dropping" in text or "namespace is in drop" in text - or "being dropped" in text and "namespace" in text + or ("being dropped" in text and "namespace" in text) ) diff --git a/tests/integration_tests/namespace_hybrid_search_helpers.py b/tests/integration_tests/namespace_hybrid_search_helpers.py index 28a94f05..6280a0de 100644 --- a/tests/integration_tests/namespace_hybrid_search_helpers.py +++ b/tests/integration_tests/namespace_hybrid_search_helpers.py @@ -136,7 +136,7 @@ def count_corpus_matches(corpus: list[CorpusRecord], where: dict[str, Any]) -> i def l2_squared(a: list[float], b: list[float]) -> float: """L2 squared.""" - return sum((x - y) ** 2 for x, y in zip(a, b)) + return sum((x - y) ** 2 for x, y in zip(a, b, strict=True)) def _vector_l2_norm(vec: list[float]) -> float: diff --git a/tests/integration_tests/test_namespace_query_where_document.py b/tests/integration_tests/test_namespace_query_where_document.py index 8561f08b..1fbd7cdf 100644 --- a/tests/integration_tests/test_namespace_query_where_document.py +++ b/tests/integration_tests/test_namespace_query_where_document.py @@ -110,6 +110,7 @@ def test_where_document_not_contains(self, db_client): include=["documents"], ) docs = result.get("documents", [[]])[0] or [] + assert docs, "expected at least one hit for $not_contains 'machine'" violations = [d for d in docs if "machine" in d.lower()] assert not violations, ( f"$not_contains 'machine' should exclude docs with 'machine', " @@ -135,6 +136,7 @@ def test_where_document_and(self, db_client): include=["documents"], ) docs = result.get("documents", [[]])[0] or [] + assert docs, "expected at least one hit for $and machine+learning" violations = [ d for d in docs if not ("machine" in d.lower() and "learning" in d.lower()) @@ -163,6 +165,7 @@ def test_where_document_or(self, db_client): include=["documents"], ) docs = result.get("documents", [[]])[0] or [] + assert docs, "expected at least one hit for $or blockchain|Kubernetes" violations = [ d for d in docs if not ("blockchain" in d.lower() or "kubernetes" in d.lower()) @@ -189,6 +192,7 @@ def test_where_document_regex(self, db_client): include=["documents"], ) docs = result.get("documents", [[]])[0] or [] + assert docs, "expected at least one hit for $regex machine.*learning" violations = [ d for d in docs if "machine" not in d.lower() or "learning" not in d.lower() @@ -225,6 +229,7 @@ def test_where_document_and_or_nested(self, db_client): include=["documents"], ) docs = result.get("documents", [[]])[0] or [] + assert docs, "expected at least one hit for nested $and($or)" violations = [ d for d in docs if not ( @@ -269,6 +274,7 @@ def test_where_document_or_and_nested(self, db_client): include=["documents"], ) docs = result.get("documents", [[]])[0] or [] + assert docs, "expected at least one hit for nested $or($and)" violations = [ d for d in docs if not ( diff --git a/tests/unit_tests/test_document_query_builder.py b/tests/unit_tests/test_document_query_builder.py index 314d92b7..2ceb96df 100644 --- a/tests/unit_tests/test_document_query_builder.py +++ b/tests/unit_tests/test_document_query_builder.py @@ -34,9 +34,35 @@ def test_and_or_nested(self): {"$contains": "c"}, ] }) - assert expr is not None - assert "bool" in expr - assert "must" in expr["bool"] + assert expr == { + "bool": { + "must": [ + { + "query_string": { + "fields": ["document"], + "query": "a b", + "default_operator": "or", + } + }, + { + "query_string": { + "fields": ["document"], + "query": "c", + } + }, + ] + } + } + + def test_query_string_escapes_reserved_characters(self): + """Reserved Lucene characters are escaped in query_string terms.""" + expr = build_document_hybrid_expression({"$contains": "C++"}) + assert expr == { + "query_string": { + "fields": ["document"], + "query": r"C\+\+", + } + } def test_regex_leaf(self): """Test regex leaf.""" @@ -52,6 +78,9 @@ def test_pure_must_not_adds_exists_filter(self): expr = build_document_hybrid_expression({"$not_contains": "x"}) wrapped = document_expr_as_knn_filter(expr) assert wrapped is not None + assert wrapped["bool"]["must_not"] == [ + {"query_string": {"fields": ["document"], "query": "x"}}, + ] assert wrapped["bool"]["filter"] == [{"exists": {"field": "document"}}] def test_contains_wrapped_in_bool_must(self): From 513a33a7de2b8b461a0a71be345a1f0752f7a62c Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 26 Jun 2026 15:14:25 +0800 Subject: [PATCH 69/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91ob=E5=86=85?= =?UTF-8?q?=E6=A0=B8=E4=B8=8D=E6=94=AF=E6=8C=81ivf=E7=B4=A2=E5=BC=95?= =?UTF-8?q?=E8=B6=85=E8=B6=8A2048=E7=BB=B4=EF=BC=8C=E4=BD=86=E6=98=AFsdk?= =?UTF-8?q?=E7=BB=99=E7=9A=84=E8=8C=83=E5=9B=B4=E6=98=AF[1,4096]?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 22 ++++++-- src/pyseekdb/client/configuration.py | 12 +++- .../test_namespace_constraints.py | 56 +++++++++++++++++++ tests/unit_tests/test_namespace.py | 1 + 4 files changed, 82 insertions(+), 9 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index a103bdc8..55dadca3 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -30,6 +30,9 @@ from .configuration import ( DEFAULT_DISTANCE_METRIC, DEFAULT_VECTOR_DIMENSION, + LOGIC_DATA_TABLE_LOB_INROW_THRESHOLD, + MAX_HNSW_VECTOR_DIMENSION, + MAX_IVF_VECTOR_DIMENSION, Configuration, ConfigurationParam, FulltextIndexConfig, @@ -1009,8 +1012,10 @@ def create_collection( ) dimension = hnsw_config.dimension - if dimension < 1 or dimension > 4096: - raise ValueError(f"Dimension must be between 1 and 4096, got {dimension}") + if dimension < 1 or dimension > MAX_HNSW_VECTOR_DIMENSION: + raise ValueError( + f"Dimension must be between 1 and {MAX_HNSW_VECTOR_DIMENSION}, got {dimension}" + ) # Extract fulltext parser configuration fulltext_index_clause = _get_fulltext_index_sql(schema.fulltext_index) @@ -1109,9 +1114,6 @@ def _create_namespace_collection(self, name: str, schema: Schema, partition_coun dimension = DEFAULT_VECTOR_DIMENSION distance = DEFAULT_DISTANCE_METRIC - if dimension < 1 or dimension > 4096: - raise ValueError(f"Dimension must be between 1 and 4096, got {dimension}") - is_ss = self._is_shared_storage_mode() settings = { "version": 2, @@ -1134,6 +1136,14 @@ def _create_namespace_collection(self, name: str, schema: Schema, partition_coun collection_meta = self._create_ns_collection_meta(name, settings) collection_id = collection_meta["collection_id"] + if isinstance(dimension, bool) or not isinstance(dimension, int): + raise TypeError(f"dimension must be an integer, got {type(dimension).__name__}") + if dimension < 1 or dimension > MAX_IVF_VECTOR_DIMENSION: + raise ValueError( + f"Dimension must be between 1 and {MAX_IVF_VECTOR_DIMENSION} for namespace " + f"IVF collections, got {dimension}" + ) + self._ensure_namespace_catalogs() try: @@ -1521,7 +1531,7 @@ def _create_namespace_physical_tables( created_by VARCHAR(64) DEFAULT '', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, {index_sql} - ) TABLEGROUP=`{tg_name}` COMMENT='逻辑表主数据' DEFAULT CHARSET=utf8mb4 ORGANIZATION HEAP IS_LOGIC_TABLE = TRUE + ) TABLEGROUP=`{tg_name}` COMMENT='逻辑表主数据' DEFAULT CHARSET=utf8mb4 ORGANIZATION HEAP IS_LOGIC_TABLE = TRUE LOB_INROW_THRESHOLD={LOGIC_DATA_TABLE_LOB_INROW_THRESHOLD} {partition_clause}""" self._execute(data_sql) diff --git a/src/pyseekdb/client/configuration.py b/src/pyseekdb/client/configuration.py index 8dec0c28..b6017ac2 100644 --- a/src/pyseekdb/client/configuration.py +++ b/src/pyseekdb/client/configuration.py @@ -14,6 +14,11 @@ # So we use 384 as the default dimension to match DEFAULT_VECTOR_DIMENSION = 384 # Matches DefaultEmbeddingFunction dimension DEFAULT_DISTANCE_METRIC = "cosine" +MAX_HNSW_VECTOR_DIMENSION = 4096 +# Namespace IVF uses logic_data_table with LOB_INROW_THRESHOLD sized for float32 vectors +# (dimension * 4 bytes). Below the default ~8KB threshold, IVF indexing rejects out-row LOB. +MAX_IVF_VECTOR_DIMENSION = MAX_HNSW_VECTOR_DIMENSION +LOGIC_DATA_TABLE_LOB_INROW_THRESHOLD = MAX_IVF_VECTOR_DIMENSION * 4 PrimitiveValue = str | int | float | bool @@ -131,7 +136,7 @@ def _validate_hnsw_base_fields(config: "HNSWConfiguration") -> None: """Validate core HNSW index fields shared across index subtypes.""" if isinstance(config.dimension, bool) or not isinstance(config.dimension, int): raise TypeError(f"dimension must be an integer, got {type(config.dimension).__name__}") - _validate_int_range(config.dimension, key="dimension", min_value=1, max_value=4096) + _validate_int_range(config.dimension, key="dimension", min_value=1, max_value=MAX_HNSW_VECTOR_DIMENSION) config.distance = _normalize_str_enum(config.distance, field_name="distance") valid_distances = [e.value for e in DistanceMetric] @@ -373,7 +378,8 @@ class IVFConfiguration: IVF (Inverted File) index configuration for Agent Database namespace-enabled collections. Args: - dimension: Vector dimension (number of elements in each vector) + dimension: Vector dimension (1..4096 on namespace IVF). logic_data_table is + created with ``LOB_INROW_THRESHOLD`` so float32 vectors stay in-row. distance: Distance metric for similarity calculation (e.g., 'l2', 'cosine', 'inner_product') type: IVF index subtype ('ivf_flat', 'ivf_sq8', 'ivf_pq') centroids_fresh_mode: SPFresh mode for online index updates (e.g., 'spfresh'). Defaults to None (disabled). @@ -391,7 +397,7 @@ def __post_init__(self): """Validate IVF configuration fields and normalize properties.""" if isinstance(self.dimension, bool) or not isinstance(self.dimension, int): raise TypeError(f"dimension must be an integer, got {type(self.dimension).__name__}") - _validate_int_range(self.dimension, key="dimension", min_value=1, max_value=4096) + _validate_int_range(self.dimension, key="dimension", min_value=1, max_value=MAX_IVF_VECTOR_DIMENSION) self.distance = _normalize_str_enum(self.distance, field_name="distance") valid_distances = [e.value for e in DistanceMetric] diff --git a/tests/integration_tests/test_namespace_constraints.py b/tests/integration_tests/test_namespace_constraints.py index 5a0c1af4..651d94f7 100644 --- a/tests/integration_tests/test_namespace_constraints.py +++ b/tests/integration_tests/test_namespace_constraints.py @@ -61,6 +61,62 @@ def test_ivf_flat_accepted(self, oceanbase_client): oceanbase_client.delete_collection(name=name) +class TestNamespaceIvfDimensionConstraint: + """Namespace IVF dimension must stay within SDK max (logic_data_table LOB in-row threshold).""" + + def test_dimension_2049_accepted(self, oceanbase_client): + """Dimensions above default LOB threshold succeed with logic_data_table DDL fix.""" + name = _unique_name("_dim_2049") + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=2049, distance="cosine", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + collection = oceanbase_client.create_collection( + name=name, schema=schema, use_namespace=True, partition_count=1, + ) + try: + assert collection.dimension == 2049 + finally: + oceanbase_client.delete_collection(name=name) + + def test_dimension_4096_accepted(self, oceanbase_client): + """Positive control: max supported IVF dimension succeeds.""" + name = _unique_name("_dim_4096") + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=4096, distance="cosine", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + collection = oceanbase_client.create_collection( + name=name, schema=schema, use_namespace=True, partition_count=1, + ) + try: + assert collection.dimension == 4096 + finally: + oceanbase_client.delete_collection(name=name) + + def test_dimension_4097_rejected_at_sdk(self, oceanbase_client): + """Dimensions above 4096 are rejected before OB DDL.""" + name = _unique_name("_dim_4097") + with pytest.raises(ValueError, match="between 1 and 4096"): + oceanbase_client.create_collection( + name=name, + schema=Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration( + dimension=4097, distance="cosine", centroids_fresh_mode="spfresh", + ), + embedding_function=None, + ), + ), + use_namespace=True, + ) + assert not oceanbase_client.has_collection(name) + + class TestNamespaceMinVersionConstraint: """TestNamespaceMinVersionConstraint class.""" def test_connected_ob_meets_min_version(self, oceanbase_client): diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 31a05010..ead00dd7 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -834,6 +834,7 @@ def test_create_namespace_physical_tables_sn_inline_vector_index(self): assert "FULLTEXT INDEX idx_fts(document) WITH PARSER ik" in data_create assert "SEARCH INDEX idx_json(data_content)" in data_create assert "centroids_fresh_mode=spfresh" in data_create + assert "LOB_INROW_THRESHOLD=16384" in data_create assert not any(s.startswith("CREATE VECTOR INDEX") for s in c.executed_sqls) assert not any("_hot_table" in s for s in c.executed_sqls if s.startswith("CREATE TABLE")) From b39abbd0f606d159bdf02c9381ee5fd1e6b80053 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 26 Jun 2026 17:28:44 +0800 Subject: [PATCH 70/84] =?UTF-8?q?fix:=E3=80=90pyseekdb=E3=80=91=20?= =?UTF-8?q?=E4=B8=8D=E4=BC=A0=E5=85=A5=E5=90=91=E9=87=8F=E7=B4=A2=E5=BC=95?= =?UTF-8?q?=E9=85=8D=E7=BD=AE=E7=9A=84=E6=97=B6=E5=80=99=EF=BC=8C=E5=90=91?= =?UTF-8?q?=E9=87=8F=E5=88=97=E6=98=AFsdk=E5=86=85=E9=83=A8=E6=8C=87?= =?UTF-8?q?=E5=AE=9A=E7=9A=84384=E7=BB=B4=EF=BC=8Cns.add=E6=8A=A5=E9=94=99?= =?UTF-8?q?=E4=BF=A1=E6=81=AF=E4=B8=8D=E6=98=8E=E7=A1=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 155 ++++++++++++++++-- src/pyseekdb/client/collection.py | 7 + src/pyseekdb/client/namespace.py | 12 ++ src/pyseekdb/client/validators.py | 37 +++++ .../test_namespace_optional_indexes.py | 20 +++ tests/unit_tests/test_namespace.py | 59 +++++++ .../test_namespace_no_index_embedding_dim.py | 47 ++++++ 7 files changed, 322 insertions(+), 15 deletions(-) create mode 100644 tests/unit_tests/test_namespace_no_index_embedding_dim.py diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 55dadca3..3c36b4cb 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -223,6 +223,8 @@ def _validate_collection_name(name: str) -> None: _quote_sql_identifier, _validate_database_name, _validate_namespace_name, + _validate_namespace_no_index_explicit_embeddings, + _validate_namespace_explicit_embedding_dimensions, _validate_record_ids, ) @@ -1178,6 +1180,7 @@ def _create_namespace_collection(self, name: str, schema: Schema, partition_coun distance=distance, use_namespace=True, partition_count=pc, + has_vector_index=settings.get("dense_index_type") == "ivf", ) def _get_embedding_function_dimension(self, embedding_function: EmbeddingFunction) -> int: @@ -2019,6 +2022,7 @@ def _build_ns_collection_from_meta(self, meta: dict, embedding_function=_NOT_PRO distance=distance, use_namespace=True, partition_count=partition_count, + has_vector_index=settings.get("dense_index_type") == "ivf", ) def _resolve_collection_metadata_from_sdk_collections(self, collection_name: str) -> _CollectionMeta | None: @@ -2872,6 +2876,7 @@ def _collection_add( # noqa: C901 """ logger.debug(f"Adding data to collection '{collection_name}'") + explicit_embeddings = embeddings is not None # Normalize inputs to lists if isinstance(ids, str): ids = [ids] @@ -2887,6 +2892,13 @@ def _collection_add( # noqa: C901 ): embeddings = [embeddings] + self._warn_explicit_embeddings_override_embedding_function( + operation="collection.add", + explicit_embeddings=explicit_embeddings, + has_documents=bool(documents), + embedding_function=embedding_function, + ) + # Handle vector generation logic: # 1. If embeddings are provided, use them directly without embedding # 2. If embeddings are not provided but documents are provided: @@ -3048,6 +3060,7 @@ def _collection_update( # noqa: C901 """ logger.debug(f"Updating data in collection '{collection_name}'") + explicit_embeddings = embeddings is not None # Normalize inputs to lists if isinstance(ids, str): ids = [ids] @@ -3063,6 +3076,13 @@ def _collection_update( # noqa: C901 ): embeddings = [embeddings] + self._warn_explicit_embeddings_override_embedding_function( + operation="collection.update", + explicit_embeddings=explicit_embeddings, + has_documents=bool(documents), + embedding_function=embedding_function, + ) + # Handle vector generation logic: # 1. If embeddings are provided, use them directly without embedding # 2. If embeddings are not provided but documents are provided: @@ -4890,6 +4910,13 @@ def _build_knn_expression( # noqa: C901 embedding_function = kwargs.get("embedding_function") + self._warn_explicit_embeddings_override_embedding_function( + operation="hybrid_search.knn", + explicit_embeddings=query_embeddings is not None, + has_documents=query_texts is not None, + embedding_function=embedding_function, + ) + def _normalize_vectors(raw_embeddings: Any) -> list[list[float]]: """Normalize input vectors to a consistent list-of-floats form.""" if raw_embeddings is None: @@ -5262,6 +5289,45 @@ def _append_namespace_filter( return f"WHERE {ns_cond} AND ({inner})", params return f"WHERE {ns_cond} AND ({where_clause})", params + @staticmethod + def _validate_namespace_explicit_embeddings_if_needed( + embeddings: list[list[float]] | None, + *, + explicit_embeddings: bool, + has_vector_index: bool, + collection_dimension: int | None, + ) -> None: + """Validate user-supplied embeddings match the collection VECTOR column dimension.""" + if not explicit_embeddings or not embeddings: + return + expected = ( + collection_dimension + if collection_dimension is not None + else DEFAULT_VECTOR_DIMENSION + ) + _validate_namespace_explicit_embedding_dimensions( + embeddings, + expected_dimension=expected, + has_vector_index=has_vector_index, + ) + + @staticmethod + def _warn_explicit_embeddings_override_embedding_function( + *, + operation: str, + explicit_embeddings: bool, + has_documents: bool, + embedding_function: EmbeddingFunction[EmbeddingDocuments] | None, + ) -> None: + """Log when explicit embeddings take priority over embedding_function.""" + if explicit_embeddings and has_documents and embedding_function is not None: + logger.warning( + "%s: explicit embeddings provided together with documents while an " + "embedding_function is configured; using explicit embeddings and " + "not calling embedding_function.", + operation, + ) + def _count_namespace_records_by_id( self, table_name: str, @@ -5383,10 +5449,9 @@ def _namespace_add( # noqa: C901 **kwargs, ) -> None: """Add records to a namespace collection.""" - ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) - self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, - ) + has_vector_index = kwargs.pop("has_vector_index", True) + collection_dimension = kwargs.pop("collection_dimension", None) + explicit_embeddings = embeddings is not None if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) @@ -5406,6 +5471,13 @@ def _namespace_add( # noqa: C901 ): embeddings = [embeddings] + self._warn_explicit_embeddings_override_embedding_function( + operation="namespace.add", + explicit_embeddings=explicit_embeddings, + has_documents=bool(documents), + embedding_function=embedding_function, + ) + if embeddings: pass elif documents: @@ -5437,6 +5509,17 @@ def _namespace_add( # noqa: C901 if embeddings and len(embeddings) != num_items: raise ValueError(f"Number of embeddings ({len(embeddings)}) does not match number of ids ({num_items})") + self._validate_namespace_explicit_embeddings_if_needed( + embeddings, + explicit_embeddings=explicit_embeddings, + has_vector_index=has_vector_index, + collection_dimension=collection_dimension, + ) + + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) @@ -5479,10 +5562,9 @@ def _namespace_update( **kwargs, ) -> None: """Update existing records in a namespace collection.""" - ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) - self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, - ) + has_vector_index = kwargs.pop("has_vector_index", True) + collection_dimension = kwargs.pop("collection_dimension", None) + explicit_embeddings = embeddings is not None if isinstance(ids, str): ids = [ids] _validate_record_ids(ids) @@ -5502,6 +5584,13 @@ def _namespace_update( ): embeddings = [embeddings] + self._warn_explicit_embeddings_override_embedding_function( + operation="namespace.update", + explicit_embeddings=explicit_embeddings, + has_documents=bool(documents), + embedding_function=embedding_function, + ) + if embeddings: # embeddings provided, use them directly without embedding pass @@ -5517,6 +5606,17 @@ def _namespace_update( " 2. Provide embedding_function to auto-generate embeddings from documents." ) + self._validate_namespace_explicit_embeddings_if_needed( + embeddings, + explicit_embeddings=explicit_embeddings, + has_vector_index=has_vector_index, + collection_dimension=collection_dimension, + ) + + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) @@ -5626,6 +5726,8 @@ def _namespace_upsert( **kwargs, ) -> None: """Insert or update records in a namespace collection.""" + has_vector_index = kwargs.pop("has_vector_index", True) + collection_dimension = kwargs.pop("collection_dimension", None) ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, @@ -5686,7 +5788,10 @@ def _namespace_upsert( collection_id=collection_id, collection_name=collection_name, namespace_id=namespace_id, namespace_name=namespace_name, ids=add_ids, embeddings=add_embs, metadatas=add_metas, - documents=add_docs, embedding_function=embedding_function, **kwargs, + documents=add_docs, embedding_function=embedding_function, + has_vector_index=has_vector_index, + collection_dimension=collection_dimension, + **kwargs, ) if update_indices: @@ -5698,7 +5803,10 @@ def _namespace_upsert( collection_id=collection_id, collection_name=collection_name, namespace_id=namespace_id, namespace_name=namespace_name, ids=upd_ids, embeddings=upd_embs, metadatas=upd_metas, - documents=upd_docs, embedding_function=embedding_function, **kwargs, + documents=upd_docs, embedding_function=embedding_function, + has_vector_index=has_vector_index, + collection_dimension=collection_dimension, + **kwargs, ) self._reconcile_namespace_duplicate_records( @@ -5797,13 +5905,19 @@ def _namespace_query( # noqa: C901 **kwargs, ) -> dict[str, Any]: """Run a vector query against a namespace collection.""" - ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) - self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, - ) + has_vector_index = kwargs.pop("has_vector_index", True) + collection_dimension = kwargs.pop("collection_dimension", kwargs.pop("dimension", None)) + explicit_query_embeddings = query_embeddings is not None embedding_function = kwargs.get("embedding_function") distance = kwargs.get("distance", DEFAULT_DISTANCE_METRIC) + self._warn_explicit_embeddings_override_embedding_function( + operation="namespace.query", + explicit_embeddings=explicit_query_embeddings, + has_documents=query_texts is not None, + embedding_function=embedding_function, + ) + if query_embeddings is not None: pass elif query_texts is not None: @@ -5815,6 +5929,17 @@ def _namespace_query( # noqa: C901 raise ValueError("Neither query_embeddings nor query_texts provided.") query_embeddings = self._normalize_query_embeddings(query_embeddings) + self._validate_namespace_explicit_embeddings_if_needed( + query_embeddings, + explicit_embeddings=explicit_query_embeddings, + has_vector_index=has_vector_index, + collection_dimension=collection_dimension, + ) + + ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) + self._set_session_ns_context( + collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + ) include_fields = self._normalize_include_fields(include) # Logical-table vector search must use hybrid_search DSL; direct SQL vector @@ -5862,7 +5987,7 @@ def _namespace_query( # noqa: C901 include=hybrid_include, embedding_function=embedding_function, distance=distance, - dimension=kwargs.get("dimension"), + dimension=collection_dimension, **hybrid_kwargs, ) diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 517fc007..2502bbae 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -48,6 +48,7 @@ def __init__( sparse_vector_index_config: Optional["SparseVectorIndexConfig"] = None, use_namespace: bool = False, partition_count: int | None = None, + has_vector_index: bool = False, **metadata, ): """Initialize a lightweight collection handle bound to a client implementation.""" @@ -60,6 +61,7 @@ def __init__( self._sparse_vector_index_config = sparse_vector_index_config self._use_namespace = use_namespace self._partition_count = partition_count + self._has_vector_index = has_vector_index self._metadata = metadata # ==================== Properties ==================== @@ -116,6 +118,11 @@ def use_namespace(self) -> bool: """Whether this collection routes data operations through namespaces.""" return self._use_namespace + @property + def has_vector_index(self) -> bool: + """Whether this collection has a dense VECTOR INDEX (namespace IVF).""" + return self._has_vector_index + @property def partition_count(self) -> int | None: """Number of partitions for the namespace physical tables. diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py index 4213f191..a9de484c 100644 --- a/src/pyseekdb/client/namespace.py +++ b/src/pyseekdb/client/namespace.py @@ -59,6 +59,14 @@ def __repr__(self) -> str: f"collection='{self._collection.name}')" ) + def _dml_collection_context(self) -> dict[str, Any]: + """Collection metadata required for namespace DML/DQL validation.""" + return { + "has_vector_index": self._collection.has_vector_index, + "collection_dimension": self._collection.dimension, + "dimension": self._collection.dimension, + } + def _guard_exists(self) -> None: """Raise if this namespace or its collection was deleted.""" if not self._client._ns_collection_exists_by_id(self._collection.id): @@ -94,6 +102,7 @@ def add( metadatas=metadatas, documents=documents, embedding_function=self._collection.embedding_function, + **self._dml_collection_context(), **kwargs, ) @@ -117,6 +126,7 @@ def update( metadatas=metadatas, documents=documents, embedding_function=self._collection.embedding_function, + **self._dml_collection_context(), **kwargs, ) @@ -140,6 +150,7 @@ def upsert( metadatas=metadatas, documents=documents, embedding_function=self._collection.embedding_function, + **self._dml_collection_context(), **kwargs, ) @@ -192,6 +203,7 @@ def query( include=include, embedding_function=self._collection.embedding_function, distance=self._collection.distance, + **self._dml_collection_context(), **kwargs, ) diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py index f9b4c1a8..403b0cb1 100644 --- a/src/pyseekdb/client/validators.py +++ b/src/pyseekdb/client/validators.py @@ -111,3 +111,40 @@ def _validate_record_ids(ids: list[str]) -> None: f"Invalid record id: '{rid}'. Record id contains invalid characters. " "Only letters, digits, and underscore are allowed: [a-zA-Z0-9_]" ) + + +def _validate_namespace_explicit_embedding_dimensions( + embeddings: list[list[float]], + *, + expected_dimension: int, + has_vector_index: bool, +) -> None: + """Reject explicit embeddings whose dimension does not match the collection VECTOR column.""" + for i, vec in enumerate(embeddings): + if vec is None: + continue + actual = len(vec) + if actual != expected_dimension: + if has_vector_index: + raise ValueError( + f"Embedding dimension mismatch: expected {expected_dimension}, " + f"got {actual} at index {i}." + ) + raise ValueError( + f"Collections without a vector index store embeddings as " + f"VECTOR({expected_dimension}); explicit embeddings must be " + f"{expected_dimension}-dimensional, got {actual} at index {i}." + ) + + +def _validate_namespace_no_index_explicit_embeddings( + embeddings: list[list[float]], + *, + expected_dimension: int, +) -> None: + """Reject explicit embeddings whose dimension does not match a no-index VECTOR column.""" + _validate_namespace_explicit_embedding_dimensions( + embeddings, + expected_dimension=expected_dimension, + has_vector_index=False, + ) diff --git a/tests/integration_tests/test_namespace_optional_indexes.py b/tests/integration_tests/test_namespace_optional_indexes.py index cdcebeb0..d9e4c1f2 100644 --- a/tests/integration_tests/test_namespace_optional_indexes.py +++ b/tests/integration_tests/test_namespace_optional_indexes.py @@ -144,6 +144,9 @@ def test_vector_only_indexes_and_queries(self, db_client): result = ns.query(query_embeddings=_unit_vector(3), n_results=3) assert result["ids"] and result["ids"][0] + with pytest.raises(ValueError, match="Embedding dimension mismatch: expected 3"): + ns.add(ids="bad_dim", embeddings=[1.0, 2.0]) + knn = ns.hybrid_search( knn={"query_embeddings": _unit_vector(3), "n_results": 3}, n_results=3, @@ -199,6 +202,23 @@ def test_search_only_indexes_and_queries(self, db_client): finally: db_client.delete_collection(name=collection.name) + def test_search_only_rejects_non_default_explicit_embeddings(self, db_client): + """Without VECTOR INDEX, explicit embeddings must match VECTOR(384).""" + collection = _create_collection(db_client, "search_dim", _schema_search_only()) + try: + assert collection.dimension == 384 + assert collection.has_vector_index is False + ns = collection.create_namespace("ns") + + with pytest.raises(ValueError, match="384-dimensional"): + ns.add(ids="bad", embeddings=[1.0, 2.0, 3.0]) + + ns.add(ids="ok", embeddings=[0.0] * 384) + got = ns.get(ids="ok", include=["embeddings"]) + assert len(got["embeddings"][0]) == 384 + finally: + db_client.delete_collection(name=collection.name) + if __name__ == "__main__": pytest.main([__file__, "-v", "-s", "-k", "oceanbase"]) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index ead00dd7..a970a0f4 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -585,6 +585,65 @@ def test_add_single_sql(self): assert '\\"id\\": \\"d1\\"' in sql or '"id": "d1"' in sql assert '\\"tag\\": \\"a\\"' in sql or '"tag": "a"' in sql + def test_add_warns_when_embeddings_and_documents_with_embedding_function(self, caplog): + """Warn when explicit embeddings override embedding_function.""" + import logging + from unittest.mock import MagicMock + + caplog.set_level(logging.WARNING) + c = self._client() + c._namespace_add( + **self._common_kwargs(), + ids="d1", + embeddings=[1.0, 2.0, 3.0], + documents="hello", + embedding_function=MagicMock(), + has_vector_index=True, + collection_dimension=3, + ) + assert any( + "explicit embeddings" in r.message and "embedding_function" in r.message + for r in caplog.records + ) + + def test_add_rejects_explicit_embeddings_wrong_dim_without_vector_index(self): + """Without VECTOR INDEX, explicit embeddings must match collection dimension.""" + c = self._client() + with pytest.raises(ValueError, match="384-dimensional"): + c._namespace_add( + **self._common_kwargs(), + ids="d1", + embeddings=[1.0, 2.0, 3.0], + has_vector_index=False, + collection_dimension=384, + ) + assert not c.executed_sqls + + def test_add_accepts_384_explicit_embeddings_without_vector_index(self): + """Explicit 384-dim embeddings are allowed when no vector index is configured.""" + c = self._client() + c._namespace_add( + **self._common_kwargs(), + ids="d1", + embeddings=[0.0] * 384, + has_vector_index=False, + collection_dimension=384, + ) + assert c.executed_sqls + + def test_add_rejects_explicit_embeddings_wrong_dim_with_vector_index(self): + """With VECTOR INDEX, explicit embeddings must match collection dimension.""" + c = self._client() + with pytest.raises(ValueError, match="Embedding dimension mismatch: expected 3"): + c._namespace_add( + **self._common_kwargs(), + ids="d1", + embeddings=[1.0, 2.0], + has_vector_index=True, + collection_dimension=3, + ) + assert not c.executed_sqls + def test_add_batch_sql(self): """Test add batch sql.""" c = self._client() diff --git a/tests/unit_tests/test_namespace_no_index_embedding_dim.py b/tests/unit_tests/test_namespace_no_index_embedding_dim.py new file mode 100644 index 00000000..502c5f6e --- /dev/null +++ b/tests/unit_tests/test_namespace_no_index_embedding_dim.py @@ -0,0 +1,47 @@ +"""Unit tests for namespace explicit embedding dimension validation.""" + +import pytest + +from pyseekdb.client.validators import ( + _validate_namespace_explicit_embedding_dimensions, + _validate_namespace_no_index_explicit_embeddings, +) + + +class TestNoIndexExplicitEmbeddingDimension: + def test_accepts_matching_dimension(self): + _validate_namespace_no_index_explicit_embeddings( + [[0.0] * 384, [1.0] * 384], + expected_dimension=384, + ) + + def test_rejects_wrong_dimension(self): + with pytest.raises(ValueError, match="384-dimensional"): + _validate_namespace_no_index_explicit_embeddings( + [[1.0, 2.0, 3.0]], + expected_dimension=384, + ) + + def test_reports_index_in_error(self): + with pytest.raises(ValueError, match="index 1"): + _validate_namespace_no_index_explicit_embeddings( + [[0.0] * 384, [1.0, 2.0]], + expected_dimension=384, + ) + + +class TestIvfExplicitEmbeddingDimension: + def test_rejects_wrong_dimension(self): + with pytest.raises(ValueError, match="Embedding dimension mismatch: expected 3"): + _validate_namespace_explicit_embedding_dimensions( + [[1.0, 2.0]], + expected_dimension=3, + has_vector_index=True, + ) + + def test_accepts_matching_dimension(self): + _validate_namespace_explicit_embedding_dimensions( + [[1.0, 0.0, 0.0]], + expected_dimension=3, + has_vector_index=True, + ) From e06c500922363bec20d32f12284c49ae35df3b7a Mon Sep 17 00:00:00 2001 From: lc533134 Date: Fri, 26 Jun 2026 22:35:30 +0800 Subject: [PATCH 71/84] issue_fix --- src/pyseekdb/client/client_base.py | 67 ++++++++++++++++--- src/pyseekdb/client/document_query_builder.py | 26 +++++-- .../test_namespace_lifecycle_concurrency.py | 6 +- .../unit_tests/test_document_query_builder.py | 14 ++++ .../test_namespace_no_index_embedding_dim.py | 13 ++++ 5 files changed, 110 insertions(+), 16 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 3c36b4cb..205c30ca 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -234,20 +234,33 @@ def _validate_collection_name(name: str) -> None: _NS_DATA_CONTENT_ID_EXPR = "JSON_UNQUOTE(JSON_EXTRACT(data_content, '$.id'))" -def _build_default_ltable_schema() -> dict: - """Return the default logical-table schema (metadata/content/embedding columns and indexes).""" +def _build_default_ltable_schema( + *, + has_fulltext: bool = False, + has_ivf: bool = False, +) -> dict: + """Return the logical-table schema matching the indexes actually provisioned.""" + index_info: list[dict[str, Any]] = [ + {"index_seq": 0, "index_type": "PRIMARY", "indexed_columns": []}, + {"index_seq": 1, "index_type": "SEARCH_INDEX", "indexed_columns": [1]}, + ] + next_seq = 2 + if has_fulltext: + index_info.append( + {"index_seq": next_seq, "index_type": "FULLTEXT", "indexed_columns": [2]} + ) + next_seq += 1 + if has_ivf: + index_info.append( + {"index_seq": next_seq, "index_type": "IVF", "indexed_columns": [3]} + ) return { "col_info": [ - {"col_idx": 1, "col_name": "metadata", "col_type": "JSON"}, - {"col_idx": 2, "col_name": "content", "col_type": "TEXT"}, + {"col_idx": 1, "col_name": "metadata", "col_type": "JSON"}, + {"col_idx": 2, "col_name": "content", "col_type": "TEXT"}, {"col_idx": 3, "col_name": "embedding", "col_type": "VECTOR"}, ], - "index_info": [ - {"index_seq": 0, "index_type": "PRIMARY", "indexed_columns": []}, - {"index_seq": 1, "index_type": "SEARCH_INDEX", "indexed_columns": [1]}, - {"index_seq": 2, "index_type": "FULLTEXT", "indexed_columns": [2]}, - {"index_seq": 3, "index_type": "IVF", "indexed_columns": [3]}, - ], + "index_info": index_info, } @@ -1129,6 +1142,8 @@ def _create_namespace_collection(self, name: str, schema: Schema, partition_coun settings["dense_index_type"] = "ivf" if ivf_config.centroids_fresh_mode is not None: settings["centroids_fresh_mode"] = ivf_config.centroids_fresh_mode + if schema.fulltext_index is not None: + settings["has_fulltext_index"] = True if dense_embedding_function is not None and EmbeddingFunction.support_persistence(dense_embedding_function): settings["embedding_function"] = { "name": dense_embedding_function.name(), @@ -1815,6 +1830,33 @@ def _insert_ns_ltable_catalog_row( raise return self._fetch_ns_ltable_id(collection_id, namespace_id, ltable_name) + def _resolve_ns_ltable_index_layout(self, collection_id: str) -> tuple[bool, bool]: + """Infer which optional indexes exist for a namespace collection.""" + collection_id_escaped = escape_string(collection_id) + rows = self._execute( + f"SELECT settings FROM {self._qtable(CollectionNames.sdk_collections_table_name())} " + f"WHERE collection_id = '{collection_id_escaped}'" + ) + settings: dict[str, Any] = {} + if rows: + raw = rows[0]["settings"] if isinstance(rows[0], dict) else rows[0][0] + settings = json.loads(raw) if raw else {} + has_ivf = settings.get("dense_index_type") == "ivf" + if "has_fulltext_index" in settings: + has_fulltext = bool(settings["has_fulltext_index"]) + else: + data_table = NamespaceCollectionNames.data_table_name(collection_id) + if self._table_exists(data_table): + index_rows = self._execute(f"SHOW INDEX FROM `{data_table}`") + index_names = { + (row.get("Key_name") if isinstance(row, dict) else row[2]) + for row in (index_rows or []) + } + has_fulltext = "idx_fts" in index_names + else: + has_fulltext = False + return has_fulltext, has_ivf + def _finalize_ns_namespace_meta( self, collection_id: str, @@ -1827,7 +1869,10 @@ def _finalize_ns_namespace_meta( schema_table = self._qtable( NamespaceCollectionNames.logic_schema_table_name(collection_id) ) - schema_content = json.dumps(_build_default_ltable_schema()) + has_fulltext, has_ivf = self._resolve_ns_ltable_index_layout(collection_id) + schema_content = json.dumps( + _build_default_ltable_schema(has_fulltext=has_fulltext, has_ivf=has_ivf) + ) with contextlib.suppress(Exception): self._execute( f"INSERT INTO {schema_table} (namespace_id, ltable_id, schema_content) " diff --git a/src/pyseekdb/client/document_query_builder.py b/src/pyseekdb/client/document_query_builder.py index 712f76d5..68a3d416 100644 --- a/src/pyseekdb/client/document_query_builder.py +++ b/src/pyseekdb/client/document_query_builder.py @@ -13,8 +13,14 @@ def _escape_query_string_term(text: str) -> str: """Escape Lucene ``query_string`` syntax characters in a user term.""" - escaped = text.replace("\\", "\\\\") - return _QUERY_STRING_RESERVED_RE.sub(r"\\\1", escaped) + return _QUERY_STRING_RESERVED_RE.sub(r"\\\1", text) + + +def _is_atomic_query_string_term(text: str) -> bool: + """Return whether *text* is a single token safe for fast-path ``$and``/``$or`` joining.""" + if not text or re.search(r"\s", text): + return False + return _QUERY_STRING_RESERVED_RE.search(text) is None def _query_string_contains(query: str, *, boost: float | None = None) -> dict[str, Any]: @@ -147,7 +153,13 @@ def _combine_document_bool( ) -> dict[str, Any] | None: """Combine child ``where_document`` clauses with AND or OR.""" if combiner == "and": - if all(isinstance(c, dict) and "$contains" in c and isinstance(c["$contains"], str) for c in conditions): + if all( + isinstance(c, dict) + and "$contains" in c + and isinstance(c["$contains"], str) + and _is_atomic_query_string_term(c["$contains"]) + for c in conditions + ): queries = [c["$contains"] for c in conditions if isinstance(c, dict)] body: dict[str, Any] = { "fields": [_DOCUMENT_FIELD], @@ -171,7 +183,13 @@ def _combine_document_bool( } if combiner == "or": - if all(isinstance(c, dict) and "$contains" in c and isinstance(c["$contains"], str) for c in conditions): + if all( + isinstance(c, dict) + and "$contains" in c + and isinstance(c["$contains"], str) + and _is_atomic_query_string_term(c["$contains"]) + for c in conditions + ): queries = [c["$contains"] for c in conditions if isinstance(c, dict)] body: dict[str, Any] = { "fields": [_DOCUMENT_FIELD], diff --git a/tests/integration_tests/test_namespace_lifecycle_concurrency.py b/tests/integration_tests/test_namespace_lifecycle_concurrency.py index e50df851..67a8ff81 100644 --- a/tests/integration_tests/test_namespace_lifecycle_concurrency.py +++ b/tests/integration_tests/test_namespace_lifecycle_concurrency.py @@ -100,6 +100,7 @@ def test_multi_client_has_does_not_error_during_drop(oceanbase_client): drop_can_finish = threading.Event() errors: list[Exception] = [] has_results: list[bool] = [] + has_completed_while_drop_blocked = threading.Event() try: collection.create_namespace(ns_name) @@ -128,6 +129,7 @@ def _has_worker() -> None: assert drop_started.wait(timeout=30), "drop did not start" for _ in range(10): has_results.append(checker.has_namespace(ns_name)) + has_completed_while_drop_blocked.set() time.sleep(0.05) except Exception as exc: # noqa: BLE001 errors.append(exc) @@ -137,12 +139,14 @@ def _has_worker() -> None: drop_thread.start() has_thread.start() + observed_during_blocked_drop = has_completed_while_drop_blocked.wait(timeout=30) drop_can_finish.set() drop_thread.join(timeout=120) has_thread.join(timeout=120) assert errors == [] - assert has_results + assert observed_during_blocked_drop, "has_namespace did not complete while drop was blocked" + assert has_results and has_results[0] is True assert checker.has_namespace(ns_name) is False finally: for client in clients: diff --git a/tests/unit_tests/test_document_query_builder.py b/tests/unit_tests/test_document_query_builder.py index 2ceb96df..72f33ea3 100644 --- a/tests/unit_tests/test_document_query_builder.py +++ b/tests/unit_tests/test_document_query_builder.py @@ -26,6 +26,20 @@ def test_not_contains_builds_must_not(self): } } + def test_and_multi_word_uses_bool_composition(self): + """Multi-word $contains terms must not use the single query_string fast path.""" + expr = build_document_hybrid_expression({ + "$and": [{"$contains": "foo bar"}, {"$contains": "baz"}], + }) + assert expr == { + "bool": { + "must": [ + {"query_string": {"fields": ["document"], "query": "foo bar"}}, + {"query_string": {"fields": ["document"], "query": "baz"}}, + ] + } + } + def test_and_or_nested(self): """Test and or nested.""" expr = build_document_hybrid_expression({ diff --git a/tests/unit_tests/test_namespace_no_index_embedding_dim.py b/tests/unit_tests/test_namespace_no_index_embedding_dim.py index 502c5f6e..1b638490 100644 --- a/tests/unit_tests/test_namespace_no_index_embedding_dim.py +++ b/tests/unit_tests/test_namespace_no_index_embedding_dim.py @@ -9,6 +9,12 @@ class TestNoIndexExplicitEmbeddingDimension: + def test_ignores_none_embeddings(self): + _validate_namespace_no_index_explicit_embeddings( + [None, [0.0] * 384], + expected_dimension=384, + ) + def test_accepts_matching_dimension(self): _validate_namespace_no_index_explicit_embeddings( [[0.0] * 384, [1.0] * 384], @@ -31,6 +37,13 @@ def test_reports_index_in_error(self): class TestIvfExplicitEmbeddingDimension: + def test_ignores_none_embeddings(self): + _validate_namespace_explicit_embedding_dimensions( + [None, [1.0, 0.0, 0.0]], + expected_dimension=3, + has_vector_index=True, + ) + def test_rejects_wrong_dimension(self): with pytest.raises(ValueError, match="Embedding dimension mismatch: expected 3"): _validate_namespace_explicit_embedding_dimensions( From 251eb056f88c8dea59c3efa062b0df2e9e4f181c Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 29 Jun 2026 14:50:26 +0800 Subject: [PATCH 72/84] =?UTF-8?q?ut=E9=80=82=E9=85=8Dcollection=5Fdimensio?= =?UTF-8?q?n?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/unit_tests/test_namespace.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index a970a0f4..0268f72e 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -565,13 +565,21 @@ def _common_kwargs(self): namespace_name="test_ns", ) + def _ivf_kwargs(self): + """Common kwargs for 3-dim IVF namespace SQL generation tests.""" + return { + **self._common_kwargs(), + "has_vector_index": True, + "collection_dimension": 3, + } + # ---- ADD ---- def test_add_single_sql(self): """Test add single sql.""" c = self._client() c._namespace_add( - **self._common_kwargs(), + **self._ivf_kwargs(), ids="d1", embeddings=[1.0, 2.0, 3.0], documents="hello", @@ -593,13 +601,11 @@ def test_add_warns_when_embeddings_and_documents_with_embedding_function(self, c caplog.set_level(logging.WARNING) c = self._client() c._namespace_add( - **self._common_kwargs(), + **self._ivf_kwargs(), ids="d1", embeddings=[1.0, 2.0, 3.0], documents="hello", embedding_function=MagicMock(), - has_vector_index=True, - collection_dimension=3, ) assert any( "explicit embeddings" in r.message and "embedding_function" in r.message @@ -648,7 +654,7 @@ def test_add_batch_sql(self): """Test add batch sql.""" c = self._client() c._namespace_add( - **self._common_kwargs(), + **self._ivf_kwargs(), ids=["d1", "d2", "d3"], embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], documents=["A", "B", "C"], @@ -664,7 +670,7 @@ def test_add_without_metadata_sql(self): """Test add without metadata sql.""" c = self._client() c._namespace_add( - **self._common_kwargs(), + **self._ivf_kwargs(), ids="d1", embeddings=[1.0, 2.0, 3.0], ) @@ -696,7 +702,7 @@ def test_update_embedding_and_document_sql(self): """Test update embedding and document sql.""" c = self._client() c._namespace_update( - **self._common_kwargs(), + **self._ivf_kwargs(), ids="d1", embeddings=[9.0, 8.0, 7.0], documents="Updated", @@ -758,7 +764,7 @@ def test_query_basic_dsl(self): c = self._client() c.query_return_value = [] c._namespace_query( - **self._common_kwargs(), + **self._ivf_kwargs(), query_embeddings=[1.0, 0.0, 0.0], n_results=3, distance="l2", @@ -776,7 +782,7 @@ def test_query_with_where_dsl(self): c = self._client() c.query_return_value = [] c._namespace_query( - **self._common_kwargs(), + **self._ivf_kwargs(), query_embeddings=[1.0, 0.0, 0.0], n_results=5, where={"category": "AI"}, @@ -1135,6 +1141,8 @@ def _common_kwargs(self): collection_name="test_coll", namespace_id="7", namespace_name="test_ns", + has_vector_index=True, + collection_dimension=3, ) def test_add_single_record_ok(self): From 7abf5c7a8483c67e6b98354034fa0d9945146ed7 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 29 Jun 2026 15:16:05 +0800 Subject: [PATCH 73/84] =?UTF-8?q?=E9=9D=9Enamespace=E8=B7=AF=E5=BE=84?= =?UTF-8?q?=E4=B8=8D=E5=8E=BB=E6=A0=A1=E9=AA=8C=5Fvalidate=5Finclude?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/collection.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 2502bbae..3e8b4ba6 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -12,7 +12,6 @@ from .validators import ( _MAX_N_RESULTS, - _validate_include, _validate_namespace_name, _validate_n_results, ) @@ -508,7 +507,6 @@ def query( """ self._guard_collection_data_api() _validate_n_results(n_results) - _validate_include(include) return self._client._collection_query( collection_id=self._id, collection_name=self._name, @@ -591,7 +589,6 @@ def get( ) """ self._guard_collection_data_api() - _validate_include(include) return self._client._collection_get( collection_id=self._id, collection_name=self._name, @@ -681,7 +678,6 @@ def hybrid_search( """ self._guard_collection_data_api() _validate_n_results(n_results) - _validate_include(include) # When no query/knn provided, return only ids/distances by default if include is None and not query and not knn: include = [] From ec803f2b2efb8a5a4089158056c75840d9bb0360 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 29 Jun 2026 15:57:13 +0800 Subject: [PATCH 74/84] dont run namespace testcases on workflow --- .github/workflows/ci.yml | 10 +- pyproject.toml | 5 +- src/pyseekdb/client/client_base.py | 397 ++++++++---------- src/pyseekdb/client/client_seekdb_server.py | 5 +- src/pyseekdb/client/collection.py | 11 +- src/pyseekdb/client/configuration.py | 22 +- src/pyseekdb/client/document_query_builder.py | 56 +-- src/pyseekdb/client/filters.py | 4 +- src/pyseekdb/client/kernel_errors.py | 16 +- src/pyseekdb/client/namespace.py | 9 +- src/pyseekdb/client/schema.py | 8 +- src/pyseekdb/client/validators.py | 48 +-- tests/integration_tests/conftest.py | 2 - .../namespace_dml_helpers.py | 27 +- .../namespace_fts_helpers.py | 59 +-- .../namespace_hybrid_search_helpers.py | 255 ++++------- .../test_collection_embedding_function.py | 5 +- ...st_get_or_create_collection_concurrency.py | 9 +- ...t_get_or_create_collection_multiprocess.py | 11 +- .../test_logic_table_monitoring_p1.py | 6 +- .../test_namespace_constraints.py | 36 +- tests/integration_tests/test_namespace_dml.py | 46 +- .../test_namespace_dml_multi_coll.py | 11 +- .../test_namespace_dml_multi_coll_multi_ns.py | 3 +- .../test_namespace_dml_multi_ns.py | 3 +- .../test_namespace_drop_validation.py | 105 ++--- .../test_namespace_get_delete_filters.py | 5 +- .../test_namespace_hybrid_search_combined.py | 10 +- .../test_namespace_hybrid_search_fulltext.py | 7 +- ...rid_search_fulltext_multi_coll_multi_ns.py | 4 +- ...space_hybrid_search_multi_coll_multi_ns.py | 28 +- ...st_namespace_hybrid_search_search_index.py | 5 +- ...t_namespace_hybrid_search_triple_branch.py | 11 +- ...earch_triple_branch_multi_coll_multi_ns.py | 50 +-- .../test_namespace_hybrid_search_vector.py | 20 +- .../test_namespace_lifecycle.py | 58 ++- .../test_namespace_lifecycle_concurrency.py | 10 +- .../test_namespace_optional_indexes.py | 10 +- .../test_namespace_prewarm.py | 6 +- .../test_namespace_prewarm_lob.py | 36 +- .../integration_tests/test_namespace_query.py | 11 +- .../test_namespace_query_where_document.py | 163 ++++--- .../test_namespace_reconnect.py | 6 +- .../test_namespace_session_vars.py | 12 +- .../test_namespace_upsert_concurrency.py | 14 +- tests/unit_tests/test_build_document_query.py | 4 +- tests/unit_tests/test_build_knn_filter.py | 13 +- .../unit_tests/test_build_query_expression.py | 2 + .../unit_tests/test_document_query_builder.py | 2 - ...st_get_or_create_collection_concurrency.py | 26 +- ...est_get_or_create_namespace_concurrency.py | 28 +- tests/unit_tests/test_namespace.py | 187 +++++---- .../test_namespace_kernel_errors.py | 59 +-- .../test_namespace_lifecycle_lock.py | 1 + .../test_namespace_upsert_reconcile.py | 39 +- ...sentence_transformer_embedding_function.py | 2 +- 56 files changed, 869 insertions(+), 1129 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 374d4dcd..31bc08d5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -95,7 +95,10 @@ jobs: log_level=INFO fi set -o pipefail - uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} -k "not server and not embedded and not oceanbase" | tee pytest.log + uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} \ + --ignore-glob='test_namespace*.py' \ + --ignore-glob='test_logic_table_monitoring_p1.py' \ + -k "not server and not embedded and not oceanbase" | tee pytest.log tail -n 1 pytest.log | grep '=======' | grep 'passed' |grep -q 'failed' && exit 1 || exit 0 integration-test: @@ -158,5 +161,8 @@ jobs: log_level=INFO fi set -o pipefail - uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} -k "${{ matrix.test_mode }}" | tee pytest.log + uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} \ + --ignore-glob='test_namespace*.py' \ + --ignore-glob='test_logic_table_monitoring_p1.py' \ + -k "${{ matrix.test_mode }}" | tee pytest.log tail -n 1 pytest.log | grep '=======' | grep 'passed' | grep -q 'failed' && exit 1 || exit 0 diff --git a/pyproject.toml b/pyproject.toml index 6be2ee99..01fc3a94 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -107,8 +107,9 @@ ignore = [ ] [tool.ruff.lint.per-file-ignores] -"tests/*" = ["S101", "S608"] -"src/pyseekdb/client/client_base.py" = ["S608"] +"tests/**" = ["S101", "S608", "C901", "S105", "S110", "TRY300", "B017"] +"src/pyseekdb/client/client_base.py" = ["S608", "C901"] +"src/pyseekdb/client/document_query_builder.py" = ["C901"] [tool.ruff.format] preview = true diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 205c30ca..d7e204c2 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -17,13 +17,6 @@ from pymysql.converters import escape_string -from .document_query_builder import ( - build_document_hybrid_expression, - doc_matches_where_document, - document_expr_as_knn_filter, - where_document_knn_prefilterable, -) -from .kernel_errors import maybe_reraise_friendly_kernel_error, namespace_kernel_error_guard from .admin_client import DEFAULT_TENANT, AdminAPI from .base_connection import BaseConnection from .collection import Collection @@ -42,6 +35,12 @@ VectorIndexConfig, ) from .database import Database +from .document_query_builder import ( + build_document_hybrid_expression, + doc_matches_where_document, + document_expr_as_knn_filter, + where_document_knn_prefilterable, +) from .embedding_function import ( Documents as EmbeddingDocuments, ) @@ -51,6 +50,7 @@ get_default_embedding_function, ) from .filters import FilterBuilder +from .kernel_errors import maybe_reraise_friendly_kernel_error, namespace_kernel_error_guard from .meta_info import CollectionFieldNames, CollectionNames, NamespaceCollectionNames, NamespaceFieldNames from .query_types import QueryHint from .schema import Schema, SparseVectorIndexConfig @@ -62,6 +62,14 @@ ) from .sql_utils import _query_hint_to_sql, is_query_sql from .types import K as FieldKey +from .validators import ( + _MAX_N_RESULTS, + _MAX_NAMESPACE_BATCH_SIZE, + _quote_sql_identifier, + _validate_database_name, + _validate_namespace_explicit_embedding_dimensions, + _validate_record_ids, +) from .version import Version # Type alias for embedding_function parameter that can be EmbeddingFunction, None, or sentinel @@ -109,13 +117,8 @@ def _is_namespace_catalog_conflict_error(exc: BaseException) -> bool: current: BaseException | None = exc while current is not None: message = str(current).lower() - if ( - "duplicate entry" in message - and ( - "uk_sdk_ns_coll_name" in message - or "uk_sdk_lt_coll_ns_name" in message - or "code=1062" in message - ) + if "duplicate entry" in message and ( + "uk_sdk_ns_coll_name" in message or "uk_sdk_lt_coll_ns_name" in message or "code=1062" in message ): return True if type(current).__name__ == "IntegrityError" and "1062" in message: @@ -129,14 +132,11 @@ def _is_sdk_collection_catalog_conflict_error(exc: BaseException) -> bool: current: BaseException | None = exc while current is not None: message = str(current).lower() - if ( - "duplicate entry" in message - and ( - "uk_sdk_coll_name" in message - or "idx_name" in message - or "collection_name" in message - or "code=1062" in message - ) + if "duplicate entry" in message and ( + "uk_sdk_coll_name" in message + or "idx_name" in message + or "collection_name" in message + or "code=1062" in message ): return True if type(current).__name__ == "IntegrityError" and "1062" in message: @@ -217,17 +217,6 @@ def _validate_collection_name(name: str) -> None: ) -from .validators import ( # noqa: F401 - _MAX_N_RESULTS, - _MAX_NAMESPACE_BATCH_SIZE, - _quote_sql_identifier, - _validate_database_name, - _validate_namespace_name, - _validate_namespace_no_index_explicit_embeddings, - _validate_namespace_explicit_embedding_dimensions, - _validate_record_ids, -) - _DEFAULT_PARTITION_COUNT = 1000 # Unquoted id for WHERE/CASE; plain JSON_EXTRACT returns a quoted JSON string and # can route through SEARCH INDEX on SS logic tables, breaking cross-namespace id lookups. @@ -246,14 +235,10 @@ def _build_default_ltable_schema( ] next_seq = 2 if has_fulltext: - index_info.append( - {"index_seq": next_seq, "index_type": "FULLTEXT", "indexed_columns": [2]} - ) + index_info.append({"index_seq": next_seq, "index_type": "FULLTEXT", "indexed_columns": [2]}) next_seq += 1 if has_ivf: - index_info.append( - {"index_seq": next_seq, "index_type": "IVF", "indexed_columns": [3]} - ) + index_info.append({"index_seq": next_seq, "index_type": "IVF", "indexed_columns": [3]}) return { "col_info": [ {"col_idx": 1, "col_name": "metadata", "col_type": "JSON"}, @@ -538,9 +523,7 @@ def _is_shared_storage_mode(self) -> bool: return cached result = False try: - rows = self._execute( - "SELECT VALUE FROM oceanbase.GV$OB_PARAMETERS WHERE name = 'ob_startup_mode'" - ) + rows = self._execute("SELECT VALUE FROM oceanbase.GV$OB_PARAMETERS WHERE name = 'ob_startup_mode'") if rows: val = rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["VALUE"] result = str(val).upper() == "SHARED_STORAGE" @@ -559,7 +542,7 @@ def _stg_cache_policy_clause(self) -> str: return 'STORAGE_CACHE_POLICY = (GLOBAL = "hot")' return "" - def detect_db_type_and_version(self) -> tuple[str, "Version"]: # noqa: C901 + def detect_db_type_and_version(self) -> tuple[str, "Version"]: """ Detect database type and version. @@ -816,7 +799,7 @@ def fork_database(self, source_name: str, destination_name: str, tenant: str = D # ==================== Collection Management (User-facing) ==================== - def _prepare_schema_parameters( # noqa: C901 + def _prepare_schema_parameters( self, configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, @@ -985,13 +968,8 @@ def create_collection( ) # Only fully initialized collections (metadata + physical table) count as existing. # Metadata without a table is treated as an incomplete create and repaired below. - if self.has_collection(name): - # A namespace collection whose catalog row exists but whose physical - # tables are incomplete (e.g. a crash interrupted creation) is allowed - # through so the create can be resumed/finished idempotently. A complete - # collection still errors as "already exists". - if not (use_namespace and self._is_incomplete_ns_collection(name)): - raise ValueError(f"Collection '{name}' already exists") + if self.has_collection(name) and not (use_namespace and self._is_incomplete_ns_collection(name)): + raise ValueError(f"Collection '{name}' already exists") # Resolve schema: either use the provided schema or build one from legacy params if schema is not None: @@ -1028,9 +1006,7 @@ def create_collection( dimension = hnsw_config.dimension if dimension < 1 or dimension > MAX_HNSW_VECTOR_DIMENSION: - raise ValueError( - f"Dimension must be between 1 and {MAX_HNSW_VECTOR_DIMENSION}, got {dimension}" - ) + raise ValueError(f"Dimension must be between 1 and {MAX_HNSW_VECTOR_DIMENSION}, got {dimension}") # Extract fulltext parser configuration fulltext_index_clause = _get_fulltext_index_sql(schema.fulltext_index) @@ -1084,7 +1060,9 @@ def create_collection( **kwargs, ) - def _create_namespace_collection(self, name: str, schema: Schema, partition_count: int | None = None, **kwargs) -> "Collection": + def _create_namespace_collection( + self, name: str, schema: Schema, partition_count: int | None = None, **kwargs + ) -> "Collection": """Create a namespace-enabled collection and its catalog/physical tables.""" dense_embedding_function = schema.vector_index.embedding_function ivf_config = schema.vector_index.ivf @@ -1237,9 +1215,7 @@ def _create_sdk_collections_if_not_exists(self) -> None: ) COMMENT='Settings of collections created by SDK' ORGANIZATION INDEX {scp};""" self._execute(create_table_sql) try: - self._execute( - f"CREATE UNIQUE INDEX uk_sdk_coll_name ON {sdk_coll} (collection_name)" - ) + self._execute(f"CREATE UNIQUE INDEX uk_sdk_coll_name ON {sdk_coll} (collection_name)") except Exception as exc: _reraise_unless_unique_index_exists(exc) except Exception as e: @@ -1380,8 +1356,7 @@ def _ensure_namespace_catalogs(self) -> None: self._execute(namespaces_stats_sql) try: self._execute( - f"CREATE UNIQUE INDEX uk_sdk_ns_coll_name ON {ns_namespaces_q} " - f"(collection_id, namespace_name)" + f"CREATE UNIQUE INDEX uk_sdk_ns_coll_name ON {ns_namespaces_q} (collection_id, namespace_name)" ) except Exception as exc: _reraise_unless_unique_index_exists(exc) @@ -1407,9 +1382,7 @@ def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> di collection_name_escaped = escape_string(collection_name) sdk_coll = self._qtable(CollectionNames.sdk_collections_table_name()) insert_sql = ( - f"INSERT INTO {sdk_coll} " - f"(collection_name, settings) " - f"VALUES ('{collection_name_escaped}', '{settings_str}')" + f"INSERT INTO {sdk_coll} (collection_name, settings) VALUES ('{collection_name_escaped}', '{settings_str}')" ) try: self._execute(insert_sql) @@ -1421,8 +1394,7 @@ def _create_ns_collection_meta(self, collection_name: str, settings: dict) -> di with contextlib.suppress(Exception): conn_getter().rollback() rows = self._execute( - f"SELECT collection_id FROM {sdk_coll} " - f"WHERE collection_name = '{collection_name_escaped}'" + f"SELECT collection_id FROM {sdk_coll} WHERE collection_name = '{collection_name_escaped}'" ) collection_id = str(rows[0][0] if isinstance(rows[0], (list, tuple)) else rows[0]["collection_id"]) self._set_session_ns_context(collection_id=collection_id) @@ -1595,7 +1567,7 @@ def _cleanup_namespace_physical_tables(self, collection_id: str) -> None: NamespaceCollectionNames.data_table_name, NamespaceCollectionNames.logic_schema_table_name, NamespaceCollectionNames.kv_data_table_name, - NamespaceCollectionNames.hot_table_name + NamespaceCollectionNames.hot_table_name, ]: with contextlib.suppress(Exception): self._execute(f"DROP TABLE IF EXISTS `{suffix_fn(collection_id)}`") @@ -1618,15 +1590,11 @@ def _tablegroup_exists(self, tablegroup_name: str) -> bool: """Whether `tablegroup_name` exists in the current OceanBase tenant.""" name_escaped = escape_string(tablegroup_name) try: - rows = self._execute( - "SELECT 1 FROM oceanbase.DBA_OB_TABLEGROUPS " - f"WHERE TABLEGROUP_NAME = '{name_escaped}'" - ) + rows = self._execute(f"SELECT 1 FROM oceanbase.DBA_OB_TABLEGROUPS WHERE TABLEGROUP_NAME = '{name_escaped}'") return bool(rows) except Exception: return False - def _ns_missing_physical_resources(self, collection_id: str, is_shared_storage: bool) -> list[str]: """Return expected tablegroup/tables that are absent for this namespace collection.""" resources: list[tuple[str, bool]] = [ @@ -1639,11 +1607,7 @@ def _ns_missing_physical_resources(self, collection_id: str, is_shared_storage: resources.append((NamespaceCollectionNames.hot_table_name(collection_id), False)) missing: list[str] = [] for resource_name, is_tablegroup in resources: - exists = ( - self._tablegroup_exists(resource_name) - if is_tablegroup - else self._table_exists(resource_name) - ) + exists = self._tablegroup_exists(resource_name) if is_tablegroup else self._table_exists(resource_name) if not exists: missing.append(resource_name) return missing @@ -1660,9 +1624,7 @@ def _is_incomplete_ns_collection(self, name: str) -> bool: self._use_catalog_database() return len(self._ns_missing_physical_resources(meta["collection_id"], is_ss)) > 0 - def _purge_broken_ns_collection_if_incomplete( - self, collection_name: str, meta: dict | None = None - ) -> bool: + def _purge_broken_ns_collection_if_incomplete(self, collection_name: str, meta: dict | None = None) -> bool: """Purge a namespace collection whose catalog row exists but physical resources are incomplete. Returns True when the collection was purged and should be treated as non-existent. @@ -1686,9 +1648,7 @@ def _purge_broken_ns_collection_if_incomplete( self._delete_ns_collection_meta(collection_name) return True - def _resolve_namespace_ltable_id( - self, collection_id: str, namespace_id: int - ) -> int: + def _resolve_namespace_ltable_id(self, collection_id: str, namespace_id: int) -> int: """Resolve the default ltable_id for (collection_id, namespace_id) from sdk_ltables. Cached per (collection_id, namespace_id) on the client instance to avoid the extra round-trip on every DML/DQL call. @@ -1712,17 +1672,14 @@ def _resolve_namespace_ltable_id( ) if not rows: raise ValueError( - f"No default ltable found for collection_id={collection_id}, " - f"namespace_id={namespace_id} in sdk_ltables" + f"No default ltable found for collection_id={collection_id}, namespace_id={namespace_id} in sdk_ltables" ) row = rows[0] lt_id = int(row[0] if isinstance(row, (list, tuple)) else row["ltable_id"]) cache[key] = lt_id return lt_id - def _cache_namespace_ltable_id( - self, collection_id: str, namespace_id: int, ltable_id: int - ) -> None: + def _cache_namespace_ltable_id(self, collection_id: str, namespace_id: int, ltable_id: int) -> None: """Cache the resolved logical-table id for a namespace.""" key = (str(collection_id), int(namespace_id)) cache = getattr(self, "_ns_ltable_id_cache", None) @@ -1848,10 +1805,7 @@ def _resolve_ns_ltable_index_layout(self, collection_id: str) -> tuple[bool, boo data_table = NamespaceCollectionNames.data_table_name(collection_id) if self._table_exists(data_table): index_rows = self._execute(f"SHOW INDEX FROM `{data_table}`") - index_names = { - (row.get("Key_name") if isinstance(row, dict) else row[2]) - for row in (index_rows or []) - } + index_names = {(row.get("Key_name") if isinstance(row, dict) else row[2]) for row in (index_rows or [])} has_fulltext = "idx_fts" in index_names else: has_fulltext = False @@ -1866,13 +1820,9 @@ def _finalize_ns_namespace_meta( ) -> dict: """Finalize namespace metadata after catalog rows and physical tables are created.""" self._cache_namespace_ltable_id(collection_id, namespace_id, ltable_id) - schema_table = self._qtable( - NamespaceCollectionNames.logic_schema_table_name(collection_id) - ) + schema_table = self._qtable(NamespaceCollectionNames.logic_schema_table_name(collection_id)) has_fulltext, has_ivf = self._resolve_ns_ltable_index_layout(collection_id) - schema_content = json.dumps( - _build_default_ltable_schema(has_fulltext=has_fulltext, has_ivf=has_ivf) - ) + schema_content = json.dumps(_build_default_ltable_schema(has_fulltext=has_fulltext, has_ivf=has_ivf)) with contextlib.suppress(Exception): self._execute( f"INSERT INTO {schema_table} (namespace_id, ltable_id, schema_content) " @@ -1886,18 +1836,12 @@ def _create_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> if self._get_ns_namespace_meta(collection_id, namespace_name) is not None: raise ValueError(f"Namespace '{namespace_name}' already exists") try: - ns_id = self._insert_ns_namespace_catalog_row( - collection_id, namespace_name, idempotent=False - ) - lt_id = self._insert_ns_ltable_catalog_row( - collection_id, ns_id, idempotent=False - ) + ns_id = self._insert_ns_namespace_catalog_row(collection_id, namespace_name, idempotent=False) + lt_id = self._insert_ns_ltable_catalog_row(collection_id, ns_id, idempotent=False) except Exception as exc: if _is_namespace_catalog_conflict_error(exc): self._rollback_connection_if_supported() - raise ValueError( - f"Namespace '{namespace_name}' already exists" - ) from exc + raise ValueError(f"Namespace '{namespace_name}' already exists") from exc raise return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) @@ -1907,18 +1851,10 @@ def _get_or_create_ns_namespace_meta(self, collection_id: str, namespace_name: s if meta is not None: if meta.get("ltable_id") is not None: return meta - lt_id = self._insert_ns_ltable_catalog_row( - collection_id, int(meta["namespace_id"]), idempotent=True - ) - return self._finalize_ns_namespace_meta( - collection_id, namespace_name, int(meta["namespace_id"]), lt_id - ) - ns_id = self._insert_ns_namespace_catalog_row( - collection_id, namespace_name, idempotent=True - ) - lt_id = self._insert_ns_ltable_catalog_row( - collection_id, ns_id, idempotent=True - ) + lt_id = self._insert_ns_ltable_catalog_row(collection_id, int(meta["namespace_id"]), idempotent=True) + return self._finalize_ns_namespace_meta(collection_id, namespace_name, int(meta["namespace_id"]), lt_id) + ns_id = self._insert_ns_namespace_catalog_row(collection_id, namespace_name, idempotent=True) + lt_id = self._insert_ns_ltable_catalog_row(collection_id, ns_id, idempotent=True) return self._finalize_ns_namespace_meta(collection_id, namespace_name, ns_id, lt_id) def _get_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> dict | None: @@ -1991,12 +1927,8 @@ def _delete_ns_namespace_meta(self, collection_id: str, namespace_name: str) -> collection_id_escaped = escape_string(collection_id) # PL reads session database_name; must match where catalog tables live. self._use_catalog_database() - self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(ns_id), ltable_id=lt_id - ) - self._execute( - f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})" - ) + self._set_session_ns_context(collection_id=collection_id, namespace_id=int(ns_id), ltable_id=lt_id) + self._execute(f"CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE('{collection_id_escaped}', {ns_id})") def _list_ns_namespaces(self, collection_id: str) -> list[dict]: """List namespaces registered for a collection.""" @@ -2027,10 +1959,8 @@ def _get_ns_namespace_id(self, collection_id: str, namespace_name: str) -> str: def get_collection(self, name: str, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED) -> "Collection": """Get an existing collection by name.""" ns_meta = None - try: + with contextlib.suppress(Exception): ns_meta = self._get_ns_collection_meta(name) - except Exception: - pass if ns_meta is not None: if self._purge_broken_ns_collection_if_incomplete(collection_name=name, meta=ns_meta): ns_meta = None @@ -2089,7 +2019,7 @@ def _resolve_collection_metadata_from_sdk_collections(self, collection_name: str raise ValueError(f"Failed to resolve collection metadata from sdk_collections table: {e}") from e return None - def _resolve_collection_metadata_from_table(self, table_name: str, collection_name: str) -> dict[str, Any]: # noqa: C901 + def _resolve_collection_metadata_from_table(self, table_name: str, collection_name: str) -> dict[str, Any]: """ Resolve collection metadata infromation from collection table (not sdk_collections table) """ @@ -2397,10 +2327,7 @@ def _list_ns_collections(self) -> list["Collection"]: check_result = self._execute(check_sql) if not check_result: return result - rows = self._execute( - f"SELECT collection_id, collection_name, settings " - f"FROM `{sdk_table}`" - ) + rows = self._execute(f"SELECT collection_id, collection_name, settings FROM `{sdk_table}`") for row in rows: try: if isinstance(row, dict): @@ -2425,7 +2352,7 @@ def _list_ns_collections(self) -> list["Collection"]: except Exception as e: logger.warning(f"Failed to build namespace collection from row: {e}") except Exception: - pass + logger.debug("Failed to list namespace collections from catalog", exc_info=True) return result def _list_collections_v2(self) -> list["Collection"]: @@ -2698,9 +2625,7 @@ def _get_or_resume_existing_collection( """Return an existing collection or resume an incomplete namespace-enabled one.""" if use_namespace: if self._get_ns_collection_meta(name) is None: - raise ValueError( - f"Collection '{name}' conflicted during create but namespace metadata is missing" - ) + raise ValueError(f"Collection '{name}' conflicted during create but namespace metadata is missing") if self._is_incomplete_ns_collection(name): return self.create_collection( name=name, @@ -2822,7 +2747,7 @@ def _collection_fork(self, collection: Collection, forked_name: str) -> None: # -------------------- DML Operations -------------------- - def _generate_sparse_embeddings( # noqa: C901 + def _generate_sparse_embeddings( self, sparse_config: SparseVectorIndexConfig, documents: list[str] | None, @@ -2892,7 +2817,7 @@ def _generate_sparse_embeddings( # noqa: C901 logger.debug(f"✅ Successfully generated {len(sparse_vectors)} sparse embeddings") return sparse_vectors - def _collection_add( # noqa: C901 + def _collection_add( self, collection_id: str | None, collection_name: str, @@ -3076,7 +3001,7 @@ def _collection_add( # noqa: C901 self._execute(sql) logger.debug(f"✅ Successfully added {num_items} item(s) to collection '{collection_name}'") - def _collection_update( # noqa: C901 + def _collection_update( self, collection_id: str | None, collection_name: str, @@ -3239,7 +3164,7 @@ def _collection_update( # noqa: C901 logger.debug(f"✅ Successfully updated {len(ids)} item(s) in collection '{collection_name}'") - def _collection_upsert( # noqa: C901 + def _collection_upsert( self, collection_id: str | None, collection_name: str, @@ -3906,7 +3831,7 @@ def _execute(self, sql: str) -> Any: # -------------------- DQL Operations (Common Implementation) -------------------- - def _collection_query( # noqa: C901 + def _collection_query( self, collection_id: str | None, collection_name: str, @@ -4128,7 +4053,7 @@ def _collection_query( # noqa: C901 ) return result - def _collection_query_sparse( # noqa: C901 + def _collection_query_sparse( self, conn, table_name: str, @@ -4287,7 +4212,7 @@ def _collection_query_sparse( # noqa: C901 ) return result - def _collection_get( # noqa: C901 + def _collection_get( self, collection_id: str | None, collection_name: str, @@ -4539,7 +4464,7 @@ def _collection_hybrid_search( # Transform SQL query results to standard format return self._transform_sql_result(result_rows, include) - def _build_search_parm( # noqa: C901 + def _build_search_parm( self, query: dict[str, Any] | list[dict[str, Any]] | None, knn: dict[str, Any] | list[dict[str, Any]] | None, @@ -4737,7 +4662,7 @@ def _build_query_expression(self, query: dict[str, Any]) -> dict[str, Any] | Non return None - def _build_document_query( # noqa: C901 + def _build_document_query( self, where_document: dict[str, Any], boost: float | None = None ) -> dict[str, Any] | None: """ @@ -4778,9 +4703,7 @@ def _apply_boost(target: Any) -> None: _apply_boost(expr) return expr - return _with_boost( - build_document_hybrid_expression(where_document, boost=boost) - ) + return _with_boost(build_document_hybrid_expression(where_document, boost=boost)) def _build_metadata_filter_for_search_parm(self, where: dict[str, Any] | None) -> list[dict[str, Any]]: """ @@ -4811,9 +4734,7 @@ def _build_search_parm_field_name(self, key: str) -> str: return CollectionFieldNames.ID return f"(JSON_EXTRACT(metadata, '$.{key}'))" - def _build_metadata_filter_conditions( # noqa: C901 - self, condition: dict[str, Any] - ) -> list[dict[str, Any]]: + def _build_metadata_filter_conditions(self, condition: dict[str, Any]) -> list[dict[str, Any]]: """ Recursively build metadata filter conditions from nested dictionary @@ -4861,9 +4782,7 @@ def _build_metadata_filter_conditions( # noqa: C901 # non-scoring (filter) context the kernel does not reliably apply the # implicit "at least one should" default, which otherwise yields an # intermittent `1210 Invalid argument`. - result.append( - {"bool": {"should": should_conditions, "minimum_should_match": 1}} - ) + result.append({"bool": {"should": should_conditions, "minimum_should_match": 1}}) return result if "$not" in condition: @@ -4917,7 +4836,7 @@ def _build_metadata_filter_conditions( # noqa: C901 return result - def _build_knn_expression( # noqa: C901 + def _build_knn_expression( self, knn: dict[str, Any], dimension: int | None = None, **kwargs ) -> dict[str, Any] | list[dict[str, Any]] | None: """ @@ -5023,14 +4942,9 @@ def _normalize_vectors(raw_embeddings: Any) -> list[list[float]]: knn_filter = pos_filters if where_document is not None and where_document_knn_prefilterable(where_document): - doc_filter = document_expr_as_knn_filter( - build_document_hybrid_expression(where_document) - ) + doc_filter = document_expr_as_knn_filter(build_document_hybrid_expression(where_document)) if doc_filter is not None: - if knn_filter is None: - knn_filter = [doc_filter] - else: - knn_filter = [*knn_filter, doc_filter] + knn_filter = [doc_filter] if knn_filter is None else [*knn_filter, doc_filter] for vector in vectors: expr = {"field": "embedding", "k": n_results, "query_vector": vector} @@ -5075,20 +4989,18 @@ def _post_filter_namespace_query_result( if len(kept_indices) >= n_results: break - def _pick(key: str) -> list[Any]: - """Pick kept elements from one result group.""" - groups = result.get(key) - if not groups or qi >= len(groups): - return [] - group = groups[qi] - return [group[i] for i in kept_indices if i < len(group)] - filtered["ids"].append([ids[i] for i in kept_indices]) - if "distances" in result and result["distances"]: - filtered["distances"].append(_pick("distances")) + if result.get("distances"): + dist_groups = result.get("distances") + if dist_groups and qi < len(dist_groups): + group = dist_groups[qi] + filtered["distances"].append([group[i] for i in kept_indices if i < len(group)]) for optional_key in ("documents", "metadatas", "embeddings"): if optional_key in filtered: - filtered[optional_key].append(_pick(optional_key)) + groups = result.get(optional_key) + if groups and qi < len(groups): + group = groups[qi] + filtered[optional_key].append([group[i] for i in kept_indices if i < len(group)]) return filtered @@ -5132,9 +5044,7 @@ def _hybrid_row_score(self, row: dict[str, Any]) -> float: return float(val) return 0.0 - def _transform_sql_result( # noqa: C901 - self, result_rows: list[dict[str, Any]], include: list[str] | None - ) -> dict[str, Any]: + def _transform_sql_result(self, result_rows: list[dict[str, Any]], include: list[str] | None) -> dict[str, Any]: """ Transform SQL query results to standard format (query-compatible format) @@ -5345,11 +5255,7 @@ def _validate_namespace_explicit_embeddings_if_needed( """Validate user-supplied embeddings match the collection VECTOR column dimension.""" if not explicit_embeddings or not embeddings: return - expected = ( - collection_dimension - if collection_dimension is not None - else DEFAULT_VECTOR_DIMENSION - ) + expected = collection_dimension if collection_dimension is not None else DEFAULT_VECTOR_DIMENSION _validate_namespace_explicit_embedding_dimensions( embeddings, expected_dimension=expected, @@ -5449,17 +5355,11 @@ def _reconcile_namespace_duplicate_records( meta_val = metadatas[i] if metadatas and i < len(metadatas) else None emb_val = embeddings[i] if embeddings and i < len(embeddings) else None for attempt in range(120): - duplicate_count = self._count_namespace_records_by_id( - table_name, ns_id, ltable_id, record_id - ) + duplicate_count = self._count_namespace_records_by_id(table_name, ns_id, ltable_id, record_id) if duplicate_count <= 1: break - self._delete_namespace_records_by_id( - table_name, ns_id, ltable_id, record_id - ) - if self._count_namespace_records_by_id( - table_name, ns_id, ltable_id, record_id - ) == 0: + self._delete_namespace_records_by_id(table_name, ns_id, ltable_id, record_id) + if self._count_namespace_records_by_id(table_name, ns_id, ltable_id, record_id) == 0: self._namespace_add( collection_id=collection_id, collection_name=collection_name, @@ -5475,12 +5375,10 @@ def _reconcile_namespace_duplicate_records( if attempt < 119: time.sleep(0.05 * min(attempt + 1, 10)) else: - raise ValueError( - f"Failed to reconcile duplicate namespace rows for record_id={record_id!r}" - ) + raise ValueError(f"Failed to reconcile duplicate namespace rows for record_id={record_id!r}") @namespace_kernel_error_guard - def _namespace_add( # noqa: C901 + def _namespace_add( self, collection_id: str | None, collection_name: str, @@ -5563,7 +5461,9 @@ def _namespace_add( # noqa: C901 ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) @@ -5660,7 +5560,9 @@ def _namespace_update( ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) @@ -5719,7 +5621,7 @@ def _namespace_update( has_doc = True doc_case_parts.append(f"WHEN {id_expr} = %s THEN %s") params.extend([record_id, documents[i]]) - if record_id not in [rid for rid in batch_ids]: + if record_id not in list(batch_ids): batch_ids.append(record_id) if metadatas and i < len(metadatas) and metadatas[i] is not None: has_meta = True @@ -5775,7 +5677,9 @@ def _namespace_upsert( collection_dimension = kwargs.pop("collection_dimension", None) ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) if isinstance(ids, str): ids = [ids] @@ -5830,10 +5734,15 @@ def _namespace_upsert( add_metas = [metadatas[i] for i in add_indices] if metadatas else None add_embs = [embeddings[i] for i in add_indices] if embeddings else None self._namespace_add( - collection_id=collection_id, collection_name=collection_name, - namespace_id=namespace_id, namespace_name=namespace_name, - ids=add_ids, embeddings=add_embs, metadatas=add_metas, - documents=add_docs, embedding_function=embedding_function, + collection_id=collection_id, + collection_name=collection_name, + namespace_id=namespace_id, + namespace_name=namespace_name, + ids=add_ids, + embeddings=add_embs, + metadatas=add_metas, + documents=add_docs, + embedding_function=embedding_function, has_vector_index=has_vector_index, collection_dimension=collection_dimension, **kwargs, @@ -5845,10 +5754,15 @@ def _namespace_upsert( upd_metas = [metadatas[i] for i in update_indices] if metadatas else None upd_embs = [embeddings[i] for i in update_indices] if embeddings else None self._namespace_update( - collection_id=collection_id, collection_name=collection_name, - namespace_id=namespace_id, namespace_name=namespace_name, - ids=upd_ids, embeddings=upd_embs, metadatas=upd_metas, - documents=upd_docs, embedding_function=embedding_function, + collection_id=collection_id, + collection_name=collection_name, + namespace_id=namespace_id, + namespace_name=namespace_name, + ids=upd_ids, + embeddings=upd_embs, + metadatas=upd_metas, + documents=upd_docs, + embedding_function=embedding_function, has_vector_index=has_vector_index, collection_dimension=collection_dimension, **kwargs, @@ -5884,7 +5798,9 @@ def _namespace_delete( """Delete records from a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) if ids is None and where is None and where_document is None: raise ValueError("At least one of ids, where, or where_document must be provided") @@ -5935,7 +5851,7 @@ def _namespace_delete( self._execute(sql) @namespace_kernel_error_guard - def _namespace_query( # noqa: C901 + def _namespace_query( self, collection_id: str | None, collection_name: str, @@ -5983,7 +5899,9 @@ def _namespace_query( # noqa: C901 ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) include_fields = self._normalize_include_fields(include) @@ -5995,10 +5913,7 @@ def _namespace_query( # noqa: C901 all_embeddings: list[list[Any]] = [] all_distances: list[list[float]] = [] - hybrid_kwargs = { - k: v for k, v in kwargs.items() - if k not in ("embedding_function", "distance", "dimension") - } + hybrid_kwargs = {k: v for k, v in kwargs.items() if k not in ("embedding_function", "distance", "dimension")} for query_vector in query_embeddings: knn_cfg: dict[str, Any] = { @@ -6038,7 +5953,9 @@ def _namespace_query( # noqa: C901 if post_filter_wd is not None: batch = self._post_filter_namespace_query_result( - batch, post_filter_wd, n_results=n_results, + batch, + post_filter_wd, + n_results=n_results, ) batch_ids = batch.get("ids") or [[]] @@ -6065,7 +5982,7 @@ def _namespace_query( # noqa: C901 return result @namespace_kernel_error_guard - def _namespace_get( # noqa: C901 + def _namespace_get( self, collection_id: str | None, collection_name: str, @@ -6082,7 +5999,9 @@ def _namespace_get( # noqa: C901 """Fetch records from a namespace collection by id/where/pagination.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) include_fields = self._normalize_include_fields(include) table_name = NamespaceCollectionNames.data_table_name(collection_id) @@ -6162,18 +6081,22 @@ def _namespace_get( # noqa: C901 result_embeddings.append(emb) elif isinstance(row, (list, tuple)): idx = 0 - rid_raw = row[idx]; idx += 1 + rid_raw = row[idx] + idx += 1 rid = json.loads(rid_raw) if isinstance(rid_raw, str) else rid_raw result_ids.append(rid) if "documents" in include_fields or "document" in include_fields or include is None: - result_documents.append(row[idx]); idx += 1 + result_documents.append(row[idx]) + idx += 1 if "metadatas" in include_fields or "metadata" in include_fields or include is None: - meta_raw = row[idx]; idx += 1 + meta_raw = row[idx] + idx += 1 if isinstance(meta_raw, str): meta_raw = json.loads(meta_raw) result_metadatas.append(meta_raw or {}) if "embeddings" in include_fields or "embedding" in include_fields: - emb = row[idx]; idx += 1 + emb = row[idx] + idx += 1 if isinstance(emb, bytes): emb = self._parse_embedding_from_bytes(emb) elif isinstance(emb, str): @@ -6201,7 +6124,9 @@ def _namespace_count( """Count records in a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) @@ -6289,7 +6214,7 @@ def _inject_filter_into_knn(node): if isinstance(existing, list): node["filter"] = existing + ns_filter else: - node["filter"] = [existing] + ns_filter + node["filter"] = [existing, *ns_filter] else: node["filter"] = list(ns_filter) return node @@ -6312,7 +6237,7 @@ def _inject_filter_into_query(node): if isinstance(existing, list): bool_node["filter"] = existing + ns_filter else: - bool_node["filter"] = [existing] + ns_filter + bool_node["filter"] = [existing, *ns_filter] else: bool_node["filter"] = list(ns_filter) return node @@ -6358,14 +6283,19 @@ def _namespace_hybrid_search( """Run a hybrid (vector + fulltext) search against a namespace collection.""" ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) conn = self._ensure_connection() table_name = NamespaceCollectionNames.data_table_name(collection_id) ns_id = int(namespace_id) search_parm = self._build_search_parm( - query, knn, rank, n_results, + query, + knn, + rank, + n_results, include=include, dimension=kwargs.get("dimension"), **{k: v for k, v in kwargs.items() if k != "dimension"}, @@ -6391,14 +6321,16 @@ def _namespace_hybrid_search( hint_sql = _query_hint_to_sql(query_hint, table_name=table_name) or "" hybrid_sql = ( - f"SELECT {hint_sql + ' ' if hint_sql else ''}* " - f"FROM hybrid_search(TABLE `{table_name}`, '{escaped_params}')" + f"SELECT {hint_sql + ' ' if hint_sql else ''}* FROM hybrid_search(TABLE `{table_name}`, '{escaped_params}')" ) result_rows = self._execute_query_with_cursor(conn, hybrid_sql, [], use_context_manager) if not result_rows: return { - "ids": [[]], "distances": [[]], "metadatas": [[]], - "documents": [[]], "embeddings": [[]], + "ids": [[]], + "distances": [[]], + "metadatas": [[]], + "documents": [[]], + "embeddings": [[]], } return self._transform_ns_hybrid_result(result_rows, include) @@ -6408,8 +6340,11 @@ def _transform_ns_hybrid_result( """Transform raw hybrid-search rows into the public result shape.""" if not result_rows: return { - "ids": [[]], "distances": [[]], "metadatas": [[]], - "documents": [[]], "embeddings": [[]], + "ids": [[]], + "distances": [[]], + "metadatas": [[]], + "documents": [[]], + "embeddings": [[]], } ids = [] diff --git a/src/pyseekdb/client/client_seekdb_server.py b/src/pyseekdb/client/client_seekdb_server.py index 90d4d033..313cb7bd 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -3,7 +3,6 @@ Supports both seekdb Server and OceanBase Server """ -import contextlib import logging from collections.abc import Sequence @@ -221,7 +220,9 @@ def _namespace_prewarm( ltable_id = self._resolve_namespace_ltable_id(collection_id, namespace_id) self._use_catalog_database() self._set_session_ns_context( - collection_id=collection_id, namespace_id=int(namespace_id), ltable_id=ltable_id, + collection_id=collection_id, + namespace_id=int(namespace_id), + ltable_id=ltable_id, ) sql = f"CALL DBMS_LOGIC_TABLE.PREWARM('{collection_id}', {namespace_id})" self._execute(sql) diff --git a/src/pyseekdb/client/collection.py b/src/pyseekdb/client/collection.py index 3e8b4ba6..16385a98 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -11,9 +11,8 @@ from typing import TYPE_CHECKING, Any, Optional from .validators import ( - _MAX_N_RESULTS, - _validate_namespace_name, _validate_n_results, + _validate_namespace_name, ) if TYPE_CHECKING: @@ -144,7 +143,9 @@ def _guard_collection_data_api(self) -> None: def _guard_namespace_enabled(self) -> None: """Raise if namespace APIs are used on a non-namespace collection.""" if not self._use_namespace: - raise ValueError("Namespace is not enabled for this collection. Use use_namespace=True when creating the collection.") + raise ValueError( + "Namespace is not enabled for this collection. Use use_namespace=True when creating the collection." + ) if not self._client._ns_collection_exists_by_id(self._id): raise ValueError( f"Collection '{self._name}' no longer exists (it may have been deleted). " @@ -163,6 +164,7 @@ def create_namespace(self, name: str) -> "Namespace": self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace + meta = self._client._create_ns_namespace_meta(self._id, name) return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) @@ -171,6 +173,7 @@ def get_namespace(self, name: str) -> "Namespace": self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace + meta = self._client._get_ns_namespace_meta(self._id, name) if meta is None: raise ValueError(f"Namespace '{name}' not found") @@ -181,6 +184,7 @@ def get_or_create_namespace(self, name: str) -> "Namespace": self._guard_namespace_enabled() _validate_namespace_name(name) from .namespace import Namespace + meta = self._client._get_or_create_ns_namespace_meta(self._id, name) return Namespace(client=self._client, collection=self, name=name, namespace_id=meta["namespace_id"]) @@ -194,6 +198,7 @@ def list_namespaces(self) -> list["Namespace"]: """List all active namespaces in this collection.""" self._guard_namespace_enabled() from .namespace import Namespace + metas = self._client._list_ns_namespaces(self._id) return [ Namespace(client=self._client, collection=self, name=m["namespace_name"], namespace_id=m["namespace_id"]) diff --git a/src/pyseekdb/client/configuration.py b/src/pyseekdb/client/configuration.py index b6017ac2..9b947d13 100644 --- a/src/pyseekdb/client/configuration.py +++ b/src/pyseekdb/client/configuration.py @@ -2,7 +2,7 @@ import warnings from dataclasses import dataclass -from enum import Enum +from enum import Enum, StrEnum from typing import Any, TypedDict from pyseekdb.client.embedding_function import DefaultEmbeddingFunction, EmbeddingFunction @@ -202,7 +202,7 @@ def _validate_hnsw_configuration(config: "HNSWConfiguration") -> None: raise TypeError(f"bq_use_fht must be a bool, got {type(config.bq_use_fht).__name__}") -class DistanceMetric(str, Enum): +class DistanceMetric(StrEnum): """ Distance metric constants for vector similarity calculation. @@ -214,7 +214,7 @@ class DistanceMetric(str, Enum): INNER_PRODUCT = "inner_product" -class HNSWIndexType(str, Enum): +class HNSWIndexType(StrEnum): """Supported HNSW index subtypes.""" HNSW = "hnsw" @@ -222,13 +222,13 @@ class HNSWIndexType(str, Enum): HNSW_BQ = "hnsw_bq" -class HNSWIndexLib(str, Enum): +class HNSWIndexLib(StrEnum): """Supported HNSW index libraries.""" VSAG = "vsag" -class IVFIndexType(str, Enum): +class IVFIndexType(StrEnum): """Supported IVF index subtypes for namespace-enabled collections.""" IVF_FLAT = "ivf_flat" @@ -236,14 +236,14 @@ class IVFIndexType(str, Enum): IVF_PQ = "ivf_pq" -class IVFIndexLib(str, Enum): +class IVFIndexLib(StrEnum): """Supported IVF index libraries.""" OB = "ob" VSAG = "vsag" -class FulltextAnalyzer(str, Enum): +class FulltextAnalyzer(StrEnum): """Supported fulltext analyzers.""" SPACE = "space" @@ -253,14 +253,14 @@ class FulltextAnalyzer(str, Enum): NGRAM2 = "ngram2" -class IKMode(str, Enum): +class IKMode(StrEnum): """Supported IK analyzer segmentation modes.""" SMART = "smart" MAX_WORD = "max_word" -class BQRefineType(str, Enum): +class BQRefineType(StrEnum): """Supported binary-quantization refine types for HNSW BQ indexes.""" SQ8 = "sq8" @@ -415,9 +415,7 @@ def __post_init__(self): raise ValueError(f"lib must be one of {valid_libs}, got {self.lib}") if self.centroids_fresh_mode is not None and not isinstance(self.centroids_fresh_mode, str): - raise TypeError( - f"centroids_fresh_mode must be a str, got {type(self.centroids_fresh_mode).__name__}" - ) + raise TypeError(f"centroids_fresh_mode must be a str, got {type(self.centroids_fresh_mode).__name__}") _ensure_primitive_properties(self.properties) diff --git a/src/pyseekdb/client/document_query_builder.py b/src/pyseekdb/client/document_query_builder.py index 68a3d416..1107d0da 100644 --- a/src/pyseekdb/client/document_query_builder.py +++ b/src/pyseekdb/client/document_query_builder.py @@ -76,10 +76,7 @@ def build_document_hybrid_expression( return _combine_document_bool(value, combiner="or", boost=boost) # Multiple top-level keys: treat as implicit AND. - parts = [ - build_document_hybrid_expression({key: val}, boost=boost) - for key, val in where_document.items() - ] + parts = [build_document_hybrid_expression({key: val}, boost=boost) for key, val in where_document.items()] parts = [part for part in parts if part is not None] if not parts: return None @@ -116,9 +113,7 @@ def where_document_knn_prefilterable(where_document: dict[str, Any] | str | None return False if _contains_operator(where_document, "$not_contains"): return False - if _contains_operator(where_document, "$regex"): - return False - return True + return not _contains_operator(where_document, "$regex") def doc_matches_where_document(document: str, where_document: dict[str, Any] | str) -> bool: @@ -133,15 +128,9 @@ def doc_matches_where_document(document: str, where_document: dict[str, Any] | s if "$regex" in where_document: return re.search(str(where_document["$regex"]), document) is not None if "$and" in where_document: - return all( - doc_matches_where_document(document, sub) - for sub in where_document["$and"] - ) + return all(doc_matches_where_document(document, sub) for sub in where_document["$and"]) if "$or" in where_document: - return any( - doc_matches_where_document(document, sub) - for sub in where_document["$or"] - ) + return any(doc_matches_where_document(document, sub) for sub in where_document["$or"]) raise ValueError(f"Unsupported where_document: {where_document!r}") @@ -170,7 +159,9 @@ def _combine_document_bool( body["boost"] = boost return {"query_string": body} - if all(isinstance(c, dict) and "$not_contains" in c and isinstance(c["$not_contains"], str) for c in conditions): + if all( + isinstance(c, dict) and "$not_contains" in c and isinstance(c["$not_contains"], str) for c in conditions + ): return { "bool": { "must_not": [ @@ -182,23 +173,22 @@ def _combine_document_bool( } } - if combiner == "or": - if all( - isinstance(c, dict) - and "$contains" in c - and isinstance(c["$contains"], str) - and _is_atomic_query_string_term(c["$contains"]) - for c in conditions - ): - queries = [c["$contains"] for c in conditions if isinstance(c, dict)] - body: dict[str, Any] = { - "fields": [_DOCUMENT_FIELD], - "query": " ".join(_escape_query_string_term(q) for q in queries), - "default_operator": "or", - } - if boost is not None: - body["boost"] = boost - return {"query_string": body} + if combiner == "or" and all( + isinstance(c, dict) + and "$contains" in c + and isinstance(c["$contains"], str) + and _is_atomic_query_string_term(c["$contains"]) + for c in conditions + ): + queries = [c["$contains"] for c in conditions if isinstance(c, dict)] + body: dict[str, Any] = { + "fields": [_DOCUMENT_FIELD], + "query": " ".join(_escape_query_string_term(q) for q in queries), + "default_operator": "or", + } + if boost is not None: + body["boost"] = boost + return {"query_string": body} child_exprs: list[dict[str, Any]] = [] for condition in conditions: diff --git a/src/pyseekdb/client/filters.py b/src/pyseekdb/client/filters.py index 501d6d1f..b9315162 100644 --- a/src/pyseekdb/client/filters.py +++ b/src/pyseekdb/client/filters.py @@ -189,9 +189,7 @@ def _build_document_condition( params.append(value) elif key == "$not_contains": - clauses.append( - f"NOT (MATCH({document_column}) AGAINST (%s IN NATURAL LANGUAGE MODE))" - ) + clauses.append(f"NOT (MATCH({document_column}) AGAINST (%s IN NATURAL LANGUAGE MODE))") params.append(value) elif key == "$regex": diff --git a/src/pyseekdb/client/kernel_errors.py b/src/pyseekdb/client/kernel_errors.py index 6bb0445e..9242dcb5 100644 --- a/src/pyseekdb/client/kernel_errors.py +++ b/src/pyseekdb/client/kernel_errors.py @@ -6,7 +6,8 @@ import contextvars import functools import re -from typing import Any, Callable, Iterator, TypeVar +from collections.abc import Callable, Iterator +from typing import Any, TypeVar _NS_KERNEL_ERROR_CTX: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( "ns_kernel_error_ctx", default=None @@ -145,13 +146,11 @@ def namespace_kernel_error_scope( collection_id: str | None = None, ): """Attach namespace/collection context used to translate kernel SQL errors.""" - token = _NS_KERNEL_ERROR_CTX.set( - { - "namespace_name": namespace_name or "", - "collection_name": collection_name or "", - "collection_id": collection_id or "", - } - ) + token = _NS_KERNEL_ERROR_CTX.set({ + "namespace_name": namespace_name or "", + "collection_name": collection_name or "", + "collection_id": collection_id or "", + }) try: yield finally: @@ -163,6 +162,7 @@ def namespace_kernel_error_scope( def namespace_kernel_error_guard(method: _F) -> _F: """Decorator for ``BaseClient._namespace_*`` methods with standard leading args.""" + @functools.wraps(method) def wrapper( self, diff --git a/src/pyseekdb/client/namespace.py b/src/pyseekdb/client/namespace.py index a9de484c..0446bdd2 100644 --- a/src/pyseekdb/client/namespace.py +++ b/src/pyseekdb/client/namespace.py @@ -22,7 +22,7 @@ class Namespace: def __init__( self, client: Any, - collection: "Collection", + collection: Collection, name: str, namespace_id: str, ): @@ -43,20 +43,19 @@ def namespace_id(self) -> str: return self._namespace_id @property - def collection(self) -> "Collection": + def collection(self) -> Collection: """Parent collection that owns this namespace.""" return self._collection @property - def embedding_function(self) -> "EmbeddingFunction | None": + def embedding_function(self) -> EmbeddingFunction | None: """Embedding function inherited from the parent collection.""" return self._collection.embedding_function def __repr__(self) -> str: """Return a debug-friendly representation of this namespace.""" return ( - f"Namespace(name='{self._name}', namespace_id={self._namespace_id}, " - f"collection='{self._collection.name}')" + f"Namespace(name='{self._name}', namespace_id={self._namespace_id}, collection='{self._collection.name}')" ) def _dml_collection_context(self) -> dict[str, Any]: diff --git a/src/pyseekdb/client/schema.py b/src/pyseekdb/client/schema.py index 41baa0fd..971388bd 100644 --- a/src/pyseekdb/client/schema.py +++ b/src/pyseekdb/client/schema.py @@ -26,7 +26,13 @@ from typing import Any -from .configuration import FulltextIndexConfig, HNSWConfiguration, IVFConfiguration, SparseVectorIndexConfig, VectorIndexConfig +from .configuration import ( + FulltextIndexConfig, + HNSWConfiguration, + IVFConfiguration, + SparseVectorIndexConfig, + VectorIndexConfig, +) class Schema: diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py index 403b0cb1..9899f936 100644 --- a/src/pyseekdb/client/validators.py +++ b/src/pyseekdb/client/validators.py @@ -7,18 +7,16 @@ _MAX_NAMESPACE_NAME_LENGTH = 256 _MAX_NAMESPACE_BATCH_SIZE = 100 _MAX_N_RESULTS = 16384 # OceanBase vector-search k upper bound -_VALID_INCLUDE_FIELDS = frozenset( - { - "documents", - "document", - "metadatas", - "metadata", - "embeddings", - "embedding", - "distances", - "distance", - } -) +_VALID_INCLUDE_FIELDS = frozenset({ + "documents", + "document", + "metadatas", + "metadata", + "embeddings", + "embedding", + "distances", + "distance", +}) def _validate_include(include: list[str] | None) -> None: @@ -26,9 +24,7 @@ def _validate_include(include: list[str] | None) -> None: if include is None: return if not isinstance(include, list): - raise TypeError( - f"include must be a list[str] or None, got {type(include).__name__}" - ) + raise TypeError(f"include must be a list[str] or None, got {type(include).__name__}") invalid = [field for field in include if field not in _VALID_INCLUDE_FIELDS] if invalid: allowed = "documents, metadatas, embeddings, distances" @@ -42,9 +38,7 @@ def _validate_include(include: list[str] | None) -> None: def _validate_namespace_name(name: str) -> None: """Validate a namespace name for type, length, and allowed characters.""" if not isinstance(name, str): - raise TypeError( - f"Invalid namespace name: '{name}'. Namespace name must be a string, got {type(name).__name__}" - ) + raise TypeError(f"Invalid namespace name: '{name}'. Namespace name must be a string, got {type(name).__name__}") if not name: raise ValueError(f"Invalid namespace name: '{name}'. Namespace name must not be empty") if len(name) > _MAX_NAMESPACE_NAME_LENGTH: @@ -61,9 +55,7 @@ def _validate_namespace_name(name: str) -> None: def _validate_database_name(name: str) -> None: """Validate a SQL database identifier for catalog table qualification.""" if not isinstance(name, str): - raise TypeError( - f"Invalid database name: '{name}'. Database name must be a string, got {type(name).__name__}" - ) + raise TypeError(f"Invalid database name: '{name}'. Database name must be a string, got {type(name).__name__}") if not name: raise ValueError(f"Invalid database name: '{name}'. Database name must not be empty") if _NAME_PATTERN.match(name) is None: @@ -84,22 +76,17 @@ def _validate_n_results(n_results: int, *, max_results: int = _MAX_N_RESULTS) -> raise ValueError(f"n_results must be an integer >= 1, got {n_results!r}") if n_results > max_results: raise ValueError( - f"n_results must be <= {max_results}, got {n_results}. " - "Use a smaller value or paginate with offset/limit." + f"n_results must be <= {max_results}, got {n_results}. Use a smaller value or paginate with offset/limit." ) def _validate_record_ids(ids: list[str]) -> None: """Validate namespace/collection record id list shape and per-id constraints.""" if not isinstance(ids, list): - raise TypeError( - f"Invalid record ids: expected list[str], got {type(ids).__name__}" - ) + raise TypeError(f"Invalid record ids: expected list[str], got {type(ids).__name__}") for rid in ids: if not isinstance(rid, str): - raise TypeError( - f"Invalid record id: '{rid}'. Record id must be a string, got {type(rid).__name__}" - ) + raise TypeError(f"Invalid record id: '{rid}'. Record id must be a string, got {type(rid).__name__}") if not rid: raise ValueError("Invalid record id: Record id must not be empty") if len(rid) > _MAX_NAME_LENGTH: @@ -127,8 +114,7 @@ def _validate_namespace_explicit_embedding_dimensions( if actual != expected_dimension: if has_vector_index: raise ValueError( - f"Embedding dimension mismatch: expected {expected_dimension}, " - f"got {actual} at index {i}." + f"Embedding dimension mismatch: expected {expected_dimension}, got {actual} at index {i}." ) raise ValueError( f"Collections without a vector index store embeddings as " diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index ad87f57e..8307f3c7 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -270,5 +270,3 @@ def oceanbase_admin_client(): with contextlib.suppress(Exception): if hasattr(client, "close"): client.close() - - diff --git a/tests/integration_tests/namespace_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py index 95ceaeec..42dcf6da 100644 --- a/tests/integration_tests/namespace_dml_helpers.py +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -4,6 +4,7 @@ from __future__ import annotations +import contextlib import time from dataclasses import dataclass from typing import Any @@ -31,6 +32,7 @@ @dataclass(frozen=True) class DmlRecord: """DmlRecord class.""" + doc_id: str embedding: list[float] document: str @@ -52,7 +54,9 @@ def create_ns_collection(client: Any, suffix: str = "") -> Any: """Create ns collection.""" name = f"test_ns_dml_{int(time.time() * 1000)}{suffix}" return client.create_collection( - name=name, schema=ns_schema(), use_namespace=True, + name=name, + schema=ns_schema(), + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) @@ -60,10 +64,8 @@ def create_ns_collection(client: Any, suffix: str = "") -> Any: def cleanup(client: Any, *collections: Any) -> None: """Cleanup.""" for collection in collections: - try: + with contextlib.suppress(Exception): client.delete_collection(name=collection.name) - except Exception: - pass def build_large_dml_corpus( @@ -132,9 +134,7 @@ def build_filter_corpus( meta: dict[str, Any] = {"tag": "keep", "category": category, "seq": i} if ns_tag is not None: meta["ns_tag"] = ns_tag - records.append( - DmlRecord(f"{id_prefix}_keep_{i:04d}", _filter_embedding(idx), document, meta) - ) + records.append(DmlRecord(f"{id_prefix}_keep_{i:04d}", _filter_embedding(idx), document, meta)) idx += 1 for i in range(purge_count): @@ -217,14 +217,9 @@ def assert_get_where_document_count( if substring and result.get("documents"): for i, doc in enumerate(result["documents"]): if doc is None: - raise AssertionError( - f"{prefix}document at index {i} is None for id={ids[i]!r}" - ) + raise AssertionError(f"{prefix}document at index {i} is None for id={ids[i]!r}") if substring not in doc.lower(): - raise AssertionError( - f"{prefix}id={ids[i]!r} document={doc!r} " - f"does not contain substring {substring!r}" - ) + raise AssertionError(f"{prefix}id={ids[i]!r} document={doc!r} does not contain substring {substring!r}") return result @@ -309,9 +304,7 @@ def load_dml_corpus( if assert_count: expected = len(corpus) actual = namespace.count() - assert actual == expected, ( - f"namespace {namespace.name!r} expected count {expected}, got {actual}" - ) + assert actual == expected, f"namespace {namespace.name!r} expected count {expected}, got {actual}" def assert_peek_result( diff --git a/tests/integration_tests/namespace_fts_helpers.py b/tests/integration_tests/namespace_fts_helpers.py index 8325f7f1..b1e8b41a 100644 --- a/tests/integration_tests/namespace_fts_helpers.py +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -14,14 +14,15 @@ from dataclasses import dataclass from typing import Any, Literal -VectorDistanceMetric = Literal["l2", "cosine"] -VECTOR_DISTANCE_METRICS: tuple[VectorDistanceMetric, ...] = ("l2", "cosine") - from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT + from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema +VectorDistanceMetric = Literal["l2", "cosine"] +VECTOR_DISTANCE_METRICS: tuple[VectorDistanceMetric, ...] = ("l2", "cosine") + # Rare ASCII tokens to reduce IK segmentation surprises in assertions. TOKEN_ZPX = "TOKENZPX" TOKEN_ALP = "TOKENALP" @@ -38,6 +39,7 @@ @dataclass(frozen=True) class CorpusRecord: """CorpusRecord class.""" + doc_id: str document: str embedding: list[float] @@ -48,6 +50,7 @@ class CorpusRecord: @dataclass(frozen=True) class FtsQueryCase: """FtsQueryCase class.""" + name: str where_document: dict[str, Any] | str n_results: int @@ -191,15 +194,9 @@ def doc_matches_where_document(document: str, where_document: dict[str, Any] | s if "$not_contains" in where_document: return where_document["$not_contains"].lower() not in text if "$and" in where_document: - return all( - doc_matches_where_document(document, sub) - for sub in where_document["$and"] - ) + return all(doc_matches_where_document(document, sub) for sub in where_document["$and"]) if "$or" in where_document: - return any( - doc_matches_where_document(document, sub) - for sub in where_document["$or"] - ) + return any(doc_matches_where_document(document, sub) for sub in where_document["$or"]) raise ValueError(f"Unsupported where_document for hybrid_search: {where_document!r}") @@ -246,9 +243,7 @@ def corpus_matches_fts( """Corpus matches fts.""" if not doc_matches_where_document(record.document, where_document): return False - if where is not None and not doc_matches_where_metadata(record.metadata, where): - return False - return True + return not (where is not None and not doc_matches_where_metadata(record.metadata, where)) def effective_rel_hint(record: CorpusRecord, where_document: dict[str, Any] | str) -> int: @@ -348,7 +343,7 @@ def assert_hybrid_fulltext_result( ) -> None: """Assert hybrid fulltext result.""" assert result is not None - assert "ids" in result and result["ids"] + assert result.get("ids") ids = result["ids"][0] distances = result.get("distances", [[]])[0] if result.get("distances") else [] documents = result.get("documents", [[]])[0] if result.get("documents") else [] @@ -364,7 +359,7 @@ def assert_hybrid_fulltext_result( id_to_meta = {rec.doc_id: rec.metadata for rec in corpus} meta_rows = metadatas if metadatas else [None] * len(ids) - for doc_id, doc_text, meta in zip(ids, documents, meta_rows): + for doc_id, doc_text, meta in zip(ids, documents, meta_rows, strict=False): assert doc_id in id_to_doc, f"unknown id {doc_id!r} in search results" body = doc_text if doc_text is not None else id_to_doc[doc_id] assert doc_matches_where_document(body, where_document), ( @@ -525,7 +520,9 @@ def setup_large_fts_collection( name = f"test_ns_hs_ft_{distance}_{int(time.time() * 1000)}" collection = db_client.create_collection( - name=name, schema=ns_schema(distance), use_namespace=True, + name=name, + schema=ns_schema(distance), + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) return corpus, collection @@ -551,14 +548,12 @@ def setup_fts_namespace_with_corpus( return namespace -def setup_large_fts_namespace(db_client: Any, *, namespace_name: str = "ns_hs_ft_large") -> tuple[ - list[CorpusRecord], Any, Any -]: +def setup_large_fts_namespace( + db_client: Any, *, namespace_name: str = "ns_hs_ft_large" +) -> tuple[list[CorpusRecord], Any, Any]: """Create collection + namespace with corpus (single-test convenience wrapper).""" corpus, collection = setup_large_fts_collection(db_client) - namespace = setup_fts_namespace_with_corpus( - collection, corpus, namespace_name=namespace_name - ) + namespace = setup_fts_namespace_with_corpus(collection, corpus, namespace_name=namespace_name) return corpus, collection, namespace @@ -621,9 +616,7 @@ def setup_multi_coll_multi_ns_fts( only the namespace that received the bulk insert is expected to serve hybrid_search FTS (see ``setup_same_collection_both_ns_fts`` for same-collection multi-ns behavior). """ - ctx = _create_multi_coll_multi_ns_layout( - db_client, name_prefix="test_ns_hs_ft_mcmn", distance=distance - ) + ctx = _create_multi_coll_multi_ns_layout(db_client, name_prefix="test_ns_hs_ft_mcmn", distance=distance) corpus = build_large_fts_corpus(CORPUS_SIZE) if len(corpus) <= 1000: raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") @@ -680,9 +673,7 @@ def setup_multi_coll_multi_ns_fts_single_loaded( if loaded_quadrant not in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: raise ValueError(f"unknown quadrant {loaded_quadrant!r}") - ctx = _create_multi_coll_multi_ns_layout( - db_client, name_prefix="test_ns_hs_ft_mcmn1", distance=distance - ) + ctx = _create_multi_coll_multi_ns_layout(db_client, name_prefix="test_ns_hs_ft_mcmn1", distance=distance) corpus = build_large_fts_corpus(CORPUS_SIZE) insert_corpus_in_batches(ctx[loaded_quadrant], corpus) @@ -717,9 +708,7 @@ def assert_hybrid_search_no_hits( include=["documents"], ) ids = result.get("ids", [[]])[0] if result.get("ids") else [] - assert len(ids) == 0, ( - f"expected no full-text hits in namespace {namespace.name!r}, got {len(ids)} ids: {ids[:5]!r}" - ) + assert len(ids) == 0, f"expected no full-text hits in namespace {namespace.name!r}, got {len(ids)} ids: {ids[:5]!r}" return result @@ -778,10 +767,6 @@ def assert_not_contains_no_token_leak( ) -> None: """Assert not contains no token leak.""" result_ids = set(result["ids"][0]) - forbidden_ids = { - rec.doc_id - for rec in corpus - if forbidden_token.lower() in rec.document.lower() - } + forbidden_ids = {rec.doc_id for rec in corpus if forbidden_token.lower() in rec.document.lower()} leaked = result_ids & forbidden_ids assert not leaked, f"$not_contains leaked forbidden ids: {sorted(leaked)[:10]}" diff --git a/tests/integration_tests/namespace_hybrid_search_helpers.py b/tests/integration_tests/namespace_hybrid_search_helpers.py index 6280a0de..e36a3e08 100644 --- a/tests/integration_tests/namespace_hybrid_search_helpers.py +++ b/tests/integration_tests/namespace_hybrid_search_helpers.py @@ -16,32 +16,28 @@ BATCH_SIZE, CORPUS_SIZE, INDEX_SETTLE_SECONDS, + MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, + MULTI_COLL_MULTI_NS_QUADRANT_KEYS, TOKEN_ALP, TOKEN_ZPX, VECTOR_DISTANCE_METRICS, - VectorDistanceMetric, CorpusRecord, + VectorDistanceMetric, assert_hybrid_fulltext_result, assert_not_contains_no_token_leak, build_large_fts_corpus, - corpus_matches_fts, doc_matches_where_document, doc_matches_where_metadata, + flat_hybrid_search_schema, get_fts_case, - insert_corpus_in_batches, insert_corpus_into_collection, - insert_corpus_into_namespace, run_hybrid_search_fts_case, - ns_schema, - flat_hybrid_search_schema, setup_fts_namespace_with_corpus, setup_large_fts_collection, setup_multi_coll_multi_ns_fts, setup_multi_coll_multi_ns_fts_single_loaded, teardown_large_fts_collection, teardown_multi_coll_multi_ns_fts, - MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, - MULTI_COLL_MULTI_NS_QUADRANT_KEYS, ) # Fixed query vector for KNN tests (same dimension as corpus embeddings). @@ -67,6 +63,7 @@ @dataclass(frozen=True) class SearchIndexQueryCase: """SearchIndexQueryCase class.""" + name: str where: dict[str, Any] n_results: int @@ -77,6 +74,7 @@ class SearchIndexQueryCase: @dataclass(frozen=True) class VectorKnnCase: """VectorKnnCase class.""" + name: str query_vector: list[float] n_results: int @@ -157,7 +155,7 @@ def knn_compare_distance( norm_b = _vector_l2_norm(b) if norm_a == 0.0 or norm_b == 0.0: return 2.0 - cos_sim = sum(x * y for x, y in zip(a, b)) / (norm_a * norm_b) + cos_sim = sum(x * y for x, y in zip(a, b, strict=False)) / (norm_a * norm_b) return 1.0 - cos_sim raise ValueError(f"unsupported vector distance metric: {metric!r}") @@ -169,16 +167,10 @@ def ensure_shared_hybrid_search_collection( vector_distance: VectorDistanceMetric, ) -> dict[str, Any]: """Return or create a shared (db mode, distance metric) corpus + collection entry.""" - mode = ( - request.node.callspec.params["db_client"] - if request.node.callspec - else "default" - ) + mode = request.node.callspec.params["db_client"] if request.node.callspec else "default" cache_key = f"{mode}:{vector_distance}" if cache_key not in shared_cache: - corpus, collection = setup_large_fts_collection( - db_client, distance=vector_distance - ) + corpus, collection = setup_large_fts_collection(db_client, distance=vector_distance) shared_cache[cache_key] = { "db_client": db_client, "corpus": corpus, @@ -203,13 +195,11 @@ def _knn_scored_rows( seq = rec.metadata.get("seq", 0) if not isinstance(seq, int): seq = 0 - rows.append( - ( - rec.doc_id, - knn_compare_distance(rec.embedding, query_vector, distance_metric), - seq, - ) - ) + rows.append(( + rec.doc_id, + knn_compare_distance(rec.embedding, query_vector, distance_metric), + seq, + )) return rows @@ -222,9 +212,7 @@ def nearest_knn_doc_ids( tol: float = 1e-9, ) -> set[str]: """Doc ids at minimum compare distance (ties allowed).""" - scored = _knn_scored_rows( - corpus, query_vector, where, distance_metric=distance_metric - ) + scored = _knn_scored_rows(corpus, query_vector, where, distance_metric=distance_metric) if not scored: return set() min_d = min(d for _, d, _ in scored) @@ -240,9 +228,7 @@ def expected_knn_ids( distance_metric: VectorDistanceMetric = "l2", ) -> list[str]: """Expected knn ids.""" - scored = _knn_scored_rows( - corpus, query_vector, where, distance_metric=distance_metric - ) + scored = _knn_scored_rows(corpus, query_vector, where, distance_metric=distance_metric) # Tie-break: distance, then seq (matches OB hybrid_search), then doc_id. scored.sort(key=lambda item: (item[1], item[2], item[0])) return [doc_id for doc_id, _, _ in scored[:n_results]] @@ -259,7 +245,7 @@ def assert_hybrid_search_index_result( ) -> None: """Assert hybrid search index result.""" assert result is not None - assert "ids" in result and result["ids"] + assert result.get("ids") ids = result["ids"][0] assert len(ids) <= n_results assert len(ids) >= min_hits, f"expected at least {min_hits} search-index hits" @@ -267,9 +253,9 @@ def assert_hybrid_search_index_result( id_to_meta = {rec.doc_id: rec.metadata for rec in corpus} for doc_id in ids: assert doc_id in id_to_meta, f"unknown id {doc_id!r}" - assert corpus_matches_where( - next(rec for rec in corpus if rec.doc_id == doc_id), where - ), f"id={doc_id!r} does not satisfy where={where!r}" + assert corpus_matches_where(next(rec for rec in corpus if rec.doc_id == doc_id), where), ( + f"id={doc_id!r} does not satisfy where={where!r}" + ) total_matches = count_corpus_matches(corpus, where) if exact_match_count is not None: @@ -278,9 +264,7 @@ def assert_hybrid_search_index_result( ) if exact_match_count is not None and exact_match_count <= n_results: assert len(ids) == exact_match_count - expected_ids = { - rec.doc_id for rec in corpus if corpus_matches_where(rec, where) - } + expected_ids = {rec.doc_id for rec in corpus if corpus_matches_where(rec, where)} assert set(ids) == expected_ids elif len(ids) == n_results and total_matches > n_results: # No higher-priority scalar key — every returned row is valid; ensure we did not @@ -306,7 +290,7 @@ def assert_hybrid_knn_result( corpus ``distance_metric`` ground truth. """ assert result is not None - assert "ids" in result and result["ids"] + assert result.get("ids") ids = result["ids"][0] scores = result.get("distances", [[]])[0] if result.get("distances") else [] assert len(ids) <= n_results @@ -316,33 +300,23 @@ def assert_hybrid_knn_result( for doc_id in ids: rec = corpus_by_id[doc_id] if where is not None: - assert corpus_matches_where(rec, where), ( - f"id={doc_id!r} metadata does not satisfy knn.where={where!r}" - ) + assert corpus_matches_where(rec, where), f"id={doc_id!r} metadata does not satisfy knn.where={where!r}" if scores: assert len(scores) == len(ids) if check_score_order: for i in range(len(scores) - 1): - assert scores[i] >= scores[i + 1] or math.isclose( - scores[i], scores[i + 1] - ), ( - f"KNN scores should be non-increasing (higher = nearer, " - f"{distance_metric} index): {scores!r}" + assert scores[i] >= scores[i + 1] or math.isclose(scores[i], scores[i + 1]), ( + f"KNN scores should be non-increasing (higher = nearer, {distance_metric} index): {scores!r}" ) - expected = expected_knn_ids( - corpus, query_vector, n_results, where, distance_metric=distance_metric - ) + expected = expected_knn_ids(corpus, query_vector, n_results, where, distance_metric=distance_metric) if len(ids) == len(expected): assert ids == expected, ( - f"KNN id order mismatch ({distance_metric} ground truth): " - f"got {ids!r} expected {expected!r}" + f"KNN id order mismatch ({distance_metric} ground truth): got {ids!r} expected {expected!r}" ) elif check_top1 and expected: - nearest = nearest_knn_doc_ids( - corpus, query_vector, where, distance_metric=distance_metric - ) + nearest = nearest_knn_doc_ids(corpus, query_vector, where, distance_metric=distance_metric) assert ids[0] in nearest, ( f"top-1 KNN must be among nearest neighbors ({distance_metric} tie set), " f"got {ids[0]!r} expected one of {sorted(nearest)[:8]!r}" @@ -350,20 +324,11 @@ def assert_hybrid_knn_result( ) if len(ids) == n_results: - scored = _knn_scored_rows( - corpus, query_vector, where, distance_metric=distance_metric - ) - if scored: - worst_d = max( - d for doc_id, d, _ in scored if doc_id in ids - ) - else: - worst_d = float("inf") + scored = _knn_scored_rows(corpus, query_vector, where, distance_metric=distance_metric) + worst_d = max(d for doc_id, d, _ in scored if doc_id in ids) if scored else float("inf") for doc_id, d, _ in scored: if doc_id not in ids and d < worst_d - 1e-9: - raise AssertionError( - f"closer match {doc_id!r} ({distance_metric}={d}) missing from top-{n_results}" - ) + raise AssertionError(f"closer match {doc_id!r} ({distance_metric}={d}) missing from top-{n_results}") def run_hybrid_search_index_case( @@ -477,9 +442,7 @@ def run_hybrid_combined_case( check_ranking=case.check_fts_ranking, ) elif case.where is not None: - assert_hybrid_search_index_result( - corpus, result, case.where, case.n_results - ) + assert_hybrid_search_index_result(corpus, result, case.where, case.n_results) return result @@ -739,12 +702,10 @@ def corpus_matches_triple_intersection( case: HybridTripleBranchCase, ) -> bool: """Corpus matches triple intersection.""" - if case.where_document is not None: - if not doc_matches_where_document(record.document, case.where_document): - return False - if case.where is not None: - if not corpus_matches_where(record, case.where): - return False + if case.where_document is not None and not doc_matches_where_document(record.document, case.where_document): + return False + if case.where is not None and not corpus_matches_where(record, case.where): + return False if case.knn is not None: knn_where = case.knn.get("where") if knn_where is not None and not corpus_matches_where(record, knn_where): @@ -768,11 +729,7 @@ def _slice_hybrid_result( """Slice hybrid result.""" corpus_by_id = {rec.doc_id: rec for rec in corpus} ids = result["ids"][0] - indices = [ - i - for i, doc_id in enumerate(ids) - if doc_id in corpus_by_id and predicate(corpus_by_id[doc_id]) - ] + indices = [i for i, doc_id in enumerate(ids) if doc_id in corpus_by_id and predicate(corpus_by_id[doc_id])] sliced: dict[str, Any] = {"ids": [[ids[i] for i in indices]]} for key in ("distances", "documents", "metadatas"): rows = result.get(key) @@ -790,12 +747,10 @@ def assert_hybrid_triple_branch_result( ) -> None: """Assert hybrid triple branch result.""" assert result is not None - assert "ids" in result and result["ids"] + assert result.get("ids") ids = result["ids"][0] assert len(ids) <= case.n_results - assert len(ids) >= case.min_hits, ( - f"case {case.name!r}: expected at least {case.min_hits} hits, got {len(ids)}" - ) + assert len(ids) >= case.min_hits, f"case {case.name!r}: expected at least {case.min_hits} hits, got {len(ids)}" corpus_by_id = {rec.doc_id: rec for rec in corpus} for doc_id in ids: @@ -806,9 +761,9 @@ def assert_hybrid_triple_branch_result( assert_not_contains_no_token_leak(corpus, result, TOKEN_ZPX) if case.where is not None: for doc_id in ids: - assert doc_matches_where_metadata( - corpus_by_id[doc_id].metadata, case.where - ), f"id={doc_id!r} violates query.where={case.where!r}" + assert doc_matches_where_metadata(corpus_by_id[doc_id].metadata, case.where), ( + f"id={doc_id!r} violates query.where={case.where!r}" + ) if case.knn and case.knn.get("where"): for doc_id in ids: assert corpus_matches_where(corpus_by_id[doc_id], case.knn["where"]), ( @@ -823,15 +778,11 @@ def _fts_row(rec: CorpusRecord) -> bool: """Fts row.""" if not doc_matches_where_document(rec.document, case.where_document): return False - if case.where is not None and not doc_matches_where_metadata(rec.metadata, case.where): - return False - return True + return not (case.where is not None and not doc_matches_where_metadata(rec.metadata, case.where)) fts_result = _slice_hybrid_result(result, corpus, _fts_row) fts_ids = fts_result["ids"][0] - assert len(fts_ids) >= 1, ( - f"case {case.name!r}: expected at least one FTS-matching row in hybrid result" - ) + assert len(fts_ids) >= 1, f"case {case.name!r}: expected at least one FTS-matching row in hybrid result" assert_hybrid_fulltext_result( corpus, fts_result, @@ -854,17 +805,13 @@ def _si_row(rec: CorpusRecord) -> bool: """Si row.""" if not corpus_matches_where(rec, case.where): return False - if case.where_document is not None and not doc_matches_where_document( - rec.document, case.where_document - ): - return False - return True + return not ( + case.where_document is not None and not doc_matches_where_document(rec.document, case.where_document) + ) si_result = _slice_hybrid_result(result, corpus, _si_row) si_ids = si_result["ids"][0] - assert len(si_ids) >= case.min_hits, ( - f"case {case.name!r}: expected at least {case.min_hits} search-index rows" - ) + assert len(si_ids) >= case.min_hits, f"case {case.name!r}: expected at least {case.min_hits} search-index rows" assert_hybrid_search_index_result( corpus, si_result, @@ -890,17 +837,13 @@ def _knn_row(rec: CorpusRecord) -> bool: return False if case.where is not None and not doc_matches_where_metadata(rec.metadata, case.where): return False - if case.where_document is not None and not doc_matches_where_document( - rec.document, case.where_document - ): - return False - return True + return not ( + case.where_document is not None and not doc_matches_where_document(rec.document, case.where_document) + ) knn_result = _slice_hybrid_result(result, corpus, _knn_row) knn_ids = knn_result["ids"][0] - assert len(knn_ids) >= 1, ( - f"case {case.name!r}: expected at least one KNN-eligible row in hybrid result" - ) + assert len(knn_ids) >= 1, f"case {case.name!r}: expected at least one KNN-eligible row in hybrid result" assert_hybrid_knn_result( corpus, knn_result, @@ -917,55 +860,37 @@ def _knn_row(rec: CorpusRecord) -> bool: total = count_triple_intersection_matches(corpus, case) if case.exact_match_count is not None: assert total == case.exact_match_count, ( - f"case {case.name!r}: expected {case.exact_match_count} intersection " - f"matches, corpus has {total}" + f"case {case.name!r}: expected {case.exact_match_count} intersection matches, corpus has {total}" ) for doc_id in ids: assert corpus_matches_triple_intersection(corpus_by_id[doc_id], case), ( f"id={doc_id!r} outside triple-branch intersection for case {case.name!r}" ) if case.exact_match_count is not None and case.exact_match_count <= case.n_results: - expected_ids = { - rec.doc_id - for rec in corpus - if corpus_matches_triple_intersection(rec, case) - } + expected_ids = {rec.doc_id for rec in corpus if corpus_matches_triple_intersection(rec, case)} assert set(ids) == expected_ids return if case.verify == "rrf_fusion": if case.where is not None: for doc_id in ids: - assert doc_matches_where_metadata( - corpus_by_id[doc_id].metadata, case.where - ), f"id={doc_id!r} violates query.where={case.where!r}" + assert doc_matches_where_metadata(corpus_by_id[doc_id].metadata, case.where), ( + f"id={doc_id!r} violates query.where={case.where!r}" + ) if case.knn and case.knn.get("where"): for doc_id in ids: assert corpus_matches_where(corpus_by_id[doc_id], case.knn["where"]), ( f"id={doc_id!r} violates knn.where={case.knn['where']!r}" ) - intersection_ids = { - rec.doc_id - for rec in corpus - if corpus_matches_triple_intersection(rec, case) - } - assert intersection_ids, ( - f"case {case.name!r}: intersection must be non-empty for RRF fusion checks" - ) + intersection_ids = {rec.doc_id for rec in corpus if corpus_matches_triple_intersection(rec, case)} + assert intersection_ids, f"case {case.name!r}: intersection must be non-empty for RRF fusion checks" returned = set(ids) assert returned & intersection_ids, ( - f"case {case.name!r}: RRF result must include at least one intersection hit, " - f"got {ids[:5]!r}" + f"case {case.name!r}: RRF result must include at least one intersection hit, got {ids[:5]!r}" ) if case.where_document is not None: - fts_hits = { - rec.doc_id - for rec in corpus - if doc_matches_where_document(rec.document, case.where_document) - } - assert returned & fts_hits, ( - f"case {case.name!r}: RRF result must include at least one FTS hit" - ) + fts_hits = {rec.doc_id for rec in corpus if doc_matches_where_document(rec.document, case.where_document)} + assert returned & fts_hits, f"case {case.name!r}: RRF result must include at least one FTS hit" return raise ValueError(f"unsupported verify mode: {case.verify!r}") @@ -1043,8 +968,7 @@ def execute_hybrid_triple_branch_search( if hasattr(target, "namespace_id"): return execute_hybrid_triple_branch_on_namespace(target, case) raise TypeError( - f"unsupported hybrid_search target {target!r}: expected flat Collection " - f"(use_namespace=False) or Namespace" + f"unsupported hybrid_search target {target!r}: expected flat Collection (use_namespace=False) or Namespace" ) @@ -1059,16 +983,12 @@ def setup_large_fts_flat_collection( raise ValueError(f"corpus must exceed 1000 rows, got {len(corpus)}") name = f"test_hs_tb_baseline_{distance}_{int(time.time() * 1000)}" - collection = db_client.create_collection( - name=name, schema=flat_hybrid_search_schema(distance), use_namespace=False - ) + collection = db_client.create_collection(name=name, schema=flat_hybrid_search_schema(distance), use_namespace=False) assert collection.use_namespace is False insert_corpus_into_collection(collection, corpus) expected = len(corpus) actual = collection.count() - assert actual == expected, ( - f"baseline collection {collection.name!r} expected {expected} rows, got {actual}" - ) + assert actual == expected, f"baseline collection {collection.name!r} expected {expected} rows, got {actual}" time.sleep(INDEX_SETTLE_SECONDS) return corpus, collection @@ -1082,9 +1002,7 @@ def _try_assert_hybrid_triple_branch_result( ) -> BaseException | None: """Try assert hybrid triple branch result.""" try: - assert_hybrid_triple_branch_result( - corpus, result, case, distance_metric=distance_metric - ) + assert_hybrid_triple_branch_result(corpus, result, case, distance_metric=distance_metric) return None except BaseException as exc: return exc @@ -1112,12 +1030,8 @@ def assert_hybrid_triple_branch_vs_collection_baseline( - ``[observer/namespace]`` — collection passes ground truth but namespace fails - ``[test-case/baseline]`` — namespace passes but collection baseline fails """ - coll_err = _try_assert_hybrid_triple_branch_result( - corpus, collection_result, case, distance_metric=distance_metric - ) - ns_err = _try_assert_hybrid_triple_branch_result( - corpus, namespace_result, case, distance_metric=distance_metric - ) + coll_err = _try_assert_hybrid_triple_branch_result(corpus, collection_result, case, distance_metric=distance_metric) + ns_err = _try_assert_hybrid_triple_branch_result(corpus, namespace_result, case, distance_metric=distance_metric) coll_ids = collection_result["ids"][0] ns_ids = namespace_result["ids"][0] @@ -1173,9 +1087,7 @@ def run_hybrid_triple_branch_case( if collection_baseline is not None: assert_flat_collection_baseline(collection_baseline) try: - collection_result = execute_hybrid_triple_branch_on_collection( - collection_baseline, case - ) + collection_result = execute_hybrid_triple_branch_on_collection(collection_baseline, case) except BaseException as coll_exec_err: raise AssertionError( f"[test-case/baseline] case {case.name!r}: collection.hybrid_search " @@ -1190,15 +1102,16 @@ def run_hybrid_triple_branch_case( f"{ns_exec_err!r}" ) from ns_exec_err assert_hybrid_triple_branch_vs_collection_baseline( - corpus, case, namespace_result, collection_result, + corpus, + case, + namespace_result, + collection_result, distance_metric=distance_metric, ) return namespace_result result = execute_hybrid_triple_branch_on_namespace(namespace, case) - assert_hybrid_triple_branch_result( - corpus, result, case, distance_metric=distance_metric - ) + assert_hybrid_triple_branch_result(corpus, result, case, distance_metric=distance_metric) return result @@ -1628,9 +1541,7 @@ def run_knn_case_on_quadrants( knn_case = case if isinstance(case, VectorKnnCase) else get_vector_knn_case(case) for key in quadrant_keys: corpus, namespace = ctx[key] - run_hybrid_knn_case( - namespace, corpus, knn_case, distance_metric=distance_metric - ) + run_hybrid_knn_case(namespace, corpus, knn_case, distance_metric=distance_metric) def assert_hybrid_search_index_no_hits( @@ -1646,9 +1557,7 @@ def assert_hybrid_search_index_no_hits( include=["metadatas"], ) ids = result.get("ids", [[]])[0] if result.get("ids") else [] - assert len(ids) == 0, ( - f"expected no search-index hits in namespace {namespace.name!r}, got {ids[:5]!r}" - ) + assert len(ids) == 0, f"expected no search-index hits in namespace {namespace.name!r}, got {ids[:5]!r}" def assert_hybrid_triple_branch_no_hits( @@ -1671,9 +1580,7 @@ def assert_hybrid_triple_branch_no_hits( include=["documents"], ) ids = result.get("ids", [[]])[0] if result.get("ids") else [] - assert len(ids) == 0, ( - f"expected no triple-branch hits in namespace {namespace.name!r}, got {ids[:5]!r}" - ) + assert len(ids) == 0, f"expected no triple-branch hits in namespace {namespace.name!r}, got {ids[:5]!r}" __all__ = [ @@ -1691,15 +1598,13 @@ def assert_hybrid_triple_branch_no_hits( "TRIPLE_BRANCH_CASES", "VECTOR_DISTANCE_METRICS", "VECTOR_KNN_CASES", - "VectorDistanceMetric", - "ensure_shared_hybrid_search_collection", - "knn_compare_distance", "WHERE_BROAD_METADATA", "WHERE_DOCUMENT_UNIVERSAL", "WHERE_FILLER_SEQ", "HybridCombinedCase", "HybridTripleBranchCase", "SearchIndexQueryCase", + "VectorDistanceMetric", "VectorKnnCase", "assert_flat_collection_baseline", "assert_hybrid_search_index_no_hits", @@ -1708,6 +1613,7 @@ def assert_hybrid_triple_branch_no_hits( "assert_hybrid_triple_branch_vs_collection_baseline", "assert_not_contains_no_token_leak", "build_large_fts_corpus", + "ensure_shared_hybrid_search_collection", "execute_hybrid_triple_branch_on_collection", "execute_hybrid_triple_branch_on_namespace", "execute_hybrid_triple_branch_search", @@ -1716,10 +1622,11 @@ def assert_hybrid_triple_branch_no_hits( "get_search_index_case", "get_triple_branch_case", "get_vector_knn_case", + "knn_compare_distance", "run_hybrid_combined_case", "run_hybrid_knn_case", - "run_hybrid_search_index_case", "run_hybrid_search_fts_case", + "run_hybrid_search_index_case", "run_hybrid_triple_branch_case", "run_knn_case_on_quadrants", "run_search_index_case_on_quadrants", diff --git a/tests/integration_tests/test_collection_embedding_function.py b/tests/integration_tests/test_collection_embedding_function.py index 0912d15e..f06afe8b 100644 --- a/tests/integration_tests/test_collection_embedding_function.py +++ b/tests/integration_tests/test_collection_embedding_function.py @@ -637,7 +637,7 @@ def build_from_config(config: dict[str, Any]) -> "UsableEmbeddingFunction": with contextlib.suppress(Exception): db_client.delete_collection(name=collection_name) - def test_multiple_custom_embedding_functions(self, db_client): # noqa: C901 + def test_multiple_custom_embedding_functions(self, db_client): """ Test multiple collections with different custom embedding functions. @@ -807,7 +807,8 @@ def build_from_config( EmbeddingFunctionRegistry._registry.pop(ef.name()) # Actually, since we registered it, getting should work - pytest.raises(ValueError, db_client.get_collection, name=collection_name) + with pytest.raises(ValueError): + db_client.get_collection(name=collection_name) finally: # Cleanup with contextlib.suppress(Exception): diff --git a/tests/integration_tests/test_get_or_create_collection_concurrency.py b/tests/integration_tests/test_get_or_create_collection_concurrency.py index 3f9498e6..b1107eb9 100644 --- a/tests/integration_tests/test_get_or_create_collection_concurrency.py +++ b/tests/integration_tests/test_get_or_create_collection_concurrency.py @@ -9,10 +9,9 @@ import uuid from concurrent.futures import ThreadPoolExecutor, as_completed -import pytest +from pymysql.converters import escape_string import pyseekdb -from pymysql.converters import escape_string def _make_oceanbase_client(): @@ -33,8 +32,7 @@ def _count_sdk_collections(client, collection_name: str) -> int: """Count sdk collections.""" name_escaped = escape_string(collection_name) rows = client._server._execute( - "SELECT COUNT(*) AS cnt FROM sdk_collections " - f"WHERE collection_name = '{name_escaped}'" + f"SELECT COUNT(*) AS cnt FROM sdk_collections WHERE collection_name = '{name_escaped}'" ) row = rows[0] return int(row["cnt"] if isinstance(row, dict) else row[0]) @@ -52,6 +50,7 @@ def _has_unique_name_index(client) -> bool: class TestGetOrCreateCollectionConcurrencyOceanBase: """TestGetOrCreateCollectionConcurrencyOceanBase class.""" + def test_concurrent_get_or_create_collection_is_idempotent(self, oceanbase_client): """8 independent clients race on get_or_create_collection with the same name.""" owner = oceanbase_client @@ -74,7 +73,7 @@ def _worker(thread_id: int) -> None: "collection_id": str(coll.id), "name": coll.name, } - except Exception as exc: # noqa: BLE001 + except Exception as exc: errors.append(f"thread {thread_id}: {exc!r}") finally: if hasattr(client, "close"): diff --git a/tests/integration_tests/test_get_or_create_collection_multiprocess.py b/tests/integration_tests/test_get_or_create_collection_multiprocess.py index 8208bb96..92c2a594 100644 --- a/tests/integration_tests/test_get_or_create_collection_multiprocess.py +++ b/tests/integration_tests/test_get_or_create_collection_multiprocess.py @@ -147,7 +147,7 @@ def _require_embedded_pylibseekdb() -> None: ) -def _run_processes( # noqa: C901 +def _run_processes( target: Callable[..., None], args_list: list[tuple[Any, ...]], expected_count: int, @@ -566,6 +566,7 @@ def _seed_collection_rows( class TestGetOrCreateCollectionMultiprocess: """TestGetOrCreateCollectionMultiprocess class.""" + def test_concurrent_get_or_create_collection(self, multiprocess_db, _mode): """Test concurrent get or create collection.""" client_config = multiprocess_db @@ -590,12 +591,9 @@ def test_concurrent_get_or_create_collection(self, multiprocess_db, _mode): collection = _get_collection(client, collection_name) assert collection.name == collection_name rows = client._server._execute( - "SELECT collection_id FROM sdk_collections " - f"WHERE collection_name = '{collection_name}'" - ) - assert len(rows) == 1, ( - f"expected exactly one sdk_collections row for {collection_name!r}, got {len(rows)}" + f"SELECT collection_id FROM sdk_collections WHERE collection_name = '{collection_name}'" ) + assert len(rows) == 1, f"expected exactly one sdk_collections row for {collection_name!r}, got {len(rows)}" row = rows[0] catalog_id = row["collection_id"] if isinstance(row, dict) else row[0] assert str(catalog_id) == str(collection.id) @@ -606,6 +604,7 @@ def test_concurrent_get_or_create_collection(self, multiprocess_db, _mode): class TestMultiprocessMultithreadCrud: """TestMultiprocessMultithreadCrud class.""" + def test_concurrent_add(self, crud_collection, _mode): """Test concurrent add.""" client_config, collection_name = crud_collection diff --git a/tests/integration_tests/test_logic_table_monitoring_p1.py b/tests/integration_tests/test_logic_table_monitoring_p1.py index 070c08b7..b668eca2 100644 --- a/tests/integration_tests/test_logic_table_monitoring_p1.py +++ b/tests/integration_tests/test_logic_table_monitoring_p1.py @@ -10,9 +10,10 @@ import uuid from concurrent.futures import ThreadPoolExecutor, as_completed -import pyseekdb from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, ns_schema +import pyseekdb + def _make_oceanbase_client(): """Make oceanbase client.""" @@ -51,6 +52,7 @@ def _count_sdk_ltables(client, collection_id: str, namespace_id: int) -> int: class TestLogicTableMonitoringP1: """TestLogicTableMonitoringP1 class.""" + def test_concurrent_get_or_create_same_namespace_is_idempotent(self, oceanbase_client): """8 threads with independent clients race on the same namespace name.""" owner = oceanbase_client @@ -78,7 +80,7 @@ def _worker(thread_id: int) -> None: "namespace_id": str(ns.namespace_id), "name": ns.name, } - except Exception as exc: # noqa: BLE001 + except Exception as exc: errors.append(f"thread {thread_id}: {exc!r}") finally: if hasattr(client, "close"): diff --git a/tests/integration_tests/test_namespace_constraints.py b/tests/integration_tests/test_namespace_constraints.py index 651d94f7..f5f7d9d6 100644 --- a/tests/integration_tests/test_namespace_constraints.py +++ b/tests/integration_tests/test_namespace_constraints.py @@ -10,8 +10,8 @@ import time import pytest - from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT + from pyseekdb import IVFConfiguration from pyseekdb.client.client_base import NAMESPACE_MIN_OB_VERSION from pyseekdb.client.configuration import VectorIndexConfig @@ -36,14 +36,13 @@ def _unique_name(suffix: str) -> str: class TestNamespaceIvfTypeConstraint: """TestNamespaceIvfTypeConstraint class.""" + @pytest.mark.parametrize("ivf_type", ["ivf_pq", "ivf_sq8"]) def test_non_ivf_flat_rejected(self, oceanbase_client, ivf_type): """Non-ivf_flat IVF types must be rejected at the SDK layer.""" name = _unique_name(f"_{ivf_type}") with pytest.raises(ValueError, match="only supports IVF index type 'ivf_flat'"): - oceanbase_client.create_collection( - name=name, schema=_make_schema(ivf_type), use_namespace=True - ) + oceanbase_client.create_collection(name=name, schema=_make_schema(ivf_type), use_namespace=True) # Rejection happens before any DDL: no collection metadata is left behind. assert not oceanbase_client.has_collection(name) @@ -51,7 +50,9 @@ def test_ivf_flat_accepted(self, oceanbase_client): """Positive control: ivf_flat namespace collection creation succeeds.""" name = _unique_name("_ivf_flat") collection = oceanbase_client.create_collection( - name=name, schema=_make_schema("ivf_flat"), use_namespace=True, + name=name, + schema=_make_schema("ivf_flat"), + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) try: @@ -74,7 +75,10 @@ def test_dimension_2049_accepted(self, oceanbase_client): ), ) collection = oceanbase_client.create_collection( - name=name, schema=schema, use_namespace=True, partition_count=1, + name=name, + schema=schema, + use_namespace=True, + partition_count=1, ) try: assert collection.dimension == 2049 @@ -91,7 +95,10 @@ def test_dimension_4096_accepted(self, oceanbase_client): ), ) collection = oceanbase_client.create_collection( - name=name, schema=schema, use_namespace=True, partition_count=1, + name=name, + schema=schema, + use_namespace=True, + partition_count=1, ) try: assert collection.dimension == 4096 @@ -107,7 +114,9 @@ def test_dimension_4097_rejected_at_sdk(self, oceanbase_client): schema=Schema( vector_index=VectorIndexConfig( ivf=IVFConfiguration( - dimension=4097, distance="cosine", centroids_fresh_mode="spfresh", + dimension=4097, + distance="cosine", + centroids_fresh_mode="spfresh", ), embedding_function=None, ), @@ -119,6 +128,7 @@ def test_dimension_4097_rejected_at_sdk(self, oceanbase_client): class TestNamespaceMinVersionConstraint: """TestNamespaceMinVersionConstraint class.""" + def test_connected_ob_meets_min_version(self, oceanbase_client): """The kernel under test must already be >= 4.6.1, and creation succeeds.""" db_type, version = oceanbase_client._server.detect_db_type_and_version() @@ -129,7 +139,9 @@ def test_connected_ob_meets_min_version(self, oceanbase_client): name = _unique_name("_ver_ok") collection = oceanbase_client.create_collection( - name=name, schema=_make_schema("ivf_flat"), use_namespace=True, + name=name, + schema=_make_schema("ivf_flat"), + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) try: @@ -147,13 +159,11 @@ def test_old_ob_version_rejected(self, oceanbase_client, monkeypatch): ) name = _unique_name("_ver_old") with pytest.raises(ValueError, match=r"requires OceanBase version >= 4\.6\.1"): - oceanbase_client.create_collection( - name=name, schema=_make_schema("ivf_flat"), use_namespace=True - ) + oceanbase_client.create_collection(name=name, schema=_make_schema("ivf_flat"), use_namespace=True) # restore before checking leftovers monkeypatch.undo() assert not oceanbase_client.has_collection(name) def test_min_version_constant(self): """Test min version constant.""" - assert NAMESPACE_MIN_OB_VERSION == Version("4.6.1.0") + assert Version("4.6.1.0") == NAMESPACE_MIN_OB_VERSION diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py index f76f56e8..4aca9138 100644 --- a/tests/integration_tests/test_namespace_dml.py +++ b/tests/integration_tests/test_namespace_dml.py @@ -8,7 +8,6 @@ from __future__ import annotations import pytest - from namespace_dml_helpers import ( LARGE_DML_CORPUS_SIZE, assert_get_absent, @@ -18,7 +17,6 @@ cleanup, corpus_id_set, create_ns_collection, - insert_dml_corpus_in_batches, load_dml_corpus, ) @@ -36,7 +34,7 @@ class TestNamespaceDML: def test_add_single(self, db_client): """Test add single.""" collection = create_ns_collection(db_client, suffix="_add1") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0], documents="Hello", metadatas={"tag": "a"}) result = ns.get(ids="d1") @@ -71,15 +69,15 @@ def test_ops_blocked_after_collection_deleted(self, db_client): ns = _create_namespace(collection, "dml_ns") ns.add(ids="d1", embeddings=[1.0, 2.0, 3.0]) db_client.delete_collection(name=collection.name) - with pytest.raises(ValueError, match="no longer exists|does not exist"): + with pytest.raises(ValueError, match=r"no longer exists|does not exist"): ns.add(ids="d2", embeddings=[4.0, 5.0, 6.0]) - with pytest.raises(ValueError, match="no longer exists|does not exist"): + with pytest.raises(ValueError, match=r"no longer exists|does not exist"): ns.query(query_embeddings=[1.0, 2.0, 3.0], n_results=1) def test_add_batch(self, db_client): """Test add batch.""" collection = create_ns_collection(db_client, suffix="_addb") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.add( ids=["d1", "d2", "d3"], @@ -95,7 +93,7 @@ def test_add_batch(self, db_client): def test_get_by_id(self, db_client): """Test get by id.""" collection = create_ns_collection(db_client, suffix="_getid") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.add( ids=["g1", "g2"], @@ -113,7 +111,7 @@ def test_get_by_id(self, db_client): def test_get_with_limit(self, db_client): """Test get with limit.""" collection = create_ns_collection(db_client, suffix="_getlim") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.add( ids=["l1", "l2", "l3"], @@ -127,7 +125,7 @@ def test_get_with_limit(self, db_client): def test_update_metadata(self, db_client): """Test update metadata.""" collection = create_ns_collection(db_client, suffix="_upd") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.add(ids="u1", embeddings=[1.0, 2.0, 3.0], metadatas={"score": 10}) ns.update(ids="u1", metadatas={"score": 99}) @@ -139,7 +137,7 @@ def test_update_metadata(self, db_client): def test_update_document_and_embedding(self, db_client): """Test update document and embedding.""" collection = create_ns_collection(db_client, suffix="_upddoc") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.add(ids="u2", embeddings=[1.0, 2.0, 3.0], documents="Original") ns.update(ids="u2", embeddings=[9.0, 8.0, 7.0], documents="Updated") @@ -151,7 +149,7 @@ def test_update_document_and_embedding(self, db_client): def test_upsert_existing(self, db_client): """Test upsert existing.""" collection = create_ns_collection(db_client, suffix="_upsex") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.add(ids="up1", embeddings=[1.0, 2.0, 3.0], metadatas={"v": 1}) ns.upsert(ids="up1", embeddings=[4.0, 5.0, 6.0], metadatas={"v": 2}) @@ -163,7 +161,7 @@ def test_upsert_existing(self, db_client): def test_upsert_new(self, db_client): """Test upsert new.""" collection = create_ns_collection(db_client, suffix="_upsnew") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.upsert(ids="up_new", embeddings=[1.0, 1.0, 1.0], metadatas={"fresh": True}) result = ns.get(ids="up_new", include=["metadatas"]) @@ -175,7 +173,7 @@ def test_upsert_new(self, db_client): def test_delete_by_ids(self, db_client): """Test delete by ids.""" collection = create_ns_collection(db_client, suffix="_del") - ns = _create_namespace(collection,"dml_ns") + ns = _create_namespace(collection, "dml_ns") try: ns.add( ids=["del1", "del2"], @@ -195,7 +193,7 @@ class TestNamespaceDMLCountPeekAtScale: def large_ns(self, db_client): """Large ns.""" collection = create_ns_collection(db_client, suffix="_large_cp") - ns = _create_namespace(collection,"large_ns") + ns = _create_namespace(collection, "large_ns") corpus = build_large_dml_corpus(LARGE_DML_CORPUS_SIZE, id_prefix="large") load_dml_corpus(ns, corpus) yield collection, ns, corpus @@ -204,7 +202,7 @@ def large_ns(self, db_client): def test_count_empty_namespace(self, db_client): """Test count empty namespace.""" collection = create_ns_collection(db_client, suffix="_cnt_empty") - ns = _create_namespace(collection,"empty_ns") + ns = _create_namespace(collection, "empty_ns") try: assert ns.count() == 0 finally: @@ -227,7 +225,7 @@ def test_count_after_partial_delete(self, large_ns): def test_peek_empty_namespace(self, db_client): """Test peek empty namespace.""" collection = create_ns_collection(db_client, suffix="_peek_empty") - ns = _create_namespace(collection,"peek_empty") + ns = _create_namespace(collection, "peek_empty") try: result = ns.peek(limit=10) assert_peek_result(result, expected_len=0) @@ -285,14 +283,10 @@ class TestNamespaceDMLCountPeekMultiNs: def dual_ns_ctx(self, db_client): """Dual ns ctx.""" collection = create_ns_collection(db_client, suffix="_dual_cp") - ns_a = _create_namespace(collection,"ns_alpha") - ns_b = _create_namespace(collection,"ns_beta") - corpus_a = build_large_dml_corpus( - self.NS_A_SIZE, id_prefix="alpha", ns_tag="alpha" - ) - corpus_b = build_large_dml_corpus( - self.NS_B_SIZE, id_prefix="beta", ns_tag="beta" - ) + ns_a = _create_namespace(collection, "ns_alpha") + ns_b = _create_namespace(collection, "ns_beta") + corpus_a = build_large_dml_corpus(self.NS_A_SIZE, id_prefix="alpha", ns_tag="alpha") + corpus_b = build_large_dml_corpus(self.NS_B_SIZE, id_prefix="beta", ns_tag="beta") load_dml_corpus(ns_a, corpus_a) load_dml_corpus(ns_b, corpus_b) ctx = { @@ -393,7 +387,7 @@ class TestNamespaceDMLFullCycle: def test_full_dml_cycle_verified_by_get(self, db_client): """Test full dml cycle verified by get.""" collection = create_ns_collection(db_client, suffix="_cycle") - ns = _create_namespace(collection,"cycle_ns") + ns = _create_namespace(collection, "cycle_ns") doc_id = "cycle_doc" try: ns.add( @@ -447,7 +441,7 @@ def test_full_dml_cycle_verified_by_get(self, db_client): def test_batch_add_then_get_each(self, db_client): """Test batch add then get each.""" collection = create_ns_collection(db_client, suffix="_batch") - ns = _create_namespace(collection,"batch_ns") + ns = _create_namespace(collection, "batch_ns") try: ns.add( ids=["b1", "b2", "b3"], diff --git a/tests/integration_tests/test_namespace_dml_multi_coll.py b/tests/integration_tests/test_namespace_dml_multi_coll.py index 0b18413c..a7dd3bb0 100644 --- a/tests/integration_tests/test_namespace_dml_multi_coll.py +++ b/tests/integration_tests/test_namespace_dml_multi_coll.py @@ -4,7 +4,6 @@ """ import pytest - from namespace_dml_helpers import ( assert_get_absent, assert_get_present, @@ -14,8 +13,8 @@ class TestNamespaceDMLMultiColl: - """TestNamespaceDMLMultiColl class.""" + def test_same_namespace_name_across_collections(self, db_client): """Test same namespace name across collections.""" coll_a = create_ns_collection(db_client, suffix="_mc_a") @@ -38,12 +37,8 @@ def test_same_namespace_name_across_collections(self, db_client): metadatas={"coll": "B"}, ) - assert_get_present( - ns_a, "doc1", documents="Data in coll A", metadatas={"coll": "A"} - ) - assert_get_present( - ns_b, "doc1", documents="Data in coll B", metadatas={"coll": "B"} - ) + assert_get_present(ns_a, "doc1", documents="Data in coll A", metadatas={"coll": "A"}) + assert_get_present(ns_b, "doc1", documents="Data in coll B", metadatas={"coll": "B"}) finally: cleanup(db_client, coll_a, coll_b) diff --git a/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py index 89d3ce91..2cff2fcb 100644 --- a/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py @@ -4,7 +4,6 @@ """ import pytest - from namespace_dml_helpers import ( assert_get_absent, assert_get_present, @@ -14,8 +13,8 @@ class TestNamespaceDMLMultiCollMultiNs: - """TestNamespaceDMLMultiCollMultiNs class.""" + def _setup_quadrants(self, db_client): """Setup quadrants.""" coll_1 = create_ns_collection(db_client, suffix="_mcmn_c1") diff --git a/tests/integration_tests/test_namespace_dml_multi_ns.py b/tests/integration_tests/test_namespace_dml_multi_ns.py index 6380326f..ea1057c2 100644 --- a/tests/integration_tests/test_namespace_dml_multi_ns.py +++ b/tests/integration_tests/test_namespace_dml_multi_ns.py @@ -4,7 +4,6 @@ """ import pytest - from namespace_dml_helpers import ( assert_get_absent, assert_get_present, @@ -14,8 +13,8 @@ class TestNamespaceDMLMultiNs: - """TestNamespaceDMLMultiNs class.""" + def test_add_isolation(self, db_client): """Test add isolation.""" collection = create_ns_collection(db_client, suffix="_mns_add") diff --git a/tests/integration_tests/test_namespace_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py index a6ed5c37..c4f679f9 100644 --- a/tests/integration_tests/test_namespace_drop_validation.py +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -1,18 +1,18 @@ """test namespace drop validation module.""" +import contextlib import threading import time import pytest +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT import pyseekdb -from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.meta_info import NamespaceCollectionNames from pyseekdb.client.schema import Schema - RECYCLEBIN_PREFIX = "__recyclebin_" # LTABLE_BG scans every 30s; poll slightly beyond two intervals. @@ -43,7 +43,9 @@ def _make_collection(client, suffix: str = ""): ), ) return client.create_collection( - name=name, schema=schema, use_namespace=True, + name=name, + schema=schema, + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) @@ -153,8 +155,9 @@ def _count_kv_data_rows(client, collection_id: str, namespace_id: int) -> int: return int(rows[0]["c"] if "c" in rows[0] else rows[0]["C"]) -def _wait_until(predicate, timeout_sec: float = BG_POLL_TIMEOUT_SEC, - interval_sec: float = BG_POLL_INTERVAL_SEC, desc: str = ""): +def _wait_until( + predicate, timeout_sec: float = BG_POLL_TIMEOUT_SEC, interval_sec: float = BG_POLL_INTERVAL_SEC, desc: str = "" +): """Wait until.""" deadline = time.time() + timeout_sec last_exc = None @@ -162,7 +165,7 @@ def _wait_until(predicate, timeout_sec: float = BG_POLL_TIMEOUT_SEC, try: if predicate(): return - except Exception as exc: # noqa: BLE001 + except Exception as exc: last_exc = exc time.sleep(interval_sec) msg = f"timeout after {timeout_sec}s waiting for {desc!r}" @@ -181,8 +184,7 @@ def _seed_kv_data(client, collection_id: str, namespace_id: int, count: int = 5) values.append(f"({namespace_id}, X'{key_hex}', X'00')") _execute( client, - f"INSERT INTO `{db}`.`{tbl}` (namespace_id, kv_key, kv_value) " - f"VALUES {', '.join(values)}", + f"INSERT INTO `{db}`.`{tbl}` (namespace_id, kv_key, kv_value) VALUES {', '.join(values)}", ) @@ -195,7 +197,7 @@ def _seed_logic_data(client, collection_id: str, namespace_id: int, ltable_id: i values.append( f"({namespace_id}, {ltable_id}, 'doc_{i}', " f"X'0000803f0000000000000000', " - f"'{{\"id\": \"id_{i}\", \"metadata\": {{\"k\": \"v\"}}}}')" + f'\'{{"id": "id_{i}", "metadata": {{"k": "v"}}}}\')' ) _execute( client, @@ -225,8 +227,7 @@ def _hot_table_exists(client, collection_id: str) -> bool: try: rows = _execute( client, - "SELECT 1 AS ok FROM information_schema.tables " - f"WHERE table_schema = '{db}' AND table_name = '{tbl}'", + f"SELECT 1 AS ok FROM information_schema.tables WHERE table_schema = '{db}' AND table_name = '{tbl}'", ) return bool(rows) except Exception: @@ -267,7 +268,7 @@ def _ensure_hot_table(client, collection_id: str) -> str | None: _execute(client, ddl) if _hot_table_exists(client, collection_id): return None - except Exception as exc: # noqa: BLE001 + except Exception as exc: last_err = str(exc) return last_err or "hot_table not visible after CREATE" @@ -281,8 +282,7 @@ def _seed_hot_table(client, collection_id: str, namespace_id: int): try: _execute( client, - f"INSERT INTO `{db}`.`{tbl}` (namespace_id, last_access_time) " - f"VALUES ({namespace_id}, NOW(6))", + f"INSERT INTO `{db}`.`{tbl}` (namespace_id, last_access_time) VALUES ({namespace_id}, NOW(6))", ) return True except Exception: @@ -377,15 +377,13 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = _create_namespace(collection,"ns_validate") + ns = _create_namespace(collection, "ns_validate") ns_id = int(ns.namespace_id) # Pre-conditions orig_name = _fetch_namespace_name(client, coll_id, ns_id) assert orig_name == "ns_validate" - assert _count_ltables(client, coll_id, ns_id) >= 1, ( - "create_namespace should have inserted a default ltable" - ) + assert _count_ltables(client, coll_id, ns_id) >= 1, "create_namespace should have inserted a default ltable" assert _count_logic_schema_rows(client, coll_id, ns_id) >= 1, ( "create_namespace should have inserted a logic_schema row" ) @@ -408,19 +406,11 @@ def test_drop_namespace_renames_and_cleans_catalog(self, oceanbase_client): assert new_name.startswith(RECYCLEBIN_PREFIX), ( f"namespace must be renamed to __recyclebin_ prefix, got {new_name!r}" ) - assert "ns_validate" in new_name, ( - f"original name should be embedded in recyclebin name, got {new_name!r}" - ) - assert _count_ltables(client, coll_id, ns_id) == 0, ( - "sdk_ltables rows must be deleted" - ) - assert _count_logic_schema_rows(client, coll_id, ns_id) == 0, ( - "logic_schema_table rows must be deleted" - ) + assert "ns_validate" in new_name, f"original name should be embedded in recyclebin name, got {new_name!r}" + assert _count_ltables(client, coll_id, ns_id) == 0, "sdk_ltables rows must be deleted" + assert _count_logic_schema_rows(client, coll_id, ns_id) == 0, "logic_schema_table rows must be deleted" if is_ss: - assert _count_hot_table_rows(client, coll_id, ns_id) == 0, ( - "hot_table row must be deleted in SS mode" - ) + assert _count_hot_table_rows(client, coll_id, ns_id) == 0, "hot_table row must be deleted in SS mode" # Logic data rows still exist — async background task handles physical cleanup. assert _count_logic_data_rows(client, coll_id, ns_id, lt_id) == 3, ( @@ -438,7 +428,7 @@ def test_drop_namespace_async_cleanup(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = _create_namespace(collection,"ns_async") + ns = _create_namespace(collection, "ns_async") ns_id = int(ns.namespace_id) lt_id = _fetch_ltable_id(client, coll_id, ns_id) assert lt_id is not None @@ -479,7 +469,7 @@ def test_async_skipped_while_active_ltable_remains(self, oceanbase_client): coll_id = collection.id blocker_lt_id = 9_000_001 try: - ns = _create_namespace(collection,"ns_block") + ns = _create_namespace(collection, "ns_block") ns_id = int(ns.namespace_id) _drop_namespace_via_pl(client, coll_id, ns_id) assert _fetch_namespace_name(client, coll_id, ns_id).startswith(RECYCLEBIN_PREFIX) @@ -487,7 +477,11 @@ def test_async_skipped_while_active_ltable_remains(self, oceanbase_client): # Simulate a stuck active ltable row (sync phase normally clears these). _insert_active_ltable_row( - client, coll_id, ns_id, "active_ltable_blocker", blocker_lt_id, + client, + coll_id, + ns_id, + "active_ltable_blocker", + blocker_lt_id, ) assert _count_active_ltables(client, coll_id, ns_id) == 1 @@ -499,8 +493,7 @@ def test_async_skipped_while_active_ltable_remains(self, oceanbase_client): _execute( client, - f"DELETE FROM {_catalog_table(client, 'sdk_ltables')} " - f"WHERE ltable_id = {blocker_lt_id}", + f"DELETE FROM {_catalog_table(client, 'sdk_ltables')} WHERE ltable_id = {blocker_lt_id}", ) assert _count_active_ltables(client, coll_id, ns_id) == 0 @@ -517,7 +510,7 @@ def test_async_no_kv_data_fast_cleanup(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = _create_namespace(collection,"ns_empty_kv") + ns = _create_namespace(collection, "ns_empty_kv") ns_id = int(ns.namespace_id) assert _count_kv_data_rows(client, coll_id, ns_id) == 0 @@ -548,9 +541,7 @@ def test_async_hot_table_cleaned_on_non_ss(self, oceanbase_client): _drop_namespace_via_pl(client, coll_id, ns_id) assert _fetch_namespace_name(client, coll_id, ns_id).startswith(RECYCLEBIN_PREFIX) - assert _count_hot_table_rows(client, coll_id, ns_id) == 1, ( - "sync DROP must not remove hot_table on non-SS" - ) + assert _count_hot_table_rows(client, coll_id, ns_id) == 1, "sync DROP must not remove hot_table on non-SS" def _hot_and_ns_gone(): """Hot and ns gone.""" @@ -569,7 +560,7 @@ def test_async_drop_history_recorded(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = _create_namespace(collection,"ns_history") + ns = _create_namespace(collection, "ns_history") ns_id = int(ns.namespace_id) _seed_kv_data(client, coll_id, ns_id, count=2) _drop_namespace_via_pl(client, coll_id, ns_id) @@ -600,16 +591,14 @@ def test_recyclebin_name_format(self, oceanbase_client): client = oceanbase_client collection = _make_collection(client) try: - ns = _create_namespace(collection,"fmt_ns") + ns = _create_namespace(collection, "fmt_ns") ns_id = int(ns.namespace_id) _drop_namespace_via_pl(client, collection.id, ns_id) new_name = _fetch_namespace_name(client, collection.id, ns_id) # Expected format: __recyclebin__ assert new_name.startswith("__recyclebin_fmt_ns_") - ts_suffix = new_name[len("__recyclebin_fmt_ns_"):] - assert ts_suffix.isdigit() and len(ts_suffix) >= 16, ( - f"trailing timestamp_us looks malformed: {new_name!r}" - ) + ts_suffix = new_name[len("__recyclebin_fmt_ns_") :] + assert ts_suffix.isdigit() and len(ts_suffix) >= 16, f"trailing timestamp_us looks malformed: {new_name!r}" finally: client.delete_collection(name=collection.name) @@ -621,7 +610,7 @@ def test_multiple_ltables_all_deleted(self, oceanbase_client): client = oceanbase_client collection = _make_collection(client) try: - ns = _create_namespace(collection,"ns_multi") + ns = _create_namespace(collection, "ns_multi") ns_id = int(ns.namespace_id) # Insert two extra ltable rows directly so we can verify bulk delete. lt_table = _catalog_table(client, "sdk_ltables") @@ -653,7 +642,7 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): schema_tbl = NamespaceCollectionNames.logic_schema_table_name(coll_id) schema_tbl_q = f"`{client._server.database}`.`{schema_tbl}`" try: - ns = _create_namespace(collection,"ns_rollback") + ns = _create_namespace(collection, "ns_rollback") ns_id = int(ns.namespace_id) orig_name = _fetch_namespace_name(client, coll_id, ns_id) assert orig_name == "ns_rollback" @@ -669,16 +658,12 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): # Everything before the failed step must have been rolled back. after_name = _fetch_namespace_name(client, coll_id, ns_id) - assert after_name == "ns_rollback", ( - f"namespace_name must be rolled back to original, got {after_name!r}" - ) - assert _count_ltables(client, coll_id, ns_id) == ltable_cnt_before, ( - "sdk_ltables delete must be rolled back" - ) + assert after_name == "ns_rollback", f"namespace_name must be rolled back to original, got {after_name!r}" + assert _count_ltables(client, coll_id, ns_id) == ltable_cnt_before, "sdk_ltables delete must be rolled back" finally: # Recreate the schema table (empty) so cleanup_namespace_physical_tables # / delete_collection runs cleanly. - try: + with contextlib.suppress(Exception): _execute( client, f"CREATE TABLE IF NOT EXISTS {schema_tbl_q} (" @@ -688,8 +673,6 @@ def test_atomic_rollback_on_logic_schema_missing(self, oceanbase_client): f" PRIMARY KEY (namespace_id, ltable_id)) " f"PARTITION BY KEY(namespace_id) PARTITIONS 8", ) - except Exception: - pass client.delete_collection(name=collection.name) # ------------------------------------------------------------- # @@ -701,7 +684,7 @@ def test_drop_namespace_is_idempotent(self, oceanbase_client): collection = _make_collection(client) coll_id = collection.id try: - ns = _create_namespace(collection,"ns_idem") + ns = _create_namespace(collection, "ns_idem") ns_id = int(ns.namespace_id) _drop_namespace_via_pl(client, coll_id, ns_id) first_name = _fetch_namespace_name(client, coll_id, ns_id) @@ -736,7 +719,7 @@ def test_concurrent_drop_from_two_sdks(self, oceanbase_client): collection = _make_collection(client_a, suffix="_conc") coll_id = collection.id try: - ns = _create_namespace(collection,"ns_concurrent") + ns = _create_namespace(collection, "ns_concurrent") ns_id = int(ns.namespace_id) results = {} @@ -748,7 +731,7 @@ def _worker(tag, c): barrier.wait(timeout=10) _drop_namespace_via_pl(c, coll_id, ns_id) results[tag] = "ok" - except Exception as exc: # noqa: BLE001 + except Exception as exc: results[tag] = f"err: {exc!r}" ta = threading.Thread(target=_worker, args=("A", client_a)) @@ -777,10 +760,8 @@ def _worker(tag, c): client_a.delete_collection(name=collection.name) finally: if hasattr(client_b, "close"): - try: + with contextlib.suppress(Exception): client_b.close() - except Exception: - pass if __name__ == "__main__": diff --git a/tests/integration_tests/test_namespace_get_delete_filters.py b/tests/integration_tests/test_namespace_get_delete_filters.py index 5a076c16..b496e472 100644 --- a/tests/integration_tests/test_namespace_get_delete_filters.py +++ b/tests/integration_tests/test_namespace_get_delete_filters.py @@ -11,7 +11,6 @@ from __future__ import annotations import pytest - from namespace_dml_helpers import ( FILTER_MULTI_NS_KEEP, FILTER_MULTI_NS_PURGE, @@ -28,7 +27,7 @@ class TestNamespaceDeleteWhere: - """P0 #1–#2: conditional delete on namespace.""" + """P0 #1-#2: conditional delete on namespace.""" def test_delete_where_metadata_tag(self, db_client): """Test delete where metadata tag.""" @@ -93,7 +92,7 @@ def test_delete_where_document_obsolete(self, db_client): class TestNamespaceGetWhere: - """P0 #3–#4: conditional get on namespace.""" + """P0 #3-#4: conditional get on namespace.""" def test_get_where_metadata_category_ai(self, db_client): """Test get where metadata category ai.""" diff --git a/tests/integration_tests/test_namespace_hybrid_search_combined.py b/tests/integration_tests/test_namespace_hybrid_search_combined.py index b80fdeae..977c59c4 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_combined.py +++ b/tests/integration_tests/test_namespace_hybrid_search_combined.py @@ -16,7 +16,6 @@ from typing import Any, ClassVar import pytest - from namespace_hybrid_search_helpers import ( HYBRID_COMBINED_CASES, VectorDistanceMetric, @@ -31,6 +30,7 @@ @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) class TestNamespaceHybridSearchCombined: """TestNamespaceHybridSearchCombined class.""" + _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} @pytest.fixture(autouse=True) @@ -41,9 +41,7 @@ def _bind_shared_collection( request: pytest.FixtureRequest, ) -> None: """Bind shared collection.""" - entry = ensure_shared_hybrid_search_collection( - self._shared_by_mode, db_client, request, vector_distance - ) + entry = ensure_shared_hybrid_search_collection(self._shared_by_mode, db_client, request, vector_distance) self._corpus = entry["corpus"] self._collection = entry["collection"] self._vector_distance = vector_distance @@ -68,9 +66,7 @@ def test_hybrid_search_combined(self, db_client, case_name: str): """Test hybrid search combined.""" case = get_hybrid_combined_case(case_name) namespace = self._new_namespace(case_name) - run_hybrid_combined_case( - namespace, self._corpus, case, distance_metric=self._vector_distance - ) + run_hybrid_combined_case(namespace, self._corpus, case, distance_metric=self._vector_distance) if __name__ == "__main__": diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext.py index 19375cc1..ae5ad05f 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_fulltext.py +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext.py @@ -14,12 +14,11 @@ from typing import Any, ClassVar import pytest - from namespace_fts_helpers import ( TOKEN_ZPX, CorpusRecord, - get_fts_case, assert_not_contains_no_token_leak, + get_fts_case, run_hybrid_search_fts_case, setup_fts_namespace_with_corpus, setup_large_fts_collection, @@ -129,9 +128,7 @@ def test_hybrid_search_fulltext_and_zpx_alp_filter_has_both(self, db_client): """``$and`` on ``where_document`` combined with ``has_both`` metadata filter.""" self._run_fts_case("and_zpx_alp_filter_has_both") - def test_hybrid_search_fulltext_contains_zpx_filter_zpx_hint_and_alp_zero( - self, db_client - ): + def test_hybrid_search_fulltext_contains_zpx_filter_zpx_hint_and_alp_zero(self, db_client): """``where_document`` + ``where`` ``$and``: ``zpx_hint`` range and ``alp_hint`` equality.""" self._run_fts_case("contains_zpx_filter_zpx_hint_and_alp_zero") diff --git a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py index 8826c4f1..7824d774 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py @@ -16,7 +16,6 @@ import time import pytest - from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from namespace_fts_helpers import ( CORPUS_SIZE, @@ -181,8 +180,7 @@ def test_hybrid_search_fulltext_top1_contains_zpx(self, db_client): include=["documents"], ) assert top_result["ids"][0][0] == f"{key}_zpx_top_5", ( - f"[{key}] most relevant TOKEN_ZPX document must rank first, " - f"got {top_result['ids'][0]}" + f"[{key}] most relevant TOKEN_ZPX document must rank first, got {top_result['ids'][0]}" ) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) diff --git a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py index 2c3fe958..7067c5cc 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py @@ -7,7 +7,6 @@ from __future__ import annotations import pytest - from namespace_fts_helpers import ( MULTI_COLL_MULTI_NS_QUADRANT_KEYS, VectorDistanceMetric, @@ -22,8 +21,8 @@ get_search_index_case, get_vector_knn_case, run_hybrid_combined_case, - run_hybrid_search_index_case, run_hybrid_knn_case, + run_hybrid_search_index_case, run_knn_case_on_quadrants, run_search_index_case_on_quadrants, setup_multi_coll_multi_ns_fts, @@ -34,6 +33,7 @@ class TestNamespaceHybridSearchMultiCollMultiNs: """TestNamespaceHybridSearchMultiCollMultiNs class.""" + def test_cross_quadrant_search_index_isolation(self, db_client): """Only the loaded quadrant returns ``has_both`` rows.""" ctx = setup_multi_coll_multi_ns_fts_single_loaded(db_client, loaded_quadrant="c1_x") @@ -60,9 +60,7 @@ def test_cross_quadrant_fts_isolation(self, db_client): if key == "c1_x": continue _, namespace = ctx[key] - assert_hybrid_search_no_hits( - namespace, case.where_document, n_results=case.n_results - ) + assert_hybrid_search_no_hits(namespace, case.where_document, n_results=case.n_results) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) @@ -122,41 +120,31 @@ def test_hybrid_search_fts_plus_search_index_multi_coll(self, db_client): try: for key in ("c1_x", "c2_y"): corpus, namespace = ctx[key] - run_hybrid_search_fts_case( - namespace, corpus, get_fts_case("contains_zpx_filter_has_both") - ) + run_hybrid_search_fts_case(namespace, corpus, get_fts_case("contains_zpx_filter_has_both")) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) - def test_hybrid_search_combined_rrf_multi_coll( - self, db_client, vector_distance: VectorDistanceMetric - ): + def test_hybrid_search_combined_rrf_multi_coll(self, db_client, vector_distance: VectorDistanceMetric): """Test hybrid search combined rrf multi coll.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: case = get_hybrid_combined_case("fts_zpx_filter_gte_knn") for key in ("c1_x", "c2_y"): corpus, namespace = ctx[key] - run_hybrid_combined_case( - namespace, corpus, case, distance_metric=vector_distance - ) + run_hybrid_combined_case(namespace, corpus, case, distance_metric=vector_distance) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) - def test_hybrid_search_vector_plus_search_index_multi_coll( - self, db_client, vector_distance: VectorDistanceMetric - ): + def test_hybrid_search_vector_plus_search_index_multi_coll(self, db_client, vector_distance: VectorDistanceMetric): """Test hybrid search vector plus search index multi coll.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: knn_case = get_vector_knn_case("knn_filter_has_both") for key in ("c1_x", "c2_y"): corpus, namespace = ctx[key] - run_hybrid_knn_case( - namespace, corpus, knn_case, distance_metric=vector_distance - ) + run_hybrid_knn_case(namespace, corpus, knn_case, distance_metric=vector_distance) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) diff --git a/tests/integration_tests/test_namespace_hybrid_search_search_index.py b/tests/integration_tests/test_namespace_hybrid_search_search_index.py index b716ac63..14a90fdb 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_search_index.py +++ b/tests/integration_tests/test_namespace_hybrid_search_search_index.py @@ -10,7 +10,6 @@ from typing import Any, ClassVar import pytest - from namespace_hybrid_search_helpers import ( SEARCH_INDEX_CASES, get_search_index_case, @@ -57,9 +56,7 @@ def teardown_class(cls) -> None: def _run_case(self, case_name: str) -> None: """Run case.""" - run_hybrid_search_index_case( - self._namespace, self._corpus, get_search_index_case(case_name) - ) + run_hybrid_search_index_case(self._namespace, self._corpus, get_search_index_case(case_name)) @pytest.mark.parametrize("case_name", [c.name for c in SEARCH_INDEX_CASES]) def test_hybrid_search_search_index_operators(self, db_client, case_name: str): diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py index 55a76796..fa6fef01 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py @@ -25,10 +25,9 @@ from typing import Any, ClassVar import pytest - from namespace_hybrid_search_helpers import ( - TRIPLE_BRANCH_CASES, TOKEN_ZPX, + TRIPLE_BRANCH_CASES, VectorDistanceMetric, ensure_shared_hybrid_search_collection, get_triple_branch_case, @@ -52,9 +51,7 @@ def _bind_shared_collection( request: pytest.FixtureRequest, ) -> None: """Bind shared collection.""" - entry = ensure_shared_hybrid_search_collection( - self._shared_by_mode, db_client, request, vector_distance - ) + entry = ensure_shared_hybrid_search_collection(self._shared_by_mode, db_client, request, vector_distance) self._corpus = entry["corpus"] self._collection = entry["collection"] self._vector_distance = vector_distance @@ -107,9 +104,7 @@ def test_hybrid_search_triple_branch_fts_hits_with_knn_active(self, db_client): assert len(ids) > 0, "expected at least one hybrid FTS+KNN hit" docs = result.get("documents", [[]])[0] for doc_id, doc_text in zip(ids, docs, strict=True): - assert TOKEN_ZPX.lower() in (doc_text or "").lower(), ( - f"id={doc_id!r} must contain {TOKEN_ZPX!r}" - ) + assert TOKEN_ZPX.lower() in (doc_text or "").lower(), f"id={doc_id!r} must contain {TOKEN_ZPX!r}" if __name__ == "__main__": diff --git a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py index 12200003..db26f88d 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py @@ -15,7 +15,6 @@ import time import pytest - from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from namespace_fts_helpers import ( CORPUS_SIZE, @@ -47,18 +46,10 @@ _FTS_TRIPLE_CASE_NAMES: tuple[str, ...] = tuple( c.name for c in TRIPLE_BRANCH_CASES if c.verify in ("fts", "not_contains") ) -_SI_TRIPLE_CASE_NAMES: tuple[str, ...] = tuple( - c.name for c in TRIPLE_BRANCH_CASES if c.verify == "search_index" -) -_KNN_TRIPLE_CASE_NAMES: tuple[str, ...] = tuple( - c.name for c in TRIPLE_BRANCH_CASES if c.verify == "knn" -) -_INTERSECTION_CASE_NAMES: tuple[str, ...] = tuple( - c.name for c in TRIPLE_BRANCH_CASES if c.verify == "intersection" -) -_RRF_CASE_NAMES: tuple[str, ...] = tuple( - c.name for c in TRIPLE_BRANCH_CASES if c.verify == "rrf_fusion" -) +_SI_TRIPLE_CASE_NAMES: tuple[str, ...] = tuple(c.name for c in TRIPLE_BRANCH_CASES if c.verify == "search_index") +_KNN_TRIPLE_CASE_NAMES: tuple[str, ...] = tuple(c.name for c in TRIPLE_BRANCH_CASES if c.verify == "knn") +_INTERSECTION_CASE_NAMES: tuple[str, ...] = tuple(c.name for c in TRIPLE_BRANCH_CASES if c.verify == "intersection") +_RRF_CASE_NAMES: tuple[str, ...] = tuple(c.name for c in TRIPLE_BRANCH_CASES if c.verify == "rrf_fusion") def _namespace_corpus(base_corpus: list[CorpusRecord], key: str) -> list[CorpusRecord]: @@ -120,14 +111,11 @@ def _setup_multi_coll_multi_ns_isolation_triple( class TestNamespaceHybridSearchTripleBranchMultiCollMultiNs: """TestNamespaceHybridSearchTripleBranchMultiCollMultiNs class.""" + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) - def test_cross_quadrant_triple_branch_isolation( - self, db_client, vector_distance: VectorDistanceMetric - ): + def test_cross_quadrant_triple_branch_isolation(self, db_client, vector_distance: VectorDistanceMetric): """Loaded namespaces return only their own rows; empty namespaces stay empty.""" - ctx = _setup_multi_coll_multi_ns_isolation_triple( - db_client, distance=vector_distance - ) + ctx = _setup_multi_coll_multi_ns_isolation_triple(db_client, distance=vector_distance) try: case = get_triple_branch_case("intersection_both_tokens") for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: @@ -166,13 +154,9 @@ def test_cross_quadrant_search_index_branch_isolation(self, db_client): teardown_multi_coll_multi_ns_fts(db_client, ctx) @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) - def test_triple_branch_top1_fts_all_quadrants( - self, db_client, vector_distance: VectorDistanceMetric - ): + def test_triple_branch_top1_fts_all_quadrants(self, db_client, vector_distance: VectorDistanceMetric): """Test triple branch top1 fts all quadrants.""" - ctx = _setup_multi_coll_multi_ns_isolation_triple( - db_client, distance=vector_distance - ) + ctx = _setup_multi_coll_multi_ns_isolation_triple(db_client, distance=vector_distance) try: for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: _, namespace = ctx[key] @@ -206,9 +190,7 @@ def test_hybrid_search_triple_branch_fts_all_loaded_quadrants(self, db_client, c teardown_multi_coll_multi_ns_fts(db_client, ctx) @pytest.mark.parametrize("case_name", _SI_TRIPLE_CASE_NAMES) - def test_hybrid_search_triple_branch_search_index_all_loaded_quadrants( - self, db_client, case_name: str - ): + def test_hybrid_search_triple_branch_search_index_all_loaded_quadrants(self, db_client, case_name: str): """Test hybrid search triple branch search index all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client) try: @@ -224,9 +206,7 @@ def test_hybrid_search_triple_branch_knn_all_loaded_quadrants( """Test hybrid search triple branch knn all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: - run_triple_branch_case_on_quadrants( - ctx, case_name, distance_metric=vector_distance - ) + run_triple_branch_case_on_quadrants(ctx, case_name, distance_metric=vector_distance) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) @@ -238,9 +218,7 @@ def test_hybrid_search_triple_branch_intersection_all_loaded_quadrants( """Test hybrid search triple branch intersection all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: - run_triple_branch_case_on_quadrants( - ctx, case_name, distance_metric=vector_distance - ) + run_triple_branch_case_on_quadrants(ctx, case_name, distance_metric=vector_distance) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) @@ -252,9 +230,7 @@ def test_hybrid_search_triple_branch_rrf_all_loaded_quadrants( """Test hybrid search triple branch rrf all loaded quadrants.""" ctx = setup_multi_coll_multi_ns_fts(db_client, distance=vector_distance) try: - run_triple_branch_case_on_quadrants( - ctx, case_name, distance_metric=vector_distance - ) + run_triple_branch_case_on_quadrants(ctx, case_name, distance_metric=vector_distance) finally: teardown_multi_coll_multi_ns_fts(db_client, ctx) diff --git a/tests/integration_tests/test_namespace_hybrid_search_vector.py b/tests/integration_tests/test_namespace_hybrid_search_vector.py index 086ce2b5..0a2b06ea 100644 --- a/tests/integration_tests/test_namespace_hybrid_search_vector.py +++ b/tests/integration_tests/test_namespace_hybrid_search_vector.py @@ -11,7 +11,6 @@ from typing import Any, ClassVar import pytest - from namespace_hybrid_search_helpers import ( KNN_QUERY_VECTOR, VECTOR_KNN_CASES, @@ -39,9 +38,7 @@ def _bind_shared_collection( request: pytest.FixtureRequest, ) -> None: """Bind shared collection.""" - entry = ensure_shared_hybrid_search_collection( - self._shared_by_mode, db_client, request, vector_distance - ) + entry = ensure_shared_hybrid_search_collection(self._shared_by_mode, db_client, request, vector_distance) self._corpus = entry["corpus"] self._collection = entry["collection"] self._vector_distance = vector_distance @@ -85,12 +82,15 @@ def test_hybrid_search_vector_top1_nearest(self, db_client): n_results=1, include=["distances"], ) - assert result["ids"][0][0] == expected_knn_ids( - self._corpus, - KNN_QUERY_VECTOR, - 1, - distance_metric=self._vector_distance, - )[0] + assert ( + result["ids"][0][0] + == expected_knn_ids( + self._corpus, + KNN_QUERY_VECTOR, + 1, + distance_metric=self._vector_distance, + )[0] + ) if __name__ == "__main__": diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index 2d9bc391..a0f7fb96 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -7,18 +7,17 @@ import time import pytest +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT import pyseekdb -from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema - class TestNamespaceLifecycle: - """TestNamespaceLifecycle class.""" + def _create_ns_collection(self, client, suffix=""): """Create ns collection.""" name = f"test_ns_lc_{int(time.time() * 1000)}{suffix}" @@ -29,7 +28,9 @@ def _create_ns_collection(self, client, suffix=""): ), ) collection = client.create_collection( - name=name, schema=schema, use_namespace=True, + name=name, + schema=schema, + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) return collection @@ -66,8 +67,10 @@ def test_get_collection_restores_embedding_function(self, db_client): embedding_function=ef, ), ) - collection = db_client.create_collection( - name=name, schema=schema, use_namespace=True, + db_client.create_collection( + name=name, + schema=schema, + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) try: @@ -165,7 +168,7 @@ def test_prewarm_after_namespace_deleted_raises_friendly_error(self, db_client): ns = collection.create_namespace("prewarm_del") try: collection.delete_namespace("prewarm_del") - with pytest.raises(ValueError, match="no longer exists|being dropped"): + with pytest.raises(ValueError, match=r"no longer exists|being dropped"): ns.prewarm() finally: db_client.delete_collection(name=collection.name) @@ -194,9 +197,7 @@ def test_create_resumes_incomplete_collection(self, oceanbase_client): embedding_function=None, ), ) - collection = client.create_collection( - name=name, schema=schema, use_namespace=True, partition_count=4 - ) + collection = client.create_collection(name=name, schema=schema, use_namespace=True, partition_count=4) try: cid = collection.id # Simulate an interrupted create: drop one physical table. @@ -205,9 +206,7 @@ def test_create_resumes_incomplete_collection(self, oceanbase_client): assert srv._is_incomplete_ns_collection(name) is True # Re-create resumes (same id, rebuilds the missing table) instead of raising. - resumed = client.create_collection( - name=name, schema=schema, use_namespace=True, partition_count=4 - ) + resumed = client.create_collection(name=name, schema=schema, use_namespace=True, partition_count=4) assert resumed.id == cid assert srv._is_incomplete_ns_collection(name) is False @@ -218,9 +217,7 @@ def test_create_resumes_incomplete_collection(self, oceanbase_client): # A complete collection still rejects a duplicate create. with pytest.raises(ValueError, match="already exists"): - client.create_collection( - name=name, schema=schema, use_namespace=True, partition_count=4 - ) + client.create_collection(name=name, schema=schema, use_namespace=True, partition_count=4) finally: client.delete_collection(name=name) @@ -238,9 +235,7 @@ def test_get_collection_purges_broken_ns_collection(self, oceanbase_client): embedding_function=None, ), ) - collection = client.create_collection( - name=name, schema=schema, use_namespace=True, partition_count=4 - ) + collection = client.create_collection(name=name, schema=schema, use_namespace=True, partition_count=4) cid = collection.id try: srv._use_catalog_database() @@ -299,14 +294,12 @@ def test_namespace_on_non_namespace_collection_raises(self, db_client): db_client.delete_collection(name=name) except (ValueError, RuntimeError): db_client._server._execute(f"DROP TABLE IF EXISTS `{name}`") - db_client._server._execute( - f"DELETE FROM `sdk_collections` WHERE COLLECTION_NAME = '{name}'" - ) - + db_client._server._execute(f"DELETE FROM `sdk_collections` WHERE COLLECTION_NAME = '{name}'") def test_create_namespace_collection_with_hnsw_raises(self, db_client): """Test create namespace collection with hnsw raises.""" from pyseekdb.client.configuration import HNSWConfiguration + name = f"test_ns_hnsw_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( @@ -321,13 +314,12 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): """Verify hot_table is created when _is_shared_storage_mode returns True (SS mode).""" import json from unittest.mock import patch + from pyseekdb.client.meta_info import NamespaceCollectionNames client = oceanbase_client - with patch.object( - type(client._server), "_is_shared_storage_mode", return_value=True - ): + with patch.object(type(client._server), "_is_shared_storage_mode", return_value=True): name = f"test_ns_ss_{int(time.time() * 1000)}" schema = Schema( vector_index=VectorIndexConfig( @@ -336,7 +328,9 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): ), ) collection = client.create_collection( - name=name, schema=schema, use_namespace=True, + name=name, + schema=schema, + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) @@ -361,11 +355,11 @@ def test_ss_mode_creates_hot_table(self, oceanbase_client): finally: client.delete_collection(name=name) - def test_custom_collection_partition_count(self, oceanbase_client): """Verify create_collection(partition_count=...) controls the PARTITIONS clause and is exposed via collection.partition_count.""" import re + from pyseekdb.client.meta_info import NamespaceCollectionNames name = f"test_ns_lc_pc_{int(time.time() * 1000)}" @@ -375,9 +369,7 @@ def test_custom_collection_partition_count(self, oceanbase_client): embedding_function=None, ), ) - collection = oceanbase_client.create_collection( - name=name, schema=schema, use_namespace=True, partition_count=4 - ) + collection = oceanbase_client.create_collection(name=name, schema=schema, use_namespace=True, partition_count=4) try: assert collection.partition_count == 4 # Reopened handle restores partition_count from settings. @@ -387,9 +379,7 @@ def test_custom_collection_partition_count(self, oceanbase_client): rows = oceanbase_client._server._execute(f"SHOW CREATE TABLE `{data_table}`") create_sql = rows[0].get("Create Table", "") if rows else "" partitions = re.findall(r"partition `p\d+`", create_sql) - assert len(partitions) == 4, ( - f"Expected 4 partitions, found {len(partitions)}: {create_sql[-300:]}" - ) + assert len(partitions) == 4, f"Expected 4 partitions, found {len(partitions)}: {create_sql[-300:]}" finally: oceanbase_client.delete_collection(name=collection.name) diff --git a/tests/integration_tests/test_namespace_lifecycle_concurrency.py b/tests/integration_tests/test_namespace_lifecycle_concurrency.py index 67a8ff81..c8e2c71c 100644 --- a/tests/integration_tests/test_namespace_lifecycle_concurrency.py +++ b/tests/integration_tests/test_namespace_lifecycle_concurrency.py @@ -9,9 +9,9 @@ from unittest.mock import patch import pytest +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, cleanup, ns_schema import pyseekdb -from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, cleanup, ns_schema def _unique_name(prefix: str) -> str: @@ -52,7 +52,7 @@ def _drop_worker() -> None: """Drop worker.""" try: clients[0].get_collection(collection.name).delete_namespace(ns_name) - except Exception as exc: # noqa: BLE001 + except Exception as exc: errors.append(exc) finally: drop_done.set() @@ -64,7 +64,7 @@ def _has_worker() -> None: drop_done.wait(timeout=60) for _ in range(30): post_drop_has.append(coll.has_namespace(ns_name)) - except Exception as exc: # noqa: BLE001 + except Exception as exc: errors.append(exc) drop_thread = threading.Thread(target=_drop_worker) @@ -120,7 +120,7 @@ def _drop_worker() -> None: try: with patch.object(clients[0]._server, "_execute", side_effect=_slow_drop_execute): dropper.delete_namespace(ns_name) - except Exception as exc: # noqa: BLE001 + except Exception as exc: errors.append(exc) def _has_worker() -> None: @@ -131,7 +131,7 @@ def _has_worker() -> None: has_results.append(checker.has_namespace(ns_name)) has_completed_while_drop_blocked.set() time.sleep(0.05) - except Exception as exc: # noqa: BLE001 + except Exception as exc: errors.append(exc) drop_thread = threading.Thread(target=_drop_worker) diff --git a/tests/integration_tests/test_namespace_optional_indexes.py b/tests/integration_tests/test_namespace_optional_indexes.py index d9e4c1f2..db064223 100644 --- a/tests/integration_tests/test_namespace_optional_indexes.py +++ b/tests/integration_tests/test_namespace_optional_indexes.py @@ -20,8 +20,8 @@ import pymysql.err import pytest - from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT + from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.meta_info import NamespaceCollectionNames @@ -60,9 +60,7 @@ def _index_names(client: Any, collection_id: str) -> set[str]: def _collection_settings(client: Any, collection_id: str) -> dict: """Collection settings.""" - rows = client._server._execute( - f"SELECT settings FROM sdk_collections WHERE collection_id = '{collection_id}'" - ) + rows = client._server._execute(f"SELECT settings FROM sdk_collections WHERE collection_id = '{collection_id}'") raw = rows[0]["settings"] if isinstance(rows[0], dict) else rows[0][0] return json.loads(raw) if raw else {} @@ -71,7 +69,9 @@ def _create_collection(client: Any, label: str, schema: Schema) -> Any: """Create collection.""" name = f"test_ns_opt_idx_{label}_{int(time.time() * 1000)}" return client.create_collection( - name=name, schema=schema, use_namespace=True, + name=name, + schema=schema, + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) diff --git a/tests/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py index 6744eb29..ff4490c2 100644 --- a/tests/integration_tests/test_namespace_prewarm.py +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -16,8 +16,8 @@ class TestNamespacePrewarm: - """TestNamespacePrewarm class.""" + def _create_ns_collection_and_namespace(self, client): """Create ns collection and namespace.""" from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT @@ -30,7 +30,9 @@ def _create_ns_collection_and_namespace(self, client): ), ) collection = client.create_collection( - name=name, schema=schema, use_namespace=True, + name=name, + schema=schema, + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) namespace = collection.create_namespace("pw_ns") diff --git a/tests/integration_tests/test_namespace_prewarm_lob.py b/tests/integration_tests/test_namespace_prewarm_lob.py index a46c03d3..9ab5d1aa 100644 --- a/tests/integration_tests/test_namespace_prewarm_lob.py +++ b/tests/integration_tests/test_namespace_prewarm_lob.py @@ -63,6 +63,7 @@ # ==================== SQL observability helpers ==================== + def _exec(client, sql): """Exec.""" return client._server._execute(sql) @@ -125,16 +126,17 @@ def _flush_tenant_macro_cache(): return False try: conn = pymysql.connect( - host=OB_HOST, port=OB_PORT, user=OB_SYS_USER, - password=OB_SYS_PASSWORD, autocommit=True, + host=OB_HOST, + port=OB_PORT, + user=OB_SYS_USER, + password=OB_SYS_PASSWORD, + autocommit=True, ) except Exception: return False try: with conn.cursor() as cur: - cur.execute( - f"ALTER SYSTEM FLUSH SS_LOCAL_CACHE TENANT = {OB_TENANT} CACHE = macro_cache" - ) + cur.execute(f"ALTER SYSTEM FLUSH SS_LOCAL_CACHE TENANT = {OB_TENANT} CACHE = macro_cache") return True except Exception: return False @@ -189,9 +191,10 @@ def _grep_lob_prewarm_log(lob_meta_tablet_ids): # ==================== Fixtures / collection helpers ==================== -class _BaseLobPrewarm: +class _BaseLobPrewarm: """BaseLobPrewarm class.""" + def _make_collection(self, client, name, partitions): """Make collection.""" schema = Schema( @@ -200,9 +203,7 @@ def _make_collection(self, client, name, partitions): embedding_function=None, ), ) - return client.create_collection( - name=name, schema=schema, use_namespace=True, partition_count=partitions - ) + return client.create_collection(name=name, schema=schema, use_namespace=True, partition_count=partitions) def _force_out_of_row_lob(self, client, collection_id, namespace_id): """Insert several >0.75MB incompressible kv_values so they spill out-of-row @@ -219,8 +220,7 @@ def _force_out_of_row_lob(self, client, collection_id, namespace_id): with raw.cursor() as cur: for i in range(LOB_PROBE_ROWS): cur.execute( - f"INSERT INTO `{kv_table}` (namespace_id, kv_key, kv_value) " - f"VALUES ({namespace_id}, %s, %s)", + f"INSERT INTO `{kv_table}` (namespace_id, kv_key, kv_value) VALUES ({namespace_id}, %s, %s)", (f"__lob_probe_{i}__", os.urandom(LARGE_BLOB_BYTES)), ) raw.commit() @@ -231,9 +231,10 @@ def _force_out_of_row_lob(self, client, collection_id, namespace_id): # ==================== Tier 1 + structural multi-namespace ==================== -class TestLobPrewarmStructural(_BaseLobPrewarm): +class TestLobPrewarmStructural(_BaseLobPrewarm): """TestLobPrewarmStructural class.""" + def test_kv_table_has_lob_meta_tablet(self, oceanbase_client): """Tier 1: the prewarm target (kv_data_table) owns an aux LOB-meta tablet.""" name = f"lob_pw_struct_{int(time.time() * 1000)}" @@ -243,8 +244,7 @@ def test_kv_table_has_lob_meta_tablet(self, oceanbase_client): kv_table_id = _resolve_table_id(oceanbase_client, kv_table) lob_tablets = _resolve_lob_meta_tablets(oceanbase_client, kv_table_id) assert len(lob_tablets) == 1, ( - f"single-partition kv table should have exactly 1 lob meta tablet, " - f"got {lob_tablets}" + f"single-partition kv table should have exactly 1 lob meta tablet, got {lob_tablets}" ) finally: oceanbase_client.delete_collection(name=collection.name) @@ -272,21 +272,19 @@ def test_multi_namespace_prewarm_is_independent(self, oceanbase_client): recorded = {int(r["namespace_id"]) for r in rows} expected = {int(ns._namespace_id) for ns in namespaces} assert expected.issubset(recorded), ( - f"every prewarmed namespace must have a hot_table row: " - f"expected {expected}, got {recorded}" + f"every prewarmed namespace must have a hot_table row: expected {expected}, got {recorded}" ) cache = _ss_cache_blocks_for_tablets(oceanbase_client, own_lob_tablets) if cache: - assert set(cache).issubset(own_lob_tablets), ( - "cached lob tablets must belong to this collection only" - ) + assert set(cache).issubset(own_lob_tablets), "cached lob tablets must belong to this collection only" finally: oceanbase_client.delete_collection(name=collection.name) # ==================== Tier 2 + Tier 3: real LOB caching ==================== + class TestLobPrewarmCaching(_BaseLobPrewarm): """TestLobPrewarmCaching class.""" diff --git a/tests/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py index 00832584..60c6e2e8 100644 --- a/tests/integration_tests/test_namespace_query.py +++ b/tests/integration_tests/test_namespace_query.py @@ -6,16 +6,16 @@ import time import pytest - from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT + from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema class TestNamespaceQuery: - """TestNamespaceQuery class.""" + def _setup(self, client, suffix=""): """Setup.""" name = f"test_ns_q_{int(time.time() * 1000)}{suffix}" @@ -27,7 +27,9 @@ def _setup(self, client, suffix=""): fulltext_index=FulltextIndexConfig(analyzer="ik"), ) collection = client.create_collection( - name=name, schema=schema, use_namespace=True, + name=name, + schema=schema, + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) return collection @@ -86,7 +88,7 @@ def test_query_with_metadata_filter(self, db_client): ) assert result is not None assert len(result["ids"][0]) > 0 - if "metadatas" in result and result["metadatas"]: + if result.get("metadatas"): for meta in result["metadatas"][0]: assert meta["category"] == "AI" finally: @@ -209,7 +211,6 @@ def test_same_id_different_namespaces(self, db_client): finally: db_client.delete_collection(name=collection.name) - # ==================== hybrid_search tests ==================== def test_hybrid_search_fulltext_only(self, db_client): diff --git a/tests/integration_tests/test_namespace_query_where_document.py b/tests/integration_tests/test_namespace_query_where_document.py index 1fbd7cdf..8e35064f 100644 --- a/tests/integration_tests/test_namespace_query_where_document.py +++ b/tests/integration_tests/test_namespace_query_where_document.py @@ -8,65 +8,79 @@ import time -import pytest - from namespace_dml_helpers import cleanup + from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import FulltextIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema def _schema() -> Schema: - """Namespace collection schema for where_document query tests.""" - return Schema( - vector_index=VectorIndexConfig( - ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), - embedding_function=None, - ), - fulltext_index=FulltextIndexConfig(analyzer="ik"), - ) + """Namespace collection schema for where_document query tests.""" + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), + ) def _create_ns(client, suffix: str): - """Create a namespace-enabled test collection.""" - name = f"test_ns_qb_{int(time.time() * 1000)}{suffix}" - return client.create_collection( - name=name, schema=_schema(), use_namespace=True, partition_count=4, - ) + """Create a namespace-enabled test collection.""" + name = f"test_ns_qb_{int(time.time() * 1000)}{suffix}" + return client.create_collection( + name=name, + schema=_schema(), + use_namespace=True, + partition_count=4, + ) def _seed(ns) -> None: - """Load deterministic corpus for where_document operator checks.""" - ns.add( - ids=["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k"], - embeddings=[ - [1.0, 0.0, 0.0], [1.0, 0.2, 0.0], [1.0, 0.4, 0.0], - [1.0, 0.6, 0.0], [1.0, 0.8, 0.0], [0.8, 1.0, 0.0], - [0.6, 1.0, 0.0], [0.4, 1.0, 0.0], [0.2, 1.0, 0.0], - [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], - ], - documents=[ - "machine learning AI research", - "Python programming language", - "OceanBase distributed database", - "deep learning neural networks", - "data science machine learning", - "web development Python Django", - "cloud computing AWS", - "database SQL optimization", - "container Docker Kubernetes", - "microservices architecture", - "blockchain distributed ledger", - ], - metadatas=[ - {"cat": "AI", "score": 95}, {"cat": "Prog", "score": 88}, - {"cat": "DB", "score": 92}, {"cat": "AI", "score": 90}, - {"cat": "DS", "score": 85}, {"cat": "Prog", "score": 75}, - {"cat": "Cloud", "score": 70}, {"cat": "DB", "score": 80}, - {"cat": "Cloud", "score": 65}, {"cat": "Arch", "score": 60}, - {"cat": "DB", "score": 55}, - ], - ) + """Load deterministic corpus for where_document operator checks.""" + ns.add( + ids=["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k"], + embeddings=[ + [1.0, 0.0, 0.0], + [1.0, 0.2, 0.0], + [1.0, 0.4, 0.0], + [1.0, 0.6, 0.0], + [1.0, 0.8, 0.0], + [0.8, 1.0, 0.0], + [0.6, 1.0, 0.0], + [0.4, 1.0, 0.0], + [0.2, 1.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ], + documents=[ + "machine learning AI research", + "Python programming language", + "OceanBase distributed database", + "deep learning neural networks", + "data science machine learning", + "web development Python Django", + "cloud computing AWS", + "database SQL optimization", + "container Docker Kubernetes", + "microservices architecture", + "blockchain distributed ledger", + ], + metadatas=[ + {"cat": "AI", "score": 95}, + {"cat": "Prog", "score": 88}, + {"cat": "DB", "score": 92}, + {"cat": "AI", "score": 90}, + {"cat": "DS", "score": 85}, + {"cat": "Prog", "score": 75}, + {"cat": "Cloud", "score": 70}, + {"cat": "DB", "score": 80}, + {"cat": "Cloud", "score": 65}, + {"cat": "Arch", "score": 60}, + {"cat": "DB", "score": 55}, + ], + ) class TestQueryWhereDocumentOperators: @@ -89,9 +103,7 @@ def test_where_document_contains_baseline(self, db_client): docs = result.get("documents", [[]])[0] or [] assert docs, "expected at least one hit for $contains 'machine'" for doc in docs: - assert "machine" in doc.lower(), ( - f"$contains 'machine' returned doc without 'machine': {doc}" - ) + assert "machine" in doc.lower(), f"$contains 'machine' returned doc without 'machine': {doc}" finally: cleanup(db_client, collection) @@ -113,8 +125,7 @@ def test_where_document_not_contains(self, db_client): assert docs, "expected at least one hit for $not_contains 'machine'" violations = [d for d in docs if "machine" in d.lower()] assert not violations, ( - f"$not_contains 'machine' should exclude docs with 'machine', " - f"got violations: {violations}" + f"$not_contains 'machine' should exclude docs with 'machine', got violations: {violations}" ) finally: cleanup(db_client, collection) @@ -129,22 +140,14 @@ def test_where_document_and(self, db_client): result = ns.query( query_embeddings=[1.0, 0.0, 0.0], - where_document={ - "$and": [{"$contains": "machine"}, {"$contains": "learning"}] - }, + where_document={"$and": [{"$contains": "machine"}, {"$contains": "learning"}]}, n_results=10, include=["documents"], ) docs = result.get("documents", [[]])[0] or [] assert docs, "expected at least one hit for $and machine+learning" - violations = [ - d for d in docs - if not ("machine" in d.lower() and "learning" in d.lower()) - ] - assert not violations, ( - f"$and should require both 'machine' and 'learning', " - f"got violations: {violations}" - ) + violations = [d for d in docs if not ("machine" in d.lower() and "learning" in d.lower())] + assert not violations, f"$and should require both 'machine' and 'learning', got violations: {violations}" finally: cleanup(db_client, collection) @@ -158,22 +161,14 @@ def test_where_document_or(self, db_client): result = ns.query( query_embeddings=[1.0, 0.0, 0.0], - where_document={ - "$or": [{"$contains": "blockchain"}, {"$contains": "Kubernetes"}] - }, + where_document={"$or": [{"$contains": "blockchain"}, {"$contains": "Kubernetes"}]}, n_results=10, include=["documents"], ) docs = result.get("documents", [[]])[0] or [] assert docs, "expected at least one hit for $or blockchain|Kubernetes" - violations = [ - d for d in docs - if not ("blockchain" in d.lower() or "kubernetes" in d.lower()) - ] - assert not violations, ( - f"$or should accept 'blockchain' or 'Kubernetes', " - f"got violations: {violations}" - ) + violations = [d for d in docs if not ("blockchain" in d.lower() or "kubernetes" in d.lower())] + assert not violations, f"$or should accept 'blockchain' or 'Kubernetes', got violations: {violations}" finally: cleanup(db_client, collection) @@ -193,14 +188,8 @@ def test_where_document_regex(self, db_client): ) docs = result.get("documents", [[]])[0] or [] assert docs, "expected at least one hit for $regex machine.*learning" - violations = [ - d for d in docs - if "machine" not in d.lower() or "learning" not in d.lower() - ] - assert not violations, ( - f"$regex should match docs with machine+learning, " - f"got violations: {violations}" - ) + violations = [d for d in docs if "machine" not in d.lower() or "learning" not in d.lower()] + assert not violations, f"$regex should match docs with machine+learning, got violations: {violations}" finally: cleanup(db_client, collection) @@ -231,15 +220,10 @@ def test_where_document_and_or_nested(self, db_client): docs = result.get("documents", [[]])[0] or [] assert docs, "expected at least one hit for nested $and($or)" violations = [ - d for d in docs - if not ( - ("machine" in d.lower() or "database" in d.lower()) - and "learning" in d.lower() - ) + d for d in docs if not (("machine" in d.lower() or "database" in d.lower()) and "learning" in d.lower()) ] assert not violations, ( - f"Nested $and($or) should filter (machine|database) + learning, " - f"got violations: {violations}" + f"Nested $and($or) should filter (machine|database) + learning, got violations: {violations}" ) finally: cleanup(db_client, collection) @@ -276,7 +260,8 @@ def test_where_document_or_and_nested(self, db_client): docs = result.get("documents", [[]])[0] or [] assert docs, "expected at least one hit for nested $or($and)" violations = [ - d for d in docs + d + for d in docs if not ( ("machine" in d.lower() and "learning" in d.lower()) or ("distributed" in d.lower() and "database" in d.lower()) diff --git a/tests/integration_tests/test_namespace_reconnect.py b/tests/integration_tests/test_namespace_reconnect.py index 487dc02d..67db17c0 100644 --- a/tests/integration_tests/test_namespace_reconnect.py +++ b/tests/integration_tests/test_namespace_reconnect.py @@ -14,10 +14,10 @@ import time from typing import Any -import pyseekdb - from namespace_dml_helpers import cleanup, create_ns_collection +import pyseekdb + def _new_oceanbase_client() -> Any: """Second pymysql-backed client (separate TCP connection).""" @@ -35,8 +35,8 @@ def _new_oceanbase_client() -> Any: class TestNamespaceReconnectAndMultiClient: - """TestNamespaceReconnectAndMultiClient class.""" + def test_same_client_reconnect_reads_namespace_data(self, oceanbase_client): """After _cleanup(), the next SDK call should reconnect and read prior rows.""" client = oceanbase_client diff --git a/tests/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py index 9c350047..05b42a58 100644 --- a/tests/integration_tests/test_namespace_session_vars.py +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -8,16 +8,16 @@ import time import pytest - from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT + from pyseekdb import IVFConfiguration from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema class TestNamespaceSessionVars: - """TestNamespaceSessionVars class.""" + def _create_ns_collection(self, client): """Create ns collection.""" name = f"test_ns_sessvar_{int(time.time() * 1000)}" @@ -28,15 +28,15 @@ def _create_ns_collection(self, client): ), ) return client.create_collection( - name=name, schema=schema, use_namespace=True, + name=name, + schema=schema, + use_namespace=True, partition_count=NAMESPACE_TEST_PARTITION_COUNT, ) def _query_session_vars(self, client): """Query session vars.""" - rows = client._server._execute( - "SELECT @collection_id AS cid, @namespace_id AS nsid, @ltable_id AS ltid" - ) + rows = client._server._execute("SELECT @collection_id AS cid, @namespace_id AS nsid, @ltable_id AS ltid") row = rows[0] if isinstance(row, (list, tuple)): return {"cid": row[0], "nsid": row[1], "ltid": row[2]} diff --git a/tests/integration_tests/test_namespace_upsert_concurrency.py b/tests/integration_tests/test_namespace_upsert_concurrency.py index 8e4b6975..f9f897a2 100644 --- a/tests/integration_tests/test_namespace_upsert_concurrency.py +++ b/tests/integration_tests/test_namespace_upsert_concurrency.py @@ -6,12 +6,10 @@ import os import threading import time -import uuid -import pytest +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, cleanup, ns_schema import pyseekdb -from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, cleanup, ns_schema def _unique_name(prefix: str) -> str: @@ -44,10 +42,7 @@ def test_multi_client_concurrent_upsert_same_new_id_should_keep_single_record( clients = [_new_oceanbase_client() for _ in range(4)] try: ns = collection.create_namespace("race_ns") - namespaces = [ - client.get_collection(collection.name).get_namespace("race_ns") - for client in clients - ] + namespaces = [client.get_collection(collection.name).get_namespace("race_ns") for client in clients] barrier = threading.Barrier(len(namespaces)) errors: list[Exception] = [] @@ -61,12 +56,11 @@ def _upsert_once(target_ns, idx: int) -> None: documents=f"written by client {idx}", metadatas={"client": idx}, ) - except Exception as exc: # noqa: BLE001 + except Exception as exc: errors.append(exc) threads = [ - threading.Thread(target=_upsert_once, args=(target_ns, idx)) - for idx, target_ns in enumerate(namespaces) + threading.Thread(target=_upsert_once, args=(target_ns, idx)) for idx, target_ns in enumerate(namespaces) ] for thread in threads: thread.start() diff --git a/tests/unit_tests/test_build_document_query.py b/tests/unit_tests/test_build_document_query.py index 347f04ce..f35d67fa 100644 --- a/tests/unit_tests/test_build_document_query.py +++ b/tests/unit_tests/test_build_document_query.py @@ -1,6 +1,6 @@ """Unit tests for BaseClient._build_document_query DSL generation.""" -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest @@ -16,8 +16,8 @@ def client(): class TestBuildDocumentQuery: - """TestBuildDocumentQuery class.""" + def test_and_contains_has_default_operator_and(self, client): """Test and contains has default operator and.""" where_document = {"$and": [{"$contains": "TOKENZPX"}, {"$contains": "TOKENALP"}]} diff --git a/tests/unit_tests/test_build_knn_filter.py b/tests/unit_tests/test_build_knn_filter.py index 43959129..f81c41f3 100644 --- a/tests/unit_tests/test_build_knn_filter.py +++ b/tests/unit_tests/test_build_knn_filter.py @@ -17,6 +17,7 @@ def client(): class TestBuildKnnFilterNe: """TestBuildKnnFilterNe class.""" + def test_ne_hoists_must_not_in_knn_filter(self, client): """Test ne hoists must not in knn filter.""" with patch.object( @@ -29,9 +30,11 @@ def test_ne_hoists_must_not_in_knn_filter(self, client): dimension=3, ) - assert knn_expr["filter"] == [{ - "bool": { - "filter": [{"range": {"data_content.metadata.zpx_hint": {"gte": -9223372036854775808}}}], - "must_not": [{"term": {"data_content.metadata.zpx_hint": 0}}], + assert knn_expr["filter"] == [ + { + "bool": { + "filter": [{"range": {"data_content.metadata.zpx_hint": {"gte": -9223372036854775808}}}], + "must_not": [{"term": {"data_content.metadata.zpx_hint": 0}}], + } } - }] + ] diff --git a/tests/unit_tests/test_build_query_expression.py b/tests/unit_tests/test_build_query_expression.py index 476e59bc..9dc95202 100644 --- a/tests/unit_tests/test_build_query_expression.py +++ b/tests/unit_tests/test_build_query_expression.py @@ -17,6 +17,7 @@ def client(): class TestBuildQueryExpressionNotContains: """TestBuildQueryExpressionNotContains class.""" + def test_not_contains_with_metadata_filter_hoists_must_not(self, client): """$not_contains + where must not nest a must_not-only bool inside must.""" where = {"seq": {"$gte": 43}} @@ -85,6 +86,7 @@ def test_contains_with_metadata_filter_still_uses_must(self, client): class TestBuildQueryExpressionMetadataNe: """TestBuildQueryExpressionMetadataNe class.""" + def test_ne_with_fts_hoists_must_not_from_filter(self, client): """$ne in where must not appear as a must_not-only bool inside filter.""" with patch.object( diff --git a/tests/unit_tests/test_document_query_builder.py b/tests/unit_tests/test_document_query_builder.py index 72f33ea3..cc7c9df1 100644 --- a/tests/unit_tests/test_document_query_builder.py +++ b/tests/unit_tests/test_document_query_builder.py @@ -3,8 +3,6 @@ import sys from pathlib import Path -import pytest - project_root = Path(__file__).parent.parent.parent sys.path.insert(0, str(project_root / "src")) diff --git a/tests/unit_tests/test_get_or_create_collection_concurrency.py b/tests/unit_tests/test_get_or_create_collection_concurrency.py index 53a863fc..1a05a52b 100644 --- a/tests/unit_tests/test_get_or_create_collection_concurrency.py +++ b/tests/unit_tests/test_get_or_create_collection_concurrency.py @@ -22,15 +22,16 @@ class TestCollectionCatalogConflictDetection: """TestCollectionCatalogConflictDetection class.""" + def test_detects_integrity_error_on_sdk_collections(self): """Test detects integrity error on sdk collections.""" + class IntegrityError(Exception): """IntegrityError class.""" + pass - exc = IntegrityError( - '(1062, "Duplicate entry \'my_coll\' for key \'uk_sdk_coll_name\'")' - ) + exc = IntegrityError("(1062, \"Duplicate entry 'my_coll' for key 'uk_sdk_coll_name'\")") assert _is_sdk_collection_catalog_conflict_error(exc) def test_ignores_unrelated_errors(self): @@ -40,6 +41,7 @@ def test_ignores_unrelated_errors(self): class TestCollectionCatalogInsertRecovery: """TestCollectionCatalogInsertRecovery class.""" + def test_insert_conflict_reuses_existing_collection_id(self): """Test insert conflict reuses existing collection id.""" client = MagicMock(spec=BaseClient) @@ -49,14 +51,13 @@ def test_insert_conflict_reuses_existing_collection_id(self): class IntegrityError(Exception): """IntegrityError class.""" + pass def execute_side_effect(sql): """Execute side effect.""" if "INSERT INTO" in sql: - raise IntegrityError( - '(1062, "Duplicate entry \'items\' for key \'uk_sdk_coll_name\'")' - ) + raise IntegrityError("(1062, \"Duplicate entry 'items' for key 'uk_sdk_coll_name'\")") return [] client._execute.side_effect = execute_side_effect @@ -75,22 +76,23 @@ def test_existing_catalog_row_is_reused_without_insert(self): result = BaseClient._create_collection_meta_v2(client, "items", None) assert result["collection_id"] == "existing_id" - insert_calls = [ - call for call in client._execute.call_args_list if "INSERT INTO" in str(call) - ] + insert_calls = [call for call in client._execute.call_args_list if "INSERT INTO" in str(call)] assert not insert_calls class TestCollectionConflictDetection: """TestCollectionConflictDetection class.""" + def test_detects_value_error_for_existing_collection(self): """Test detects value error for existing collection.""" assert _is_collection_conflict_error(ValueError("Collection 'items' already exists")) def test_detects_seekdb_table_exists_error(self): """Test detects seekdb table exists error.""" + class SeekdbError(Exception): """SeekdbError class.""" + pass exc = SeekdbError("Table 'c$v2$abc' already exists failed: code=1050") @@ -116,11 +118,12 @@ def test_detects_conflict_in_cause_chain(self): class TestGetOrCreateCollectionRecovery: """TestGetOrCreateCollectionRecovery class.""" + @staticmethod def _bind_resume_helper(client): """Bind resume helper.""" - client._get_or_resume_existing_collection = ( - BaseClient._get_or_resume_existing_collection.__get__(client, BaseClient) + client._get_or_resume_existing_collection = BaseClient._get_or_resume_existing_collection.__get__( + client, BaseClient ) def test_returns_existing_collection_after_create_conflict(self): @@ -214,6 +217,7 @@ def test_does_not_mask_unrelated_create_errors(self): class TestListNsNamespacesRecyclebinFilter: """TestListNsNamespacesRecyclebinFilter class.""" + def test_sql_excludes_recyclebin_rows(self): """Test sql excludes recyclebin rows.""" client = MagicMock(spec=BaseClient) diff --git a/tests/unit_tests/test_get_or_create_namespace_concurrency.py b/tests/unit_tests/test_get_or_create_namespace_concurrency.py index 5f2935e7..cb29b191 100644 --- a/tests/unit_tests/test_get_or_create_namespace_concurrency.py +++ b/tests/unit_tests/test_get_or_create_namespace_concurrency.py @@ -4,7 +4,7 @@ import sys from pathlib import Path -from unittest.mock import MagicMock, call +from unittest.mock import MagicMock import pytest @@ -20,26 +20,27 @@ class TestNamespaceCatalogConflictDetection: """TestNamespaceCatalogConflictDetection class.""" + def test_detects_integrity_error_on_sdk_namespaces(self): """Test detects integrity error on sdk namespaces.""" + class IntegrityError(Exception): """IntegrityError class.""" + pass - exc = IntegrityError( - '(1062, "Duplicate entry \'cid-ns_shared\' for key \'uk_sdk_ns_coll_name\'")' - ) + exc = IntegrityError("(1062, \"Duplicate entry 'cid-ns_shared' for key 'uk_sdk_ns_coll_name'\")") assert _is_namespace_catalog_conflict_error(exc) def test_detects_integrity_error_on_sdk_ltables(self): """Test detects integrity error on sdk ltables.""" + class IntegrityError(Exception): """IntegrityError class.""" + pass - exc = IntegrityError( - '(1062, "Duplicate entry \'cid-1-default\' for key \'uk_sdk_lt_coll_ns_name\'")' - ) + exc = IntegrityError("(1062, \"Duplicate entry 'cid-1-default' for key 'uk_sdk_lt_coll_ns_name'\")") assert _is_namespace_catalog_conflict_error(exc) def test_ignores_unrelated_errors(self): @@ -49,6 +50,7 @@ def test_ignores_unrelated_errors(self): class TestGetOrCreateNamespaceMetaRecovery: """TestGetOrCreateNamespaceMetaRecovery class.""" + def test_idempotent_insert_reuses_existing_namespace(self): """Test idempotent insert reuses existing namespace.""" client = MagicMock(spec=BaseClient) @@ -64,12 +66,8 @@ def test_idempotent_insert_reuses_existing_namespace(self): result = BaseClient._get_or_create_ns_namespace_meta(client, "cid", "ns_shared") assert result["namespace_id"] == "42" - client._insert_ns_namespace_catalog_row.assert_called_once_with( - "cid", "ns_shared", idempotent=True - ) - client._insert_ns_ltable_catalog_row.assert_called_once_with( - "cid", 42, idempotent=True - ) + client._insert_ns_namespace_catalog_row.assert_called_once_with("cid", "ns_shared", idempotent=True) + client._insert_ns_ltable_catalog_row.assert_called_once_with("cid", 42, idempotent=True) def test_existing_namespace_without_ltable_creates_default_ltable(self): """Test existing namespace without ltable creates default ltable.""" @@ -89,9 +87,7 @@ def test_existing_namespace_without_ltable_creates_default_ltable(self): assert result["ltable_id"] == "7" client._insert_ns_namespace_catalog_row.assert_not_called() - client._insert_ns_ltable_catalog_row.assert_called_once_with( - "cid", 42, idempotent=True - ) + client._insert_ns_ltable_catalog_row.assert_called_once_with("cid", 42, idempotent=True) if __name__ == "__main__": diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 0268f72e..a3a72aa1 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -23,13 +23,12 @@ from pyseekdb.client.namespace import Namespace # noqa: E402 from pyseekdb.client.validators import _validate_include, _validate_n_results # noqa: E402 - # ==================== IVFConfiguration Tests ==================== class TestIVFConfiguration: - """TestIVFConfiguration class.""" + def test_valid_defaults(self): """Test valid defaults.""" config = IVFConfiguration() @@ -138,6 +137,7 @@ def test_valid_lib_vsag(self): def test_valid_lib_enum(self): """Test valid lib enum.""" from pyseekdb.client.configuration import IVFIndexLib + config = IVFConfiguration(lib=IVFIndexLib.VSAG) assert config.lib == "vsag" @@ -156,8 +156,8 @@ def test_properties_invalid_nested(self): class TestVectorIndexConfig: - """TestVectorIndexConfig class.""" + def test_ivf_only(self): """Test ivf only.""" ivf = IVFConfiguration(dimension=128) @@ -190,8 +190,8 @@ def test_neither_ivf_nor_hnsw(self): class TestCollectionNamespaceGuard: - """TestCollectionNamespaceGuard class.""" + def _make_collection(self, use_namespace: bool) -> Collection: """Make collection.""" mock_client = MagicMock() @@ -256,8 +256,8 @@ def test_guard_blocks_peek_when_namespace_enabled(self): class TestCollectionNamespaceManagementGuard: - """TestCollectionNamespaceManagementGuard class.""" + def _make_collection(self, use_namespace: bool) -> Collection: """Make collection.""" mock_client = MagicMock() @@ -310,8 +310,8 @@ def test_has_namespace_requires_namespace_enabled(self): class TestNamespaceObject: - """TestNamespaceObject class.""" + def _make_namespace(self) -> Namespace: """Make namespace.""" mock_client = MagicMock() @@ -432,6 +432,7 @@ def _ensure_connection(self): class CaptureCursor: """CaptureCursor class.""" + def execute(self, sql, params=None): """Execute.""" resolved = sql @@ -558,12 +559,12 @@ def _client(self): def _common_kwargs(self): """Common kwargs.""" - return dict( - collection_id=self.COLLECTION_ID, - collection_name="test_coll", - namespace_id=self.NAMESPACE_ID, - namespace_name="test_ns", - ) + return { + "collection_id": self.COLLECTION_ID, + "collection_name": "test_coll", + "namespace_id": self.NAMESPACE_ID, + "namespace_name": "test_ns", + } def _ivf_kwargs(self): """Common kwargs for 3-dim IVF namespace SQL generation tests.""" @@ -607,10 +608,7 @@ def test_add_warns_when_embeddings_and_documents_with_embedding_function(self, c documents="hello", embedding_function=MagicMock(), ) - assert any( - "explicit embeddings" in r.message and "embedding_function" in r.message - for r in caplog.records - ) + assert any("explicit embeddings" in r.message and "embedding_function" in r.message for r in caplog.records) def test_add_rejects_explicit_embeddings_wrong_dim_without_vector_index(self): """Without VECTOR INDEX, explicit embeddings must match collection dimension.""" @@ -834,7 +832,6 @@ def test_count_sql(self): assert "ltable_id = 1" in sql assert result == 0 - # ---- IVF TYPE in SQL ---- def test_ivf_type_default_in_sql(self): @@ -892,9 +889,7 @@ def test_create_namespace_physical_tables_sn_inline_vector_index(self): fulltext_config=FulltextIndexConfig(analyzer="ik"), is_shared_storage=False, ) - data_create = next( - s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s - ) + data_create = next(s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s) assert "VECTOR INDEX idx_vec(embedding)" in data_create assert "FULLTEXT INDEX idx_fts(document) WITH PARSER ik" in data_create assert "SEARCH INDEX idx_json(data_content)" in data_create @@ -913,17 +908,13 @@ def test_create_namespace_physical_tables_ss_inline_vector_index(self): ivf_config=ivf_config, is_shared_storage=True, ) - data_create = next( - s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s - ) + data_create = next(s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s) assert "VECTOR INDEX idx_vec(embedding)" in data_create assert "FULLTEXT INDEX" not in data_create assert "SEARCH INDEX idx_json(data_content)" in data_create assert "centroids_fresh_mode=spfresh" in data_create assert not any(s.startswith("CREATE VECTOR INDEX") for s in c.executed_sqls) - hot_create = next( - s for s in c.executed_sqls if "CREATE TABLE" in s and "_hot_table" in s - ) + hot_create = next(s for s in c.executed_sqls if "CREATE TABLE" in s and "_hot_table" in s) assert hot_create def test_create_namespace_physical_tables_search_index_only(self): @@ -934,9 +925,7 @@ def test_create_namespace_physical_tables_search_index_only(self): dimension=3, is_shared_storage=False, ) - data_create = next( - s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s - ) + data_create = next(s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s) assert "SEARCH INDEX idx_json(data_content)" in data_create assert "VECTOR INDEX" not in data_create assert "FULLTEXT INDEX" not in data_create @@ -951,9 +940,7 @@ def test_create_namespace_physical_tables_ivf_without_centroids_fresh_mode(self) ivf_config=ivf_config, is_shared_storage=False, ) - data_create = next( - s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s - ) + data_create = next(s for s in c.executed_sqls if "CREATE TABLE" in s and self.TABLE in s) assert "VECTOR INDEX idx_vec(embedding)" in data_create assert "centroids_fresh_mode" not in data_create @@ -962,67 +949,78 @@ def test_create_namespace_physical_tables_ivf_without_centroids_fresh_mode(self) class TestValidateNamespaceName: - """TestValidateNamespaceName class.""" + def test_valid_simple_name(self): """Test valid simple name.""" from pyseekdb.client.client_base import _validate_namespace_name + _validate_namespace_name("my_namespace") def test_valid_with_digits_and_underscore(self): """Test valid with digits and underscore.""" from pyseekdb.client.client_base import _validate_namespace_name + _validate_namespace_name("tenant_123_abc") def test_valid_single_char(self): """Test valid single char.""" from pyseekdb.client.client_base import _validate_namespace_name + _validate_namespace_name("a") def test_valid_boundary_256_chars(self): """Test valid boundary 256 chars.""" from pyseekdb.client.client_base import _validate_namespace_name + _validate_namespace_name("a" * 256) def test_empty_name_raises(self): """Test empty name raises.""" from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="must not be empty"): _validate_namespace_name("") def test_non_string_raises(self): """Test non string raises.""" from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(TypeError, match="must be a string"): _validate_namespace_name(123) def test_too_long_raises(self): """Test too long raises.""" from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="too long"): _validate_namespace_name("a" * 257) def test_hyphen_raises(self): """Test hyphen raises.""" from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my-namespace") def test_space_raises(self): """Test space raises.""" from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my namespace") def test_dot_raises(self): """Test dot raises.""" from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my.namespace") def test_chinese_raises(self): """Test chinese raises.""" from pyseekdb.client.client_base import _validate_namespace_name + with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("命名空间") @@ -1051,8 +1049,12 @@ def test_partition_count_property(self): """Test partition count property.""" mock_client = MagicMock() ns_coll = Collection( - client=mock_client, name="c", collection_id="1", dimension=3, - use_namespace=True, partition_count=4, + client=mock_client, + name="c", + collection_id="1", + dimension=3, + use_namespace=True, + partition_count=4, ) assert ns_coll.partition_count == 4 # Non-namespace collections are not partitioned. @@ -1064,62 +1066,72 @@ def test_partition_count_property(self): class TestValidateRecordIds: - """TestValidateRecordIds class.""" + def test_valid_single_id(self): """Test valid single id.""" from pyseekdb.client.client_base import _validate_record_ids + _validate_record_ids(["doc_1"]) def test_valid_multiple_ids(self): """Test valid multiple ids.""" from pyseekdb.client.client_base import _validate_record_ids + _validate_record_ids(["id1", "id2", "id_3"]) def test_valid_boundary_512_chars(self): """Test valid boundary 512 chars.""" from pyseekdb.client.client_base import _validate_record_ids + _validate_record_ids(["a" * 512]) def test_empty_id_raises(self): """Test empty id raises.""" from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="must not be empty"): _validate_record_ids([""]) def test_non_string_id_raises(self): """Test non string id raises.""" from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(TypeError, match="must be a string"): _validate_record_ids([123]) def test_too_long_id_raises(self): """Test too long id raises.""" from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="too long"): _validate_record_ids(["a" * 513]) def test_invalid_chars_raises(self): """Test invalid chars raises.""" from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="invalid characters"): _validate_record_ids(["doc-1"]) def test_space_in_id_raises(self): """Test space in id raises.""" from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="invalid characters"): _validate_record_ids(["doc 1"]) def test_mixed_valid_and_invalid_raises_on_first_bad(self): """Test mixed valid and invalid raises on first bad.""" from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(ValueError, match="invalid characters"): _validate_record_ids(["good_id", "bad-id"]) def test_non_list_ids_raises(self): """Test non list ids raises.""" from pyseekdb.client.client_base import _validate_record_ids + with pytest.raises(TypeError, match="expected list\\[str\\]"): _validate_record_ids("doc_1") @@ -1128,22 +1140,22 @@ def test_non_list_ids_raises(self): class TestNamespaceBatchLimit: - """TestNamespaceBatchLimit class.""" + def _client(self): """Client.""" return FakeClient() def _common_kwargs(self): """Common kwargs.""" - return dict( - collection_id="42", - collection_name="test_coll", - namespace_id="7", - namespace_name="test_ns", - has_vector_index=True, - collection_dimension=3, - ) + return { + "collection_id": "42", + "collection_name": "test_coll", + "namespace_id": "7", + "namespace_name": "test_ns", + "has_vector_index": True, + "collection_dimension": 3, + } def test_add_single_record_ok(self): """Test add single record ok.""" @@ -1197,36 +1209,42 @@ def test_delete_invalid_id_raises(self): class TestPhysicalTableNames: - """TestPhysicalTableNames class.""" + def test_data_table_name(self): """Test data table name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.data_table_name("my_coll") == "my_coll_logic_data_table" def test_hot_table_name(self): """Test hot table name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.hot_table_name("my_coll") == "my_coll_hot_table" def test_kv_data_table_name(self): """Test kv data table name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.kv_data_table_name("my_coll") == "my_coll_kv_data_table" def test_logic_schema_table_name(self): """Test logic schema table name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.logic_schema_table_name("my_coll") == "my_coll_logic_schema_table" def test_tablegroup_name(self): """Test tablegroup name.""" from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.tablegroup_name("my_coll") == "my_coll_tg" def test_namespace_catalog_table_names(self): """Test namespace catalog table names.""" from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.sdk_namespaces_table() == "sdk_namespaces" assert NamespaceCollectionNames.sdk_ltables_table() == "sdk_ltables" assert NamespaceCollectionNames.sdk_namespaces_stats_table() == "sdk_namespaces_stats" @@ -1234,11 +1252,13 @@ def test_namespace_catalog_table_names(self): def test_is_ns_data_table_true(self): """Test is ns data table true.""" from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.is_ns_data_table("my_coll_logic_data_table") is True def test_is_ns_data_table_false(self): """Test is ns data table false.""" from pyseekdb.client.meta_info import NamespaceCollectionNames + assert NamespaceCollectionNames.is_ns_data_table("my_coll_hot_table") is False @@ -1246,8 +1266,8 @@ def test_is_ns_data_table_false(self): class TestNamespaceCatalogs: - """TestNamespaceCatalogs class.""" + def test_ensure_namespace_catalogs_creates_all_catalog_tables(self): """Test ensure namespace catalogs creates all catalog tables.""" c = FakeClient() @@ -1298,8 +1318,8 @@ def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): class TestBrokenNsCollectionPurge: - """TestBrokenNsCollectionPurge class.""" + def test_purge_broken_ns_collection_if_incomplete_calls_delete(self): """Test purge broken ns collection if incomplete calls delete.""" c = FakeClient() @@ -1328,8 +1348,10 @@ def test_purge_skips_complete_collection(self): def test_get_collection_purges_incomplete_namespace_collection(self): """Test get collection purges incomplete namespace collection.""" + class GetClient(FakeClient): """GetClient class.""" + get_collection = BaseClient.get_collection c = GetClient() @@ -1345,14 +1367,12 @@ class GetClient(FakeClient): with pytest.raises(ValueError, match="does not exist"): c.get_collection("coll") - c._purge_broken_ns_collection_if_incomplete.assert_called_once_with( - collection_name="coll", meta=meta - ) + c._purge_broken_ns_collection_if_incomplete.assert_called_once_with(collection_name="coll", meta=meta) class TestValidateNResults: - """TestValidateNResults class.""" + def test_rejects_boolean_values(self): """Test rejects boolean values.""" with pytest.raises(ValueError, match="n_results must be an integer"): @@ -1362,7 +1382,6 @@ def test_rejects_boolean_values(self): class TestValidateInclude: - """TestValidateInclude class.""" def test_accepts_none_and_valid_fields(self): @@ -1389,13 +1408,14 @@ def test_rejects_non_list(self): class TestUseNamespaceValidation: - """TestUseNamespaceValidation class.""" + def test_hnsw_raises(self): """Test hnsw raises.""" c = FakeClient() + from pyseekdb.client.configuration import HNSWConfiguration, VectorIndexConfig from pyseekdb.client.schema import Schema - from pyseekdb.client.configuration import VectorIndexConfig, HNSWConfiguration + hnsw = HNSWConfiguration(dimension=128) schema = Schema(vector_index=VectorIndexConfig(hnsw=hnsw, embedding_function=None)) with pytest.raises(ValueError, match="HNSW is not allowed"): @@ -1405,9 +1425,14 @@ def test_ob_type_validation(self): """Test ob type validation.""" c = FakeClient() c.detect_db_type_and_version = MagicMock(return_value=("mysql", "8.0")) - from pyseekdb.client.schema import Schema from pyseekdb.client.configuration import VectorIndexConfig - schema = Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=3, centroids_fresh_mode="spfresh"), embedding_function=None)) + from pyseekdb.client.schema import Schema + + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, centroids_fresh_mode="spfresh"), embedding_function=None + ) + ) with pytest.raises(ValueError, match="only supported on OceanBase"): c._create_namespace_collection("test", schema) @@ -1420,14 +1445,12 @@ def test_create_namespace_collection_without_ivf_skips_vector_index(self): c._is_shared_storage_mode = MagicMock(return_value=False) c._create_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) c._ensure_namespace_catalogs = MagicMock() - from pyseekdb.client.schema import Schema from pyseekdb.client.configuration import VectorIndexConfig + from pyseekdb.client.schema import Schema schema = Schema(vector_index=VectorIndexConfig(embedding_function=None)) c._create_namespace_collection("test", schema) - data_create = next( - s for s in c.executed_sqls if "CREATE TABLE" in s and "logic_data_table" in s - ) + data_create = next(s for s in c.executed_sqls if "CREATE TABLE" in s and "logic_data_table" in s) assert "VECTOR INDEX" not in data_create assert "SEARCH INDEX idx_json(data_content)" in data_create settings = c._create_ns_collection_meta.call_args[0][1] @@ -1443,8 +1466,8 @@ def test_create_namespace_collection_ivf_without_centroids_fresh_mode(self): c._is_shared_storage_mode = MagicMock(return_value=False) c._create_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) c._ensure_namespace_catalogs = MagicMock() - from pyseekdb.client.schema import Schema from pyseekdb.client.configuration import VectorIndexConfig + from pyseekdb.client.schema import Schema schema = Schema( vector_index=VectorIndexConfig( @@ -1453,9 +1476,7 @@ def test_create_namespace_collection_ivf_without_centroids_fresh_mode(self): ) ) c._create_namespace_collection("test", schema) - data_create = next( - s for s in c.executed_sqls if "CREATE TABLE" in s and "logic_data_table" in s - ) + data_create = next(s for s in c.executed_sqls if "CREATE TABLE" in s and "logic_data_table" in s) assert "VECTOR INDEX idx_vec(embedding)" in data_create assert "centroids_fresh_mode" not in data_create settings = c._create_ns_collection_meta.call_args[0][1] @@ -1467,21 +1488,23 @@ def test_create_namespace_collection_ivf_without_centroids_fresh_mode(self): class TestDeleteNamespaceUsesKernel: - """TestDeleteNamespaceUsesKernel class.""" + def test_delete_namespace_calls_dbms_logic_table(self): """Test delete namespace calls dbms logic table.""" c = FakeClient() - c._execute = MagicMock(side_effect=[ - [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], - None, # _get_ns_namespace_meta -> SET @namespace_id - None, # _get_ns_namespace_meta -> SET @ltable_id - None, # _delete_ns_namespace_meta -> USE catalog database - None, # _delete_ns_namespace_meta -> SET @collection_id - None, # _delete_ns_namespace_meta -> SET @namespace_id - None, # _delete_ns_namespace_meta -> SET @ltable_id - None, # CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE - ]) + c._execute = MagicMock( + side_effect=[ + [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], + None, # _get_ns_namespace_meta -> SET @namespace_id + None, # _get_ns_namespace_meta -> SET @ltable_id + None, # _delete_ns_namespace_meta -> USE catalog database + None, # _delete_ns_namespace_meta -> SET @collection_id + None, # _delete_ns_namespace_meta -> SET @namespace_id + None, # _delete_ns_namespace_meta -> SET @ltable_id + None, # CALL DBMS_LOGIC_TABLE.DROP_NAMESPACE + ] + ) c._delete_ns_namespace_meta("abc123", "ns1") calls = [str(call) for call in c._execute.call_args_list] assert not any("GET_LOCK" in s for s in calls) @@ -1497,11 +1520,13 @@ def test_delete_namespace_calls_dbms_logic_table(self): def test_get_ns_namespace_meta_sets_ltable_id(self): """get_namespace must populate @ltable_id along with @namespace_id.""" c = FakeClient() - c._execute = MagicMock(side_effect=[ - [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], - None, # SET @namespace_id - None, # SET @ltable_id - ]) + c._execute = MagicMock( + side_effect=[ + [{"namespace_id": 10, "namespace_name": "ns1", "ltable_id": 7}], + None, # SET @namespace_id + None, # SET @ltable_id + ] + ) meta = c._get_ns_namespace_meta("abc123", "ns1") assert meta == {"namespace_id": "10", "namespace_name": "ns1", "ltable_id": "7"} calls = [str(call) for call in c._execute.call_args_list] diff --git a/tests/unit_tests/test_namespace_kernel_errors.py b/tests/unit_tests/test_namespace_kernel_errors.py index 7d7a7bb4..03534da8 100644 --- a/tests/unit_tests/test_namespace_kernel_errors.py +++ b/tests/unit_tests/test_namespace_kernel_errors.py @@ -45,15 +45,16 @@ def test_detects_missing_logic_data_table(self): def test_maps_duplicate_namespace_to_friendly_message(self): """Duplicate sdk_namespaces insert is handled in create path, not here.""" - exc = IntegrityError( - '(1062, "Duplicate entry \'cid-demo_ns\' for key \'uk_sdk_ns_coll_name\'")' + exc = IntegrityError("(1062, \"Duplicate entry 'cid-demo_ns' for key 'uk_sdk_ns_coll_name'\")") + assert ( + _friendly_kernel_error_message( + exc, + namespace_name="demo_ns", + collection_name="coll", + collection_id="cid", + ) + is None ) - assert _friendly_kernel_error_message( - exc, - namespace_name="demo_ns", - collection_name="coll", - collection_id="cid", - ) is None class TestNamespaceKernelErrorTranslation: @@ -62,35 +63,41 @@ class TestNamespaceKernelErrorTranslation: def test_wraps_namespace_dropping_error(self): """Test wraps namespace dropping error.""" exc = Exception("prewarm failed, errcode=-4109") - with namespace_kernel_error_scope( - namespace_name="demo_ns", - collection_name="coll", - collection_id="cid", + with ( + namespace_kernel_error_scope( + namespace_name="demo_ns", + collection_name="coll", + collection_id="cid", + ), + pytest.raises(ValueError, match="being dropped"), ): - with pytest.raises(ValueError, match="being dropped"): - maybe_reraise_friendly_kernel_error(exc) + maybe_reraise_friendly_kernel_error(exc) def test_wraps_namespace_missing_error(self): """Test wraps namespace missing error.""" exc = Exception("namespace does not exist, code=4018") - with namespace_kernel_error_scope( - namespace_name="demo_ns", - collection_name="coll", - collection_id="cid", + with ( + namespace_kernel_error_scope( + namespace_name="demo_ns", + collection_name="coll", + collection_id="cid", + ), + pytest.raises(ValueError, match="no longer exists"), ): - with pytest.raises(ValueError, match="no longer exists"): - maybe_reraise_friendly_kernel_error(exc) + maybe_reraise_friendly_kernel_error(exc) def test_wraps_missing_collection_physical_error(self): """Test wraps missing collection physical error.""" exc = Exception("Table 'abc123_logic_data_table' doesn't exist") - with namespace_kernel_error_scope( - namespace_name="demo_ns", - collection_name="coll", - collection_id="abc123", + with ( + namespace_kernel_error_scope( + namespace_name="demo_ns", + collection_name="coll", + collection_id="abc123", + ), + pytest.raises(ValueError, match="Collection 'coll' does not exist"), ): - with pytest.raises(ValueError, match="Collection 'coll' does not exist"): - maybe_reraise_friendly_kernel_error(exc) + maybe_reraise_friendly_kernel_error(exc) def test_leaves_unrelated_errors_untouched(self): """Test leaves unrelated errors untouched.""" diff --git a/tests/unit_tests/test_namespace_lifecycle_lock.py b/tests/unit_tests/test_namespace_lifecycle_lock.py index 38ea7dcf..69638055 100644 --- a/tests/unit_tests/test_namespace_lifecycle_lock.py +++ b/tests/unit_tests/test_namespace_lifecycle_lock.py @@ -15,6 +15,7 @@ class TestNamespaceLifecycleWithoutGetLock: """TestNamespaceLifecycleWithoutGetLock class.""" + def test_has_namespace_queries_catalog_directly(self): """Test has namespace queries catalog directly.""" client = MagicMock(spec=BaseClient) diff --git a/tests/unit_tests/test_namespace_upsert_reconcile.py b/tests/unit_tests/test_namespace_upsert_reconcile.py index 72552139..2ba79e38 100644 --- a/tests/unit_tests/test_namespace_upsert_reconcile.py +++ b/tests/unit_tests/test_namespace_upsert_reconcile.py @@ -17,6 +17,7 @@ class TestNamespaceUpsertReconcile: """TestNamespaceUpsertReconcile class.""" + def test_reconcile_skips_when_single_row_exists(self): """Test reconcile skips when single row exists.""" client = MagicMock(spec=BaseClient) @@ -60,9 +61,7 @@ def test_reconcile_collapses_duplicate_rows(self): embedding_function=None, ) - client._delete_namespace_records_by_id.assert_called_once_with( - "logic_data_table", 7, 9, "same_new_id" - ) + client._delete_namespace_records_by_id.assert_called_once_with("logic_data_table", 7, 9, "same_new_id") client._namespace_add.assert_called_once() add_kwargs = client._namespace_add.call_args.kwargs assert add_kwargs["ids"] == ["same_new_id"] @@ -75,22 +74,24 @@ def test_reconcile_raises_when_retries_exhausted(self): client = MagicMock(spec=BaseClient) client._count_namespace_records_by_id.return_value = 4 - with patch("pyseekdb.client.client_base.time.sleep"): - with pytest.raises(ValueError, match="Failed to reconcile duplicate namespace rows"): - BaseClient._reconcile_namespace_duplicate_records( - client, - collection_id="c" * 32, - collection_name="items", - namespace_id="7", - namespace_name="race_ns", - ltable_id=9, - table_name="logic_data_table", - ids=["same_new_id"], - documents=["doc"], - metadatas=None, - embeddings=None, - embedding_function=None, - ) + with ( + patch("pyseekdb.client.client_base.time.sleep"), + pytest.raises(ValueError, match="Failed to reconcile duplicate namespace rows"), + ): + BaseClient._reconcile_namespace_duplicate_records( + client, + collection_id="c" * 32, + collection_name="items", + namespace_id="7", + namespace_name="race_ns", + ltable_id=9, + table_name="logic_data_table", + ids=["same_new_id"], + documents=["doc"], + metadatas=None, + embeddings=None, + embedding_function=None, + ) if __name__ == "__main__": diff --git a/tests/unit_tests/test_sentence_transformer_embedding_function.py b/tests/unit_tests/test_sentence_transformer_embedding_function.py index 71ded070..df6914a2 100644 --- a/tests/unit_tests/test_sentence_transformer_embedding_function.py +++ b/tests/unit_tests/test_sentence_transformer_embedding_function.py @@ -97,7 +97,7 @@ def test_get_config_with_kwargs(self): config = ef.get_config() assert config["kwargs"]["trust_remote_code"] is True - assert config["kwargs"]["use_auth_token"] == "test-token" # noqa: S105 + assert config["kwargs"]["use_auth_token"] == "test-token" def test_build_from_config_with_defaults(self): """Test that build_from_config() restores instance with default values""" From b59159a217c04c456758ecdbe9f7967b237ceed9 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 29 Jun 2026 16:42:44 +0800 Subject: [PATCH 75/84] import:_validate_namespace_name --- tests/unit_tests/test_namespace.py | 28 +++++----------------------- 1 file changed, 5 insertions(+), 23 deletions(-) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index a3a72aa1..ad5d6259 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -21,7 +21,11 @@ from pyseekdb.client.collection import Collection # noqa: E402 from pyseekdb.client.configuration import HNSWConfiguration, IVFIndexType, VectorIndexConfig # noqa: E402 from pyseekdb.client.namespace import Namespace # noqa: E402 -from pyseekdb.client.validators import _validate_include, _validate_n_results # noqa: E402 +from pyseekdb.client.validators import ( # noqa: E402 + _validate_include, + _validate_namespace_name, + _validate_n_results, +) # ==================== IVFConfiguration Tests ==================== @@ -953,74 +957,52 @@ class TestValidateNamespaceName: def test_valid_simple_name(self): """Test valid simple name.""" - from pyseekdb.client.client_base import _validate_namespace_name - _validate_namespace_name("my_namespace") def test_valid_with_digits_and_underscore(self): """Test valid with digits and underscore.""" - from pyseekdb.client.client_base import _validate_namespace_name - _validate_namespace_name("tenant_123_abc") def test_valid_single_char(self): """Test valid single char.""" - from pyseekdb.client.client_base import _validate_namespace_name - _validate_namespace_name("a") def test_valid_boundary_256_chars(self): """Test valid boundary 256 chars.""" - from pyseekdb.client.client_base import _validate_namespace_name - _validate_namespace_name("a" * 256) def test_empty_name_raises(self): """Test empty name raises.""" - from pyseekdb.client.client_base import _validate_namespace_name - with pytest.raises(ValueError, match="must not be empty"): _validate_namespace_name("") def test_non_string_raises(self): """Test non string raises.""" - from pyseekdb.client.client_base import _validate_namespace_name - with pytest.raises(TypeError, match="must be a string"): _validate_namespace_name(123) def test_too_long_raises(self): """Test too long raises.""" - from pyseekdb.client.client_base import _validate_namespace_name - with pytest.raises(ValueError, match="too long"): _validate_namespace_name("a" * 257) def test_hyphen_raises(self): """Test hyphen raises.""" - from pyseekdb.client.client_base import _validate_namespace_name - with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my-namespace") def test_space_raises(self): """Test space raises.""" - from pyseekdb.client.client_base import _validate_namespace_name - with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my namespace") def test_dot_raises(self): """Test dot raises.""" - from pyseekdb.client.client_base import _validate_namespace_name - with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("my.namespace") def test_chinese_raises(self): """Test chinese raises.""" - from pyseekdb.client.client_base import _validate_namespace_name - with pytest.raises(ValueError, match="invalid characters"): _validate_namespace_name("命名空间") From 843e8bd02dbb59462245873a95388f40fcdebe78 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 29 Jun 2026 16:57:30 +0800 Subject: [PATCH 76/84] fix_quality --- tests/unit_tests/test_namespace.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index ad5d6259..efae69c2 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -21,11 +21,7 @@ from pyseekdb.client.collection import Collection # noqa: E402 from pyseekdb.client.configuration import HNSWConfiguration, IVFIndexType, VectorIndexConfig # noqa: E402 from pyseekdb.client.namespace import Namespace # noqa: E402 -from pyseekdb.client.validators import ( # noqa: E402 - _validate_include, - _validate_namespace_name, - _validate_n_results, -) +from pyseekdb.client.validators import _validate_include, _validate_namespace_name, _validate_n_results # noqa: E402 # ==================== IVFConfiguration Tests ==================== From 6d2f3574da576e7f94cedf8a37a5f1a36a510c3e Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 29 Jun 2026 17:12:58 +0800 Subject: [PATCH 77/84] skip_lakebase --- .github/workflows/ci.yml | 4 - tests/integration_tests/conftest.py | 7 + .../namespace_test_support.py | 120 ++++++++++++++++++ tests/unit_tests/test_namespace.py | 2 +- 4 files changed, 128 insertions(+), 5 deletions(-) create mode 100644 tests/integration_tests/namespace_test_support.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 31bc08d5..be2e5e2f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -96,8 +96,6 @@ jobs: fi set -o pipefail uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} \ - --ignore-glob='test_namespace*.py' \ - --ignore-glob='test_logic_table_monitoring_p1.py' \ -k "not server and not embedded and not oceanbase" | tee pytest.log tail -n 1 pytest.log | grep '=======' | grep 'passed' |grep -q 'failed' && exit 1 || exit 0 @@ -162,7 +160,5 @@ jobs: fi set -o pipefail uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} \ - --ignore-glob='test_namespace*.py' \ - --ignore-glob='test_logic_table_monitoring_p1.py' \ -k "${{ matrix.test_mode }}" | tee pytest.log tail -n 1 pytest.log | grep '=======' | grep 'passed' | grep -q 'failed' && exit 1 || exit 0 diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index 8307f3c7..eecd049d 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -17,6 +17,8 @@ sys.path.insert(0, str(integration_tests_root)) sys.path.insert(0, str(src_root)) +from namespace_test_support import maybe_skip_namespace_integration_test # noqa: E402 + import pyseekdb # noqa: E402 # ==================== Environment Variable Configuration ==================== @@ -270,3 +272,8 @@ def oceanbase_admin_client(): with contextlib.suppress(Exception): if hasattr(client, "close"): client.close() + + +def pytest_runtest_setup(item): + """Skip namespace integration tests on unsupported backends or OB versions.""" + maybe_skip_namespace_integration_test(item) diff --git a/tests/integration_tests/namespace_test_support.py b/tests/integration_tests/namespace_test_support.py new file mode 100644 index 00000000..4e562197 --- /dev/null +++ b/tests/integration_tests/namespace_test_support.py @@ -0,0 +1,120 @@ +"""Helpers to skip namespace integration tests on unsupported backends/versions.""" + +from __future__ import annotations + +import contextlib +import os +from pathlib import Path + +import pytest + +from pyseekdb.client.client_base import NAMESPACE_MIN_OB_VERSION + +_OB_NAMESPACE_SUPPORT: tuple[bool, str] | None = None + +# Tests that must not be short-circuited by the OB version gate. +_VERSION_SKIP_EXEMPT_TEST_NAMES = frozenset( + { + "test_old_ob_version_rejected", + "test_min_version_constant", + } +) + + +def is_namespace_integration_test(nodeid: str, fspath: str | Path) -> bool: + """Return whether the collected item belongs to the namespace integration suite.""" + stem = Path(fspath).stem + return stem.startswith("test_namespace") or stem == "test_logic_table_monitoring_p1" + + +def should_exempt_from_namespace_version_skip(item) -> bool: + """Return whether the test validates version handling itself.""" + return item.name in _VERSION_SKIP_EXEMPT_TEST_NAMES + + +def infer_integration_test_mode(item) -> str | None: + """Infer embedded/server/oceanbase mode from parametrized fixtures.""" + callspec = getattr(item, "callspec", None) + if callspec is not None: + for key in ("db_client", "admin_client", "_mode"): + if key in callspec.params: + return str(callspec.params[key]) + + fixturenames = getattr(item, "fixturenames", ()) + if "oceanbase_client" in fixturenames or "oceanbase_admin_client" in fixturenames: + return "oceanbase" + if "embedded_client" in fixturenames or "embedded_admin_client" in fixturenames: + return "embedded" + if "server_client" in fixturenames or "server_admin_client" in fixturenames: + return "server" + + # Files like test_logic_table_monitoring_p1.py build OB clients inline. + if is_namespace_integration_test(item.nodeid, item.fspath): + return "oceanbase" + return None + + +def probe_oceanbase_namespace_support() -> tuple[bool, str]: + """Probe whether the configured OceanBase instance supports namespace collections.""" + global _OB_NAMESPACE_SUPPORT + if _OB_NAMESPACE_SUPPORT is not None: + return _OB_NAMESPACE_SUPPORT + + import pyseekdb + + host = os.environ.get("OB_HOST", "localhost") + port = int(os.environ.get("OB_PORT", "11202")) + tenant = os.environ.get("OB_TENANT", "mysql") + database = os.environ.get("OB_DATABASE", "test") + user = os.environ.get("OB_USER", "root") + password = os.environ.get("OB_PASSWORD", "") + + client = None + try: + client = pyseekdb.Client( + host=host, + port=port, + tenant=tenant, + database=database, + user=user, + password=password, + ) + db_type, version = client._server.detect_db_type_and_version() + except Exception as exc: + _OB_NAMESPACE_SUPPORT = (False, f"OceanBase unavailable for namespace tests ({host}:{port}): {exc}") + return _OB_NAMESPACE_SUPPORT + finally: + if client is not None: + with contextlib.suppress(Exception): + client.close() + + if db_type.lower() != "oceanbase": + _OB_NAMESPACE_SUPPORT = ( + False, + f"namespace tests require OceanBase, got {db_type}", + ) + elif version < NAMESPACE_MIN_OB_VERSION: + _OB_NAMESPACE_SUPPORT = ( + False, + f"namespace tests require OceanBase >= {NAMESPACE_MIN_OB_VERSION}, current {version}", + ) + else: + _OB_NAMESPACE_SUPPORT = (True, "") + return _OB_NAMESPACE_SUPPORT + + +def maybe_skip_namespace_integration_test(item) -> None: + """Skip namespace integration tests when the active mode/kernel cannot run them.""" + if not is_namespace_integration_test(item.nodeid, item.fspath): + return + + mode = infer_integration_test_mode(item) + if mode in ("embedded", "server"): + pytest.skip("namespace collections require OceanBase (skip embedded/server integration modes)") + + if should_exempt_from_namespace_version_skip(item): + return + + supported, reason = probe_oceanbase_namespace_support() + if not supported: + pytest.skip(reason) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index efae69c2..284ac473 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -21,7 +21,7 @@ from pyseekdb.client.collection import Collection # noqa: E402 from pyseekdb.client.configuration import HNSWConfiguration, IVFIndexType, VectorIndexConfig # noqa: E402 from pyseekdb.client.namespace import Namespace # noqa: E402 -from pyseekdb.client.validators import _validate_include, _validate_namespace_name, _validate_n_results # noqa: E402 +from pyseekdb.client.validators import _validate_include, _validate_n_results, _validate_namespace_name # noqa: E402 # ==================== IVFConfiguration Tests ==================== From 15ab7fc9ad63d0f65cdab36ea3cfe491e46547db Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 29 Jun 2026 17:31:07 +0800 Subject: [PATCH 78/84] =?UTF-8?q?namespace=E6=A8=A1=E5=BC=8F=E5=8F=AA?= =?UTF-8?q?=E6=94=AF=E6=8C=81lakebase?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/client_base.py | 48 ++++++++++++++++--- .../namespace_test_support.py | 25 ++++++---- .../test_namespace_constraints.py | 31 ++++++++---- tests/unit_tests/test_namespace.py | 4 +- 4 files changed, 81 insertions(+), 27 deletions(-) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index d7e204c2..43cde519 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -80,14 +80,23 @@ # Maximum allowed length for user-facing collection names. _MAX_COLLECTION_NAME_LENGTH = 512 -# Minimum OceanBase version that supports namespace-enabled collections. -NAMESPACE_MIN_OB_VERSION = Version("4.6.1.0") +# Minimum LakeBase (OceanBase Database AI) version for namespace-enabled collections. +NAMESPACE_MIN_LAKEBASE_VERSION = Version("4.6.1.0") +# Backward-compatible alias used by existing tests and skip helpers. +NAMESPACE_MIN_OB_VERSION = NAMESPACE_MIN_LAKEBASE_VERSION + +_LAKEBASE_VERSION_MARKER = "database ai" logger = logging.getLogger(__name__) from .types import _NOT_PROVIDED, _NotProvided # noqa: E402, F401 +def is_lakebase_version_string(version_str: str) -> bool: + """Return whether a ``SELECT version()`` string identifies a LakeBase cluster.""" + return _LAKEBASE_VERSION_MARKER in version_str.lower() + + def _extract_collection_id_from_sdk_row(row: Any) -> str: """Extract the collection_id from an sdk_collections row (dict/tuple/scalar).""" if isinstance(row, dict): @@ -506,16 +515,43 @@ class BaseClient(BaseConnection, AdminAPI): # ==================== Database Type Detection ==================== def _validate_ob_database_type(self) -> None: - """Validate that the backend is OceanBase and meets the minimum version for namespaces.""" + """Validate that the backend is LakeBase and meets the minimum version for namespaces.""" db_type, version = self.detect_db_type_and_version() if db_type.lower() != "oceanbase": - raise ValueError("use_namespace=True is only supported on OceanBase") - if version < NAMESPACE_MIN_OB_VERSION: + raise ValueError("use_namespace=True is only supported on LakeBase (OceanBase Database AI)") + if not self._is_lakebase_cluster(): + raise ValueError( + "use_namespace=True is only supported on LakeBase (OceanBase Database AI); " + "the connected cluster is standard OceanBase" + ) + if version < NAMESPACE_MIN_LAKEBASE_VERSION: raise ValueError( - f"use_namespace=True requires OceanBase version >= {NAMESPACE_MIN_OB_VERSION}, " + f"use_namespace=True requires LakeBase version >= {NAMESPACE_MIN_LAKEBASE_VERSION}, " f"current version is {version}" ) + def _is_lakebase_cluster(self) -> bool: + """Return whether the connected OceanBase cluster is LakeBase (OceanBase Database AI).""" + cached = getattr(self, "_lakebase_cluster", None) + if cached is not None: + return cached + result = False + try: + rows = self._execute("SELECT version() AS version") + if rows: + row = rows[0] + if isinstance(row, dict): + version_str = row.get("version") or row.get("VERSION") or "" + elif isinstance(row, (tuple, list)) and row: + version_str = row[0] + else: + version_str = str(row) + result = is_lakebase_version_string(str(version_str)) + except Exception: + result = False + self._lakebase_cluster = result + return result + def _is_shared_storage_mode(self) -> bool: """Return whether the OceanBase deployment runs in shared-storage mode.""" cached = getattr(self, "_shared_storage", None) diff --git a/tests/integration_tests/namespace_test_support.py b/tests/integration_tests/namespace_test_support.py index 4e562197..7d8e1654 100644 --- a/tests/integration_tests/namespace_test_support.py +++ b/tests/integration_tests/namespace_test_support.py @@ -8,17 +8,16 @@ import pytest -from pyseekdb.client.client_base import NAMESPACE_MIN_OB_VERSION +from pyseekdb.client.client_base import NAMESPACE_MIN_LAKEBASE_VERSION _OB_NAMESPACE_SUPPORT: tuple[bool, str] | None = None # Tests that must not be short-circuited by the OB version gate. -_VERSION_SKIP_EXEMPT_TEST_NAMES = frozenset( - { - "test_old_ob_version_rejected", - "test_min_version_constant", - } -) +_VERSION_SKIP_EXEMPT_TEST_NAMES = frozenset({ + "test_old_lakebase_version_rejected", + "test_standard_oceanbase_rejected", + "test_min_version_constant", +}) def is_namespace_integration_test(nodeid: str, fspath: str | Path) -> bool: @@ -80,6 +79,7 @@ def probe_oceanbase_namespace_support() -> tuple[bool, str]: password=password, ) db_type, version = client._server.detect_db_type_and_version() + is_lakebase = client._server._is_lakebase_cluster() except Exception as exc: _OB_NAMESPACE_SUPPORT = (False, f"OceanBase unavailable for namespace tests ({host}:{port}): {exc}") return _OB_NAMESPACE_SUPPORT @@ -91,12 +91,17 @@ def probe_oceanbase_namespace_support() -> tuple[bool, str]: if db_type.lower() != "oceanbase": _OB_NAMESPACE_SUPPORT = ( False, - f"namespace tests require OceanBase, got {db_type}", + f"namespace tests require LakeBase (OceanBase Database AI), got {db_type}", ) - elif version < NAMESPACE_MIN_OB_VERSION: + elif not is_lakebase: _OB_NAMESPACE_SUPPORT = ( False, - f"namespace tests require OceanBase >= {NAMESPACE_MIN_OB_VERSION}, current {version}", + "namespace tests require LakeBase (OceanBase Database AI); connected cluster is standard OceanBase", + ) + elif version < NAMESPACE_MIN_LAKEBASE_VERSION: + _OB_NAMESPACE_SUPPORT = ( + False, + f"namespace tests require LakeBase >= {NAMESPACE_MIN_LAKEBASE_VERSION}, current {version}", ) else: _OB_NAMESPACE_SUPPORT = (True, "") diff --git a/tests/integration_tests/test_namespace_constraints.py b/tests/integration_tests/test_namespace_constraints.py index f5f7d9d6..b8dd93a1 100644 --- a/tests/integration_tests/test_namespace_constraints.py +++ b/tests/integration_tests/test_namespace_constraints.py @@ -2,7 +2,7 @@ Namespace constraint integration tests. Covers two SDK-level constraints for use_namespace=True collections: -1. Minimum OceanBase version: namespace-enabled collections require OB >= 4.6.1. +1. Minimum LakeBase version: namespace-enabled collections require LakeBase (OceanBase Database AI) >= 4.6.1. 2. Vector index type: only IVF_FLAT is currently supported; ivf_sq8 / ivf_pq are rejected at the SDK layer before any DDL is issued. """ @@ -13,7 +13,7 @@ from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT from pyseekdb import IVFConfiguration -from pyseekdb.client.client_base import NAMESPACE_MIN_OB_VERSION +from pyseekdb.client.client_base import NAMESPACE_MIN_LAKEBASE_VERSION, NAMESPACE_MIN_OB_VERSION from pyseekdb.client.configuration import VectorIndexConfig from pyseekdb.client.schema import Schema from pyseekdb.client.version import Version @@ -129,12 +129,13 @@ def test_dimension_4097_rejected_at_sdk(self, oceanbase_client): class TestNamespaceMinVersionConstraint: """TestNamespaceMinVersionConstraint class.""" - def test_connected_ob_meets_min_version(self, oceanbase_client): - """The kernel under test must already be >= 4.6.1, and creation succeeds.""" + def test_connected_lakebase_meets_min_version(self, oceanbase_client): + """The kernel under test must be LakeBase >= 4.6.1, and creation succeeds.""" db_type, version = oceanbase_client._server.detect_db_type_and_version() assert db_type.lower() == "oceanbase" - assert version >= NAMESPACE_MIN_OB_VERSION, ( - f"connected OB version {version} < required {NAMESPACE_MIN_OB_VERSION}" + assert oceanbase_client._server._is_lakebase_cluster() is True + assert version >= NAMESPACE_MIN_LAKEBASE_VERSION, ( + f"connected LakeBase version {version} < required {NAMESPACE_MIN_LAKEBASE_VERSION}" ) name = _unique_name("_ver_ok") @@ -149,16 +150,26 @@ def test_connected_ob_meets_min_version(self, oceanbase_client): finally: oceanbase_client.delete_collection(name=name) - def test_old_ob_version_rejected(self, oceanbase_client, monkeypatch): - """Simulate an older OB kernel: namespace creation must fail with a clear error.""" + def test_standard_oceanbase_rejected(self, oceanbase_client, monkeypatch): + """Standard OceanBase (without Database AI) must be rejected even when version is new enough.""" server = oceanbase_client._server + monkeypatch.setattr(type(server), "_is_lakebase_cluster", lambda self: False) + name = _unique_name("_not_lakebase") + with pytest.raises(ValueError, match="only supported on LakeBase"): + oceanbase_client.create_collection(name=name, schema=_make_schema("ivf_flat"), use_namespace=True) + assert not oceanbase_client.has_collection(name) + + def test_old_lakebase_version_rejected(self, oceanbase_client, monkeypatch): + """Simulate an older LakeBase kernel: namespace creation must fail with a clear error.""" + server = oceanbase_client._server + monkeypatch.setattr(type(server), "_is_lakebase_cluster", lambda self: True) monkeypatch.setattr( type(server), "detect_db_type_and_version", lambda self: ("oceanbase", Version("4.6.0.0")), ) name = _unique_name("_ver_old") - with pytest.raises(ValueError, match=r"requires OceanBase version >= 4\.6\.1"): + with pytest.raises(ValueError, match=r"requires LakeBase version >= 4\.6\.1"): oceanbase_client.create_collection(name=name, schema=_make_schema("ivf_flat"), use_namespace=True) # restore before checking leftovers monkeypatch.undo() @@ -166,4 +177,4 @@ def test_old_ob_version_rejected(self, oceanbase_client, monkeypatch): def test_min_version_constant(self): """Test min version constant.""" - assert Version("4.6.1.0") == NAMESPACE_MIN_OB_VERSION + assert Version("4.6.1.0") == NAMESPACE_MIN_LAKEBASE_VERSION == NAMESPACE_MIN_OB_VERSION diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 284ac473..0e158bba 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -1411,7 +1411,7 @@ def test_ob_type_validation(self): ivf=IVFConfiguration(dimension=3, centroids_fresh_mode="spfresh"), embedding_function=None ) ) - with pytest.raises(ValueError, match="only supported on OceanBase"): + with pytest.raises(ValueError, match="only supported on LakeBase"): c._create_namespace_collection("test", schema) def test_create_namespace_collection_without_ivf_skips_vector_index(self): @@ -1420,6 +1420,7 @@ def test_create_namespace_collection_without_ivf_skips_vector_index(self): c = FakeClient() c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", Version("4.6.1.0"))) + c._is_lakebase_cluster = MagicMock(return_value=True) c._is_shared_storage_mode = MagicMock(return_value=False) c._create_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) c._ensure_namespace_catalogs = MagicMock() @@ -1441,6 +1442,7 @@ def test_create_namespace_collection_ivf_without_centroids_fresh_mode(self): c = FakeClient() c.detect_db_type_and_version = MagicMock(return_value=("oceanbase", Version("4.6.1.0"))) + c._is_lakebase_cluster = MagicMock(return_value=True) c._is_shared_storage_mode = MagicMock(return_value=False) c._create_ns_collection_meta = MagicMock(return_value={"collection_id": "abc123"}) c._ensure_namespace_catalogs = MagicMock() From dc95767e1709f3740b8b4ab8fe814042cc977d24 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Mon, 29 Jun 2026 19:08:56 +0800 Subject: [PATCH 79/84] =?UTF-8?q?=E9=80=82=E9=85=8D=E5=86=85=E6=A0=B8?= =?UTF-8?q?=E8=AE=A1=E7=AE=97lob=E5=A4=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pyseekdb/client/configuration.py | 6 ++++-- tests/unit_tests/test_namespace.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/pyseekdb/client/configuration.py b/src/pyseekdb/client/configuration.py index 9b947d13..e62a3b80 100644 --- a/src/pyseekdb/client/configuration.py +++ b/src/pyseekdb/client/configuration.py @@ -16,9 +16,11 @@ DEFAULT_DISTANCE_METRIC = "cosine" MAX_HNSW_VECTOR_DIMENSION = 4096 # Namespace IVF uses logic_data_table with LOB_INROW_THRESHOLD sized for float32 vectors -# (dimension * 4 bytes). Below the default ~8KB threshold, IVF indexing rejects out-row LOB. +# plus ObLobCommon header (see OB ob_vector_index_util.cpp IVF in-row check). MAX_IVF_VECTOR_DIMENSION = MAX_HNSW_VECTOR_DIMENSION -LOGIC_DATA_TABLE_LOB_INROW_THRESHOLD = MAX_IVF_VECTOR_DIMENSION * 4 +# Align with OB IVF in-row validation: dim * sizeof(float) + sizeof(ObLobCommon). +_OB_LOB_COMMON_HEADER_BYTES = 4 +LOGIC_DATA_TABLE_LOB_INROW_THRESHOLD = MAX_IVF_VECTOR_DIMENSION * 4 + _OB_LOB_COMMON_HEADER_BYTES PrimitiveValue = str | int | float | bool diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 0e158bba..5c98fc26 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -894,7 +894,7 @@ def test_create_namespace_physical_tables_sn_inline_vector_index(self): assert "FULLTEXT INDEX idx_fts(document) WITH PARSER ik" in data_create assert "SEARCH INDEX idx_json(data_content)" in data_create assert "centroids_fresh_mode=spfresh" in data_create - assert "LOB_INROW_THRESHOLD=16384" in data_create + assert "LOB_INROW_THRESHOLD=16388" in data_create assert not any(s.startswith("CREATE VECTOR INDEX") for s in c.executed_sqls) assert not any("_hot_table" in s for s in c.executed_sqls if s.startswith("CREATE TABLE")) From 6860dbd0ae5b1b4b3bd3a0ec2f3656ddbcdc0347 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 30 Jun 2026 10:28:45 +0800 Subject: [PATCH 80/84] fix_workflow --- .github/workflows/ci.yml | 2 + .../namespace_test_support.py | 91 +++++++++++++++---- 2 files changed, 74 insertions(+), 19 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index be2e5e2f..a8ca209b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -96,6 +96,8 @@ jobs: fi set -o pipefail uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} \ + --ignore-glob='test_namespace*.py' \ + --ignore-glob='test_logic_table_monitoring_p1.py' \ -k "not server and not embedded and not oceanbase" | tee pytest.log tail -n 1 pytest.log | grep '=======' | grep 'passed' |grep -q 'failed' && exit 1 || exit 0 diff --git a/tests/integration_tests/namespace_test_support.py b/tests/integration_tests/namespace_test_support.py index 7d8e1654..53b6d56f 100644 --- a/tests/integration_tests/namespace_test_support.py +++ b/tests/integration_tests/namespace_test_support.py @@ -11,12 +11,15 @@ from pyseekdb.client.client_base import NAMESPACE_MIN_LAKEBASE_VERSION _OB_NAMESPACE_SUPPORT: tuple[bool, str] | None = None +_OB_CONNECTION_AVAILABLE: tuple[bool, str] | None = None -# Tests that must not be short-circuited by the OB version gate. -_VERSION_SKIP_EXEMPT_TEST_NAMES = frozenset({ +# Pure SDK checks; no database connection required. +_VERSION_CONSTANT_ONLY_TESTS = frozenset({"test_min_version_constant"}) + +# Version/LakeBase validation tests that mock kernel behavior but still need a live OB connection. +_OB_CONNECTION_EXEMPT_TESTS = frozenset({ "test_old_lakebase_version_rejected", "test_standard_oceanbase_rejected", - "test_min_version_constant", }) @@ -28,7 +31,54 @@ def is_namespace_integration_test(nodeid: str, fspath: str | Path) -> bool: def should_exempt_from_namespace_version_skip(item) -> bool: """Return whether the test validates version handling itself.""" - return item.name in _VERSION_SKIP_EXEMPT_TEST_NAMES + return item.name in _VERSION_CONSTANT_ONLY_TESTS or item.name in _OB_CONNECTION_EXEMPT_TESTS + + +def _ob_connection_env() -> dict[str, str | int]: + """Read OceanBase connection settings from the environment.""" + return { + "host": os.environ.get("OB_HOST", "localhost"), + "port": int(os.environ.get("OB_PORT", "11202")), + "tenant": os.environ.get("OB_TENANT", "mysql"), + "database": os.environ.get("OB_DATABASE", "test"), + "user": os.environ.get("OB_USER", "root"), + "password": os.environ.get("OB_PASSWORD", ""), + } + + +def probe_oceanbase_connection() -> tuple[bool, str]: + """Return whether an OceanBase instance is reachable (regardless of LakeBase/version).""" + global _OB_CONNECTION_AVAILABLE + if _OB_CONNECTION_AVAILABLE is not None: + return _OB_CONNECTION_AVAILABLE + + import pyseekdb + + env = _ob_connection_env() + client = None + try: + client = pyseekdb.Client( + host=env["host"], + port=env["port"], + tenant=env["tenant"], + database=env["database"], + user=env["user"], + password=env["password"], + ) + rows = client._server._execute("SELECT 1 AS ok") + if not rows: + raise RuntimeError("empty result from SELECT 1") + _OB_CONNECTION_AVAILABLE = (True, "") + except Exception as exc: + _OB_CONNECTION_AVAILABLE = ( + False, + f"OceanBase unavailable for namespace validation tests ({env['host']}:{env['port']}): {exc}", + ) + finally: + if client is not None: + with contextlib.suppress(Exception): + client.close() + return _OB_CONNECTION_AVAILABLE def infer_integration_test_mode(item) -> str | None: @@ -61,27 +111,24 @@ def probe_oceanbase_namespace_support() -> tuple[bool, str]: import pyseekdb - host = os.environ.get("OB_HOST", "localhost") - port = int(os.environ.get("OB_PORT", "11202")) - tenant = os.environ.get("OB_TENANT", "mysql") - database = os.environ.get("OB_DATABASE", "test") - user = os.environ.get("OB_USER", "root") - password = os.environ.get("OB_PASSWORD", "") - + env = _ob_connection_env() client = None try: client = pyseekdb.Client( - host=host, - port=port, - tenant=tenant, - database=database, - user=user, - password=password, + host=env["host"], + port=env["port"], + tenant=env["tenant"], + database=env["database"], + user=env["user"], + password=env["password"], ) db_type, version = client._server.detect_db_type_and_version() is_lakebase = client._server._is_lakebase_cluster() except Exception as exc: - _OB_NAMESPACE_SUPPORT = (False, f"OceanBase unavailable for namespace tests ({host}:{port}): {exc}") + _OB_NAMESPACE_SUPPORT = ( + False, + f"OceanBase unavailable for namespace tests ({env['host']}:{env['port']}): {exc}", + ) return _OB_NAMESPACE_SUPPORT finally: if client is not None: @@ -117,7 +164,13 @@ def maybe_skip_namespace_integration_test(item) -> None: if mode in ("embedded", "server"): pytest.skip("namespace collections require OceanBase (skip embedded/server integration modes)") - if should_exempt_from_namespace_version_skip(item): + if item.name in _VERSION_CONSTANT_ONLY_TESTS: + return + + if item.name in _OB_CONNECTION_EXEMPT_TESTS: + available, reason = probe_oceanbase_connection() + if not available: + pytest.skip(reason) return supported, reason = probe_oceanbase_namespace_support() From b071f5d3103be769eb3c6cc9098b00c4d0ced064 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 30 Jun 2026 10:53:21 +0800 Subject: [PATCH 81/84] fix_issue_hnwyllmm --- docs/guide/index.md | 1 + docs/guide/namespace.md | 99 ++++++++++++++ examples/namespace_example.py | 52 +++++++ src/pyseekdb/client/client_base.py | 34 ++++- .../test_namespace_lifecycle.py | 2 +- .../test_namespace_public_api.py | 78 +++++++++++ tests/unit_tests/test_namespace.py | 18 ++- tests/unit_tests/test_namespace_public_api.py | 127 ++++++++++++++++++ 8 files changed, 408 insertions(+), 3 deletions(-) create mode 100644 docs/guide/namespace.md create mode 100644 examples/namespace_example.py create mode 100644 tests/integration_tests/test_namespace_public_api.py create mode 100644 tests/unit_tests/test_namespace_public_api.py diff --git a/docs/guide/index.md b/docs/guide/index.md index a1274a40..8a1a2a3b 100644 --- a/docs/guide/index.md +++ b/docs/guide/index.md @@ -9,6 +9,7 @@ This guide covers installation, connections, data modeling, and common operation client-connection admin-client collection-management +namespace dml dql embedding-functions diff --git a/docs/guide/namespace.md b/docs/guide/namespace.md new file mode 100644 index 00000000..873e5797 --- /dev/null +++ b/docs/guide/namespace.md @@ -0,0 +1,99 @@ +# Namespace Collections + +Namespace collections partition data inside a single physical collection. Each **namespace** is a logical slice (similar to a tenant or dataset) that shares catalog tables and IVF-backed vector indexes, while keeping documents isolated by `namespace_id`. + +**Requirements** + +- LakeBase **4.6.1.0** or newer (`OceanBase Database AI` in `SELECT version()`) +- `use_namespace=True` when creating the collection +- An explicit **IVF** `Schema` (the default non-namespace path builds **HNSW**, which namespace collections do not support) + +## 1. Create a namespace-enabled collection + +```python +import pyseekdb +from pyseekdb import IVFConfiguration, FulltextIndexConfig, VectorIndexConfig, Schema + +client = pyseekdb.Client(host="127.0.0.1", port=2881, database="test") + +schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=384, distance="cosine", centroids_fresh_mode="spfresh"), + embedding_function=None, # or pass your embedding function + ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), +) + +collection = client.create_collection( + name="products", + schema=schema, + use_namespace=True, + partition_count=4, # optional; controls table partitioning +) +``` + +If you omit `schema` with `use_namespace=True`, pyseekdb raises a clear error instead of silently building the default HNSW schema. + +**Not supported yet** + +- Default / HNSW schemas (`configuration=HNSWConfiguration(...)` without an IVF `Schema`) +- `SparseVectorIndexConfig` on namespace collections + +## 2. Manage namespaces + +```python +ns = collection.create_namespace("electronics") +# or idempotently: +ns = collection.get_or_create_namespace("electronics") + +collection.list_namespaces() # ["electronics", ...] +collection.delete_namespace("electronics") +``` + +## 3. Add and query data inside a namespace + +Namespace DML uses **record ids** — per-document business keys you pass as `ids` in `namespace.add()`, `get()`, `update()`, and `delete()`. They are stored in `data_content.id` (JSON), not as the table primary key. + +```python +ns.add( + ids=["doc_1", "doc_2"], + embeddings=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], + documents=["phone", "laptop"], + metadatas=[{"brand": "A"}, {"brand": "B"}], +) + +rows = ns.get(ids=["doc_1"], include=["documents", "metadatas"]) +hits = ns.query(query_embeddings=[[0.1, 0.2, 0.3]], n_results=5) +``` + +**Record id rules (namespace path)** + +- Must be non-empty strings +- Currently limited to `[A-Za-z0-9_]` (letters, digits, underscore) +- Maximum length 512 characters + +Standard (non-namespace) collections accept broader `_id` values (for example UUIDs with hyphens). If you need UUID-style ids on the namespace path, normalize them (for example `doc_a1b2c3d4`) or wait for a future relaxation. + +## 4. `get_or_create_collection` compatibility + +`get_or_create_collection(..., use_namespace=...)` must match how the collection was originally created: + +```python +# OK — reuses the existing namespace collection +client.get_or_create_collection("products", schema=schema, use_namespace=True) + +# Raises ValueError — name already used by a standard HNSW collection +client.get_or_create_collection("legacy_coll", use_namespace=True) +``` + +Delete the old collection or pick a new name when switching between standard and namespace modes. + +## 5. Collection-level hybrid search + +Namespace collections also expose `collection.hybrid_search(...)` across namespaces. See integration examples under `tests/integration_tests/test_namespace_*.py` for FTS + vector patterns. + +## See also + +- [Collection management](collection-management.md) — standard HNSW collections +- [DML operations](dml.md) — `add` / `update` / `delete` on standard collections +- `examples/namespace_example.py` — minimal runnable sample diff --git a/examples/namespace_example.py b/examples/namespace_example.py new file mode 100644 index 00000000..e22c1a39 --- /dev/null +++ b/examples/namespace_example.py @@ -0,0 +1,52 @@ +""" +Namespace collection example (LakeBase 4.6.1.0+). + +Demonstrates: +1. Creating a namespace-enabled collection with an explicit IVF schema +2. Creating a namespace and adding documents with record ids +3. Vector query inside the namespace +""" + +import pyseekdb +from pyseekdb import FulltextIndexConfig, IVFConfiguration, Schema, VectorIndexConfig + +# Connect to LakeBase / OceanBase (adjust host/port/credentials) +client = pyseekdb.Client( + host="127.0.0.1", + port=2881, + database="test", + user="root@test", + password="", +) + +schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), +) + +collection_name = "demo_namespace_collection" +collection = client.create_collection( + name=collection_name, + schema=schema, + use_namespace=True, + partition_count=4, +) + +ns = collection.get_or_create_namespace("demo") + +# Record ids are per-document keys inside this namespace (see docs/guide/namespace.md). +ns.add( + ids=["item_alpha", "item_beta"], + embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + documents=["alpha product", "beta product"], + metadatas=[{"tier": "a"}, {"tier": "b"}], +) + +results = ns.query(query_embeddings=[[1.0, 0.0, 0.0]], n_results=2, include=["documents", "metadatas"]) +print("Top match:", results["ids"][0][0], results["documents"][0][0]) + +collection.delete_namespace("demo") +client.delete_collection(collection_name) diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index 43cde519..658d1234 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -1014,6 +1014,15 @@ def create_collection( "schema and configuration/embedding_function are both provided, schema will be used", stacklevel=2, ) + elif use_namespace: + raise ValueError( + "use_namespace=True requires an explicit Schema with an IVF vector index. " + "When schema is omitted, create_collection builds the default non-namespace " + "schema (HNSW), which namespace collections do not support. " + "Pass schema=Schema(vector_index=VectorIndexConfig(" + "ivf=IVFConfiguration(dimension=..., distance=...)), ...). " + "Or set use_namespace=False for a standard HNSW collection." + ) else: # Legacy path: convert configuration + embedding_function into a Schema schema = self._prepare_schema_parameters(configuration, embedding_function) @@ -1104,8 +1113,17 @@ def _create_namespace_collection( ivf_config = schema.vector_index.ivf hnsw_config = schema.vector_index.hnsw + if schema.sparse_vector_index is not None: + raise ValueError( + "use_namespace=True does not support SparseVectorIndexConfig yet. " + "Remove sparse_vector_index from the schema or set use_namespace=False." + ) if hnsw_config is not None: - raise ValueError("use_namespace=True only supports IVF index type, HNSW is not allowed") + raise ValueError( + "use_namespace=True does not support HNSW. Namespace collections require an IVF " + "schema: Schema(vector_index=VectorIndexConfig(ivf=IVFConfiguration(dimension=..., " + "distance=...)), ...)." + ) if ivf_config is not None and ivf_config.type != IVFIndexType.IVF_FLAT.value: raise ValueError( f"use_namespace=True currently only supports IVF index type '{IVFIndexType.IVF_FLAT.value}', " @@ -2625,6 +2643,7 @@ def get_or_create_collection( use_namespace=use_namespace, **kwargs, ) + self._assert_get_or_create_namespace_mode_matches(name, use_namespace) return self.get_collection(name, embedding_function=embedding_function) try: @@ -2671,8 +2690,21 @@ def _get_or_resume_existing_collection( use_namespace=use_namespace, **kwargs, ) + self._assert_get_or_create_namespace_mode_matches(name, use_namespace) return self.get_collection(name, embedding_function=embedding_function) + def _assert_get_or_create_namespace_mode_matches(self, name: str, use_namespace: bool) -> None: + """Reject get_or_create when an existing collection's namespace mode differs.""" + existing_use_namespace = self._get_ns_collection_meta(name) is not None + if existing_use_namespace == use_namespace: + return + kind = "namespace-enabled" if existing_use_namespace else "standard" + raise ValueError( + f"Collection '{name}' already exists as a {kind} collection " + f"(use_namespace={existing_use_namespace}), but get_or_create_collection was called with " + f"use_namespace={use_namespace}. Delete the collection or use a different name." + ) + def _get_collection_table_name(self, collection_id: str | None, collection_name: str) -> str: """ Get collection table name diff --git a/tests/integration_tests/test_namespace_lifecycle.py b/tests/integration_tests/test_namespace_lifecycle.py index a0f7fb96..0ceb3e8e 100644 --- a/tests/integration_tests/test_namespace_lifecycle.py +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -307,7 +307,7 @@ def test_create_namespace_collection_with_hnsw_raises(self, db_client): embedding_function=None, ), ) - with pytest.raises(ValueError, match="HNSW is not allowed"): + with pytest.raises(ValueError, match="does not support HNSW"): db_client.create_collection(name=name, schema=schema, use_namespace=True) def test_ss_mode_creates_hot_table(self, oceanbase_client): diff --git a/tests/integration_tests/test_namespace_public_api.py b/tests/integration_tests/test_namespace_public_api.py new file mode 100644 index 00000000..ae5ac75d --- /dev/null +++ b/tests/integration_tests/test_namespace_public_api.py @@ -0,0 +1,78 @@ +""" +Integration tests for namespace public API compatibility (create / get_or_create). +""" + +from __future__ import annotations + +import time + +import pyseekdb +import pytest +from pyseekdb.client.configuration import HNSWConfiguration, SparseVectorIndexConfig, VectorIndexConfig +from pyseekdb.client.schema import Schema +from unittest.mock import MagicMock + +from tests.integration_tests.namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, ns_schema + + +class TestNamespacePublicAPIIntegration: + """End-to-end checks for namespace create/get_or_create guardrails.""" + + def test_create_without_schema_raises_clear_error(self, oceanbase_client): + """Omitting schema must fail before building default HNSW.""" + name = f"test_ns_no_schema_{int(time.time() * 1000)}" + with pytest.raises(ValueError, match="requires an explicit Schema with an IVF vector index"): + oceanbase_client.create_collection(name=name, use_namespace=True) + + def test_create_with_sparse_vector_raises(self, oceanbase_client): + """Sparse vector schema is rejected for namespace collections.""" + name = f"test_ns_sparse_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=pyseekdb.IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + sparse_vector_index=SparseVectorIndexConfig(embedding_function=MagicMock()), + ) + with pytest.raises(ValueError, match="does not support SparseVectorIndexConfig"): + oceanbase_client.create_collection( + name=name, + schema=schema, + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + + def test_get_or_create_mode_mismatch(self, oceanbase_client): + """Standard and namespace collections with the same name are incompatible.""" + name = f"test_ns_mode_{int(time.time() * 1000)}" + oceanbase_client.create_collection( + name=name, + configuration=HNSWConfiguration(dimension=3), + embedding_function=None, + ) + try: + with pytest.raises(ValueError, match="already exists as a standard collection"): + oceanbase_client.get_or_create_collection(name=name, schema=ns_schema(), use_namespace=True) + finally: + oceanbase_client.delete_collection(name=name) + + def test_get_or_create_reuses_matching_namespace_collection(self, oceanbase_client): + """get_or_create with matching namespace mode returns the same handle.""" + name = f"test_ns_goc_{int(time.time() * 1000)}" + created = oceanbase_client.create_collection( + name=name, + schema=ns_schema(), + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + try: + reused = oceanbase_client.get_or_create_collection( + name=name, + schema=ns_schema(), + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + assert reused.id == created.id + assert reused.use_namespace is True + finally: + oceanbase_client.delete_collection(name=name) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index 5c98fc26..b3d63c8e 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -1396,7 +1396,23 @@ def test_hnsw_raises(self): hnsw = HNSWConfiguration(dimension=128) schema = Schema(vector_index=VectorIndexConfig(hnsw=hnsw, embedding_function=None)) - with pytest.raises(ValueError, match="HNSW is not allowed"): + with pytest.raises(ValueError, match="does not support HNSW"): + c._create_namespace_collection("test", schema) + + def test_sparse_vector_raises(self): + """Test sparse vector index raises.""" + c = FakeClient() + from pyseekdb.client.configuration import SparseVectorIndexConfig, VectorIndexConfig + from pyseekdb.client.schema import Schema + + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + sparse_vector_index=SparseVectorIndexConfig(embedding_function=MagicMock()), + ) + with pytest.raises(ValueError, match="does not support SparseVectorIndexConfig"): c._create_namespace_collection("test", schema) def test_ob_type_validation(self): diff --git a/tests/unit_tests/test_namespace_public_api.py b/tests/unit_tests/test_namespace_public_api.py new file mode 100644 index 00000000..ad208edc --- /dev/null +++ b/tests/unit_tests/test_namespace_public_api.py @@ -0,0 +1,127 @@ +""" +Public API tests for namespace collection creation and get_or_create behavior. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +project_root = Path(__file__).parent.parent.parent +src_root = project_root / "src" +sys.path.insert(0, str(src_root)) + +from pyseekdb.client.client_base import BaseClient # noqa: E402 +from pyseekdb.client.configuration import ( # noqa: E402 + HNSWConfiguration, + SparseVectorIndexConfig, + VectorIndexConfig, +) +from pyseekdb.client.schema import Schema # noqa: E402 +from pyseekdb.client.types import _NOT_PROVIDED # noqa: E402 +from pyseekdb import IVFConfiguration # noqa: E402 +from tests.unit_tests.test_namespace import FakeClient # noqa: E402 + + +class TestNamespaceCreateCollectionPublicAPI: + """create_collection(use_namespace=True) validation via public entry points.""" + + @staticmethod + def _client() -> FakeClient: + client = FakeClient() + client.create_collection = BaseClient.create_collection.__get__(client, BaseClient) + return client + + def test_create_without_schema_raises_clear_error(self): + """Omitting schema must not fall through to default HNSW.""" + client = self._client() + with pytest.raises(ValueError, match="requires an explicit Schema with an IVF vector index"): + client.create_collection(name="ns_coll", use_namespace=True) + + def test_create_with_explicit_hnsw_schema_raises(self): + """Explicit HNSW schema is rejected with an IVF-oriented message.""" + client = self._client() + schema = Schema( + vector_index=VectorIndexConfig( + hnsw=HNSWConfiguration(dimension=3), + embedding_function=None, + ), + ) + with pytest.raises(ValueError, match="does not support HNSW"): + client.create_collection(name="ns_coll", schema=schema, use_namespace=True) + + def test_create_with_sparse_vector_schema_raises(self): + """Sparse vector index in schema is rejected for namespace collections.""" + client = self._client() + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + sparse_vector_index=SparseVectorIndexConfig(embedding_function=MagicMock()), + ) + with pytest.raises(ValueError, match="does not support SparseVectorIndexConfig"): + client.create_collection(name="ns_coll", schema=schema, use_namespace=True) + + +class TestNamespaceGetOrCreatePublicAPI: + """get_or_create_collection namespace mode compatibility.""" + + @staticmethod + def _bind_helpers(client): + client._assert_get_or_create_namespace_mode_matches = ( + BaseClient._assert_get_or_create_namespace_mode_matches.__get__(client, BaseClient) + ) + client.get_or_create_collection = BaseClient.get_or_create_collection.__get__(client, BaseClient) + + def test_mode_mismatch_standard_then_namespace(self): + """Existing standard collection cannot be fetched as namespace.""" + client = MagicMock(spec=BaseClient) + self._bind_helpers(client) + client.has_collection.return_value = True + client._is_incomplete_ns_collection.return_value = False + client._get_ns_collection_meta.return_value = None + + with pytest.raises(ValueError, match="already exists as a standard collection"): + client.get_or_create_collection("items", use_namespace=True) + + client.get_collection.assert_not_called() + + def test_mode_mismatch_namespace_then_standard(self): + """Existing namespace collection cannot be fetched as standard.""" + client = MagicMock(spec=BaseClient) + self._bind_helpers(client) + client.has_collection.return_value = True + client._is_incomplete_ns_collection.return_value = False + client._get_ns_collection_meta.return_value = { + "collection_id": "1", + "collection_name": "items", + "settings": {"use_namespace": True}, + } + + with pytest.raises(ValueError, match="already exists as a namespace-enabled collection"): + client.get_or_create_collection("items", use_namespace=False) + + client.get_collection.assert_not_called() + + def test_matching_namespace_mode_returns_existing(self): + """Same namespace mode reuses the existing collection.""" + client = MagicMock(spec=BaseClient) + self._bind_helpers(client) + client.has_collection.return_value = True + client._is_incomplete_ns_collection.return_value = False + client._get_ns_collection_meta.return_value = { + "collection_id": "1", + "collection_name": "items", + "settings": {"use_namespace": True}, + } + existing = object() + client.get_collection.return_value = existing + + result = client.get_or_create_collection("items", use_namespace=True) + + assert result is existing + client.get_collection.assert_called_once_with("items", embedding_function=_NOT_PROVIDED) From 7811b2b6b7f977959d8b5443452cf791ce3e5caf Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 30 Jun 2026 11:45:09 +0800 Subject: [PATCH 82/84] without ignore --- .github/workflows/ci.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a8ca209b..be2e5e2f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -96,8 +96,6 @@ jobs: fi set -o pipefail uv run pytest tests/integration_tests/ -v --log-cli-level=${log_level} \ - --ignore-glob='test_namespace*.py' \ - --ignore-glob='test_logic_table_monitoring_p1.py' \ -k "not server and not embedded and not oceanbase" | tee pytest.log tail -n 1 pytest.log | grep '=======' | grep 'passed' |grep -q 'failed' && exit 1 || exit 0 From 3546102f708979d2de573d29a21c7c4108e59d2f Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 30 Jun 2026 11:50:36 +0800 Subject: [PATCH 83/84] quality_fix --- tests/integration_tests/namespace_test_support.py | 10 +++++++--- tests/integration_tests/test_namespace_public_api.py | 6 +++--- tests/unit_tests/test_namespace_public_api.py | 2 +- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/integration_tests/namespace_test_support.py b/tests/integration_tests/namespace_test_support.py index 53b6d56f..35274995 100644 --- a/tests/integration_tests/namespace_test_support.py +++ b/tests/integration_tests/namespace_test_support.py @@ -66,9 +66,13 @@ def probe_oceanbase_connection() -> tuple[bool, str]: password=env["password"], ) rows = client._server._execute("SELECT 1 AS ok") - if not rows: - raise RuntimeError("empty result from SELECT 1") - _OB_CONNECTION_AVAILABLE = (True, "") + if rows: + _OB_CONNECTION_AVAILABLE = (True, "") + else: + _OB_CONNECTION_AVAILABLE = ( + False, + f"OceanBase unavailable for namespace validation tests ({env['host']}:{env['port']}): empty result from SELECT 1", + ) except Exception as exc: _OB_CONNECTION_AVAILABLE = ( False, diff --git a/tests/integration_tests/test_namespace_public_api.py b/tests/integration_tests/test_namespace_public_api.py index ae5ac75d..5d0f560b 100644 --- a/tests/integration_tests/test_namespace_public_api.py +++ b/tests/integration_tests/test_namespace_public_api.py @@ -5,13 +5,13 @@ from __future__ import annotations import time +from unittest.mock import MagicMock -import pyseekdb import pytest + +import pyseekdb from pyseekdb.client.configuration import HNSWConfiguration, SparseVectorIndexConfig, VectorIndexConfig from pyseekdb.client.schema import Schema -from unittest.mock import MagicMock - from tests.integration_tests.namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, ns_schema diff --git a/tests/unit_tests/test_namespace_public_api.py b/tests/unit_tests/test_namespace_public_api.py index ad208edc..30404284 100644 --- a/tests/unit_tests/test_namespace_public_api.py +++ b/tests/unit_tests/test_namespace_public_api.py @@ -14,6 +14,7 @@ src_root = project_root / "src" sys.path.insert(0, str(src_root)) +from pyseekdb import IVFConfiguration # noqa: E402 from pyseekdb.client.client_base import BaseClient # noqa: E402 from pyseekdb.client.configuration import ( # noqa: E402 HNSWConfiguration, @@ -22,7 +23,6 @@ ) from pyseekdb.client.schema import Schema # noqa: E402 from pyseekdb.client.types import _NOT_PROVIDED # noqa: E402 -from pyseekdb import IVFConfiguration # noqa: E402 from tests.unit_tests.test_namespace import FakeClient # noqa: E402 From fb2c7cffd51289f96d0bcc8a0188312d0f8951a3 Mon Sep 17 00:00:00 2001 From: lc533134 Date: Tue, 30 Jun 2026 13:00:34 +0800 Subject: [PATCH 84/84] fix_test --- ...est_get_or_create_collection_multiprocess.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tests/integration_tests/test_get_or_create_collection_multiprocess.py b/tests/integration_tests/test_get_or_create_collection_multiprocess.py index 92c2a594..5c511995 100644 --- a/tests/integration_tests/test_get_or_create_collection_multiprocess.py +++ b/tests/integration_tests/test_get_or_create_collection_multiprocess.py @@ -40,6 +40,8 @@ THREADS_PER_PROCESS = 3 ITEMS_PER_THREAD = 5 WORKER_TIMEOUT_SECONDS = 60 +# OB CE mini in CI defaults to 10s; bulk HNSW insert/refresh can exceed that under load. +OB_QUERY_TIMEOUT_MICROSECONDS = 60_000_000 MIN_PYLIBSEEKDB_VERSION = Version("1.3.0.post1") _MP_CONTEXT = mp.get_context("spawn") @@ -82,6 +84,11 @@ def _import_pyseekdb(): return importlib.import_module("pyseekdb") +def _relax_oceanbase_query_timeout(client) -> None: + """Raise OB session query timeout for slow CI mini clusters.""" + client._server._execute(f"SET ob_query_timeout = {OB_QUERY_TIMEOUT_MICROSECONDS}") + + def _make_client(client_config: dict[str, Any]): """Make client.""" pyseekdb = _import_pyseekdb() @@ -89,7 +96,7 @@ def _make_client(client_config: dict[str, Any]): if mode == "embedded": return pyseekdb.Client(path=client_config["path"], database=client_config["database"]) if mode in ("server", "oceanbase"): - return pyseekdb.Client( + client = pyseekdb.Client( host=client_config["host"], port=client_config["port"], tenant=client_config["tenant"], @@ -97,6 +104,9 @@ def _make_client(client_config: dict[str, Any]): user=client_config["user"], password=client_config["password"], ) + if mode == "oceanbase": + _relax_oceanbase_query_timeout(client) + return client raise ValueError(f"unsupported client mode: {mode}") @@ -107,13 +117,16 @@ def _make_admin_client(client_config: dict[str, Any]): if mode == "embedded": return pyseekdb.AdminClient(path=client_config["path"]) if mode in ("server", "oceanbase"): - return pyseekdb.AdminClient( + admin = pyseekdb.AdminClient( host=client_config["host"], port=client_config["port"], tenant=client_config["tenant"], user=client_config["user"], password=client_config["password"], ) + if mode == "oceanbase": + _relax_oceanbase_query_timeout(admin) + return admin raise ValueError(f"unsupported client mode: {mode}")