sqlspec 0.25.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

sqlspec/adapters/asyncpg/data_dictionary.py

@@ -0,0 +1,173 @@
+"""PostgreSQL-specific data dictionary for metadata queries via asyncpg."""
+
+import re
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.driver import AsyncDataDictionaryBase, AsyncDriverAdapterBase, VersionInfo
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from sqlspec.adapters.asyncpg.driver import AsyncpgDriver
+
+logger = get_logger("adapters.asyncpg.data_dictionary")
+
+# Compiled regex patterns
+POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?")
+
+__all__ = ("PostgresAsyncDataDictionary",)
+
+
+class PostgresAsyncDataDictionary(AsyncDataDictionaryBase):
+    """PostgreSQL-specific async data dictionary."""
+
+    async def get_version(self, driver: AsyncDriverAdapterBase) -> "VersionInfo | None":
+        """Get PostgreSQL database version information.
+
+        Args:
+            driver: Async database driver instance
+
+        Returns:
+            PostgreSQL version information or None if detection fails
+        """
+        asyncpg_driver = cast("AsyncpgDriver", driver)
+        version_str = await asyncpg_driver.select_value("SELECT version()")
+        if not version_str:
+            logger.warning("No PostgreSQL version information found")
+            return None
+
+        # Parse version like "PostgreSQL 15.3 on x86_64-pc-linux-gnu..."
+        version_match = POSTGRES_VERSION_PATTERN.search(str(version_str))
+        if not version_match:
+            logger.warning("Could not parse PostgreSQL version: %s", version_str)
+            return None
+
+        major = int(version_match.group(1))
+        minor = int(version_match.group(2))
+        patch = int(version_match.group(3)) if version_match.group(3) else 0
+
+        version_info = VersionInfo(major, minor, patch)
+        logger.debug("Detected PostgreSQL version: %s", version_info)
+        return version_info
+
+    async def get_feature_flag(self, driver: AsyncDriverAdapterBase, feature: str) -> bool:
+        """Check if PostgreSQL database supports a specific feature.
+
+        Args:
+            driver: Async database driver instance
+            feature: Feature name to check
+
+        Returns:
+            True if feature is supported, False otherwise
+        """
+        version_info = await self.get_version(driver)
+        if not version_info:
+            return False
+
+        feature_checks: dict[str, Callable[..., bool]] = {
+            "supports_json": lambda v: v >= VersionInfo(9, 2, 0),
+            "supports_jsonb": lambda v: v >= VersionInfo(9, 4, 0),
+            "supports_uuid": lambda _: True,  # UUID extension widely available
+            "supports_arrays": lambda _: True,  # PostgreSQL has excellent array support
+            "supports_returning": lambda v: v >= VersionInfo(8, 2, 0),
+            "supports_upsert": lambda v: v >= VersionInfo(9, 5, 0),  # ON CONFLICT
+            "supports_window_functions": lambda v: v >= VersionInfo(8, 4, 0),
+            "supports_cte": lambda v: v >= VersionInfo(8, 4, 0),
+            "supports_transactions": lambda _: True,
+            "supports_prepared_statements": lambda _: True,
+            "supports_schemas": lambda _: True,
+            "supports_partitioning": lambda v: v >= VersionInfo(10, 0, 0),
+        }
+
+        if feature in feature_checks:
+            return bool(feature_checks[feature](version_info))
+
+        return False
+
+    async def get_optimal_type(self, driver: AsyncDriverAdapterBase, type_category: str) -> str:
+        """Get optimal PostgreSQL type for a category.
+
+        Args:
+            driver: Async database driver instance
+            type_category: Type category
+
+        Returns:
+            PostgreSQL-specific type name
+        """
+        version_info = await self.get_version(driver)
+
+        if type_category == "json":
+            if version_info and version_info >= VersionInfo(9, 4, 0):
+                return "JSONB"  # Prefer JSONB over JSON
+            if version_info and version_info >= VersionInfo(9, 2, 0):
+                return "JSON"
+            return "TEXT"
+
+        type_map = {
+            "uuid": "UUID",
+            "boolean": "BOOLEAN",
+            "timestamp": "TIMESTAMP WITH TIME ZONE",
+            "text": "TEXT",
+            "blob": "BYTEA",
+            "array": "ARRAY",
+        }
+        return type_map.get(type_category, "TEXT")
+
+    async def get_columns(
+        self, driver: AsyncDriverAdapterBase, table: str, schema: "str | None" = None
+    ) -> "list[dict[str, Any]]":
+        """Get column information for a table using information_schema.
+
+        Args:
+            driver: AsyncPG driver instance
+            table: Table name to query columns for
+            schema: Schema name (None for default 'public')
+
+        Returns:
+            List of column metadata dictionaries with keys:
+                - column_name: Name of the column
+                - data_type: PostgreSQL data type
+                - is_nullable: Whether column allows NULL (YES/NO)
+                - column_default: Default value if any
+        """
+        asyncpg_driver = cast("AsyncpgDriver", driver)
+
+        if schema:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM information_schema.columns
+                WHERE table_name = '{table}' AND table_schema = '{schema}'
+                ORDER BY ordinal_position
+            """
+        else:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM information_schema.columns
+                WHERE table_name = '{table}' AND table_schema = 'public'
+                ORDER BY ordinal_position
+            """
+
+        result = await asyncpg_driver.execute(sql)
+        return result.data or []
+
+    def list_available_features(self) -> "list[str]":
+        """List available PostgreSQL feature flags.
+
+        Returns:
+            List of supported feature names
+        """
+        return [
+            "supports_json",
+            "supports_jsonb",
+            "supports_uuid",
+            "supports_arrays",
+            "supports_returning",
+            "supports_upsert",
+            "supports_window_functions",
+            "supports_cte",
+            "supports_transactions",
+            "supports_prepared_statements",
+            "supports_schemas",
+            "supports_partitioning",
+        ]

