diff --git a/haiku_rag_slim/haiku/rag/client/scope.py b/haiku_rag_slim/haiku/rag/client/scope.py new file mode 100644 index 00000000..f8453e23 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/client/scope.py @@ -0,0 +1,138 @@ +from dataclasses import dataclass +from pathlib import Path + +from haiku.rag.config import AppConfig +from haiku.rag.store.exceptions import AmbiguousDatabaseError +from haiku.rag.utils import locate_database + + +@dataclass(frozen=True) +class DatabaseRef: + """One database, and the name it answers to. + + ``name`` is the key from ``lancedb.databases``, and the only identity that + leaves the configuration: it travels in results, citations and the errors an + operator or a model sees, where a location must not. The invariant below is + the exception, and deliberately so — a malformed ref is a programming error + raised in the caller's own process, where naming what it was given is what + makes it fixable. None where nothing names the database — a path given + directly, or the legacy single ``uri``. + + Location is resolved once, on construction, into exactly one of ``uri`` and + ``db_path``. Keeping the configured string and re-reading it later would let a + path the caller gave be reinterpreted as a URI because it happens to carry a + scheme. + """ + + name: str | None + uri: str + db_path: Path | None + + def __post_init__(self) -> None: + if bool(self.uri) == (self.db_path is not None): + raise ValueError( + "a database is either a URI or a local path: " + f"got uri={self.uri!r} and db_path={self.db_path!r}" + ) + + @classmethod + def at(cls, path: Path | str, *, name: str | None = None) -> "DatabaseRef": + """A database at a path the caller named, taken as given.""" + return cls(name=name, uri="", db_path=Path(path)) + + @classmethod + def configured(cls, name: str | None, location: str) -> "DatabaseRef": + """A database from ``lancedb.databases``, where the configured value is a + URI or a local path depending on whether it carries a scheme.""" + uri, db_path = locate_database(location) + return cls(name=name, uri=uri, db_path=db_path) + + def connection(self, config: AppConfig) -> tuple[AppConfig, Path | None]: + """The configuration and path to open this one database with. + + A copy. The scope is resolved once from the caller's configuration, and + rewriting that configuration in place is what left downstream code unable + to tell that a set had been named. + """ + one = config.model_copy(deep=True) + one.lancedb.databases = {} + one.lancedb.uri = self.uri + return one, self.db_path + + +@dataclass(frozen=True) +class DatabaseScope: + """The databases an operation covers. + + Resolved once, from configuration plus at most one selector, and passed down + rather than re-derived. Never empty: every resolution reaches a database, and + the sessions built from a scope have no meaning without one. + + Nothing here reads the environment. ``HAIKU_RAG_DB`` is honoured by the + capability entry point alone, which passes it as ``database_path``, so + resolving a scope cannot quietly change what any other caller opens. + """ + + databases: tuple[DatabaseRef, ...] + + def __post_init__(self) -> None: + if not self.databases: + raise ValueError("a scope covers at least one database") + + @classmethod + def resolve( + cls, + config: AppConfig, + *, + database_name: str | None = None, + database_path: Path | str | None = None, + ) -> "DatabaseScope": + """The databases named by `config` and at most one selector. + + A path names one database that nothing calls anything; a name selects one + of the configured set and keeps its name. With no selector the configured + set is covered in configuration order, a set of one included. + """ + if database_name is not None and database_path is not None: + raise AmbiguousDatabaseError( + "a database name and a database path both name one database; " + "pass one of them" + ) + + declared = config.lancedb.databases + + if database_path is not None: + return cls((DatabaseRef.at(database_path),)) + + if database_name is not None: + if database_name not in declared: + raise AmbiguousDatabaseError( + f"unknown database {database_name!r}; lancedb.databases names " + f"{', '.join(sorted(declared)) or 'nothing'}" + ) + return cls( + (DatabaseRef.configured(database_name, declared[database_name]),) + ) + + if declared: + return cls( + tuple( + DatabaseRef.configured(name, location) + for name, location in declared.items() + ) + ) + + if config.lancedb.uri: + return cls((DatabaseRef(name=None, uri=config.lancedb.uri, db_path=None),)) + + return cls((DatabaseRef.at(config.storage.data_dir / "haiku.rag.lancedb"),)) + + @property + def covers_multiple(self) -> bool: + """Whether this scope covers more than one database.""" + return len(self.databases) > 1 + + @property + def names(self) -> tuple[str, ...]: + """The configured names covered, in order. Empty where none is named.""" + return tuple(ref.name for ref in self.databases if ref.name is not None) diff --git a/tests/test_database_scope.py b/tests/test_database_scope.py new file mode 100644 index 00000000..bd845f6e --- /dev/null +++ b/tests/test_database_scope.py @@ -0,0 +1,178 @@ +from pathlib import Path + +import pytest + +from haiku.rag.client.scope import DatabaseRef, DatabaseScope +from haiku.rag.config.models import AppConfig, LanceDBConfig, StorageConfig +from haiku.rag.store.exceptions import AmbiguousDatabaseError + + +def _config(**kwargs) -> AppConfig: + return AppConfig(lancedb=LanceDBConfig(**kwargs)) + + +class TestResolution: + """One selector at most, and the same answer wherever it is asked.""" + + def test_a_name_and_a_path_together_are_refused(self): + config = _config(databases={"alpha": "/data/alpha.lancedb"}) + + with pytest.raises(AmbiguousDatabaseError, match="pass one of them"): + DatabaseScope.resolve( + config, database_name="alpha", database_path=Path("/data/other.lancedb") + ) + + def test_a_path_names_one_unnamed_database(self): + """A path says which database, not what it is called, even where the + configuration would have named it.""" + config = _config(databases={"alpha": "/data/alpha.lancedb"}) + + scope = DatabaseScope.resolve(config, database_path=Path("/data/other.lancedb")) + + assert scope.databases == (DatabaseRef.at("/data/other.lancedb"),) + assert scope.names == () + assert not scope.covers_multiple + + def test_a_named_database_keeps_its_name(self): + config = _config(databases={"alpha": "/data/alpha.lancedb", "beta": "b://b"}) + + scope = DatabaseScope.resolve(config, database_name="beta") + + assert scope.databases == (DatabaseRef("beta", "b://b", None),) + assert scope.names == ("beta",) + + def test_an_unknown_name_is_refused(self): + config = _config(databases={"alpha": "/data/alpha.lancedb"}) + + with pytest.raises(AmbiguousDatabaseError, match="unknown database 'nope'"): + DatabaseScope.resolve(config, database_name="nope") + + def test_no_selector_covers_the_configured_set_in_order(self): + config = _config( + databases={"beta": "/data/b.lancedb", "alpha": "/data/a.lancedb"} + ) + + scope = DatabaseScope.resolve(config) + + assert scope.names == ("beta", "alpha") + assert scope.covers_multiple + + def test_a_configured_set_of_one_is_still_a_named_database(self): + """Its name is what results and citations carry, so it survives.""" + config = _config(databases={"alpha": "/data/alpha.lancedb"}) + + scope = DatabaseScope.resolve(config) + + assert scope.databases == ( + DatabaseRef.configured("alpha", "/data/alpha.lancedb"), + ) + assert not scope.covers_multiple + + def test_a_bare_uri_is_one_unnamed_database(self): + scope = DatabaseScope.resolve(_config(uri="s3://bucket/one.lancedb")) + + assert scope.databases == (DatabaseRef(None, "s3://bucket/one.lancedb", None),) + + def test_nothing_configured_falls_back_to_the_data_directory(self, tmp_path): + config = AppConfig(storage=StorageConfig(data_dir=tmp_path)) + + scope = DatabaseScope.resolve(config) + + assert scope.databases == (DatabaseRef.at(tmp_path / "haiku.rag.lancedb"),) + + def test_the_environment_is_not_consulted(self, monkeypatch, tmp_path): + """HAIKU_RAG_DB is honoured by the capability entry point alone. Reading it + here would change what every other caller opens.""" + monkeypatch.setenv("HAIKU_RAG_DB", "/data/from-the-environment.lancedb") + config = _config(databases={"alpha": "/data/alpha.lancedb"}) + + scope = DatabaseScope.resolve(config) + + assert scope.names == ("alpha",) + + def test_a_path_is_never_reinterpreted_as_a_uri(self): + """A caller naming a path means that path. Sending it back through the + configured-location rules would let a scheme turn it into a URI.""" + scope = DatabaseScope.resolve( + _config(), database_path="s3://bucket/looks-like-a-uri.lancedb" + ) + + [ref] = scope.databases + assert ref.db_path == Path("s3://bucket/looks-like-a-uri.lancedb") + assert ref.uri == "" + + def test_a_configured_location_with_a_scheme_is_a_uri(self): + """A configured value is a URI or a path depending on its scheme, which is + what makes it different from a path the caller gave.""" + config = _config(databases={"alpha": "s3://bucket/alpha.lancedb"}) + + [ref] = DatabaseScope.resolve(config).databases + + assert ref.uri == "s3://bucket/alpha.lancedb" + assert ref.db_path is None + + def test_a_database_is_a_uri_or_a_path(self): + """Both would silently ignore the path; neither fails later, when the + connection is derived and there is nothing to open. + + The message names what it was given: this is a programming error raised + in the caller's own process, not one an operator or a model ever sees. + """ + with pytest.raises(ValueError, match="either a URI or a local path") as both: + DatabaseRef(None, "s3://bucket/a.lancedb", Path("/data/a.lancedb")) + assert "s3://bucket/a.lancedb" in str(both.value) + + with pytest.raises(ValueError, match="either a URI or a local path") as neither: + DatabaseRef(None, "", None) + assert "db_path=None" in str(neither.value) + + def test_a_scope_covers_at_least_one_database(self): + """Every resolution reaches a database, and the sessions built from a + scope have no meaning without one.""" + with pytest.raises(ValueError, match="at least one database"): + DatabaseScope(()) + + +class TestConnectionDerivation: + """Opening one of a set must not disturb the configuration it came from.""" + + def test_a_local_location_becomes_a_path(self): + config = _config(databases={"alpha": "/data/alpha.lancedb"}) + [ref] = DatabaseScope.resolve(config).databases + + one, db_path = ref.connection(config) + + assert db_path == Path("/data/alpha.lancedb") + assert one.lancedb.uri == "" + assert one.lancedb.databases == {} + + def test_a_uri_location_stays_a_uri(self): + config = _config(databases={"alpha": "s3://bucket/alpha.lancedb"}) + [ref] = DatabaseScope.resolve(config).databases + + one, db_path = ref.connection(config) + + assert db_path is None + assert one.lancedb.uri == "s3://bucket/alpha.lancedb" + + def test_the_original_configuration_is_untouched(self): + """Rewriting it in place is what left downstream code unable to tell a set + had been named.""" + config = _config(databases={"alpha": "/a.lancedb", "beta": "/b.lancedb"}) + + for ref in DatabaseScope.resolve(config).databases: + ref.connection(config) + + assert config.lancedb.databases == {"alpha": "/a.lancedb", "beta": "/b.lancedb"} + assert config.lancedb.uri == "" + + def test_each_derived_configuration_is_its_own_copy(self): + config = _config(databases={"alpha": "/a.lancedb", "beta": "s3://b/b.lancedb"}) + alpha, beta = DatabaseScope.resolve(config).databases + + one, _ = alpha.connection(config) + other, _ = beta.connection(config) + + assert one is not other + assert one.lancedb.uri == "" + assert other.lancedb.uri == "s3://b/b.lancedb"