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

sqlspec/adapters/psycopg/driver.py

@@ -14,8 +14,9 @@ PostgreSQL Features:
 - PostgreSQL-specific error handling
 """
 
+import datetime
 import io
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 import psycopg
 
@@ -25,14 +26,28 @@ from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.result import SQLResult
 from sqlspec.core.statement import SQL, StatementConfig
 from sqlspec.driver import AsyncDriverAdapterBase, SyncDriverAdapterBase
-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
 from sqlspec.utils.serializers import to_json
 
 if TYPE_CHECKING:
     from contextlib import AbstractAsyncContextManager, AbstractContextManager
 
+    from sqlspec.driver._async import AsyncDataDictionaryBase
     from sqlspec.driver._common import ExecutionResult
+    from sqlspec.driver._sync import SyncDataDictionaryBase
 
 logger = get_logger("adapters.psycopg")
 
@@ -94,7 +109,12 @@ psycopg_statement_config = StatementConfig(
             ParameterStyle.NAMED_PYFORMAT,
             ParameterStyle.NUMERIC,
         },
-        type_coercion_map={dict: to_json},
+        type_coercion_map={
+            dict: to_json,
+            datetime.datetime: lambda x: x,
+            datetime.date: lambda x: x,
+            datetime.time: lambda x: x,
+        },
         has_native_list_expansion=True,
         needs_static_script_compilation=False,
         preserve_parameter_format=True,
@@ -119,20 +139,23 @@ class PsycopgSyncCursor:
 
     def __init__(self, connection: PsycopgSyncConnection) -> None:
         self.connection = connection
-        self.cursor: Optional[Any] = None
+        self.cursor: Any | None = None
 
     def __enter__(self) -> Any:
         self.cursor = self.connection.cursor()
         return self.cursor
 
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        _ = (exc_type, exc_val, exc_tb)
+    def __exit__(self, *_: Any) -> None:
         if self.cursor is not None:
             self.cursor.close()
 
 
 class PsycopgSyncExceptionHandler:
-    """Context manager for handling PostgreSQL psycopg database exceptions."""
+    """Context manager for handling PostgreSQL psycopg database exceptions.
+
+    Maps PostgreSQL SQLSTATE error codes to specific SQLSpec exceptions
+    for better error handling in application code.
+    """
 
     __slots__ = ()
 
@@ -142,39 +165,90 @@ class PsycopgSyncExceptionHandler:
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         if exc_type is None:
             return
-
-        if issubclass(exc_type, psycopg.IntegrityError):
-            e = exc_val
-            msg = f"PostgreSQL psycopg integrity constraint violation: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, psycopg.ProgrammingError):
-            e = exc_val
-            error_msg = str(e).lower()
-            if "syntax" in error_msg or "parse" in error_msg:
-                msg = f"PostgreSQL psycopg SQL syntax error: {e}"
-                raise SQLParsingError(msg) from e
-            msg = f"PostgreSQL psycopg programming error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, psycopg.OperationalError):
-            e = exc_val
-            msg = f"PostgreSQL psycopg operational error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, psycopg.DatabaseError):
-            e = exc_val
-            msg = f"PostgreSQL psycopg database error: {e}"
-            raise SQLSpecError(msg) from e
         if issubclass(exc_type, psycopg.Error):
-            e = exc_val
-            msg = f"PostgreSQL psycopg 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 database operation 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: psycopg.Error 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 PsycopgSyncDriver(SyncDriverAdapterBase):
@@ -187,14 +261,14 @@ class PsycopgSyncDriver(SyncDriverAdapterBase):
     bulk data transfer, and PostgreSQL-specific error handling.
     """
 
