Merge pull request #804 from morluto/agent/fix-schema-registry-transactions

Preserve schema registry transaction invariants

Author: morluto

src/jacobian/schema_registry.py

@@ -131,12 +131,11 @@ def _register_canonical(
 
         registration = (name, version, canonical_schema)
         transaction_identity = self.store.transaction_identity
-        registrations = self._registrations
-        if transaction_identity is not None:
-            registrations = self._pending.setdefault(
-                transaction_identity,
-                _PendingRegistrations(),
-            ).registrations
+        if transaction_identity is None:
+            registrations = self._registrations
+        else:
+            pending = self._pending.get(transaction_identity)
+            registrations = pending.registrations if pending is not None else {}
         cached_uri = registrations.get(registration)
         if cached_uri is not None:
             return cached_uri
@@ -155,11 +154,16 @@ def _register_canonical(
             version=version,
             definition=schema,
         )
-        registrations[registration] = schema_uri
         if transaction_identity is None:
+            self._registrations[registration] = schema_uri
             self._schema_bytes[schema_uri] = canonical_schema
         else:
-            self._pending[transaction_identity].schemas[schema_uri] = canonical_schema
+            pending = self._pending.setdefault(
+                transaction_identity,
+                _PendingRegistrations(),
+            )
+            pending.registrations[registration] = schema_uri
+            pending.schemas[schema_uri] = canonical_schema
         return schema_uri
 
     def register_model(
@@ -179,10 +183,21 @@ def register_model(
         """
 
         self._reconcile_pending()
+        canonical_schema = _model_schema_bytes(model)
+        schema = cast(dict[str, Any], loads_strict_json(canonical_schema))
+        self._ensure_model_contract_available(
+            self.store.descriptor_uri(
+                kind="schema",
+                name=name,
+                version=version,
+                definition=schema,
+            ),
+            model,
+        )
         schema_uri = self._register_canonical(
             name=name,
             version=version,
-            canonical_schema=_model_schema_bytes(model),
+            canonical_schema=canonical_schema,
         )
         self._bind_model_contract(schema_uri, model)
         if producer_only:
@@ -214,16 +229,35 @@ def _bind_model_contract(
         schema_uri: str,
         model: type[BaseModel],
     ) -> None:
+        self._ensure_model_contract_available(schema_uri, model)
         transaction_identity = self.store.transaction_identity
-        model_contracts = self._model_contracts
         if transaction_identity is not None:
             model_contracts = self._pending[transaction_identity].model_contracts
-        registered = model_contracts.get(schema_uri)
+        else:
+            model_contracts = self._model_contracts
+        model_contracts[schema_uri] = model
+
+    def _ensure_model_contract_available(
+        self,
+        schema_uri: str,
+        model: type[BaseModel],
+    ) -> None:
+        committed = self._model_contracts.get(schema_uri)
+        if committed is not None and committed is not model:
+            raise SchemaRegistryError(
+                "one schema URI cannot use multiple model-backed contracts"
+            )
+        transaction_identity = self.store.transaction_identity
+        if transaction_identity is None:
+            return
+        pending = self._pending.get(transaction_identity)
+        registered = (
+            pending.model_contracts.get(schema_uri) if pending is not None else None
+        )
         if registered is not None and registered is not model:
             raise SchemaRegistryError(
                 "one schema URI cannot use multiple model-backed contracts"
             )
-        model_contracts[schema_uri] = model
 
     def resolve(self, schema_uri: str) -> dict[str, Any]:
         """Load a previously registered schema definition."""
@@ -263,6 +297,9 @@ def _reconcile_pending(self) -> None:
         for transaction_identity, pending in tuple(self._pending.items()):
             if transaction_identity == active_identity:
                 continue
+            if not pending.schemas:
+                del self._pending[transaction_identity]
+                continue
             witness_uri = next(iter(pending.schemas))
             try:
                 self.store.get_descriptor(witness_uri, expected_kind="schema")

tests/component/schemas/test_schema_registry.py

@@ -326,3 +326,110 @@ def unexpected_blob_write(_data: bytes) -> str:
 
     with pytest.raises(SchemaValidationError, match="pair must be ordered"):
         runtime_registry.validate(schema_uri, {"first": 2, "second": 1})
+
+
+def test_failed_transactional_registration_does_not_leave_empty_pending_state(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    store = ArtifactRepository(tmp_path)
+    registry = SchemaRegistry(store)
+    register_descriptor = store.register_descriptor
+
+    def fail_first_registration(
+        *,
+        kind: str,
+        name: str,
+        version: str,
+        definition: dict[str, Any],
+    ) -> str:
+        if name == "will-fail":
+            raise RuntimeError("simulated descriptor failure")
+        return register_descriptor(
+            kind=kind,
+            name=name,
+            version=version,
+            definition=definition,
+        )
+
+    monkeypatch.setattr(store, "register_descriptor", fail_first_registration)
+
+    with (
+        store.transaction(),
+        pytest.raises(
+            RuntimeError,
+            match="simulated descriptor failure",
+        ),
+    ):
+        registry.register(
+            name="will-fail",
+            version="1",
+            schema={"type": "object"},
+        )
+
+    assert registry.register(
+        name="still-usable",
+        version="1",
+        schema={"type": "object"},
+    ).startswith("artifact://sha256/")
+
+
+def test_transaction_cannot_replace_committed_model_contract(tmp_path: Path) -> None:
+    store = ArtifactRepository(tmp_path)
+    registry = SchemaRegistry(store)
+    schema_uri = registry.register_model(
+        name="shared-model-contract",
+        version="1",
+        model=_CachedSchemaModel,
+    )
+
+    with store.transaction():
+        with pytest.raises(SchemaRegistryError, match="one schema URI"):
+            registry.register_model(
+                name="shared-model-contract",
+                version="1",
+                model=_EquivalentCachedSchemaModel,
+            )
+        assert registry._pending == {}
+
+    assert registry._model_contracts[schema_uri] is _CachedSchemaModel
+
+
+def test_transaction_can_reattach_the_same_model_contract(tmp_path: Path) -> None:
+    store = ArtifactRepository(tmp_path)
+    registry = SchemaRegistry(store)
+    schema_uri = registry.register_model(
+        name="same-model-contract",
+        version="1",
+        model=_CachedSchemaModel,
+    )
+
+    with store.transaction():
+        assert (
+            registry.register_model(
+                name="same-model-contract",
+                version="1",
+                model=_CachedSchemaModel,
+            )
+            == schema_uri
+        )
+
+    assert registry._model_contracts[schema_uri] is _CachedSchemaModel
+
+
+def test_transaction_cannot_bind_two_models_to_the_same_schema(tmp_path: Path) -> None:
+    store = ArtifactRepository(tmp_path)
+    registry = SchemaRegistry(store)
+
+    with store.transaction():
+        registry.register_model(
+            name="intra-transaction-model-conflict",
+            version="1",
+            model=_CachedSchemaModel,
+        )
+        with pytest.raises(SchemaRegistryError, match="one schema URI"):
+            registry.register_model(
+                name="intra-transaction-model-conflict",
+                version="1",
+                model=_EquivalentCachedSchemaModel,
+            )