feat(spanner): add Litestar integration with session store and shard support

Author: cofin

sqlspec/adapters/spanner/adk/__init__.py

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

sqlspec/adapters/spanner/adk/store.py

@@ -0,0 +1,219 @@
+"""Spanner session store for Litestar integration."""
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any, cast
+
+from google.api_core import exceptions as api_exceptions
+
+from sqlspec.extensions.litestar.store import BaseSQLSpecStore
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.sync_tools import async_
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.spanner.config import SpannerSyncConfig
+
+logger = get_logger("adapters.spanner.litestar.store")
+
+__all__ = ("SpannerSyncStore",)
+
+
+class SpannerSyncStore(BaseSQLSpecStore["SpannerSyncConfig"]):
+    """Spanner-backed Litestar session store using sync driver wrapped as async."""
+
+    __slots__ = ("_index_options", "_shard_count", "_table_options")
+
+    def __init__(self, config: "SpannerSyncConfig") -> None:
+        super().__init__(config)
+        litestar_cfg = cast("dict[str, Any]", getattr(config, "extension_config", {}).get("litestar", {}))
+        self._shard_count: int = int(litestar_cfg.get("shard_count", 0)) if litestar_cfg.get("shard_count") else 0
+        self._table_options: str | None = litestar_cfg.get("table_options")
+        self._index_options: str | None = litestar_cfg.get("index_options")
+
+    def _datetime_to_timestamp(self, dt: "datetime | None") -> "datetime | None":
+        if dt is None:
+            return None
+        if dt.tzinfo is None:
+            return dt.replace(tzinfo=timezone.utc)
+        return dt
+
+    def _timestamp_to_datetime(self, ts: "datetime | None") -> "datetime | None":
+        if ts is None:
+            return None
+        if ts.tzinfo is None:
+            return ts.replace(tzinfo=timezone.utc)
+        return ts
+
+    def _build_params(
+        self, key: str, expires_at: "datetime | None" = None, data: "bytes | None" = None
+    ) -> dict[str, Any]:
+        return {"session_id": key, "data": data, "expires_at": self._datetime_to_timestamp(expires_at)}
+
+    async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        return await async_(self._get)(key, renew_for)
+
+    def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None":
+        with self._config.provide_session() as driver:
+            sql = f"""
+            SELECT data, expires_at
+            FROM {self._table_name}
+            WHERE session_id = @session_id
+            AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP())
+            """
+            if self._shard_count > 1:
+                sql = f"""
+                SELECT data, expires_at
+                FROM {self._table_name}
+                WHERE shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})
+                AND session_id = @session_id
+                AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP())
+                """
+            result = driver.select_one(sql, {"session_id": key})
+            if result is None:
+                return None
+
+            data = result.get("data")
+            expires_at = self._timestamp_to_datetime(result.get("expires_at"))
+
+            if renew_for is not None and expires_at is not None:
+                new_expires = self._calculate_expires_at(renew_for)
+                update_sql = f"""
+                UPDATE {self._table_name}
+                SET expires_at = @expires_at, updated_at = PENDING_COMMIT_TIMESTAMP()
+                WHERE session_id = @session_id
+                """
+                if self._shard_count > 1:
+                    update_sql = f"{update_sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})"
+                driver.execute(update_sql, self._build_params(key, new_expires))
+
+            return bytes(data) if data is not None else None
+
+    async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None:
+        await async_(self._set)(key, value, expires_in)
+
+    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)
+        params = self._build_params(key, expires_at, data)
+
+        upsert_sql = f"""
+        UPDATE {self._table_name}
+        SET data = @data,
+            expires_at = @expires_at,
+            updated_at = PENDING_COMMIT_TIMESTAMP()
+        WHERE session_id = @session_id
+        """
+        if self._shard_count > 1:
+            upsert_sql = f"{upsert_sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})"
+        insert_sql = f"""
+        INSERT {self._table_name} (session_id, data, expires_at, created_at, updated_at)
+        VALUES (@session_id, @data, @expires_at, PENDING_COMMIT_TIMESTAMP(), PENDING_COMMIT_TIMESTAMP())
+        """
+        with self._config.provide_session() as driver:
+            update_result = driver.execute(upsert_sql, params)
+            if update_result.rows_affected == 0:
+                driver.execute(insert_sql, params)
+
+    async def delete(self, key: str) -> None:
+        await async_(self._delete)(key)
+
+    def _delete(self, key: str) -> None:
+        sql = f"DELETE FROM {self._table_name} WHERE session_id = @session_id"
+        if self._shard_count > 1:
+            sql = f"{sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})"
+        with self._config.provide_session() as driver:
+            driver.execute(sql, {"session_id": key})
+
+    async def delete_all(self) -> None:
+        await async_(self._delete_all)()
+
+    def _delete_all(self) -> None:
+        sql = f"DELETE FROM {self._table_name}"
+        with self._config.provide_session() as driver:
+            driver.execute(sql)
+
+    async def exists(self, key: str) -> bool:
+        return await async_(self._exists)(key)
+
+    def _exists(self, key: str) -> bool:
+        sql = f"""
+        SELECT 1 FROM {self._table_name}
+        WHERE session_id = @session_id
+        AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP())
+        LIMIT 1
+        """
+        with self._config.provide_session() as driver:
+            row = driver.select_one(sql, {"session_id": key})
+            return row is not None
+
+    async def expires_in(self, key: str) -> "int | None":
+        return await async_(self._expires_in)(key)
+
+    def _expires_in(self, key: str) -> "int | None":
+        sql = f"""
+        SELECT expires_at FROM {self._table_name}
+        WHERE session_id = @session_id
+        """
+        if self._shard_count > 1:
+            sql = f"{sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})"
+        with self._config.provide_session() as driver:
+            row = driver.select_one(sql, {"session_id": key})
+            if row is None:
+                return None
+            expires_at = self._timestamp_to_datetime(row.get("expires_at"))
+            if expires_at is None:
+                return None
+            delta = expires_at - datetime.now(timezone.utc)
+            return max(int(delta.total_seconds()), 0)
+
+    async def delete_expired(self) -> int:
+        return await async_(self._delete_expired)()
+
+    def _delete_expired(self) -> int:
+        sql = f"""
+        DELETE FROM {self._table_name}
+        WHERE expires_at IS NOT NULL AND expires_at <= CURRENT_TIMESTAMP()
+        """
+        with self._config.provide_session() as driver:
+            result = driver.execute(sql)
+            return result.rows_affected or 0
+
+    async def create_table(self) -> None:
+        await async_(self._create_table)()
+
+    def _create_table(self) -> None:
+        ddl_statements = [self._get_create_table_sql(), self._get_create_index_sql()]
+        try:
+            self._config.get_database().update_ddl(ddl_statements).result(300)  # type: ignore[no-untyped-call]
+        except api_exceptions.AlreadyExists:
+            return
+
+    def _get_create_table_sql(self) -> str:
+        shard_column = ""
+        pk = "PRIMARY KEY (session_id)"
+        if self._shard_count > 1:
+            shard_column = f",\n  shard_id INT64 AS (MOD(FARM_FINGERPRINT(session_id), {self._shard_count})) STORED"
+            pk = "PRIMARY KEY (shard_id, session_id)"
+        options = ""
+        if self._table_options:
+            options = f"\nOPTIONS ({self._table_options})"
+        return f"""
+CREATE TABLE {self._table_name} (
+  session_id STRING(128) NOT NULL,
+  data BYTES(MAX) NOT NULL,
+  expires_at TIMESTAMP,
+  created_at TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true),
+  updated_at TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true){shard_column}
+) {pk}{options}
+"""
+
+    def _get_create_index_sql(self) -> str:
+        leading = "expires_at"
+        if self._shard_count > 1:
+            leading = "shard_id, expires_at"
+        opts = ""
+        if self._index_options:
+            opts = f" OPTIONS ({self._index_options})"
+        return f"CREATE INDEX idx_{self._table_name}_expires_at ON {self._table_name}({leading}){opts}"
+
+    def _get_drop_table_sql(self) -> "list[str]":
+        return [f"DROP INDEX idx_{self._table_name}_expires_at", f"DROP TABLE {self._table_name}"]

