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

sqlspec/adapters/psqlpy/data_dictionary.py

@@ -1,12 +1,14 @@
 """PostgreSQL-specific data dictionary for metadata queries via psqlpy."""
 
 import re
-from typing import TYPE_CHECKING, Callable, Optional, cast
+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")
@@ -20,7 +22,7 @@ __all__ = ("PsqlpyAsyncDataDictionary",)
 class PsqlpyAsyncDataDictionary(AsyncDataDictionaryBase):
     """PostgreSQL-specific async data dictionary via psqlpy."""
 
-    async def get_version(self, driver: AsyncDriverAdapterBase) -> "Optional[VersionInfo]":
+    async def get_version(self, driver: AsyncDriverAdapterBase) -> "VersionInfo | None":
         """Get PostgreSQL database version information.
 
         Args:
@@ -111,6 +113,43 @@ class PsqlpyAsyncDataDictionary(AsyncDataDictionaryBase):
         }
         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.
 

sqlspec/adapters/psqlpy/driver.py

@@ -6,7 +6,7 @@ and transaction management.
 
 import decimal
 import re
-from typing import TYPE_CHECKING, Any, Final, Optional
+from typing import TYPE_CHECKING, Any, Final
 
 import psqlpy
 import psqlpy.exceptions
@@ -16,7 +16,19 @@ 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:
@@ -93,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__ = ()
 
@@ -103,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):
@@ -143,8 +213,8 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
     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()
@@ -156,7 +226,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase):
             )
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
-        self._data_dictionary: Optional[AsyncDataDictionaryBase] = None
+        self._data_dictionary: AsyncDataDictionaryBase | None = None
 
     def with_cursor(self, connection: "PsqlpyConnection") -> "PsqlpyCursor":
         """Create context manager for psqlpy cursor.
@@ -177,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:

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",)

sqlspec/adapters/psqlpy/litestar/store.py

