6f976ef2a9
Every `Store` built its own connection with its own caches and discarded them on close, so the index a vector query loads was refetched by the next connection. On object storage that first fetch dominates: measured on a ~500k-chunk 2560-dim corpus over a ~200ms link, the first query cost ~41s and the second ~3s, and a new connection reusing the session cost ~7s instead of ~47s. `connect_lancedb` now passes a process-wide session, keyed on the configured cache sizes so a caller asking for different sizes gets its own. Also sets `read_consistency_interval`, defaulting to 30s. It was None, meaning a connection never re-checked for other processes' writes. Per-call connections hid that; a shared session makes connections long-lived enough for a reader to go stale against the ingester. All three settings reject negatives at the config boundary. A negative cache size raises OverflowError and a negative interval panics inside Lance, so neither is catchable further in. Zero stays valid for both: no cache, and check on every read. The routing tests now assert the kwargs they care about rather than the full call signature, since every connection carries the two new kwargs.
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, is_new_db):
|
|
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._get_stored_vector_dim() is None
|
|
|
|
@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
|