sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/adapters/cockroach_psycopg/litestar/store.py

@@ -0,0 +1,327 @@
+"""CockroachDB session stores for Litestar integration using psycopg."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from psycopg.rows import dict_row
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+from sqlspec.utils.sync_tools import async_
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.cockroach_psycopg.config import CockroachPsycopgAsyncConfig, CockroachPsycopgSyncConfig
+
+
+__all__ = ("CockroachPsycopgAsyncStore", "CockroachPsycopgSyncStore")
+
+
+class CockroachPsycopgAsyncStore(BaseSQLSpecStore["CockroachPsycopgAsyncConfig"]):
+    """CockroachDB session store using psycopg async driver."""
+
+    __slots__ = ()
+
+    def __init__(self, config: "CockroachPsycopgAsyncConfig") -> None:
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        """Get CockroachDB CREATE TABLE SQL with optimized schema."""
+        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
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at
+        ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL;
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        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:
+        sql = self._get_create_table_sql()
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(sql)
+            await driver.commit()
+        self._log_table_created()
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
+        """
+
+        conn_context = self._config.provide_connection()
+        async with conn_context as conn:
+            async with conn.cursor(row_factory=dict_row) as cur:
+                await cur.execute(sql.encode(), (key,))
+                row = await cur.fetchone()
+
+            if row is None:
+                return None
+
+            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 = %s, updated_at = CURRENT_TIMESTAMP
+                    WHERE session_id = %s
+                    """
+                    await conn.execute(update_sql.encode(), (new_expires_at, key))
+                    await conn.commit()
+
+            return bytes(row["data"])
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        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 (%s, %s, %s)
+        ON CONFLICT (session_id)
+        DO UPDATE SET
+            data = EXCLUDED.data,
+            expires_at = EXCLUDED.expires_at,
+            updated_at = CURRENT_TIMESTAMP
+        """
+
+        conn_context = self._config.provide_connection()
+        async with conn_context as conn:
+            await conn.execute(sql.encode(), (key, data, expires_at))
+            await conn.commit()
+
+    async def delete(self, key: str) -> None:
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = %s"
+
+        conn_context = self._config.provide_connection()
+        async with conn_context as conn:
+            await conn.execute(sql.encode(), (key,))
+            await conn.commit()
+
+    async def delete_all(self) -> None:
+        sql = f"DELETE FROM {self._table_name}"
+
+        conn_context = self._config.provide_connection()
+        async with conn_context as conn:
+            await conn.execute(sql.encode())
+            await conn.commit()
+        self._log_delete_all()
+
+    async def exists(self, key: str) -> bool:
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
+        """
+
+        conn_context = self._config.provide_connection()
+        async with conn_context as conn, conn.cursor() as cur:
+            await cur.execute(sql.encode(), (key,))
+            row = await cur.fetchone()
+            return row is not None
+
+    async def expires_in(self, key: str) -> "int | None":
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        """
+
+        conn_context = self._config.provide_connection()
+        async with conn_context as conn:
+            async with conn.cursor(row_factory=dict_row) as cur:
+                await cur.execute(sql.encode(), (key,))
+                row = await cur.fetchone()
+
+            if row is None or row["expires_at"] is None:
+                return None
+
+            expires_at = row["expires_at"]
+            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:
+        sql = f"DELETE FROM {self._table_name} WHERE expires_at <= CURRENT_TIMESTAMP"
+
+        conn_context = self._config.provide_connection()
+        async with conn_context as conn, conn.cursor() as cur:
+            await cur.execute(sql.encode())
+            await conn.commit()
+            count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0
+            if count > 0:
+                self._log_delete_expired(count)
+            return count
+
+
+class CockroachPsycopgSyncStore(BaseSQLSpecStore["CockroachPsycopgSyncConfig"]):
+    """CockroachDB session store using psycopg sync driver."""
+
+    __slots__ = ()
+
+    def __init__(self, config: "CockroachPsycopgSyncConfig") -> None:
+        super().__init__(config)
+
+    def _get_create_table_sql(self) -> str:
+        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
+        );
+
+        CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at
+        ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL;
+        """
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"]
+
+    def _create_table(self) -> None:
+        sql = self._get_create_table_sql()
+        with self._config.provide_session() as driver:
+            driver.execute_script(sql)
+            driver.commit()
+        self._log_table_created()
+
+    async def create_table(self) -> None:
+        await async_(self._create_table)()
+
+    def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        sql = f"""
+        SELECT data, expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
+        """
+
+        with self._config.provide_connection() as conn:
+            with conn.cursor(row_factory=dict_row) as cur:
+                cur.execute(sql.encode(), (key,))
+                row = cur.fetchone()
+
+            if row is None:
+                return None
+
+            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 = %s, updated_at = CURRENT_TIMESTAMP
+                    WHERE session_id = %s
+                    """
+                    conn.execute(update_sql.encode(), (new_expires_at, key))
+                    conn.commit()
+
+            return bytes(row["data"])
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        return await async_(self._get)(key, renew_for)
+
+    def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        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 (%s, %s, %s)
+        ON CONFLICT (session_id)
+        DO UPDATE SET
+            data = EXCLUDED.data,
+            expires_at = EXCLUDED.expires_at,
+            updated_at = CURRENT_TIMESTAMP
+        """
+
+        with self._config.provide_connection() as conn:
+            conn.execute(sql.encode(), (key, data, expires_at))
+            conn.commit()
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        await async_(self._set)(key, value, expires_in=expires_in)
+
+    def _delete(self, key: str) -> None:
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = %s"
+
+        with self._config.provide_connection() as conn:
+            conn.execute(sql.encode(), (key,))
+            conn.commit()
+
+    async def delete(self, key: str) -> None:
+        await async_(self._delete)(key)
+
+    def _delete_all(self) -> None:
+        sql = f"DELETE FROM {self._table_name}"
+
+        with self._config.provide_connection() as conn:
+            conn.execute(sql.encode())
+            conn.commit()
+        self._log_delete_all()
+
+    async def delete_all(self) -> None:
+        await async_(self._delete_all)()
+
+    def _exists(self, key: str) -> bool:
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = %s
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
+        """
+
+        with self._config.provide_connection() as conn, conn.cursor() as cur:
+            cur.execute(sql.encode(), (key,))
+            row = cur.fetchone()
+            return row is not None
+
+    async def exists(self, key: str) -> bool:
+        return await async_(self._exists)(key)
+
+    def _expires_in(self, key: str) -> "int | None":
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = %s
+        """
+
+        with self._config.provide_connection() as conn:
+            with conn.cursor(row_factory=dict_row) as cur:
+                cur.execute(sql.encode(), (key,))
+                row = cur.fetchone()
+
+            if row is None or row["expires_at"] is None:
+                return None
+
+            expires_at = row["expires_at"]
+            now = datetime.now(timezone.utc)
+
+            if expires_at <= now:
+                return 0
+
+            delta = expires_at - now
+            return int(delta.total_seconds())
+
+    async def expires_in(self, key: str) -> "int | None":
+        return await async_(self._expires_in)(key)
+
+    def _delete_expired(self) -> int:
+        sql = f"DELETE FROM {self._table_name} WHERE expires_at <= CURRENT_TIMESTAMP"
+
+        with self._config.provide_connection() as conn, conn.cursor() as cur:
+            cur.execute(sql.encode())
+            conn.commit()
+            count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0
+            if count > 0:
+                self._log_delete_expired(count)
+            return count
+
+    async def delete_expired(self) -> int:
+        return await async_(self._delete_expired)()

sqlspec/adapters/duckdb/__init__.py

@@ -0,0 +1,29 @@
+"""DuckDB adapter for SQLSpec."""
+
+from sqlspec.adapters.duckdb._typing import DuckDBConnection, DuckDBCursor
+from sqlspec.adapters.duckdb.config import (
+    DuckDBConfig,
+    DuckDBConnectionParams,
+    DuckDBDriverFeatures,
+    DuckDBExtensionConfig,
+    DuckDBPoolParams,
+    DuckDBSecretConfig,
+)
+from sqlspec.adapters.duckdb.core import default_statement_config
+from sqlspec.adapters.duckdb.driver import DuckDBDriver, DuckDBExceptionHandler
+from sqlspec.adapters.duckdb.pool import DuckDBConnectionPool
+
+__all__ = (
+    "DuckDBConfig",
+    "DuckDBConnection",
+    "DuckDBConnectionParams",
+    "DuckDBConnectionPool",
+    "DuckDBCursor",
+    "DuckDBDriver",
+    "DuckDBDriverFeatures",
+    "DuckDBExceptionHandler",
+    "DuckDBExtensionConfig",
+    "DuckDBPoolParams",
+    "DuckDBSecretConfig",
+    "default_statement_config",
+)

sqlspec/adapters/duckdb/_typing.py

@@ -0,0 +1,104 @@
+"""DuckDB 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
+
+from duckdb import DuckDBPyConnection
+
+_DuckDBConnection = DuckDBPyConnection
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import TracebackType
+    from typing import TypeAlias
+
+    from sqlspec.adapters.duckdb.driver import DuckDBDriver
+    from sqlspec.core import StatementConfig
+
+    DuckDBConnection: TypeAlias = _DuckDBConnection
+
+if not TYPE_CHECKING:
+    DuckDBConnection = _DuckDBConnection
+
+
+class DuckDBCursor:
+    """Context manager for DuckDB connection-as-cursor.
+
+    DuckDB connections implement the cursor interface and preserve
+    variable state. Using connection directly avoids cursor overhead
+    and fixes SET VARIABLE persistence.
+
+    See: https://github.com/litestar-org/sqlspec/issues/341
+    """
+
+    __slots__ = ("connection",)
+
+    def __init__(self, connection: "DuckDBConnection") -> None:
+        self.connection = connection
+
+    def __enter__(self) -> "DuckDBConnection":
+        return self.connection
+
+    def __exit__(self, *_: Any) -> None:
+        pass  # Connection lifecycle managed by pool/session
+
+
+class DuckDBSessionContext:
+    """Sync context manager for DuckDB 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",
+        driver_features: "dict[str, Any]",
+        prepare_driver: "Callable[[DuckDBDriver], DuckDBDriver]",
+    ) -> 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: DuckDBDriver | None = None
+
+    def __enter__(self) -> "DuckDBDriver":
+        from sqlspec.adapters.duckdb.driver import DuckDBDriver
+
+        self._connection = self._acquire_connection()
+        self._driver = DuckDBDriver(
+            connection=self._connection, statement_config=self._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
+
+
+__all__ = ("DuckDBConnection", "DuckDBCursor", "DuckDBSessionContext")

sqlspec/adapters/duckdb/adk/__init__.py

@@ -0,0 +1,14 @@
+"""DuckDB ADK store for Google Agent Development Kit.
+
+DuckDB is an OLAP database optimized for analytical queries. This adapter provides
+embedded session storage with zero-configuration setup, excellent for development,
+testing, and analytical workloads.
+
+Notes:
+    For highly concurrent DML operations, consider PostgreSQL or other
+    OLTP-optimized databases.
+"""
+
+from sqlspec.adapters.duckdb.adk.store import DuckdbADKMemoryStore, DuckdbADKStore
+
+__all__ = ("DuckdbADKMemoryStore", "DuckdbADKStore")