wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
22 lines
646 B
Python
22 lines
646 B
Python
"""Shared utility functions for the Postgres checkpoint & storage classes."""
|
|
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
|
|
from psycopg import Connection
|
|
from psycopg.rows import DictRow
|
|
from psycopg_pool import ConnectionPool
|
|
|
|
Conn = Connection[DictRow] | ConnectionPool[Connection[DictRow]]
|
|
|
|
|
|
@contextmanager
|
|
def get_connection(conn: Conn) -> Iterator[Connection[DictRow]]:
|
|
if isinstance(conn, Connection):
|
|
yield conn
|
|
elif isinstance(conn, ConnectionPool):
|
|
with conn.connection() as conn:
|
|
yield conn
|
|
else:
|
|
raise TypeError(f"Invalid connection type: {type(conn)}")
|