sqlspec 0.13.1__py3-none-any.whl → 0.16.2__py3-none-any.whl

sqlspec/adapters/psqlpy/__init__.py

@@ -1,6 +1,16 @@
 """Psqlpy adapter for SQLSpec."""
 
-from sqlspec.adapters.psqlpy.config import CONNECTION_FIELDS, POOL_FIELDS, PsqlpyConfig
-from sqlspec.adapters.psqlpy.driver import PsqlpyConnection, PsqlpyDriver
+from sqlspec.adapters.psqlpy._types import PsqlpyConnection
+from sqlspec.adapters.psqlpy.config import PsqlpyConfig, PsqlpyConnectionParams, PsqlpyPoolParams
+from sqlspec.adapters.psqlpy.driver import PsqlpyCursor, PsqlpyDriver, PsqlpyExceptionHandler, psqlpy_statement_config
 
-__all__ = ("CONNECTION_FIELDS", "POOL_FIELDS", "PsqlpyConfig", "PsqlpyConnection", "PsqlpyDriver")
+__all__ = (
+    "PsqlpyConfig",
+    "PsqlpyConnection",
+    "PsqlpyConnectionParams",
+    "PsqlpyCursor",
+    "PsqlpyDriver",
+    "PsqlpyExceptionHandler",
+    "PsqlpyPoolParams",
+    "psqlpy_statement_config",
+)

sqlspec/adapters/psqlpy/_types.py

@@ -0,0 +1,11 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from psqlpy import Connection
+    from typing_extensions import TypeAlias
+
+    PsqlpyConnection: TypeAlias = Connection
+else:
+    from psqlpy import Connection as PsqlpyConnection
+
+__all__ = ("PsqlpyConnection",)

sqlspec/adapters/psqlpy/config.py

@@ -3,334 +3,128 @@
 import logging
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, ClassVar, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, TypedDict, Union
 
 from psqlpy import ConnectionPool
+from typing_extensions import NotRequired
 
-from sqlspec.adapters.psqlpy.driver import PsqlpyConnection, PsqlpyDriver
+from sqlspec.adapters.psqlpy._types import PsqlpyConnection
+from sqlspec.adapters.psqlpy.driver import PsqlpyCursor, PsqlpyDriver, psqlpy_statement_config
 from sqlspec.config import AsyncDatabaseConfig
-from sqlspec.statement.sql import SQLConfig
-from sqlspec.typing import DictRow, Empty
+from sqlspec.core.statement import StatementConfig
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from sqlglot.dialects.dialect import DialectType
 
 logger = logging.getLogger("sqlspec.adapters.psqlpy")
 
