"""PostgreSQL psycopg driver implementation."""
from collections.abc import Sized
from contextlib import AsyncExitStack, ExitStack
from typing import TYPE_CHECKING, Any, cast
import psycopg
from typing_extensions import LiteralString
from sqlspec.adapters.psycopg._typing import (
PsycopgAsyncConnection,
PsycopgAsyncCursor,
PsycopgAsyncSessionContext,
PsycopgComposed,
PsycopgSQL,
PsycopgSyncConnection,
PsycopgSyncCursor,
PsycopgSyncSessionContext,
)
from sqlspec.adapters.psycopg.core import (
TRANSACTION_STATUS_IDLE,
PipelineCursorEntry,
PreparedStackOperation,
PsycopgAsyncStreamSource,
PsycopgSyncStreamSource,
build_async_pipeline_execution_result,
build_copy_from_command,
build_pipeline_execution_result,
build_truncate_command,
create_mapped_exception,
default_statement_config,
driver_profile,
execute_with_optional_parameters,
execute_with_optional_parameters_async,
pipeline_supported,
resolve_many_rowcount,
resolve_rowcount,
)
from sqlspec.adapters.psycopg.data_dictionary import PsycopgAsyncDataDictionary, PsycopgSyncDataDictionary
from sqlspec.core import (
SQL,
SQLResult,
StackResult,
StatementConfig,
StatementStack,
get_cache_config,
is_copy_from_operation,
is_copy_operation,
is_copy_to_operation,
register_driver_profile,
)
from sqlspec.driver import (
AsyncDriverAdapterBase,
AsyncRowStream,
BaseAsyncExceptionHandler,
BaseSyncExceptionHandler,
StackExecutionObserver,
SyncDriverAdapterBase,
SyncRowStream,
describe_stack_statement,
)
from sqlspec.exceptions import SQLSpecError, StackExecutionError
from sqlspec.utils.logging import get_logger
from sqlspec.utils.text import normalize_identifier, quote_identifier
from sqlspec.utils.type_guards import is_readable, resolve_row_format
if TYPE_CHECKING:
from collections import abc
from sqlspec.adapters.psycopg._typing import PsycopgPipelineDriver
from sqlspec.core import ArrowResult
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
__all__ = (
"PsycopgAsyncCursor",
"PsycopgAsyncDriver",
"PsycopgAsyncExceptionHandler",
"PsycopgAsyncSessionContext",
"PsycopgSyncCursor",
"PsycopgSyncDriver",
"PsycopgSyncExceptionHandler",
"PsycopgSyncSessionContext",
)
logger = get_logger("sqlspec.adapters.psycopg")
_PSYCOPG_OID_TOKENS: "dict[int, str]" = {
16: "bool",
17: "binary",
20: "int64",
21: "int16",
23: "int32",
25: "string",
114: "string",
700: "float32",
701: "float64",
1043: "string",
1082: "date",
1083: "time",
1114: "timestamp",
1184: "timestamptz",
1700: "decimal",
2950: "string",
3802: "string",
}
def _resolve_column_types(description: Any) -> "dict[str, str] | None":
"""Map psycopg cursor column OIDs to neutral Arrow type tokens.
Returns ``None`` when the cursor has no description or reports no
recognizable OIDs, so callers can pass the result straight through
without adding a code path for the empty case.
"""
if not description:
return None
column_types: dict[str, str] = {}
for col in description:
token = _PSYCOPG_OID_TOKENS.get(col.type_code)
if token is not None:
column_types[col.name] = token
return column_types or None
def pipeline_operation_failed(cursor: Any, statement: "SQL") -> bool:
"""Return True when a synced pipeline cursor reflects a failed non-select operation.
After a pipeline sync raises, the failing operation and every operation queued
behind it report a negative rowcount, while operations that committed report a
non-negative one. Row-returning operations surface their failure when the result
is fetched, so they are excluded here.
"""
if statement.returns_rows():
return False
try:
rowcount = cursor.rowcount
except Exception:
return True
return isinstance(rowcount, int) and rowcount < 0
class PsycopgPipelineMixin:
"""Shared helpers for psycopg sync/async pipeline execution."""
__slots__ = ()
def _prepare_records_for_arrow(
self, records: "abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]"
) -> "abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]":
driver = cast("PsycopgPipelineDriver", self)
return cast(
"abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]",
driver.prepare_driver_parameters(records, driver.statement_config, is_many=True),
)
def _prepare_pipeline_operations(self, stack: "StatementStack") -> "list[PreparedStackOperation] | None":
prepared: list[PreparedStackOperation] = []
for index, operation in enumerate(stack.operations):
if operation.method != "execute":
return None
kwargs = dict(operation.keyword_arguments) if operation.keyword_arguments else {}
statement_config = kwargs.pop("statement_config", None)
driver = cast("PsycopgPipelineDriver", self)
config = statement_config or driver.statement_config
sql_statement = driver.prepare_statement(
operation.statement, operation.arguments, statement_config=config, kwargs=kwargs
)
if sql_statement.is_script or sql_statement.is_many:
return None
sql_text, prepared_parameters = driver._compiled_sql( # pyright: ignore[reportPrivateUsage]
sql_statement, config
)
prepared.append(
PreparedStackOperation(
operation_index=index,
operation=operation,
statement=sql_statement,
sql=cast("LiteralString | PsycopgSQL | PsycopgComposed", sql_text),
parameters=prepared_parameters,
)
)
return prepared
class PsycopgSyncExceptionHandler(BaseSyncExceptionHandler):
"""Context manager for handling PostgreSQL psycopg database exceptions.
Maps PostgreSQL SQLSTATE error codes to specific SQLSpec exceptions
for better error handling in application code.
Uses deferred exception pattern for mypyc compatibility: exceptions
are stored in pending_exception rather than raised from __exit__
to avoid ABI boundary violations with compiled code.
"""
__slots__ = ()
def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool:
if exc_type is None:
return False
if issubclass(exc_type, psycopg.Error):
self.pending_exception = create_mapped_exception(exc_val)
return True
return False
class PsycopgSyncDriver(PsycopgPipelineMixin, SyncDriverAdapterBase):
"""PostgreSQL psycopg synchronous driver.
Provides synchronous database operations for PostgreSQL using psycopg3.
Supports SQL statement execution with parameter binding, transaction
management, result processing with column metadata, parameter style
conversion, PostgreSQL arrays and JSON handling, COPY operations for
bulk data transfer, and PostgreSQL-specific error handling.
"""
__slots__ = ("_data_dictionary", "_restore_autocommit", "_transaction_active")
dialect = "postgres"
def __init__(
self,
connection: PsycopgSyncConnection,
statement_config: "StatementConfig | None" = None,
driver_features: "dict[str, Any] | None" = None,
) -> None:
if statement_config is None:
statement_config = default_statement_config.replace(
enable_caching=get_cache_config().compiled_cache_enabled
)
super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
self._data_dictionary: PsycopgSyncDataDictionary | None = None
self._restore_autocommit = False
self._transaction_active = False