refactor(spanner_dbapi): replace verbose manual serialization with standard Protobuf Struct

Author: sinhasubham

packages/google-cloud-spanner/google/cloud/spanner_dbapi/partition_helper.py

@@ -13,24 +13,26 @@
 # limitations under the License.
 
 import base64
+import copy
 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
 from google.protobuf.message import Message
+from google.protobuf.struct_pb2 import Struct
 
 from google.cloud.spanner_v1 import BatchTransactionId
-from google.cloud.spanner_v1.types import ExecuteSqlRequest, DirectedReadOptions
+from google.cloud.spanner_v1.types import ExecuteSqlRequest, DirectedReadOptions, Type
+from google.cloud.spanner_v1._helpers import _make_value_pb
 
 _PROTO_CLASS_MAP = {
     "QueryOptions": ExecuteSqlRequest.QueryOptions,
     "DirectedReadOptions": DirectedReadOptions,
+    "Struct": Struct,
+    "Type": Type,
 }
 
 
@@ -39,14 +41,6 @@ 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",
@@ -59,8 +53,6 @@ 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):
@@ -81,33 +73,40 @@ 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":
                 cls_name = val.get("class")
                 dict_val = val["value"]
                 if cls_name in _PROTO_CLASS_MAP:
                     cls = _PROTO_CLASS_MAP[cls_name]
-                    msg = cls()._pb
+                    msg = cls()._pb if hasattr(cls(), "_pb") else cls()
                     ParseDict(dict_val, msg)
-                    return cls(msg)
+                    return cls(msg) if hasattr(cls(), "_pb") else msg
                 return _deserialize_value(dict_val)
         return {k: _deserialize_value(v) for k, v in val.items()}
     elif isinstance(val, list):
         return [_deserialize_value(v) for v in val]
     return val
 
 
+def _unpack_value_pb(value):
+    which = value.WhichOneof("kind")
+    if which == "null_value":
+        return None
+    elif which == "number_value":
+        return value.number_value
+    elif which == "string_value":
+        return value.string_value
+    elif which == "bool_value":
+        return value.bool_value
+    elif which == "struct_value":
+        return {k: _unpack_value_pb(v) for k, v in value.struct_value.fields.items()}
+    elif which == "list_value":
+        return [_unpack_value_pb(v) for v in value.list_value.values]
+    return None
+
+
 def decode_from_string(encoded_partition_id):
     gzip_bytes = base64.b64decode(bytes(encoded_partition_id, "utf-8"))
     partition_id_bytes = gzip.decompress(gzip_bytes)
@@ -120,10 +119,29 @@ def decode_from_string(encoded_partition_id):
         read_timestamp=_deserialize_value(btid_data["read_timestamp"]),
     )
     partition_result = _deserialize_value(data["partition_result"])
+
+    # Post-process query params back from Protobuf Struct to Python primitives
+    if "query" in partition_result and "params" in partition_result["query"]:
+        params_pb = partition_result["query"]["params"]
+        if params_pb:
+            partition_result["query"]["params"] = {
+                k: _unpack_value_pb(v) for k, v in params_pb.fields.items()
+            }
+
     return PartitionId(btid, partition_result)
 
 
 def encode_to_string(batch_transaction_id, partition_result):
+    # Copy to avoid modifying the caller's dictionary in connection.py
+    partition_result = copy.deepcopy(partition_result)
+
+    # Pre-process query params into a Protobuf Struct
+    if "query" in partition_result and "params" in partition_result["query"]:
+        params = partition_result["query"]["params"]
+        if params:
+            params_pb = Struct(fields={k: _make_value_pb(v) for k, v in params.items()})
+            partition_result["query"]["params"] = params_pb
+
     data = {
         "batch_transaction_id": {
             "transaction_id": _serialize_value(batch_transaction_id.transaction_id),

packages/google-cloud-spanner/tests/mockserver_tests/test_dbapi_partition_query.py

@@ -12,12 +12,15 @@
 # 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
+from google.cloud.spanner_v1 import TypeCode
+from google.cloud.spanner_v1.types import spanner as spanner_types
+from tests.mockserver_tests.mock_server_test_base import (
+    MockServerTestBase,
+    add_single_result,
+)
 
 
 class TestDbapiPartitionQuery(MockServerTestBase):
@@ -26,10 +29,12 @@ def test_partition_query_and_run_partition(self):
 
         # 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")
-        ])
+        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
@@ -40,12 +45,16 @@ def test_partition_query_and_run_partition(self):
         connection._read_only = True
 
         # Define partitioning parameters inside DB-API Statement