-CONNECTION_FIELDS = frozenset(
-    {
-        "dsn",
-        "username",
-        "password",
-        "db_name",
-        "host",
-        "port",
-        "connect_timeout_sec",
-        "connect_timeout_nanosec",
-        "tcp_user_timeout_sec",
-        "tcp_user_timeout_nanosec",
-        "keepalives",
-        "keepalives_idle_sec",
-        "keepalives_idle_nanosec",
-        "keepalives_interval_sec",
-        "keepalives_interval_nanosec",
-        "keepalives_retries",
-        "ssl_mode",
-        "ca_file",
-        "target_session_attrs",
-        "options",
-        "application_name",
-        "client_encoding",
-        "gssencmode",
-        "sslnegotiation",
-        "sslcompression",
-        "sslcert",
-        "sslkey",
-        "sslpassword",
-        "sslrootcert",
-        "sslcrl",
-        "require_auth",
-        "channel_binding",
-        "krbsrvname",
-        "gsslib",
-        "gssdelegation",
-        "service",
-        "load_balance_hosts",
-    }
-)
-
-POOL_FIELDS = CONNECTION_FIELDS.union({"hosts", "ports", "conn_recycling_method", "max_db_pool_size", "configure"})
-
-__all__ = ("CONNECTION_FIELDS", "POOL_FIELDS", "PsqlpyConfig")
+
+class PsqlpyConnectionParams(TypedDict, total=False):
+    """Psqlpy connection parameters."""
+
+    dsn: NotRequired[str]
+    username: NotRequired[str]
+    password: NotRequired[str]
+    db_name: NotRequired[str]
+    host: NotRequired[str]
+    port: NotRequired[int]
+    connect_timeout_sec: NotRequired[int]
+    connect_timeout_nanosec: NotRequired[int]
+    tcp_user_timeout_sec: NotRequired[int]
+    tcp_user_timeout_nanosec: NotRequired[int]
+    keepalives: NotRequired[bool]
+    keepalives_idle_sec: NotRequired[int]
+    keepalives_idle_nanosec: NotRequired[int]
+    keepalives_interval_sec: NotRequired[int]
+    keepalives_interval_nanosec: NotRequired[int]
+    keepalives_retries: NotRequired[int]
+    ssl_mode: NotRequired[str]
+    ca_file: NotRequired[str]
+    target_session_attrs: NotRequired[str]
+    options: NotRequired[str]
+    application_name: NotRequired[str]
+    client_encoding: NotRequired[str]
+    gssencmode: NotRequired[str]
+    sslnegotiation: NotRequired[str]
+    sslcompression: NotRequired[str]
+    sslcert: NotRequired[str]
+    sslkey: NotRequired[str]
+    sslpassword: NotRequired[str]
+    sslrootcert: NotRequired[str]
+    sslcrl: NotRequired[str]
+    require_auth: NotRequired[str]
+    channel_binding: NotRequired[str]
+    krbsrvname: NotRequired[str]
+    gsslib: NotRequired[str]
+    gssdelegation: NotRequired[str]
+    service: NotRequired[str]
+    load_balance_hosts: NotRequired[str]
+
+
+class PsqlpyPoolParams(PsqlpyConnectionParams, total=False):
+    """Psqlpy pool parameters."""
+
+    hosts: NotRequired[list[str]]
+    ports: NotRequired[list[int]]
+    conn_recycling_method: NotRequired[str]
+    max_db_pool_size: NotRequired[int]
+    configure: NotRequired["Callable[..., Any]"]
+    extra: NotRequired[dict[str, Any]]
+
+
+__all__ = ("PsqlpyConfig", "PsqlpyConnectionParams", "PsqlpyCursor", "PsqlpyPoolParams")
 
 
 class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyDriver]):
     """Configuration for Psqlpy asynchronous database connections with direct field-based configuration."""
 
