sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/adapters/psqlpy/data_dictionary.py

@@ -0,0 +1,121 @@
+"""PostgreSQL-specific data dictionary for metadata queries via psqlpy."""
+
+from typing import TYPE_CHECKING, ClassVar
+
+from mypy_extensions import mypyc_attr
+
+from sqlspec.data_dictionary.dialects.postgres import resolve_postgres_json_type
+from sqlspec.driver import AsyncDataDictionaryBase
+from sqlspec.typing import ColumnMetadata, ForeignKeyMetadata, IndexMetadata, TableMetadata, VersionInfo
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.psqlpy.driver import PsqlpyDriver
+
+__all__ = ("PsqlpyDataDictionary",)
+
+
+@mypyc_attr(allow_interpreted_subclasses=True, native_class=False)
+class PsqlpyDataDictionary(AsyncDataDictionaryBase):
+    """PostgreSQL-specific async data dictionary via psqlpy."""
+
+    dialect: ClassVar[str] = "postgres"
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    async def get_version(self, driver: "PsqlpyDriver") -> "VersionInfo | None":
+        """Get PostgreSQL database version information."""
+        driver_id = id(driver)
+        # Inline cache check to avoid cross-module method call that causes mypyc segfault
+        if driver_id in self._version_fetch_attempted:
+            return self._version_cache.get(driver_id)
+        # Not cached, fetch from database
+
+        version_value = await driver.select_value(self.get_query("version"))
+        if not version_value:
+            self._log_version_unavailable(type(self).dialect, "missing")
+            self.cache_version(driver_id, None)
+            return None
+
+        version_info = self.parse_version_with_pattern(self.get_dialect_config().version_pattern, str(version_value))
+        if version_info is None:
+            self._log_version_unavailable(type(self).dialect, "parse_failed")
+            self.cache_version(driver_id, None)
+            return None
+
+        self._log_version_detected(type(self).dialect, version_info)
+        self.cache_version(driver_id, version_info)
+        return version_info
+
+    async def get_feature_flag(self, driver: "PsqlpyDriver", feature: str) -> bool:
+        """Check if PostgreSQL database supports a specific feature."""
+        version_info = await self.get_version(driver)
+        return self.resolve_feature_flag(feature, version_info)
+
+    async def get_optimal_type(self, driver: "PsqlpyDriver", type_category: str) -> str:
+        """Get optimal PostgreSQL type for a category."""
+        config = self.get_dialect_config()
+        version_info = await self.get_version(driver)
+
+        if type_category == "json":
+            return resolve_postgres_json_type(version_info)
+
+        return config.get_optimal_type(type_category)
+
+    async def get_tables(self, driver: "PsqlpyDriver", schema: "str | None" = None) -> "list[TableMetadata]":
+        """Get tables sorted by topological dependency order using Recursive CTE."""
+        schema_name = self.resolve_schema(schema)
+        self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables")
+        return await driver.select(
+            self.get_query("tables_by_schema"), schema_name=schema_name, schema_type=TableMetadata
+        )
+
+    async def get_columns(
+        self, driver: "PsqlpyDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[ColumnMetadata]":
+        """Get column information for a table or schema."""
+        schema_name = self.resolve_schema(schema)
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="columns")
+            return await driver.select(
+                self.get_query("columns_by_schema"), schema_name=schema_name, schema_type=ColumnMetadata
+            )
+
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="columns")
+        return await driver.select(
+            self.get_query("columns_by_table"), schema_name=schema_name, table_name=table, schema_type=ColumnMetadata
+        )
+
+    async def get_indexes(
+        self, driver: "PsqlpyDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[IndexMetadata]":
+        """Get index metadata for a table or schema."""
+        schema_name = self.resolve_schema(schema)
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes")
+            return await driver.select(
+                self.get_query("indexes_by_schema"), schema_name=schema_name, schema_type=IndexMetadata
+            )
+
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="indexes")
+        return await driver.select(
+            self.get_query("indexes_by_table"), schema_name=schema_name, table_name=table, schema_type=IndexMetadata
+        )
+
+    async def get_foreign_keys(
+        self, driver: "PsqlpyDriver", table: "str | None" = None, schema: "str | None" = None
+    ) -> "list[ForeignKeyMetadata]":
+        """Get foreign key metadata."""
+        schema_name = self.resolve_schema(schema)
+        if table is None:
+            self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="foreign_keys")
+            return await driver.select(
+                self.get_query("foreign_keys_by_schema"), schema_name=schema_name, schema_type=ForeignKeyMetadata
+            )
+        self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="foreign_keys")
+        return await driver.select(
+            self.get_query("foreign_keys_by_table"),
+            schema_name=schema_name,
+            table_name=table,
+            schema_type=ForeignKeyMetadata,
+        )

