diff --git a/src/pyseekdb/client/client_base.py b/src/pyseekdb/client/client_base.py index f7b9acd2..35222adc 100644 --- a/src/pyseekdb/client/client_base.py +++ b/src/pyseekdb/client/client_base.py @@ -95,6 +95,12 @@ _LAKEBASE_VERSION_MARKER = "database ai" +# Catalog DDL is shared by every client process. Keep retries short and +# bounded: success is determined by re-reading the catalog state, never by +# merely suppressing a concurrent DDL exception. +_CATALOG_BOOTSTRAP_BACKOFF_SECONDS = (0.02, 0.05, 0.1, 0.2, 0.4) +_CATALOG_BOOTSTRAP_TRANSIENT_ERROR_CODES = (1050, 1061, 1146, 1205, 1213) + logger = logging.getLogger(__name__) from .types import _NOT_PROVIDED, _NotProvided # noqa: E402, F401 @@ -167,17 +173,54 @@ def _is_sdk_collection_catalog_conflict_error(exc: BaseException) -> bool: return False -def _reraise_unless_unique_index_exists(exc: BaseException) -> None: - """Re-raise unless the exception indicates the unique index is already present.""" - message = str(exc).lower() - if ( - "already exists" in message - or "duplicate key name" in message - or "code=1061" in message - or ("1061" in message and "duplicate" in message) - ): - return - raise exc +def _has_database_error_code(exc: BaseException, error_codes: tuple[int, ...]) -> bool: + """Return whether an exception cause chain contains one of the database error codes.""" + current: BaseException | None = exc + while current is not None: + if current.args and isinstance(current.args[0], int) and current.args[0] in error_codes: + return True + message = str(current).lower() + for error_code in error_codes: + if ( + f"({error_code}," in message + or f"({error_code})" in message + or f"code={error_code}" in message + or f"code: {error_code}" in message + or f"errno {error_code}" in message + ): + return True + current = current.__cause__ + return False + + +def _is_catalog_table_missing_error(exc: BaseException) -> bool: + """Return whether a catalog state probe failed because the table is not visible yet.""" + if _has_database_error_code(exc, (1146,)): + return True + current: BaseException | None = exc + while current is not None: + message = str(current).lower() + if "ret=-5019" in message or "table not exist" in message or "table doesn't exist" in message: + return True + current = current.__cause__ + return False + + +def _is_catalog_bootstrap_transient_error(exc: BaseException) -> bool: + """Return whether catalog bootstrap may safely retry after checking real state again.""" + return _has_database_error_code(exc, _CATALOG_BOOTSTRAP_TRANSIENT_ERROR_CODES) or _is_catalog_table_missing_error( + exc + ) + + +class _CatalogBootstrapStateNotReady(RuntimeError): + """Internal signal used when DDL returned but the current connection cannot verify the result.""" + + +def _require_catalog_bootstrap_state(is_ready: bool, message: str) -> None: + """Raise the internal retry signal unless a catalog state probe succeeded.""" + if not is_ready: + raise _CatalogBootstrapStateNotReady(message) def _extract_hnsw_config(config: ConfigurationParam) -> HNSWConfiguration | None: @@ -1269,8 +1312,8 @@ def _get_embedding_function_dimension(self, embedding_function: EmbeddingFunctio def _create_sdk_collections_if_not_exists(self) -> None: """Create the sdk_collections catalog table if it does not already exist.""" try: - self._use_catalog_database() - sdk_coll = self._qtable(CollectionNames.sdk_collections_table_name()) + sdk_coll_name = CollectionNames.sdk_collections_table_name() + sdk_coll = self._qtable(sdk_coll_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(), '-', '')), @@ -1280,11 +1323,11 @@ def _create_sdk_collections_if_not_exists(self) -> None: updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, 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) + self._ensure_catalog_table( + table_name=sdk_coll_name, + create_table_sql=create_table_sql, + unique_indexes={"uk_sdk_coll_name": ("collection_name",)}, + ) except Exception as e: raise ValueError(f"Failed to create sdk_collections table: {e}") from e @@ -1375,11 +1418,125 @@ 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 _catalog_table_is_visible(self, table_name: str) -> bool: + """Check whether the current connection can resolve a catalog table.""" + try: + self._execute(f"SELECT 1 FROM {self._qtable(table_name)} LIMIT 0") + except Exception as exc: + if _is_catalog_table_missing_error(exc): + return False + raise + else: + return True + + @staticmethod + def _index_row_value(row: Any, key: str, position: int) -> Any: + """Read a SHOW INDEX field from either a dict cursor row or a tuple row.""" + if isinstance(row, dict): + row_by_lower_key = {str(row_key).lower(): value for row_key, value in row.items()} + return row_by_lower_key.get(key.lower()) + if isinstance(row, (tuple, list)) and len(row) > position: + return row[position] + return None + + def _catalog_unique_index_is_ready( + self, + table_name: str, + index_name: str, + expected_columns: tuple[str, ...], + ) -> bool: + """Verify that a named catalog index exists, is unique, and has the expected columns.""" + rows = self._execute(f"SHOW INDEX FROM {self._qtable(table_name)}") or [] + matching_rows = [ + row for row in rows if str(self._index_row_value(row, "Key_name", 2) or "").lower() == index_name.lower() + ] + if not matching_rows: + return False + + ordered_columns: list[tuple[int, str]] = [] + for row in matching_rows: + non_unique = self._index_row_value(row, "Non_unique", 1) + if str(non_unique) != "0": + raise ValueError(f"Catalog index {index_name} on {table_name} exists but is not UNIQUE") + sequence = self._index_row_value(row, "Seq_in_index", 3) + column = self._index_row_value(row, "Column_name", 4) + if sequence is None or column is None: + raise ValueError(f"Unable to verify catalog index {index_name} on {table_name}") + ordered_columns.append((int(sequence), str(column).lower())) + + actual_columns = tuple(column for _, column in sorted(ordered_columns)) + normalized_expected_columns = tuple(column.lower() for column in expected_columns) + if actual_columns != normalized_expected_columns: + raise ValueError( + f"Catalog index {index_name} on {table_name} has columns {actual_columns}, " + f"expected {normalized_expected_columns}" + ) + return True + + def _ensure_catalog_table( + self, + table_name: str, + create_table_sql: str, + unique_indexes: dict[str, tuple[str, ...]], + ) -> None: + """Ensure a shared catalog table is visible and has the required unique indexes.""" + attempts = len(_CATALOG_BOOTSTRAP_BACKOFF_SECONDS) + 1 + last_error: BaseException | None = None + + for attempt in range(attempts): + try: + self._use_catalog_database() + + # Fast path: do not issue DDL once the table and indexes are ready. + table_is_visible = self._catalog_table_is_visible(table_name) + if not table_is_visible: + self._execute(create_table_sql) + _require_catalog_bootstrap_state( + self._catalog_table_is_visible(table_name), + f"Catalog table {table_name} is not visible after CREATE TABLE", + ) + + for index_name, columns in unique_indexes.items(): + if self._catalog_unique_index_is_ready(table_name, index_name, columns): + continue + column_sql = ", ".join(_quote_sql_identifier(column) for column in columns) + self._execute( + f"CREATE UNIQUE INDEX {_quote_sql_identifier(index_name)} " + f"ON {self._qtable(table_name)} ({column_sql})" + ) + _require_catalog_bootstrap_state( + self._catalog_unique_index_is_ready(table_name, index_name, columns), + f"Catalog index {index_name} on {table_name} is not visible after CREATE INDEX", + ) + + # The current connection must still see the final state before success. + _require_catalog_bootstrap_state( + self._catalog_table_is_visible(table_name), + f"Catalog table {table_name} is not visible after bootstrap", + ) + except Exception as exc: + if not isinstance(exc, _CatalogBootstrapStateNotReady) and not _is_catalog_bootstrap_transient_error( + exc + ): + raise + last_error = exc + self._rollback_connection_if_supported() + if attempt < len(_CATALOG_BOOTSTRAP_BACKOFF_SECONDS): + time.sleep(_CATALOG_BOOTSTRAP_BACKOFF_SECONDS[attempt]) + else: + return + + if last_error is not None: + raise last_error + raise RuntimeError(f"Failed to ensure catalog table {table_name}") + 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()) + ns_namespaces_name = NamespaceCollectionNames.sdk_namespaces_table() + ns_ltables_name = NamespaceCollectionNames.sdk_ltables_table() + ns_stats_name = NamespaceCollectionNames.sdk_namespaces_stats_table() + ns_namespaces_q = self._qtable(ns_namespaces_name) + ns_ltables_q = self._qtable(ns_ltables_name) 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, @@ -1404,7 +1561,7 @@ def _ensure_namespace_catalogs(self) -> None: UNIQUE KEY uk_sdk_lt_coll_ns_name (collection_id, namespace_id, ltable_name), KEY idx_sdk_lt_by_ns (collection_id, namespace_id) ) COMMENT='LTable catalog' ORGANIZATION INDEX {scp};""" - namespaces_stats_sql = f"""CREATE TABLE IF NOT EXISTS {self._qtable(NamespaceCollectionNames.sdk_namespaces_stats_table())} ( + namespaces_stats_sql = f"""CREATE TABLE IF NOT EXISTS {self._qtable(ns_stats_name)} ( 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', @@ -1418,22 +1575,21 @@ def _ensure_namespace_catalogs(self) -> None: 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) + self._ensure_catalog_table( + table_name=ns_namespaces_name, + create_table_sql=ns_namespaces_sql, + unique_indexes={"uk_sdk_ns_coll_name": ("collection_id", "namespace_name")}, + ) + self._ensure_catalog_table( + table_name=ns_ltables_name, + create_table_sql=ns_ltables_sql, + unique_indexes={"uk_sdk_lt_coll_ns_name": ("collection_id", "namespace_id", "ltable_name")}, + ) + self._ensure_catalog_table( + table_name=ns_stats_name, + create_table_sql=namespaces_stats_sql, + unique_indexes={}, + ) def _rollback_connection_if_supported(self) -> None: """Roll back the current connection transaction if the backend supports it.""" 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 1a05a52b..55960aa6 100644 --- a/tests/unit_tests/test_get_or_create_collection_concurrency.py +++ b/tests/unit_tests/test_get_or_create_collection_concurrency.py @@ -4,7 +4,7 @@ import sys from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -14,6 +14,8 @@ from pyseekdb.client.client_base import ( # noqa: E402 BaseClient, + _is_catalog_bootstrap_transient_error, + _is_catalog_table_missing_error, _is_collection_conflict_error, _is_sdk_collection_catalog_conflict_error, ) @@ -80,6 +82,207 @@ def test_existing_catalog_row_is_reused_without_insert(self): assert not insert_calls +class TestCollectionCatalogBootstrap: + """Tests for state-driven, retry-bounded catalog initialization.""" + + @staticmethod + def _client(execute_side_effect): + """Build a client mock with the real catalog state helpers bound.""" + client = MagicMock(spec=BaseClient) + client._qtable.side_effect = lambda table: f"`test_db`.`{table}`" + client._use_catalog_database = MagicMock() + client._rollback_connection_if_supported = MagicMock() + client._execute.side_effect = execute_side_effect + client._catalog_table_is_visible = BaseClient._catalog_table_is_visible.__get__(client, BaseClient) + client._catalog_unique_index_is_ready = BaseClient._catalog_unique_index_is_ready.__get__(client, BaseClient) + client._index_row_value = BaseClient._index_row_value + return client + + @staticmethod + def _unique_index_row(column="collection_name", *, non_unique=0): + """Build one tuple-shaped SHOW INDEX row.""" + return ("sdk_collections", non_unique, "uk_sdk_coll_name", 1, column) + + def test_fast_path_uses_no_ddl_when_catalog_is_ready(self): + """A ready catalog is only probed and never receives redundant DDL.""" + + def execute(sql): + if sql.startswith("SHOW INDEX"): + return [self._unique_index_row()] + return [] + + client = self._client(execute) + + BaseClient._ensure_catalog_table( + client, + "sdk_collections", + "CREATE TABLE IF NOT EXISTS sdk_collections (...) ", + {"uk_sdk_coll_name": ("collection_name",)}, + ) + + sql_calls = [call.args[0] for call in client._execute.call_args_list] + assert not any(sql.startswith("CREATE") for sql in sql_calls) + assert sum(sql.startswith("SHOW INDEX") for sql in sql_calls) == 1 + + def test_missing_table_is_created_and_inline_unique_index_is_verified(self): + """A missing table is created once and its inline unique key avoids extra index DDL.""" + state = {"table": False} + + def execute(sql): + if sql.startswith("SELECT 1") and not state["table"]: + raise RuntimeError('(1146, "Table test_db.sdk_collections doesn\'t exist")') + if sql.startswith("CREATE TABLE"): + state["table"] = True + return None + if sql.startswith("SHOW INDEX"): + return [self._unique_index_row()] + return [] + + client = self._client(execute) + + BaseClient._ensure_catalog_table( + client, + "sdk_collections", + "CREATE TABLE IF NOT EXISTS sdk_collections (...) ", + {"uk_sdk_coll_name": ("collection_name",)}, + ) + + sql_calls = [call.args[0] for call in client._execute.call_args_list] + assert sum(sql.startswith("CREATE TABLE") for sql in sql_calls) == 1 + assert not any(sql.startswith("CREATE UNIQUE INDEX") for sql in sql_calls) + + def test_existing_legacy_table_adds_only_the_missing_unique_index(self): + """An old table without the required index receives exactly one index DDL.""" + state = {"index": False} + + def execute(sql): + if sql.startswith("SHOW INDEX"): + return [self._unique_index_row()] if state["index"] else [] + if sql.startswith("CREATE UNIQUE INDEX"): + state["index"] = True + return [] + + client = self._client(execute) + + BaseClient._ensure_catalog_table( + client, + "sdk_collections", + "CREATE TABLE IF NOT EXISTS sdk_collections (...) ", + {"uk_sdk_coll_name": ("collection_name",)}, + ) + + sql_calls = [call.args[0] for call in client._execute.call_args_list] + assert sum(sql.startswith("CREATE UNIQUE INDEX") for sql in sql_calls) == 1 + assert not any(sql.startswith("CREATE TABLE") for sql in sql_calls) + + def test_duplicate_index_race_rechecks_state_before_succeeding(self): + """A concurrent 1061 is retried and accepted only after the index is visible.""" + state = {"create_attempts": 0} + + def execute(sql): + if sql.startswith("SHOW INDEX"): + return [self._unique_index_row()] if state["create_attempts"] else [] + if sql.startswith("CREATE UNIQUE INDEX"): + state["create_attempts"] += 1 + raise RuntimeError('(1061, "Duplicate key name uk_sdk_coll_name")') + return [] + + client = self._client(execute) + + with patch("pyseekdb.client.client_base.time.sleep") as sleep: + BaseClient._ensure_catalog_table( + client, + "sdk_collections", + "CREATE TABLE IF NOT EXISTS sdk_collections (...) ", + {"uk_sdk_coll_name": ("collection_name",)}, + ) + + assert state["create_attempts"] == 1 + client._rollback_connection_if_supported.assert_called_once() + sleep.assert_called_once() + + def test_persistent_schema_visibility_race_exhausts_bounded_retries(self): + """Persistent 1146 errors are not swallowed as successful initialization.""" + + def execute(sql): + if sql.startswith("SELECT 1"): + raise RuntimeError('(1146, "Table test_db.sdk_collections doesn\'t exist")') + return [] + + client = self._client(execute) + + with ( + patch("pyseekdb.client.client_base.time.sleep") as sleep, + pytest.raises(RuntimeError, match="not visible after CREATE TABLE"), + ): + BaseClient._ensure_catalog_table( + client, + "sdk_collections", + "CREATE TABLE IF NOT EXISTS sdk_collections (...) ", + {"uk_sdk_coll_name": ("collection_name",)}, + ) + + assert sleep.call_count == 5 + assert client._rollback_connection_if_supported.call_count == 6 + + def test_non_transient_errors_are_not_retried(self): + """Permission and other unrelated failures remain visible to callers.""" + + def execute(_sql): + raise RuntimeError('(1142, "SELECT command denied")') + + client = self._client(execute) + + with ( + patch("pyseekdb.client.client_base.time.sleep") as sleep, + pytest.raises(RuntimeError, match="1142"), + ): + BaseClient._ensure_catalog_table( + client, + "sdk_collections", + "CREATE TABLE IF NOT EXISTS sdk_collections (...) ", + {"uk_sdk_coll_name": ("collection_name",)}, + ) + + sleep.assert_not_called() + client._rollback_connection_if_supported.assert_not_called() + + def test_wrong_existing_index_definition_is_rejected(self): + """A same-named non-unique or wrong-column index is never treated as ready.""" + + def execute(sql): + if sql.startswith("SHOW INDEX"): + return [self._unique_index_row(non_unique=1)] + return [] + + client = self._client(execute) + + with pytest.raises(ValueError, match="not UNIQUE"): + BaseClient._ensure_catalog_table( + client, + "sdk_collections", + "CREATE TABLE IF NOT EXISTS sdk_collections (...) ", + {"uk_sdk_coll_name": ("collection_name",)}, + ) + + def test_transient_error_detection_walks_cause_chain(self): + """Wrapped database error codes remain classifiable.""" + inner = RuntimeError('(1146, "Table test_db.sdk_collections doesn\'t exist")') + outer = ValueError("catalog bootstrap failed") + outer.__cause__ = inner + + assert _is_catalog_table_missing_error(outer) + assert _is_catalog_bootstrap_transient_error(outer) + assert not _is_catalog_bootstrap_transient_error(RuntimeError('(1142, "SELECT command denied")')) + + def test_embedded_table_missing_error_format_is_transient(self): + """Embedded reports symbolic names followed by a parenthesized numeric code.""" + exc = RuntimeError("execute sql failed OB_TABLE_NOT_EXIST(1146): Table '%s.%s' doesn't exist") + + assert _is_catalog_table_missing_error(exc) + assert _is_catalog_bootstrap_transient_error(exc) + + class TestCollectionConflictDetection: """TestCollectionConflictDetection class.""" diff --git a/tests/unit_tests/test_namespace.py b/tests/unit_tests/test_namespace.py index b3d63c8e..2947fa83 100644 --- a/tests/unit_tests/test_namespace.py +++ b/tests/unit_tests/test_namespace.py @@ -7,6 +7,7 @@ """ import os +import re import sys from pathlib import Path from unittest.mock import MagicMock @@ -424,6 +425,8 @@ def __init__(self): self.executed_sqls = [] self.query_sqls = [] self.query_return_value = [] + self.catalog_tables = set() + self.catalog_indexes = {} def _ensure_connection(self): """Ensure connection.""" @@ -470,6 +473,23 @@ def _execute(self, sql, params=None): if os.environ.get("PYSEEKDB_PRINT_SQL", "").lower() in ("1", "true", "yes"): print(f"[pyseekdb SQL] {sql}", flush=True) self.executed_sqls.append(sql) + + table_match = re.search(r"`[^`]+`\.`([^`]+)`", sql) + table_name = table_match.group(1) if table_match else None + if sql.startswith("SELECT 1 FROM") and table_name not in self.catalog_tables: + raise RuntimeError(f'(1146, "Table test.{table_name} doesn\'t exist")') + if sql.startswith("CREATE TABLE") and table_name is not None: + self.catalog_tables.add(table_name) + for index_match in re.finditer(r"UNIQUE KEY\s+`?([^\s`(]+)`?\s*\(([^)]+)\)", sql, re.IGNORECASE): + index_name = index_match.group(1) + columns = tuple(column.strip().strip("`").lower() for column in index_match.group(2).split(",")) + self.catalog_indexes.setdefault(table_name, {})[index_name] = columns + if sql.startswith("SHOW INDEX") and table_name is not None: + return [ + (table_name, 0, index_name, position, column) + for index_name, columns in self.catalog_indexes.get(table_name, {}).items() + for position, column in enumerate(columns, start=1) + ] return None # Bypass the sdk_ltables lookup in unit tests: SQL-generation tests don't @@ -1268,18 +1288,17 @@ def test_ensure_namespace_catalogs_creates_all_catalog_tables(self): 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.""" + def test_ensure_namespace_catalogs_creates_only_missing_tables_in_order(self): + """Catalog bootstrap creates missing tables in order without redundant index DDL.""" 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] + create_table_sqls = [sql for sql in c.executed_sqls if sql.startswith("CREATE TABLE")] + assert len(create_table_sqls) == 3 + assert "`test`.`sdk_namespaces`" in create_table_sqls[0] + assert "`test`.`sdk_ltables`" in create_table_sqls[1] + assert "`test`.`sdk_namespaces_stats`" in create_table_sqls[2] + assert not any(sql.startswith("CREATE UNIQUE INDEX") for sql in c.executed_sqls) def test_delete_ns_collection_meta_cleans_namespaces_stats_table(self): """Test delete ns collection meta cleans namespaces stats table."""