fix: add support for more parameter types

olavloite · olavloite · commit 5ef63941799f · 2026-05-12T15:51:55.000+02:00
diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/partition_helper.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/partition_helper.py
@@ -14,9 +14,12 @@
 
 import base64
 import datetime
+import decimal
 import gzip
 import json
+import uuid
 from dataclasses import dataclass
+from google.cloud.spanner_v1.data_types import Interval, JsonObject
 from typing import Any
 
 from google.protobuf.json_format import MessageToDict, ParseDict
@@ -36,6 +39,14 @@ def _serialize_value(val: Any) -> Any:
         return {"__type__": "bytes", "value": base64.b64encode(val).decode("utf-8")}
     elif isinstance(val, datetime.datetime):
         return {"__type__": "datetime", "value": val.isoformat()}
+    elif isinstance(val, datetime.date):
+        return {"__type__": "date", "value": val.isoformat()}
+    elif isinstance(val, decimal.Decimal):
+        return {"__type__": "decimal", "value": str(val)}
+    elif isinstance(val, uuid.UUID):
+        return {"__type__": "uuid", "value": str(val)}
+    elif isinstance(val, Interval):
+        return {"__type__": "interval", "value": str(val)}
     elif hasattr(val, "_pb"):
         return {
             "__type__": "protobuf",
@@ -48,6 +59,8 @@ def _serialize_value(val: Any) -> Any:
             "class": val.__class__.__name__,
             "value": MessageToDict(val, preserving_proto_field_name=True),
         }
+    elif isinstance(val, JsonObject):
+        return {"__type__": "json_object", "value": val.serialize()}
     elif isinstance(val, dict):
         return {k: _serialize_value(v) for k, v in val.items()}
     elif isinstance(val, list):
@@ -68,6 +81,16 @@ def _deserialize_value(val: Any) -> Any:
                 if dt_str.endswith("Z"):
                     dt_str = dt_str[:-1] + "+00:00"
                 return datetime.datetime.fromisoformat(dt_str)
+            elif t == "date":
+                return datetime.date.fromisoformat(val["value"])
+            elif t == "decimal":
+                return decimal.Decimal(val["value"])
+            elif t == "uuid":
+                return uuid.UUID(val["value"])
+            elif t == "interval":
+                return Interval.from_str(val["value"])
+            elif t == "json_object":
+                return JsonObject.from_str(val["value"])
             elif t == "tuple":
                 return tuple(_deserialize_value(x) for x in val["value"])
             elif t == "protobuf":
diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/testing/mock_spanner.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/testing/mock_spanner.py
@@ -36,16 +36,21 @@ class MockSpanner:
     def __init__(self):
         self.results = {}
         self.execute_streaming_sql_results = {}
+        self.partition_results = {}
         self.errors = {}
 
     def clear_results(self):
         self.results = {}
         self.execute_streaming_sql_results = {}
+        self.partition_results = {}
         self.errors = {}
 
     def add_result(self, sql: str, result: result_set.ResultSet):
         self.results[sql.lower().strip()] = result
 
+    def add_partition_result(self, sql: str, result: spanner.PartitionResponse):
+        self.partition_results[sql.lower().strip()] = result
+
     def add_execute_streaming_sql_results(
         self, sql: str, partial_result_sets: list[result_set.PartialResultSet]
     ):
@@ -57,6 +62,12 @@ def get_result(self, sql: str) -> result_set.ResultSet:
             raise ValueError(f"No result found for {sql}")
         return result
 
+    def get_partition_result(self, sql: str) -> spanner.PartitionResponse:
+        result = self.partition_results.get(sql.lower().strip())
+        if result is None:
+            return spanner.PartitionResponse()
+        return result
+
     def add_error(self, method: str, error: _Status):
         if not hasattr(self, "_errors_list"):
             self._errors_list = {}
@@ -300,11 +311,12 @@ def Rollback(self, request, context):
 
     def PartitionQuery(self, request, context):
         self._requests.append(request)
-        return spanner.PartitionResponse()
+        return self.mock_spanner.get_partition_result(request.sql)
 
     def PartitionRead(self, request, context):
         self._requests.append(request)
