feat(spanner): enhance Spanner configuration and add integration tests for ADK and Litestar session store

Author: cofin

sqlspec/adapters/spanner/config.py

@@ -156,11 +156,20 @@ def _create_pool(self) -> AbstractSessionPool:
         database = instance.database(database_id, pool=None)
 
         pool_type = cast("type[AbstractSessionPool]", self.pool_config.get("pool_type", FixedSizePool))
-        connection_keys = {"project", "instance_id", "database_id", "credentials", "client_options", "pool_type"}
-        pool_kwargs = {k: v for k, v in self.pool_config.items() if k not in connection_keys and v is not None}
 
-        if pool_type is FixedSizePool and "size" not in pool_kwargs and "max_sessions" in self.pool_config:
-            pool_kwargs["size"] = self.pool_config["max_sessions"]
+        pool_kwargs: dict[str, Any] = {}
+        if pool_type is FixedSizePool:
+            if "size" in self.pool_config:
+                pool_kwargs["size"] = self.pool_config["size"]
+            elif "max_sessions" in self.pool_config:
+                pool_kwargs["size"] = self.pool_config["max_sessions"]
+            if "labels" in self.pool_config:
+                pool_kwargs["labels"] = self.pool_config["labels"]
+        else:
+            valid_pool_keys = {"size", "labels", "ping_interval"}
+            pool_kwargs = {k: v for k, v in self.pool_config.items() if k in valid_pool_keys and v is not None}
+            if "size" not in pool_kwargs and "max_sessions" in self.pool_config:
+                pool_kwargs["size"] = self.pool_config["max_sessions"]
 
         pool_factory = cast("Callable[..., AbstractSessionPool]", pool_type)
         return pool_factory(database, **pool_kwargs)
@@ -176,7 +185,7 @@ def provide_connection(
         """Yield a Snapshot (default) or Transaction from the configured pool."""
         database = self.get_database()
         if transaction:
-            with cast("Any", database).transaction() as txn:  # type: ignore[no-untyped-call]
+            with cast("Any", database).transaction() as txn:
                 yield cast("SpannerConnection", txn)
         else:
             with cast("Any", database).snapshot() as snapshot:

sqlspec/adapters/spanner/dialect/__init__.py

@@ -1,6 +1,6 @@
 """Spanner dialect submodule."""
 
-from sqlspec.adapters.spanner.dialect._spanner import Spanner
 from sqlspec.adapters.spanner.dialect._spangres import Spangres
+from sqlspec.adapters.spanner.dialect._spanner import Spanner
 
-__all__ = ("Spanner", "Spangres")
+__all__ = ("Spangres", "Spanner")

sqlspec/adapters/spanner/dialect/_spangres.py

@@ -1,4 +1,4 @@
-"""Google Cloud Spanner PostgreSQL-interface dialect (\"Spangres\")."""
+r"""Google Cloud Spanner PostgreSQL-interface dialect ("Spangres")."""
 
 from typing import Any, cast
 

sqlspec/adapters/spanner/dialect/_spanner.py

@@ -76,7 +76,7 @@ def _parse_property(self) -> exp.Expression:
                     this=exp.Literal.string(_ROW_DELETION_NAME), value=exp.Tuple(expressions=[column, interval])
                 )
 
-            if self._match_text_seq("TTL"):  # PostgreSQL-dialect style, keep for compatibility
+            if self._match_text_seq("TTL"):  # type: ignore[no-untyped-call]  # PostgreSQL-dialect style, keep for compatibility
                 self._match_text_seq("INTERVAL")  # type: ignore[no-untyped-call]
                 interval = cast("exp.Expression", self._parse_expression())
                 self._match_text_seq("ON")  # type: ignore[no-untyped-call]

sqlspec/builder/_ddl.py

@@ -1155,8 +1155,8 @@ def _create_base_expression(self) -> exp.Expression:
         if self._select_query is None:
             self._raise_sql_builder_error("SELECT query must be set for CREATE MATERIALIZED VIEW.")
 
-        select_expr = None
-        select_parameters = None
+        select_expr: exp.Expression | None = None
+        select_parameters: dict[str, Any] | None = None
 
         if isinstance(self._select_query, SQL):
             select_expr = self._select_query.expression
