diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 374d4dcd..be2e5e2f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -95,7 +95,8 @@ 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} \ + -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 +159,6 @@ 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} \ + -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/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/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/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/__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..03827b28 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, @@ -68,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") @@ -167,6 +171,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..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): @@ -194,6 +195,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 +204,7 @@ def create_collection( schema=schema, configuration=configuration, embedding_function=embedding_function, + use_namespace=use_namespace, **kwargs, ) @@ -227,6 +230,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 +239,7 @@ def get_or_create_collection( schema=schema, configuration=configuration, embedding_function=embedding_function, + use_namespace=use_namespace, **kwargs, ) @@ -243,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 5e153d6f..658d1234 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -5,8 +5,10 @@ import contextlib import json import logging +import os import re import struct +import time import warnings from abc import ABC, abstractmethod from collections.abc import Sequence @@ -21,13 +23,24 @@ 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, HNSWConfiguration, + IVFConfiguration, + IVFIndexType, 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, ) @@ -37,7 +50,8 @@ get_default_embedding_function, ) from .filters import FilterBuilder -from .meta_info import CollectionFieldNames, CollectionNames +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 from .sparse_embedding_function import ( @@ -48,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 @@ -58,14 +80,27 @@ # Maximum allowed length for user-facing collection names. _MAX_COLLECTION_NAME_LENGTH = 512 +# 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): - 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: @@ -74,6 +109,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() @@ -85,7 +121,54 @@ def _is_collection_conflict_error(exc: BaseException) -> bool: return False +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() + 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 _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() + 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 _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: return None elif isinstance(config, HNSWConfiguration): @@ -99,6 +182,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): @@ -142,6 +226,38 @@ def _validate_collection_name(name: str) -> None: ) +_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 _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": 3, "col_name": "embedding", "col_type": "VECTOR"}, + ], + "index_info": index_info, + } + + def _get_fulltext_index_sql( fulltext_config: FulltextIndexConfig | None = None, ) -> str: @@ -211,6 +327,22 @@ 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: + """Build the IVF vector index DDL fragment from an IVFConfiguration.""" + 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.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 "" + 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. @@ -274,6 +406,7 @@ def create_collection( schema: Schema | None = None, configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, + use_namespace: bool = False, **kwargs, ) -> "Collection": """ @@ -292,6 +425,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 @@ -345,6 +479,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"] @@ -379,7 +514,71 @@ class BaseClient(BaseConnection, AdminAPI): # ==================== Database Type Detection ==================== - def detect_db_type_and_version(self) -> tuple[str, "Version"]: # noqa: C901 + def _validate_ob_database_type(self) -> None: + """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 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 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) + 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"] + result = str(val).upper() == "SHARED_STORAGE" + except Exception: + 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"]: """ Detect database type and version. @@ -464,6 +663,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) @@ -481,9 +681,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"), @@ -633,13 +835,14 @@ 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, ) -> 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() @@ -730,8 +933,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, @@ -743,6 +945,8 @@ def create_collection( schema: Schema | None = None, configuration: ConfigurationParam = _NOT_PROVIDED, embedding_function: EmbeddingFunctionParam = _NOT_PROVIDED, + use_namespace: bool = False, + partition_count: int | None = None, **kwargs, ) -> "Collection": """Create a new collection. @@ -758,6 +962,10 @@ 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. + 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: @@ -790,9 +998,13 @@ 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)." + ) # 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): + 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 @@ -802,12 +1014,24 @@ 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) logger.debug(f"schema: {schema}") + if use_namespace: + 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 dense_embedding_function = schema.vector_index.embedding_function @@ -826,8 +1050,8 @@ 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) @@ -881,6 +1105,131 @@ def create_collection( **kwargs, ) + 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 + + 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 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}', " + f"got '{ivf_config.type}'" + ) + self._validate_ob_database_type() + + # 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: + 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: + 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 + + 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.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(), + "properties": dense_embedding_function.get_config(), + } + + 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: + 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, + 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: + 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( + client=self, + name=name, + collection_id=collection_id, + dimension=dimension, + embedding_function=dense_embedding_function, + 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: """Get the dimension from an embedding function.""" try: @@ -905,16 +1254,24 @@ 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: - create_table_sql = """CREATE TABLE IF NOT EXISTS sdk_collections ( + 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, 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) - ) COMMENT='Settings of collections created by SDK';""" + UNIQUE KEY uk_sdk_coll_name (collection_name) + ) 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)") + 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 @@ -924,6 +1281,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} @@ -959,7 +1317,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 @@ -969,13 +1335,691 @@ 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 except Exception as e: raise ValueError(f"Failed to create collection metadata: {e}") from e + # ==================== 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") + _validate_database_name(db) + return db + + def _qtable(self, table: str) -> str: + """Fully-qualified catalog table: `{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 {_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()) + 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, + 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' 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, + 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' 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', + 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 ORGANIZATION INDEX + PARTITION BY KEY(namespace_id) PARTITIONS 8;""" + self._execute(ns_namespaces_sql) + self._execute(ns_ltables_sql) + self._execute(namespaces_stats_sql) + try: + self._execute( + 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) + 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.""" + 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) + sdk_coll = self._qtable(CollectionNames.sdk_collections_table_name()) + insert_sql = ( + f"INSERT INTO {sdk_coll} (collection_name, settings) VALUES ('{collection_name_escaped}', '{settings_str}')" + ) + 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} 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: + """Fetch namespace collection metadata from the catalog.""" + collection_name_escaped = escape_string(collection_name) + try: + rows = self._execute( + f"SELECT collection_id, collection_name, settings " + f"FROM {self._qtable(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: + """Return whether a namespace collection with the given name exists.""" + try: + return self._get_ns_collection_meta(collection_name) is not None + 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: + """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") + 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_ltables_table()}` " + f"WHERE collection_id = '{collection_id_escaped}'" + ) + with contextlib.suppress(Exception): + self._execute( + f"DELETE FROM `{NamespaceCollectionNames.sdk_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( + self, + collection_id: str, + dimension: int, + ivf_config=None, + fulltext_config=None, + is_shared_storage: bool = False, + 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) + schema_table = NamespaceCollectionNames.logic_schema_table_name(collection_id) + + 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 {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 IF NOT EXISTS `{tg_name}` SHARDING='ADAPTIVE'") + + data_sql = f"""CREATE TABLE IF NOT EXISTS `{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, + {index_sql} + ) 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) + + if is_shared_storage: + hot_table = NamespaceCollectionNames.hot_table_name(collection_id) + 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, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + PRIMARY KEY(namespace_id) + ) TABLEGROUP=`{tg_name}` COMMENT='热点/TTL附属表' DEFAULT CHARSET=utf8mb4 ORGANIZATION INDEX + {partition_clause}""") + + 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, + PRIMARY KEY(namespace_id, kv_key) + ) 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}` ( + 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 ORGANIZATION INDEX + {partition_clause}""") + + except Exception: + if cleanup_on_error: + self._cleanup_namespace_physical_tables(collection_id) + 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, + NamespaceCollectionNames.kv_data_table_name, + NamespaceCollectionNames.hot_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 _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 _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(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]] = [ + (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: + 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 + 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_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) -> 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. + + 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 {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" + ) + if not rows: + raise ValueError( + 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: + """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: + cache = {} + self._ns_ltable_id_cache = cache + cache[key] = int(ltable_id) + + def _set_session_ns_context( + self, + collection_id: str | None = None, + 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: + 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 _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()) + rows = self._execute( + f"SELECT namespace_id FROM {ns_table} " + f"WHERE collection_id = '{collection_id_escaped}' AND namespace_name = '{namespace_name_escaped}'" + ) + 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: + """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()) + lt_rows = self._execute( + f"SELECT ltable_id FROM {lt_table} " + f"WHERE collection_id = '{collection_id_escaped}' AND namespace_id = {int(namespace_id)} " + f"AND ltable_name = '{ltable_name_escaped}'" + ) + 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: + """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()) + 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): + self._rollback_connection_if_supported() + 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: + """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()) + 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): + self._rollback_connection_if_supported() + else: + 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, + namespace_name: str, + 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)) + 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) " + 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: + """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: + 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 + 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: + 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) + 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()) + 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"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 + row = rows[0] + 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"] + 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: + meta["ltable_id"] = str(lt_id) + 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: + """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: + """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") + 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]: + """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())} " + f"WHERE collection_id = '{collection_id_escaped}' " + f"AND LEFT(namespace_name, 13) <> '__recyclebin_' " + f"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: + """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}") + return meta["namespace_id"] + + # ==================== 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 + with contextlib.suppress(Exception): + ns_meta = self._get_ns_collection_meta(name) + if ns_meta is not None: + 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: @@ -983,6 +2027,33 @@ 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": + """Build a namespace Collection facade from catalog metadata.""" + 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 + elif "embedding_function" in settings: + ef_info = settings["embedding_function"] + 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"], + collection_id=meta["collection_id"], + dimension=dimension, + embedding_function=ef, + 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: """ Resolve collection metadata infromation from sdk_collections table @@ -1002,7 +2073,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) """ @@ -1086,6 +2157,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) @@ -1169,6 +2241,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") @@ -1236,6 +2309,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") @@ -1290,11 +2367,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"]: + """List namespace-enabled collections.""" + 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 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: + logger.debug("Failed to list namespace collections from catalog", exc_info=True) + 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 @@ -1310,19 +2426,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: @@ -1419,9 +2542,10 @@ 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 _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 @@ -1429,8 +2553,14 @@ 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: - 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 {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: return False @@ -1471,6 +2601,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. @@ -1488,6 +2619,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: @@ -1499,10 +2631,19 @@ 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) 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, + ) + self._assert_get_or_create_namespace_mode_matches(name, use_namespace) return self.get_collection(name, embedding_function=embedding_function) try: @@ -1511,34 +2652,83 @@ def get_or_create_collection( schema=schema, configuration=configuration, embedding_function=embedding_function, + use_namespace=use_namespace, **kwargs, ) except Exception as exc: - if _is_collection_conflict_error(exc) or self.has_collection(name): - return self.get_collection(name, embedding_function=embedding_function) + if _is_collection_conflict_error(exc): + return self._get_or_resume_existing_collection( + name, + schema=schema, + configuration=configuration, + embedding_function=embedding_function, + use_namespace=use_namespace, + **kwargs, + ) raise - def _get_collection_table_name(self, collection_id: str | None, collection_name: str) -> str: - """ - Get collection table name - """ - if collection_id: - return CollectionNames.table_name_v2(collection_id) - return CollectionNames.table_name(collection_name) + 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, + ) + 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 + """ + if collection_id: + return CollectionNames.table_name_v2(collection_id) + 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}") @@ -1557,7 +2747,13 @@ 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}'" + """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())} " + 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}'") @@ -1619,7 +2815,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, @@ -1689,7 +2885,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, @@ -1718,6 +2914,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] @@ -1733,6 +2930,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: @@ -1865,7 +3069,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, @@ -1894,6 +3098,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] @@ -1909,6 +3114,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: @@ -2020,7 +3232,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, @@ -2400,18 +3612,20 @@ def _execute_query_with_cursor( Returns: List of normalized row dictionaries """ - 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: + if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): + print(f"[pyseekdb SQL] {sql} -- params={params}", flush=True) + 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) @@ -2425,6 +3639,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: """ @@ -2647,34 +3864,42 @@ 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() 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) -------------------- - def _collection_query( # noqa: C901 + def _collection_query( self, collection_id: str | None, collection_name: str, @@ -2896,7 +4121,7 @@ def _collection_query( # noqa: C901 ) return result - def _collection_query_sparse( # noqa: C901 + def _collection_query_sparse( self, conn, table_name: str, @@ -3055,7 +4280,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, @@ -3282,6 +4507,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"): @@ -3296,7 +4532,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, @@ -3365,6 +4601,75 @@ 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 + + @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 + 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): + 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 @@ -3382,33 +4687,50 @@ 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, negative = self._hoist_must_not_from_filters(filter_conditions) + 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: # 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) + 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 filter_conditions: - # Full-text search with metadata filtering - return {"bool": {"must": [doc_query], "filter": filter_conditions}} - else: - # Full-text search only + 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 + return {"bool": bool_q} + 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: """ @@ -3425,10 +4747,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): @@ -3447,69 +4771,7 @@ 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: - return _with_boost({ - "bool": { - "must_not": [ - { - "query_string": { - "fields": ["document"], - "query": where_document["$not_contains"], - } - } - ] - } - }) - - # 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), - } - }) - - # 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": " OR ".join(escaped_queries), - } - }) - - # 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]]: """ @@ -3531,15 +4793,16 @@ 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 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 @@ -3561,7 +4824,20 @@ 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`. 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: @@ -3570,7 +4846,11 @@ 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: @@ -3624,7 +4904,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: """ @@ -3649,12 +4929,28 @@ 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}") + 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") + 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: return [] if isinstance(raw_embeddings, list) and raw_embeddings and isinstance(raw_embeddings[0], list): @@ -3702,19 +4998,80 @@ 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 + + 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: + 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} 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) 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 + + filtered["ids"].append([ids[i] for i in kept_indices]) + 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: + 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 + def _build_source_fields(self, include: list[str] | None) -> list[str]: """ Infer OceanBase GET_SQL `_source` allowlist from include. @@ -3737,9 +5094,25 @@ def _build_source_fields(self, include: list[str] | None) -> list[str]: return source - def _transform_sql_result( # noqa: C901 - self, result_rows: list[dict[str, Any]], include: list[str] | None - ) -> dict[str, Any]: + 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(self, result_rows: list[dict[str, Any]], include: list[str] | None) -> dict[str, Any]: """ Transform SQL query results to standard format (query-compatible format) @@ -3773,18 +5146,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: @@ -3916,3 +5278,1212 @@ 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] | None) -> dict[str, Any] | None: + """Rewrite a WHERE clause to target namespace-scoped columns.""" + 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, + ) -> 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 + 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 + + @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, + namespace_id: int, + 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}` " + 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: + """Delete namespace records matching the given id.""" + 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 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, + 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, + ) + 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}") + + @namespace_kernel_error_guard + def _namespace_add( + 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: + """Add records to a namespace collection.""" + 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) + 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] + + 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: + 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 add(), or\n" + " 2. Provide embedding_function to auto-generate embeddings from documents." + ) + elif metadatas: + pass + else: + raise ValueError( + "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) + 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})") + + 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) + + 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) + + @namespace_kernel_error_guard + 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: + """Update existing records in a namespace collection.""" + 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) + 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] + + 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 + 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." + ) + + 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) + + id_expr = _NS_DATA_CONTENT_ID_EXPR + active_ids = [] + for i, record_id in enumerate(ids): + 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: + 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 list(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) + meta_case_parts.append( + f"WHEN {id_expr} = %s THEN JSON_SET(data_content, '$.metadata', CAST(%s AS JSON))" + ) + 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() + + @namespace_kernel_error_guard + 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: + """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, + ) + 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) + + existing_ids = set() + id_expr = _NS_DATA_CONTENT_ID_EXPR + id_placeholders = ", ".join(["%s"] * len(ids)) + check_sql = ( + 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})" + ) + 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 = [] + 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, + has_vector_index=has_vector_index, + collection_dimension=collection_dimension, + **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, + has_vector_index=has_vector_index, + collection_dimension=collection_dimension, + **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, + ) + + @namespace_kernel_error_guard + 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: + """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, + ) + 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_placeholders = " OR ".join([f"{_NS_DATA_CONTENT_ID_EXPR} = %s"] * len(ids)) + conditions.append(f"({id_placeholders})") + params.extend(ids) + + 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, ltable_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) + + @namespace_kernel_error_guard + def _namespace_query( + 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]: + """Run a vector query against a namespace collection.""" + 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: + 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) + 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 + # 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")} + + 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 + + post_filter_wd: dict[str, Any] | str | None = None + hybrid_include = include + if where_document is not None: + 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=None, + knn=knn_cfg, + n_results=knn_cfg["n_results"], + include=hybrid_include, + embedding_function=embedding_function, + distance=distance, + dimension=collection_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]) + + if "documents" in include_fields or "document" in include_fields or include is None: + 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: + 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: + batch_emb = batch.get("embeddings") or [[]] + all_embeddings.append(batch_emb[0] if batch_emb else []) + + 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: + result["metadatas"] = all_metadatas + if "embeddings" in include_fields or "embedding" in include_fields: + result["embeddings"] = all_embeddings + return result + + @namespace_kernel_error_guard + def _namespace_get( + 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]: + """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, + ) + 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"{_NS_DATA_CONTENT_ID_EXPR} = '{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, 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: + 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 + + @namespace_kernel_error_guard + def _namespace_count( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + 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, + ) + table_name = NamespaceCollectionNames.data_table_name(collection_id) + ns_id = int(namespace_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() + 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 + + @namespace_kernel_error_guard + 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 a small sample of records from a namespace collection.""" + 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, 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(): + new_key = k + if k == "_id": + 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): + 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): + """Inject the namespace filter into a KNN search expression.""" + 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 + + # 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): + """Inject the namespace filter into a query expression.""" + 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 + 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"] + 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 + + @namespace_kernel_error_guard + 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]: + """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, + ) + 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, ltable_id) + + 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() + + escaped_params = search_parm_json.replace("'", "''") + + hint_sql = _query_hint_to_sql(query_hint, table_name=table_name) or "" + hybrid_sql = ( + 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": [[]], + } + 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]: + """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 = [] + + 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) + + distances.append(self._hybrid_row_score(row)) + + 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: + """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 48532580..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: @@ -270,6 +271,18 @@ 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: + """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 4dcc1a0e..313cb7bd 100644 --- a/src/pyseekdb/client/client_seekdb_server.py +++ b/src/pyseekdb/client/client_seekdb_server.py @@ -12,6 +12,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__) @@ -79,6 +80,10 @@ def _ensure_connection(self) -> pymysql.Connection: **self.kwargs, ) logger.info(f"✅ Connected to remote server: {self.host}:{self.port}/{self.database}") + try: + self._use_catalog_database() + except Exception as exc: + logger.warning("Failed to initialize catalog database on connect: %s", exc) return self._connection @@ -99,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) ==================== @@ -194,12 +200,34 @@ 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" ) return self.tenant + @namespace_kernel_error_guard + def _namespace_prewarm( + self, + collection_id: str | None, + collection_name: str, + namespace_id: str, + 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, + ) + sql = f"CALL DBMS_LOGIC_TABLE.PREWARM('{collection_id}', {namespace_id})" + 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 db5ce46d..16385a98 100644 --- a/src/pyseekdb/client/collection.py +++ b/src/pyseekdb/client/collection.py @@ -10,9 +10,15 @@ from typing import TYPE_CHECKING, Any, Optional +from .validators import ( + _validate_n_results, + _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 +44,22 @@ def __init__( embedding_function: Optional["EmbeddingFunction[EmbeddingDocuments]"] = None, distance: str | None = None, sparse_vector_index_config: Optional["SparseVectorIndexConfig"] = None, + use_namespace: bool = False, + partition_count: int | None = None, + has_vector_index: 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 + """Initialize a lightweight collection handle bound to a client implementation.""" + 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._partition_count = partition_count + self._has_vector_index = has_vector_index self._metadata = metadata # ==================== Properties ==================== @@ -112,8 +111,107 @@ def sparse_embedding_function(self) -> Optional["SparseEmbeddingFunction"]: return self._sparse_vector_index_config.embedding_function return None + @property + 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. + + 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: + """Raise if collection-level DML/DQL is used on a namespace-enabled collection.""" + 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 _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): + 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: - return f"Collection(name='{self._name}', dimension={self._dimension}, client={self._client.mode})" + """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 + + 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": + """Return an existing namespace handle or raise if it does not exist.""" + 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") + 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 + + 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: + """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) + 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: + """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) + + # ==================== End Namespace Management ==================== def fork(self, forked_name: str) -> "Collection": """ @@ -151,6 +249,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 +294,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 +338,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 +382,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 +424,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 +510,8 @@ def query( query_hint=QueryHint(parallel=8, query_timeout=10.0) ) """ + self._guard_collection_data_api() + _validate_n_results(n_results) return self._client._collection_query( collection_id=self._id, collection_name=self._name, @@ -488,6 +593,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 +681,8 @@ def hybrid_search( query_hint=QueryHint(parallel=6, query_timeout=15.0) ) """ + self._guard_collection_data_api() + _validate_n_results(n_results) # When no query/knn provided, return only ids/distances by default if include is None and not query and not knn: include = [] @@ -619,6 +727,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]: @@ -642,6 +751,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..e62a3b80 100644 --- a/src/pyseekdb/client/configuration.py +++ b/src/pyseekdb/client/configuration.py @@ -1,6 +1,8 @@ +"""Index and analyzer configuration types for collection and schema creation.""" + 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 @@ -12,10 +14,18 @@ # 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 +# plus ObLobCommon header (see OB ob_vector_index_util.cpp IVF in-row check). +MAX_IVF_VECTOR_DIMENSION = MAX_HNSW_VECTOR_DIMENSION +# 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 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 +36,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 +45,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 +53,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 +62,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 +79,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 +87,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 +104,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 +121,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,9 +135,10 @@ 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) + _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] @@ -138,6 +157,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 +182,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: @@ -183,7 +204,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. @@ -195,17 +216,38 @@ class DistanceMetric(str, Enum): INNER_PRODUCT = "inner_product" -class HNSWIndexType(str, Enum): +class HNSWIndexType(StrEnum): + """Supported HNSW index subtypes.""" + HNSW = "hnsw" HNSW_SQ = "hnsw_sq" HNSW_BQ = "hnsw_bq" -class HNSWIndexLib(str, Enum): +class HNSWIndexLib(StrEnum): + """Supported HNSW index libraries.""" + + VSAG = "vsag" + + +class IVFIndexType(StrEnum): + """Supported IVF index subtypes for namespace-enabled collections.""" + + IVF_FLAT = "ivf_flat" + IVF_SQ8 = "ivf_sq8" + IVF_PQ = "ivf_pq" + + +class IVFIndexLib(StrEnum): + """Supported IVF index libraries.""" + + OB = "ob" VSAG = "vsag" -class FulltextAnalyzer(str, Enum): +class FulltextAnalyzer(StrEnum): + """Supported fulltext analyzers.""" + SPACE = "space" NGRAM = "ngram" BENG = "beng" @@ -213,12 +255,16 @@ 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" FP32 = "fp32" @@ -237,6 +283,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) @@ -283,6 +330,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) @@ -294,36 +342,102 @@ 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 +@dataclass +class IVFConfiguration: + """ + IVF (Inverted File) index configuration for Agent Database namespace-enabled collections. + + Args: + 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). + 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 + centroids_fresh_mode: str | None = None + 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=MAX_IVF_VECTOR_DIMENSION) + + 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.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) + + @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: self.embedding_function = DefaultEmbeddingFunction() + if self.ivf is not None: + self.ivf.__post_init__() if self.hnsw is not None: self.hnsw.__post_init__() @@ -385,6 +499,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}'" @@ -417,6 +532,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 @@ -464,6 +580,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/document_query_builder.py b/src/pyseekdb/client/document_query_builder.py new file mode 100644 index 00000000..1107d0da --- /dev/null +++ b/src/pyseekdb/client/document_query_builder.py @@ -0,0 +1,265 @@ +"""Build hybrid-search and SQL document filter expressions from ``where_document``.""" + +from __future__ import annotations + +import re +from typing import Any + +_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.""" + 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]: + """Build a ``query_string`` leaf for document full-text match.""" + body: dict[str, Any] = { + "fields": [_DOCUMENT_FIELD], + "query": _escape_query_string_term(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 + return not _contains_operator(where_document, "$regex") + + +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) + 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": "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" 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: + 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 diff --git a/src/pyseekdb/client/filters.py b/src/pyseekdb/client/filters.py index 4e73926b..b9315162 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,10 @@ 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/src/pyseekdb/client/kernel_errors.py b/src/pyseekdb/client/kernel_errors.py new file mode 100644 index 00000000..9242dcb5 --- /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 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 +) + +_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/meta_info.py b/src/pyseekdb/client/meta_info.py index 24d1e0f2..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,4 +58,70 @@ 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" + _LOGIC_SCHEMA_SUFFIX = "_logic_schema_table" + _TG_SUFFIX = "_tg" + + @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" + 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..0446bdd2 --- /dev/null +++ b/src/pyseekdb/client/namespace.py @@ -0,0 +1,295 @@ +""" +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 + +from .validators import _validate_include, _validate_n_results + +if TYPE_CHECKING: + from .collection import Collection + from .embedding_function import EmbeddingFunction + + +class Namespace: + """Scoped view of a single namespace within a namespace-enabled collection.""" + + def __init__( + self, + client: Any, + collection: Collection, + 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 + self._namespace_id = namespace_id + + @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}, 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): + 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). " + "Operations are not allowed on a deleted namespace." + ) + + # ==================== 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: + """Insert records into this namespace.""" + self._guard_exists() + 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, + **self._dml_collection_context(), + **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: + """Update existing records in this namespace.""" + self._guard_exists() + 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, + **self._dml_collection_context(), + **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: + """Insert or update records in this namespace.""" + self._guard_exists() + 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, + **self._dml_collection_context(), + **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: + """Delete records from this namespace by ids or filters.""" + self._guard_exists() + 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]: + """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, + 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, + **self._dml_collection_context(), + **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]: + """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( + 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]: + """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, + 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 the number of records in this namespace.""" + self._guard_exists() + 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 up to ``limit`` records from this namespace.""" + self._guard_exists() + 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: + """Prewarm namespace physical tables to reduce cold-start latency.""" + self._guard_exists() + 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..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, SparseVectorIndexConfig, VectorIndexConfig +from .configuration import ( + FulltextIndexConfig, + HNSWConfiguration, + IVFConfiguration, + SparseVectorIndexConfig, + VectorIndexConfig, +) class Schema: @@ -84,21 +90,24 @@ 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, ): + """Init.""" if isinstance(vector_index, VectorIndexConfig): 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 +140,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,11 +151,12 @@ 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 def __repr__(self) -> str: + """Repr.""" parts = [] if self.vector_index is not None: parts.append(f"vector_index={self.vector_index}") diff --git a/src/pyseekdb/client/validators.py b/src/pyseekdb/client/validators.py new file mode 100644 index 00000000..9899f936 --- /dev/null +++ b/src/pyseekdb/client/validators.py @@ -0,0 +1,136 @@ +"""Input validators for namespace names and record identifiers.""" + +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 +_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: + """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__}") + 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_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 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( + 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__}") + 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_]" + ) + + +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}, 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/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..eecd049d 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -10,11 +10,15 @@ 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)) +from namespace_test_support import maybe_skip_namespace_integration_test # noqa: E402 + import pyseekdb # noqa: E402 # ==================== Environment Variable Configuration ==================== @@ -268,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_dml_helpers.py b/tests/integration_tests/namespace_dml_helpers.py new file mode 100644 index 00000000..42dcf6da --- /dev/null +++ b/tests/integration_tests/namespace_dml_helpers.py @@ -0,0 +1,381 @@ +""" +Shared helpers for namespace DML integration tests. +""" + +from __future__ import annotations + +import contextlib +import time +from dataclasses import dataclass +from typing import Any + +from pyseekdb import IVFConfiguration +from pyseekdb.client.configuration import FulltextIndexConfig, 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). +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: + """DmlRecord class.""" + + doc_id: str + embedding: list[float] + document: str + metadata: dict[str, Any] + + +def ns_schema() -> Schema: + """Ns schema.""" + 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_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, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + + +def cleanup(client: Any, *collections: Any) -> None: + """Cleanup.""" + for collection in collections: + with contextlib.suppress(Exception): + client.delete_collection(name=collection.name) + + +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]: + """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, + ((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]: + """Assert get where document count.""" + 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} 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: + """Assert count after document delete.""" + 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: + """Insert dml corpus in batches.""" + 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: + """Load dml corpus.""" + 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 peek result.""" + 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, + embeddings: list[float] | None, + include: list[str] | None, +) -> list[str] | None: + """Default include.""" + 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]: + """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']}" + 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: + """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 new file mode 100644 index 00000000..b1e8b41a --- /dev/null +++ b/tests/integration_tests/namespace_fts_helpers.py @@ -0,0 +1,772 @@ +""" +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, Literal + +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" +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: + """CorpusRecord class.""" + + 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: + """FtsQueryCase class.""" + + name: str + where_document: dict[str, Any] | str + n_results: int + check_ranking: bool = True + where: dict[str, Any] | None = None + + +def ns_schema(distance: VectorDistanceMetric = "l2") -> Schema: + """Ns schema.""" + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance=distance, centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), + ) + + +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=distance), + 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: + """Add.""" + 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]: + """Embedding for index.""" + 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: + """Doc matches where document.""" + 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: + """Corpus matches fts.""" + if not doc_matches_where_document(record.document, where_document): + return False + 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: + """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: + """Expected best id.""" + 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(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] + target.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 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: + 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 hybrid fulltext result.""" + assert result is not None + 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 [] + 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, 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), ( + 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: + """Get fts case.""" + 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, + *, + 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_{distance}_{int(time.time() * 1000)}" + collection = db_client.create_collection( + name=name, + schema=ns_schema(distance), + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + 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) + namespace.prewarm() + insert_corpus_into_namespace(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: + """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) + + +# 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, + 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", + 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)): + for ns_suffix in ("x", "y"): + ns = collection.create_namespace(f"ns_{ns_suffix}") + ns.prewarm() + ctx[f"{coll_tag}_{ns_suffix}"] = ns + return ctx + + +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. + + 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", distance=distance) + 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, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + 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) + 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", + 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", distance=distance) + 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: + """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: + 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]: + """Run hybrid search fts case.""" + 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: + """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()} + 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..e36a3e08 --- /dev/null +++ b/tests/integration_tests/namespace_hybrid_search_helpers.py @@ -0,0 +1,1641 @@ +""" +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 +import time +from dataclasses import dataclass +from typing import Any, Literal + +from namespace_fts_helpers import ( + 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, + CorpusRecord, + VectorDistanceMetric, + assert_hybrid_fulltext_result, + assert_not_contains_no_token_leak, + build_large_fts_corpus, + doc_matches_where_document, + doc_matches_where_metadata, + flat_hybrid_search_schema, + get_fts_case, + insert_corpus_into_collection, + 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, +) + +# 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}} + +# 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: + """SearchIndexQueryCase class.""" + + name: str + where: dict[str, Any] + n_results: int + min_hits: int = 1 + exact_match_count: int | None = None + + +@dataclass(frozen=True) +class VectorKnnCase: + """VectorKnnCase class.""" + + 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 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: + """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, strict=True)) + + +def _vector_l2_norm(vec: list[float]) -> float: + """Vector l2 norm.""" + 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, strict=False)) / (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], + where: dict[str, Any] | None = None, + *, + 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 + 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]: + """Expected knn ids.""" + 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( + 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 hybrid search index result.""" + assert result is not None + 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" + + 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, + *, + 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 result.get("ids") + ids = result["ids"][0] + 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" + + 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 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, {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): 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: + 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}") + + +def run_hybrid_search_index_case( + namespace: Any, + 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, + 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, + *, + 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, + } + 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, + distance_metric=distance_metric, + check_top1=case.check_top1, + ) + return result + + +def run_hybrid_combined_case( + namespace: Any, + corpus: list[CorpusRecord], + case: HybridCombinedCase, + *, + 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} + 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"), + distance_metric=distance_metric, + 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: + """Get search index case.""" + 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: + """Get vector knn case.""" + 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: + """Get hybrid combined case.""" + for case in HYBRID_COMBINED_CASES: + if case.name == name: + return case + 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 + # 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 + + +def _default_knn( + n_results: int, + *, + 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), + "where": where or WHERE_BROAD_METADATA, + } + + +def corpus_matches_triple_intersection( + record: CorpusRecord, + case: HybridTripleBranchCase, +) -> bool: + """Corpus matches triple intersection.""" + 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): + return False + return True + + +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)) + + +def _slice_hybrid_result( + result: dict[str, Any], + 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 = [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, + *, + distance_metric: VectorDistanceMetric = "l2", +) -> None: + """Assert hybrid triple branch result.""" + assert result is not None + 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)}" + + 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: + """Fts row.""" + if not doc_matches_where_document(rec.document, case.where_document): + return False + 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_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: + """Si row.""" + if not corpus_matches_where(rec, case.where): + return False + 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_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: + """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): + return False + 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_hybrid_knn_result( + corpus, + knn_result, + 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 + + 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 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, 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]: + """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} + 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 (use_namespace=False) or Namespace" + ) + + +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_{distance}_{int(time.time() * 1000)}" + 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}" + time.sleep(INDEX_SETTLE_SECONDS) + return corpus, collection + + +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.""" + try: + assert_hybrid_triple_branch_result(corpus, result, case, distance_metric=distance_metric) + 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], + *, + distance_metric: VectorDistanceMetric = "l2", +) -> 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, 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] + + 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, + 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: + 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, + 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) + 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, + ), + 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: + """Get triple branch case.""" + 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, + 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] + run_hybrid_triple_branch_case( + namespace, + corpus, + tb_case, + collection_baseline=collection_baseline, + distance_metric=distance_metric, + ) + + +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: + """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] + 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, + *, + 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] + run_hybrid_knn_case(namespace, corpus, knn_case, distance_metric=distance_metric) + + +def assert_hybrid_search_index_no_hits( + namespace: Any, + where: dict[str, Any], + *, + 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, + 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}" + + +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", + "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", + "TRIPLE_BRANCH_CASES", + "VECTOR_DISTANCE_METRICS", + "VECTOR_KNN_CASES", + "WHERE_BROAD_METADATA", + "WHERE_DOCUMENT_UNIVERSAL", + "WHERE_FILLER_SEQ", + "HybridCombinedCase", + "HybridTripleBranchCase", + "SearchIndexQueryCase", + "VectorDistanceMetric", + "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", + "ensure_shared_hybrid_search_collection", + "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", + "knn_compare_distance", + "run_hybrid_combined_case", + "run_hybrid_knn_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", + "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", + "teardown_large_fts_collection", + "teardown_multi_coll_multi_ns_fts", +] diff --git a/tests/integration_tests/namespace_test_support.py b/tests/integration_tests/namespace_test_support.py new file mode 100644 index 00000000..35274995 --- /dev/null +++ b/tests/integration_tests/namespace_test_support.py @@ -0,0 +1,182 @@ +"""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_LAKEBASE_VERSION + +_OB_NAMESPACE_SUPPORT: tuple[bool, str] | None = None +_OB_CONNECTION_AVAILABLE: tuple[bool, str] | None = None + +# 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", +}) + + +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_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 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, + 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: + """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 + + 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"], + ) + 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 ({env['host']}:{env['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 LakeBase (OceanBase Database AI), got {db_type}", + ) + elif not is_lakebase: + _OB_NAMESPACE_SUPPORT = ( + False, + "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, "") + 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 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() + if not supported: + pytest.skip(reason) 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 new file mode 100644 index 00000000..b1107eb9 --- /dev/null +++ b/tests/integration_tests/test_get_or_create_collection_concurrency.py @@ -0,0 +1,96 @@ +""" +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 + +from pymysql.converters import escape_string + +import pyseekdb + + +def _make_oceanbase_client(): + """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: + """Count sdk collections.""" + name_escaped = escape_string(collection_name) + rows = client._server._execute( + 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]) + + +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] + if key_name == "uk_sdk_coll_name": + return True + return False + + +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 + 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: + """Worker.""" + 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: + 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..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") @@ -60,31 +62,41 @@ 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 _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() mode = client_config["mode"] 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"], @@ -92,36 +104,46 @@ 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}") def _make_admin_client(client_config: dict[str, Any]): + """Make admin client.""" pyseekdb = _import_pyseekdb() mode = client_config["mode"] 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}") 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: @@ -138,12 +160,13 @@ def _require_embedded_pylibseekdb() -> None: ) -def _run_processes( # noqa: C901 +def _run_processes( target: Callable[..., None], args_list: list[tuple[Any, ...]], 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 +207,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 +231,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 +264,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 +296,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 +335,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 +375,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 +408,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 +479,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 +519,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 +529,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 +554,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 +578,10 @@ 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}" @@ -560,13 +603,23 @@ 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( + 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) finally: with contextlib.suppress(Exception): client.delete_collection(collection_name) class TestMultiprocessMultithreadCrud: + """TestMultiprocessMultithreadCrud class.""" + def test_concurrent_add(self, crud_collection, _mode): + """Test concurrent add.""" client_config, collection_name = crud_collection results = _run_processes( @@ -590,6 +643,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 @@ -609,6 +663,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) @@ -627,6 +682,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 @@ -646,6 +702,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 @@ -669,6 +726,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 new file mode 100644 index 00000000..b668eca2 --- /dev/null +++ b/tests/integration_tests/test_logic_table_monitoring_p1.py @@ -0,0 +1,103 @@ +""" +Logic table monitoring P1 integration tests. + +Includes concurrent get_or_create_namespace idempotency coverage. +""" + +from __future__ import annotations + +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor, as_completed + +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, ns_schema + +import pyseekdb + + +def _make_oceanbase_client(): + """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: + """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}'" + ) + 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: + """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)} " + "AND ltable_name = 'default'" + ) + row = rows[0] + return int(row["cnt"] if isinstance(row, dict) else row[0]) + + +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 + 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: + """Worker.""" + 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: + 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/integration_tests/test_namespace_constraints.py b/tests/integration_tests/test_namespace_constraints.py new file mode 100644 index 00000000..b8dd93a1 --- /dev/null +++ b/tests/integration_tests/test_namespace_constraints.py @@ -0,0 +1,180 @@ +""" +Namespace constraint integration tests. + +Covers two SDK-level constraints for use_namespace=True collections: +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. +""" + +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_LAKEBASE_VERSION, 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: + """Make schema.""" + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine", type=ivf_type, centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + + +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.""" + 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, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + try: + assert collection.use_namespace is True + assert collection.dimension == 3 + finally: + 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_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 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") + collection = oceanbase_client.create_collection( + name=name, + schema=_make_schema("ivf_flat"), + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + try: + assert collection.use_namespace is True + finally: + oceanbase_client.delete_collection(name=name) + + 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 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() + assert not oceanbase_client.has_collection(name) + + def test_min_version_constant(self): + """Test min version constant.""" + assert Version("4.6.1.0") == NAMESPACE_MIN_LAKEBASE_VERSION == NAMESPACE_MIN_OB_VERSION diff --git a/tests/integration_tests/test_namespace_dml.py b/tests/integration_tests/test_namespace_dml.py new file mode 100644 index 00000000..4aca9138 --- /dev/null +++ b/tests/integration_tests/test_namespace_dml.py @@ -0,0 +1,462 @@ +""" +Namespace DML integration tests. + +Small-scale: add / update / upsert / delete / get. +Large-scale (>=1000 rows): count / peek boundaries and multi-namespace orthogonality. +""" + +from __future__ import annotations + +import pytest +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, + load_dml_corpus, +) + + +def _create_namespace(collection, name: str): + """Create namespace.""" + ns = collection.create_namespace(name) + ns.prewarm() + return ns + + +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: + 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: + 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: + 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() + with pytest.raises(ValueError, match="no longer exists"): + ns.prewarm() + finally: + 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]) + db_client.delete_collection(name=collection.name) + 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=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") + 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: + 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: + 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: + 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: + 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: + 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: + 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: + 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: + 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: + 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: + 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: + 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: + 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: + 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: + ns.add( + ids=["del1", "del2"], + embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + ) + ns.delete(ids="del1") + assert len(ns.get(ids="del1")["ids"]) == 0 + assert len(ns.get(ids="del2")["ids"]) == 1 + finally: + cleanup(db_client, collection) + + +class TestNamespaceDMLCountPeekAtScale: + """count / peek with >=1000 rows in a single namespace.""" + + @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") + load_dml_corpus(ns, corpus) + yield collection, ns, corpus + 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: + assert ns.count() == 0 + finally: + 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) + 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): + """Test peek empty namespace.""" + collection = create_ns_collection(db_client, suffix="_peek_empty") + ns = _create_namespace(collection, "peek_empty") + try: + result = ns.peek(limit=10) + assert_peek_result(result, expected_len=0) + finally: + 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): + """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( + 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): + """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") + 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): + """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 + + @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): + """Test peek alpha boundaries.""" + 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): + """Test peek beta boundaries.""" + 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): + """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( + 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): + """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"] + 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): + """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" + 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) + assert ns.count() == 0 + finally: + 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: + 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..a7dd3bb0 --- /dev/null +++ b/tests/integration_tests/test_namespace_dml_multi_coll.py @@ -0,0 +1,69 @@ +""" +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: + """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") + 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="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): + """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") + 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") + + 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") + + 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..2cff2fcb --- /dev/null +++ b/tests/integration_tests/test_namespace_dml_multi_coll_multi_ns.py @@ -0,0 +1,85 @@ +""" +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: + """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 = { + "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"), + } + for key in ("c1_x", "c1_y", "c2_x", "c2_y"): + quadrants[key].prewarm() + 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: + 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): + """Test delete in one quadrant.""" + 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..ea1057c2 --- /dev/null +++ b/tests/integration_tests/test_namespace_dml_multi_ns.py @@ -0,0 +1,114 @@ +""" +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: + """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") + 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"}) + + 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): + """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") + 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"}) + + 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): + """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") + 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") + 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, 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") + finally: + 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") + ns_a.prewarm() + ns_b.prewarm() + 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): + """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") + 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"}) + + 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_drop_validation.py b/tests/integration_tests/test_namespace_drop_validation.py new file mode 100644 index 00000000..c4f679f9 --- /dev/null +++ b/tests/integration_tests/test_namespace_drop_validation.py @@ -0,0 +1,768 @@ +"""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 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. +BG_POLL_TIMEOUT_SEC = 65.0 +BG_POLL_INTERVAL_SEC = 1.0 + + +# --------------------------------------------------------------------------- # +# helpers # +# --------------------------------------------------------------------------- # + + +def _create_namespace(collection, name: str): + """Create namespace.""" + 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}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + return client.create_collection( + name=name, + schema=schema, + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + + +def _execute(client, sql: str): + """Execute.""" + 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: + 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): + """Fetch namespace name.""" + ns_table = _catalog_table(client, "sdk_namespaces") + rows = _execute( + client, + f"SELECT namespace_name FROM {ns_table} " + 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: + """Count ltables.""" + 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}", + ) + 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: + """Count logic schema rows.""" + tbl = NamespaceCollectionNames.logic_schema_table_name(collection_id) + db = client._server.database + rows = _execute( + client, + 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: + """Count hot table rows.""" + tbl = NamespaceCollectionNames.hot_table_name(collection_id) + db = client._server.database + try: + rows = _execute( + client, + f"SELECT COUNT(*) AS c FROM `{db}`.`{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 _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: + 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 `{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: + """Count kv data rows.""" + tbl = NamespaceCollectionNames.kv_data_table_name(collection_id) + db = client._server.database + try: + rows = _execute( + client, + 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 = "" +): + """Wait until.""" + deadline = time.time() + timeout_sec + last_exc = None + while time.time() < deadline: + try: + if predicate(): + return + except Exception as exc: + 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) VALUES {', '.join(values)}", + ) + + +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 = [] + 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 `{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: + """Fetch ltable id.""" + lt_table = _catalog_table(client, "sdk_ltables") + rows = _execute( + client, + 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", + ) + if not rows: + return None + 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: + """Hot table exists.""" + tbl = NamespaceCollectionNames.hot_table_name(collection_id) + db = client._server.database + try: + rows = _execute( + client, + f"SELECT 1 AS ok FROM information_schema.tables 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: + 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): + """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 `{db}`.`{tbl}` (namespace_id, last_access_time) VALUES ({namespace_id}, NOW(6))", + ) + return True + except Exception: + 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, +): + """Insert active ltable row.""" + 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): + """Fetch namespace drop history.""" + 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).""" + _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: + """DROP_NAMESPACE sync transaction + LTABLE_BG async namespace deletion.""" + + # ------------------------------------------------------------- # + # 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) + coll_id = collection.id + try: + 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_logic_schema_rows(client, coll_id, ns_id) >= 1, ( + "create_namespace should have inserted a logic_schema row" + ) + if is_ss: + 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 + # 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) + 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: 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), ( + 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" + if is_ss: + 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): + """LTABLE_BG NAMESPACE_DELETE: batch-delete kv_data, remove sdk_namespaces row.""" + client = oceanbase_client + collection = _make_collection(client) + coll_id = collection.id + try: + 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 + + 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_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 _count_kv_data_rows(client, coll_id, ns_id) == kv_after_seed + 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 + ) + + _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 = _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) + 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')} 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 = _create_namespace(collection, "ns_empty_kv") + ns_id = int(ns.namespace_id) + assert _count_kv_data_rows(client, coll_id, ns_id) == 0 + + _drop_namespace_via_pl(client, coll_id, ns_id) + assert _fetch_namespace_name(client, coll_id, ns_id).startswith(RECYCLEBIN_PREFIX) + + _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 = _create_namespace(collection, "ns_hot_async") + ns_id = int(ns.namespace_id) + assert _count_hot_table_rows(client, coll_id, ns_id) == 1 + + _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(): + """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 = _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) + + 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") + + 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) + + # ------------------------------------------------------------- # + # 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: + 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}" + 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): + """Test multiple ltables all deleted.""" + client = oceanbase_client + collection = _make_collection(client) + try: + 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") + _execute( + client, + 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')", + ) + 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_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) + schema_tbl_q = f"`{client._server.database}`.`{schema_tbl}`" + try: + 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" + 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_q}") + + 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_ltables delete must be rolled back" + finally: + # Recreate the schema table (empty) so cleanup_namespace_physical_tables + # / delete_collection runs cleanly. + with contextlib.suppress(Exception): + _execute( + client, + 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," + f" PRIMARY KEY (namespace_id, ltable_id)) " + f"PARTITION BY KEY(namespace_id) PARTITIONS 8", + ) + 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): + """Test drop namespace is idempotent.""" + client = oceanbase_client + collection = _make_collection(client) + coll_id = collection.id + try: + 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) + 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 = _create_namespace(collection, "ns_concurrent") + ns_id = int(ns.namespace_id) + + results = {} + barrier = threading.Barrier(2) + + def _worker(tag, c): + """Worker.""" + try: + barrier.wait(timeout=10) + _drop_namespace_via_pl(c, coll_id, ns_id) + results[tag] = "ok" + except Exception as exc: + 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"): + with contextlib.suppress(Exception): + client_b.close() + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) 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..b496e472 --- /dev/null +++ b/tests/integration_tests/test_namespace_get_delete_filters.py @@ -0,0 +1,262 @@ +""" +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 + - cross-namespace get(where=...) on a shared connection (regression) +""" + +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): + """Test delete where metadata tag.""" + 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) + 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): + """Test delete where document obsolete.""" + 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: + 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=gt["keep"]) + finally: + cleanup(db_client, collection) + + +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() + 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): + """Test get where document contains python.""" + 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: + 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): + """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") + ns_a.prewarm() + ns_b.prewarm() + 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"}, + 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) + + +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): + """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") + ns_a.prewarm() + ns_b.prewarm() + 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) + + +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..977c59c4 --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_combined.py @@ -0,0 +1,73 @@ +""" +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) + +Vector branches run under both L2 and cosine IVF index metrics. +""" + +from __future__ import annotations + +import time +from typing import Any, ClassVar + +import pytest +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, + teardown_large_fts_collection, +) + + +@pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) +class TestNamespaceHybridSearchCombined: + """TestNamespaceHybridSearchCombined class.""" + + _shared_by_mode: ClassVar[dict[str, dict[str, Any]]] = {} + + @pytest.fixture(autouse=True) + def _bind_shared_collection( + self, + db_client: Any, + 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) + self._corpus = entry["corpus"] + self._collection = entry["collection"] + self._vector_distance = vector_distance + + @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, + 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, 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) + + +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..ae5ad05f --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext.py @@ -0,0 +1,137 @@ +""" +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, + assert_not_contains_no_token_leak, + get_fts_case, + 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: + """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) + 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: + """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, + self._corpus, + namespace_name=ns_name, + ) + + 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)) + + 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..7824d774 --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_fulltext_multi_coll_multi_ns.py @@ -0,0 +1,190 @@ +""" +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 loaded namespace, cross-namespace data +isolation, and empty namespace isolation. + +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 time + +import pytest +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT +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, + 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, + 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)): + for ns_suffix in ("x", "y"): + 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: + 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 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) + 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) + + def test_cross_quadrant_fts_isolation(self, db_client): + """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") + for key in MULTI_COLL_MULTI_NS_QUADRANT_KEYS: + 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, + 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): + """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: + 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): + """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: + _, 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] == f"{key}_zpx_top_5", ( + 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) + + +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..7067c5cc --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_multi_coll_multi_ns.py @@ -0,0 +1,153 @@ +""" +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, + VectorDistanceMetric, + 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 ( + assert_hybrid_search_index_no_hits, + 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_knn_case_on_quadrants, + run_search_index_case_on_quadrants, + setup_multi_coll_multi_ns_fts, + setup_multi_coll_multi_ns_fts_single_loaded, + teardown_multi_coll_multi_ns_fts, +) + + +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") + 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] + 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): + """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") + 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): + """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) + 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): + """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) + 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, 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) + finally: + 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"): + 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) + + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) + 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) + 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): + """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) + 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..14a90fdb --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_search_index.py @@ -0,0 +1,68 @@ +""" +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 + +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: + """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) + 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: + """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)) + + @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_triple_branch.py b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py new file mode 100644 index 00000000..fa6fef01 --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch.py @@ -0,0 +1,111 @@ +""" +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. For ``verify="fts"``, hits and +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 and l2" -v -s +""" + +from __future__ import annotations + +import time +from typing import Any, ClassVar + +import pytest +from namespace_hybrid_search_helpers import ( + TOKEN_ZPX, + TRIPLE_BRANCH_CASES, + VectorDistanceMetric, + ensure_shared_hybrid_search_collection, + get_triple_branch_case, + run_hybrid_triple_branch_case, + setup_fts_namespace_with_corpus, + 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, + 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) + self._corpus = entry["corpus"] + self._collection = entry["collection"] + self._vector_distance = vector_distance + + @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, + 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): + """Test hybrid search triple branch.""" + case = get_triple_branch_case(case_name) + namespace = self._new_namespace(case_name) + 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.""" + 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": 5, + }, + knn={ + "query_embeddings": [1.0, 1.0, 0.0], + "where": {"rel_hint": {"$gte": 0}}, + "n_results": 20, + }, + 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, strict=True): + assert TOKEN_ZPX.lower() in (doc_text or "").lower(), f"id={doc_id!r} must contain {TOKEN_ZPX!r}" + + +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..db26f88d --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_triple_branch_multi_coll_multi_ns.py @@ -0,0 +1,263 @@ +""" +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_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT +from namespace_fts_helpers import ( + CORPUS_SIZE, + MULTI_COLL_MULTI_NS_FTS_LOADED_QUADRANTS, + MULTI_COLL_MULTI_NS_QUADRANT_KEYS, + TOKEN_ZPX, + CorpusRecord, + VectorDistanceMetric, + 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]: + """Namespace corpus.""" + 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, + *, + 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_{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)): + for ns_suffix in ("x", "y"): + 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: + 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: + """TestNamespaceHybridSearchTripleBranchMultiCollMultiNs class.""" + + @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, 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, + distance_metric=vector_distance, + ) + 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): + """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") + 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) + + @pytest.mark.parametrize("vector_distance", ["l2", "cosine"]) + 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) + 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): + """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) + 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): + """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) + 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, 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(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, 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(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, 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(ctx, case_name, 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_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, 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, + 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), ( + 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 new file mode 100644 index 00000000..0a2b06ea --- /dev/null +++ b/tests/integration_tests/test_namespace_hybrid_search_vector.py @@ -0,0 +1,97 @@ +""" +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 + +import time +from typing import Any, ClassVar + +import pytest +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, + 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, + 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) + self._corpus = entry["corpus"] + self._collection = entry["collection"] + self._vector_distance = vector_distance + + @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, + self._corpus, + namespace_name=ns_name, + ) + + def _run_knn_case(self, case_name: str) -> None: + """Run knn case.""" + namespace = self._new_namespace(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): + """Test hybrid search vector knn cases.""" + self._run_knn_case(case_name) + + 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") + 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, + distance_metric=self._vector_distance, + )[0] + ) + + +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..0ceb3e8e --- /dev/null +++ b/tests/integration_tests/test_namespace_lifecycle.py @@ -0,0 +1,400 @@ +""" +Namespace lifecycle integration tests. +Tests collection creation with use_namespace=True, namespace CRUD, and collection deletion. +""" + +import contextlib +import time + +import pytest +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT + +import pyseekdb +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}" + 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=NAMESPACE_TEST_PARTITION_COUNT, + ) + 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 + 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): + """Test get collection preserves namespace flag.""" + 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_get_collection_restores_embedding_function(self, db_client): + """Test get collection restores embedding function.""" + from pyseekdb import DefaultEmbeddingFunction + + ef = DefaultEmbeddingFunction() + name = f"test_ns_ef_{int(time.time() * 1000)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=ef.dimension, distance="cosine", centroids_fresh_mode="spfresh"), + embedding_function=ef, + ), + ) + 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 + 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): + """Test create and get namespace.""" + 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): + """Test get or create namespace.""" + 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): + """Test has namespace.""" + 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): + """Test list namespaces.""" + 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): + """Test delete namespace.""" + 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): + """Test get nonexistent namespace raises.""" + 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_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=r"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) + 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_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", centroids_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_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): + """Test namespace ops blocked after collection deleted.""" + 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): + """Test collection data api blocked when namespace enabled.""" + 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): + """Test namespace on non namespace collection raises.""" + 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 (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}'") + + 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( + hnsw=HNSWConfiguration(dimension=3), + embedding_function=None, + ), + ) + 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): + """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", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + collection = client.create_collection( + name=name, + schema=schema, + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + + 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) + + 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)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="cosine", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + 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. + 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: + 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( + name=name, + configuration=pyseekdb.HNSWConfiguration(dimension=3), + embedding_function=None, + partition_count=4, + ) + assert not db_client.has_collection(name) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) 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..c8e2c71c --- /dev/null +++ b/tests/integration_tests/test_namespace_lifecycle_concurrency.py @@ -0,0 +1,159 @@ +"""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 +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, cleanup, ns_schema + +import pyseekdb + + +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")), + 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: + """Drop worker.""" + try: + clients[0].get_collection(collection.name).delete_namespace(ns_name) + except Exception as exc: + errors.append(exc) + finally: + drop_done.set() + + def _has_worker() -> None: + """Has worker.""" + 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: + 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_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(), + 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() + errors: list[Exception] = [] + has_results: list[bool] = [] + has_completed_while_drop_blocked = threading.Event() + + 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): + """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) + except Exception as exc: + errors.append(exc) + + def _has_worker() -> None: + """Has worker.""" + try: + 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: + errors.append(exc) + + drop_thread = threading.Thread(target=_drop_worker) + has_thread = threading.Thread(target=_has_worker) + 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 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: + with contextlib.suppress(Exception): + client.close() + cleanup(oceanbase_client, collection) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) 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..db064223 --- /dev/null +++ b/tests/integration_tests/test_namespace_optional_indexes.py @@ -0,0 +1,224 @@ +""" +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 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 +from pyseekdb.client.schema import Schema + + +def _schema_fts_only() -> Schema: + """Schema fts only.""" + return Schema( + vector_index=VectorIndexConfig(embedding_function=None), + fulltext_index=FulltextIndexConfig(analyzer="ik"), + ) + + +def _schema_vector_only() -> Schema: + """Schema vector only.""" + return Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + + +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}'") + 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: + """Create collection.""" + name = f"test_ns_opt_idx_{label}_{int(time.time() * 1000)}" + 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]: + """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)], + 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): + """Test fts only indexes and queries.""" + 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): + """Test vector only indexes and queries.""" + 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("centroids_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] + + 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, + ) + 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): + """Test search only indexes and queries.""" + 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) + + 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/integration_tests/test_namespace_prewarm.py b/tests/integration_tests/test_namespace_prewarm.py new file mode 100644 index 00000000..ff4490c2 --- /dev/null +++ b/tests/integration_tests/test_namespace_prewarm.py @@ -0,0 +1,86 @@ +""" +Namespace prewarm integration tests. + +- embedded mode: prewarm raises ValueError +- oceanbase SS mode: prewarm inserts hot_table record, repeat prewarm updates last_access_time +""" + +import time + +import pytest + +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: + """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)}" + 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=NAMESPACE_TEST_PARTITION_COUNT, + ) + namespace = collection.create_namespace("pw_ns") + 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( + f"SELECT namespace_id, last_access_time FROM `{hot_table}` WHERE namespace_id = {namespace_id}" + ) + return rows + + 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() + + 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: + oceanbase_client.delete_collection(name=collection.name) + + 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() + 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) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) 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..9ab5d1aa --- /dev/null +++ b/tests/integration_tests/test_namespace_prewarm_lob.py @@ -0,0 +1,377 @@ +""" +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 +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): + """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}'", + ) + 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: + """BaseLobPrewarm class.""" + + def _make_collection(self, client, name, partitions): + """Make collection.""" + 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, 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 + 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) 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): + """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)}" + 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, 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: 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): + """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): + """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"]) 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..5d0f560b --- /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 +from unittest.mock import MagicMock + +import pytest + +import pyseekdb +from pyseekdb.client.configuration import HNSWConfiguration, SparseVectorIndexConfig, VectorIndexConfig +from pyseekdb.client.schema import Schema +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/integration_tests/test_namespace_query.py b/tests/integration_tests/test_namespace_query.py new file mode 100644 index 00000000..60c6e2e8 --- /dev/null +++ b/tests/integration_tests/test_namespace_query.py @@ -0,0 +1,423 @@ +""" +Namespace query integration tests. +Tests vector similarity query, metadata filtering, include control, and multi-namespace isolation. +""" + +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}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + fulltext_index=FulltextIndexConfig(analyzer="ik"), + ) + collection = client.create_collection( + name=name, + schema=schema, + use_namespace=True, + partition_count=NAMESPACE_TEST_PARTITION_COUNT, + ) + return collection + + def _insert_data(self, ns): + """Insert data.""" + 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): + """Test basic vector query.""" + 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): + """Test query with metadata filter.""" + 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 result.get("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): + """Test query with score filter.""" + 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): + """Test query with include.""" + 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): + """Test query include embeddings.""" + 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_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") + 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): + """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") + 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): + """Test hybrid search fulltext only.""" + 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): + """Test hybrid search vector only.""" + 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): + """Test hybrid search combined.""" + 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): + """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") + 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) + + 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): + """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: + 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) + + 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"]) 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..8e35064f --- /dev/null +++ b/tests/integration_tests/test_namespace_query_where_document.py @@ -0,0 +1,275 @@ +""" +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 + +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 [] + 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', 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 [] + 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', 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 [] + 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', 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 [] + 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, 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 [] + 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()) + ] + assert not violations, ( + f"Nested $and($or) should filter (machine|database) + learning, 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 [] + assert docs, "expected at least one hit for nested $or($and)" + 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/integration_tests/test_namespace_reconnect.py b/tests/integration_tests/test_namespace_reconnect.py new file mode 100644 index 00000000..67db17c0 --- /dev/null +++ b/tests/integration_tests/test_namespace_reconnect.py @@ -0,0 +1,147 @@ +""" +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 + +from namespace_dml_helpers import cleanup, create_ns_collection + +import pyseekdb + + +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: + """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 + 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/integration_tests/test_namespace_session_vars.py b/tests/integration_tests/test_namespace_session_vars.py new file mode 100644 index 00000000..05b42a58 --- /dev/null +++ b/tests/integration_tests/test_namespace_session_vars.py @@ -0,0 +1,97 @@ +""" +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 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)}" + schema = Schema( + vector_index=VectorIndexConfig( + ivf=IVFConfiguration(dimension=3, distance="l2", centroids_fresh_mode="spfresh"), + embedding_function=None, + ), + ) + return client.create_collection( + 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") + 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): + """Test session collection id after create.""" + 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): + """Test session ns ids after create namespace.""" + 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): + """Test session ns id after get namespace.""" + collection = self._create_ns_collection(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 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 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) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) 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..f9f897a2 --- /dev/null +++ b/tests/integration_tests/test_namespace_upsert_concurrency.py @@ -0,0 +1,79 @@ +"""Concurrent namespace.upsert() repro for duplicate business ids.""" + +from __future__ import annotations + +import contextlib +import os +import threading +import time + +from namespace_dml_helpers import NAMESPACE_TEST_PARTITION_COUNT, cleanup, ns_schema + +import pyseekdb + + +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")), + 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, +): + """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(), + 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: + """Upsert once.""" + try: + barrier.wait(timeout=30) + 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: + 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(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"]) + 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_build_document_query.py b/tests/unit_tests/test_build_document_query.py new file mode 100644 index 00000000..f35d67fa --- /dev/null +++ b/tests/unit_tests/test_build_document_query.py @@ -0,0 +1,74 @@ +"""Unit tests for BaseClient._build_document_query DSL generation.""" + +from unittest.mock import 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: + """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 + 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): + """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 + 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): + """Test single contains no default operator.""" + 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): + """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 + qs = result["query_string"] + assert qs["default_operator"] == "and" + assert qs["boost"] == 2.0 + + def test_and_contains_three_terms(self, client): + """Test and contains three terms.""" + 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_build_knn_filter.py b/tests/unit_tests/test_build_knn_filter.py new file mode 100644 index 00000000..f81c41f3 --- /dev/null +++ b/tests/unit_tests/test_build_knn_filter.py @@ -0,0 +1,40 @@ +"""Unit tests for KNN filter DSL generation with hoisted must_not.""" + +from unittest.mock import patch + +import pytest + + +@pytest.fixture() +def client(): + """Client.""" + from pyseekdb.client.client_base import BaseClient + + with patch.multiple(BaseClient, __abstractmethods__=set()): + instance = BaseClient.__new__(BaseClient) + return instance + + +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", + 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 new file mode 100644 index 00000000..9dc95202 --- /dev/null +++ b/tests/unit_tests/test_build_query_expression.py @@ -0,0 +1,130 @@ +"""Unit tests for BaseClient._build_query_expression DSL generation.""" + +from unittest.mock import patch + +import pytest + + +@pytest.fixture() +def client(): + """Client.""" + from pyseekdb.client.client_base import BaseClient + + with patch.multiple(BaseClient, __abstractmethods__=set()): + instance = BaseClient.__new__(BaseClient) + return instance + + +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}} + 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_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"}, + }) + + assert expr == { + "bool": { + "must_not": [ + { + "query_string": { + "fields": ["document"], + "query": "TOKENZPX", + } + } + ], + } + } + 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", + 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"] + + +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( + 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): + """Test ne with not contains and positive where.""" + 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"] 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..cc7c9df1 --- /dev/null +++ b/tests/unit_tests/test_document_query_builder.py @@ -0,0 +1,102 @@ +"""Unit tests for document_query_builder.""" + +import sys +from pathlib import Path + +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_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({ + "$and": [ + {"$or": [{"$contains": "a"}, {"$contains": "b"}]}, + {"$contains": "c"}, + ] + }) + 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.""" + 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"]["must_not"] == [ + {"query_string": {"fields": ["document"], "query": "x"}}, + ] + 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]}} 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..1a05a52b 100644 --- a/tests/unit_tests/test_get_or_create_collection_concurrency.py +++ b/tests/unit_tests/test_get_or_create_collection_concurrency.py @@ -15,30 +15,101 @@ 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: + """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'\")") + 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'\")") + 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): + """Test existing catalog row is reused without insert.""" + 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: + """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 @@ -46,9 +117,20 @@ 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) - 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") @@ -59,8 +141,10 @@ 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) - 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") @@ -73,6 +157,82 @@ 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): + """Test conflict on namespace collection resumes incomplete handle.""" + 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): + """Test conflict on namespace collection without metadata reraises.""" + 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): + """Test resumes incomplete namespace collection when present.""" + 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): + """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") + + with pytest.raises(ValueError, match="invalid dimension"): + BaseClient.get_or_create_collection(client, "items") + + +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 = [ + ("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_get_or_create_namespace_concurrency.py b/tests/unit_tests/test_get_or_create_namespace_concurrency.py new file mode 100644 index 00000000..cb29b191 --- /dev/null +++ b/tests/unit_tests/test_get_or_create_namespace_concurrency.py @@ -0,0 +1,94 @@ +""" +Unit tests for concurrent-safe get_or_create_namespace helpers. +""" + +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 ( # noqa: E402 + BaseClient, + _is_namespace_catalog_conflict_error, +) + + +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'\")") + 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'\")") + 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 + 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): + """Test existing namespace without ltable creates default ltable.""" + 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"]) diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py new file mode 100644 index 00000000..b3d63c8e --- /dev/null +++ b/tests/unit_tests/test_namespace.py @@ -0,0 +1,1536 @@ +""" +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 +from pyseekdb.client.validators import _validate_include, _validate_n_results, _validate_namespace_name # noqa: E402 + +# ==================== IVFConfiguration Tests ==================== + + +class TestIVFConfiguration: + """TestIVFConfiguration class.""" + + def test_valid_defaults(self): + """Test valid defaults.""" + config = IVFConfiguration() + assert config.dimension == 384 + assert config.distance == "cosine" + assert config.centroids_fresh_mode is None + 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}}) + + +# ==================== VectorIndexConfig Tests ==================== + + +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 + + +# ==================== Collection Namespace Guard Tests ==================== + + +class TestCollectionNamespaceGuard: + """TestCollectionNamespaceGuard class.""" + + def _make_collection(self, use_namespace: bool) -> Collection: + """Make 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): + """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() + + +# ==================== Collection Namespace Management Guard Tests ==================== + + +class TestCollectionNamespaceManagementGuard: + """TestCollectionNamespaceManagementGuard class.""" + + def _make_collection(self, use_namespace: bool) -> Collection: + """Make 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): + """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") + + +# ==================== Namespace Object Tests ==================== + + +class TestNamespaceObject: + """TestNamespaceObject class.""" + + def _make_namespace(self) -> Namespace: + """Make 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): + """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 + assert "100" in r + 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() + + +# ==================== Namespace SQL Generation Tests ==================== + + +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: + 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): + """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) + 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 + + # 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) + 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): + """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 + + +class TestNamespaceSQLGeneration: + """Verify SQL statements generated by BaseClient._namespace_* methods.""" + + COLLECTION_ID = "42" + NAMESPACE_ID = "7" + TABLE = "42_logic_data_table" + + def _client(self): + """Client.""" + return FakeClient() + + def _common_kwargs(self): + """Common kwargs.""" + 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.""" + 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._ivf_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_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._ivf_kwargs(), + ids="d1", + embeddings=[1.0, 2.0, 3.0], + documents="hello", + embedding_function=MagicMock(), + ) + 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() + c._namespace_add( + **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"], + 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): + """Test add without metadata sql.""" + c = self._client() + c._namespace_add( + **self._ivf_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): + """Test update metadata sql.""" + 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 "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 "'d1'" in sql + + def test_update_embedding_and_document_sql(self): + """Test update embedding and document sql.""" + c = self._client() + c._namespace_update( + **self._ivf_kwargs(), + ids="d1", + embeddings=[9.0, 8.0, 7.0], + documents="Updated", + ) + 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 ---- + + def test_delete_by_ids_sql(self): + """Test delete by ids sql.""" + 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_UNQUOTE(JSON_EXTRACT(data_content, '$.id')) = " in sql + 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(), + 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 + + def test_delete_by_where_document_sql(self): + """Test delete by where document sql.""" + 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 "MATCH(document) AGAINST" in sql + assert "obsolete" in sql + + # ---- QUERY ---- + + def test_query_basic_dsl(self): + """Test query basic dsl.""" + c = self._client() + c.query_return_value = [] + c._namespace_query( + **self._ivf_kwargs(), + query_embeddings=[1.0, 0.0, 0.0], + n_results=3, + distance="l2", + ) + sql = c.query_sqls[-1] + assert "hybrid_search(TABLE" in sql + assert f"`{self.TABLE}`" 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_dsl(self): + """Test query with where dsl.""" + c = self._client() + c.query_return_value = [] + c._namespace_query( + **self._ivf_kwargs(), + query_embeddings=[1.0, 0.0, 0.0], + n_results=5, + where={"category": "AI"}, + distance="cosine", + ) + sql = c.query_sqls[-1] + assert "hybrid_search(TABLE" in sql + assert "data_content.metadata.category" in sql + assert "cosine_distance(embedding," not in sql + + # ---- GET ---- + + def test_get_by_ids_sql(self): + """Test get by ids sql.""" + 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_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( + **self._common_kwargs(), + limit=10, + ) + sql = c.query_sqls[-1] + assert "LIMIT 10" in sql + + # ---- 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()) + 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 + + 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", centroids_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(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 + assert "centroids_fresh_mode=spfresh" 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")) + + 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", centroids_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 "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) + 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_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( + 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 "centroids_fresh_mode" not in data_create + + +# ==================== Namespace Name Validation Tests ==================== + + +class TestValidateNamespaceName: + """TestValidateNamespaceName class.""" + + def test_valid_simple_name(self): + """Test valid simple name.""" + _validate_namespace_name("my_namespace") + + def test_valid_with_digits_and_underscore(self): + """Test valid with digits and underscore.""" + _validate_namespace_name("tenant_123_abc") + + def test_valid_single_char(self): + """Test valid single char.""" + _validate_namespace_name("a") + + def test_valid_boundary_256_chars(self): + """Test valid boundary 256 chars.""" + _validate_namespace_name("a" * 256) + + def test_empty_name_raises(self): + """Test empty name raises.""" + with pytest.raises(ValueError, match="must not be empty"): + _validate_namespace_name("") + + def test_non_string_raises(self): + """Test non string raises.""" + with pytest.raises(TypeError, match="must be a string"): + _validate_namespace_name(123) + + def test_too_long_raises(self): + """Test too long raises.""" + with pytest.raises(ValueError, match="too long"): + _validate_namespace_name("a" * 257) + + def test_hyphen_raises(self): + """Test hyphen raises.""" + with pytest.raises(ValueError, match="invalid characters"): + _validate_namespace_name("my-namespace") + + def test_space_raises(self): + """Test space raises.""" + with pytest.raises(ValueError, match="invalid characters"): + _validate_namespace_name("my namespace") + + def test_dot_raises(self): + """Test dot raises.""" + with pytest.raises(ValueError, match="invalid characters"): + _validate_namespace_name("my.namespace") + + def test_chinese_raises(self): + """Test chinese raises.""" + 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, + 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 ==================== + + +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") + + +# ==================== Namespace Batch Limit Tests ==================== + + +class TestNamespaceBatchLimit: + """TestNamespaceBatchLimit class.""" + + def _client(self): + """Client.""" + return FakeClient() + + def _common_kwargs(self): + """Common kwargs.""" + 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.""" + 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)] + 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): + """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)] + 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): + """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") + + +# ==================== Physical Table Names Tests ==================== + + +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 + + +# ==================== Namespace Catalog Tests ==================== + + +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() + + sql = "\n".join(c.executed_sqls) + 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 `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 `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 + assert "included_index BOOL" in sql + assert "PRIMARY KEY (namespace_id, ltable_id, included_index)" in sql + 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() + + 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): + """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() + 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) + + +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", + "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): + """Test purge skips complete collection.""" + 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): + """Test get collection purges incomplete namespace collection.""" + + class GetClient(FakeClient): + """GetClient class.""" + + 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) + + +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"): + _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 ==================== + + +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 + + hnsw = HNSWConfiguration(dimension=128) + schema = Schema(vector_index=VectorIndexConfig(hnsw=hnsw, embedding_function=None)) + 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): + """Test ob type validation.""" + c = FakeClient() + c.detect_db_type_and_version = MagicMock(return_value=("mysql", "8.0")) + from pyseekdb.client.configuration import VectorIndexConfig + 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 LakeBase"): + 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() + 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 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) + 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 "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() + 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 pyseekdb.client.configuration import VectorIndexConfig + from pyseekdb.client.schema import Schema + + 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 "centroids_fresh_mode" not in data_create + settings = c._create_ns_collection_meta.call_args[0][1] + assert settings["dense_index_type"] == "ivf" + assert "centroids_fresh_mode" not in settings + + +# ==================== Delete Namespace Uses Kernel ==================== + + +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._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) + 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_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) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) 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..03534da8 --- /dev/null +++ b/tests/unit_tests/test_namespace_kernel_errors.py @@ -0,0 +1,110 @@ +""" +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", + ), + 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", + ), + 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", + ), + 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) 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..69638055 --- /dev/null +++ b/tests/unit_tests/test_namespace_lifecycle_lock.py @@ -0,0 +1,43 @@ +"""Unit tests for namespace catalog concurrency without GET_LOCK.""" + +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 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 + + assert BaseClient._has_ns_namespace(client, "cid", "ns1") is False + 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", + "namespace_name": "ns1", + "ltable_id": "7", + } + + with pytest.raises(ValueError, match="already exists"): + BaseClient._create_ns_namespace_meta(client, "cid", "ns1") + + client._insert_ns_namespace_catalog_row.assert_not_called() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) 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..1b638490 --- /dev/null +++ b/tests/unit_tests/test_namespace_no_index_embedding_dim.py @@ -0,0 +1,60 @@ +"""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_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], + 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_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( + [[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, + ) 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..30404284 --- /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 import IVFConfiguration # noqa: E402 +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 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) 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..2ba79e38 --- /dev/null +++ b/tests/unit_tests/test_namespace_upsert_reconcile.py @@ -0,0 +1,98 @@ +""" +Unit tests for namespace upsert duplicate-record reconciliation. +""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +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: + """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 + + 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): + """Test reconcile collapses duplicate rows.""" + client = MagicMock(spec=BaseClient) + client._count_namespace_records_by_id.side_effect = [4, 0, 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=["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]] + + 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 + + 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__": + pytest.main([__file__, "-v"]) 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"""