feat(db): PostgresTransaction.find_one_and_update for work-kernel CAS

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

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:

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"

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 ------------------------------------
 
 

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()