-        return spanner.PartitionResponse()
+        # For reads, look up by target table name
+        return self.mock_spanner.get_partition_result(request.table)
 
     def BatchWrite(self, request, context):
         self._requests.append(request)
diff --git a/packages/google-cloud-spanner/tests/_helpers.py b/packages/google-cloud-spanner/tests/_helpers.py
@@ -1,7 +1,10 @@
 from os import getenv
 from unittest import IsolatedAsyncioTestCase
 
-import mock
+try:
+    import mock
+except ImportError:
+    import unittest.mock as mock
 
 from google.cloud.spanner_v1 import gapic_version
 from google.cloud.spanner_v1.database_sessions_manager import TransactionType
diff --git a/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_partition_query.py b/packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_partition_query.py
@@ -0,0 +1,118 @@
+# Copyright 2024 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from google.cloud.spanner_dbapi.connection import Connection
+from google.cloud.spanner_v1.types import spanner as spanner_types
+from google.cloud.spanner_v1 import TypeCode
+from tests.mockserver_tests.mock_server_test_base import MockServerTestBase, add_single_result
+from google.cloud.spanner_dbapi.parsed_statement import ParsedStatement, Statement
+
+
+class TestDbapiPartitionQuery(MockServerTestBase):
+    def test_partition_query_and_run_partition(self):
+        sql = "SELECT name FROM users WHERE active = true"
+
+        # 1. Set up mock results for PartitionQuery RPC in the mock servicer
+        partition_response = spanner_types.PartitionResponse()
+        partition_response.partitions.extend([
+            spanner_types.Partition(partition_token=b"mock-token-1"),
+            spanner_types.Partition(partition_token=b"mock-token-2")
+        ])
+        self.spanner_service.mock_spanner.add_partition_result(sql, partition_response)
+
+        # 2. Set up mock results for ExecuteSql when executing the partitions
+        add_single_result(sql, "name", TypeCode.STRING, [("Alice",), ("Bob",)])
+
+        # 3. Connect via DB-API and mark connection as read-only (required for partitioning)
+        connection = Connection(self.instance, self.database)
+        connection._read_only = True
+
+        # Define partitioning parameters inside DB-API Statement
+        from google.cloud.spanner_dbapi.parsed_statement import StatementType, ClientSideStatementType
+        parsed = ParsedStatement(
+            statement_type=StatementType.CLIENT_SIDE,
+            statement=Statement(sql),
+            client_side_statement_type=ClientSideStatementType.PARTITION_QUERY,
+            client_side_statement_params=["SELECT name FROM users WHERE active = true"]
+        )
+
+        # Generate serialized token strings (Base64 + GZip JSON)
+        partition_ids = connection.partition_query(parsed)
+        self.assertEqual(2, len(partition_ids))
+
+        # 4. Reconstruct & Execute the partitions by deserializing their tokens
+        all_names = []
+        for token in partition_ids:
+            result_stream = connection.run_partition(token)
+            for row in result_stream:
+                all_names.append(row[0])
+
+        # Verify results are successfully round-tripped and parsed
+        self.assertIn("Alice", all_names)
+        self.assertIn("Bob", all_names)
+
+    def test_partition_query_with_complex_parameters(self):
+        import decimal
+        import datetime
+
+        sql = "SELECT name FROM users WHERE active = @active AND salary > @salary AND signup_time = @signup_time"
+
+        # Set up complex parameter values (bool, Decimal, datetime)
+        params = {
+            "active": True,
+            "salary": decimal.Decimal("75000.50"),
+            "signup_time": datetime.datetime(2026, 5, 10, 12, 34, 56, tzinfo=datetime.timezone.utc)
+        }
+        from google.cloud.spanner_v1 import Type
+        param_types = {
+            "active": Type(code=TypeCode.BOOL),
+            "salary": Type(code=TypeCode.NUMERIC),
+            "signup_time": Type(code=TypeCode.TIMESTAMP)
+        }
+
+        # 1. Mock results for the partition generation RPC
+        partition_response = spanner_types.PartitionResponse()
+        partition_response.partitions.extend([
+            spanner_types.Partition(partition_token=b"complex-mock-token-1")
+        ])
+        self.spanner_service.mock_spanner.add_partition_result(sql, partition_response)
+
+        # 2. Mock results for execution of partition streaming SQL
+        add_single_result(sql, "name", TypeCode.STRING, [("Charlie",)])
+
+        # 3. Establish Connection
+        connection = Connection(self.instance, self.database)
+        connection._read_only = True
+
+        from google.cloud.spanner_dbapi.parsed_statement import StatementType, ClientSideStatementType
+        parsed = ParsedStatement(
+            statement_type=StatementType.CLIENT_SIDE,
+            statement=Statement(sql, params=params, param_types=param_types),
+            client_side_statement_type=ClientSideStatementType.PARTITION_QUERY,
+            client_side_statement_params=[sql]
+        )
+
+        # Execute partition generation - this serializes query parameters!
+        partition_ids = connection.partition_query(parsed)
+        self.assertEqual(1, len(partition_ids))
+
+        # 4. Reconstruct and run the partition E2E
+        all_names = []
+        for token in partition_ids:
+            result_stream = connection.run_partition(token)
+            for row in result_stream:
+                all_names.append(row[0])
+
+        self.assertEqual(["Charlie"], all_names)
diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_partition_helper.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_partition_helper.py
@@ -122,3 +122,151 @@ def test_insecure_deserialization_failure(self):
         # Since we now use json.loads, a pickle payload will fail to decode as UTF-8 / JSON
         with self.assertRaises((json.JSONDecodeError, UnicodeDecodeError)):
             partition_helper.decode_from_string(encoded_pickle)
