Make DrivineAssembly pi-pilgrim 0.4's AsyncPilgrimAssembly, so await close() closes the pools again

docs/09-with-pi-pilgrim.md

@@ -49,12 +49,35 @@ app = FastAPI(lifespan=lifespan)
   `await AssemblyHolder.aclose()`.
 - **Entering your assembly doesn't start Drivine**, as entering `DrivineAssembly` itself does: call
   `await assembly.drivine.start()`, as above.
-- **The sync assembly**, `drivine.sync.DrivineAssembly`, closes with `with` or `close()`. The async one closes only
-  with `async with` or `await aclose()`, as its drivers close asynchronously: `close()` refuses, and says so.
+- **The async assembly is an `AsyncPilgrimAssembly`**, as its drivers close asynchronously: `await close()`,
+  `asyncio.run(close())`, `await aclose()` and `async with` all close it, and `with` refuses, and says to use
+  `async with`. Await `close()`: left unawaited, it closes nothing, and Python warns that a coroutine was never
+  awaited. Your assembly closes it whether it holds it with `AssemblyHolder.shared(DrivineAssembly)` or as a component,
+  `self.shared(DrivineAssembly)`. The sync assembly, `drivine.sync.DrivineAssembly`, is a plain `PilgrimAssembly`,
+  which closes with `with` or `close()`.
 - It registers the database the environment describes. To register yours in code, subclass it and override
   `register_databases(registry)`.
 - `Assembled` injects a manager into a class something else makes:
   `pm: PersistenceManager = Assembled(assembly=DrivineAssembly)`.
 
 The managers are plain classes, so an application with another container, or none, makes them itself, and closes the
 registry with `async with DatabaseRegistry() as registry:`.
+
+### Upgrading from 0.2.0
+
+0.3.0 moved `DrivineAssembly` onto pi-pilgrim 0.3's lifecycle: the assembly closes when you leave `async with`, or when
+the application assembly that holds it closes, and closing it closes every database's driver and pool.
+
+- **Shutdown.** 0.2.0 said to call `await assembly.drivine.close()` when the application stops. That closes everything
+  again in every release after 0.3.0, where `DrivineAssembly` is pi-pilgrim's `AsyncPilgrimAssembly`, whose `close()`
+  is `aclose()` (0.3.0 raised `PilgrimError`, or `TypeError` when Drivine had made no registry). Await it: called
+  without `await`, it closes nothing, and Python warns that a coroutine was never awaited. The 0.3 way is to leave your
+  assembly's `async with`, as in the lifespan above, or to call `await assembly.aclose()`, which closes your components
+  and then Drivine. To close only Drivine, `await assembly.drivine.aclose()`.
+- **`with DrivineAssembly()`** raises `AsyncTeardownRequired`, a `PilgrimError`: use `async with`, or
+  `drivine.sync.DrivineAssembly` for the sync API, which closes with `with` or `close()`.
+- **Sync `close()` on your assembly** closes Drivine while it has made no registry, whether your assembly holds it with
+  `AssemblyHolder.shared` or with `self.shared(DrivineAssembly)`. Once Drivine has made its registry it refuses with
+  `AsyncTeardownRequired`, before closing anything, since the drivers close asynchronously: close with
+  `await aclose()`. It refuses from the start if your subclass of `DrivineAssembly` overrides `aclose()`, which a sync
+  close can't run.

docs/design.md

@@ -563,12 +563,13 @@ class AppAssembly(PilgrimAssembly):
 - `Assembled` injects Drivine's components into classes made by a framework, by type or by key:
   `gom: StatelessGraphObjectManager = Assembled()`.
 - Closing is pi-pilgrim's, as Spring closes drivine4j's beans with its context. The registry is a `shared` component,
-  and its `close` is the teardown pi-pilgrim infers, so closing the assembly (`await aclose()`, or leaving
-  `async with`) closes every driver and pool, and an application's assembly closes the `DrivineAssembly` it asked
-  `AssemblyHolder` for. `DrivineAssembly` has no `close` of its own: an async one would hide pi-pilgrim's sync
-  `close()`, and `with` would get a coroutine it never awaits. Entering `DrivineAssembly` calls `start()`, as Spring
-  starts drivine4j on refresh; entering an application's assembly doesn't, so the application calls
-  `drivine.start()`.
+  and its `close` is the teardown pi-pilgrim infers, so closing the assembly (`await aclose()`, or leaving `async with`)
+  closes every driver and pool, and an application's assembly closes the `DrivineAssembly` it asked `AssemblyHolder`
+  for, or holds as a component. The async `DrivineAssembly` is pi-pilgrim's `AsyncPilgrimAssembly`, whose `close()` is
+  `aclose()`, so 0.2.0's `await drivine.close()` closes it, and whose `with` refuses, as its exit couldn't await; it
+  overrides `__enter__` only to name the sync API in that refusal. The sync one is a plain `PilgrimAssembly`. Entering
+  `DrivineAssembly` calls `start()`, as Spring starts drivine4j on refresh; entering an application's assembly doesn't,
+  so the application calls `drivine.start()`.
 - The managers are plain classes, so an application with another container, or none, constructs them itself.
   `DatabaseRegistry` is a context manager for it: `async with DatabaseRegistry() as registry:` closes it.
 

drivine/drivine_assembly.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from typing import Self
+from typing import NoReturn, Self
 
