mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 11:19:54 +02:00
98935e1ffd
Signed-off-by: Tyler Ball <tyleraball@gmail.com> Co-authored-by: Phoenix Logan <plogan@chanzuckerberg.com> Co-authored-by: Tyler Ball <2481463+tyler-ball@users.noreply.github.com>
22 lines
650 B
Python
22 lines
650 B
Python
"""Shared utility functions for the Postgres checkpoint & storage classes."""
|
|
|
|
from contextlib import contextmanager
|
|
from typing import Iterator, Union
|
|
|
|
from psycopg import Connection
|
|
from psycopg.rows import DictRow
|
|
from psycopg_pool import ConnectionPool
|
|
|
|
Conn = Union[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)}")
|