@@ -1251,8 +1251,8 @@ def _create_base_expression(self) -> exp.Expression:
         if self._select_query is None:
             self._raise_sql_builder_error("SELECT query must be set for CREATE VIEW.")
 
-        select_expr = None
-        select_parameters = None
+        select_expr: exp.Expression | None = None
+        select_parameters: dict[str, Any] | None = None
 
         if isinstance(self._select_query, SQL):
             select_expr = self._select_query.expression

sqlspec/builder/_vector_expressions.py

@@ -1,9 +1,10 @@
 """Custom SQLGlot expressions for vector distance operations.
 
 Provides dialect-specific SQL generation for vector similarity search
-across PostgreSQL (pgvector), MySQL 9+, and Oracle 23ai+.
+across PostgreSQL (pgvector), MySQL 9+, Oracle 23ai+, BigQuery, and Spanner.
 """
 
+from contextlib import suppress
 from typing import Any
 
 from sqlglot import exp
@@ -194,11 +195,13 @@ def _register_with_sqlglot() -> None:
     from sqlglot.dialects.postgres import Postgres
     from sqlglot.generator import Generator
 
-    try:  # optional, only when Spanner dialects are present
-        from sqlspec.adapters.spanner.dialect import Spanner, Spangres
-    except Exception:  # pragma: no cover - optional import
-        Spanner = None
-        Spangres = None
+    spanner_dialect: type | None = None
+    spangres_dialect: type | None = None
+    with suppress(ImportError):
+        from sqlspec.adapters.spanner.dialect import Spangres, Spanner
+
+        spanner_dialect = Spanner
+        spangres_dialect = Spangres
 
     def vector_distance_sql_base(generator: "Generator", expression: "VectorDistance") -> str:
         """Base generator for VectorDistance expressions."""
@@ -247,10 +250,10 @@ def vector_distance_sql_duckdb(generator: "Generator", expression: "VectorDistan
     Oracle.Generator.TRANSFORMS[VectorDistance] = vector_distance_sql_oracle
     BigQuery.Generator.TRANSFORMS[VectorDistance] = vector_distance_sql_bigquery
     DuckDB.Generator.TRANSFORMS[VectorDistance] = vector_distance_sql_duckdb
-    if Spanner is not None:
-        Spanner.Generator.TRANSFORMS[VectorDistance] = vector_distance_sql_spanner
-    if Spangres is not None:
-        Spangres.Generator.TRANSFORMS[VectorDistance] = vector_distance_sql_postgres
+    if spanner_dialect is not None:
+        spanner_dialect.Generator.TRANSFORMS[VectorDistance] = vector_distance_sql_spanner  # type: ignore[attr-defined]
+    if spangres_dialect is not None:
+        spangres_dialect.Generator.TRANSFORMS[VectorDistance] = vector_distance_sql_postgres  # type: ignore[attr-defined]
 
 
 _register_with_sqlglot()

sqlspec/extensions/litestar/cli.py

@@ -52,7 +52,7 @@ def add_sessions_delete_expired_command() -> None:
     except ImportError:
         return
 
-    @sessions_group.command("delete-expired")  # type: ignore[misc]
+    @sessions_group.command("delete-expired")  # type: ignore[untyped-decorator]
     @click.option(
         "--verbose", is_flag=True, default=False, help="Show detailed information about the cleanup operation"
     )

tests/integration/test_adapters/test_spanner/conftest.py

@@ -1,5 +1,5 @@
 from collections.abc import Generator
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING
 
 import pytest
 from google.api_core import exceptions as api_exceptions
@@ -17,7 +17,9 @@
 
 
 @pytest.fixture(scope="session")
-def spanner_database(spanner_service: SpannerService, spanner_connection: spanner.Client) -> Generator["Database", None, None]:
+def spanner_database(
+    spanner_service: SpannerService, spanner_connection: spanner.Client
+) -> Generator["Database", None, None]:
     """Ensure emulator instance and database exist, yield Database."""
     instance = spanner_connection.instance(spanner_service.instance_name)
     if not instance.exists():