diff --git a/pg_llm_batch/reconciliation_single_flight.py b/pg_llm_batch/reconciliation_single_flight.py new file mode 100644 index 000000000..dd3427e7a --- /dev/null +++ b/pg_llm_batch/reconciliation_single_flight.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) ContextualWisdomLab. +"""Tenant-qualified PostgreSQL single-flight authority for reconciliation.""" + +from __future__ import annotations + +from contextlib import contextmanager +from hashlib import sha256 +from typing import Any, Iterator + +from .db import ( + validate_endpoint_alias, + validate_remote_resource_id, + validate_tenant_scope, +) +from .exceptions import PgLlmBatchError, ValidationError +from .reconciliation import ReconciliationCandidate + +_LOCK_DOMAIN = b"pg-llm-batch:reconciliation-single-flight:v1" + + +class ReconciliationSingleFlightError(PgLlmBatchError): + """Raised when PostgreSQL cannot prove single-flight lock state safely.""" + + def __init__(self, phase: str, reason: str) -> None: + """Create one bounded content-free advisory-lock failure.""" + super().__init__( + message="Reconciliation single-flight database operation failed", + error_code="RECONCILIATION_SINGLE_FLIGHT_FAILED", + details={"phase": phase, "reason": reason}, + ) + + +def _invalid_identity() -> ValidationError: + """Build the fixed redacted single-flight identity error.""" + return ValidationError( + field="reconciliation_single_flight_identity", + value="", + reason=( + "must contain a valid trusted tenant scope, endpoint alias, and " + "remote batch identifier" + ), + message="Reconciliation single-flight identity is invalid", + ) + + +def _validated_identity( + tenant_scope: Any, + candidate: Any, +) -> tuple[str, str, str]: + """Return one canonical exact-type lock identity without reflecting bad input. + + Caller-owned identity evidence must use an exact built-in tenant string and + the exact package-owned candidate dataclass. Candidate members must also be + exact built-in strings. Subclasses are rejected before attribute methods, + regex, normalization, hashing, or encoding authority can execute. + """ + if ( + type(tenant_scope) is not str + or type(candidate) is not ReconciliationCandidate + ): + raise _invalid_identity() from None + endpoint_value = candidate.endpoint_alias + remote_value = candidate.remote_batch_id + if type(endpoint_value) is not str or type(remote_value) is not str: + raise _invalid_identity() from None + try: + tenant = validate_tenant_scope(tenant_scope) + endpoint_alias = validate_endpoint_alias(endpoint_value) + remote_batch_id = validate_remote_resource_id( + remote_value, + "remote_batch_id", + ) + except ValidationError: + raise _invalid_identity() from None + return tenant, endpoint_alias, remote_batch_id + + +def _lock_key(tenant_scope: str, endpoint_alias: str, remote_batch_id: str) -> int: + """Derive one domain-separated signed PostgreSQL advisory-lock key.""" + digest = sha256() + digest.update(_LOCK_DOMAIN) + for value in (tenant_scope, endpoint_alias, remote_batch_id): + encoded = value.encode("utf-8") + digest.update(len(encoded).to_bytes(4, byteorder="big")) + digest.update(encoded) + return int.from_bytes(digest.digest()[:8], byteorder="big", signed=True) + + +def _execute_boolean_lock_operation( + cursor: Any, + sql: str, + lock_key: int, + *, + phase: str, +) -> bool: + """Execute one parameterized advisory-lock operation with bounded evidence.""" + try: + cursor.execute(sql, (lock_key,)) + row = cursor.fetchone() + except Exception: + raise ReconciliationSingleFlightError( + phase, + "database_operation_failed", + ) from None + + if ( + type(row) not in (tuple, list) + or len(row) != 1 + or type(row[0]) is not bool + ): + reason = ( + "invalid_database_result" + if phase == "acquire" + else "lock_release_not_confirmed" + ) + raise ReconciliationSingleFlightError(phase, reason) from None + return row[0] + + +@contextmanager +def reconciliation_single_flight( + cursor: Any, + tenant_scope: str, + candidate: ReconciliationCandidate, +) -> Iterator[bool]: + """Hold a non-blocking cross-process lock for one reconciliation identity. + + The caller owns the PostgreSQL connection and must dedicate that database + session to at most one concurrent or nested single-flight attempt while this + context is active. PostgreSQL session advisory locks are re-entrant within a + session, so sharing the same session between concurrent attempts would not + provide mutual exclusion. The same session must remain alive for the entire + context lifetime; process or session loss releases the advisory lock. + + ``True`` means this session acquired authority to reconcile the validated + tenant/endpoint/remote-batch identity. ``False`` means another database + session currently holds that authority and the caller should defer. This is + a transient single-flight primitive, not a durable lease or an exactly-once + delivery guarantee. + """ + tenant, endpoint_alias, remote_batch_id = _validated_identity( + tenant_scope, + candidate, + ) + lock_key = _lock_key(tenant, endpoint_alias, remote_batch_id) + acquired = _execute_boolean_lock_operation( + cursor, + "SELECT pg_try_advisory_lock(%s)", + lock_key, + phase="acquire", + ) + if not acquired: + yield False + return + + try: + yield True + finally: + released = _execute_boolean_lock_operation( + cursor, + "SELECT pg_advisory_unlock(%s)", + lock_key, + phase="release", + ) + if not released: + raise ReconciliationSingleFlightError( + "release", + "lock_release_not_confirmed", + ) from None diff --git a/tests/test_reconciliation_single_flight.py b/tests/test_reconciliation_single_flight.py new file mode 100644 index 000000000..eef1007fb --- /dev/null +++ b/tests/test_reconciliation_single_flight.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for tenant-qualified cross-process reconciliation single-flight locks.""" + +from __future__ import annotations + +import traceback +from typing import Any + +import pytest + +from pg_llm_batch.exceptions import ValidationError +from pg_llm_batch.reconciliation import ReconciliationCandidate +from pg_llm_batch.reconciliation_single_flight import ( + ReconciliationSingleFlightError, + reconciliation_single_flight, +) + + +class RecordingCursor: + """Minimal cursor double recording advisory-lock SQL and bounded results.""" + + def __init__( + self, + results: list[Any], + *, + fail_execute_at: int | None = None, + execute_error: Exception | None = None, + ) -> None: + """Store fetch results plus an optional deterministic execute failure.""" + self.results = list(results) + self.fail_execute_at = fail_execute_at + self.execute_error = execute_error or RuntimeError("database sentinel secret") + self.calls: list[tuple[str, tuple[Any, ...]]] = [] + + def execute(self, sql: str, params: tuple[Any, ...]) -> None: + """Record one parameterized SQL execution or raise the configured failure.""" + self.calls.append((sql, params)) + if self.fail_execute_at == len(self.calls): + raise self.execute_error + + def fetchone(self) -> Any: + """Return the next configured database result.""" + return self.results.pop(0) if self.results else None + + +def test_single_flight_acquires_and_releases_same_parameterized_lock() -> None: + """A free identity must hold one session advisory lock for the context body.""" + cursor = RecordingCursor([(True,), (True,)]) + candidate = ReconciliationCandidate("gateway-a", "batch-1") + + with reconciliation_single_flight(cursor, "tenant-a", candidate) as acquired: + assert acquired is True + assert len(cursor.calls) == 1 + + assert len(cursor.calls) == 2 + acquire_sql, acquire_params = cursor.calls[0] + release_sql, release_params = cursor.calls[1] + assert "pg_try_advisory_lock(%s)" in acquire_sql + assert "pg_advisory_unlock(%s)" in release_sql + assert acquire_params == release_params + assert len(acquire_params) == 1 + assert type(acquire_params[0]) is int + assert -(1 << 63) <= acquire_params[0] < (1 << 63) + assert "tenant-a" not in acquire_sql + assert "gateway-a" not in acquire_sql + assert "batch-1" not in acquire_sql + + +def test_single_flight_contention_returns_false_without_unlock() -> None: + """A lock owned by another database session must defer without false release.""" + cursor = RecordingCursor([(False,)]) + + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ) as acquired: + assert acquired is False + + assert len(cursor.calls) == 1 + assert "pg_try_advisory_lock(%s)" in cursor.calls[0][0] + + +def test_single_flight_releases_when_context_body_raises() -> None: + """Caller failure must not leak an acquired session advisory lock.""" + cursor = RecordingCursor([(True,), (True,)]) + + with pytest.raises(ValueError, match="caller failure"): + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ) as acquired: + assert acquired is True + raise ValueError("caller failure") + + assert len(cursor.calls) == 2 + assert "pg_advisory_unlock(%s)" in cursor.calls[1][0] + + +@pytest.mark.parametrize( + ("tenant_scope", "candidate"), + [ + ("../tenant", ReconciliationCandidate("gateway-a", "batch-1")), + ("tenant-a", ReconciliationCandidate("\x00gateway", "batch-1")), + ("tenant-a", ReconciliationCandidate("gateway-a", "secret/provider/path")), + ], +) +def test_single_flight_invalid_identity_fails_before_database_work( + tenant_scope: str, + candidate: ReconciliationCandidate, +) -> None: + """Untrusted identity text must fail closed without cursor or reflected content.""" + cursor = RecordingCursor([]) + + with pytest.raises(ValidationError) as caught: + with reconciliation_single_flight(cursor, tenant_scope, candidate): + pytest.fail("invalid identity must never enter the context body") + + assert cursor.calls == [] + assert caught.value.details["value"] == "" + assert tenant_scope not in str(caught.value) + assert candidate.endpoint_alias not in str(caught.value) + assert candidate.remote_batch_id not in str(caught.value) + + +def test_single_flight_key_is_tenant_and_provider_identity_qualified() -> None: + """Distinct trusted identities must not intentionally share advisory-lock keys.""" + candidates = ( + ("tenant-a", ReconciliationCandidate("gateway-a", "batch-1")), + ("tenant-b", ReconciliationCandidate("gateway-a", "batch-1")), + ("tenant-a", ReconciliationCandidate("gateway-b", "batch-1")), + ("tenant-a", ReconciliationCandidate("gateway-a", "batch-2")), + ) + keys: list[int] = [] + + for tenant_scope, candidate in candidates: + cursor = RecordingCursor([(False,)]) + with reconciliation_single_flight(cursor, tenant_scope, candidate) as acquired: + assert acquired is False + keys.append(cursor.calls[0][1][0]) + + assert len(set(keys)) == len(keys) + + +def test_single_flight_normalizes_endpoint_alias_before_keying() -> None: + """Equivalent endpoint aliases must contend on one canonical lock identity.""" + keys: list[int] = [] + for endpoint_alias in ("gateway-a", " gateway-a "): + cursor = RecordingCursor([(False,)]) + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate(endpoint_alias, "batch-1"), + ) as acquired: + assert acquired is False + keys.append(cursor.calls[0][1][0]) + + assert keys[0] == keys[1] + + +@pytest.mark.parametrize("result", [None, (), (1,), (True, False), "true"]) +def test_single_flight_invalid_acquire_result_fails_closed(result: Any) -> None: + """Malformed database lock evidence must never be interpreted as acquisition.""" + cursor = RecordingCursor([result]) + + with pytest.raises(ReconciliationSingleFlightError) as caught: + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ): + pytest.fail("malformed lock evidence must not enter the context body") + + assert caught.value.details == { + "phase": "acquire", + "reason": "invalid_database_result", + } + + +def test_single_flight_redacts_database_acquire_failure() -> None: + """Lower-layer acquisition diagnostics must not escape package evidence.""" + sentinel = "postgres password=secret" + cursor = RecordingCursor( + [], + fail_execute_at=1, + execute_error=RuntimeError(sentinel), + ) + + with pytest.raises(ReconciliationSingleFlightError) as caught: + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ): + pytest.fail("failed acquisition must not enter the context body") + + assert caught.value.details == { + "phase": "acquire", + "reason": "database_operation_failed", + } + assert sentinel not in str(caught.value) + + +@pytest.mark.parametrize("release_result", [None, (), (False,), (1,), (True, False)]) +def test_single_flight_unconfirmed_release_fails_closed(release_result: Any) -> None: + """An acquired session lock needs explicit positive release evidence.""" + cursor = RecordingCursor([(True,), release_result]) + + with pytest.raises(ReconciliationSingleFlightError) as caught: + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ) as acquired: + assert acquired is True + + assert caught.value.details == { + "phase": "release", + "reason": "lock_release_not_confirmed", + } + + +def test_unconfirmed_release_suppresses_sensitive_caller_traceback() -> None: + """A failed unlock must not expose a sensitive caller exception in evidence.""" + sentinel = "caller payload secret" + cursor = RecordingCursor([(True,), (False,)]) + + with pytest.raises(ReconciliationSingleFlightError) as caught: + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ) as acquired: + assert acquired is True + raise RuntimeError(sentinel) + + rendered = "".join( + traceback.format_exception( + type(caught.value), + caught.value, + caught.value.__traceback__, + ) + ) + assert caught.value.__suppress_context__ is True + assert sentinel not in rendered + + +def test_malformed_release_suppresses_sensitive_caller_traceback() -> None: + """Malformed unlock evidence must not retain a sensitive caller exception.""" + sentinel = "caller payload secret from malformed release" + cursor = RecordingCursor([(True,), (1,)]) + + with pytest.raises(ReconciliationSingleFlightError) as caught: + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ) as acquired: + assert acquired is True + raise RuntimeError(sentinel) + + rendered = "".join( + traceback.format_exception( + type(caught.value), + caught.value, + caught.value.__traceback__, + ) + ) + assert caught.value.__suppress_context__ is True + assert caught.value.__cause__ is None + assert sentinel not in rendered + + +def test_single_flight_redacts_database_release_failure() -> None: + """Lower-layer release diagnostics must not escape package evidence.""" + sentinel = "postgres release secret" + cursor = RecordingCursor( + [(True,)], + fail_execute_at=2, + execute_error=RuntimeError(sentinel), + ) + + with pytest.raises(ReconciliationSingleFlightError) as caught: + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ) as acquired: + assert acquired is True + + assert caught.value.details == { + "phase": "release", + "reason": "database_operation_failed", + } + assert sentinel not in str(caught.value) diff --git a/tests/test_reconciliation_single_flight_exact_types.py b/tests/test_reconciliation_single_flight_exact_types.py new file mode 100644 index 000000000..9de53e04b --- /dev/null +++ b/tests/test_reconciliation_single_flight_exact_types.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Hostile-subclass regressions for reconciliation single-flight evidence.""" + +from __future__ import annotations + +import traceback +from typing import Any + +import pytest + +from pg_llm_batch.exceptions import ValidationError +from pg_llm_batch.reconciliation import ReconciliationCandidate +from pg_llm_batch.reconciliation_single_flight import ( + ReconciliationSingleFlightError, + reconciliation_single_flight, +) + +_SECRET_SENTINEL = "SECRET-SENTINEL hostile single-flight evidence" + + +class _RecordingCursor: + """Record lock operations while returning configured database evidence.""" + + def __init__(self, results: list[Any]) -> None: + self.results = list(results) + self.calls: list[tuple[str, tuple[Any, ...]]] = [] + + def execute(self, sql: str, params: tuple[Any, ...]) -> None: + """Record one parameterized advisory-lock operation.""" + self.calls.append((sql, params)) + + def fetchone(self) -> Any: + """Return the next configured database result.""" + return self.results.pop(0) if self.results else None + + +class _HostileTenantScope(str): + """Represent caller-controlled tenant text with executable behavior.""" + + def __hash__(self) -> int: + """Raise if the subclass reaches hashing or set membership.""" + raise RuntimeError(_SECRET_SENTINEL) + + +class _HostileCandidateText(str): + """Represent forged candidate text that must be refused before use.""" + + def __hash__(self) -> int: + """Raise if the subclass reaches hashing or set membership.""" + raise RuntimeError(_SECRET_SENTINEL) + + +class _HostileCandidate(ReconciliationCandidate): + """Execute caller code if candidate attributes are read before refusal.""" + + def __getattribute__(self, name: str) -> Any: + """Raise instead of supplying a trustworthy endpoint alias.""" + if name == "endpoint_alias": + raise RuntimeError(_SECRET_SENTINEL) + return super().__getattribute__(name) + + +class _HostileLockRow(tuple[Any, ...]): + """Execute database-row subclass code during result-shape validation.""" + + def __len__(self) -> int: + """Raise instead of supplying a trustworthy result shape.""" + raise RuntimeError(_SECRET_SENTINEL) + + +def _rendered_exception(error: BaseException) -> str: + """Render one traceback for confidentiality assertions.""" + return "".join(traceback.format_exception(type(error), error, error.__traceback__)) + + +@pytest.mark.parametrize( + ("tenant_scope", "candidate"), + [ + ( + _HostileTenantScope("tenant-a"), + ReconciliationCandidate("gateway-a", "batch-1"), + ), + ( + "tenant-a", + _HostileCandidate("gateway-a", "batch-1"), + ), + ], +) +def test_hostile_identity_subclasses_fail_before_database_work( + tenant_scope: Any, + candidate: Any, +) -> None: + """Identity subclasses must not execute before bounded validation.""" + cursor = _RecordingCursor([]) + + with pytest.raises(ValidationError) as caught: + with reconciliation_single_flight(cursor, tenant_scope, candidate): + pytest.fail("invalid identity must not enter the context body") + + assert caught.value.details["field"] == "reconciliation_single_flight_identity" + assert caught.value.details["value"] == "" + assert _SECRET_SENTINEL not in _rendered_exception(caught.value) + assert caught.value.__cause__ is None + assert caught.value.__context__ is None + assert cursor.calls == [] + + +@pytest.mark.parametrize( + "candidate", + [ + ReconciliationCandidate(_HostileCandidateText("gateway-a"), "batch-1"), + ReconciliationCandidate("gateway-a", _HostileCandidateText("batch-1")), + ], +) +def test_hostile_candidate_member_subclasses_fail_before_database_work( + candidate: ReconciliationCandidate, +) -> None: + """Candidate text subclasses must fail before validation or lock work.""" + cursor = _RecordingCursor([]) + + with pytest.raises(ValidationError) as caught: + with reconciliation_single_flight(cursor, "tenant-a", candidate): + pytest.fail("invalid identity must not enter the context body") + + assert caught.value.details["field"] == "reconciliation_single_flight_identity" + assert caught.value.details["value"] == "" + assert _SECRET_SENTINEL not in _rendered_exception(caught.value) + assert caught.value.__cause__ is None + assert caught.value.__context__ is None + assert cursor.calls == [] + + +def test_hostile_lock_row_subclass_is_bounded_database_evidence() -> None: + """Database result subclasses must not execute during shape validation.""" + cursor = _RecordingCursor([_HostileLockRow((True,))]) + + with pytest.raises(ReconciliationSingleFlightError) as caught: + with reconciliation_single_flight( + cursor, + "tenant-a", + ReconciliationCandidate("gateway-a", "batch-1"), + ): + pytest.fail("invalid lock evidence must not enter the context body") + + assert caught.value.details == { + "phase": "acquire", + "reason": "invalid_database_result", + } + assert _SECRET_SENTINEL not in _rendered_exception(caught.value) + assert caught.value.__cause__ is None + assert caught.value.__context__ is None + assert len(cursor.calls) == 1