perf(spanner): optimize query option merging and prevent in-place mutation (#18358)

Short-circuit query option merging when options are unset to eliminate
allocations on the hot path, make field handling generic, and merge into
a fresh protobuf to avoid in-place mutation of base options.

Author: olavloite

packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py

@@ -160,6 +160,43 @@ def _try_to_coerce_bytes(bytestring):
         )
 
 
+_VALID_QUERY_OPTIONS_KEYS = frozenset(
+    ExecuteSqlRequest.QueryOptions._meta.fields.keys()
+)
+
+
+def _to_query_options(options):
+    """Normalize dict or QueryOptions to a non-empty QueryOptions, or None.
+
+    :type options:
+        :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions`
+        or :class:`dict` or None
+    :param options: Query options to normalize.
+
+    :rtype:
+        :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions`
+        or None
+    :returns:
+        A non-empty QueryOptions instance, or None if options is empty or None.
+
+    :raises TypeError:
+        If options is not a QueryOptions, dict, or None.
+    :raises ValueError:
+        If options is a dict containing unknown fields.
+    """
+    if options is None:
+        return None
+    if isinstance(options, dict):
+        if options.keys() <= _VALID_QUERY_OPTIONS_KEYS and not any(options.values()):
+            return None
+        options = ExecuteSqlRequest.QueryOptions(options)
+    elif not isinstance(options, ExecuteSqlRequest.QueryOptions):
+        raise TypeError(
+            f"query_options must be a QueryOptions or dict, got {type(options).__name__}"
+        )
+    return options if type(options).pb(options).ByteSize() > 0 else None
+
+
 def _merge_query_options(base, merge):
     """Merge higher precedence QueryOptions with current QueryOptions.
 
@@ -182,23 +219,20 @@ def _merge_query_options(base, merge):
         QueryOptions object formed by merging the two given QueryOptions.
         If the resultant object only has empty fields, returns None.
     """
-    combined = base or ExecuteSqlRequest.QueryOptions()
-    if isinstance(combined, dict):
-        combined = ExecuteSqlRequest.QueryOptions(
-            optimizer_version=combined.get("optimizer_version", ""),
-            optimizer_statistics_package=combined.get(
-                "optimizer_statistics_package", ""
-            ),
-        )
-    merge = merge or ExecuteSqlRequest.QueryOptions()
-    if isinstance(merge, dict):
-        merge = ExecuteSqlRequest.QueryOptions(
-            optimizer_version=merge.get("optimizer_version", ""),
-            optimizer_statistics_package=merge.get("optimizer_statistics_package", ""),
-        )
-    type(combined).pb(combined).MergeFrom(type(merge).pb(merge))
-    if not combined.optimizer_version and not combined.optimizer_statistics_package:
+    if base is None and merge is None:
         return None
+
+    base = _to_query_options(base)
+    merge = _to_query_options(merge)
+    if base is None:
+        return merge
+    if merge is None:
+        return base
+
+    combined = ExecuteSqlRequest.QueryOptions()
+    combined_pb = type(combined).pb(combined)
+    combined_pb.CopyFrom(type(base).pb(base))
+    combined_pb.MergeFrom(type(merge).pb(merge))
     return combined
 
 

packages/google-cloud-spanner/tests/unit/test__helpers.py

@@ -22,7 +22,50 @@
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.semconv.resource import ResourceAttributes
 
-from google.cloud.spanner_v1 import TransactionOptions, _helpers
+from google.cloud.spanner_v1 import ExecuteSqlRequest, TransactionOptions, _helpers
+
+
+class Test_to_query_options(unittest.TestCase):
+    def _callFUT(self, *args, **kw):
+        from google.cloud.spanner_v1._helpers import _to_query_options
+
+        return _to_query_options(*args, **kw)
+
+    def test_none(self):
+        self.assertIsNone(self._callFUT(None))
+
+    def test_empty_dict(self):
+        self.assertIsNone(self._callFUT({}))
+
+    def test_dict_with_empty_values(self):
+        self.assertIsNone(self._callFUT({"optimizer_version": ""}))
+
+    def test_valid_dict(self):
+        expected = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
+        result = self._callFUT({"optimizer_version": "1"})
+        self.assertEqual(result, expected)
+
+    def test_empty_proto_object(self):
+        self.assertIsNone(self._callFUT(ExecuteSqlRequest.QueryOptions()))
+
+    def test_populated_proto_object(self):
+        options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
+        result = self._callFUT(options)
+        self.assertEqual(result, options)
+
+    def test_invalid_type(self):
+        for invalid_value in ("invalid", 123, "", [], False):
+            with self.subTest(invalid_value=invalid_value):
+                with self.assertRaises(TypeError):
+                    self._callFUT(invalid_value)
+
+    def test_unknown_key_with_empty_value_raises_error(self):
+        with self.assertRaises(ValueError):
+            self._callFUT({"optmizer_version": ""})
+
+    def test_unknown_key_with_non_empty_value_raises_error(self):
+        with self.assertRaises(ValueError):
+            self._callFUT({"optmizer_version": "1"})
 
 
 class Test_merge_query_options(unittest.TestCase):
@@ -37,8 +80,6 @@ def test_base_none_and_merge_none(self):
         self.assertIsNone(result)
 
     def test_base_dict_and_merge_none(self):
