"""Psycopg database configuration with direct field-based configuration."""
from contextlib import suppress
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypedDict, cast
from mypy_extensions import mypyc_attr
from psycopg.adapt import AdaptersMap
from psycopg.types.json import set_json_dumps, set_json_loads
from typing_extensions import NotRequired, Self
import sqlspec.adapters.psycopg._typing as _psycopg_typing
from sqlspec.adapters.psycopg._typing import (
PsycopgAsyncConnection,
PsycopgAsyncCursor,
PsycopgAsyncSessionContext,
PsycopgConnection,
PsycopgSyncConnection,
PsycopgSyncCursor,
PsycopgSyncSessionContext,
)
from sqlspec.adapters.psycopg._typing import PsycopgAsyncConnectionPool as AsyncConnectionPool
from sqlspec.adapters.psycopg._typing import PsycopgAsyncNullConnectionPool as AsyncNullConnectionPool
from sqlspec.adapters.psycopg._typing import PsycopgConnectionPool as ConnectionPool
from sqlspec.adapters.psycopg._typing import PsycopgNullConnectionPool as NullConnectionPool
from sqlspec.adapters.psycopg.core import apply_driver_features, default_statement_config
from sqlspec.adapters.psycopg.driver import (
PsycopgAsyncDriver,
PsycopgAsyncExceptionHandler,
PsycopgSyncDriver,
PsycopgSyncExceptionHandler,
)
from sqlspec.adapters.psycopg.type_converter import register_pgvector_async, register_pgvector_sync
from sqlspec.config import AsyncDatabaseConfig, ExtensionConfigs, SyncDatabaseConfig
from sqlspec.core.capabilities import TypeCoercionCapabilities
from sqlspec.core.config_runtime import (
build_postgres_extension_probe_names,
is_postgres_extension_active,
resolve_postgres_extension_state,
resolve_runtime_statement_config,
)
from sqlspec.driver import (
AsyncPoolConnectionContext,
AsyncPoolSessionFactory,
SyncPoolConnectionContext,
SyncPoolSessionFactory,
)
from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError
from sqlspec.extensions.events import EventRuntimeHints
from sqlspec.typing import ALLOYDB_CONNECTOR_INSTALLED
from sqlspec.utils.config_tools import normalize_connection_config
from sqlspec.utils.serializers import from_json, to_json
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable, Mapping
from types import TracebackType
from sqlspec.adapters.psycopg._typing import PsycopgAdaptContext as AdaptContext
from sqlspec.adapters.psycopg._typing import PsycopgAsyncConnectFailedCB as AsyncConnectFailedCB
from sqlspec.adapters.psycopg._typing import PsycopgAsyncConnectionCB as AsyncConnectionCB
from sqlspec.adapters.psycopg._typing import PsycopgAsyncRowFactory as AsyncRowFactory
from sqlspec.adapters.psycopg._typing import PsycopgConnectFailedCB as ConnectFailedCB
from sqlspec.adapters.psycopg._typing import PsycopgConnection as Connection
from sqlspec.adapters.psycopg._typing import PsycopgConnectionCB as ConnectionCB
from sqlspec.adapters.psycopg._typing import PsycopgCursor as Cursor
from sqlspec.adapters.psycopg._typing import PsycopgNativeAsyncConnection as AsyncConnection
from sqlspec.adapters.psycopg._typing import PsycopgNativeAsyncCursor as AsyncCursor
from sqlspec.adapters.psycopg._typing import PsycopgRowFactory as RowFactory
from sqlspec.core import StatementConfig
__all__ = (
"PsycopgAsyncConfig",
"PsycopgAsyncCursor",
"PsycopgConnectionParams",
"PsycopgDriverFeatures",
"PsycopgPoolParams",
"PsycopgSyncConfig",
"PsycopgSyncCursor",
"build_connection_config",
)
PsycopgSSLMode = Literal["disable", "allow", "prefer", "require", "verify-ca", "verify-full"]
_ALLOYDB_DIRECT_CONNECTION_KEYS = ("conninfo", "host", "hostaddr", "port", "user", "password", "dbname", "db")
_POOL_ONLY_CONFIG_KEYS: Final[frozenset[str]] = frozenset({
"check",
"close_returns",
"configure",
"max_idle",
"max_lifetime",
"max_size",
"max_waiting",
"min_size",
"name",
"num_workers",
"null_pool",
"open",
"reconnect_failed",
"reconnect_timeout",
"reset",
"timeout",
})
_ALLOYDB_CONNECTOR_PACKAGE = "google-cloud-alloydb-connector"
[docs]
class PsycopgConnectionParams(TypedDict):
"""Psycopg connection parameters."""
conninfo: NotRequired[str]
host: NotRequired[str]
port: NotRequired[int]
user: NotRequired[str]
password: NotRequired[str]
dbname: NotRequired[str]
connect_timeout: NotRequired[int]
options: NotRequired[str]
application_name: NotRequired[str]
sslmode: NotRequired[PsycopgSSLMode]
sslcert: NotRequired[str]
sslkey: NotRequired[str]
sslrootcert: NotRequired[str]
autocommit: NotRequired[bool]
prepare_threshold: NotRequired[int | None]
context: NotRequired["AdaptContext | None"]
row_factory: NotRequired["RowFactory[Any] | AsyncRowFactory[Any] | None"]
cursor_factory: NotRequired["type[Cursor[Any] | AsyncCursor[Any]] | None"]
extra: NotRequired["dict[str, Any]"]
[docs]
class PsycopgPoolParams(PsycopgConnectionParams):
"""Psycopg pool parameters."""
min_size: NotRequired[int]
max_size: NotRequired[int]
connection_class: NotRequired["type[Connection[Any] | AsyncConnection[Any]]"]
name: NotRequired[str]
timeout: NotRequired[float]
max_waiting: NotRequired[int]
max_lifetime: NotRequired[float]
max_idle: NotRequired[float]
reconnect_timeout: NotRequired[float]
num_workers: NotRequired[int]
open: NotRequired[bool | None]
configure: NotRequired["ConnectionCB[PsycopgSyncConnection] | AsyncConnectionCB[PsycopgAsyncConnection] | None"]
check: NotRequired["ConnectionCB[PsycopgSyncConnection] | AsyncConnectionCB[PsycopgAsyncConnection] | None"]
reset: NotRequired["ConnectionCB[PsycopgSyncConnection] | AsyncConnectionCB[PsycopgAsyncConnection] | None"]
close_returns: NotRequired[bool]
reconnect_failed: NotRequired["ConnectFailedCB | AsyncConnectFailedCB | None"]
kwargs: NotRequired["dict[str, Any]"]
null_pool: NotRequired[bool]
[docs]
class PsycopgDriverFeatures(TypedDict):
"""Psycopg driver feature flags.
enable_pgvector: Enable automatic pgvector extension support for vector similarity search.
Requires pgvector-python package (pip install pgvector) and PostgreSQL with pgvector extension.
Defaults to True when pgvector-python is installed.
Provides automatic conversion between Python objects and PostgreSQL vector types.
Enables vector similarity operations and index support.
Set to False to disable pgvector support even when package is available.
enable_paradedb: Enable ParadeDB (pg_search) extension detection.
When enabled and the pg_search extension is detected, the SQL dialect
switches to "paradedb" which supports search operators (@@@, &&&, etc.)
and inherits all pgvector distance operators.
Defaults to True. Independent of enable_pgvector.
enable_pg_textsearch: Enable pg_textsearch extension detection for BM25 search.
When enabled and the pg_textsearch extension is detected, the SQL dialect
switches to "pg_textsearch" which supports the BM25 <@> relevance ranking operator.
Defaults to True.
json_serializer: Custom JSON serializer for StatementConfig parameter handling.
json_deserializer: Custom JSON deserializer reference stored alongside the serializer for parity with asyncpg.
on_connection_create: Callback executed when a connection is created/acquired from the pool.
Receives the raw psycopg connection for low-level driver configuration.
Runs after internal setup (pgvector registration).
For sync config: Callable[[PsycopgSyncConnection], None]
For async config: Callable[[PsycopgAsyncConnection], Awaitable[None]]
enable_events: Enable database event channel support.
Defaults to True when extension_config["events"] is configured.
Provides pub/sub capabilities via LISTEN/NOTIFY or table-backed fallback.
Requires extension_config["events"] for migration setup when using poll_queue or notify_queue.
events_backend: Event channel backend selection.
Options: "notify", "notify_queue", "poll_queue"
- "notify": Transient PostgreSQL LISTEN/NOTIFY with no replay or retry
- "notify_queue": Durable queue plus a PostgreSQL notification wakeup hint
- "poll_queue": Durable queue discovered by polling
Defaults to "notify".
enable_alloydb: Enable sync-only Google AlloyDB connector integration.
Requires google-cloud-alloydb-connector. Defaults to False.
alloydb_instance_uri: AlloyDB instance URI in projects/... resource path form.
Required when enable_alloydb is True.
enable_alloydb_iam_auth: Enable AlloyDB IAM database authentication for sync connector connections.
Defaults to False.
alloydb_ip_type: AlloyDB connector IP type. Defaults to PRIVATE.
null_pool: Enable NullConnectionPool / AsyncNullConnectionPool for serverless / PgBouncer environments.
"""
enable_pgvector: NotRequired[bool]
enable_paradedb: NotRequired[bool]
enable_pg_textsearch: NotRequired[bool]
json_serializer: NotRequired["Callable[[Any], str]"]
json_deserializer: NotRequired["Callable[[str], Any]"]
on_connection_create: NotRequired["Callable[..., Any]"]
enable_events: NotRequired[bool]
events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]]
enable_alloydb: NotRequired[bool]
alloydb_instance_uri: NotRequired[str]
enable_alloydb_iam_auth: NotRequired[bool]
alloydb_ip_type: NotRequired[str]
null_pool: NotRequired[bool]
def build_connection_config(connection_config: "PsycopgPoolParams | Mapping[str, Any] | None") -> dict[str, Any]:
"""Build normalized psycopg connection configuration, resolving aliases for libpq compatibility.
Maps connection string aliases (dsn, url, connection_string) to conninfo, database aliases
(database, db) to dbname, and user aliases (username) to user, while discarding redundant keys
that libpq rejects.
"""
config = normalize_connection_config(connection_config)
conninfo = (
config.pop("conninfo", None)
or config.pop("dsn", None)
or config.pop("url", None)
or config.pop("connection_string", None)
)
if conninfo is not None:
config["conninfo"] = conninfo
dbname = config.pop("dbname", None) or config.pop("database", None) or config.pop("db", None)
if dbname is not None:
config["dbname"] = dbname
user = config.pop("user", None) or config.pop("username", None)
if user is not None:
config["user"] = user
return config
class PsycopgSyncConnectionContext(SyncPoolConnectionContext):
"""Context manager for Psycopg connections."""
__slots__ = ()
def __init__(self, config: "PsycopgSyncConfig") -> None:
super().__init__(config)
def __enter__(self) -> "PsycopgSyncConnection":
self._ctx = self._config.provide_pool().connection()
return cast("PsycopgSyncConnection", self._ctx.__enter__())
def __exit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> bool | None:
if self._ctx is None:
return None
return cast("bool | None", self._ctx.__exit__(exc_type, exc_val, exc_tb))
class _PsycopgSyncSessionConnectionHandler(SyncPoolSessionFactory):
__slots__ = ()
def acquire_connection(self) -> "PsycopgSyncConnection":
self._ctx = self._config.provide_pool().connection()
return cast("PsycopgSyncConnection", self._ctx.__enter__())
def release_connection(self, _conn: "PsycopgSyncConnection", **kwargs: Any) -> None:
if self._ctx is not None:
self._ctx.__exit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb"))
self._ctx = None
[docs]
@mypyc_attr(native_class=False)
class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool, PsycopgSyncDriver]):
"""Configuration for Psycopg synchronous database connections with direct field-based configuration."""
driver_type: "ClassVar[type[PsycopgSyncDriver]]" = PsycopgSyncDriver
connection_type: "ClassVar[type[PsycopgSyncConnection]]" = PsycopgSyncConnection
supports_transactional_ddl: "ClassVar[bool]" = True
supports_migration_schemas: "ClassVar[bool]" = True
supports_native_arrow_export: "ClassVar[bool]" = True
supports_native_arrow_import: "ClassVar[bool]" = True
supports_native_parquet_export: "ClassVar[bool]" = True
supports_native_parquet_import: "ClassVar[bool]" = True
supports_native_row_streaming: "ClassVar[bool]" = True
type_coercion_capabilities: "ClassVar[TypeCoercionCapabilities]" = TypeCoercionCapabilities(
datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=True, uuid_binding="native"
)
_connection_context_class: "ClassVar[type[PsycopgSyncConnectionContext]]" = PsycopgSyncConnectionContext
_session_factory_class: "ClassVar[type[_PsycopgSyncSessionConnectionHandler]]" = (
_PsycopgSyncSessionConnectionHandler
)
_session_context_class: "ClassVar[type[PsycopgSyncSessionContext]]" = PsycopgSyncSessionContext
_default_statement_config = default_statement_config
[docs]
def __init__(
self,
*,
connection_config: "PsycopgPoolParams | dict[str, Any] | None" = None,
connection_instance: "ConnectionPool | None" = None,
migration_config: "dict[str, Any] | None" = None,
statement_config: "StatementConfig | None" = None,
driver_features: "PsycopgDriverFeatures | dict[str, Any] | None" = None,
bind_key: "str | None" = None,
extension_config: "ExtensionConfigs | None" = None,
**kwargs: Any,
) -> None:
"""Initialize Psycopg synchronous configuration.
Args:
connection_config: Connection and pool configuration parameters (TypedDict or dict)
connection_instance: Existing pool instance to use
migration_config: Migration configuration
statement_config: Default SQL statement configuration
driver_features: Optional driver feature configuration
bind_key: Optional unique identifier for this configuration
extension_config: Extension-specific configuration
**kwargs: Additional keyword arguments
"""
connection_config = build_connection_config(connection_config)
statement_config = statement_config or default_statement_config
statement_config, driver_features = apply_driver_features(statement_config, driver_features)
features_dict = dict(driver_features) if driver_features else {}
self._user_connection_hook: Callable[[PsycopgSyncConnection], None] | None = features_dict.pop(
"on_connection_create", None
)
features_dict.setdefault("enable_alloydb", False)
features_dict.setdefault("enable_alloydb_iam_auth", False)
features_dict.setdefault("alloydb_ip_type", "PRIVATE")
self._pgvector_available: bool | None = None
self._paradedb_available: bool | None = None
self._pg_textsearch_available: bool | None = None
self._pg_textsearch_probe_error: Exception | None = None
self._alloydb_connector: Any | None = None
super().__init__(
connection_config=connection_config,
connection_instance=connection_instance,
migration_config=migration_config,
statement_config=statement_config,
driver_features=features_dict,
bind_key=bind_key,
extension_config=extension_config,
**kwargs,
)
self._validate_alloydb_config()
def _validate_alloydb_config(self) -> None:
"""Validate sync-only AlloyDB connector configuration."""
if not self.driver_features.get("enable_alloydb", False):
return
if not ALLOYDB_CONNECTOR_INSTALLED:
raise MissingDependencyError(package=_ALLOYDB_CONNECTOR_PACKAGE, install_package=_ALLOYDB_CONNECTOR_PACKAGE)
instance_uri = self.driver_features.get("alloydb_instance_uri")
if not instance_uri:
msg = (
"alloydb_instance_uri required when enable_alloydb is True. "
"Format: 'projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE'"
)
raise ImproperConfigurationError(msg)
if not str(instance_uri).startswith("projects/"):
msg = (
f"Invalid AlloyDB instance URI format: {instance_uri}. "
"Expected format: 'projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE'"
)
raise ImproperConfigurationError(msg)
def _setup_alloydb_connector(
self, config: "dict[str, Any]", pool_parameters: "dict[str, Any] | None" = None
) -> None:
"""Setup AlloyDB connector and configure psycopg-pool connection_class."""
if self._alloydb_connector is None:
self._alloydb_connector = _psycopg_typing.PsycopgAlloydbConnector()
user = config.get("user")
password = config.get("password")
database = config.get("dbname") or config.get("db")
for key in _ALLOYDB_DIRECT_CONNECTION_KEYS:
config.pop(key, None)
if pool_parameters is not None:
pool_parameters["connection_class"] = _make_alloydb_connection_class(
connector=self._alloydb_connector,
instance_uri=str(self.driver_features["alloydb_instance_uri"]),
enable_iam_auth=bool(self.driver_features.get("enable_alloydb_iam_auth", False)),
ip_type=str(self.driver_features.get("alloydb_ip_type", "PRIVATE")),
user=cast("str | None", user),
password=cast("str | None", password),
database=cast("str | None", database),
)
def _create_pool(self) -> "ConnectionPool":
"""Create the actual connection pool."""
all_config = dict(self.connection_config)
all_config.pop("null_pool", None)
pool_parameters = {
"connection_class": all_config.pop("connection_class", None),
"min_size": all_config.pop("min_size", 4),
"max_size": all_config.pop("max_size", None),
"name": all_config.pop("name", None),
"timeout": all_config.pop("timeout", 30.0),
"max_waiting": all_config.pop("max_waiting", 0),
"max_lifetime": all_config.pop("max_lifetime", 3600.0),
"max_idle": all_config.pop("max_idle", 600.0),
"reconnect_timeout": all_config.pop("reconnect_timeout", 300.0),
"reconnect_failed": all_config.pop("reconnect_failed", None),
"num_workers": all_config.pop("num_workers", 3),
"open": all_config.pop("open", True),
"check": all_config.pop("check", None),
"reset": all_config.pop("reset", None),
"close_returns": all_config.pop("close_returns", False),
}
pool_parameters["configure"] = all_config.pop("configure", self._configure_connection)
pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None}
conninfo = all_config.pop("conninfo", None)
kwargs = all_config.pop("kwargs", {})
all_config.update(kwargs)
if self.driver_features.get("enable_alloydb", False):
if conninfo is not None:
all_config["conninfo"] = conninfo
self._setup_alloydb_connector(all_config, pool_parameters)
conninfo = None
is_null_pool = bool(self.connection_config.get("null_pool") or self.driver_features.get("null_pool"))
pool_cls = NullConnectionPool if is_null_pool else ConnectionPool
if is_null_pool:
pool_parameters.pop("min_size", None)
if conninfo:
pool = pool_cls(conninfo, kwargs=all_config, **pool_parameters)
else:
pool = pool_cls("", kwargs=all_config, **pool_parameters)
return pool
def _configure_connection(self, conn: "PsycopgSyncConnection") -> None:
autocommit_setting = self.connection_config.get("autocommit")
if autocommit_setting is not None:
conn.autocommit = autocommit_setting
serializer = self.driver_features.get("json_serializer", to_json)
deserializer = self.driver_features.get("json_deserializer", from_json)
if isinstance(getattr(conn, "adapters", None), AdaptersMap):
set_json_dumps(serializer, conn)
set_json_loads(deserializer, conn)
if self._pgvector_available is None:
detected_extensions: set[str] = set()
extensions = build_postgres_extension_probe_names(self.driver_features)
if extensions:
try:
cursor = conn.execute(
"SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,)
)
results = cursor.fetchall()
detected_extensions = {r[0] for r in results} # type: ignore[index]
except Exception as exc:
detected_extensions = set()
self._pg_textsearch_probe_error = exc
self.statement_config, self._pgvector_available, self._paradedb_available = (
resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions)
)
self._pg_textsearch_available = is_postgres_extension_active(self.driver_features, "pg_textsearch")
if self._pgvector_available:
register_pgvector_sync(conn)
if not conn.autocommit:
conn.rollback()
if self._user_connection_hook is not None:
self._user_connection_hook(conn)
@property
def pg_textsearch_available(self) -> bool:
"""Return True if the pg_textsearch extension is available."""
return bool(self._pg_textsearch_available)
def _ensure_pg_textsearch_available(self) -> None:
if self._pg_textsearch_available:
return
msg = "ADK memory enable_bm25 requires the pg_textsearch PostgreSQL extension"
raise ImproperConfigurationError(msg) from self._pg_textsearch_probe_error
def _close_pool(self) -> None:
"""Close the actual connection pool and cleanup connectors."""
try:
if self.connection_instance:
self.connection_instance.close()
finally:
self.connection_instance = None
if self._alloydb_connector is not None:
self._alloydb_connector.close()
self._alloydb_connector = None
def _connection_kwargs(self) -> "tuple[type[Any], str, dict[str, Any]]":
"""Split the pooled configuration into standalone connect arguments.
Returns:
The connection class, the conninfo string, and the connect kwargs.
"""
all_config = {k: v for k, v in self.connection_config.items() if k not in _POOL_ONLY_CONFIG_KEYS}
connection_class = all_config.pop("connection_class", None) or PsycopgConnection
conninfo = all_config.pop("conninfo", None)
all_config.update(all_config.pop("kwargs", {}))
if self.driver_features.get("enable_alloydb", False):
if conninfo is not None:
all_config["conninfo"] = conninfo
alloydb_parameters: dict[str, Any] = {}
self._setup_alloydb_connector(all_config, alloydb_parameters)
connection_class = alloydb_parameters["connection_class"]
conninfo = None
return cast("type[Any]", connection_class), str(conninfo or ""), all_config
[docs]
def create_connection(self) -> "PsycopgSyncConnection":
"""Open a standalone connection owned by the caller.
The connection is not bound to the pool, so closing it actually closes
it and consumes no pool slot.
Returns:
A psycopg Connection instance.
"""
connection_class, conninfo, connection_kwargs = self._connection_kwargs()
connection = connection_class.connect(conninfo, **connection_kwargs)
configure = self.connection_config.get("configure", self._configure_connection)
try:
configure(cast("PsycopgSyncConnection", connection))
except BaseException:
with suppress(Exception):
connection.close()
raise
return cast("PsycopgSyncConnection", connection)
[docs]
def provide_session(
self, *_args: Any, statement_config: "StatementConfig | None" = None, **_kwargs: Any
) -> "PsycopgSyncSessionContext":
"""Provide a driver session context manager.
Args:
*_args: Additional arguments.
statement_config: Optional statement configuration override.
**_kwargs: Additional keyword arguments.
Returns:
A PsycopgSyncDriver session context manager.
"""
handler = _PsycopgSyncSessionConnectionHandler(self)
return PsycopgSyncSessionContext(
acquire_connection=handler.acquire_connection,
release_connection=handler.release_connection,
statement_config=statement_config
or (lambda: resolve_runtime_statement_config(None, self.statement_config, default_statement_config)),
driver_features=self.driver_features,
prepare_driver=self._prepare_driver,
)
[docs]
def provide_pool(self, *args: Any, **kwargs: Any) -> "ConnectionPool":
"""Provide pool instance.
Returns:
The connection pool.
"""
if not self.connection_instance:
self.connection_instance = self.create_pool()
return self.connection_instance
[docs]
def get_signature_namespace(self) -> "dict[str, Any]":
"""Get the signature namespace for Psycopg types.
This provides all Psycopg-specific types that Litestar needs to recognize
to avoid serialization attempts.
Returns:
Dictionary mapping type names to types.
"""
namespace = super().get_signature_namespace()
namespace.update({
"PsycopgConnectionParams": PsycopgConnectionParams,
"PsycopgPoolParams": PsycopgPoolParams,
"PsycopgSyncConnectionContext": PsycopgSyncConnectionContext,
"PsycopgSyncConnection": PsycopgSyncConnection,
"PsycopgSyncCursor": PsycopgSyncCursor,
"PsycopgSyncDriver": PsycopgSyncDriver,
"PsycopgSyncExceptionHandler": PsycopgSyncExceptionHandler,
"PsycopgSyncSessionContext": PsycopgSyncSessionContext,
})
return namespace
[docs]
def get_event_runtime_hints(self) -> "EventRuntimeHints":
"""Return polling defaults for PostgreSQL queue fallback."""
return EventRuntimeHints(poll_interval=0.5, select_for_update=True, skip_locked=True)
class PsycopgAsyncConnectionContext(AsyncPoolConnectionContext):
"""Async context manager for Psycopg connections."""
__slots__ = ("_ctx",)
def __init__(self, config: "PsycopgAsyncConfig") -> None:
super().__init__(config)
self._ctx: Any = None
async def __aenter__(self) -> "PsycopgAsyncConnection":
if self._config.connection_instance is None:
self._config.connection_instance = await self._config.create_pool()
if self._config.connection_instance:
self._ctx = self._config.connection_instance.connection()
return cast("PsycopgAsyncConnection", await self._ctx.__aenter__())
msg = "Connection pool not initialized"
raise ImproperConfigurationError(msg)
async def __aexit__(
self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
) -> bool | None:
if self._ctx:
return cast("bool | None", await self._ctx.__aexit__(exc_type, exc_val, exc_tb))
return None
class _PsycopgAsyncSessionConnectionHandler(AsyncPoolSessionFactory):
__slots__ = ("_ctx",)
def __init__(self, config: "PsycopgAsyncConfig") -> None:
super().__init__(config)
self._ctx: Any = None
async def acquire_connection(self) -> "PsycopgAsyncConnection":
if self._config.connection_instance is None:
self._config.connection_instance = await self._config.create_pool()
self._ctx = self._config.connection_instance.connection()
return cast("PsycopgAsyncConnection", await self._ctx.__aenter__())
async def release_connection(self, _conn: "PsycopgAsyncConnection", **kwargs: Any) -> None:
if self._ctx is None:
return
await self._ctx.__aexit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb"))
self._ctx = None
[docs]
@mypyc_attr(native_class=False)
class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnectionPool, PsycopgAsyncDriver]):
"""Configuration for Psycopg asynchronous database connections with direct field-based configuration."""
driver_type: ClassVar[type[PsycopgAsyncDriver]] = PsycopgAsyncDriver
connection_type: "ClassVar[type[PsycopgAsyncConnection]]" = PsycopgAsyncConnection
supports_transactional_ddl: "ClassVar[bool]" = True
supports_migration_schemas: "ClassVar[bool]" = True
supports_native_arrow_export: ClassVar[bool] = True
supports_native_arrow_import: ClassVar[bool] = True
supports_native_parquet_export: ClassVar[bool] = True
supports_native_parquet_import: ClassVar[bool] = True
supports_native_row_streaming: ClassVar[bool] = True
type_coercion_capabilities: "ClassVar[TypeCoercionCapabilities]" = TypeCoercionCapabilities(
datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=True, uuid_binding="native"
)
_connection_context_class: "ClassVar[type[PsycopgAsyncConnectionContext]]" = PsycopgAsyncConnectionContext
_session_factory_class: "ClassVar[type[_PsycopgAsyncSessionConnectionHandler]]" = (
_PsycopgAsyncSessionConnectionHandler
)
_session_context_class: "ClassVar[type[PsycopgAsyncSessionContext]]" = PsycopgAsyncSessionContext
_default_statement_config = default_statement_config
[docs]
def __init__(
self,
*,
connection_config: "PsycopgPoolParams | dict[str, Any] | None" = None,
connection_instance: "AsyncConnectionPool | None" = None,
migration_config: "dict[str, Any] | None" = None,
statement_config: "StatementConfig | None" = None,
driver_features: "PsycopgDriverFeatures | dict[str, Any] | None" = None,
bind_key: "str | None" = None,
extension_config: "ExtensionConfigs | None" = None,
**kwargs: Any,
) -> None:
"""Initialize Psycopg asynchronous configuration.
Args:
connection_config: Connection and pool configuration parameters (TypedDict or dict)
connection_instance: Existing pool instance to use
migration_config: Migration configuration
statement_config: Default SQL statement configuration
driver_features: Optional driver feature configuration
bind_key: Optional unique identifier for this configuration
extension_config: Extension-specific configuration
**kwargs: Additional keyword arguments
"""
connection_config = build_connection_config(connection_config)
statement_config = statement_config or default_statement_config
statement_config, driver_features = apply_driver_features(statement_config, driver_features)
features_dict = dict(driver_features) if driver_features else {}
self._user_connection_hook: Callable[[PsycopgAsyncConnection], Awaitable[None]] | None = features_dict.pop(
"on_connection_create", None
)
self._pgvector_available: bool | None = None
self._paradedb_available: bool | None = None
self._pg_textsearch_available: bool | None = None
self._pg_textsearch_probe_error: Exception | None = None
super().__init__(
connection_config=connection_config,
connection_instance=connection_instance,
migration_config=migration_config,
statement_config=statement_config,
driver_features=features_dict,
bind_key=bind_key,
extension_config=extension_config,
**kwargs,
)
async def _create_pool(self) -> "AsyncConnectionPool":
"""Create the actual async connection pool."""
all_config = dict(self.connection_config)
all_config.pop("null_pool", None)
pool_parameters = {
"connection_class": all_config.pop("connection_class", None),
"min_size": all_config.pop("min_size", 4),
"max_size": all_config.pop("max_size", None),
"name": all_config.pop("name", None),
"timeout": all_config.pop("timeout", 30.0),
"max_waiting": all_config.pop("max_waiting", 0),
"max_lifetime": all_config.pop("max_lifetime", 3600.0),
"max_idle": all_config.pop("max_idle", 600.0),
"reconnect_timeout": all_config.pop("reconnect_timeout", 300.0),
"reconnect_failed": all_config.pop("reconnect_failed", None),
"num_workers": all_config.pop("num_workers", 3),
"check": all_config.pop("check", None),
"reset": all_config.pop("reset", None),
"close_returns": all_config.pop("close_returns", False),
}
open_pool = all_config.pop("open", True)
pool_parameters["open"] = False if open_pool is True else open_pool
pool_parameters["configure"] = all_config.pop("configure", self._configure_async_connection)
pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None}
conninfo = all_config.pop("conninfo", None)
kwargs = all_config.pop("kwargs", {})
all_config.update(kwargs)
is_null_pool = bool(self.connection_config.get("null_pool") or self.driver_features.get("null_pool"))
pool_cls = AsyncNullConnectionPool if is_null_pool else AsyncConnectionPool
if is_null_pool:
pool_parameters.pop("min_size", None)
if conninfo:
pool = pool_cls(conninfo, kwargs=all_config, **pool_parameters)
else:
pool = pool_cls("", kwargs=all_config, **pool_parameters)
if open_pool is True:
await pool.open()
return pool
async def _configure_async_connection(self, conn: "PsycopgAsyncConnection") -> None:
autocommit_setting = self.connection_config.get("autocommit")
if autocommit_setting is not None:
await conn.set_autocommit(autocommit_setting)
serializer = self.driver_features.get("json_serializer", to_json)
deserializer = self.driver_features.get("json_deserializer", from_json)
if isinstance(getattr(conn, "adapters", None), AdaptersMap):
set_json_dumps(serializer, conn)
set_json_loads(deserializer, conn)
if self._pgvector_available is None:
detected_extensions: set[str] = set()
extensions = build_postgres_extension_probe_names(self.driver_features)
if extensions:
try:
cursor = await conn.execute(
"SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,)
)
results = await cursor.fetchall()
detected_extensions = {r[0] for r in results} # type: ignore[index]
except Exception as exc:
detected_extensions = set()
self._pg_textsearch_probe_error = exc
self.statement_config, self._pgvector_available, self._paradedb_available = (
resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions)
)
self._pg_textsearch_available = is_postgres_extension_active(self.driver_features, "pg_textsearch")
if self._pgvector_available:
await register_pgvector_async(conn)
if not conn.autocommit:
await conn.rollback()
if self._user_connection_hook is not None:
await self._user_connection_hook(conn)
@property
def pg_textsearch_available(self) -> bool:
"""Return True if the pg_textsearch extension is available."""
return bool(self._pg_textsearch_available)
def _ensure_pg_textsearch_available(self) -> None:
if self._pg_textsearch_available:
return
msg = "ADK memory enable_bm25 requires the pg_textsearch PostgreSQL extension"
raise ImproperConfigurationError(msg) from self._pg_textsearch_probe_error
async def _close_pool(self) -> None:
"""Close the actual async connection pool."""
if not self.connection_instance:
return
try:
await self.connection_instance.close()
finally:
self.connection_instance = None
def _connection_kwargs(self) -> "tuple[type[Any], str, dict[str, Any]]":
"""Split the pooled configuration into standalone connect arguments.
Returns:
The connection class, the conninfo string, and the connect kwargs.
"""
all_config = {k: v for k, v in self.connection_config.items() if k not in _POOL_ONLY_CONFIG_KEYS}
connection_class = all_config.pop("connection_class", None) or PsycopgAsyncConnection
conninfo = all_config.pop("conninfo", None)
all_config.update(all_config.pop("kwargs", {}))
return cast("type[Any]", connection_class), str(conninfo or ""), all_config
[docs]
async def create_connection(self) -> "PsycopgAsyncConnection": # pyright: ignore
"""Open a standalone connection owned by the caller.
The connection is not bound to the pool, so closing it actually closes
it and consumes no pool slot.
Returns:
A psycopg AsyncConnection instance.
"""
connection_class, conninfo, connection_kwargs = self._connection_kwargs()
connection = await connection_class.connect(conninfo, **connection_kwargs)
configure = self.connection_config.get("configure", self._configure_async_connection)
try:
await configure(cast("PsycopgAsyncConnection", connection))
except BaseException:
with suppress(Exception):
await connection.close()
raise
return cast("PsycopgAsyncConnection", connection)
[docs]
def get_signature_namespace(self) -> "dict[str, Any]":
"""Get the signature namespace for PsycopgAsyncConfig types.
Returns:
Dictionary mapping type names to types.
"""
namespace = super().get_signature_namespace()
namespace.update({
"PsycopgAsyncConnectionContext": PsycopgAsyncConnectionContext,
"PsycopgAsyncConnection": PsycopgAsyncConnection,
"PsycopgAsyncCursor": PsycopgAsyncCursor,
"PsycopgAsyncDriver": PsycopgAsyncDriver,
"PsycopgAsyncExceptionHandler": PsycopgAsyncExceptionHandler,
"PsycopgAsyncSessionContext": PsycopgAsyncSessionContext,
"PsycopgConnectionParams": PsycopgConnectionParams,
"PsycopgPoolParams": PsycopgPoolParams,
})
return namespace
[docs]
def provide_session(
self, *_args: Any, statement_config: "StatementConfig | None" = None, **_kwargs: Any
) -> "PsycopgAsyncSessionContext":
"""Provide an async driver session context manager.
Args:
*_args: Additional arguments.
statement_config: Optional statement configuration override.
**_kwargs: Additional keyword arguments.
Returns:
A PsycopgAsyncDriver session context manager.
"""
handler = _PsycopgAsyncSessionConnectionHandler(self)
return PsycopgAsyncSessionContext(
acquire_connection=handler.acquire_connection,
release_connection=handler.release_connection,
statement_config=statement_config
or (lambda: resolve_runtime_statement_config(None, self.statement_config, default_statement_config)),
driver_features=self.driver_features,
prepare_driver=self._prepare_driver,
)
[docs]
async def provide_pool(self, *args: Any, **kwargs: Any) -> "AsyncConnectionPool":
"""Provide async pool instance.
Returns:
The async connection pool.
"""
if not self.connection_instance:
self.connection_instance = await self.create_pool()
return self.connection_instance
[docs]
def get_event_runtime_hints(self) -> "EventRuntimeHints":
"""Return polling defaults for PostgreSQL queue fallback."""
return EventRuntimeHints(poll_interval=0.5, select_for_update=True, skip_locked=True)
def _make_alloydb_connection_class(
*,
connector: Any,
instance_uri: str,
enable_iam_auth: bool,
ip_type: str,
user: str | None,
password: str | None,
database: str | None,
) -> "type[PsycopgConnection[Any]]":
"""Build a psycopg connection class backed by the AlloyDB connector."""
class _AlloyDBPsycopgConnection(PsycopgConnection[Any]):
@classmethod
def connect(
cls,
conninfo: str = "",
*,
autocommit: bool = False,
prepare_threshold: int | None = 5,
context: "AdaptContext | None" = None,
row_factory: "RowFactory[Any] | None" = None,
cursor_factory: "type[Cursor[Any]] | None" = None,
**kwargs: str | int | None,
) -> Self:
_ = (cls, conninfo, autocommit, prepare_threshold, context, row_factory, cursor_factory)
connector_kwargs = dict(kwargs)
connector_kwargs["enable_iam_auth"] = enable_iam_auth
connector_kwargs["ip_type"] = ip_type
if user is not None:
connector_kwargs["user"] = user
if password is not None:
connector_kwargs["password"] = password
if database is not None:
connector_kwargs["db"] = database
return cast("Self", connector.connect(instance_uri, "psycopg", **connector_kwargs))
return _AlloyDBPsycopgConnection
