Give every client shape a way to be released

`close()` refused a client covering a set, which left one with no method to
call: `async with` was the only lifecycle it had. `aclose()` runs that teardown
for a caller that owns the client some other way, whatever it covers, and
nothing to release is not an error, so it is safe before entering and after
closing.

`close()` stays what it is, one connection and nothing else, and says so:
draining the background vacuum and releasing the embedder and reranker are
awaitable.

docs/python.md

@@ -24,6 +24,12 @@ async with HaikuRAG("path/to/database.lancedb", read_only=True) as client:
     # await client.create_document(...)  # Would raise ReadOnlyError
 ```
 
+`async with` is the lifecycle. A caller that owns the client some other way
+releases it with `await client.aclose()`, which does the same work for every
+client shape. `client.close()` closes the connection to one database and nothing
+else, since draining the background vacuum and releasing the embedder and
+reranker are awaitable; it refuses a client covering several.
+
 !!! note
     Databases must be explicitly created with `create=True` or via `haiku-rag init` before use. Operations on non-existent databases will raise `FileNotFoundError`.
 

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -169,6 +169,7 @@ class HaikuRAG:
         self._scope: DatabaseScope | None = None
         self._session: SingleDatabaseSession | FederatedSession | None = None
         self._owns_session = True
+        self._closed = False
 
     @property
     def covers_multiple(self) -> bool:
@@ -321,6 +322,7 @@ class HaikuRAG:
         covering several opens their sessions lazily, so `store` and the
         repositories stay unset until one database is named.
         """
+        self._closed = False
         if not self._owns_session:
             assert self._session is not None
             return self
@@ -452,7 +454,14 @@ class HaikuRAG:
                 )
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):  # noqa: ARG002
-        """Async context manager exit."""
+        """Async context manager exit.
+
+        Nothing to release is not an error: exiting before entering and exiting
+        twice both do nothing. The session stays readable afterwards, so what a
+        client covered can still be asked.
+        """
+        if self._session is None or self._closed:
+            return False
         # Branch on what this client covers, not on what it happened to open:
         # a federating client that answered no query has nothing open and no
         # store either.
@@ -468,15 +477,16 @@ class HaikuRAG:
             # only place they are closed — and only if anything built them.
             await self._aclose_cached("embedder")
             await self._aclose_cached("_own_reranker")
+            self._closed = True
             return False
         if not self._owns_session:
             await self._release_own()
             return False
-        assert self._session is not None
         # The session drains, releases its store's embedder and closes; the
         # cached reference here is only discarded, never closed twice.
         await self._release_own()
         await self._session.aclose()
+        self._closed = True
         return False
 
     async def _release_own(self) -> None:
@@ -944,13 +954,32 @@ class HaikuRAG:
         """Optimize and clean up old versions across all tables."""
         await self._single_session("vacuum").store.vacuum()
 
-    def close(self):
-        """Close the underlying store connection.
+    async def aclose(self) -> None:
+        """Release everything this client opened, whatever it covers.
+
+        The teardown `async with` runs, for a caller that owns the client's
+        lifetime some other way. Nothing to release is not an error, so this is
+        safe before entering and after closing.
+        """
+        await self.__aexit__(None, None, None)
+
+    def close(self) -> None:
+        """Close the connection to the one database this client opened.
+
+        The connection and nothing else: draining the background vacuum and
+        releasing the embedder and reranker are awaitable, so `aclose` is what
+        does all of it, and `async with` is the usual way to ask for it.
 
         A client covering one of a set borrows that database and never closes
-        it: the set opened it and the set closes it.
+        it: the set opened it and the set closes it. A client covering a set has
+        no single connection to close and refuses.
         """
-        session = self._single_session("close")
+        if not isinstance(self._session, SingleDatabaseSession):
+            raise AmbiguousDatabaseError(
+                "close works on one connection, and this client covers "
+                f"{', '.join(sorted(self.source_names))}; await aclose() to "
+                "release every database it opened"
+            )
         if not self._owns_session:
             return
-        session.close()
+        self._session.close()

tests/multi_db/test_lifecycle.py

@@ -245,6 +245,83 @@ class TestBorrowedDatabases:
         assert closed == ["reranker"]
 
 
+class TestReleasingAClient:
+    """`async with` is the usual lifecycle, and `aclose` is it for a caller that
+    owns the client some other way. `close` is a connection, not a lifecycle."""
+
+    @pytest.mark.asyncio
+    async def test_aclose_releases_a_set(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about cats"])
+
+        rag = HaikuRAG(config=config)
+        await rag.__aenter__()
+        alpha, beta = await rag.clients_for(["alpha", "beta"])
+        assert alpha.store.db.is_open()
+
+        await rag.aclose()
+
+        assert not alpha.store.db.is_open()
+        assert not beta.store.db.is_open()
+
+    @pytest.mark.asyncio
+    async def test_aclose_releases_one_database(self, tmp_path):
+        config = _config(tmp_path, ["alpha"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+
+        rag = HaikuRAG(config=config)
+        await rag.__aenter__()
+        assert rag.store.db.is_open()
+
+        await rag.aclose()
+
+        assert not rag.store.db.is_open()
+
+    @pytest.mark.asyncio
+    async def test_aclose_before_entering_does_nothing(self, tmp_path):
+        """Nothing was opened, so there is nothing to release and no error."""
+        config = _config(tmp_path, ["alpha"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+
+        await HaikuRAG(config=config).aclose()
+
+    @pytest.mark.asyncio
+    async def test_aclose_twice_releases_once(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+
+        closed: list[str] = []
+
+        rag = HaikuRAG(config=config)
+        await rag.__aenter__()
+        (alpha,) = await rag.clients_for(["alpha"])
+        session = rag._session
+        assert isinstance(session, FederatedSession)
+        real = session.aclose
+
+        async def counting():
+            closed.append("set")
+            await real()
+
+        session.aclose = counting  # ty: ignore[invalid-assignment]
+
+        await rag.aclose()
+        await rag.aclose()
+
+        assert closed == ["set"]
+        assert not alpha.store.db.is_open()
+
+    @pytest.mark.asyncio
+    async def test_close_refuses_a_set_and_names_aclose(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+
+        async with HaikuRAG(config=config) as rag:
+            with pytest.raises(AmbiguousDatabaseError, match="aclose"):
+                rag.close()
+
+
 class TestSharingTheReranker:
     @pytest.mark.asyncio
     async def test_the_set_builds_and_closes_one_reranker(self, tmp_path, monkeypatch):