guard app/UI filesystem checks for remote storage

This commit is contained in:
Yiorgis Gozadinos 2026-04-03 12:02:20 +03:00
parent 110accb8e7
commit 343bfd7199
No known key found for this signature in database
3 changed files with 110 additions and 10 deletions

View file

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

View file

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

View file

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