-        from google.cloud.spanner_v1 import ExecuteSqlRequest
-
         base = {
             "optimizer_version": "2",
             "optimizer_statistics_package": "auto_20191128_14_47_22UTC",
@@ -52,16 +93,12 @@ def test_base_dict_and_merge_none(self):
         self.assertEqual(result, expected)
 
     def test_base_empty_and_merge_empty(self):
-        from google.cloud.spanner_v1 import ExecuteSqlRequest
-
         base = ExecuteSqlRequest.QueryOptions()
         merge = ExecuteSqlRequest.QueryOptions()
         result = self._callFUT(base, merge)
         self.assertIsNone(result)
 
     def test_base_none_merge_object(self):
-        from google.cloud.spanner_v1 import ExecuteSqlRequest
-
         base = None
         merge = ExecuteSqlRequest.QueryOptions(
             optimizer_version="3",
@@ -71,28 +108,138 @@ def test_base_none_merge_object(self):
         self.assertEqual(result, merge)
 
     def test_base_none_merge_dict(self):
-        from google.cloud.spanner_v1 import ExecuteSqlRequest
-
         base = None
         merge = {"optimizer_version": "3"}
         expected = ExecuteSqlRequest.QueryOptions(optimizer_version="3")
         result = self._callFUT(base, merge)
         self.assertEqual(result, expected)
 
     def test_base_object_merge_dict(self):
-        from google.cloud.spanner_v1 import ExecuteSqlRequest
+        base = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="1",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
+        merge = {"optimizer_version": "3"}
+        expected = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="3",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
+        result = self._callFUT(base, merge)
+        self.assertEqual(result, expected)
+
+    def test_base_object_and_merge_none(self):
+        base = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="2",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
+        result = self._callFUT(base, None)
+        self.assertEqual(result, base)
+
+    def test_base_empty_object_and_merge_none(self):
+        base = ExecuteSqlRequest.QueryOptions()
+        result = self._callFUT(base, None)
+        self.assertIsNone(result)
+
+    def test_base_none_merge_empty_object(self):
+        merge = ExecuteSqlRequest.QueryOptions()
+        result = self._callFUT(None, merge)
+        self.assertIsNone(result)
 
+    def test_base_object_not_mutated_on_merge(self):
         base = ExecuteSqlRequest.QueryOptions(
             optimizer_version="1",
             optimizer_statistics_package="auto_20191128_14_47_22UTC",
         )
         merge = {"optimizer_version": "3"}
+        result = self._callFUT(base, merge)
         expected = ExecuteSqlRequest.QueryOptions(
             optimizer_version="3",
             optimizer_statistics_package="auto_20191128_14_47_22UTC",
         )
+        self.assertEqual(result, expected)
+        self.assertEqual(base.optimizer_version, "1")
+
+    def test_base_dict_merge_dict(self):
+        base = {"optimizer_version": "1"}
+        merge = {"optimizer_statistics_package": "auto_20191128_14_47_22UTC"}
+        expected = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="1",
+            optimizer_statistics_package="auto_20191128_14_47_22UTC",
+        )
+        result = self._callFUT(base, merge)
+        self.assertEqual(result, expected)
+
+    def test_base_dict_override_dict(self):
+        base = {
+            "optimizer_version": "1",
+            "optimizer_statistics_package": "pkg1",
+        }
+        merge = {"optimizer_version": "2"}
+        expected = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="2",
+            optimizer_statistics_package="pkg1",
+        )
+        result = self._callFUT(base, merge)
+        self.assertEqual(result, expected)
+
+    def test_base_dict_empty_merge_none(self):
+        result = self._callFUT({}, None)
+        self.assertIsNone(result)
+
+    def test_base_none_merge_dict_empty(self):
+        result = self._callFUT(None, {})
+        self.assertIsNone(result)
+
+    def test_base_empty_dict_merge_empty_dict(self):
+        result = self._callFUT({}, {})
+        self.assertIsNone(result)
+
+    def test_base_empty_dict_merge_object(self):
+        merge = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
+        result = self._callFUT({}, merge)
+        self.assertEqual(result, merge)
+
+    def test_base_object_merge_empty_dict(self):
+        base = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
+        result = self._callFUT(base, {})
+        self.assertEqual(result, base)
+
+    def test_base_object_merge_object(self):
+        base = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="1",
+            optimizer_statistics_package="pkg1",
+        )
+        merge = ExecuteSqlRequest.QueryOptions(optimizer_version="2")
         result = self._callFUT(base, merge)
+        expected = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="2",
+            optimizer_statistics_package="pkg1",
+        )
         self.assertEqual(result, expected)
+        self.assertEqual(base.optimizer_version, "1")
+        self.assertEqual(base.optimizer_statistics_package, "pkg1")
+        self.assertEqual(merge.optimizer_version, "2")
+        self.assertEqual(merge.optimizer_statistics_package, "")
+
+    def test_invalid_type_raises_error(self):
+        for invalid_value in ("invalid", 123, "", [], False):
+            with self.subTest(invalid_value=invalid_value):
+                with self.assertRaises(TypeError):
+                    self._callFUT(invalid_value, None)
+                with self.assertRaises(TypeError):
+                    self._callFUT(None, invalid_value)
+
+    def test_unknown_key_in_base_raises_error(self):
+        with self.assertRaises(ValueError):
+            self._callFUT({"optmizer_version": ""}, None)
+        with self.assertRaises(ValueError):
+            self._callFUT({"optmizer_version": "1"}, None)
+
+    def test_unknown_key_in_merge_raises_error(self):
+        with self.assertRaises(ValueError):
+            self._callFUT(None, {"optmizer_version": ""})
+        with self.assertRaises(ValueError):
+            self._callFUT(None, {"optmizer_version": "1"})
 
 
 class Test_get_cloud_region(unittest.TestCase):