sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/adapters/psqlpy/litestar/store.py

@@ -0,0 +1,270 @@
+"""Psqlpy session store for Litestar integration."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.psqlpy.config import PsqlpyConfig
+
+
+__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(connection_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)
+        self._log_table_created()
+
+    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)
+        self._log_delete_all()
+
+    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:
+                self._log_delete_expired(count)
+            return count

sqlspec/adapters/psqlpy/type_converter.py

@@ -0,0 +1,113 @@
+"""PostgreSQL-specific type conversion for psqlpy adapter.
+
+Combines output conversion (database results → Python) and input conversion
+(Python params → PostgreSQL format) in a single module. Designed for mypyc
+compilation with no nested functions.
+
+Output conversion handles:
+- PostgreSQL-specific types like intervals and arrays
+- Standard type detection (UUID, JSON, datetime, etc.)
+
+Input conversion handles:
+- pgvector type handlers (placeholder for future support)
+"""
+
+import re
+from typing import TYPE_CHECKING, Any, Final
+
+from sqlspec.core.type_converter import CachedOutputConverter
+from sqlspec.typing import PGVECTOR_INSTALLED
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from psqlpy import Connection
+
+__all__ = ("PG_SPECIAL_CHARS", "PG_SPECIFIC_REGEX", "PostgreSQLOutputConverter", "register_pgvector")
+
+logger = get_logger(__name__)
+
+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)?)?))|"
+    r"(?P<pg_array>\{(?:[^{}]+|\{[^{}]*\})*\})"
+    r")$",
+    re.IGNORECASE,
+)
+
+PG_SPECIAL_CHARS: Final[frozenset[str]] = frozenset({"{", "-", ":", "T", ".", "P", "[", "Y", "M", "D", "H", "S"})
+
+
+class PostgreSQLOutputConverter(CachedOutputConverter):
+    """PostgreSQL-specific output conversion with interval and array support.
+
+    Extends CachedOutputConverter with PostgreSQL-specific functionality
+    for interval and array type handling.
+    """
+
+    __slots__ = ()
+
+    def __init__(self, cache_size: int = 5000) -> None:
+        """Initialize converter with PostgreSQL-specific options.
+
+        Args:
+            cache_size: Maximum number of string values to cache (default: 5000)
+        """
+        super().__init__(special_chars=PG_SPECIAL_CHARS, cache_size=cache_size)
+
+    def _convert_detected(self, value: str, detected_type: str) -> Any:
+        """Convert value with PostgreSQL-specific handling.
+
+        Args:
+            value: String value to convert.
+            detected_type: Detected type name.
+
+        Returns:
+            Converted value or original for PostgreSQL-specific types.
+        """
+        if detected_type in {"interval", "pg_array"}:
+            return value
+        try:
+            return self.convert_value(value, detected_type)
+        except Exception:
+            return value
+
+    def detect_type(self, value: str) -> "str | None":
+        """Detect types including PostgreSQL-specific types.
+
+        Args:
+            value: String value to analyze.
+
+        Returns:
+            Type name if detected, None otherwise.
+        """
+        detected_type = super().detect_type(value)
+        if detected_type:
+            return detected_type
+
+        match = PG_SPECIFIC_REGEX.match(value)
+        if match:
+            for group_name in ["interval", "pg_array"]:
+                if match.group(group_name):
+                    return group_name
+
+        return None
+
+
+def register_pgvector(connection: "Connection") -> None:
+    """Register pgvector type handlers on psqlpy connection.
+
+    Currently a placeholder for future implementation. The psqlpy library
+    does not yet expose a type handler registration API compatible with
+    pgvector's automatic conversion system.
+
+    Args:
+        connection: Psqlpy connection instance.
+
+    Note:
+        When psqlpy adds type handler support, this function will:
+        - Register pgvector extension on the connection
+        - Enable automatic NumPy array <-> PostgreSQL vector conversion
+        - Support vector similarity search operations
+    """
+    if not PGVECTOR_INSTALLED:
+        return

sqlspec/adapters/psycopg/__init__.py

@@ -0,0 +1,38 @@
+from sqlspec.adapters.psycopg._typing import (
+    PsycopgAsyncConnection,
+    PsycopgAsyncCursor,
+    PsycopgSyncConnection,
+    PsycopgSyncCursor,
+)
+from sqlspec.adapters.psycopg.config import (
+    PsycopgAsyncConfig,
+    PsycopgConnectionParams,
+    PsycopgDriverFeatures,
+    PsycopgPoolParams,
+    PsycopgSyncConfig,
+)
+from sqlspec.adapters.psycopg.core import default_statement_config
+from sqlspec.adapters.psycopg.driver import (
+    PsycopgAsyncDriver,
+    PsycopgAsyncExceptionHandler,
+    PsycopgSyncDriver,
+    PsycopgSyncExceptionHandler,
+)
+from sqlspec.dialects import postgres  # noqa: F401
+
+__all__ = (
+    "PsycopgAsyncConfig",
+    "PsycopgAsyncConnection",
+    "PsycopgAsyncCursor",
+    "PsycopgAsyncDriver",
+    "PsycopgAsyncExceptionHandler",
+    "PsycopgConnectionParams",
+    "PsycopgDriverFeatures",
+    "PsycopgPoolParams",
+    "PsycopgSyncConfig",
+    "PsycopgSyncConnection",
+    "PsycopgSyncCursor",
+    "PsycopgSyncDriver",
+    "PsycopgSyncExceptionHandler",
+    "default_statement_config",
+)

sqlspec/adapters/psycopg/_typing.py

@@ -0,0 +1,218 @@
+"""Psycopg adapter type definitions.
+
+This module contains type aliases and classes that are excluded from mypyc
+compilation to avoid ABI boundary issues.
+"""
+
+from typing import TYPE_CHECKING, Any, Protocol
+
+from psycopg import AsyncConnection, AsyncCursor, Connection, Cursor
+from psycopg.rows import DictRow as PsycopgDictRow
+from psycopg.sql import SQL as PsycopgSQL  # noqa: N811
+from psycopg.sql import Composed as PsycopgComposed
+from psycopg.sql import Identifier as PsycopgIdentifier
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import TracebackType
+    from typing import TypeAlias
+
+    from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver
+    from sqlspec.builder import QueryBuilder
+    from sqlspec.core import SQL, Statement, StatementConfig
+
+    PsycopgSyncConnection: TypeAlias = Connection[PsycopgDictRow]
+    PsycopgAsyncConnection: TypeAlias = AsyncConnection[PsycopgDictRow]
+    PsycopgSyncRawCursor: TypeAlias = Cursor[PsycopgDictRow]
+    PsycopgAsyncRawCursor: TypeAlias = AsyncCursor[PsycopgDictRow]
+
+if not TYPE_CHECKING:
+    PsycopgSyncConnection = Connection
+    PsycopgAsyncConnection = AsyncConnection
+    PsycopgSyncRawCursor = Cursor
+    PsycopgAsyncRawCursor = AsyncCursor
+
+
+class PsycopgSyncCursor:
+    """Context manager for PostgreSQL psycopg cursor management."""
+
+    __slots__ = ("connection", "cursor")
+
+    def __init__(self, connection: "PsycopgSyncConnection") -> None:
+        self.connection = connection
+        self.cursor: PsycopgSyncRawCursor | None = None
+
+    def __enter__(self) -> "PsycopgSyncRawCursor":
+        self.cursor = self.connection.cursor()
+        return self.cursor
+
+    def __exit__(self, *_: Any) -> None:
+        if self.cursor is not None:
+            self.cursor.close()
+
+
+class PsycopgAsyncCursor:
+    """Async context manager for PostgreSQL psycopg cursor management."""
+
+    __slots__ = ("connection", "cursor")
+
+    def __init__(self, connection: "PsycopgAsyncConnection") -> None:
+        self.connection = connection
+        self.cursor: PsycopgAsyncRawCursor | None = None
+
+    async def __aenter__(self) -> "PsycopgAsyncRawCursor":
+        self.cursor = self.connection.cursor()
+        return self.cursor
+
+    async def __aexit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
+    ) -> None:
+        _ = (exc_type, exc_val, exc_tb)
+        if self.cursor is not None:
+            await self.cursor.close()
+
+
+class PsycopgPipelineDriver(Protocol):
+    """Protocol for psycopg pipeline driver methods used in stack execution."""
+
+    statement_config: "StatementConfig"
+
+    def prepare_statement(
+        self,
+        statement: "SQL | Statement | QueryBuilder",
+        parameters: Any,
+        *,
+        statement_config: "StatementConfig | None" = None,
+        kwargs: "dict[str, Any] | None" = None,
+    ) -> "SQL": ...
+
+    def _get_compiled_sql(self, statement: "SQL", statement_config: "StatementConfig") -> "tuple[str, Any]": ...
+
+
+class PsycopgSyncSessionContext:
+    """Sync context manager for psycopg sessions.
+
+    This class is intentionally excluded from mypyc compilation to avoid ABI
+    boundary issues. It receives callables from uncompiled config classes and
+    instantiates compiled Driver objects, acting as a bridge between compiled
+    and uncompiled code.
+
+    Uses callable-based connection management to decouple from config implementation.
+    """
+
+    __slots__ = (
+        "_acquire_connection",
+        "_connection",
+        "_driver",
+        "_driver_features",
+        "_prepare_driver",
+        "_release_connection",
+        "_statement_config",
+    )
+
+    def __init__(
+        self,
+        acquire_connection: "Callable[[], Any]",
+        release_connection: "Callable[[Any], Any]",
+        statement_config: "StatementConfig | Callable[[], StatementConfig]",
+        driver_features: "dict[str, Any]",
+        prepare_driver: "Callable[[PsycopgSyncDriver], PsycopgSyncDriver]",
+    ) -> None:
+        self._acquire_connection = acquire_connection
+        self._release_connection = release_connection
+        self._statement_config = statement_config
+        self._driver_features = driver_features
+        self._prepare_driver = prepare_driver
+        self._connection: Any = None
+        self._driver: PsycopgSyncDriver | None = None
+
+    def __enter__(self) -> "PsycopgSyncDriver":
+        from sqlspec.adapters.psycopg.driver import PsycopgSyncDriver
+
+        self._connection = self._acquire_connection()
+        statement_config = self._statement_config() if callable(self._statement_config) else self._statement_config
+        self._driver = PsycopgSyncDriver(
+            connection=self._connection, statement_config=statement_config, driver_features=self._driver_features
+        )
+        return self._prepare_driver(self._driver)
+
+    def __exit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
+    ) -> "bool | None":
+        if self._connection is not None:
+            self._release_connection(self._connection)
+            self._connection = None
+        return None
+
+
+class PsycopgAsyncSessionContext:
+    """Async context manager for psycopg sessions.
+
+    This class is intentionally excluded from mypyc compilation to avoid ABI
+    boundary issues. It receives callables from uncompiled config classes and
+    instantiates compiled Driver objects, acting as a bridge between compiled
+    and uncompiled code.
+
+    Uses callable-based connection management to decouple from config implementation.
+    """
+
+    __slots__ = (
+        "_acquire_connection",
+        "_connection",
+        "_driver",
+        "_driver_features",
+        "_prepare_driver",
+        "_release_connection",
+        "_statement_config",
+    )
+
+    def __init__(
+        self,
+        acquire_connection: "Callable[[], Any]",
+        release_connection: "Callable[[Any], Any]",
+        statement_config: "StatementConfig | Callable[[], StatementConfig]",
+        driver_features: "dict[str, Any]",
+        prepare_driver: "Callable[[PsycopgAsyncDriver], PsycopgAsyncDriver]",
+    ) -> None:
+        self._acquire_connection = acquire_connection
+        self._release_connection = release_connection
+        self._statement_config = statement_config
+        self._driver_features = driver_features
+        self._prepare_driver = prepare_driver
+        self._connection: Any = None
+        self._driver: PsycopgAsyncDriver | None = None
+
+    async def __aenter__(self) -> "PsycopgAsyncDriver":
+        from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver
+
+        self._connection = await self._acquire_connection()
+        statement_config = self._statement_config() if callable(self._statement_config) else self._statement_config
+        self._driver = PsycopgAsyncDriver(
+            connection=self._connection, statement_config=statement_config, driver_features=self._driver_features
+        )
+        return self._prepare_driver(self._driver)
+
+    async def __aexit__(
+        self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None"
+    ) -> "bool | None":
+        if self._connection is not None:
+            await self._release_connection(self._connection)
+            self._connection = None
+        return None
+
+
+__all__ = (
+    "PsycopgAsyncConnection",
+    "PsycopgAsyncCursor",
+    "PsycopgAsyncRawCursor",
+    "PsycopgAsyncSessionContext",
+    "PsycopgComposed",
+    "PsycopgDictRow",
+    "PsycopgIdentifier",
+    "PsycopgPipelineDriver",
+    "PsycopgSQL",
+    "PsycopgSyncConnection",
+    "PsycopgSyncCursor",
+    "PsycopgSyncRawCursor",
+    "PsycopgSyncSessionContext",
+)

sqlspec/adapters/psycopg/adk/__init__.py

@@ -0,0 +1,10 @@
+"""Psycopg ADK store module."""
+
+from sqlspec.adapters.psycopg.adk.store import (
+    PsycopgAsyncADKMemoryStore,
+    PsycopgAsyncADKStore,
+    PsycopgSyncADKMemoryStore,
+    PsycopgSyncADKStore,
+)
+
+__all__ = ("PsycopgAsyncADKMemoryStore", "PsycopgAsyncADKStore", "PsycopgSyncADKMemoryStore", "PsycopgSyncADKStore")