From 641ae8f4f77c7ec176113d69fb9435a038118231 Mon Sep 17 00:00:00 2001 From: 35C4n0r Date: Sun, 24 May 2026 09:01:34 +0530 Subject: [PATCH] fix(coordination): split AdvisoryLock into sync and async variants - the single AdvisoryLock class used sync with for try_acquire/release and async with for atry_acquire/arelease, but PostgresSaver._cursor() is a sync context manager and AsyncPostgresSaver._cursor() is async - mixing them caused TypeError at runtime - split into SyncAdvisoryLock (PostgresSaver) and AsyncAdvisoryLock (AsyncPostgresSaver); each class only uses the cursor protocol that matches its checkpointer - SyncAdvisoryLock.atry_acquire/arelease delegate to the sync methods so it remains usable from asweep() when paired with PostgresSaver - AsyncAdvisoryLock.try_acquire/release raise NotImplementedError since AsyncPostgresSaver has no sync cursor path - update get_advisory_lock() return type and sweeper.py annotation --- .../_coordination.py | 53 ++++++++++++++----- .../sweeper.py | 4 +- 2 files changed, 42 insertions(+), 15 deletions(-) diff --git a/src/langgraph_ephemeral_checkpointer/_coordination.py b/src/langgraph_ephemeral_checkpointer/_coordination.py index 68682c7..b2a516c 100644 --- a/src/langgraph_ephemeral_checkpointer/_coordination.py +++ b/src/langgraph_ephemeral_checkpointer/_coordination.py @@ -12,43 +12,70 @@ signed=True, ) +_ACQUIRE_SQL = "SELECT pg_try_advisory_lock(%s) AS acquired" +_RELEASE_SQL = "SELECT pg_advisory_unlock(%s)" -class AdvisoryLock: - """PostgreSQL session-level advisory lock. Released automatically when the connection closes.""" + +class SyncAdvisoryLock: + """Advisory lock for PostgresSaver (sync). Uses a sync cursor.""" def __init__(self, checkpointer) -> None: self._checkpointer = checkpointer def try_acquire(self) -> bool: with self._checkpointer._cursor() as cur: - cur.execute("SELECT pg_try_advisory_lock(%s) AS acquired", (_LOCK_KEY,)) + cur.execute(_ACQUIRE_SQL, (_LOCK_KEY,)) return bool(cur.fetchone()["acquired"]) def release(self) -> None: with self._checkpointer._cursor() as cur: - cur.execute("SELECT pg_advisory_unlock(%s)", (_LOCK_KEY,)) + cur.execute(_RELEASE_SQL, (_LOCK_KEY,)) + + async def atry_acquire(self) -> bool: + return self.try_acquire() + + async def arelease(self) -> None: + self.release() + + +class AsyncAdvisoryLock: + """Advisory lock for AsyncPostgresSaver (async). Uses an async cursor.""" + + def __init__(self, checkpointer) -> None: + self._checkpointer = checkpointer + + def try_acquire(self) -> bool: + raise NotImplementedError("Use atry_acquire() with AsyncPostgresSaver") + + def release(self) -> None: + raise NotImplementedError("Use arelease() with AsyncPostgresSaver") async def atry_acquire(self) -> bool: async with self._checkpointer._cursor() as cur: - await cur.execute( - "SELECT pg_try_advisory_lock(%s) AS acquired", (_LOCK_KEY,) - ) + await cur.execute(_ACQUIRE_SQL, (_LOCK_KEY,)) row = await cur.fetchone() return bool(row["acquired"]) async def arelease(self) -> None: async with self._checkpointer._cursor() as cur: - await cur.execute("SELECT pg_advisory_unlock(%s)", (_LOCK_KEY,)) + await cur.execute(_RELEASE_SQL, (_LOCK_KEY,)) + + +# Union type for type annotations in sweeper.py +AdvisoryLock = SyncAdvisoryLock | AsyncAdvisoryLock + -def get_advisory_lock(checkpointer, enable: bool) -> AdvisoryLock | None: - """Return an AdvisoryLock for the given checkpointer, or None. +def get_advisory_lock(checkpointer, enable: bool) -> SyncAdvisoryLock | AsyncAdvisoryLock | None: + """Return an advisory lock for the given checkpointer, or None. Args: checkpointer: The LangGraph checkpointer to lock against. enable: If False, always returns None without checking the backend. Returns: - AdvisoryLock if the backend is PostgreSQL and enable is True, else None. + SyncAdvisoryLock for PostgresSaver, AsyncAdvisoryLock for + AsyncPostgresSaver, or None if the backend is unsupported or + enable is False. """ if not enable: return None @@ -56,14 +83,14 @@ def get_advisory_lock(checkpointer, enable: bool) -> AdvisoryLock | None: try: from langgraph.checkpoint.postgres import PostgresSaver if isinstance(checkpointer, PostgresSaver): - return AdvisoryLock(checkpointer) + return SyncAdvisoryLock(checkpointer) except ImportError: pass try: from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver if isinstance(checkpointer, AsyncPostgresSaver): - return AdvisoryLock(checkpointer) + return AsyncAdvisoryLock(checkpointer) except ImportError: pass diff --git a/src/langgraph_ephemeral_checkpointer/sweeper.py b/src/langgraph_ephemeral_checkpointer/sweeper.py index 07af6f4..793075e 100644 --- a/src/langgraph_ephemeral_checkpointer/sweeper.py +++ b/src/langgraph_ephemeral_checkpointer/sweeper.py @@ -5,7 +5,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver from . import _strategies -from ._coordination import AdvisoryLock, get_advisory_lock +from ._coordination import AsyncAdvisoryLock, SyncAdvisoryLock, get_advisory_lock from ._strategies import Strategy, ThreadTimestamps from ._uuid6 import unix_to_uuid6, uuid6_to_unix from .policy import TTLPolicy @@ -62,7 +62,7 @@ def __init__( self._on_before_delete = on_before_delete self._on_sweep_complete = on_sweep_complete self._strategy: Strategy = _strategies.detect(checkpointer) if _strategy is None else _strategy - self._advisory_lock: AdvisoryLock | None = get_advisory_lock(checkpointer, enable_coordination) + self._advisory_lock: SyncAdvisoryLock | AsyncAdvisoryLock | None = get_advisory_lock(checkpointer, enable_coordination) self._task: asyncio.Task | None = None self._stop_event: asyncio.Event | None = None self._interval_seconds: int = 300