sqlspec/adapters/asyncpg/driver.py

@@ -4,8 +4,9 @@ Provides async PostgreSQL connectivity with parameter processing, resource manag
 PostgreSQL COPY operation support, and transaction management.
 """
 
+import datetime
 import re
-from typing import TYPE_CHECKING, Any, Final, Optional
+from typing import TYPE_CHECKING, Any, Final
 
 import asyncpg
 
@@ -13,7 +14,19 @@ from sqlspec.core.cache import get_cache_config
 from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.statement import StatementConfig
 from sqlspec.driver import AsyncDriverAdapterBase
-from sqlspec.exceptions import SQLParsingError, SQLSpecError
+from sqlspec.exceptions import (
+    CheckViolationError,
+    DatabaseConnectionError,
+    DataError,
+    ForeignKeyViolationError,
+    IntegrityError,
+    NotNullViolationError,
+    OperationalError,
+    SQLParsingError,
+    SQLSpecError,
+    TransactionError,
+    UniqueViolationError,
+)
 from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
@@ -23,12 +36,55 @@ if TYPE_CHECKING:
     from sqlspec.core.result import SQLResult
     from sqlspec.core.statement import SQL
     from sqlspec.driver import ExecutionResult
+    from sqlspec.driver._async import AsyncDataDictionaryBase
 
 __all__ = ("AsyncpgCursor", "AsyncpgDriver", "AsyncpgExceptionHandler", "asyncpg_statement_config")
 
 logger = get_logger("adapters.asyncpg")
 
 
+def _convert_datetime_param(value: Any) -> Any:
+    """Convert datetime parameter, handling ISO strings.
+
+    Args:
+        value: datetime object or ISO format string
+
+    Returns:
+        datetime object for asyncpg
+    """
+    if isinstance(value, str):
+        return datetime.datetime.fromisoformat(value)
+    return value
+
+
+def _convert_date_param(value: Any) -> Any:
+    """Convert date parameter, handling ISO strings.
+
+    Args:
+        value: date object or ISO format string
+
+    Returns:
+        date object for asyncpg
+    """
+    if isinstance(value, str):
+        return datetime.date.fromisoformat(value)
+    return value
+
+
+def _convert_time_param(value: Any) -> Any:
+    """Convert time parameter, handling ISO strings.
+
+    Args:
+        value: time object or ISO format string
+
+    Returns:
+        time object for asyncpg
+    """
+    if isinstance(value, str):
+        return datetime.time.fromisoformat(value)
+    return value
+
+
 asyncpg_statement_config = StatementConfig(
     dialect="postgres",
     parameter_config=ParameterStyleConfig(
@@ -36,7 +92,11 @@ asyncpg_statement_config = StatementConfig(
         supported_parameter_styles={ParameterStyle.NUMERIC, ParameterStyle.POSITIONAL_PYFORMAT},
         default_execution_parameter_style=ParameterStyle.NUMERIC,
         supported_execution_parameter_styles={ParameterStyle.NUMERIC},
-        type_coercion_map={},
+        type_coercion_map={
+            datetime.datetime: _convert_datetime_param,
+            datetime.date: _convert_date_param,
+            datetime.time: _convert_time_param,
+        },
         has_native_list_expansion=True,
         needs_static_script_compilation=False,
         preserve_parameter_format=True,
@@ -63,12 +123,15 @@ class AsyncpgCursor:
     async def __aenter__(self) -> "AsyncpgConnection":
         return self.connection
 
-    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        _ = (exc_type, exc_val, exc_tb)
+    async def __aexit__(self, *_: Any) -> None: ...
 
 
 class AsyncpgExceptionHandler:
-    """Async context manager for handling AsyncPG database exceptions."""
+    """Async context manager for handling AsyncPG database exceptions.
+
+    Maps PostgreSQL SQLSTATE error codes to specific SQLSpec exceptions
+    for better error handling in application code.
+    """
 
     __slots__ = ()
 
@@ -79,21 +142,89 @@ class AsyncpgExceptionHandler:
         if exc_type is None:
             return
         if issubclass(exc_type, asyncpg.PostgresError):
-            e = exc_val
-            error_code = getattr(e, "sqlstate", None)
-            if error_code:
-                if error_code.startswith("23"):
-                    msg = f"PostgreSQL integrity constraint violation [{error_code}]: {e}"
-                elif error_code.startswith("42"):
-                    msg = f"PostgreSQL SQL syntax error [{error_code}]: {e}"
-                    raise SQLParsingError(msg) from e
-                elif error_code.startswith("08"):
-                    msg = f"PostgreSQL connection error [{error_code}]: {e}"
-                else:
-                    msg = f"PostgreSQL database error [{error_code}]: {e}"
-            else:
-                msg = f"PostgreSQL database error: {e}"
-            raise SQLSpecError(msg) from e
+            self._map_postgres_exception(exc_val)
+
+    def _map_postgres_exception(self, e: Any) -> None:
+        """Map PostgreSQL exception to SQLSpec exception.
+
+        Args:
+            e: asyncpg.PostgresError instance
+
+        Raises:
+            Specific SQLSpec exception based on SQLSTATE code
+        """
+        error_code = getattr(e, "sqlstate", None)
+
+        if not error_code:
+            self._raise_generic_error(e, None)
+            return
+
+        if error_code == "23505":
+            self._raise_unique_violation(e, error_code)
+        elif error_code == "23503":
+            self._raise_foreign_key_violation(e, error_code)
+        elif error_code == "23502":
+            self._raise_not_null_violation(e, error_code)
+        elif error_code == "23514":
+            self._raise_check_violation(e, error_code)
+        elif error_code.startswith("23"):
+            self._raise_integrity_error(e, error_code)
+        elif error_code.startswith("42"):
+            self._raise_parsing_error(e, error_code)
+        elif error_code.startswith("08"):
+            self._raise_connection_error(e, error_code)
+        elif error_code.startswith("40"):
+            self._raise_transaction_error(e, error_code)
+        elif error_code.startswith("22"):
+            self._raise_data_error(e, error_code)
+        elif error_code.startswith(("53", "54", "55", "57", "58")):
+            self._raise_operational_error(e, error_code)
+        else:
+            self._raise_generic_error(e, error_code)
+
+    def _raise_unique_violation(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL unique constraint violation [{code}]: {e}"
+        raise UniqueViolationError(msg) from e
+
+    def _raise_foreign_key_violation(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL foreign key constraint violation [{code}]: {e}"
+        raise ForeignKeyViolationError(msg) from e
+
+    def _raise_not_null_violation(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL not-null constraint violation [{code}]: {e}"
+        raise NotNullViolationError(msg) from e
+
+    def _raise_check_violation(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL check constraint violation [{code}]: {e}"
+        raise CheckViolationError(msg) from e
+
+    def _raise_integrity_error(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL integrity constraint violation [{code}]: {e}"
+        raise IntegrityError(msg) from e
+
+    def _raise_parsing_error(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL SQL syntax error [{code}]: {e}"
+        raise SQLParsingError(msg) from e
+
+    def _raise_connection_error(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL connection error [{code}]: {e}"
+        raise DatabaseConnectionError(msg) from e
+
+    def _raise_transaction_error(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL transaction error [{code}]: {e}"
+        raise TransactionError(msg) from e
+
+    def _raise_data_error(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL data error [{code}]: {e}"
+        raise DataError(msg) from e
+
+    def _raise_operational_error(self, e: Any, code: str) -> None:
+        msg = f"PostgreSQL operational error [{code}]: {e}"
+        raise OperationalError(msg) from e
+
+    def _raise_generic_error(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL database error [{code}]: {e}" if code else f"PostgreSQL database error: {e}"
+        raise SQLSpecError(msg) from e
 
 
 class AsyncpgDriver(AsyncDriverAdapterBase):
@@ -104,14 +235,14 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
     and caching, and parameter processing with type coercion.
     """
 