-from pipilgrim import PilgrimAssembly
+from pipilgrim import AsyncPilgrimAssembly, AsyncTeardownRequired
 
 from drivine.connection.database_registry import DEFAULT, DatabaseRegistry
 from drivine.dsl.index_advice_policy import IndexAdvicePolicy
@@ -15,7 +15,7 @@
 from drivine.schema.schema_options import SchemaOptions
 
 
-class DrivineAssembly(PilgrimAssembly):
+class DrivineAssembly(AsyncPilgrimAssembly):
     """
     Drivine's components: the registry and the managers. On its own:
 
@@ -124,3 +124,10 @@ async def __aenter__(self) -> Self:
             await self.aclose()
             raise
         return self
+
+    def __enter__(self) -> NoReturn:
+        """Refuses, as pi-pilgrim's does, and names the sync API, whose assembly `with` closes."""
+        raise AsyncTeardownRequired(
+            f"{type(self).__name__} closes its drivers and pools asynchronously: use async with, or "
+            f"drivine.sync.DrivineAssembly for the sync API."
+        )

pyproject.toml

@@ -9,7 +9,7 @@ authors = [{ name = "Liberation Data" }]
 requires-python = ">=3.11"
 dependencies = [
     "pydantic>=2.7",
-    "pi-pilgrim>=0.3.0",
+    "pi-pilgrim>=0.4.0,<0.5",
 ]
 keywords = ["graph", "database", "neo4j", "memgraph", "falkordb", "neptune", "oracle", "postgres", "cypher", "ogm"]
 classifiers = [

scripts/unasync.py

@@ -8,7 +8,8 @@
 
 A module is written again when it has async code, or imports one that does; the rest (models, the query DSL, the
 backends) is shared by both APIs. drivine/concurrency.py, what the async API asks of asyncio, has a counterpart written
-by hand: drivine/sync/concurrency.py. A test marked `async_only` (one that runs tasks concurrently) is left out.
+by hand: drivine/sync/concurrency.py. A test marked `async_only` (one that runs tasks concurrently) is left out, as
+is a definition ASYNC_ONLY names.
 """
 
 from __future__ import annotations
@@ -48,6 +49,13 @@ class Tree:
 EXTERNAL = {"falkordb.asyncio": "falkordb"}
 """Driver modules whose sync counterpart is another module."""
 
+ASYNC_ONLY = {"drivine.drivine_assembly": ("DrivineAssembly.__enter__",)}
+"""
+Definitions left out of the sync code, by module: the async API's own. The async DrivineAssembly's `__enter__` refuses
+`with`, naming the sync API. Made sync, it would replace the `__enter__` that `__aenter__` becomes, with which the sync
+assembly, a plain PilgrimAssembly, starts.
+"""
+
 RENAMES = {
     "AsyncIterator": "Iterator",
     "AsyncIterable": "Iterable",
@@ -73,6 +81,7 @@ class Tree:
     "AsyncConnection": "Connection",
     "AsyncCursor": "Cursor",
     "AsyncIteratorFixture": "IteratorFixture",
+    "AsyncPilgrimAssembly": "PilgrimAssembly",
 }
 """Names whose sync counterpart is another name: the language's, and each driver's."""
 
@@ -137,16 +146,34 @@ def written_again(sources: dict[str, ast.Module]) -> set[str]:
         marked |= more
 
 
-def without_async_only(text: str, source: ast.Module) -> str:
-    """The text without the tests marked `async_only`."""
+def without_async_only(text: str, source: ast.Module, module: str = "") -> str:
+    """The text without the tests marked `async_only`, or the module's definitions that ASYNC_ONLY names."""
     lines = text.splitlines(keepends=True)
     dropped: set[int] = set()
-    for it in ast.walk(source):
-        if isinstance(it, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and any(
-            ast.unparse(decorator).endswith("mark.async_only") for decorator in it.decorator_list
-        ):
-            first = min(decorator.lineno for decorator in it.decorator_list)
-            dropped.update(range(first - 1, (it.end_lineno or it.lineno)))
+    named = ASYNC_ONLY.get(module, ())
+    found: set[str] = set()
+    definitions = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
+
+    def visit(body: list[ast.stmt], prefix: str) -> None:
+        for it in body:
+            if not isinstance(it, definitions):
+                continue
+            name = f"{prefix}{it.name}"
+            if name in named:
+                found.add(name)
+            if name in named or any(
+                ast.unparse(decorator).endswith("mark.async_only") for decorator in it.decorator_list
+            ):
+                first = min([it.lineno, *(decorator.lineno for decorator in it.decorator_list)])
+                dropped.update(range(first - 1, (it.end_lineno or it.lineno)))
+            elif isinstance(it, ast.ClassDef):
+                visit(it.body, f"{name}.")
+
+    visit(source.body, "")
+    missing = [it for it in named if it not in found]
+    if missing:
+        # A renamed definition would otherwise be written into the sync code without a word.
+        raise SystemExit(f"ASYNC_ONLY names {', '.join(missing)} in {module}, which has no such definition.")
     return "".join(line for i, line in enumerate(lines) if i not in dropped)
 
 
@@ -252,7 +279,7 @@ def write(into: Path) -> list[Path]:
     written: list[Path] = []
     for name in sorted(generated):
         path = paths[name]
-        text = rewritten(without_async_only(path.read_text(), sources[name]), generated)
+        text = rewritten(without_async_only(path.read_text(), sources[name], name), generated)
         left = leftovers(text)
         if left:
             raise SystemExit(f"{path.relative_to(ROOT)} still holds {', '.join(left)} once made sync.")