function_storage: shard_key routing for CF DO + tx_cached_value fixture

Threads a shard_key through every FunctionStorage method so the CF DO
backend can route storage ops to a stable Durable Object instance per
attach (or per authenticated user, or anonymous fallback). BoundStorage
and TransactionBoundStorage derive the key from request.attach_id /
auth.principal via _derive_shard_key and pass it to every backend call.

- function_storage.py: add _derive_shard_key (attach_id → "att-<hex>",
  authenticated principal → "prn-<sha256>", anonymous → "loc-anon");
  shard_key=str kwarg on every Protocol method; BoundStorage and
  TransactionBoundStorage accept request= / attach_id= / auth= and
  thread shard_key through. SQLite backend ignores shard_key (no
  sharding).
- function_storage_cf_do.py / function_storage_azure_sql.py: shard_key
  threaded through every method (CF DO uses it as the DO routing key;
  Azure ignores it).
- protocol.py, scalar_function.py, table_function.py,
  table_in_out_function.py: pass request= / auth= to BoundStorage and
  TransactionBoundStorage at every construction site.
- tests/test_function_storage*.py: cover the new kwarg and the derive
  precedence.
- vgi/_test_fixtures/table/transaction_storage.py: new tx_cached_value
  fixture exercising bind→init→process opaque-data round-trip via
  BindResponse.opaque_data + transaction_storage caching. on_init
  declares max_workers=1 because the function emits exactly one row
  with no work-queue coordination.
- vgi/_test_fixtures/table/__init__.py + worker.py: register the new
  fixture.
- profiling_example.py: minor adjustment.
- worker.py: minor adjustment to init dispatch.
- docs/shared-storage.md: point at the relocated CF DO repo.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: rustyconover

docs/shared-storage.md

@@ -110,10 +110,12 @@ all operations are inherently atomic without locking.
 
 ### Setup
 