-    __slots__ = (
-        "_dialect",
-        "application_name",
-        "ca_file",
-        "channel_binding",
-        "client_encoding",
-        "configure",
-        "conn_recycling_method",
-        "connect_timeout_nanosec",
-        "connect_timeout_sec",
-        "db_name",
-        "default_row_type",
-        "dsn",
-        "extras",
-        "gssdelegation",
-        "gssencmode",
-        "gsslib",
-        "host",
-        "hosts",
-        "keepalives",
-        "keepalives_idle_nanosec",
-        "keepalives_idle_sec",
-        "keepalives_interval_nanosec",
-        "keepalives_interval_sec",
-        "keepalives_retries",
-        "krbsrvname",
-        "load_balance_hosts",
-        "max_db_pool_size",
-        "options",
-        "password",
-        "pool_instance",
-        "port",
-        "ports",
-        "require_auth",
-        "service",
-        "ssl_mode",
-        "sslcert",
-        "sslcompression",
-        "sslcrl",
-        "sslkey",
-        "sslnegotiation",
-        "sslpassword",
-        "sslrootcert",
-        "statement_config",
-        "target_session_attrs",
-        "tcp_user_timeout_nanosec",
-        "tcp_user_timeout_sec",
-        "username",
-    )
-
-    is_async: ClassVar[bool] = True
-    supports_connection_pooling: ClassVar[bool] = True
-
-    driver_type: type[PsqlpyDriver] = PsqlpyDriver
-    connection_type: type[PsqlpyConnection] = PsqlpyConnection
-    # Parameter style support information
-    supported_parameter_styles: ClassVar[tuple[str, ...]] = ("numeric",)
-    """Psqlpy only supports $1, $2, ... (numeric) parameter style."""
-
-    preferred_parameter_style: ClassVar[str] = "numeric"
-    """Psqlpy's native parameter style is $1, $2, ... (numeric)."""
+    driver_type: ClassVar[type[PsqlpyDriver]] = PsqlpyDriver
+    connection_type: "ClassVar[type[PsqlpyConnection]]" = PsqlpyConnection
 
     def __init__(
         self,
-        statement_config: Optional[SQLConfig] = None,
-        default_row_type: type[DictRow] = DictRow,
-        # Connection parameters
-        dsn: Optional[str] = None,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
-        db_name: Optional[str] = None,
-        host: Optional[str] = None,
-        port: Optional[int] = None,
-        hosts: Optional[list[str]] = None,
-        ports: Optional[list[int]] = None,
-        connect_timeout_sec: Optional[int] = None,
-        connect_timeout_nanosec: Optional[int] = None,
-        tcp_user_timeout_sec: Optional[int] = None,
-        tcp_user_timeout_nanosec: Optional[int] = None,
-        keepalives: Optional[bool] = None,
-        keepalives_idle_sec: Optional[int] = None,
-        keepalives_idle_nanosec: Optional[int] = None,
-        keepalives_interval_sec: Optional[int] = None,
-        keepalives_interval_nanosec: Optional[int] = None,
-        keepalives_retries: Optional[int] = None,
-        ssl_mode: Optional[str] = None,
-        ca_file: Optional[str] = None,
-        target_session_attrs: Optional[str] = None,
-        options: Optional[str] = None,
-        application_name: Optional[str] = None,
-        client_encoding: Optional[str] = None,
-        gssencmode: Optional[str] = None,
-        sslnegotiation: Optional[str] = None,
-        sslcompression: Optional[bool] = None,
-        sslcert: Optional[str] = None,
-        sslkey: Optional[str] = None,
-        sslpassword: Optional[str] = None,
-        sslrootcert: Optional[str] = None,
-        sslcrl: Optional[str] = None,
-        require_auth: Optional[str] = None,
-        channel_binding: Optional[str] = None,
-        krbsrvname: Optional[str] = None,
-        gsslib: Optional[str] = None,
-        gssdelegation: Optional[bool] = None,
-        service: Optional[str] = None,
-        load_balance_hosts: Optional[str] = None,
-        # Pool parameters
-        conn_recycling_method: Optional[str] = None,
-        max_db_pool_size: Optional[int] = None,
-        configure: Optional["Callable[[ConnectionPool], None]"] = None,
+        *,
+        pool_config: Optional[Union[PsqlpyPoolParams, dict[str, Any]]] = None,
+        statement_config: Optional[StatementConfig] = None,
         pool_instance: Optional[ConnectionPool] = None,
-        **kwargs: Any,
+        migration_config: Optional[dict[str, Any]] = None,
     ) -> None:
         """Initialize Psqlpy asynchronous configuration.
 
         Args:
-            statement_config: Default SQL statement configuration
-            default_row_type: Default row type for results
-            dsn: DSN of the PostgreSQL database
-            username: Username of the user in the PostgreSQL
-            password: Password of the user in the PostgreSQL
-            db_name: Name of the database in PostgreSQL
-            host: Host of the PostgreSQL (use for single host)
-            port: Port of the PostgreSQL (use for single host)
-            hosts: List of hosts of the PostgreSQL (use for multiple hosts)
-            ports: List of ports of the PostgreSQL (use for multiple hosts)
-            connect_timeout_sec: The time limit in seconds applied to each socket-level connection attempt
-            connect_timeout_nanosec: Nanoseconds for connection timeout, can be used only with connect_timeout_sec
-            tcp_user_timeout_sec: The time limit that transmitted data may remain unacknowledged before a connection is forcibly closed
-            tcp_user_timeout_nanosec: Nanoseconds for tcp_user_timeout, can be used only with tcp_user_timeout_sec
-            keepalives: Controls the use of TCP keepalive. Defaults to True (on)
-            keepalives_idle_sec: The number of seconds of inactivity after which a keepalive message is sent to the server
-            keepalives_idle_nanosec: Nanoseconds for keepalives_idle_sec
-            keepalives_interval_sec: The time interval between TCP keepalive probes
-            keepalives_interval_nanosec: Nanoseconds for keepalives_interval_sec
-            keepalives_retries: The maximum number of TCP keepalive probes that will be sent before dropping a connection
-            ssl_mode: SSL mode (disable, prefer, require, verify-ca, verify-full)
-            ca_file: Path to ca_file for SSL
-            target_session_attrs: Specifies requirements of the session (e.g., 'read-write', 'read-only', 'primary', 'standby')
-            options: Command line options used to configure the server
-            application_name: Sets the application_name parameter on the server
-            client_encoding: Sets the client_encoding parameter
-            gssencmode: GSS encryption mode (disable, prefer, require)
-            sslnegotiation: SSL negotiation mode (postgres, direct)
-            sslcompression: Whether to use SSL compression
-            sslcert: Client SSL certificate file
-            sslkey: Client SSL private key file
-            sslpassword: Password for the SSL private key
-            sslrootcert: SSL root certificate file
-            sslcrl: SSL certificate revocation list file
-            require_auth: Authentication method requirements
-            channel_binding: Channel binding preference (disable, prefer, require)
-            krbsrvname: Kerberos service name
-            gsslib: GSS library to use
-            gssdelegation: Forward GSS credentials to server
-            service: Service name for additional parameters
-            load_balance_hosts: Controls the order in which the client tries to connect to the available hosts and addresses ('disable' or 'random')
-            conn_recycling_method: How a connection is recycled
-            max_db_pool_size: Maximum size of the connection pool. Defaults to 10
-            configure: Callback to configure new connections
+            pool_config: Pool configuration parameters (TypedDict or dict)
             pool_instance: Existing connection pool instance to use
-            **kwargs: Additional parameters (stored in extras)
-        """
-        # Store connection parameters as instance attributes
-        self.dsn = dsn
-        self.username = username
-        self.password = password
-        self.db_name = db_name
-        self.host = host
-        self.port = port
-        self.hosts = hosts
-        self.ports = ports
-        self.connect_timeout_sec = connect_timeout_sec
-        self.connect_timeout_nanosec = connect_timeout_nanosec
-        self.tcp_user_timeout_sec = tcp_user_timeout_sec
-        self.tcp_user_timeout_nanosec = tcp_user_timeout_nanosec
-        self.keepalives = keepalives
-        self.keepalives_idle_sec = keepalives_idle_sec
-        self.keepalives_idle_nanosec = keepalives_idle_nanosec
-        self.keepalives_interval_sec = keepalives_interval_sec
-        self.keepalives_interval_nanosec = keepalives_interval_nanosec
-        self.keepalives_retries = keepalives_retries
-        self.ssl_mode = ssl_mode
-        self.ca_file = ca_file
-        self.target_session_attrs = target_session_attrs
-        self.options = options
-        self.application_name = application_name
-        self.client_encoding = client_encoding
-        self.gssencmode = gssencmode
-        self.sslnegotiation = sslnegotiation
-        self.sslcompression = sslcompression
-        self.sslcert = sslcert
-        self.sslkey = sslkey
-        self.sslpassword = sslpassword
-        self.sslrootcert = sslrootcert
-        self.sslcrl = sslcrl
-        self.require_auth = require_auth
-        self.channel_binding = channel_binding
-        self.krbsrvname = krbsrvname
-        self.gsslib = gsslib
-        self.gssdelegation = gssdelegation
-        self.service = service
-        self.load_balance_hosts = load_balance_hosts
-
-        # Store pool parameters as instance attributes
-        self.conn_recycling_method = conn_recycling_method
-        self.max_db_pool_size = max_db_pool_size
-        self.configure = configure
-
-        self.extras = kwargs or {}
-
-        # Store other config
-        self.statement_config = statement_config or SQLConfig()
-        self.default_row_type = default_row_type
-        self.pool_instance: Optional[ConnectionPool] = pool_instance
-        self._dialect: DialectType = None
-
-        super().__init__()
-
-    @property
-    def connection_config_dict(self) -> dict[str, Any]:
-        """Return the connection configuration as a dict for psqlpy.Connection.
-
-        This method filters out pool-specific parameters that are not valid for psqlpy.Connection.
+            statement_config: Default SQL statement configuration
+            migration_config: Migration configuration
         """