@@ -0,0 +1,272 @@
+"""Psqlpy session store for Litestar integration."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.psqlpy.config import PsqlpyConfig
+
+logger = get_logger("adapters.psqlpy.litestar.store")
+
+__all__ = ("PsqlpyStore",)
+
+
+class PsqlpyStore(BaseSQLSpecStore["PsqlpyConfig"]):
+    """PostgreSQL session store using Psqlpy driver.
+
+    Implements server-side session storage for Litestar using PostgreSQL
+    via the Psqlpy driver (Rust-based async driver). Provides efficient
+    session management with:
+    - Native async PostgreSQL operations via Rust
+    - UPSERT support using ON CONFLICT
+    - Automatic expiration handling
+    - Efficient cleanup of expired sessions
+
+    Args:
+        config: PsqlpyConfig instance.
+
+    Example:
+        from sqlspec.adapters.psqlpy import PsqlpyConfig
+        from sqlspec.adapters.psqlpy.litestar.store import PsqlpyStore
+
+        config = PsqlpyConfig(pool_config={"dsn": "postgresql://..."})
+        store = PsqlpyStore(config)
+        await store.create_table()
+    """
+
+    __slots__ = ()
+
+    def __init__(self, config: "PsqlpyConfig") -> None:
+        """Initialize Psqlpy session store.
+
+        Args:
+            config: PsqlpyConfig instance.
+
+        Notes:
+            Table name is read from config.extension_config["litestar"]["session_table"].
+        """
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        """Get PostgreSQL CREATE TABLE SQL with optimized schema.
+
+        Returns:
+            SQL statement to create the sessions table with proper indexes.
+
+        Notes:
+            - Uses TIMESTAMPTZ for timezone-aware expiration timestamps
+            - Partial index WHERE expires_at IS NOT NULL reduces index size/maintenance
+            - FILLFACTOR 80 leaves space for HOT updates, reducing table bloat
+            - Audit columns (created_at, updated_at) help with debugging
+            - Table name is internally controlled, not user input (S608 suppressed)
+        """
+        return f"""
+        CREATE TABLE IF NOT EXISTS {self._table_name} (
+            session_id TEXT PRIMARY KEY,
+            data BYTEA NOT NULL,
+            expires_at TIMESTAMPTZ,
+            created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+        ) WITH (fillfactor = 80);
+
+        CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at
+        ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL;
+
+        ALTER TABLE {self._table_name} SET (
+            autovacuum_vacuum_scale_factor = 0.05,
+            autovacuum_analyze_scale_factor = 0.02
+        );
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        """Get PostgreSQL DROP TABLE SQL statements.
+
+        Returns:
+            List of SQL statements to drop indexes and table.
+        """
+        return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"]
+
+    async def create_table(self) -> None:
+        """Create the session table if it doesn't exist."""
+        sql = self._get_create_table_sql()
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(sql)
+        logger.debug("Created session table: %s", self._table_name)
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        """Get a session value by key.
+
+        Args:
+            key: Session ID to retrieve.
+            renew_for: If given, renew the expiry time for this duration.
+
+        Returns:
+            Session data as bytes if found and not expired, None otherwise.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP instead of NOW() for SQL standard compliance.
+            The query planner can use the partial index for expires_at > CURRENT_TIMESTAMP.
+        """
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = $1
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
+        """
+
+        async with self._config.provide_connection() as conn:
+            query_result = await conn.fetch(sql, [key])
+            rows = query_result.result()
+
+            if not rows:
+                return None
+
+            row = rows[0]
+
+            if renew_for is not None and row["expires_at"] is not None:
+                new_expires_at = self._calculate_expires_at(renew_for)
+                if new_expires_at is not None:
+                    update_sql = f"""
+                    UPDATE {self._table_name}
+                    SET expires_at = $1, updated_at = CURRENT_TIMESTAMP
+                    WHERE session_id = $2
+                    """
+                    await conn.execute(update_sql, [new_expires_at, key])
+
+            return bytes(row["data"])
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        """Store a session value.
+
+        Args:
+            key: Session ID.
+            value: Session data.
+            expires_in: Time until expiration.
+
+        Notes:
+            Uses EXCLUDED to reference the proposed insert values in ON CONFLICT.
+            Updates updated_at timestamp on every write for audit trail.
+        """
+        data = self._value_to_bytes(value)
+        expires_at = self._calculate_expires_at(expires_in)
+
+        sql = f"""
+        INSERT INTO {self._table_name} (session_id, data, expires_at)
+        VALUES ($1, $2, $3)
+        ON CONFLICT (session_id)
+        DO UPDATE SET
+            data = EXCLUDED.data,
+            expires_at = EXCLUDED.expires_at,
+            updated_at = CURRENT_TIMESTAMP
+        """
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql, [key, data, expires_at])
+
+    async def delete(self, key: str) -> None:
+        """Delete a session by key.
+
+        Args:
+            key: Session ID to delete.
+        """
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = $1"
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql, [key])
+
+    async def delete_all(self) -> None:
+        """Delete all sessions from the store."""
+        sql = f"DELETE FROM {self._table_name}"
+
+        async with self._config.provide_connection() as conn:
+            await conn.execute(sql)
+        logger.debug("Deleted all sessions from table: %s", self._table_name)
+
+    async def exists(self, key: str) -> bool:
+        """Check if a session key exists and is not expired.
+
+        Args:
+            key: Session ID to check.
+
+        Returns:
+            True if the session exists and is not expired.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP for consistency with get() method.
+            Uses fetch() instead of fetch_val() to handle zero-row case.
+        """
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = $1
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
+        """
+
+        async with self._config.provide_connection() as conn:
+            query_result = await conn.fetch(sql, [key])
+            rows = query_result.result()
+            return len(rows) > 0
+
+    async def expires_in(self, key: str) -> "int | None":
+        """Get the time in seconds until the session expires.
+
+        Args:
+            key: Session ID to check.
+
+        Returns:
+            Seconds until expiration, or None if no expiry or key doesn't exist.
+
+        Notes:
+            Uses fetch() to handle the case where the key doesn't exist.
+        """
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = $1
+        """
+
+        async with self._config.provide_connection() as conn:
+            query_result = await conn.fetch(sql, [key])
+            rows = query_result.result()
+
+            if not rows:
+                return None
+
+            expires_at = rows[0]["expires_at"]
+
+            if expires_at is None:
+                return None
+
+            now = datetime.now(timezone.utc)
+            if expires_at <= now:
+                return 0
+
+            delta = expires_at - now
+            return int(delta.total_seconds())
+
+    async def delete_expired(self) -> int:
+        """Delete all expired sessions.
+
+        Returns:
+            Number of sessions deleted.
+
+        Notes:
+            Uses CURRENT_TIMESTAMP for consistency.
+            Uses RETURNING to get deleted row count since psqlpy QueryResult
+            doesn't expose command tags.
+            For very large tables (10M+ rows), consider batching deletes
+            to avoid holding locks too long.
+        """
+        sql = f"""
+        DELETE FROM {self._table_name}
+        WHERE expires_at <= CURRENT_TIMESTAMP
+        RETURNING session_id
+        """
+
+        async with self._config.provide_connection() as conn:
+            query_result = await conn.fetch(sql, [])
+            rows = query_result.result()
+            count = len(rows)
+            if count > 0:
+                logger.debug("Cleaned up %d expired sessions", count)
+            return count

sqlspec/adapters/psqlpy/type_converter.py

@@ -6,11 +6,11 @@ backward compatibility.
 """
 
 import re
