diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index a0213a6f..e9d968d5 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -45,9 +45,15 @@ class HaikuRAGApp: # pragma: no cover self.before = before self.console = Console() + @property + def _is_local(self) -> bool: + from haiku.rag.store.engine import ConnectionMode + + return ConnectionMode.from_config(self.config) == ConnectionMode.LOCAL + async def init(self): """Initialize a new database.""" - if self.db_path.exists(): + if self._is_local and self.db_path.exists(): self.console.print( f"[yellow]Database already exists at {self.db_path}[/yellow]" ) @@ -63,7 +69,7 @@ class HaikuRAGApp: # pragma: no cover async def info(self): """Display read-only information about the database without modifying it.""" - import lancedb + from haiku.rag.store.engine import Store, connect_lancedb # Basic: show path self.console.print("[bold]haiku.rag database info[/bold]") @@ -71,17 +77,15 @@ class HaikuRAGApp: # pragma: no cover f" [repr.attrib_name]path[/repr.attrib_name]: {self.db_path}" ) - if not self.db_path.exists(): + if self._is_local and not self.db_path.exists(): self.console.print("[red]Database path does not exist.[/red]") return # Connect without going through Store to avoid upgrades/validation writes - db = lancedb.connect(self.db_path) + db = connect_lancedb(self.config, self.db_path) versions = get_package_versions() - from haiku.rag.store.engine import Store - store = Store( self.db_path, config=self.config, @@ -201,7 +205,7 @@ class HaikuRAGApp: # pragma: no cover """ from haiku.rag.store.engine import Store - if not self.db_path.exists(): + if self._is_local and not self.db_path.exists(): self.console.print("[red]Database path does not exist.[/red]") return diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py index 4621c0e3..23dc4f8f 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py @@ -64,21 +64,25 @@ class InfoModal(ModalScreen): async def on_mount(self) -> None: """Load and display database info.""" - import lancedb + from haiku.rag.store.engine import ConnectionMode, connect_lancedb lines: list[str] = [] # Path lines.append(f"[bold $accent]path[/bold $accent]: {self.db_path}") - if not self.db_path.exists(): + is_local = ( + ConnectionMode.from_config(self.client.store._config) + == ConnectionMode.LOCAL + ) + if is_local and not self.db_path.exists(): lines.append("[red]Database path does not exist.[/red]") self._content_widget.update("\n".join(lines)) return # Connect to get table info try: - db = lancedb.connect(self.db_path) + db = connect_lancedb(self.client.store._config, self.db_path) table_names = set(db.list_tables().tables) except Exception as e: lines.append(f"[red]Failed to open database: {e}[/red]") diff --git a/tests/test_info.py b/tests/test_info.py index cd24bac4..ffbdf2ca 100644 --- a/tests/test_info.py +++ b/tests/test_info.py @@ -1,8 +1,10 @@ import json +from unittest.mock import patch import pytest from haiku.rag.app import HaikuRAGApp +from haiku.rag.config.models import AppConfig, LanceDBConfig @pytest.mark.asyncio @@ -153,3 +155,93 @@ async def test_app_info_with_vector_index(temp_db_path, capsys): # Check basic info still present assert "documents: 1" in out assert "chunks: 512" in out + + +@pytest.mark.asyncio +async def test_app_info_uses_connect_lancedb_for_remote(tmp_path, capsys): + """info() should use connect_lancedb() instead of direct lancedb.connect() for remote URIs.""" + nonexistent = tmp_path / "does_not_exist" / "db.lancedb" + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", + storage_options={"endpoint": "http://localhost:9000"}, + ) + ) + app = HaikuRAGApp(db_path=nonexistent, config=config) + + with patch("haiku.rag.store.engine.connect_lancedb") as mock_connect: + # Make the mock return something that lets info() proceed minimally + mock_db = mock_connect.return_value + mock_table = mock_db.open_table.return_value + mock_table.search.return_value.where.return_value.limit.return_value.to_arrow.return_value.to_pylist.return_value = [ + { + "settings": json.dumps( + { + "version": "1.0.0", + "embeddings": { + "model": { + "provider": "test", + "name": "test", + "vector_dim": 3, + } + }, + } + ) + } + ] + + with patch("haiku.rag.store.engine.Store") as mock_store_cls: + mock_store = mock_store_cls.return_value + mock_store.get_stats.return_value = { + "documents": {"exists": True, "num_rows": 0, "total_bytes": 0}, + "chunks": { + "exists": True, + "num_rows": 0, + "total_bytes": 0, + "has_vector_index": False, + "num_indexed_rows": 0, + "num_unindexed_rows": 0, + }, + } + await app.info() + + mock_connect.assert_called_once_with(config, nonexistent) + + +@pytest.mark.asyncio +async def test_app_init_skips_exists_check_for_remote(tmp_path): + """init() should not check db_path.exists() for remote URIs.""" + nonexistent = tmp_path / "does_not_exist" / "db.lancedb" + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", + storage_options={"endpoint": "http://localhost:9000"}, + ) + ) + app = HaikuRAGApp(db_path=nonexistent, config=config) + + with patch("haiku.rag.app.HaikuRAG") as mock_client_cls: + await app.init() + # Should have called HaikuRAG to create, not returned early + mock_client_cls.assert_called_once() + + +@pytest.mark.asyncio +async def test_app_history_skips_exists_check_for_remote(tmp_path): + """history() should not check db_path.exists() for remote URIs.""" + nonexistent = tmp_path / "does_not_exist" / "db.lancedb" + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", + storage_options={"endpoint": "http://localhost:9000"}, + ) + ) + app = HaikuRAGApp(db_path=nonexistent, config=config) + + with patch("haiku.rag.store.engine.Store") as mock_store_cls: + mock_store = mock_store_cls.return_value + mock_store.documents_table.list_versions.return_value = [] + mock_store.chunks_table.list_versions.return_value = [] + mock_store.settings_table.list_versions.return_value = [] + await app.history() + mock_store_cls.assert_called_once()