-1. Deploy the Cloudflare Worker (from `cloudflare/vgi-storage/`):
+1. Deploy the Cloudflare Worker. The Worker source lives in a separate
+   repository: [`vgi-cloudflare-durable-object-storage`](https://github.com/query-farm/vgi-cloudflare-durable-object-storage).
 
 ```bash
-cd cloudflare/vgi-storage
+git clone https://github.com/query-farm/vgi-cloudflare-durable-object-storage
+cd vgi-cloudflare-durable-object-storage
 npm install
 npx wrangler deploy
 ```

tests/test_function_storage.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from vgi.function_storage import FunctionStorageSqlite, UnknownInvocationError
+from vgi.function_storage import FunctionStorageSqlite
 
 
 class TestFunctionStorageSqlite:
@@ -106,11 +106,13 @@ def test_queue_pop_empty_queue(self, storage: FunctionStorageSqlite) -> None:
         result = storage.queue_pop(invocation_id)
         assert result is None
 
-    def test_queue_pop_unknown_invocation_raises(self, storage: FunctionStorageSqlite) -> None:
-        """Test popping from unknown invocation raises error."""
+    def test_queue_pop_never_pushed_returns_none(self, storage: FunctionStorageSqlite) -> None:
+        """Test popping an id that was never pushed returns None.
+
+        No distinction from drained queue per the contract.
+        """
         unknown_id = b"never_seen_before"
-        with pytest.raises(UnknownInvocationError):
-            storage.queue_pop(unknown_id)
+        assert storage.queue_pop(unknown_id) is None
 
     def test_queue_clear(self, storage: FunctionStorageSqlite) -> None:
         """Test clearing the work queue."""
@@ -123,26 +125,8 @@ def test_queue_clear(self, storage: FunctionStorageSqlite) -> None:
         cleared = storage.queue_clear(invocation_id)
         assert cleared == 3
 
-        # After clear, invocation is unregistered so pop should raise
-        with pytest.raises(UnknownInvocationError):
-            storage.queue_pop(invocation_id)
-
-    def test_queue_clear_unregisters_invocation(self, storage: FunctionStorageSqlite) -> None:
-        """Test that queue_clear unregisters the invocation_id."""
-        invocation_id = b"inv123"
-
-        # Register by pushing
-        storage.queue_push(invocation_id, [b"item1"])
-
-        # Pop should work (known invocation)
-        storage.queue_pop(invocation_id)
-
-        # Clear unregisters
-        storage.queue_clear(invocation_id)
-
-        # Now pop should raise (unknown after clear)
-        with pytest.raises(UnknownInvocationError):
-            storage.queue_pop(invocation_id)
+        # After clear, pop returns None (queue is empty/unknown)
+        assert storage.queue_pop(invocation_id) is None
 
     def test_queue_push_empty_still_registers(self, storage: FunctionStorageSqlite) -> None:
         """Test that pushing empty list still registers invocation_id."""

tests/test_function_storage_azure_sql.py

@@ -8,7 +8,6 @@
 
 pymssql = pytest.importorskip("pymssql")
 
-from vgi.function_storage import UnknownInvocationError  # noqa: E402
 from vgi.function_storage_azure_sql import FunctionStorageAzureSql  # noqa: E402
 
 
@@ -155,8 +154,8 @@ def test_queue_pop_returns_item(self, storage: FunctionStorageAzureSql, mock_cur
         """Test that queue_pop returns an item when available."""
         execution_id = b"\x02" * 16
 
-        # Combined SQL returns (registered, work_item)
-        mock_cursor.set_results([(1, b"item1")])
+        # OUTPUT deleted.work_item returns one row with the item bytes.
+        mock_cursor.set_results([(b"item1",)])
 
         result = storage.queue_pop(execution_id)
         assert result == b"item1"
@@ -169,20 +168,22 @@ def test_queue_pop_empty_queue(self, storage: FunctionStorageAzureSql, mock_curs
         """Test popping from registered but empty queue returns None."""
         execution_id = b"\x02" * 16
 
-        # Combined SQL returns (registered=1, work_item=NULL) for empty queue
-        mock_cursor.set_results([(1, None)])
+        # No row deleted → no OUTPUT row.
+        mock_cursor.set_results([])
 
         result = storage.queue_pop(execution_id)
         assert result is None
 
-    def test_queue_pop_unknown_invocation_raises(
+    def test_queue_pop_never_pushed_returns_none(
         self, storage: FunctionStorageAzureSql, mock_cursor: _MockCursor
     ) -> None:
-        """Test popping from unknown invocation raises error."""
-        # Combined SQL returns (registered=0, NULL) for unregistered invocation
-        mock_cursor.set_results([(0, None)])
-        with pytest.raises(UnknownInvocationError):
-            storage.queue_pop(b"\xff" * 16)
+        """Popping an id that was never pushed returns None.
+
+        No distinction from drained queue per the contract.
+        """
+        # Same wire shape as empty queue — no OUTPUT row.
+        mock_cursor.set_results([])
+        assert storage.queue_pop(b"\xff" * 16) is None
 
     def test_queue_push_empty_list(self, storage: FunctionStorageAzureSql, mock_cursor: _MockCursor) -> None:
         """Test pushing empty list returns 0."""

tests/test_function_storage_cf_do.py

@@ -11,7 +11,6 @@
 import httpx
 import pytest
 
-from vgi.function_storage import UnknownInvocationError
 from vgi.function_storage_cf_do import FunctionStorageCfDo
 
 
@@ -162,14 +161,6 @@ def test_queue_pop_empty_queue(self, storage: FunctionStorageCfDo, mock_transpor
         result = storage.queue_pop(b"\x02" * 16)
         assert result is None
 
-    def test_queue_pop_unknown_invocation_raises(
-        self, storage: FunctionStorageCfDo, mock_transport: _MockTransport
-    ) -> None:
-        """Test popping from unknown invocation raises error."""
-        mock_transport.queue_response(404, {"error": "unknown_invocation"})
-        with pytest.raises(UnknownInvocationError):
-            storage.queue_pop(b"\xff" * 16)
-
     def test_queue_push_empty_list(self, storage: FunctionStorageCfDo, mock_transport: _MockTransport) -> None:
         """Test pushing empty list returns 0."""
         mock_transport.queue_response(200, {"count": 0})
@@ -192,17 +183,6 @@ def test_queue_clear(self, storage: FunctionStorageCfDo, mock_transport: _MockTr
         cleared = storage.queue_clear(b"\x02" * 16)
         assert cleared == 3
 
-    def test_queue_clear_unregisters_invocation(
-        self, storage: FunctionStorageCfDo, mock_transport: _MockTransport
-    ) -> None:
-        """Test that queue_clear causes subsequent pop to raise."""
-        mock_transport.queue_response(200, {"cleared": 0})
-        storage.queue_clear(b"\x02" * 16)
-
-        mock_transport.queue_response(404, {"error": "unknown_invocation"})
-        with pytest.raises(UnknownInvocationError):
-            storage.queue_pop(b"\x02" * 16)
-
     def test_queue_clear_empty_queue(self, storage: FunctionStorageCfDo, mock_transport: _MockTransport) -> None:
         """Test clearing an empty queue returns 0."""
         mock_transport.queue_response(200, {"cleared": 0})
@@ -423,7 +403,7 @@ def _destructive_call_cases(
         return [
             ("worker_put", "worker_put"),
             ("worker_collect", "worker_collect"),
-            ("scan_worker_put", "scan_worker_put"),
+            ("stream_state_put", "stream_state_put"),
             ("queue_push", "queue_push"),
             ("queue_pop", "queue_pop"),
             ("queue_clear", "queue_clear"),
@@ -449,8 +429,8 @@ def _invoke(
             storage.worker_put(eid, worker_id=1, state=b"x")
         elif method_name == "worker_collect":
             storage.worker_collect(eid)
-        elif method_name == "scan_worker_put":
-            storage.scan_worker_put(eid, sid, b"x")
+        elif method_name == "stream_state_put":
+            storage.stream_state_put(eid, sid, b"x")
         elif method_name == "queue_push":
             storage.queue_push(eid, [b"a", b"b"])
         elif method_name == "queue_pop":
@@ -507,7 +487,7 @@ def test_read_only_calls_do_not_send_attempt_id(
         mock_transport.queue_response(200, {"rows": []})
         storage.worker_scan(b"\x01" * 16)
         mock_transport.queue_response(200, {"rows": []})
-        storage.scan_worker_scan(b"\x01" * 16)
+        storage.stream_state_scan(b"\x01" * 16)
         mock_transport.queue_response(200, {"values": [None]})
         storage.transaction_state_get(b"\xaa" * 16, [b"k"])
         mock_transport.queue_response(200, {"rows": [None]})
@@ -561,20 +541,22 @@ def test_attempt_id_unique_per_logical_call(
 
     # --- Scan Worker State Tests ---
 
-    def test_scan_worker_put_round_trip(self, storage: FunctionStorageCfDo, mock_transport: _MockTransport) -> None:
-        """scan_worker_put sends execution_id/stream_id/state and an attempt_id."""
+    def test_stream_state_put_round_trip(self, storage: FunctionStorageCfDo, mock_transport: _MockTransport) -> None:
+        """stream_state_put sends execution_id/stream_id/state and an attempt_id."""
         eid = b"\x01" * 16
         sid = b"\xde" * 16
         mock_transport.queue_response(200, {})
-        storage.scan_worker_put(eid, sid, b"payload")
+        storage.stream_state_put(eid, sid, b"payload")
         body = _body(mock_transport.requests[0])
         assert base64.b64decode(body["execution_id"]) == eid
         assert base64.b64decode(body["stream_id"]) == sid
         assert base64.b64decode(body["state"]) == b"payload"
         assert self._ATTEMPT_RE.match(body["attempt_id"])
 
-    def test_scan_worker_scan_returns_pairs(self, storage: FunctionStorageCfDo, mock_transport: _MockTransport) -> None:
-        """scan_worker_scan returns (stream_id, state) tuples and carries no attempt_id."""
+    def test_stream_state_scan_returns_pairs(
+        self, storage: FunctionStorageCfDo, mock_transport: _MockTransport
+    ) -> None:
+        """stream_state_scan returns (stream_id, state) tuples and carries no attempt_id."""
         sid_a = b"\xde" * 16
         sid_b = b"\xad" * 16
         mock_transport.queue_response(
@@ -586,7 +568,7 @@ def test_scan_worker_scan_returns_pairs(self, storage: FunctionStorageCfDo, mock
                 ],
             },
         )
-        rows = storage.scan_worker_scan(b"\x01" * 16)
+        rows = storage.stream_state_scan(b"\x01" * 16)
         assert rows == [(sid_a, b"alpha"), (sid_b, b"beta")]
 
     # --- Aggregate State Tests ---

tests/test_function_storage_cf_do_integration.py

@@ -22,7 +22,6 @@
 import httpx
 import pytest
 
-from vgi.function_storage import UnknownInvocationError
 from vgi.function_storage_cf_do import FunctionStorageCfDo
 
 # Skip the entire module unless explicitly enabled.
@@ -193,37 +192,38 @@ def test_queue_pop_replay_returns_same_item(storage: FunctionStorageCfDo) -> Non
     assert storage.queue_pop(eid) == b"b"
 
 
-def test_queue_pop_unknown_invocation(storage: FunctionStorageCfDo) -> None:
-    """Verify pop against a never-pushed execution_id raises UnknownInvocationError."""
-    with pytest.raises(UnknownInvocationError):
-        storage.queue_pop(_eid())
+def test_queue_pop_never_pushed_returns_none(storage: FunctionStorageCfDo) -> None:
+    """Pop against a never-pushed execution_id returns None.
+
+    No distinction from drained queue per the contract.
+    """
+    assert storage.queue_pop(_eid()) is None
 
 
-def test_queue_clear_unregisters(storage: FunctionStorageCfDo) -> None:
-    """Clear must remove queue rows and unregister the invocation."""
+def test_queue_clear_then_pop_returns_none(storage: FunctionStorageCfDo) -> None:
+    """Clear removes queue rows; subsequent pop returns None (queue empty)."""
     eid = _eid()
     storage.queue_push(eid, [b"a", b"b"])
     cleared = storage.queue_clear(eid)
     assert cleared == 2
 
-    with pytest.raises(UnknownInvocationError):
-        storage.queue_pop(eid)
+    assert storage.queue_pop(eid) is None
 
 
 # ---------------------------------------------------------------------------
 # Scan worker state
 # ---------------------------------------------------------------------------
 
 
-def test_scan_worker_put_and_scan_round_trip(storage: FunctionStorageCfDo) -> None:
+def test_stream_state_put_and_scan_round_trip(storage: FunctionStorageCfDo) -> None:
     """Round-trip scan worker state for two distinct stream_ids."""
     eid = _eid()
     sid_a = secrets.token_bytes(16)
     sid_b = secrets.token_bytes(16)
-    storage.scan_worker_put(eid, sid_a, b"alpha")
-    storage.scan_worker_put(eid, sid_b, b"beta")
+    storage.stream_state_put(eid, sid_a, b"alpha")
+    storage.stream_state_put(eid, sid_b, b"beta")
 
-    rows = dict(storage.scan_worker_scan(eid))
+    rows = dict(storage.stream_state_scan(eid))
     assert rows == {sid_a: b"alpha", sid_b: b"beta"}
 
 
@@ -289,7 +289,13 @@ def test_invalid_attempt_id_format_returns_400(storage: FunctionStorageCfDo) ->
     import base64
 
     eid = _eid()
-    body = {"execution_id": base64.b64encode(eid).decode(), "attempt_id": "not-hex"}
+    body = {
+        "execution_id": base64.b64encode(eid).decode(),
+        "attempt_id": "not-hex",
+        # Worker rejects missing shard_key first; supply a placeholder so
+        # the attempt_id validation downstream is what we actually exercise.
+        "shard_key": "loc-anon",
+    }
     # Drive raw httpx so we can ship an arbitrary attempt_id string.
     resp = storage._client.post("/queue_pop", json=body)
     assert resp.status_code == 400
@@ -312,10 +318,7 @@ def test_concurrent_pops_no_dupes(storage: FunctionStorageCfDo) -> None:
     lock = threading.Lock()
 
     def pop() -> None:
-        try:
-            r = storage.queue_pop(eid)
-        except UnknownInvocationError:
-            r = None
+        r = storage.queue_pop(eid)
         with lock:
             results.append(r)
 
@@ -330,10 +333,14 @@ def pop() -> None:
     assert len(set(items)) == n_items, "no item must be popped twice"
 
 
-def test_health_check_get_returns_405(storage: FunctionStorageCfDo) -> None:
-    """Sanity: GET is rejected at the router."""
+def test_health_check_get_is_rejected(storage: FunctionStorageCfDo) -> None:
+    """Sanity: GET is rejected at the router.
+
+    Accept either 405 (when no auth token is configured) or 401 (when one
+    is — the router validates the bearer token before method).
+    """
     resp = httpx.get(_URL or "")
-    assert resp.status_code == 405
+    assert resp.status_code in (401, 405)
 
 
 # ---------------------------------------------------------------------------