+
+    def test_protobuf_round_trip_reversibility(self):
+        # Any Protobuf message returned by Spanner options must be fully and reversibly
+        # round-trip deserializable to its original class, not falling back to a dict.
+        for cls_name, cls in partition_helper._PROTO_CLASS_MAP.items():
+            instance = cls()
+            serialized = partition_helper._serialize_value(instance)
+            deserialized = partition_helper._deserialize_value(serialized)
+            self.assertIsInstance(deserialized, cls)
+            self.assertEqual(deserialized.__class__.__name__, cls_name)
+
+    def test_dynamic_partition_options_registered(self):
+        # Dynamically verify that any Protobuf message class generated inside query_info or read_info
+        # during partitioning is registered in _PROTO_CLASS_MAP.
+        from google.protobuf.message import Message
+        from unittest.mock import MagicMock
+        from google.cloud.spanner_v1.database import BatchSnapshot
+        from google.cloud.spanner_v1.types import ExecuteSqlRequest, DirectedReadOptions
+
+        db = MagicMock()
+        db.observability_options = {}
+        db._instance._client._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
+
+        snapshot = BatchSnapshot(db)
+        snapshot._snapshot = MagicMock()
+        snapshot._snapshot.partition_query.return_value = [b"token-123"]
+        snapshot._snapshot.partition_read.return_value = [b"token-456"]
+
+        query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="2")
+        directed_read_options = DirectedReadOptions()
+
+        query_batches = list(snapshot.generate_query_batches(
+            sql="SELECT 1",
+            query_options=query_options,
+            directed_read_options=directed_read_options
+        ))
+
+        from google.cloud.spanner_v1.keyset import KeySet
+        read_batches = list(snapshot.generate_read_batches(
+            table="users",
+            columns=["name"],
+            keyset=KeySet(all_=True),
+            directed_read_options=directed_read_options
+        ))
+
+        discovered_protobuf_classes = set()
+
+        def collect_protobufs(val):
+            if isinstance(val, dict):
+                for v in val.values():
+                    collect_protobufs(v)
+            elif isinstance(val, list):
+                for v in val:
+                    collect_protobufs(v)
+            elif hasattr(val, "_pb") or isinstance(val, Message):
+                discovered_protobuf_classes.add(val.__class__)
+
+        for batch in query_batches + read_batches:
+            collect_protobufs(batch)
+
+        registered_classes = set(partition_helper._PROTO_CLASS_MAP.values())
+        for cls in discovered_protobuf_classes:
+            with self.subTest(cls=cls):
+                self.assertIn(
+                    cls,
+                    registered_classes,
+                    f"Protobuf class '{cls.__name__}' is generated in partition batch details "
+                    f"but is not registered in partition_helper._PROTO_CLASS_MAP! "
+                    f"Please add it to _PROTO_CLASS_MAP to prevent silent deserialization failures."
+                )
+
+    def test_all_spanner_param_types_round_trip(self):
+        import uuid
+        import datetime
+        import decimal
+        from google.cloud.spanner_v1.data_types import Interval, JsonObject
+        from google.api_core.datetime_helpers import DatetimeWithNanoseconds
+
+        complex_params = {
+            "uuid_val": uuid.UUID("12345678-1234-5678-1234-567812345678"),
+            "date_val": datetime.date(2026, 5, 12),
+            "decimal_val": decimal.Decimal("99999.99"),
+            "interval_val": Interval(months=35, days=12, nanos=54321000),
+            "json_val": JsonObject({"name": "Alice", "active": True}),
+            "timestamp_val": datetime.datetime(2026, 5, 12, 12, 34, 56, tzinfo=datetime.timezone.utc),
+            "timestamp_nanos_val": DatetimeWithNanoseconds(2026, 5, 12, 12, 34, 56, 123456, tzinfo=datetime.timezone.utc),
+            "bytes_val": b"binary-data",
+            "bool_val": True,
+            "int_val": 100,
+            "float_val": 123.45,
+            "str_val": "hello-world",
+            "none_val": None
+        }
+
+        # DYNAMIC AST VERIFICATION OF CORE SDK SUPPORTED TYPES
+        # Dynamically discover all classes checked by `isinstance` inside Spanner's `_make_value_pb`.
+        import ast
+        import inspect
+        from google.cloud.spanner_v1._helpers import _make_value_pb
+
+        source = inspect.getsource(_make_value_pb)
+        tree = ast.parse(source)
+
+        discovered_types = set()
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "isinstance":
+                if len(node.args) >= 2 and isinstance(node.args[0], ast.Name) and node.args[0].id == "value":
+                    type_arg = node.args[1]
+                    if isinstance(type_arg, ast.Tuple):
+                        for elt in type_arg.elts:
+                            if isinstance(elt, ast.Name):
+                                discovered_types.add(elt.id)
+                    elif isinstance(type_arg, ast.Name):
+                        discovered_types.add(type_arg.id)
+                    elif isinstance(type_arg, ast.Attribute):
+                        discovered_types.add(type_arg.attr)
+
+        # Map our test's `complex_params` actual class/instance types to the class name strings
+        test_param_class_names = set()
+        for val in complex_params.values():
+            if val is not None:
+                test_param_class_names.add(val.__class__.__name__)
+                # Also map base classes (e.g., DatetimeWithNanoseconds inherits from datetime)
+                for base in val.__class__.__mro__:
+                    test_param_class_names.add(base.__name__)
+
+        # Special mappings for primitive built-ins checked in standard library or tuple serialization
+        test_param_class_names.update({"list", "tuple", "Message", "ListValue"})
+
+        # Assert that every type validated in Spanner SDK's _make_value_pb
+        # has a corresponding test parameter type implemented in our round-trip check
+        for sdk_type in discovered_types:
+            with self.subTest(sdk_type=sdk_type):
+                self.assertIn(
+                    sdk_type,
+                    test_param_class_names,
+                    f"Spanner SDK parameter helper (_make_value_pb) supports type '{sdk_type}', "
+                    f"but this type has not been implemented/mapped in the DB-API partition "
+                    f"helper tests! Please add a verification case for it."
+                )
+
+        # For each parameter type, try round-trip serialization
+        for key, val in complex_params.items():
+            with self.subTest(type_key=key):
+                serialized = partition_helper._serialize_value(val)
+                json_str = json.dumps(serialized)
+                deserialized = partition_helper._deserialize_value(json.loads(json_str))
+                self.assertEqual(deserialized, val, f"Round-trip failed for {key}! Original: {val}, Deserialized: {deserialized}")