-    __slots__ = ()
+    __slots__ = ("_data_dictionary",)
     dialect = "postgres"
 
     def __init__(
         self,
         connection: PsycopgSyncConnection,
-        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()
@@ -207,6 +281,7 @@ class PsycopgSyncDriver(SyncDriverAdapterBase):
             statement_config = default_config
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
+        self._data_dictionary: SyncDataDictionaryBase | None = None
 
     def with_cursor(self, connection: PsycopgSyncConnection) -> PsycopgSyncCursor:
         """Create context manager for PostgreSQL cursor."""
@@ -255,7 +330,7 @@ class PsycopgSyncDriver(SyncDriverAdapterBase):
         except Exception as cleanup_error:
             logger.warning("Failed to cleanup transaction state: %s", cleanup_error)
 
-    def _try_special_handling(self, cursor: Any, statement: "SQL") -> "Optional[SQLResult]":
+    def _try_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult | None":
         """Hook for PostgreSQL-specific special operations.
 
         Args:
@@ -411,6 +486,19 @@ class PsycopgSyncDriver(SyncDriverAdapterBase):
         affected_rows = cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
         return self.create_execution_result(cursor, rowcount_override=affected_rows)
 
+    @property
+    def data_dictionary(self) -> "SyncDataDictionaryBase":
+        """Get the data dictionary for this driver.
+
+        Returns:
+            Data dictionary instance for metadata queries
+        """
+        if self._data_dictionary is None:
+            from sqlspec.adapters.psycopg.data_dictionary import PostgresSyncDataDictionary
+
+            self._data_dictionary = PostgresSyncDataDictionary()
+        return self._data_dictionary
+
 
 class PsycopgAsyncCursor:
     """Async context manager for PostgreSQL psycopg cursor management."""
@@ -419,7 +507,7 @@ class PsycopgAsyncCursor:
 
     def __init__(self, connection: "PsycopgAsyncConnection") -> None:
         self.connection = connection
-        self.cursor: Optional[Any] = None
+        self.cursor: Any | None = None
 
     async def __aenter__(self) -> Any:
         self.cursor = self.connection.cursor()
@@ -432,7 +520,11 @@ class PsycopgAsyncCursor:
 
 
 class PsycopgAsyncExceptionHandler:
-    """Async context manager for handling PostgreSQL psycopg database exceptions."""
+    """Async context manager for handling PostgreSQL psycopg database exceptions.
+
+    Maps PostgreSQL SQLSTATE error codes to specific SQLSpec exceptions
+    for better error handling in application code.
+    """
 
     __slots__ = ()
 
@@ -442,39 +534,90 @@ class PsycopgAsyncExceptionHandler:
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         if exc_type is None:
             return
-
-        if issubclass(exc_type, psycopg.IntegrityError):
-            e = exc_val
-            msg = f"PostgreSQL psycopg integrity constraint violation: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, psycopg.ProgrammingError):
-            e = exc_val
-            error_msg = str(e).lower()
-            if "syntax" in error_msg or "parse" in error_msg:
-                msg = f"PostgreSQL psycopg SQL syntax error: {e}"
-                raise SQLParsingError(msg) from e
-            msg = f"PostgreSQL psycopg programming error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, psycopg.OperationalError):
-            e = exc_val
-            msg = f"PostgreSQL psycopg operational error: {e}"
-            raise SQLSpecError(msg) from e
-        if issubclass(exc_type, psycopg.DatabaseError):
-            e = exc_val
-            msg = f"PostgreSQL psycopg database error: {e}"
-            raise SQLSpecError(msg) from e
         if issubclass(exc_type, psycopg.Error):
-            e = exc_val
-            msg = f"PostgreSQL psycopg 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
+            self._map_postgres_exception(exc_val)
+
+    def _map_postgres_exception(self, e: Any) -> None:
+        """Map PostgreSQL exception to SQLSpec exception.
+
+        Args:
+            e: psycopg.Error 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 PsycopgAsyncDriver(AsyncDriverAdapterBase):
@@ -488,14 +631,14 @@ class PsycopgAsyncDriver(AsyncDriverAdapterBase):
     and async pub/sub support.
     """
 
-    __slots__ = ()
+    __slots__ = ("_data_dictionary",)
     dialect = "postgres"
 
     def __init__(
         self,
         connection: "PsycopgAsyncConnection",
-        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()
@@ -508,6 +651,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterBase):
             statement_config = default_config
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
+        self._data_dictionary: AsyncDataDictionaryBase | None = None
 
     def with_cursor(self, connection: "PsycopgAsyncConnection") -> "PsycopgAsyncCursor":
         """Create async context manager for PostgreSQL cursor."""
@@ -556,7 +700,7 @@ class PsycopgAsyncDriver(AsyncDriverAdapterBase):
         except Exception as cleanup_error:
             logger.warning("Failed to cleanup transaction state: %s", cleanup_error)
 
-    async def _try_special_handling(self, cursor: Any, statement: "SQL") -> "Optional[SQLResult]":
+    async def _try_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult | None":
         """Hook for PostgreSQL-specific special operations.
 
         Args:
@@ -714,3 +858,16 @@ class PsycopgAsyncDriver(AsyncDriverAdapterBase):
 
         affected_rows = cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
         return self.create_execution_result(cursor, rowcount_override=affected_rows)
+
+    @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.psycopg.data_dictionary import PostgresAsyncDataDictionary
+
+            self._data_dictionary = PostgresAsyncDataDictionary()
+        return self._data_dictionary

sqlspec/adapters/psycopg/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""Litestar integration for Psycopg adapter."""
+
+from sqlspec.adapters.psycopg.litestar.store import PsycopgAsyncStore, PsycopgSyncStore
+
+__all__ = ("PsycopgAsyncStore", "PsycopgSyncStore")