fix: ops functions write through the wrapper's own __setitem__/__delitem__ (#3)

cosmodol/routing.py

@@ -0,0 +1,118 @@
+"""Route an operation through a wrapped store's own ``__setitem__``/``__delitem__``.
+
+The functions in ``cosmodol.ops`` need Cosmos semantics a mapping write does not have (a
+conditional replace, a transactional batch, a cascading delete, a parameterised create).
+Calling the leaf directly skips every wrapping layer's own write logic: write guards,
+write-through caches, ``preset=`` hooks, scope filters on new keys (i2mint/cosmodol#3).
+Instead, an operation *redirects* the leaf: it installs a one-shot handler for one leaf
+method, then performs the plain mapping write (``store[k] = v`` / ``del store[k]``) on the
+caller's outer store. Every layer runs exactly as for that write; when the call reaches the
+leaf, the leaf hands its (fully mapped) arguments to the handler instead of doing its usual
+upsert or delete.
+
+>>> class Leaf(dict):
+...     def __setitem__(self, k, v):
+...         if (handler := claim(self, "__setitem__", k)) is not None:
+...             return handler(k, v)
+...         super().__setitem__(k, v)
+>>> leaf = Leaf()
+>>> with redirect(leaf, "__setitem__", "a", lambda k, v: f"handled {k}={v}") as r:
+...     leaf["a"] = 1
+>>> r.result, dict(leaf)
+('handled a=1', {})
+
+The redirect is pinned to the leaf key the operation computed (``dol.inner_most_key``), so a
+layer that writes some *other* key first (an index, an audit record) cannot have the
+operation's semantics applied to it: that write is refused instead.
+
+The redirect is held in a ``ContextVar``, so it is invisible to other threads and tasks, and
+it only applies to the leaf object it names. The flip side: a layer that performs the write
+in another thread, or later (a write-behind buffer), escapes it, and that write happens as
+the leaf's ordinary upsert or delete.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import Any, Callable, Iterator, Optional
+
+
+class RedirectError(RuntimeError):
+    """A wrapping layer did not turn one mapping write into exactly one leaf write."""
+
+
+#: The leaf methods a redirect guards: while one is pending on a leaf, any of them called
+#: with anything but the expected method and key is refused.
+WRITE_METHODS = ("__setitem__", "__delitem__")
+
+
+@dataclass
+class Redirect:
+    """One pending redirect of ``leaf.<method>(key, ...)`` to ``handler``.
+
+    After the block, ``claimed`` says whether the write reached the leaf, ``result`` is
+    what the handler returned, and ``error`` what it raised (re-raised by the caller
+    even if a layer swallowed it).
+    """
+
+    leaf: Any
+    method: str
+    key: Any
+    handler: Callable
+    claimed: bool = False
+    result: Any = None
+    error: Optional[BaseException] = None
+
+
+_active: ContextVar[Optional[Redirect]] = ContextVar("cosmodol_redirect", default=None)
+
+
+@contextmanager
+def redirect(leaf, method: str, key, handler: Callable) -> Iterator[Redirect]:
+    """Within the block, the call ``leaf.<method>(key, ...)`` goes to ``handler``.
+
+    Any other write call on ``leaf`` within the block (another method, another key, or
+    the same call a second time) raises ``RedirectError`` before touching anything: one
+    mapping write was expected to reach the leaf once, as ``key``.
+    """
+    r = Redirect(leaf, method, key, handler)
+    token = _active.set(r)
+    try:
+        yield r
+    finally:
+        _active.reset(token)
+
+
+def claim(leaf, method: str, key) -> Optional[Callable]:
+    """The handler redirecting ``leaf.<method>(key, ...)`` right now, or ``None``.
+
+    Called by a leaf at the top of each of its ``WRITE_METHODS``. The returned callable
+    runs the handler and records its result or error.
+    """
+    r = _active.get()
+    if r is None or r.leaf is not leaf:
+        return None
+    if r.claimed:
+        outcome = "failed" if r.error is not None else "was performed"
+        raise RedirectError(
+            f"A wrapping layer made a second {method} call ({key!r}) on the cosmodol "
+            f"store during one operation; refused. The first call ({r.key!r}) {outcome}."
+        )
+    if (method, key) != (r.method, r.key):
+        raise RedirectError(
+            f"A wrapping layer called {method}({key!r}) on the cosmodol store where "
+            f"the operation expected {r.method}({r.key!r}); refused before acting."
+        )
+
+    def run(*args):
+        r.claimed = True
+        try:
+            r.result = r.handler(*args)
+        except BaseException as e:
+            r.error = e
+            raise
+        return r.result
+
+    return run

cosmodol/stores.py

@@ -28,6 +28,7 @@
     batch as _batch,
 )
 from cosmodol.connection import CosmosConnection
+from cosmodol.routing import claim
 from cosmodol.errors import (
     KeyMismatchError,
     MovedToFunction,
@@ -181,11 +182,15 @@ def __getitem__(self, k: str) -> dict:
         return self._maybe_strip(item)
 
     def __setitem__(self, k: str, v) -> None:
+        if (handler := claim(self, "__setitem__", k)) is not None:
+            return handler(k, v)  # an ops function, routed through every wrapper
         body = self._prepare_body(k, v)
         _item, hdrs = point_upsert(self.container, body)
         self._observe("set", hdrs)
 
     def __delitem__(self, k: str) -> None:
+        if (handler := claim(self, "__delitem__", k)) is not None:
+            return handler(k)
         if self.strict_keys:
             validate_cosmos_id(k)
         hdrs = point_delete(self.container, k, self.partition_key_value)

cosmodol/tests/test_ops.py

@@ -468,13 +468,22 @@ def test_scope_filter_stops_forced_container_delete():
     assert sorted(db.containers) == ["app_events", "events"]
 
 
-def test_add_database_does_not_raise_after_creating():
+def test_add_database_does_not_read_back_through_the_wrapper():
+    acct, client = _account(True)
+    db = add_database(acct, "extra")
+    assert db.name == "app_extra" and "app_extra" in client.dbs
+
+
+def test_add_database_refuses_a_name_outside_the_scope():
+    # Was test_add_database_does_not_raise_after_creating: before #3 the out-of-scope
+    # database was created; now the scope's own __setitem__ refuses it first.
     from dol import filt_iter
 
     acct, client = _account(True)
     hidden = filt_iter(acct, filt=lambda k: k.startswith("x"))
-    db = add_database(hidden, "extra")
-    assert db.name == "app_extra" and "app_extra" in client.dbs
+    with pytest.raises(KeyError):
+        add_database(hidden, "extra")
+    assert "app_extra" not in client.dbs
 
 
 def test_forced_delete_handles_hierarchical_and_missing_partition_keys():