"""Psqlpy driver implementation for PostgreSQL connectivity.
Provides parameter style conversion, type coercion, error handling,
and transaction management.
"""
import inspect
from typing import TYPE_CHECKING, Any, cast
from sqlspec.adapters.psqlpy._typing import PsqlpyCursor, PsqlpyDatabaseError, PsqlpyError, PsqlpySessionContext
from sqlspec.adapters.psqlpy.core import (
_DML_COUNT_COLUMN,
PsqlpyStreamSource,
_dml_count_query,
build_insert_statement,
coerce_numeric_for_write,
coerce_records_for_execute_many,
collect_rows,
create_mapped_exception,
default_statement_config,
driver_profile,
encode_records_for_binary_copy,
extract_rows_affected,
format_execute_many_parameters,
format_table_identifier,
get_parameter_casts,
prepare_parameters_with_casts,
split_schema_and_table,
)
from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyDataDictionary
from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile
from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler
from sqlspec.exceptions import SQLSpecError
from sqlspec.utils.logging import get_logger
from sqlspec.utils.text import normalize_identifier, quote_identifier
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from sqlspec.adapters.psqlpy._typing import PsqlpyConnection
from sqlspec.core import ArrowResult, SQLResult
from sqlspec.driver import ExecutionResult
from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
__all__ = ("PsqlpyCursor", "PsqlpyDriver", "PsqlpyExceptionHandler", "PsqlpySessionContext")
logger = get_logger("sqlspec.adapters.psqlpy")
class PsqlpyExceptionHandler(BaseAsyncExceptionHandler):
"""Async context manager for handling psqlpy 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 __aexit__
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, (PsqlpyDatabaseError, PsqlpyError)):
self.pending_exception = create_mapped_exception(exc_val)
return True
return False
class PsqlpyDriver(AsyncDriverAdapterBase):
"""PostgreSQL driver implementation using psqlpy.
Provides parameter style conversion, type coercion, error handling,
and transaction management.
"""
__slots__ = ("_data_dictionary",)
dialect = "postgres"
def __init__(
self,
connection: "PsqlpyConnection",
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: PsqlpyDataDictionary | None = None