-from typing import Any, Final, Optional
+from functools import lru_cache
+from typing import Any, Final
 
 from sqlspec.core.type_conversion import BaseTypeConverter
 
-# PostgreSQL-specific regex patterns for types not covered by base BaseTypeConverter
 PG_SPECIFIC_REGEX: Final[re.Pattern[str]] = re.compile(
     r"^(?:"
     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)?)?))|"
@@ -19,17 +19,50 @@ PG_SPECIFIC_REGEX: Final[re.Pattern[str]] = re.compile(
     re.IGNORECASE,
 )
 
+PG_SPECIAL_CHARS: Final[frozenset[str]] = frozenset({"{", "-", ":", "T", ".", "P", "[", "Y", "M", "D", "H", "S"})
+
 
 class PostgreSQLTypeConverter(BaseTypeConverter):
     """PostgreSQL-specific type converter with interval and array support.
 
     Extends the base BaseTypeConverter with PostgreSQL-specific functionality
     while maintaining backward compatibility for interval and array types.
+    Includes per-instance LRU cache for improved performance.
     """
 
-    __slots__ = ()
+    __slots__ = ("_convert_cache",)
+
+    def __init__(self, cache_size: int = 5000) -> None:
+        """Initialize converter with per-instance conversion cache.
+
+        Args:
+            cache_size: Maximum number of string values to cache (default: 5000)
+        """
+        super().__init__()
+
+        @lru_cache(maxsize=cache_size)
+        def _cached_convert(value: str) -> Any:
+            if not value or not any(c in value for c in PG_SPECIAL_CHARS):
+                return value
+            detected_type = self.detect_type(value)
+            return self.convert_value(value, detected_type) if detected_type else value
+
+        self._convert_cache = _cached_convert
+
+    def convert_if_detected(self, value: Any) -> Any:
+        """Convert string if special type detected (cached).
+
+        Args:
+            value: Value to potentially convert
+
+        Returns:
+            Converted value or original value
+        """
+        if not isinstance(value, str):
+            return value
+        return self._convert_cache(value)
 
-    def detect_type(self, value: str) -> Optional[str]:
+    def detect_type(self, value: str) -> str | None:
         """Detect types including PostgreSQL-specific types.
 
         Args:
@@ -38,12 +71,10 @@ class PostgreSQLTypeConverter(BaseTypeConverter):
         Returns:
             Type name if detected, None otherwise.
         """
-        # First try generic types (UUID, JSON, datetime, etc.)
         detected_type = super().detect_type(value)
         if detected_type:
             return detected_type
 
-        # Then check PostgreSQL-specific types
         match = PG_SPECIFIC_REGEX.match(value)
         if match:
             for group_name in ["interval", "pg_array"]:
@@ -62,12 +93,10 @@ class PostgreSQLTypeConverter(BaseTypeConverter):
         Returns:
             Converted value or original string for PostgreSQL-specific types.
         """
-        # For PostgreSQL-specific types, preserve as strings for backward compatibility
-        if detected_type in ("interval", "pg_array"):
-            return value  # Pass through as strings - psqlpy will handle casting
+        if detected_type in {"interval", "pg_array"}:
+            return value
 
-        # Use base converter for standard types
         return super().convert_value(value, detected_type)
 
 
-__all__ = ("PG_SPECIFIC_REGEX", "PostgreSQLTypeConverter")
+__all__ = ("PG_SPECIAL_CHARS", "PG_SPECIFIC_REGEX", "PostgreSQLTypeConverter")