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

sqlspec/adapters/psqlpy/data_dictionary.py

@@ -0,0 +1,172 @@
+"""PostgreSQL-specific data dictionary for metadata queries via psqlpy."""
+
+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.psqlpy.driver import PsqlpyDriver
+
+logger = get_logger("adapters.psqlpy.data_dictionary")
+
+# Compiled regex patterns
+POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?")
+
+__all__ = ("PsqlpyAsyncDataDictionary",)
+
+
+class PsqlpyAsyncDataDictionary(AsyncDataDictionaryBase):
+    """PostgreSQL-specific async data dictionary via psqlpy."""
+
+    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
+        """
+        version_str = await cast("PsqlpyDriver", 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[[VersionInfo], 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: Psqlpy async 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
+        """
+        psqlpy_driver = cast("PsqlpyDriver", 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 psqlpy_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/psqlpy/driver.py

@@ -4,20 +4,31 @@ Provides parameter style conversion, type coercion, error handling,
 and transaction management.
 """
 
-import datetime
 import decimal
 import re
-import uuid
-from typing import TYPE_CHECKING, Any, Final, Optional
+from typing import TYPE_CHECKING, Any, Final
 
 import psqlpy
 import psqlpy.exceptions
 
+from sqlspec.adapters.psqlpy.type_converter import PostgreSQLTypeConverter
 from sqlspec.core.cache import get_cache_config
 from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.statement import SQL, 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:
@@ -26,11 +37,14 @@ if TYPE_CHECKING:
     from sqlspec.adapters.psqlpy._types import PsqlpyConnection
     from sqlspec.core.result import SQLResult
     from sqlspec.driver import ExecutionResult
+    from sqlspec.driver._async import AsyncDataDictionaryBase
 
 __all__ = ("PsqlpyCursor", "PsqlpyDriver", "PsqlpyExceptionHandler", "psqlpy_statement_config")
 
 logger = get_logger("adapters.psqlpy")
 
+_type_converter = PostgreSQLTypeConverter()
+
 psqlpy_statement_config = StatementConfig(
     dialect="postgres",
     parameter_config=ParameterStyleConfig(
@@ -38,7 +52,7 @@ psqlpy_statement_config = StatementConfig(
         supported_parameter_styles={ParameterStyle.NUMERIC, ParameterStyle.NAMED_DOLLAR, ParameterStyle.QMARK},
         default_execution_parameter_style=ParameterStyle.NUMERIC,
         supported_execution_parameter_styles={ParameterStyle.NUMERIC},
-        type_coercion_map={tuple: list, decimal.Decimal: float},
+        type_coercion_map={tuple: list, decimal.Decimal: float, str: _type_converter.convert_if_detected},
         has_native_list_expansion=False,
         needs_static_script_compilation=False,
         allow_mixed_parameter_styles=False,
@@ -52,173 +66,6 @@ psqlpy_statement_config = StatementConfig(
 
 PSQLPY_STATUS_REGEX: Final[re.Pattern[str]] = re.compile(r"^([A-Z]+)(?:\s+(\d+))?\s+(\d+)$", re.IGNORECASE)
 
-SPECIAL_TYPE_REGEX: Final[re.Pattern[str]] = re.compile(
-    r"^(?:"
-    r"(?P<uuid>[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})|"
-    r"(?P<ipv4>(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)(?:/(?:3[0-2]|[12]?[0-9]))?)|"
-    r"(?P<ipv6>(?:(?:[0-9a-f]{1,4}:){7}[0-9a-f]{1,4}|(?:[0-9a-f]{1,4}:){1,7}:|:(?::[0-9a-f]{1,4}){1,7}|(?:[0-9a-f]{1,4}:){1,6}:[0-9a-f]{1,4}|::(?:ffff:)?(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?))(?:/(?:12[0-8]|1[01][0-9]|[1-9]?[0-9]))?)|"
-    r"(?P<mac>(?:[0-9a-f]{2}[:-]){5}[0-9a-f]{2}|[0-9a-f]{12})|"
-    r"(?P<iso_datetime>\d{4}-\d{2}-\d{2}[T\s]\d{2}:\d{2}:\d{2}(?:\.\d{1,6})?(?:Z|[+-]\d{2}:?\d{2})?)|"
-    r"(?P<iso_date>\d{4}-\d{2}-\d{2})|"
-    r"(?P<iso_time>\d{2}:\d{2}:\d{2}(?:\.\d{1,6})?(?:Z|[+-]\d{2}:?\d{2})?)|"
-    r"(?P<interval>(?:(?:\d+\s+(?:year|month|day|hour|minute|second)s?\s*)+)|(?:P(?:\d+Y)?(?:\d+M)?(?:\d+D)?(?:T(?:\d+H)?(?:\d+M)?(?:\d+(?:\.\d+)?S)?)?))|"
-    r"(?P<json>\{[\s\S]*\}|\[[\s\S]*\])|"
-    r"(?P<pg_array>\{(?:[^{}]+|\{[^{}]*\})*\})"
-    r")$",
-    re.IGNORECASE,
-)
-
-
-def _detect_postgresql_type(value: str) -> Optional[str]:
-    """Detect PostgreSQL data type from string value.
-
-    Args:
-        value: String value to analyze
-
-    Returns:
-        Type name if detected, None otherwise.
-    """
-    match = SPECIAL_TYPE_REGEX.match(value)
-    if not match:
-        return None
-
-    for group_name in [
-        "uuid",
-        "ipv4",
-        "ipv6",
-        "mac",
-        "iso_datetime",
-        "iso_date",
-        "iso_time",
-        "interval",
-        "json",
-        "pg_array",
-    ]:
-        if match.group(group_name):
-            return group_name
-
-    return None
-
-
-def _convert_uuid(value: str) -> Any:
-    """Convert UUID string to UUID object.
-
-    Args:
-        value: UUID string to convert
-
-    Returns:
-        UUID object or original value if conversion fails
-    """
-    try:
-        clean_uuid = value.replace("-", "").lower()
-        uuid_length = 32
-        if len(clean_uuid) == uuid_length:
-            formatted = f"{clean_uuid[:8]}-{clean_uuid[8:12]}-{clean_uuid[12:16]}-{clean_uuid[16:20]}-{clean_uuid[20:]}"
-            return uuid.UUID(formatted)
-        return uuid.UUID(value)
-    except (ValueError, AttributeError):
-        return value
-
-
-def _convert_iso_datetime(value: str) -> Any:
-    """Convert ISO datetime string to datetime object.
-
-    Args:
-        value: ISO datetime string to convert
-
-    Returns:
-        datetime object or original value if conversion fails
-    """
-    try:
-        normalized = value.replace("Z", "+00:00")
-        return datetime.datetime.fromisoformat(normalized)
-    except ValueError:
-        return value
-
-
-def _convert_iso_date(value: str) -> Any:
-    """Convert ISO date string to date object.
-
-    Args:
-        value: ISO date string to convert
-
-    Returns:
-        date object or original value if conversion fails
-    """
-    try:
-        return datetime.date.fromisoformat(value)
-    except ValueError:
-        return value
-
-
-def _validate_json(value: str) -> str:
-    """Validate JSON string format.
-
-    Args:
-        value: JSON string to validate
-
-    Returns:
-        Original string value
-    """
-    from sqlspec.utils.serializers import from_json
-
-    try:
-        from_json(value)
-    except (ValueError, TypeError):
-        return value
-    return value
-
-
-def _passthrough(value: str) -> str:
-    """Pass value through unchanged.
-
-    Args:
-        value: String value to pass through
-
-    Returns:
-        Original value unchanged
-    """
-    return value
-
-
-_PSQLPY_TYPE_CONVERTERS: dict[str, Any] = {
-    "uuid": _convert_uuid,
-    "iso_datetime": _convert_iso_datetime,
-    "iso_date": _convert_iso_date,
-    "iso_time": _passthrough,
-    "json": _validate_json,
-    "pg_array": _passthrough,
-    "ipv4": _passthrough,
-    "ipv6": _passthrough,
-    "mac": _passthrough,
-    "interval": _passthrough,
-}
-
-
-def _convert_psqlpy_parameters(value: Any) -> Any:
-    """Convert parameters for psqlpy compatibility.
-
-    Args:
-        value: Parameter value to convert
-
-    Returns:
-        Converted value suitable for psqlpy execution
-    """
-    if isinstance(value, str):
-        detected_type = _detect_postgresql_type(value)
-
-        if detected_type:
-            converter = _PSQLPY_TYPE_CONVERTERS.get(detected_type)
-            if converter:
-                return converter(value)
-
-        return value
-
-    if isinstance(value, (dict, list, tuple, uuid.UUID, datetime.datetime, datetime.date)):
-        return value
-
-    return value
-
 
 class PsqlpyCursor:
     """Context manager for psqlpy cursor management."""
@@ -238,7 +85,7 @@ class PsqlpyCursor:
         self._in_use = True
         return self.connection
 
-    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+    async def __aexit__(self, *_: Any) -> None:
         """Exit cursor context.
 
         Args:
@@ -246,7 +93,6 @@ class PsqlpyCursor:
             exc_val: Exception value
             exc_tb: Exception traceback
         """
-        _ = (exc_type, exc_val, exc_tb)
         self._in_use = False
 
     def is_in_use(self) -> bool:
@@ -259,7 +105,11 @@ class PsqlpyCursor:
 
 
 class PsqlpyExceptionHandler:
-    """Async context manager for handling psqlpy database exceptions."""
+    """Async context manager for handling psqlpy database exceptions.
+
+    Maps PostgreSQL SQLSTATE error codes to specific SQLSpec exceptions
+    for better error handling in application code.
+    """
 
     __slots__ = ()
 
@@ -269,31 +119,85 @@ class PsqlpyExceptionHandler:
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         if exc_type is None:
             return
+        if issubclass(exc_type, (psqlpy.exceptions.DatabaseError, psqlpy.exceptions.Error)):
+            self._map_postgres_exception(exc_val)
 
-        if issubclass(exc_type, psqlpy.exceptions.DatabaseError):
-            e = exc_val
-            msg = f"Psqlpy database error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, psqlpy.exceptions.InterfaceError):
-            e = exc_val
-            msg = f"Psqlpy interface error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, psqlpy.exceptions.Error):
-            e = exc_val
-            error_msg = str(e).lower()
-            if "syntax" in error_msg or "parse" in error_msg:
-                msg = f"Psqlpy SQL syntax error: {e}"
-                raise SQLParsingError(msg) from e
-            msg = f"Psqlpy error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, Exception):
-            e = exc_val
-            error_msg = str(e).lower()
-            if "parse" in error_msg or "syntax" in error_msg:
-                msg = f"SQL parsing failed: {e}"
-                raise SQLParsingError(msg) from e
-            msg = f"Unexpected async database operation error: {e}"
-            raise SQLSpecError(msg) from e
+    def _map_postgres_exception(self, e: Any) -> None:
+        """Map PostgreSQL exception to SQLSpec exception.
+
+        psqlpy does not expose SQLSTATE codes directly, so we use message-based
+        detection to map exceptions.
+
+        Args:
+            e: psqlpy exception instance
+
+        Raises:
+            Specific SQLSpec exception based on error message patterns
+        """
+        error_msg = str(e).lower()
+
+        if "unique" in error_msg or "duplicate key" in error_msg:
+            self._raise_unique_violation(e, None)
+        elif "foreign key" in error_msg or "violates foreign key" in error_msg:
+            self._raise_foreign_key_violation(e, None)
+        elif "not null" in error_msg or ("null value" in error_msg and "violates not-null" in error_msg):
+            self._raise_not_null_violation(e, None)
+        elif "check constraint" in error_msg or "violates check constraint" in error_msg:
+            self._raise_check_violation(e, None)
+        elif "constraint" in error_msg:
+            self._raise_integrity_error(e, None)
+        elif "syntax error" in error_msg or "parse" in error_msg:
+            self._raise_parsing_error(e, None)
+        elif "connection" in error_msg or "could not connect" in error_msg:
+            self._raise_connection_error(e, None)
+        elif "deadlock" in error_msg or "serialization failure" in error_msg:
+            self._raise_transaction_error(e, None)
+        else:
+            self._raise_generic_error(e, None)
+
+    def _raise_unique_violation(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL unique constraint violation: {e}"
+        raise UniqueViolationError(msg) from e
+
+    def _raise_foreign_key_violation(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL foreign key constraint violation: {e}"
+        raise ForeignKeyViolationError(msg) from e
+
+    def _raise_not_null_violation(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL not-null constraint violation: {e}"
+        raise NotNullViolationError(msg) from e
+
+    def _raise_check_violation(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL check constraint violation: {e}"
+        raise CheckViolationError(msg) from e
+
+    def _raise_integrity_error(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL integrity constraint violation: {e}"
+        raise IntegrityError(msg) from e
+
+    def _raise_parsing_error(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL SQL syntax error: {e}"
+        raise SQLParsingError(msg) from e
+
+    def _raise_connection_error(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL connection error: {e}"
+        raise DatabaseConnectionError(msg) from e
+
+    def _raise_transaction_error(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL transaction error: {e}"
+        raise TransactionError(msg) from e
+
+    def _raise_data_error(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL data error: {e}"
+        raise DataError(msg) from e
+
+    def _raise_operational_error(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL operational error: {e}"
+        raise OperationalError(msg) from e
+
+    def _raise_generic_error(self, e: Any, code: "str | None") -> None:
+        msg = f"PostgreSQL database error: {e}"
+        raise SQLSpecError(msg) from e
 
 
 class PsqlpyDriver(AsyncDriverAdapterBase):
@@ -303,14 +207,14 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
     and transaction management.
     """
 
-    __slots__ = ()
+    __slots__ = ("_data_dictionary",)
     dialect = "postgres"
 
     def __init__(
         self,
         connection: "PsqlpyConnection",
-        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()
@@ -322,6 +226,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
             )
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
+        self._data_dictionary: AsyncDataDictionaryBase | None = None
 
     def with_cursor(self, connection: "PsqlpyConnection") -> "PsqlpyCursor":
         """Create context manager for psqlpy cursor.
@@ -342,7 +247,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
         """
         return PsqlpyExceptionHandler()
 
-    async def _try_special_handling(self, cursor: "PsqlpyConnection", statement: SQL) -> "Optional[SQLResult]":
+    async def _try_special_handling(self, cursor: "PsqlpyConnection", statement: SQL) -> "SQLResult | None":
         """Hook for psqlpy-specific special operations.
 
         Args:
@@ -404,10 +309,9 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
         formatted_parameters = []
         for param_set in prepared_parameters:
             if isinstance(param_set, (list, tuple)):
-                converted_params = [_convert_psqlpy_parameters(param) for param in param_set]
-                formatted_parameters.append(converted_params)
+                formatted_parameters.append(list(param_set))
             else:
-                formatted_parameters.append([_convert_psqlpy_parameters(param_set)])
+                formatted_parameters.append([param_set])
 
         await cursor.execute_many(sql, formatted_parameters)
 
@@ -427,9 +331,6 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
-        if prepared_parameters:
-            prepared_parameters = [_convert_psqlpy_parameters(param) for param in prepared_parameters]
-
         if statement.returns_rows():
             query_result = await cursor.fetch(sql, prepared_parameters or [])
             dict_rows: list[dict[str, Any]] = query_result.result() if query_result else []
@@ -511,3 +412,16 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
         except psqlpy.exceptions.DatabaseError as e:
             msg = f"Failed to commit psqlpy 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.psqlpy.data_dictionary import PsqlpyAsyncDataDictionary
+
+            self._data_dictionary = PsqlpyAsyncDataDictionary()
+        return self._data_dictionary

sqlspec/adapters/psqlpy/litestar/__init__.py

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