-        # Gather non-None connection parameters
-        config = {
-            field: getattr(self, field)
-            for field in CONNECTION_FIELDS
-            if getattr(self, field, None) is not None and getattr(self, field) is not Empty
-        }
-
-        config.update(self.extras)
-
-        return config
-
-    @property
-    def pool_config_dict(self) -> dict[str, Any]:
-        """Return the full pool configuration as a dict for psqlpy.ConnectionPool.
+        processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
+        if "extra" in processed_pool_config:
+            extras = processed_pool_config.pop("extra")
+            processed_pool_config.update(extras)
+        super().__init__(
+            pool_config=processed_pool_config,
+            pool_instance=pool_instance,
+            migration_config=migration_config,
+            statement_config=statement_config or psqlpy_statement_config,
+        )
+
+    def _get_pool_config_dict(self) -> dict[str, Any]:
+        """Get pool configuration as plain dict for external library.
 
         Returns:
-            A dictionary containing all pool configuration parameters.
+            Dictionary with pool parameters, filtering out None values.
         """
-        # Gather non-None parameters from all fields (connection + pool)
-        config = {
-            field: getattr(self, field)
-            for field in POOL_FIELDS
-            if getattr(self, field, None) is not None and getattr(self, field) is not Empty
-        }
-
-        # Merge extras parameters
-        config.update(self.extras)
-
-        return config
+        return {k: v for k, v in self.pool_config.items() if v is not None}
 
     async def _create_pool(self) -> "ConnectionPool":
         """Create the actual async connection pool."""
         logger.info("Creating psqlpy connection pool", extra={"adapter": "psqlpy"})
 
         try:
-            config = self.pool_config_dict
-            pool = ConnectionPool(**config)  # pyright: ignore
+            config = self._get_pool_config_dict()
+
+            pool = ConnectionPool(**config)
             logger.info("Psqlpy connection pool created successfully", extra={"adapter": "psqlpy"})
         except Exception as e:
             logger.exception("Failed to create psqlpy connection pool", extra={"adapter": "psqlpy", "error": str(e)})
@@ -380,29 +174,21 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
             yield conn
 
     @asynccontextmanager
-    async def provide_session(self, *args: Any, **kwargs: Any) -> AsyncGenerator[PsqlpyDriver, None]:
+    async def provide_session(
+        self, *args: Any, statement_config: "Optional[StatementConfig]" = None, **kwargs: Any
+    ) -> AsyncGenerator[PsqlpyDriver, None]:
         """Provide an async driver session context manager.
 
         Args:
             *args: Additional arguments.
+            statement_config: Optional statement configuration override.
             **kwargs: Additional keyword arguments.
 
         Yields:
             A PsqlpyDriver instance.
         """
         async with self.provide_connection(*args, **kwargs) as conn:
-            statement_config = self.statement_config
-            # Inject parameter style info if not already set
-            if statement_config.allowed_parameter_styles is None:
-                from dataclasses import replace
-
-                statement_config = replace(
-                    statement_config,
-                    allowed_parameter_styles=self.supported_parameter_styles,
-                    target_parameter_style=self.preferred_parameter_style,
-                )
-            driver = self.driver_type(connection=conn, config=statement_config)
-            yield driver
+            yield self.driver_type(connection=conn, statement_config=statement_config or self.statement_config)
 
     async def provide_pool(self, *args: Any, **kwargs: Any) -> ConnectionPool:
         """Provide async pool instance.
@@ -424,5 +210,5 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
             Dictionary mapping type names to types.
         """
         namespace = super().get_signature_namespace()
-        namespace.update({"PsqlpyConnection": PsqlpyConnection})
+        namespace.update({"PsqlpyConnection": PsqlpyConnection, "PsqlpyCursor": PsqlpyCursor})
         return namespace