danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Read the table list and settings once per open

Browse source 
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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-18 13:25:19 +03:00
parent 6f976ef2a9
commit da207da106
No known key found for this signature in database
4 changed files with 125 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

99
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -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)

7
haiku_rag_slim/haiku/rag/store/repositories/settings.py
View file

@ -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:

71
tests/store/test_open_path.py Normal file
View file

@ -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() == {}

4
tests/test_lancedb_connection.py
View file

@ -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):