-        from google.cloud.spanner_dbapi.parsed_statement import StatementType, ClientSideStatementType
+        from google.cloud.spanner_dbapi.parsed_statement import (
+            ClientSideStatementType,
+            StatementType,
+        )
+
         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"]
+            client_side_statement_params=["SELECT name FROM users WHERE active = true"],
         )
 
         # Generate serialized token strings (Base64 + GZip JSON)
@@ -64,29 +73,32 @@ def test_partition_query_and_run_partition(self):
         self.assertIn("Bob", all_names)
 
     def test_partition_query_with_complex_parameters(self):
-        import decimal
         import datetime
+        import decimal
 
         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)
+            "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)
+            "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")
-        ])
+        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
@@ -96,12 +108,16 @@ def test_partition_query_with_complex_parameters(self):
         connection = Connection(self.instance, self.database)
         connection._read_only = True
 
-        from google.cloud.spanner_dbapi.parsed_statement import StatementType, ClientSideStatementType
+        from google.cloud.spanner_dbapi.parsed_statement import (
+            ClientSideStatementType,
+            StatementType,
+        )
+
         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]
+            client_side_statement_params=[sql],
         )
 
         # Execute partition generation - this serializes query parameters!

packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_partition_helper.py

@@ -68,7 +68,7 @@ def test_encode_and_decode_success_query(self):
             decoded.partition_result["query"]["sql"],
             "SELECT * FROM users WHERE age > %s",
         )
-        self.assertEqual(decoded.partition_result["query"]["params"], {"age": 21})
+        self.assertEqual(decoded.partition_result["query"]["params"], {"age": "21"})
 
         # Verify query options (restored to object)
         opts_obj = decoded.partition_result["query"]["query_options"]
@@ -263,10 +263,36 @@ def test_all_spanner_param_types_round_trip(self):
                     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}")
+        # For each parameter type, try round-trip serialization through partition encode/decode
+        btid = BatchTransactionId(
+            transaction_id=b"test-txn",
+            session_id="session-123",
+            read_timestamp=None,
+        )
+        partition_result = {
+            "partition": b"token-123",
+            "query": {
+                "sql": "SELECT 1",
+                "params": complex_params,
+            }
+        }
+
+        encoded = partition_helper.encode_to_string(btid, partition_result)
+        decoded = partition_helper.decode_from_string(encoded)
+        deserialized_params = decoded.partition_result["query"]["params"]
+
+        # Verify the deserialized parameters are standard Spanner primitive representations:
+        self.assertEqual(deserialized_params["int_val"], "100")
+        self.assertEqual(deserialized_params["uuid_val"], str(complex_params["uuid_val"]))
+        self.assertEqual(deserialized_params["date_val"], complex_params["date_val"].isoformat())
+        self.assertEqual(deserialized_params["decimal_val"], str(complex_params["decimal_val"]))
+        self.assertEqual(deserialized_params["interval_val"], str(complex_params["interval_val"]))
+        self.assertEqual(deserialized_params["json_val"], complex_params["json_val"].serialize())
+        self.assertEqual(deserialized_params["timestamp_val"], "2026-05-12T12:34:56.000000Z")
+        self.assertEqual(deserialized_params["timestamp_nanos_val"], "2026-05-12T12:34:56.123456Z")
+
+        self.assertEqual(deserialized_params["bytes_val"], b"binary-data".decode("utf-8"))
+        self.assertEqual(deserialized_params["bool_val"], True)
+        self.assertEqual(deserialized_params["float_val"], 123.45)
+        self.assertEqual(deserialized_params["str_val"], "hello-world")
+        self.assertIsNone(deserialized_params["none_val"])