From da207da106593dfa8ef33e682fa2885ac882abb1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 13:25:19 +0300 Subject: [PATCH] Read the table list and settings once per open Opening a database ran `list_tables` three times, opened the settings table twice, and read and parsed the same settings row three times: once for the stored vector dimension, once for the version behind the migration check, and once for config validation. On object storage each of those is a round trip. `_initialize` now reads both once and threads them down. `_init_tables` and `_check_migrations` take what it read instead of fetching their own copy, and `validate_config_compatibility` accepts the settings it should compare against, still reading for itself when called directly. Passing the pre-init read to validation is equivalent: nothing between the read and the validation rewrites `embeddings`, which is all it compares. The settings read no longer swallows every exception. It did before, when the only consequence was falling back to the configured vector dimension; now the same empty result feeds the migration check, where it would read as version 0.0.0 and declare every migration pending. Only decode failures are tolerated, and a decoded non-object normalizes to {} rather than reaching callers that expect a mapping. --- haiku_rag_slim/haiku/rag/store/engine.py | 99 +++++++++---------- .../haiku/rag/store/repositories/settings.py | 7 +- tests/store/test_open_path.py | 71 +++++++++++++ tests/test_lancedb_connection.py | 4 +- 4 files changed, 125 insertions(+), 56 deletions(-) create mode 100644 tests/store/test_open_path.py diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index f47a5d82..08867fa6 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -207,6 +207,11 @@ def get_document_items_arrow_schema() -> pa.Schema: return pa.schema(fields) +def _stored_vector_dim(settings: dict) -> int | None: + """The vector dimension a database's chunks were written at.""" + return settings.get("embeddings", {}).get("model", {}).get("vector_dim") + + def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]: """The index set each table carries.""" match table_name: @@ -533,29 +538,27 @@ class Store: self._config, self.db_path ) - # For remote stores (and as a safety net for local paths that exist but - # have no tables — e.g. a previously failed init), detect new DB by - # checking whether any tables exist. - is_new_db = self._is_new_db - if not is_new_db: - existing_tables = (await self.db.list_tables()).tables - if not existing_tables: - is_new_db = True + # Read once and thread onward: on object storage each of these is a + # round trip. A local path that exists with no tables is a failed init, + # so treat it as new. + existing_tables = (await self.db.list_tables()).tables + is_new_db = self._is_new_db or not existing_tables - # For existing databases, read stored vector dimension to create ChunkRecord - # that can read existing chunks. For new databases, use config's dimension. - stored_vector_dim = None - if not is_new_db: - stored_vector_dim = await self._get_stored_vector_dim() + stored_settings: dict = {} + if not is_new_db and "settings" in existing_tables: + self.settings_table = await self.db.open_table("settings") + stored_settings = await self._read_stored_settings() - # Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB) + # An existing database's chunks can only be read with the dimension they + # were written at. + stored_vector_dim = _stored_vector_dim(stored_settings) chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim) # Initialize tables (creates them if they don't exist). For an existing # DB this raises MigrationRequiredError up front when migrations are # pending, before creating any newly-introduced table. - await self._init_tables(is_new_db) + await self._init_tables(is_new_db, existing_tables, stored_settings) # Set version for new databases. if is_new_db and not self._read_only: @@ -563,7 +566,7 @@ class Store: # Validate config compatibility after connection is established if not self._skip_validation: - await self._validate_configuration() + await self._validate_configuration(stored_settings) async def __aenter__(self): # If _initialize connects to LanceDB but then fails (e.g. migration @@ -585,33 +588,26 @@ class Store: """Whether the store is in read-only mode.""" return self._read_only - async def _get_stored_vector_dim(self) -> int | None: - """Read the stored vector dimension from the settings table. + async def _read_stored_settings(self) -> dict: + """The stored settings blob, or {} if it is absent or not a JSON object. - Returns: - The stored vector dimension, or None if not found. + Only decoding failures are tolerated. A storage failure must propagate: + read as empty settings it would look like version 0.0.0, and the + migration check would declare every migration pending. """ + rows = ( + await self.settings_table.query() + .where("id = 'settings'") + .limit(1) + .to_arrow() + ).to_pylist() + if not rows or not rows[0].get("settings"): + return {} try: - existing_tables = (await self.db.list_tables()).tables - if "settings" not in existing_tables: - return None - - settings_table = await self.db.open_table("settings") - rows = ( - await settings_table.query() - .where("id = 'settings'") - .limit(1) - .to_arrow() - ).to_pylist() - if not rows or not rows[0].get("settings"): - return None - - settings = json.loads(rows[0]["settings"]) - embeddings = settings.get("embeddings", {}) - model = embeddings.get("model", {}) - return model.get("vector_dim") - except Exception: - return None + decoded = json.loads(rows[0]["settings"]) + except (json.JSONDecodeError, TypeError): + return {} + return decoded if isinstance(decoded, dict) else {} def _assert_writable(self) -> None: """Raise ReadOnlyError if the store is in read-only mode.""" @@ -742,16 +738,19 @@ class Store: except Exception as e: logger.warning(f"Could not create vector index: {e}") - async def _validate_configuration(self) -> None: + async def _validate_configuration( + self, stored_settings: dict | None = None + ) -> None: """Validate that the configuration is compatible with the database.""" from haiku.rag.store.repositories.settings import SettingsRepository settings_repo = SettingsRepository(self) - await settings_repo.validate_config_compatibility() + await settings_repo.validate_config_compatibility(stored_settings) - async def _init_tables(self, is_new_db: bool): + async def _init_tables( + self, is_new_db: bool, existing_tables: list[str], stored_settings: dict + ): """Initialize database tables (create if they don't exist).""" - existing_tables = (await self.db.list_tables()).tables # Surface pending migrations BEFORE creating any newly-introduced table. # Otherwise opening a legacy DB would either mutate it (creating an empty @@ -763,8 +762,7 @@ class Store: and not self._skip_migration_check and "settings" in existing_tables ): - self.settings_table = await self.db.open_table("settings") - await self._check_migrations() + await self._check_migrations(stored_settings.get("version", "0.0.0")) missing_tables = set(REQUIRED_TABLES) - set(existing_tables) @@ -811,10 +809,8 @@ class Store: ) await ensure_indexes(self.document_items_table, "document_items") - # Create or open settings table - if "settings" in existing_tables: - self.settings_table = await self.db.open_table("settings") - else: + # _initialize opened the settings table when the database had one. + if "settings" not in existing_tables: self.settings_table = await self.db.create_table( "settings", schema=SettingsRecord ) @@ -828,7 +824,7 @@ class Store: """Set the initial version for a new database.""" await self.set_haiku_version(metadata.version("haiku.rag-slim")) - async def _check_migrations(self) -> None: + async def _check_migrations(self, db_version: str) -> None: """Raise if migrations are pending. Opening never writes the version. Raises: @@ -837,7 +833,6 @@ class Store: from haiku.rag.store.upgrades import get_pending_upgrades current_version = metadata.version("haiku.rag-slim") - db_version = await self.get_haiku_version() pending = get_pending_upgrades(db_version) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/settings.py b/haiku_rag_slim/haiku/rag/store/repositories/settings.py index d7b2986c..e38c17cf 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/settings.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/settings.py @@ -60,7 +60,9 @@ class SettingsRepository: ) await self.store.settings_table.add([settings_record]) - async def validate_config_compatibility(self) -> None: + async def validate_config_compatibility( + self, stored_settings: dict | None = None + ) -> None: """Validate the current configuration against stored settings without writing. Opening a database never modifies it. ``vector_dim`` mismatches raise — @@ -72,7 +74,8 @@ class SettingsRepository: while a read-only open continues. Stored settings are reconciled explicitly via ``haiku-rag rebuild --set-embedder``, never on open. """ - stored_settings = await self.get_current_settings() + if stored_settings is None: + stored_settings = await self.get_current_settings() # Nothing stored to validate against — never write on open. if not stored_settings: diff --git a/tests/store/test_open_path.py b/tests/store/test_open_path.py new file mode 100644 index 00000000..887db575 --- /dev/null +++ b/tests/store/test_open_path.py @@ -0,0 +1,71 @@ +import lancedb +import pytest + +from haiku.rag.store.engine import Store + + +@pytest.fixture +def counts(monkeypatch): + """Count the connection-level calls an open makes.""" + tally: dict[str, int] = {"list_tables": 0, "open_settings": 0, "settings_query": 0} + + list_tables = lancedb.AsyncConnection.list_tables + open_table = lancedb.AsyncConnection.open_table + query = lancedb.AsyncTable.query + + async def counted_list_tables(self, *args, **kwargs): + tally["list_tables"] += 1 + return await list_tables(self, *args, **kwargs) + + async def counted_open_table(self, name, *args, **kwargs): + if name == "settings": + tally["open_settings"] += 1 + return await open_table(self, name, *args, **kwargs) + + def counted_query(self): + if self.name == "settings": + tally["settings_query"] += 1 + return query(self) + + monkeypatch.setattr(lancedb.AsyncConnection, "list_tables", counted_list_tables) + monkeypatch.setattr(lancedb.AsyncConnection, "open_table", counted_open_table) + monkeypatch.setattr(lancedb.AsyncTable, "query", counted_query) + return tally + + +@pytest.mark.asyncio +async def test_reopening_reads_the_table_list_and_settings_once(temp_db_path, counts): + async with Store(temp_db_path, create=True): + pass + for key in counts: + counts[key] = 0 + + async with Store(temp_db_path): + pass + + assert counts["list_tables"] == 1 + assert counts["open_settings"] == 1 + assert counts["settings_query"] == 1 + + +@pytest.mark.asyncio +async def test_storage_failures_propagate(temp_db_path): + """A read failure must not read as empty settings: the migration check would + then see version 0.0.0 and declare every migration pending.""" + async with Store(temp_db_path, create=True) as store: + + def boom(): + raise RuntimeError("s3 is having a day") + + store.settings_table.query = boom + + with pytest.raises(RuntimeError, match="s3 is having a day"): + await store._read_stored_settings() + + +@pytest.mark.asyncio +async def test_non_dict_settings_read_as_empty(temp_db_path): + async with Store(temp_db_path, create=True) as store: + await store.settings_table.update({"settings": "[]"}, where="id = 'settings'") + + assert await store._read_stored_settings() == {} diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index 33fa349c..ed65a5f9 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -276,7 +276,7 @@ class TestInitFailureCleanup: async def fake_connect(*args, **kwargs): return mock_conn - async def failing_init_tables(self, is_new_db): + async def failing_init_tables(self, *args): raise RuntimeError("simulated table init failure") monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect) @@ -390,7 +390,7 @@ class TestStoreMiscellany: {"settings": "not json at all"}, where="id = 'settings'" ) - assert await store._get_stored_vector_dim() is None + assert await store._read_stored_settings() == {} @pytest.mark.asyncio async def test_vacuum_skips_when_already_running(self, temp_db_path):