From fb6e3aeaeea32d7e59d14f74558e9b9654c10c2e Mon Sep 17 00:00:00 2001 From: maximilliangrand <214999687+maximilliangrand@users.noreply.github.com> Date: Fri, 16 Oct 2026 23:42:03 +0200 Subject: [PATCH] fix(spanner): preserve caller-owned request options --- .../google/cloud/spanner_v1/_helpers.py | 5 +- .../tests/unit/_async/test_snapshot.py | 40 +++++++++++++++ .../tests/unit/test_client_context.py | 26 +++++++++- .../tests/unit/test_snapshot.py | 50 ++++++++++++++----- 4 files changed, 104 insertions(+), 17 deletions(-) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py index f258b7884d9b..8d63f6bb910b 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py @@ -429,10 +429,7 @@ def _merge_request_options(request_options, client_context): if request_options is None and client_context is None: return None - if request_options is None: - request_options = RequestOptions() - elif isinstance(request_options, dict): - request_options = RequestOptions(request_options) + request_options = RequestOptions(request_options) if client_context: request_options.client_context = _merge_client_context( diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py index 1a37508fffb4..335eec572d14 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py @@ -1542,6 +1542,46 @@ async def test_read_w_request_options_dict(self): call_args = api.streaming_read.call_args self.assertEqual(call_args.kwargs["request"].request_options.priority, 1) + async def test_read_and_execute_sql_preserve_request_options(self): + from google.cloud.spanner_v1.keyset import KeySet + from google.cloud.spanner_v1.types import RequestOptions + + for method, args, rpc in ( + ("read", (TABLE_NAME, COLUMNS, KeySet(all_=True)), "streaming_read"), + ("execute_sql", (SQL_QUERY,), "execute_streaming_sql"), + ): + for options in ( + {"request_tag": "r-tag", "transaction_tag": "caller-tx-tag"}, + RequestOptions(request_tag="r-tag", transaction_tag="caller-tx-tag"), + ): + with self.subTest(method=method, options=options): + original = RequestOptions(options) + for read_only, transaction_tag, expected_tag in ( + (True, None, ""), + (False, None, "caller-tx-tag"), + (False, "snapshot-tx-tag", "snapshot-tx-tag"), + ): + database = _Database() + snapshot = self._make_snapshot(_Session(database)) + snapshot._transaction_id = TXN_ID + snapshot._read_only = read_only + snapshot.transaction_tag = transaction_tag + api_method = getattr(database.spanner_api, rpc) + api_method.return_value = _MockIterator(PartialResultSet()) + + results = await getattr(snapshot, method)( + *args, request_options=options + ) + async for _ in results: + pass + + sent_options = api_method.call_args.kwargs[ + "request" + ].request_options + self.assertEqual(sent_options.request_tag, "r-tag") + self.assertEqual(sent_options.transaction_tag, expected_tag) + self.assertEqual(RequestOptions(options), original) + async def test_read_w_directed_read_options(self): from google.cloud.spanner_v1.types import DirectedReadOptions diff --git a/packages/google-cloud-spanner/tests/unit/test_client_context.py b/packages/google-cloud-spanner/tests/unit/test_client_context.py index b371c1070f97..d3228fd1938a 100644 --- a/packages/google-cloud-spanner/tests/unit/test_client_context.py +++ b/packages/google-cloud-spanner/tests/unit/test_client_context.py @@ -64,7 +64,11 @@ def test__merge_client_context_both_set(self): self.assertEqual(result.secure_context["c"], "C") def test__merge_request_options_with_client_context(self): - request_options = RequestOptions(priority=RequestOptions.Priority.PRIORITY_LOW) + request_options = RequestOptions( + priority=RequestOptions.Priority.PRIORITY_LOW, + client_context={"secure_context": {"original": "value"}}, + ) + original = RequestOptions(request_options) client_context = RequestOptions.ClientContext( secure_context={"a": struct_pb2.Value(string_value="A")} ) @@ -73,6 +77,26 @@ def test__merge_request_options_with_client_context(self): self.assertEqual(result.priority, RequestOptions.Priority.PRIORITY_LOW) self.assertEqual(result.client_context.secure_context["a"], "A") + self.assertEqual(result.client_context.secure_context["original"], "value") + self.assertEqual(request_options, original) + + result.client_context.secure_context["original"] = "changed" + self.assertEqual(request_options, original) + + def test__merge_request_options_without_client_context_copies_options(self): + for options in ( + {"request_tag": "tag", "transaction_tag": "transaction"}, + RequestOptions(request_tag="tag", transaction_tag="transaction"), + ): + with self.subTest(options=options): + original = RequestOptions(options) + result = _merge_request_options(options, None) + self.assertEqual(result, original) + result.transaction_tag = "changed" + self.assertEqual(RequestOptions(options), original) + + def test__merge_request_options_both_none(self): + self.assertIsNone(_merge_request_options(None, None)) def test_client_init_with_client_context(self): from google.cloud.spanner_v1.client import Client diff --git a/packages/google-cloud-spanner/tests/unit/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/test_snapshot.py index 0729a4734d6e..585898653f07 100644 --- a/packages/google-cloud-spanner/tests/unit/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/test_snapshot.py @@ -1866,12 +1866,25 @@ def test_read_w_transaction_tag_success(self, mock_region): return_value="global", ) def test_read_w_request_and_transaction_tag_success(self, mock_region): - request_options = {"request_tag": "tag-1", "transaction_tag": "tag-1-1"} - self._execute_read( - multi_use=False, - request_options=request_options, - expected_request_options=RequestOptions(request_tag="tag-1"), - ) + for request_options in ( + {"request_tag": "tag-1", "transaction_tag": "tag-1-1"}, + RequestOptions(request_tag="tag-1", transaction_tag="tag-1-1"), + ): + with self.subTest(request_options=request_options): + original = RequestOptions(request_options) + self._execute_read( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(request_tag="tag-1"), + ) + self.assertEqual(RequestOptions(request_options), original) + self._execute_read( + multi_use=False, + read_only=False, + request_options=request_options, + expected_request_options=original, + ) + self.assertEqual(RequestOptions(request_options), original) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", @@ -2372,12 +2385,25 @@ def test_execute_sql_w_transaction_tag_success(self, mock_region): return_value="global", ) def test_execute_sql_w_request_and_transaction_tag_success(self, mock_region): - request_options = {"request_tag": "tag-1", "transaction_tag": "tag-1-1"} - self._execute_sql_helper( - multi_use=False, - request_options=request_options, - expected_request_options=RequestOptions(request_tag="tag-1"), - ) + for request_options in ( + {"request_tag": "tag-1", "transaction_tag": "tag-1-1"}, + RequestOptions(request_tag="tag-1", transaction_tag="tag-1-1"), + ): + with self.subTest(request_options=request_options): + original = RequestOptions(request_options) + self._execute_sql_helper( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(request_tag="tag-1"), + ) + self.assertEqual(RequestOptions(request_options), original) + self._execute_sql_helper( + multi_use=False, + read_only=False, + request_options=request_options, + expected_request_options=original, + ) + self.assertEqual(RequestOptions(request_options), original) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region",