da207da106
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.
525 lines
20 KiB
Python
525 lines
20 KiB
Python
from datetime import timedelta
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
|
from haiku.rag.store.engine import ConnectionMode, Store, connect_lancedb
|
|
|
|
|
|
class TestConnectionMode:
|
|
def test_local_when_uri_empty(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri=""))
|
|
assert ConnectionMode.from_config(config) == ConnectionMode.LOCAL
|
|
|
|
def test_cloud_when_db_uri(self):
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="db://my-database", api_key="key", region="us-east-1"
|
|
)
|
|
)
|
|
assert ConnectionMode.from_config(config) == ConnectionMode.CLOUD
|
|
|
|
def test_object_storage_s3(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
|
|
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
|
|
|
def test_object_storage_gs(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri="gs://bucket/path"))
|
|
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
|
|
|
def test_object_storage_az(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri="az://container/path"))
|
|
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
|
|
|
def test_object_storage_hdfs(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri="hdfs://namenode/path"))
|
|
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
|
|
|
def test_unknown_uri_treated_as_object_storage(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri="custom://something"))
|
|
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
|
|
|
|
|
class TestConnectLancedb:
|
|
@pytest.mark.asyncio
|
|
async def test_local_passes_absolute_db_path(self, temp_db_path):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri=""))
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config, db_path=temp_db_path)
|
|
mock_connect.assert_awaited_once()
|
|
assert mock_connect.call_args.args == (temp_db_path.absolute(),)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_resolves_relative_db_path(self, tmp_path, monkeypatch):
|
|
from pathlib import Path
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
relative = Path("db/rag.lancedb")
|
|
config = AppConfig(lancedb=LanceDBConfig(uri=""))
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config, db_path=relative)
|
|
mock_connect.assert_awaited_once()
|
|
assert mock_connect.call_args.args == (relative.absolute(),)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cloud_passes_uri_api_key_region(self):
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="db://my-database", api_key="test-key", region="us-west-2"
|
|
)
|
|
)
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config)
|
|
mock_connect.assert_awaited_once()
|
|
kwargs = mock_connect.call_args.kwargs
|
|
assert kwargs["uri"] == "db://my-database"
|
|
assert kwargs["api_key"] == "test-key"
|
|
assert kwargs["region"] == "us-west-2"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_object_storage_passes_uri_and_storage_options(self):
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path",
|
|
storage_options={
|
|
"endpoint": "http://minio:9000",
|
|
"region": "us-east-1",
|
|
},
|
|
)
|
|
)
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config)
|
|
mock_connect.assert_awaited_once()
|
|
kwargs = mock_connect.call_args.kwargs
|
|
assert kwargs["uri"] == "s3://bucket/path"
|
|
assert kwargs["storage_options"] == {
|
|
"endpoint": "http://minio:9000",
|
|
"region": "us-east-1",
|
|
}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_object_storage_without_storage_options(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config)
|
|
mock_connect.assert_awaited_once()
|
|
kwargs = mock_connect.call_args.kwargs
|
|
assert kwargs["uri"] == "s3://bucket/path"
|
|
assert "storage_options" not in kwargs
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_without_db_path_raises(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri=""))
|
|
with pytest.raises(
|
|
ValueError, match="No lancedb.uri configured and no db_path provided"
|
|
):
|
|
await connect_lancedb(config)
|
|
|
|
|
|
class TestStoreConnectionMode:
|
|
@pytest.mark.asyncio
|
|
async def test_store_connection_mode_local(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
assert store._connection_mode == ConnectionMode.LOCAL
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_connection_mode_cloud(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
with (
|
|
patch.object(Config.lancedb, "uri", "db://test-database"),
|
|
patch.object(Config.lancedb, "api_key", "test-api-key"),
|
|
patch.object(Config.lancedb, "region", "us-east-1"),
|
|
):
|
|
assert store._connection_mode == ConnectionMode.CLOUD
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_connection_mode_object_storage(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
|
|
assert store._connection_mode == ConnectionMode.OBJECT_STORAGE
|
|
|
|
|
|
class TestVacuumByConnectionMode:
|
|
@pytest.mark.asyncio
|
|
async def test_cloud_skips_vacuum(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
with (
|
|
patch.object(Config.lancedb, "uri", "db://test-database"),
|
|
patch.object(Config.lancedb, "api_key", "test-api-key"),
|
|
patch.object(Config.lancedb, "region", "us-east-1"),
|
|
):
|
|
with patch.object(
|
|
store.chunks_table, "optimize", new_callable=AsyncMock
|
|
) as mock_optimize:
|
|
await store.vacuum()
|
|
mock_optimize.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_object_storage_runs_vacuum(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
|
|
with patch.object(
|
|
store.chunks_table, "optimize", new_callable=AsyncMock
|
|
) as mock_optimize:
|
|
await store.vacuum()
|
|
mock_optimize.assert_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_runs_vacuum(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
with patch.object(Config.lancedb, "uri", ""):
|
|
with patch.object(
|
|
store.chunks_table, "optimize", new_callable=AsyncMock
|
|
) as mock_optimize:
|
|
await store.vacuum()
|
|
mock_optimize.assert_called()
|
|
|
|
|
|
class TestVectorIndexByConnectionMode:
|
|
@pytest.mark.asyncio
|
|
async def test_cloud_skips_index_creation(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
with (
|
|
patch.object(Config.lancedb, "uri", "db://test-database"),
|
|
patch.object(Config.lancedb, "api_key", "test-api-key"),
|
|
patch.object(Config.lancedb, "region", "us-east-1"),
|
|
):
|
|
with patch.object(
|
|
store.chunks_table, "count_rows", new_callable=AsyncMock
|
|
) as mock_count:
|
|
await store._ensure_vector_index()
|
|
mock_count.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_object_storage_runs_index_creation(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
|
|
with patch.object(
|
|
store.chunks_table,
|
|
"count_rows",
|
|
new_callable=AsyncMock,
|
|
return_value=0,
|
|
) as mock_count:
|
|
await store._ensure_vector_index()
|
|
mock_count.assert_called()
|
|
|
|
|
|
class TestStoreSkipsPathValidationForRemote:
|
|
@pytest.mark.asyncio
|
|
async def test_skips_path_check_for_cloud(self, tmp_path):
|
|
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="db://test-database", api_key="key", region="us-east-1"
|
|
)
|
|
)
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
):
|
|
with patch.object(Store, "_init_tables", new_callable=AsyncMock):
|
|
async with Store(
|
|
nonexistent,
|
|
config=config,
|
|
create=True,
|
|
skip_validation=True,
|
|
skip_migration_check=True,
|
|
) as store:
|
|
assert store is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skips_path_check_for_object_storage(self, tmp_path):
|
|
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path",
|
|
storage_options={"endpoint": "http://localhost:9000"},
|
|
)
|
|
)
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
):
|
|
with patch.object(Store, "_init_tables", new_callable=AsyncMock):
|
|
async with Store(
|
|
nonexistent,
|
|
config=config,
|
|
create=True,
|
|
skip_validation=True,
|
|
skip_migration_check=True,
|
|
) as store:
|
|
assert store is not None
|
|
|
|
|
|
class TestInitFailureCleanup:
|
|
@pytest.mark.asyncio
|
|
async def test_store_aenter_closes_connection_on_init_failure(
|
|
self, temp_db_path, monkeypatch
|
|
):
|
|
"""If _initialize raises after connect, __aenter__ must close the
|
|
AsyncConnection so it doesn't leak (no __aexit__ runs in that case)."""
|
|
mock_conn = AsyncMock()
|
|
mock_conn.close = lambda: mock_conn.close_calls.append(True) # type: ignore[attr-defined]
|
|
mock_conn.close_calls = [] # type: ignore[attr-defined]
|
|
|
|
async def fake_connect(*args, **kwargs):
|
|
return mock_conn
|
|
|
|
async def failing_init_tables(self, *args):
|
|
raise RuntimeError("simulated table init failure")
|
|
|
|
monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect)
|
|
monkeypatch.setattr(Store, "_init_tables", failing_init_tables)
|
|
|
|
with pytest.raises(RuntimeError, match="simulated table init failure"):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
assert store is not None
|
|
|
|
assert mock_conn.close_calls == [True], (
|
|
"AsyncConnection.close() was not called on init failure"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_aenter_closes_store_on_init_failure(
|
|
self, temp_db_path, monkeypatch
|
|
):
|
|
"""HaikuRAG.__aenter__ must close the store if _initialize fails."""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
close_calls: list[bool] = []
|
|
|
|
original_close = Store.close
|
|
|
|
def tracking_close(self):
|
|
close_calls.append(True)
|
|
original_close(self)
|
|
|
|
async def failing_init(self):
|
|
# Set db so close() has something to close
|
|
self.db = AsyncMock()
|
|
self.db.close = lambda: None
|
|
raise RuntimeError("simulated initialize failure")
|
|
|
|
monkeypatch.setattr(Store, "_initialize", failing_init)
|
|
monkeypatch.setattr(Store, "close", tracking_close)
|
|
|
|
with pytest.raises(RuntimeError, match="simulated initialize failure"):
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
pass
|
|
|
|
assert close_calls, "Store.close() was not called when _initialize raised"
|
|
|
|
|
|
class TestVectorIndexCreation:
|
|
"""_ensure_vector_index needs 256 rows of training data before it builds."""
|
|
|
|
@staticmethod
|
|
async def _seed_chunks(store, count: int) -> None:
|
|
import random
|
|
|
|
records = [
|
|
store.ChunkRecord(
|
|
document_id="doc-1",
|
|
content=f"row {i}",
|
|
content_fts=f"row {i}",
|
|
metadata="{}",
|
|
order=i,
|
|
vector=[random.random() for _ in range(store.embedder.vector_dim)],
|
|
)
|
|
for i in range(count)
|
|
]
|
|
await store.chunks_table.add(records)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_builds_index_once_enough_rows_exist(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
await self._seed_chunks(store, 256)
|
|
|
|
await store._ensure_vector_index()
|
|
|
|
indexes = await store.chunks_table.list_indices()
|
|
assert any("vector" in idx.columns for idx in indexes)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_index_failure_is_warned_not_raised(self, temp_db_path):
|
|
import logging
|
|
|
|
from haiku.rag.store import engine as engine_module
|
|
from tests.conftest import capture_logs
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
await self._seed_chunks(store, 256)
|
|
|
|
async def boom(*_args, **_kwargs):
|
|
raise RuntimeError("index build failed")
|
|
|
|
with patch.object(store.chunks_table, "create_index", boom):
|
|
with capture_logs(engine_module.logger, logging.WARNING) as records:
|
|
await store._ensure_vector_index()
|
|
|
|
assert any("index build failed" in r.getMessage() for r in records)
|
|
indexes = await store.chunks_table.list_indices()
|
|
assert not any("vector" in idx.columns for idx in indexes)
|
|
|
|
|
|
class TestStoreMiscellany:
|
|
@pytest.mark.asyncio
|
|
async def test_create_makes_missing_parent_directories(self, tmp_path):
|
|
nested = tmp_path / "a" / "b" / "db.lancedb"
|
|
|
|
async with Store(nested, create=True) as store:
|
|
assert store._is_new_db is True
|
|
|
|
assert nested.exists()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stored_vector_dim_is_none_for_corrupt_settings(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
await store.settings_table.update(
|
|
{"settings": "not json at all"}, where="id = 'settings'"
|
|
)
|
|
|
|
assert await store._read_stored_settings() == {}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vacuum_skips_when_already_running(self, temp_db_path):
|
|
import asyncio
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
async with store._vacuum_lock:
|
|
# Bounded: a regression here blocks on the held lock, and the
|
|
# timeout turns that deadlock into a clean failure.
|
|
await asyncio.wait_for(store.vacuum(), timeout=5)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_rejects_unknown_table(self, temp_db_path):
|
|
async with Store(temp_db_path, create=True) as store:
|
|
with pytest.raises(ValueError, match="Unknown table"):
|
|
await store.list_table_versions("not_a_table")
|
|
|
|
|
|
class TestSessionAndConsistency:
|
|
@pytest.mark.asyncio
|
|
async def test_session_is_shared_across_connections(self):
|
|
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config)
|
|
await connect_lancedb(config)
|
|
|
|
sessions = [c.kwargs["session"] for c in mock_connect.call_args_list]
|
|
assert sessions[0] is sessions[1]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_sizes_select_distinct_sessions(self):
|
|
small = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path", index_cache_size_bytes=1 << 20
|
|
)
|
|
)
|
|
large = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path", index_cache_size_bytes=1 << 30
|
|
)
|
|
)
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(small)
|
|
await connect_lancedb(large)
|
|
|
|
sessions = [c.kwargs["session"] for c in mock_connect.call_args_list]
|
|
assert sessions[0] is not sessions[1]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_both_cache_sizes_are_applied(self):
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path",
|
|
index_cache_size_bytes=2 << 20,
|
|
metadata_cache_size_bytes=4 << 20,
|
|
)
|
|
)
|
|
with (
|
|
patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
),
|
|
patch("haiku.rag.store.engine.lancedb.Session") as mock_session,
|
|
):
|
|
await connect_lancedb(config)
|
|
|
|
mock_session.assert_called_once_with(
|
|
index_cache_size_bytes=2 << 20, metadata_cache_size_bytes=4 << 20
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_consistency_interval_is_forwarded(self):
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path", read_consistency_interval_seconds=5
|
|
)
|
|
)
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config)
|
|
|
|
assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta(
|
|
seconds=5
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_consistency_interval_omitted_when_disabled(self):
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path", read_consistency_interval_seconds=None
|
|
)
|
|
)
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config)
|
|
|
|
assert mock_connect.call_args.kwargs["read_consistency_interval"] is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_connection_also_gets_session_and_consistency(self, tmp_path):
|
|
config = AppConfig()
|
|
with patch(
|
|
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
await connect_lancedb(config, tmp_path / "db.lancedb")
|
|
|
|
assert mock_connect.call_args.kwargs["session"] is not None
|
|
assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta(
|
|
seconds=30
|
|
)
|
|
|
|
|
|
class TestLanceDBConfigValidation:
|
|
def test_negative_values_are_rejected(self):
|
|
"""Negatives overflow or panic inside Lance, so reject them here."""
|
|
with pytest.raises(ValidationError):
|
|
LanceDBConfig(read_consistency_interval_seconds=-1)
|
|
with pytest.raises(ValidationError):
|
|
LanceDBConfig(index_cache_size_bytes=-1)
|
|
with pytest.raises(ValidationError):
|
|
LanceDBConfig(metadata_cache_size_bytes=-1)
|
|
|
|
def test_zero_is_allowed(self):
|
|
config = LanceDBConfig(
|
|
read_consistency_interval_seconds=0, index_cache_size_bytes=0
|
|
)
|
|
assert config.read_consistency_interval_seconds == 0
|