fix: improve write handling in Spanner

Author: cofin

pyproject.toml

@@ -501,6 +501,7 @@ split-on-trailing-comma = false
 "sqlspec/builder/mixins/**/*.*" = ["SLF001"]
 "sqlspec/extensions/adk/converters.py" = ["S403"]
 "sqlspec/migrations/utils.py" = ["S404"]
+"sqlspec/adapters/spanner/config.py" = ["PLC2801"]
 "tests/**/*.*" = [
   "A",
   "ARG",

sqlspec/adapters/spanner/_type_handlers.py

@@ -8,7 +8,7 @@
 """
 
 import base64
-from datetime import datetime, timezone
+from datetime import date, datetime, timezone
 from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
@@ -167,6 +167,8 @@ def infer_spanner_param_types(params: "dict[str, Any] | None") -> "dict[str, Any
             types[key] = param_types.BYTES
         elif isinstance(value, datetime):
             types[key] = param_types.TIMESTAMP
+        elif isinstance(value, date):
+            types[key] = param_types.DATE
         elif isinstance(value, dict) and hasattr(param_types, "JSON"):
             types[key] = param_types.JSON
         elif isinstance(value, list):

sqlspec/adapters/spanner/config.py

@@ -178,19 +178,37 @@ def _close_pool(self) -> None:
     def provide_connection(
         self, *args: Any, transaction: "bool" = False, **kwargs: Any
     ) -> Generator[SpannerConnection, None, None]:
-        """Yield a Snapshot (default) or Batch context from the configured pool.
+        """Yield a Snapshot (default) or Transaction context from the configured pool.
 
-        Note: Spanner does not support database.transaction() as a context manager.
-        For write operations requiring conditional logic, use database.run_in_transaction()
-        directly. The `transaction=True` option here uses database.batch() which is
-        suitable for simple insert/update/delete mutations.
+        Args:
+            *args: Additional positional arguments (unused, for interface compatibility).
+            transaction: If True, yields a Transaction context that supports
+                execute_update() for DML statements. If False (default), yields
+                a read-only Snapshot context for SELECT queries.
+            **kwargs: Additional keyword arguments (unused, for interface compatibility).
+
+        Note: For complex transactional logic with retries, use database.run_in_transaction()
+        directly. The Transaction context here auto-commits on successful exit.
         """
         database = self.get_database()
         if transaction:
-            with cast("Any", database).batch() as batch:
-                yield cast("SpannerConnection", batch)
+            session = cast("Any", database).session()
+            session.create()
+            try:
+                txn = session.transaction()
+                txn.__enter__()
+                try:
+                    yield cast("SpannerConnection", txn)
+                    if hasattr(txn, "_transaction_id") and txn._transaction_id is not None:
+                        txn.commit()
+                except Exception:
+                    if hasattr(txn, "_transaction_id") and txn._transaction_id is not None:
+                        txn.rollback()
+                    raise
+            finally:
+                session.delete()
         else:
-            with cast("Any", database).snapshot() as snapshot:
+            with cast("Any", database).snapshot(multi_use=True) as snapshot:
                 yield cast("SpannerConnection", snapshot)
 
     @contextmanager
@@ -209,7 +227,6 @@ def provide_session(
     def provide_write_session(
         self, *args: Any, statement_config: "StatementConfig | None" = None, **kwargs: Any
     ) -> Generator[SpannerSyncDriver, None, None]:
-        """Convenience wrapper that always yields a write-capable transaction session."""
         with self.provide_session(*args, statement_config=statement_config, transaction=True, **kwargs) as driver:
             yield driver
 

sqlspec/adapters/spanner/driver.py

@@ -204,32 +204,21 @@ def _execute_many(self, cursor: "SpannerConnection", statement: "SQL") -> Execut
             raise SQLConversionError(msg)
         conn = cast("Any", cursor)
 
-        parameter_sets = statement.parameters if isinstance(statement.parameters, list) else []
-        if not parameter_sets:
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+
+        if not prepared_parameters or not isinstance(prepared_parameters, list):
             msg = "execute_many requires at least one parameter set"
             raise SQLConversionError(msg)
 
-        base_params = parameter_sets[0]
-        base_statement = self.prepare_statement(
-            statement.raw_sql, *[base_params], statement_config=statement.statement_config
-        )
-        compiled_sql, _ = self._get_compiled_sql(base_statement, self.statement_config)
-
-        batch_inputs: list[dict[str, Any]] = []
-        for params in parameter_sets:
-            per_statement = self.prepare_statement(
-                statement.raw_sql, *[params], statement_config=statement.statement_config
-            )
-            _, processed_params = self._get_compiled_sql(per_statement, self.statement_config)
-            coerced_params = self._coerce_params(processed_params)
+        batch_args: list[tuple[str, dict[str, Any] | None, dict[str, Any]]] = []
+        for params in prepared_parameters:
+            coerced_params = self._coerce_params(params)
             if coerced_params is None:
                 coerced_params = {}
-            batch_inputs.append(coerced_params)
-
-        batch_args = [(compiled_sql, p, self._infer_param_types(p)) for p in batch_inputs]
+            batch_args.append((sql, coerced_params, self._infer_param_types(coerced_params)))
 
-        row_counts = conn.batch_update(batch_args)
-        total_rows = int(sum(int(count) for count in row_counts))
+        _status, row_counts = conn.batch_update(batch_args)
+        total_rows = sum(row_counts) if row_counts else 0
 
         return self.create_execution_result(cursor, rowcount_override=total_rows, is_many_result=True)
 
@@ -350,28 +339,28 @@ def _truncate_table_sync(self, table: str) -> None:
 
 
 def _build_spanner_profile() -> DriverParameterProfile:
-    type_coercions: dict[type, Any] = {dict: to_json, list: to_json, tuple: to_json}
+    type_coercions: dict[type, Any] = {dict: to_json}
     return DriverParameterProfile(
         name="Spanner",
         default_style=ParameterStyle.NAMED_AT,
         supported_styles={ParameterStyle.NAMED_AT},
         default_execution_style=ParameterStyle.NAMED_AT,
         supported_execution_styles={ParameterStyle.NAMED_AT},
         has_native_list_expansion=True,
-        json_serializer_strategy="helper",
+        json_serializer_strategy="none",
         default_dialect="spanner",
         preserve_parameter_format=True,
         needs_static_script_compilation=False,
         allow_mixed_parameter_styles=False,
         preserve_original_params_for_many=True,
         custom_type_coercions=type_coercions,
-        extras={"type_coercion_overrides": type_coercions},
+        extras={},
     )
 
 
 _SPANNER_PROFILE = _build_spanner_profile()
 register_driver_profile("spanner", _SPANNER_PROFILE)
 
 spanner_statement_config = build_statement_config_from_profile(
-    _SPANNER_PROFILE, statement_overrides={"dialect": "spanner"}, json_serializer=to_json
+    _SPANNER_PROFILE, statement_overrides={"dialect": "spanner"}
 )

tests/integration/test_adapters/test_spanner/conftest.py

@@ -35,7 +35,11 @@ def spanner_database(
 
 
 @pytest.fixture
-def spanner_config(spanner_service: SpannerService, spanner_connection: spanner.Client) -> SpannerSyncConfig:
+def spanner_config(
+    spanner_service: SpannerService, spanner_connection: spanner.Client, spanner_database: "Database"
+) -> SpannerSyncConfig:
+    """Create SpannerSyncConfig after ensuring database exists."""
+    _ = spanner_database  # Ensure database is created before config
     api_endpoint = f"{spanner_service.host}:{spanner_service.port}"
 
     return SpannerSyncConfig(
@@ -53,12 +57,27 @@ def spanner_config(spanner_service: SpannerService, spanner_connection: spanner.
 
 @pytest.fixture
 def spanner_session(spanner_config: SpannerSyncConfig) -> Generator[SpannerSyncDriver, None, None]:
+    """Read-only session for SELECT operations."""
     sql = SQLSpec()
     c = sql.add_config(spanner_config)
     with sql.provide_session(c) as session:
         yield session
 
 
+@pytest.fixture
+def spanner_write_session(spanner_config: SpannerSyncConfig) -> Generator[SpannerSyncDriver, None, None]:
+    """Write-capable session for DML operations (INSERT/UPDATE/DELETE)."""
+    with spanner_config.provide_write_session() as session:
+        yield session
+
+
+@pytest.fixture
+def spanner_read_session(spanner_config: SpannerSyncConfig) -> Generator[SpannerSyncDriver, None, None]:
+    """Read-only session for SELECT operations."""
+    with spanner_config.provide_session() as session:
+        yield session
+
+
 def run_ddl(database: "Database", statements: "list[str]", timeout: int = 300) -> None:
     """Execute DDL statements on Spanner database."""
     operation = database.update_ddl(statements)  # type: ignore[no-untyped-call]

tests/integration/test_adapters/test_spanner/test_arrow.py

@@ -1,36 +1,27 @@
-"""Integration tests for Spanner Arrow support."""
+"""Integration tests for Spanner Arrow support.
 
-from typing import TYPE_CHECKING, Any
+All operations use SQLSpec interface, not raw SDK calls.
+"""
 
 import pytest
 
 from sqlspec._typing import PYARROW_INSTALLED
 from sqlspec.adapters.spanner import SpannerSyncConfig
 
-if TYPE_CHECKING:
-    from google.cloud.spanner_v1.database import Database
-
 pytestmark = [pytest.mark.spanner, pytest.mark.skipif(not PYARROW_INSTALLED, reason="pyarrow not installed")]
 
 
-def test_select_to_arrow_basic(
-    spanner_config: SpannerSyncConfig, spanner_database: "Database", test_arrow_table: str
-) -> None:
+def test_select_to_arrow_basic(spanner_config: SpannerSyncConfig, test_arrow_table: str) -> None:
     """Test basic select_to_arrow functionality."""
     import pyarrow as pa
 
-    database = spanner_database
-
-    def insert_data(transaction: "Any") -> None:
+    with spanner_config.provide_write_session() as session:
         for i, name in enumerate(["Alice", "Bob", "Charlie"], start=1):
-            transaction.execute_update(
+            session.execute(
                 f"INSERT INTO {test_arrow_table} (id, name, value) VALUES (@id, @name, @value)",
-                params={"id": i, "name": name, "value": i * 10},
-                param_types={"id": {"code": "INT64"}, "name": {"code": "STRING"}, "value": {"code": "INT64"}},
+                {"id": i, "name": name, "value": i * 10},
             )
 
-    database.run_in_transaction(insert_data)  # type: ignore[no-untyped-call]
-
     with spanner_config.provide_session() as session:
         result = session.select_to_arrow(f"SELECT * FROM {test_arrow_table} ORDER BY id")
 
@@ -43,28 +34,19 @@ def insert_data(transaction: "Any") -> None:
         assert list(df["name"]) == ["Alice", "Bob", "Charlie"]
         assert list(df["value"]) == [10, 20, 30]
 
-    def cleanup(transaction: "Any") -> None:
-        transaction.execute_update(f"DELETE FROM {test_arrow_table} WHERE TRUE")
-
-    database.run_in_transaction(cleanup)  # type: ignore[no-untyped-call]
+    with spanner_config.provide_write_session() as session:
+        session.execute(f"DELETE FROM {test_arrow_table} WHERE TRUE")
 
 
-def test_select_to_arrow_with_parameters(
-    spanner_config: SpannerSyncConfig, spanner_database: "Database", test_arrow_table: str
-) -> None:
+def test_select_to_arrow_with_parameters(spanner_config: SpannerSyncConfig, test_arrow_table: str) -> None:
     """Test select_to_arrow with parameterized query."""
-    database = spanner_database
-
-    def insert_data(transaction: "Any") -> None:
+    with spanner_config.provide_write_session() as session:
         for i in range(1, 6):
-            transaction.execute_update(
+            session.execute(
                 f"INSERT INTO {test_arrow_table} (id, name, value) VALUES (@id, @name, @value)",
-                params={"id": i, "name": f"Item {i}", "value": i * 100},
-                param_types={"id": {"code": "INT64"}, "name": {"code": "STRING"}, "value": {"code": "INT64"}},
+                {"id": i, "name": f"Item {i}", "value": i * 100},
             )
 
-    database.run_in_transaction(insert_data)  # type: ignore[no-untyped-call]
-
     with spanner_config.provide_session() as session:
         result = session.select_to_arrow(
             f"SELECT * FROM {test_arrow_table} WHERE value > @min_value ORDER BY id", {"min_value": 200}
@@ -74,10 +56,8 @@ def insert_data(transaction: "Any") -> None:
         df = result.to_pandas()
         assert list(df["value"]) == [300, 400, 500]
 
-    def cleanup(transaction: "Any") -> None:
-        transaction.execute_update(f"DELETE FROM {test_arrow_table} WHERE TRUE")
-
-    database.run_in_transaction(cleanup)  # type: ignore[no-untyped-call]
+    with spanner_config.provide_write_session() as session:
+        session.execute(f"DELETE FROM {test_arrow_table} WHERE TRUE")
 
 
 def test_select_to_arrow_empty_result(spanner_config: SpannerSyncConfig, test_arrow_table: str) -> None:
@@ -89,91 +69,65 @@ def test_select_to_arrow_empty_result(spanner_config: SpannerSyncConfig, test_ar
         assert len(result.to_pandas()) == 0
 
 
-def test_select_to_arrow_table_format(
-    spanner_config: SpannerSyncConfig, spanner_database: "Database", test_arrow_table: str
-) -> None:
+def test_select_to_arrow_table_format(spanner_config: SpannerSyncConfig, test_arrow_table: str) -> None:
     """Test select_to_arrow with table return format (default)."""
     import pyarrow as pa
 
-    database = spanner_database
-
-    def insert_data(transaction: "Any") -> None:
+    with spanner_config.provide_write_session() as session:
         for i in range(1, 4):
-            transaction.execute_update(
+            session.execute(
                 f"INSERT INTO {test_arrow_table} (id, name, value) VALUES (@id, @name, @value)",
-                params={"id": i, "name": f"Row {i}", "value": i},
-                param_types={"id": {"code": "INT64"}, "name": {"code": "STRING"}, "value": {"code": "INT64"}},
+                {"id": i, "name": f"Row {i}", "value": i},
             )
 
-    database.run_in_transaction(insert_data)  # type: ignore[no-untyped-call]
-
     with spanner_config.provide_session() as session:
         result = session.select_to_arrow(f"SELECT * FROM {test_arrow_table} ORDER BY id", return_format="table")
 
         assert isinstance(result.data, pa.Table)
         assert result.rows_affected == 3
 
-    def cleanup(transaction: "Any") -> None:
-        transaction.execute_update(f"DELETE FROM {test_arrow_table} WHERE TRUE")
-
-    database.run_in_transaction(cleanup)  # type: ignore[no-untyped-call]
+    with spanner_config.provide_write_session() as session:
+        session.execute(f"DELETE FROM {test_arrow_table} WHERE TRUE")
 
 
-def test_select_to_arrow_batch_format(
-    spanner_config: SpannerSyncConfig, spanner_database: "Database", test_arrow_table: str
-) -> None:
+def test_select_to_arrow_batch_format(spanner_config: SpannerSyncConfig, test_arrow_table: str) -> None:
     """Test select_to_arrow with batch return format."""
     import pyarrow as pa
 
-    database = spanner_database
-
-    def insert_data(transaction: "Any") -> None:
+    with spanner_config.provide_write_session() as session:
         for i in range(1, 3):
-            transaction.execute_update(
+            session.execute(
                 f"INSERT INTO {test_arrow_table} (id, name, value) VALUES (@id, @name, @value)",
-                params={"id": i, "name": f"Batch {i}", "value": i * 5},
-                param_types={"id": {"code": "INT64"}, "name": {"code": "STRING"}, "value": {"code": "INT64"}},
+                {"id": i, "name": f"Batch {i}", "value": i * 5},
             )
 
-    database.run_in_transaction(insert_data)  # type: ignore[no-untyped-call]
-
     with spanner_config.provide_session() as session:
         result = session.select_to_arrow(f"SELECT * FROM {test_arrow_table} ORDER BY id", return_format="batch")
 
         assert isinstance(result.data, pa.RecordBatch)
         assert result.rows_affected == 2
 
-    def cleanup(transaction: "Any") -> None:
-        transaction.execute_update(f"DELETE FROM {test_arrow_table} WHERE TRUE")
+    with spanner_config.provide_write_session() as session:
+        session.execute(f"DELETE FROM {test_arrow_table} WHERE TRUE")
 
-    database.run_in_transaction(cleanup)  # type: ignore[no-untyped-call]
 
-
-def test_select_to_arrow_to_polars(
-    spanner_config: SpannerSyncConfig, spanner_database: "Database", test_arrow_table: str
-) -> None:
+def test_select_to_arrow_to_polars(spanner_config: SpannerSyncConfig, test_arrow_table: str) -> None:
     """Test select_to_arrow conversion to Polars DataFrame."""
     pytest.importorskip("polars")
-    database = spanner_database
 
-    def insert_data(transaction: "Any") -> None:
+    with spanner_config.provide_write_session() as session:
         for i in range(1, 3):
-            transaction.execute_update(
+            session.execute(
                 f"INSERT INTO {test_arrow_table} (id, name, value) VALUES (@id, @name, @value)",
-                params={"id": i, "name": f"Polars {i}", "value": i * 7},
-                param_types={"id": {"code": "INT64"}, "name": {"code": "STRING"}, "value": {"code": "INT64"}},
+                {"id": i, "name": f"Polars {i}", "value": i * 7},
             )
 
-    database.run_in_transaction(insert_data)  # type: ignore[no-untyped-call]
-
     with spanner_config.provide_session() as session:
         result = session.select_to_arrow(f"SELECT * FROM {test_arrow_table} ORDER BY id")
         df = result.to_polars()
 
         assert len(df) == 2
         assert df["name"].to_list() == ["Polars 1", "Polars 2"]
 
-    def cleanup(transaction: "Any") -> None:
-        transaction.execute_update(f"DELETE FROM {test_arrow_table} WHERE TRUE")
-
-    database.run_in_transaction(cleanup)  # type: ignore[no-untyped-call]
+    with spanner_config.provide_write_session() as session:
+        session.execute(f"DELETE FROM {test_arrow_table} WHERE TRUE")