Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 49 additions & 12 deletions src/jacobian/schema_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Avoid using committed schemas as transaction witnesses

When a transaction re-registers a schema already present in _registrations, this branch ignores the committed cache and adds that existing URI to the pending set. If the same transaction then registers a new schema and rolls back, _reconcile_pending() sees the pre-existing URI as its witness and incorrectly promotes the entire pending set, including the rolled-back schema; subsequent registration returns the cached URI without recreating its missing descriptor. Preserve committed cache hits inside transactions, or track whether each witness was created by the transaction.

Useful? React with 👍 / 👎.

cached_uri = registrations.get(registration)
if cached_uri is not None:
return cached_uri
Expand All @@ -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(
Expand All @@ -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:
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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")
Expand Down
107 changes: 107 additions & 0 deletions tests/component/schemas/test_schema_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Loading