From 110accb8e75f97c79a093c084a3454c7b3f96cae Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 3 Apr 2026 11:43:51 +0300 Subject: [PATCH] Add ConnectionMode enum (LOCAL/CLOUD/OBJECT_STORAGE) and connect_lancedb() utility to support S3, GCS, Azure, and HDFS backends via storage_options. --- haiku_rag_slim/haiku/rag/config/models.py | 1 + haiku_rag_slim/haiku/rag/store/engine.py | 70 ++++--- tests/test_lancedb_connection.py | 231 ++++++++++++++++++---- 3 files changed, 243 insertions(+), 59 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index de1d6cad..d4727f8f 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -60,6 +60,7 @@ class LanceDBConfig(BaseModel): uri: str = "" api_key: str = "" region: str = "" + storage_options: dict[str, str] = Field(default_factory=dict) class EmbeddingsConfig(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index f97e35de..5ca06037 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -2,6 +2,7 @@ import asyncio import json import logging from datetime import datetime, timedelta +from enum import Enum from importlib import metadata from pathlib import Path from typing import Any @@ -18,6 +19,42 @@ from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError logger = logging.getLogger(__name__) +OBJECT_STORAGE_PREFIXES = ("s3://", "gs://", "az://", "hdfs://") + + +class ConnectionMode(Enum): + LOCAL = "local" + CLOUD = "cloud" + OBJECT_STORAGE = "object_storage" + + @staticmethod + def from_config(config: AppConfig) -> "ConnectionMode": + uri = config.lancedb.uri + if not uri: + return ConnectionMode.LOCAL + if uri.startswith("db://"): + return ConnectionMode.CLOUD + return ConnectionMode.OBJECT_STORAGE + + +def connect_lancedb(config: AppConfig, db_path: Path | None = None): + mode = ConnectionMode.from_config(config) + if mode == ConnectionMode.CLOUD: + return lancedb.connect( + uri=config.lancedb.uri, + api_key=config.lancedb.api_key, + region=config.lancedb.region, + ) + elif mode == ConnectionMode.OBJECT_STORAGE: + kwargs: dict[str, Any] = {"uri": config.lancedb.uri} + if config.lancedb.storage_options: + kwargs["storage_options"] = config.lancedb.storage_options + return lancedb.connect(**kwargs) + else: + if db_path is None: + raise ValueError("No lancedb.uri configured and no db_path provided") + return lancedb.connect(db_path) + class DocumentRecord(LanceModel): id: str = Field(default_factory=lambda: str(uuid4())) @@ -97,7 +134,7 @@ class Store: # Check if database exists (for local filesystem only) is_new_db = False - if not self._has_cloud_config(): + if self._connection_mode == ConnectionMode.LOCAL: if not db_path.exists(): if not create: raise FileNotFoundError( @@ -110,7 +147,7 @@ class Store: Path.mkdir(db_path.parent, parents=True) # Connect to LanceDB - self.db = self._connect_to_lancedb(db_path) + self.db = connect_lancedb(self._config, db_path) # For existing databases, read stored vector dimension to create ChunkRecord # that can read existing chunks. For new databases, use config's dimension. @@ -198,9 +235,7 @@ class Store: """ self._assert_writable() - if self._has_cloud_config() and str(self._config.lancedb.uri).startswith( - "db://" - ): + if self._connection_mode == ConnectionMode.CLOUD: return # Skip if already running (non-blocking) @@ -224,26 +259,9 @@ class Store: # Handle resource errors gracefully logger.debug(f"Vacuum skipped due to resource constraints: {e}") - def _connect_to_lancedb(self, db_path: Path): - """Establish connection to LanceDB (local, cloud, or object storage).""" - # Check if we have cloud configuration - if self._has_cloud_config(): - return lancedb.connect( - uri=self._config.lancedb.uri, - api_key=self._config.lancedb.api_key, - region=self._config.lancedb.region, - ) - else: - # Local file system connection - return lancedb.connect(db_path) - - def _has_cloud_config(self) -> bool: - """Check if cloud configuration is complete.""" - return bool( - self._config.lancedb.uri - and self._config.lancedb.api_key - and self._config.lancedb.region - ) + @property + def _connection_mode(self) -> ConnectionMode: + return ConnectionMode.from_config(self._config) def get_stats(self) -> dict: """Get comprehensive table statistics. @@ -298,7 +316,7 @@ class Store: it will be replaced (using replace=True parameter). Note: Index creation requires sufficient training data. """ - if self._has_cloud_config(): + if self._connection_mode == ConnectionMode.CLOUD: return try: diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index e9b21ad5..a181b1f1 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -3,46 +3,211 @@ from unittest.mock import patch import pytest from haiku.rag.config import Config -from haiku.rag.store.engine import Store +from haiku.rag.config.models import AppConfig, LanceDBConfig +from haiku.rag.store.engine import ConnectionMode, Store, connect_lancedb -@pytest.mark.asyncio -async def test_lancedb_cloud_skips_optimization(temp_db_path): - """Test that vacuum is skipped when using LanceDB Cloud (db:// URI).""" - # Create a store - store = Store(temp_db_path, create=True) +class TestConnectionMode: + def test_local_when_uri_empty(self): + config = AppConfig(lancedb=LanceDBConfig(uri="")) + assert ConnectionMode.from_config(config) == ConnectionMode.LOCAL - # Mock all cloud config to simulate LanceDB Cloud usage - 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"), - ): - # Mock the optimize method to track if it's called - with patch.object(store.chunks_table, "optimize") as mock_optimize: - # Call vacuum - this should skip optimization for LanceDB Cloud - await store.vacuum() + 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 - # The optimize method should NOT have been called for LanceDB Cloud - mock_optimize.assert_not_called() + def test_object_storage_s3(self): + config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path")) + assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE - store.close() + 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 -@pytest.mark.asyncio -async def test_local_storage_calls_optimization(temp_db_path): - """Test that vacuum calls optimization for local storage.""" - # Create a store - store = Store(temp_db_path, create=True) +class TestConnectLancedb: + def test_local_passes_db_path(self, temp_db_path): + config = AppConfig(lancedb=LanceDBConfig(uri="")) + with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect: + connect_lancedb(config, db_path=temp_db_path) + mock_connect.assert_called_once_with(temp_db_path) - # Ensure uri is empty (local storage) - with patch.object(Config.lancedb, "uri", ""): - # Mock the optimize method to track if it's called - with patch.object(store.chunks_table, "optimize") as mock_optimize: - # Call vacuum - this should optimize all tables for local storage - await store.vacuum() + 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") as mock_connect: + connect_lancedb(config) + mock_connect.assert_called_once_with( + uri="db://my-database", api_key="test-key", region="us-west-2" + ) - # The optimize method SHOULD have been called for local storage - mock_optimize.assert_called() + 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") as mock_connect: + connect_lancedb(config) + mock_connect.assert_called_once_with( + uri="s3://bucket/path", + storage_options={ + "endpoint": "http://minio:9000", + "region": "us-east-1", + }, + ) - store.close() + def test_object_storage_without_storage_options(self): + config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path")) + with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect: + connect_lancedb(config) + mock_connect.assert_called_once_with(uri="s3://bucket/path") + + 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" + ): + connect_lancedb(config) + + +class TestStoreConnectionMode: + def test_store_connection_mode_local(self, temp_db_path): + store = Store(temp_db_path, create=True) + assert store._connection_mode == ConnectionMode.LOCAL + store.close() + + def test_store_connection_mode_cloud(self, temp_db_path): + store = Store(temp_db_path, create=True) + 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 + store.close() + + def test_store_connection_mode_object_storage(self, temp_db_path): + store = Store(temp_db_path, create=True) + with patch.object(Config.lancedb, "uri", "s3://bucket/path"): + assert store._connection_mode == ConnectionMode.OBJECT_STORAGE + store.close() + + +class TestVacuumByConnectionMode: + @pytest.mark.asyncio + async def test_cloud_skips_vacuum(self, temp_db_path): + store = Store(temp_db_path, create=True) + 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") as mock_optimize: + await store.vacuum() + mock_optimize.assert_not_called() + store.close() + + @pytest.mark.asyncio + async def test_object_storage_runs_vacuum(self, temp_db_path): + store = Store(temp_db_path, create=True) + with patch.object(Config.lancedb, "uri", "s3://bucket/path"): + with patch.object(store.chunks_table, "optimize") as mock_optimize: + await store.vacuum() + mock_optimize.assert_called() + store.close() + + @pytest.mark.asyncio + async def test_local_runs_vacuum(self, temp_db_path): + store = Store(temp_db_path, create=True) + with patch.object(Config.lancedb, "uri", ""): + with patch.object(store.chunks_table, "optimize") as mock_optimize: + await store.vacuum() + mock_optimize.assert_called() + store.close() + + +class TestVectorIndexByConnectionMode: + def test_cloud_skips_index_creation(self, temp_db_path): + store = Store(temp_db_path, create=True) + 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") as mock_count: + store._ensure_vector_index() + mock_count.assert_not_called() + store.close() + + def test_object_storage_runs_index_creation(self, temp_db_path): + store = Store(temp_db_path, create=True) + with patch.object(Config.lancedb, "uri", "s3://bucket/path"): + with patch.object( + store.chunks_table, "count_rows", return_value=0 + ) as mock_count: + store._ensure_vector_index() + mock_count.assert_called() + store.close() + + +class TestStoreSkipsPathValidationForRemote: + 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"): + with patch.object(Store, "_init_tables"): + store = Store( + nonexistent, + config=config, + create=True, + skip_validation=True, + skip_migration_check=True, + ) + store.close() + + 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"): + with patch.object(Store, "_init_tables"): + store = Store( + nonexistent, + config=config, + create=True, + skip_validation=True, + skip_migration_check=True, + ) + store.close()