-    __slots__ = ()
+    __slots__ = ("_data_dictionary",)
     dialect = "postgres"
 
     def __init__(
         self,
         connection: "AsyncpgConnection",
-        statement_config: "Optional[StatementConfig]" = None,
-        driver_features: "Optional[dict[str, Any]]" = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "dict[str, Any] | None" = None,
     ) -> None:
         if statement_config is None:
             cache_config = get_cache_config()
@@ -123,6 +254,7 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
             )
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
+        self._data_dictionary: AsyncDataDictionaryBase | None = None
 
     def with_cursor(self, connection: "AsyncpgConnection") -> "AsyncpgCursor":
         """Create context manager for AsyncPG cursor."""
@@ -132,7 +264,7 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         """Handle database exceptions with PostgreSQL error codes."""
         return AsyncpgExceptionHandler()
 
-    async def _try_special_handling(self, cursor: "AsyncpgConnection", statement: "SQL") -> "Optional[SQLResult]":
+    async def _try_special_handling(self, cursor: "AsyncpgConnection", statement: "SQL") -> "SQLResult | None":
         """Handle PostgreSQL COPY operations and other special cases.
 
         Args:
@@ -311,3 +443,16 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         except asyncpg.PostgresError as e:
             msg = f"Failed to commit async transaction: {e}"
             raise SQLSpecError(msg) from e
+
+    @property
+    def data_dictionary(self) -> "AsyncDataDictionaryBase":
+        """Get the data dictionary for this driver.
+
+        Returns:
+            Data dictionary instance for metadata queries
+        """
+        if self._data_dictionary is None:
+            from sqlspec.adapters.asyncpg.data_dictionary import PostgresAsyncDataDictionary
+
+            self._data_dictionary = PostgresAsyncDataDictionary()
+        return self._data_dictionary

sqlspec/adapters/asyncpg/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""Litestar integration for AsyncPG adapter."""
+
+from sqlspec.adapters.asyncpg.litestar.store import AsyncpgStore
+
+__all__ = ("AsyncpgStore",)