sqlspec/adapters/psqlpy/driver.py

@@ -0,0 +1,411 @@
+"""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
+
+import psqlpy.exceptions
+
+from sqlspec.adapters.psqlpy._typing import PsqlpyCursor, PsqlpySessionContext
+from sqlspec.adapters.psqlpy.core import (
+    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,
+    normalize_scalar_parameter,
+    prepare_parameters_with_casts,
+    split_schema_and_table,
+)
+from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyDataDictionary
+from sqlspec.adapters.psqlpy.type_converter import PostgreSQLOutputConverter
+from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile
+from sqlspec.driver import AsyncDriverAdapterBase, BaseAsyncExceptionHandler
+from sqlspec.exceptions import SQLSpecError
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+    from sqlspec.adapters.psqlpy._typing import PsqlpyConnection
+    from sqlspec.core import ArrowResult
+    from sqlspec.driver import ExecutionResult
+    from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry
+
+__all__ = ("PsqlpyCursor", "PsqlpyDriver", "PsqlpyExceptionHandler", "PsqlpySessionContext")
+
+logger = get_logger("sqlspec.adapters.psqlpy")
+
+_type_converter = PostgreSQLOutputConverter()
+
+
+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, (psqlpy.exceptions.DatabaseError, psqlpy.exceptions.Error)):
+            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
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # CORE DISPATCH METHODS
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> "ExecutionResult":
+        """Execute single SQL statement.
+
+        Args:
+            cursor: Psqlpy connection object
+            statement: SQL statement to execute
+
+        Returns:
+            ExecutionResult with execution metadata
+        """
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+
+        driver_parameters = prepared_parameters
+        operation_type = statement.operation_type
+        should_coerce = operation_type != "SELECT"
+        effective_parameters = coerce_numeric_for_write(driver_parameters) if should_coerce else driver_parameters
+        params = cast("Sequence[Any] | Mapping[str, Any] | None", effective_parameters) or []
+
+        if statement.returns_rows():
+            query_result = await cursor.fetch(sql, params)
+            dict_rows, column_names = collect_rows(query_result)
+
+            return self.create_execution_result(
+                cursor,
+                selected_data=dict_rows,
+                column_names=column_names,
+                data_row_count=len(dict_rows),
+                is_select_result=True,
+                row_format="dict",
+            )
+
+        result = await cursor.execute(sql, params)
+        rows_affected = extract_rows_affected(result)
+
+        return self.create_execution_result(cursor, rowcount_override=rows_affected)
+
+    async def dispatch_execute_many(self, cursor: "PsqlpyConnection", statement: SQL) -> "ExecutionResult":
+        """Execute SQL with multiple parameter sets.
+
+        Args:
+            cursor: Psqlpy connection object
+            statement: SQL statement with multiple parameter sets
+
+        Returns:
+            ExecutionResult with batch execution metadata
+        """
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+
+        if not prepared_parameters:
+            return self.create_execution_result(cursor, rowcount_override=0, is_many_result=True)
+
+        operation_type = statement.operation_type
+        should_coerce = operation_type != "SELECT"
+        formatted_parameters = format_execute_many_parameters(prepared_parameters, coerce_numeric=should_coerce)
+
+        await cursor.execute_many(sql, formatted_parameters)
+
+        rows_affected = len(formatted_parameters)
+
+        return self.create_execution_result(cursor, rowcount_override=rows_affected, is_many_result=True)
+
+    async def dispatch_execute_script(self, cursor: "PsqlpyConnection", statement: SQL) -> "ExecutionResult":
+        """Execute SQL script with statement splitting.
+
+        Args:
+            cursor: Psqlpy connection object
+            statement: SQL statement with script content
+
+        Returns:
+            ExecutionResult with script execution metadata
+
+        Notes:
+            Uses execute() with empty parameters for each statement instead of execute_batch().
+            execute_batch() uses simple query protocol which can break subsequent queries
+            that rely on extended protocol (e.g., information_schema queries with name type).
+        """
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+        prepared_parameters = cast("Sequence[Any] | Mapping[str, Any] | None", prepared_parameters)
+        statement_config = statement.statement_config
+        statements = self.split_script_statements(sql, statement_config, strip_trailing_semicolon=True)
+
+        successful_count = 0
+        last_result = None
+
+        for stmt in statements:
+            last_result = await cursor.execute(stmt, prepared_parameters or [])
+            successful_count += 1
+
+        return self.create_execution_result(
+            last_result, statement_count=len(statements), successful_statements=successful_count, is_script_result=True
+        )
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # TRANSACTION MANAGEMENT
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    async def begin(self) -> None:
+        """Begin a database transaction."""
+        try:
+            await self.connection.execute("BEGIN")
+        except psqlpy.exceptions.DatabaseError as e:
+            msg = f"Failed to begin psqlpy transaction: {e}"
+            raise SQLSpecError(msg) from e
+
+    async def commit(self) -> None:
+        """Commit the current transaction."""
+        try:
+            await self.connection.execute("COMMIT")
+        except psqlpy.exceptions.DatabaseError as e:
+            msg = f"Failed to commit psqlpy transaction: {e}"
+            raise SQLSpecError(msg) from e
+
+    async def rollback(self) -> None:
+        """Rollback the current transaction."""
+        try:
+            await self.connection.execute("ROLLBACK")
+        except psqlpy.exceptions.DatabaseError as e:
+            msg = f"Failed to rollback psqlpy transaction: {e}"
+            raise SQLSpecError(msg) from e
+
+    def with_cursor(self, connection: "PsqlpyConnection") -> "PsqlpyCursor":
+        """Create context manager for psqlpy cursor.
+
+        Args:
+            connection: Psqlpy connection object
+
+        Returns:
+            PsqlpyCursor context manager
+        """
+        return PsqlpyCursor(connection)
+
+    def handle_database_exceptions(self) -> "PsqlpyExceptionHandler":
+        """Handle database-specific exceptions.
+
+        Returns:
+            Exception handler context manager
+        """
+        return PsqlpyExceptionHandler()
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # STORAGE API METHODS
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    async def select_to_storage(
+        self,
+        statement: "SQL | str",
+        destination: "StorageDestination",
+        /,
+        *parameters: Any,
+        statement_config: "StatementConfig | None" = None,
+        partitioner: "dict[str, object] | None" = None,
+        format_hint: "StorageFormat | None" = None,
+        telemetry: "StorageTelemetry | None" = None,
+        **kwargs: Any,
+    ) -> "StorageBridgeJob":
+        """Execute a query and stream Arrow results to a storage backend."""
+
+        self._require_capability("arrow_export_enabled")
+        arrow_result = await self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs)
+        async_pipeline = self._storage_pipeline()
+        telemetry_payload = await self._write_result_to_storage_async(
+            arrow_result, destination, format_hint=format_hint, pipeline=async_pipeline
+        )
+        self._attach_partition_telemetry(telemetry_payload, partitioner)
+        return self._create_storage_job(telemetry_payload, telemetry)
+
+    async def load_from_arrow(
+        self,
+        table: str,
+        source: "ArrowResult | Any",
+        *,
+        partitioner: "dict[str, object] | None" = None,
+        overwrite: bool = False,
+        telemetry: "StorageTelemetry | None" = None,
+    ) -> "StorageBridgeJob":
+        """Load Arrow-formatted data into PostgreSQL via psqlpy binary COPY."""
+
+        self._require_capability("arrow_import_enabled")
+        arrow_table = self._coerce_arrow_table(source)
+        if overwrite:
+            qualified = format_table_identifier(table)
+            exc_handler = self.handle_database_exceptions()
+            async with exc_handler, self.with_cursor(self.connection) as cursor:
+                await cursor.execute(f"TRUNCATE TABLE {qualified}")
+            if exc_handler.pending_exception is not None:
+                raise exc_handler.pending_exception from None
+
+        columns, records = self._arrow_table_to_rows(arrow_table)
+        if records:
+            schema_name, table_name = split_schema_and_table(table)
+            exc_handler = self.handle_database_exceptions()
+            async with exc_handler, self.with_cursor(self.connection) as cursor:
+                copy_kwargs: dict[str, Any] = {"columns": columns}
+                if schema_name:
+                    copy_kwargs["schema_name"] = schema_name
+                try:
+                    copy_payload = encode_records_for_binary_copy(records)
+                    copy_operation = cursor.binary_copy_to_table(copy_payload, table_name, **copy_kwargs)
+                    if inspect.isawaitable(copy_operation):
+                        await copy_operation
+                except (TypeError, psqlpy.exceptions.DatabaseError) as exc:
+                    logger.debug("Binary COPY not available for psqlpy; falling back to INSERT statements: %s", exc)
+                    insert_sql = build_insert_statement(table, columns)
+                    formatted_records = coerce_records_for_execute_many(records)
+                    try:
+                        insert_operation = cursor.execute_many(insert_sql, formatted_records)
+                        if inspect.isawaitable(insert_operation):
+                            await insert_operation
+                    except (psqlpy.exceptions.DatabaseError, psqlpy.exceptions.Error) as fallback_exc:
+                        if "PyJSON must be dict, list, or tuple" not in str(fallback_exc):
+                            raise
+                        formatted_records = coerce_records_for_execute_many(records, parse_json_text=True)
+                        insert_operation = cursor.execute_many(insert_sql, formatted_records)
+                        if inspect.isawaitable(insert_operation):
+                            await insert_operation
+            if exc_handler.pending_exception is not None:
+                raise exc_handler.pending_exception from None
+
+        telemetry_payload = self._build_ingest_telemetry(arrow_table)
+        telemetry_payload["destination"] = table
+        self._attach_partition_telemetry(telemetry_payload, partitioner)
+        return self._create_storage_job(telemetry_payload, telemetry)
+
+    async def load_from_storage(
+        self,
+        table: str,
+        source: "StorageDestination",
+        *,
+        file_format: "StorageFormat",
+        partitioner: "dict[str, object] | None" = None,
+        overwrite: bool = False,
+    ) -> "StorageBridgeJob":
+        """Load staged artifacts from storage using the storage bridge pipeline."""
+
+        arrow_table, inbound = await self._read_arrow_from_storage_async(source, file_format=file_format)
+        return await self.load_from_arrow(
+            table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound
+        )
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # UTILITY METHODS
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    def prepare_driver_parameters(
+        self,
+        parameters: Any,
+        statement_config: "StatementConfig",
+        is_many: bool = False,
+        prepared_statement: Any | None = None,
+    ) -> Any:
+        """Prepare parameters with cast-aware type coercion for psqlpy.
+
+        Args:
+            parameters: Parameters in any format
+            statement_config: Statement configuration
+            is_many: Whether this is for execute_many operation
+            prepared_statement: Prepared statement containing the original SQL statement
+
+        Returns:
+            Parameters with cast-aware type coercion applied
+        """
+        enable_cast_detection = self.driver_features.get("enable_cast_detection", True)
+
+        if enable_cast_detection and prepared_statement and self.dialect in {"postgres", "postgresql"} and not is_many:
+            parameter_casts = get_parameter_casts(prepared_statement)
+            prepared = prepare_parameters_with_casts(parameters, parameter_casts, statement_config)
+        else:
+            prepared = super().prepare_driver_parameters(parameters, statement_config, is_many, prepared_statement)
+
+        if not is_many and isinstance(prepared, list):
+            prepared = tuple(prepared)
+
+        if not is_many and isinstance(prepared, tuple):
+            return tuple(normalize_scalar_parameter(item) for item in prepared)
+
+        return prepared
+
+    @property
+    def data_dictionary(self) -> "PsqlpyDataDictionary":
+        """Get the data dictionary for this driver.
+
+        Returns:
+            Data dictionary instance for metadata queries
+        """
+        if self._data_dictionary is None:
+            self._data_dictionary = PsqlpyDataDictionary()
+        return self._data_dictionary
+
+    # ─────────────────────────────────────────────────────────────────────────────
+    # PRIVATE/INTERNAL METHODS
+    # ─────────────────────────────────────────────────────────────────────────────
+
+    def collect_rows(self, cursor: "PsqlpyConnection", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]":
+        """Collect psqlpy rows for the direct execution path.
+
+        The ``fetched`` argument may be a psqlpy query result or a plain list.
+        """
+        dict_rows, column_names = collect_rows(fetched)
+        return dict_rows, column_names, len(dict_rows)
+
+    def resolve_rowcount(self, cursor: "PsqlpyConnection") -> int:
+        """Resolve rowcount from psqlpy result for the direct execution path."""
+        return extract_rows_affected(cursor)
+
+    def _connection_in_transaction(self) -> bool:
+        """Check if connection is in transaction."""
+        return bool(self.connection.in_transaction())
+
+
+register_driver_profile("psqlpy", driver_profile)

sqlspec/adapters/psqlpy/events/__init__.py

@@ -0,0 +1,6 @@
+"""Events helpers for the psqlpy adapter."""
+
+from sqlspec.adapters.psqlpy.events.backend import PsqlpyEventsBackend, PsqlpyHybridEventsBackend, create_event_backend
+from sqlspec.adapters.psqlpy.events.store import PsqlpyEventQueueStore
+
+__all__ = ("PsqlpyEventQueueStore", "PsqlpyEventsBackend", "PsqlpyHybridEventsBackend", "create_event_backend")