sqlspec/adapters/spanner/litestar/__init__.py

@@ -403,6 +403,28 @@ class ADKConfig(TypedDict):
         - Ignored by non-Oracle adapters
     """
 
+    shard_count: NotRequired[int]
+    """Optional hash shard count for session/event tables to reduce hotspotting.
+
+    When set (>1), adapters that support computed shard columns will create a
+    generated shard_id using MOD(FARM_FINGERPRINT(primary_key), shard_count) and
+    include it in the primary key and filters. Ignored by adapters that do not
+    support computed shards.
+    """
+
+    session_table_options: NotRequired[str]
+    """Adapter-specific table OPTIONS/clauses for the sessions table.
+
+    Passed verbatim when supported (e.g., Spanner columnar/tiered storage). Ignored by
+    adapters without table OPTIONS support.
+    """
+
+    events_table_options: NotRequired[str]
+    """Adapter-specific table OPTIONS/clauses for the events table."""
+
+    expires_index_options: NotRequired[str]
+    """Adapter-specific options for the expires/index used in ADK stores."""
+
 
 class OpenTelemetryConfig(TypedDict):
     """Configuration options for OpenTelemetry integration.

sqlspec/adapters/spanner/litestar/store.py

@@ -64,3 +64,27 @@ class LitestarConfig(TypedDict):
         - Tables created with INMEMORY PRIORITY HIGH clause
         - Ignored by unsupported adapters
     """
+
+    shard_count: NotRequired[int]
+    """Optional hash shard count for session table primary key.
+
+    When set (>1), adapters that support computed shard columns (e.g., Spanner)
+    will create a generated shard_id using MOD(FARM_FINGERPRINT(session_id), shard_count)
+    and include it in the primary key to reduce hotspotting. Ignored by adapters
+    that do not support computed shards.
+    """
+
+    table_options: NotRequired[str]
+    """Optional raw OPTIONS/engine-specific table options string.
+
+    Passed verbatim when the adapter supports table-level OPTIONS/clauses
+    (e.g., Spanner columnar/tiered storage). Ignored by adapters that do not
+    support table options.
+    """
+
+    index_options: NotRequired[str]
+    """Optional raw OPTIONS/engine-specific options for the expires_at index.
+
+    Passed verbatim to the index definition for adapters that support index
+    OPTIONS/clauses. Ignored by adapters that do not support index options.
+    """

sqlspec/config.py

@@ -0,0 +1,41 @@
+from collections.abc import AsyncGenerator
+from typing import Any, cast
+
+import pytest
+from google.auth.credentials import AnonymousCredentials
+
+from sqlspec.adapters.spanner import SpannerSyncConfig
+from sqlspec.adapters.spanner.litestar import SpannerSyncStore
+
+
+@pytest.fixture(scope="session")
+def spanner_litestar_config(spanner_service: Any) -> SpannerSyncConfig:
+    host = getattr(spanner_service, "host", "localhost")
+    port = getattr(spanner_service, "port", 9010)
+    project_id = getattr(spanner_service, "project", "test-project")
+    instance_id = getattr(spanner_service, "instance_id", getattr(spanner_service, "instance", "test-instance"))
+    database_id = getattr(spanner_service, "database_id", getattr(spanner_service, "database", "test-database"))
+    api_endpoint = f"{host}:{port}"
+
+    return SpannerSyncConfig(
+        pool_config={
+            "project": project_id,
+            "instance_id": instance_id,
+            "database_id": database_id,
+            "credentials": cast(Any, AnonymousCredentials()),  # type: ignore[no-untyped-call]
+            "client_options": {"api_endpoint": api_endpoint},
+            "min_sessions": 1,
+            "max_sessions": 5,
+        },
+        extension_config={"litestar": {"session_table": "litestar_sessions"}},
+    )
+
+
+@pytest.fixture
+async def spanner_store(spanner_litestar_config: SpannerSyncConfig) -> AsyncGenerator[SpannerSyncStore, None]:
+    store = SpannerSyncStore(spanner_litestar_config)
+    await store.create_table()
+    try:
+        yield store
+    finally:
+        await store.delete_all()

sqlspec/extensions/litestar/config.py

@@ -0,0 +1,44 @@
+"""Integration tests for Spanner session store."""
+
+import asyncio
+
+import pytest
+
+from sqlspec.adapters.spanner.litestar import SpannerSyncStore
+
+pytestmark = [pytest.mark.spanner, pytest.mark.integration]
+
+
+async def test_store_set_and_get(spanner_store: SpannerSyncStore) -> None:
+    data = b"payload"
+    await spanner_store.set("s1", data)
+    assert await spanner_store.get("s1") == data
+
+
+async def test_store_expiration(spanner_store: SpannerSyncStore) -> None:
+    await spanner_store.set("expiring", b"data", expires_in=1)
+    assert await spanner_store.exists("expiring")
+    await asyncio.sleep(1.1)
+    assert await spanner_store.get("expiring") is None
+
+
+async def test_store_delete(spanner_store: SpannerSyncStore) -> None:
+    await spanner_store.set("todelete", b"d")
+    await spanner_store.delete("todelete")
+    assert await spanner_store.get("todelete") is None
+
+
+async def test_store_renew(spanner_store: SpannerSyncStore) -> None:
+    await spanner_store.set("renew", b"r", expires_in=1)
+    await asyncio.sleep(0.5)
+    await spanner_store.get("renew", renew_for=1)
+    await asyncio.sleep(0.7)
+    assert await spanner_store.get("renew") == b"r"
+
+
+async def test_delete_expired_returns_count(spanner_store: SpannerSyncStore) -> None:
+    await spanner_store.set("exp1", b"x", expires_in=1)
+    await spanner_store.set("exp2", b"x", expires_in=1)
+    await asyncio.sleep(1.1)
+    deleted = await spanner_store.delete_expired()
+    assert deleted >= 2