diff --git a/CHANGELOG.md b/CHANGELOG.md index 36b0171..1dcab5d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.0.21] - 2026-09-19 + +### Added + +- **`PostgresTransaction.find_one_and_update`.** Public compare-and-set on the + held asyncpg connection (`SELECT … FOR UPDATE` + upsert) so CAS and + `insert_if_absent` share one transaction handle. Enables atomic work-kernel + transition + outbox units without private `_connection` access. + ## [0.0.20] - 2026-09-18 ### Added diff --git a/jvspatial/db/postgres.py b/jvspatial/db/postgres.py index f42d7ae..7c2eae3 100644 --- a/jvspatial/db/postgres.py +++ b/jvspatial/db/postgres.py @@ -2326,6 +2326,94 @@ async def find( records = finalize_find_results(records, sort=sort, limit=limit) return records + async def find_one_and_update( + self, + collection: str, + query: Dict[str, Any], + update: Dict[str, Any], + upsert: bool = False, + ) -> Optional[Dict[str, Any]]: + """Atomically find + update the first match on this transaction's connection. + + Uses ``SELECT ... FOR UPDATE`` on the held connection so the row lock + participates in the outer transaction. Callers that also + ``insert_if_absent`` on the same handle get one atomic unit of work + on commit/rollback. Does not open a nested pool transaction. + """ + from .database import _normalize_id_query + + await self._db._bootstrap_collection(collection) + col = _safe_collection(collection) + schema = _safe_collection(self._db.schema_name) + + q = _normalize_id_query(query) + translated = translate_query(q) if q else ("", []) + if translated is None: + # Fall back to in-transaction find + match (still on held conn). + candidates = await self.find(collection, q, limit=None) + matched = next( + (r for r in candidates if QueryEngine.match(r, q)), + None, + ) + if matched is None: + if not upsert: + return None + doc: Dict[str, Any] = {} + doc_id = query.get("_id", query.get("id")) + if doc_id is not None: + doc["_id"] = doc_id + doc["id"] = str(doc_id) + QueryEngine.apply_update(doc, update, apply_set_on_insert=True) + else: + doc = dict(matched) + QueryEngine.apply_update(doc, update, apply_set_on_insert=False) + record_id = doc.get("id", doc.get("_id")) + if record_id is not None: + doc["id"] = str(record_id) + await self.save(collection, doc) + return doc + + where_sql, params = translated + clause = f" WHERE {where_sql}" if where_sql else "" + row = await self._connection.fetchrow( + f"SELECT ctid, data FROM {schema}.{col}{clause} " f"LIMIT 1 FOR UPDATE", + *params, + ) + if row is None: + if not upsert: + return None + doc = {} + doc_id = query.get("_id", query.get("id")) + if doc_id is not None: + doc["_id"] = doc_id + doc["id"] = str(doc_id) + QueryEngine.apply_update(doc, update, apply_set_on_insert=True) + else: + doc = self._db._record_from_row(row) + QueryEngine.apply_update(doc, update, apply_set_on_insert=False) + + record_id = doc.get("id", doc.get("_id")) + if record_id is not None: + doc["id"] = str(record_id) + + rec_id, entity, tenant, data_json = self._db._split_payload(doc) + await self._connection.execute( + f""" + INSERT INTO {schema}.{col} (id, entity, tenant_id, data, updated_at) + VALUES ($1, $2, $3, $4::jsonb, NOW()) + ON CONFLICT (id) DO UPDATE SET + entity = EXCLUDED.entity, + tenant_id = EXCLUDED.tenant_id, + data = EXCLUDED.data, + updated_at = NOW() + """, + rec_id, + entity, + tenant, + data_json, + ) + return doc + async def commit(self) -> None: """Commit the wrapped asyncpg transaction. Idempotent.""" if not self.is_active or self.is_committed or self.is_rolled_back: diff --git a/jvspatial/version.py b/jvspatial/version.py index b5c8c4a..f53f5bf 100644 --- a/jvspatial/version.py +++ b/jvspatial/version.py @@ -9,4 +9,4 @@ # - MAJOR: Breaking changes # - MINOR: New features, backward compatible # - PATCH: Bug fixes, backward compatible -__version__ = "0.0.20" +__version__ = "0.0.21" diff --git a/tests/db/test_postgres_integration.py b/tests/db/test_postgres_integration.py index 629a876..b985ef8 100644 --- a/tests/db/test_postgres_integration.py +++ b/tests/db/test_postgres_integration.py @@ -355,6 +355,160 @@ async def test_find_one_and_update_upsert(self, pg_db: "PostgresDB") -> None: assert loaded["context"]["created"] is True +class TestPostgresTransactionAtomicOps: + """CAS + insert_if_absent must share one public transaction handle.""" + + async def test_transaction_find_one_and_update_inc( + self, pg_db: "PostgresDB" + ) -> None: + """Lease-style CAS on the transaction handle persists after commit.""" + await pg_db.save( + "o", + { + "id": "o.WorkItem.1", + "entity": "WorkItem", + "context": {"status": "queued", "lease_fence": 0}, + }, + ) + txn = await pg_db.begin_transaction() + try: + updated = await txn.find_one_and_update( + "o", + { + "id": "o.WorkItem.1", + "context.status": "queued", + "context.lease_fence": 0, + }, + { + "$set": { + "context.status": "running", + "context.lease_fence": 1, + "context.lease_token": "tok-a", + } + }, + ) + assert updated is not None + assert updated["context"]["status"] == "running" + assert updated["context"]["lease_fence"] == 1 + await pg_db.commit_transaction(txn) + except Exception: + await pg_db.rollback_transaction(txn) + raise + + loaded = await pg_db.get("o", "o.WorkItem.1") + assert loaded["context"]["status"] == "running" + assert loaded["context"]["lease_token"] == "tok-a" + + async def test_transaction_cas_and_insert_commit_together( + self, pg_db: "PostgresDB" + ) -> None: + """CAS + outbox insert commit as one unit.""" + await pg_db.save( + "o", + { + "id": "o.WorkItem.2", + "entity": "WorkItem", + "context": {"status": "queued", "lease_fence": 0}, + }, + ) + txn = await pg_db.begin_transaction() + try: + updated = await txn.find_one_and_update( + "o", + {"id": "o.WorkItem.2", "context.status": "queued"}, + {"$set": {"context.status": "running", "context.lease_fence": 1}}, + ) + assert updated is not None + outbox = await txn.insert_if_absent( + "o", + { + "id": "o.WorkOutbox.2", + "entity": "WorkOutboxEntry", + "context": { + "work_item_id": "o.WorkItem.2", + "topic": "work.transitioned", + "status": "pending", + }, + }, + ) + assert outbox.created is True + await pg_db.commit_transaction(txn) + except Exception: + await pg_db.rollback_transaction(txn) + raise + + assert (await pg_db.get("o", "o.WorkItem.2"))["context"]["status"] == "running" + assert await pg_db.get("o", "o.WorkOutbox.2") is not None + + async def test_transaction_cas_and_insert_rollback_together( + self, pg_db: "PostgresDB" + ) -> None: + """CAS + outbox insert roll back together.""" + await pg_db.save( + "o", + { + "id": "o.WorkItem.3", + "entity": "WorkItem", + "context": {"status": "queued", "lease_fence": 0}, + }, + ) + txn = await pg_db.begin_transaction() + updated = await txn.find_one_and_update( + "o", + {"id": "o.WorkItem.3", "context.status": "queued"}, + {"$set": {"context.status": "running", "context.lease_fence": 1}}, + ) + assert updated is not None + outbox = await txn.insert_if_absent( + "o", + { + "id": "o.WorkOutbox.3", + "entity": "WorkOutboxEntry", + "context": {"work_item_id": "o.WorkItem.3", "status": "pending"}, + }, + ) + assert outbox.created is True + await pg_db.rollback_transaction(txn) + + loaded = await pg_db.get("o", "o.WorkItem.3") + assert loaded["context"]["status"] == "queued" + assert loaded["context"]["lease_fence"] == 0 + assert await pg_db.get("o", "o.WorkOutbox.3") is None + + async def test_transaction_stale_cas_returns_none( + self, pg_db: "PostgresDB" + ) -> None: + """Mismatched fence must leave the row unchanged.""" + await pg_db.save( + "o", + { + "id": "o.WorkItem.4", + "entity": "WorkItem", + "context": {"status": "running", "lease_fence": 2}, + }, + ) + txn = await pg_db.begin_transaction() + try: + updated = await txn.find_one_and_update( + "o", + { + "id": "o.WorkItem.4", + "context.status": "running", + "context.lease_fence": 1, + }, + {"$set": {"context.status": "succeeded"}}, + ) + assert updated is None + await pg_db.commit_transaction(txn) + except Exception: + await pg_db.rollback_transaction(txn) + raise + + loaded = await pg_db.get("o", "o.WorkItem.4") + assert loaded["context"]["status"] == "running" + assert loaded["context"]["lease_fence"] == 2 + + # ---- Walker traversal via recursive CTE ------------------------------------ diff --git a/tests/db/test_postgres_transaction_cas.py b/tests/db/test_postgres_transaction_cas.py new file mode 100644 index 0000000..68b7b17 --- /dev/null +++ b/tests/db/test_postgres_transaction_cas.py @@ -0,0 +1,102 @@ +"""Unit tests for PostgresTransaction.find_one_and_update without a live DB.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from jvspatial.db.postgres import PostgresTransaction + + +class _FakeRow(dict): + def __getitem__(self, key: Any) -> Any: + if key == "data": + return dict(self) + if key == "ctid": + return "(0,1)" + return super().__getitem__(key) + + +@pytest.mark.asyncio +async def test_transaction_exposes_find_one_and_update() -> None: + """Public transaction handle must expose compare-and-set.""" + assert hasattr(PostgresTransaction, "find_one_and_update") + assert callable(PostgresTransaction.find_one_and_update) + + +@pytest.mark.asyncio +async def test_transaction_find_one_and_update_uses_held_connection() -> None: + """CAS must run on the held connection with FOR UPDATE, no nested txn.""" + db = MagicMock() + db.schema_name = "public" + db._bootstrap_collection = AsyncMock() + db._split_payload = MagicMock( + return_value=("o.WorkItem.1", "WorkItem", None, '{"id":"o.WorkItem.1"}') + ) + db._record_from_row = MagicMock( + side_effect=lambda row: { + "id": "o.WorkItem.1", + "entity": "WorkItem", + "context": {"status": "queued", "lease_fence": 0}, + } + ) + db._validate_insert_if_absent = MagicMock() + + conn = AsyncMock() + conn.fetchrow = AsyncMock( + return_value=_FakeRow( + id="o.WorkItem.1", + entity="WorkItem", + context={"status": "queued", "lease_fence": 0}, + ) + ) + conn.execute = AsyncMock() + + txn = PostgresTransaction(db, conn, transaction=MagicMock()) + updated = await txn.find_one_and_update( + "o", + {"id": "o.WorkItem.1", "context.status": "queued", "context.lease_fence": 0}, + { + "$set": { + "context.status": "running", + "context.lease_fence": 1, + "context.lease_token": "tok-a", + } + }, + ) + + assert updated is not None + assert updated["context"]["status"] == "running" + assert updated["context"]["lease_fence"] == 1 + assert updated["context"]["lease_token"] == "tok-a" + assert conn.fetchrow.await_count == 1 + sql = conn.fetchrow.await_args.args[0] + assert "FOR UPDATE" in sql + assert conn.execute.await_count == 1 + assert conn.transaction.await_count == 0 + assert conn.transaction.call_count == 0 + + +@pytest.mark.asyncio +async def test_transaction_find_one_and_update_stale_match_returns_none() -> None: + """Stale expected-state query must not write.""" + db = MagicMock() + db.schema_name = "public" + db._bootstrap_collection = AsyncMock() + db._record_from_row = MagicMock() + + conn = AsyncMock() + conn.fetchrow = AsyncMock(return_value=None) + conn.execute = AsyncMock() + + txn = PostgresTransaction(db, conn, transaction=MagicMock()) + updated = await txn.find_one_and_update( + "o", + {"id": "o.WorkItem.1", "context.lease_fence": 1}, + {"$set": {"context.status": "succeeded"}}, + ) + + assert updated is None + conn.execute.assert_not_awaited()