From f53eaeb5713a29fc0191ff16d4164f69d571c8e7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Fri, 18 Sep 2026 14:08:35 +0200 Subject: [PATCH] perf(spanner): avoid allocating empty RequestOptions on read and query paths Previously, Snapshot.read(), Snapshot.execute_sql(), Transaction.execute_update(), and Transaction.batch_update() unconditionally initialized an empty `RequestOptions` message when no options were supplied, even when no transaction tag or client context was present. This empty submessage was then copied into the outgoing ExecuteSqlRequest or ReadRequest. Avoid allocating `RequestOptions` on the default execution path when both `request_options` and `transaction_tag` are unset. Only construct `RequestOptions` when an explicit option or transaction tag is provided, while preserving the existing behavior of overriding or clearing tags when `request_options` is supplied. --- .../cloud/spanner_v1/_async/snapshot.py | 56 ++--- .../cloud/spanner_v1/_async/transaction.py | 18 +- .../google/cloud/spanner_v1/client.py | 13 +- .../google/cloud/spanner_v1/snapshot.py | 52 ++-- .../google/cloud/spanner_v1/transaction.py | 18 +- .../tests/unit/_async/test_snapshot.py | 165 +++++++++++++ .../tests/unit/_async/test_transaction.py | 64 +++-- .../tests/unit/test_snapshot.py | 222 ++++++++++++++---- .../tests/unit/test_transaction.py | 58 +++-- 9 files changed, 497 insertions(+), 169 deletions(-) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py index 5e64e106db99..c1d5e763ca56 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py @@ -409,20 +409,20 @@ async def read( ) request_options = _merge_request_options(request_options, client_context) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - - if self._read_only: - request_options.transaction_tag = None - if ( - directed_read_options is None - and database._directed_read_options is not None - ): - directed_read_options = database._directed_read_options - elif self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag + if request_options is not None: + if self._read_only: + request_options.transaction_tag = None + elif self.transaction_tag is not None: + request_options.transaction_tag = self.transaction_tag + elif not self._read_only and self.transaction_tag is not None: + request_options = RequestOptions(transaction_tag=self.transaction_tag) + + if ( + self._read_only + and directed_read_options is None + and database._directed_read_options is not None + ): + directed_read_options = database._directed_read_options read_request = ReadRequest( session=session.name, @@ -603,20 +603,20 @@ async def execute_sql( ) request_options = _merge_request_options(request_options, client_context) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - - if self._read_only: - request_options.transaction_tag = None - if ( - directed_read_options is None - and database._directed_read_options is not None - ): - directed_read_options = database._directed_read_options - elif self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag + if request_options is not None: + if self._read_only: + request_options.transaction_tag = None + elif self.transaction_tag is not None: + request_options.transaction_tag = self.transaction_tag + elif not self._read_only and self.transaction_tag is not None: + request_options = RequestOptions(transaction_tag=self.transaction_tag) + + if ( + self._read_only + and directed_read_options is None + and database._directed_read_options is not None + ): + directed_read_options = database._directed_read_options execute_sql_request = ExecuteSqlRequest( session=session.name, diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py index 425857e45095..c27b9bff5a88 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py @@ -524,11 +524,10 @@ async def execute_update( ) request_options = _merge_request_options(request_options, client_context) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - request_options.transaction_tag = self.transaction_tag + if request_options is not None: + request_options.transaction_tag = self.transaction_tag + elif self.transaction_tag is not None: + request_options = RequestOptions(transaction_tag=self.transaction_tag) trace_attributes = { "db.statement": dml, @@ -684,11 +683,10 @@ async def batch_update( ) request_options = _merge_request_options(request_options, client_context) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - request_options.transaction_tag = self.transaction_tag + if request_options is not None: + request_options.transaction_tag = self.transaction_tag + elif self.transaction_tag is not None: + request_options = RequestOptions(transaction_tag=self.transaction_tag) trace_attributes = { # Get just the queries from the DML statement batch diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py index 7ae7f58a0e0d..5fb7d019212d 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py @@ -320,9 +320,7 @@ def __init__( raise ValueError( "Both username and password must be specified for Omni authentication" ) - from google.cloud.spanner_v1.omni.credentials import ( - SpannerOmniCredentials, - ) + from google.cloud.spanner_v1.omni.credentials import SpannerOmniCredentials if has_username and has_password: credentials = SpannerOmniCredentials( @@ -339,11 +337,10 @@ def __init__( disable_builtin_metrics = True elif isinstance(credentials, AnonymousCredentials): self._emulator_host = self._client_options.api_endpoint - else: - if username is not None or password is not None: - raise ValueError( - "username and password can only be used when instance_type='omni'." - ) + elif username is not None or password is not None: + raise ValueError( + "username and password can only be used when instance_type='omni'." + ) super(Client, self).__init__( project=project, credentials=credentials, diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py index ce2bd07306d8..33bfdce69092 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py @@ -362,19 +362,19 @@ def read( database._instance._client._client_context, self._client_context ) request_options = _merge_request_options(request_options, client_context) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - if self._read_only: - request_options.transaction_tag = None - if ( - directed_read_options is None - and database._directed_read_options is not None - ): - directed_read_options = database._directed_read_options - elif self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag + if request_options is not None: + if self._read_only: + request_options.transaction_tag = None + elif self.transaction_tag is not None: + request_options.transaction_tag = self.transaction_tag + elif not self._read_only and self.transaction_tag is not None: + request_options = RequestOptions(transaction_tag=self.transaction_tag) + if ( + self._read_only + and directed_read_options is None + and (database._directed_read_options is not None) + ): + directed_read_options = database._directed_read_options read_request = ReadRequest( session=session.name, table=table, @@ -543,19 +543,19 @@ def execute_sql( database._instance._client._client_context, self._client_context ) request_options = _merge_request_options(request_options, client_context) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - if self._read_only: - request_options.transaction_tag = None - if ( - directed_read_options is None - and database._directed_read_options is not None - ): - directed_read_options = database._directed_read_options - elif self.transaction_tag is not None: - request_options.transaction_tag = self.transaction_tag + if request_options is not None: + if self._read_only: + request_options.transaction_tag = None + elif self.transaction_tag is not None: + request_options.transaction_tag = self.transaction_tag + elif not self._read_only and self.transaction_tag is not None: + request_options = RequestOptions(transaction_tag=self.transaction_tag) + if ( + self._read_only + and directed_read_options is None + and (database._directed_read_options is not None) + ): + directed_read_options = database._directed_read_options execute_sql_request = ExecuteSqlRequest( session=session.name, sql=sql, diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py index 1bd8d01d348d..221aa23765ff 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py @@ -436,11 +436,10 @@ def execute_update( database._instance._client._client_context, self._client_context ) request_options = _merge_request_options(request_options, client_context) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - request_options.transaction_tag = self.transaction_tag + if request_options is not None: + request_options.transaction_tag = self.transaction_tag + elif self.transaction_tag is not None: + request_options = RequestOptions(transaction_tag=self.transaction_tag) trace_attributes = {"db.statement": dml, "request_options": request_options} is_inline_begin = False if self._transaction_id is None: @@ -575,11 +574,10 @@ def batch_update( database._instance._client._client_context, self._client_context ) request_options = _merge_request_options(request_options, client_context) - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - request_options.transaction_tag = self.transaction_tag + if request_options is not None: + request_options.transaction_tag = self.transaction_tag + elif self.transaction_tag is not None: + request_options = RequestOptions(transaction_tag=self.transaction_tag) trace_attributes = { "db.statement": ";".join([statement.sql for statement in parsed]), "request_options": request_options, 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 e4e8785f2c43..9e74039ee867 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py @@ -701,6 +701,93 @@ async def test_read_w_transaction_tag(self): call_args.kwargs["request"].request_options.transaction_tag, "tag" ) + async def test_read_w_transaction_tag_default_request_options(self): + database = _Database() + api = database.spanner_api + session = _Session(database) + snapshot = self._make_snapshot(session) + snapshot._transaction_id = TXN_ID + snapshot.transaction_tag = "tag" + snapshot._read_only = False + + api.streaming_read.return_value = _MockIterator(PartialResultSet()) + from google.cloud.spanner_v1.keyset import KeySet + + results = await snapshot.read(TABLE_NAME, COLUMNS, KeySet(all_=True)) + async for _ in results: + pass + + call_args = api.streaming_read.call_args + self.assertIsNotNone(call_args, "streaming_read should have been called") + self.assertEqual( + call_args.kwargs["request"].request_options.transaction_tag, "tag" + ) + + async def test_read_w_request_options_dict_with_transaction_tag(self): + database = _Database() + api = database.spanner_api + session = _Session(database) + snapshot = self._make_snapshot(session) + snapshot._transaction_id = TXN_ID + snapshot.transaction_tag = "tag" + snapshot._read_only = False + + api.streaming_read.return_value = _MockIterator(PartialResultSet()) + from google.cloud.spanner_v1.keyset import KeySet + + raw_dict_options = { + "request_tag": "r-tag", + "transaction_tag": "caller-tx-tag", + } + results = await snapshot.read( + TABLE_NAME, COLUMNS, KeySet(all_=True), request_options=raw_dict_options + ) + async for _ in results: + pass + + call_args = api.streaming_read.call_args + self.assertIsNotNone(call_args, "streaming_read should have been called") + self.assertEqual( + call_args.kwargs["request"].request_options.request_tag, "r-tag" + ) + self.assertEqual( + call_args.kwargs["request"].request_options.transaction_tag, "tag" + ) + + async def test_read_w_request_options_caller_tag_preserved_when_snapshot_tag_is_none( + self, + ): + database = _Database() + api = database.spanner_api + session = _Session(database) + snapshot = self._make_snapshot(session) + snapshot._transaction_id = TXN_ID + snapshot.transaction_tag = None + snapshot._read_only = False + + api.streaming_read.return_value = _MockIterator(PartialResultSet()) + from google.cloud.spanner_v1.keyset import KeySet + + raw_dict_options = { + "request_tag": "r-tag", + "transaction_tag": "caller-tx-tag", + } + results = await snapshot.read( + TABLE_NAME, COLUMNS, KeySet(all_=True), request_options=raw_dict_options + ) + async for _ in results: + pass + + call_args = api.streaming_read.call_args + self.assertIsNotNone(call_args, "streaming_read should have been called") + self.assertEqual( + call_args.kwargs["request"].request_options.request_tag, "r-tag" + ) + self.assertEqual( + call_args.kwargs["request"].request_options.transaction_tag, + "caller-tx-tag", + ) + async def test_execute_sql_w_request_options_dict(self): database = _Database() api = database.spanner_api @@ -827,6 +914,84 @@ async def test_execute_sql_w_transaction_tag(self): call_args.kwargs["request"].request_options.transaction_tag, "tag" ) + async def test_execute_sql_w_transaction_tag_default_request_options(self): + database = _Database() + api = database.spanner_api + session = _Session(database) + snapshot = self._make_snapshot(session) + snapshot._transaction_id = TXN_ID + snapshot.transaction_tag = "tag" + snapshot._read_only = False + + api.execute_streaming_sql.return_value = _MockIterator(PartialResultSet()) + results = await snapshot.execute_sql(SQL_QUERY) + async for _ in results: + pass + + call_args = api.execute_streaming_sql.call_args + self.assertEqual( + call_args.kwargs["request"].request_options.transaction_tag, "tag" + ) + + async def test_execute_sql_w_request_options_dict_with_transaction_tag(self): + database = _Database() + api = database.spanner_api + session = _Session(database) + snapshot = self._make_snapshot(session) + snapshot._transaction_id = TXN_ID + snapshot.transaction_tag = "tag" + snapshot._read_only = False + + api.execute_streaming_sql.return_value = _MockIterator(PartialResultSet()) + raw_dict_options = { + "request_tag": "r-tag", + "transaction_tag": "caller-tx-tag", + } + results = await snapshot.execute_sql( + SQL_QUERY, request_options=raw_dict_options + ) + async for _ in results: + pass + + call_args = api.execute_streaming_sql.call_args + self.assertEqual( + call_args.kwargs["request"].request_options.request_tag, "r-tag" + ) + self.assertEqual( + call_args.kwargs["request"].request_options.transaction_tag, "tag" + ) + + async def test_execute_sql_w_request_options_caller_tag_preserved_when_snapshot_tag_is_none( + self, + ): + database = _Database() + api = database.spanner_api + session = _Session(database) + snapshot = self._make_snapshot(session) + snapshot._transaction_id = TXN_ID + snapshot.transaction_tag = None + snapshot._read_only = False + + api.execute_streaming_sql.return_value = _MockIterator(PartialResultSet()) + raw_dict_options = { + "request_tag": "r-tag", + "transaction_tag": "caller-tx-tag", + } + results = await snapshot.execute_sql( + SQL_QUERY, request_options=raw_dict_options + ) + async for _ in results: + pass + + call_args = api.execute_streaming_sql.call_args + self.assertEqual( + call_args.kwargs["request"].request_options.request_tag, "r-tag" + ) + self.assertEqual( + call_args.kwargs["request"].request_options.transaction_tag, + "caller-tx-tag", + ) + def test_ctor_incompatible_options(self): import datetime diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_transaction.py b/packages/google-cloud-spanner/tests/unit/_async/test_transaction.py index 7b0edeb4d171..0262ac7e468f 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_transaction.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_transaction.py @@ -799,6 +799,7 @@ async def _execute_update_helper( timeout=gapic_v1.method.DEFAULT, begin=True, use_multiplexed=False, + transaction_tag=TRANSACTION_TAG, ): from google.protobuf.struct_pb2 import Struct @@ -832,17 +833,12 @@ async def _execute_update_helper( session = _Session(database) transaction = self._make_one(session) - transaction.transaction_tag = TRANSACTION_TAG + transaction.transaction_tag = transaction_tag transaction._execute_sql_request_count = count if begin: transaction._transaction_id = TRANSACTION_ID - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - row_count = await transaction.execute_update( DML_QUERY_WITH_PARAM, PARAMS, @@ -873,9 +869,13 @@ async def _execute_update_helper( expected_query_options = _merge_query_options( expected_query_options, query_options ) - expected_request_options = RequestOptions(request_options) - if request_options.request_tag: - expected_request_options.request_tag = request_options.request_tag + if request_options is not None: + expected_request_options = RequestOptions(request_options) + expected_request_options.transaction_tag = transaction_tag + elif transaction_tag is not None: + expected_request_options = RequestOptions(transaction_tag=transaction_tag) + else: + expected_request_options = None expected_request = ExecuteSqlRequest( session=self.SESSION_NAME, @@ -905,8 +905,13 @@ async def _execute_update_helper( expected_attributes = self._build_span_attributes( database, **{"db.statement": DML_QUERY_WITH_PARAM} ) - if request_options.request_tag: - expected_attributes["request.tag"] = request_options.request_tag + request_tag = ( + request_options.get("request_tag") + if isinstance(request_options, dict) + else getattr(request_options, "request_tag", None) + ) + if request_tag: + expected_attributes["request.tag"] = request_tag self.assertSpanAttributes( "CloudSpanner.Transaction.execute_update", attributes=expected_attributes ) @@ -1044,6 +1049,16 @@ async def test_execute_update_w_query_options(self, mock_region): async def test_execute_update_wo_begin(self, mock_region): await self._execute_update_helper(begin=False) + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + @CrossSync.pytest + async def test_execute_update_wo_request_options_and_transaction_tag( + self, mock_region + ): + await self._execute_update_helper(transaction_tag=None) + @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", return_value="global", @@ -1108,6 +1123,7 @@ async def _batch_update_helper( timeout=gapic_v1.method.DEFAULT, begin=True, use_multiplexed=False, + transaction_tag=TRANSACTION_TAG, ): from google.protobuf.struct_pb2 import Struct from google.rpc.status_pb2 import Status @@ -1172,17 +1188,12 @@ async def _batch_update_helper( session = _Session(database) transaction = self._make_one(session) - transaction.transaction_tag = TRANSACTION_TAG + transaction.transaction_tag = transaction_tag transaction._execute_sql_request_count = count if begin: transaction._transaction_id = TRANSACTION_ID - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - status, row_counts = await transaction.batch_update( dml_statements, request_options=request_options, @@ -1217,8 +1228,13 @@ async def _batch_update_helper( ExecuteBatchDmlRequest.Statement(sql=update_dml), ExecuteBatchDmlRequest.Statement(sql=delete_dml), ] - expected_request_options = request_options - expected_request_options.transaction_tag = TRANSACTION_TAG + if request_options is not None: + expected_request_options = RequestOptions(request_options) + expected_request_options.transaction_tag = transaction_tag + elif transaction_tag is not None: + expected_request_options = RequestOptions(transaction_tag=transaction_tag) + else: + expected_request_options = None expected_request = ExecuteBatchDmlRequest( session=self.SESSION_NAME, @@ -1255,6 +1271,16 @@ async def _batch_update_helper( async def test_batch_update_wo_begin(self, mock_region): await self._batch_update_helper(begin=False) + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + @CrossSync.pytest + async def test_batch_update_wo_request_options_and_transaction_tag( + self, mock_region + ): + await self._batch_update_helper(transaction_tag=None) + @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", return_value="global", diff --git a/packages/google-cloud-spanner/tests/unit/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/test_snapshot.py index 4a65d46aea8c..47b30c8914b0 100644 --- a/packages/google-cloud-spanner/tests/unit/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/test_snapshot.py @@ -58,6 +58,7 @@ TABLE_NAME = "citizens" COLUMNS = ["email", "first_name", "last_name", "age"] +_DEFAULT_EXPECTED_REQUEST_OPTIONS = object() SQL_QUERY = """\ SELECT first_name, last_name, age FROM citizens ORDER BY age""" SQL_QUERY_WITH_PARAM = """ @@ -1000,6 +1001,9 @@ def _execute_read( directed_read_options=None, directed_read_options_at_client_level=None, use_multiplexed=False, + read_only=True, + transaction_tag=None, + expected_request_options=_DEFAULT_EXPECTED_REQUEST_OPTIONS, ): """Helper for testing _SnapshotBase.read(). Executes method and verifies transaction state, begin transaction API call, and span attributes and events. @@ -1067,17 +1071,28 @@ def _execute_read( api = database.spanner_api = build_spanner_api() api.streaming_read.return_value = _MockIterator(*result_sets) session = _Session(database) - derived = _build_snapshot_derived(session) + derived = _build_snapshot_derived(session, read_only=read_only) + derived.transaction_tag = transaction_tag derived._multi_use = multi_use derived._read_request_count = count if not first: derived._transaction_id = TXN_ID - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) + if expected_request_options is not _DEFAULT_EXPECTED_REQUEST_OPTIONS: + actual_expected_request_options = expected_request_options + elif request_options is not None: + actual_expected_request_options = RequestOptions(request_options) + if read_only: + actual_expected_request_options.transaction_tag = None + elif transaction_tag is not None: + actual_expected_request_options.transaction_tag = transaction_tag + elif not read_only and transaction_tag is not None: + actual_expected_request_options = RequestOptions( + transaction_tag=transaction_tag + ) + else: + actual_expected_request_options = None transaction_selector_pb = derived._build_transaction_selector_pb() @@ -1117,11 +1132,6 @@ def _execute_read( else: expected_limit = LIMIT - # Transaction tag is ignored for read request. - expected_request_options = RequestOptions(request_options) - if derived.transaction_tag: - expected_request_options.transaction_tag = derived.transaction_tag - expected_directed_read_options = ( directed_read_options if directed_read_options is not None @@ -1137,19 +1147,19 @@ def _execute_read( index=INDEX, limit=expected_limit, partition_token=partition, - request_options=expected_request_options, + request_options=actual_expected_request_options, directed_read_options=expected_directed_read_options, ) req_id = f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1" + expected_metadata = [ + ("google-cloud-resource-prefix", database.name), + ] + if not read_only: + expected_metadata.append(("x-goog-spanner-route-to-leader", "true")) + expected_metadata.append(("x-goog-spanner-request-id", req_id)) api.streaming_read.assert_called_once_with( request=expected_request, - metadata=[ - ("google-cloud-resource-prefix", database.name), - ( - "x-goog-spanner-request-id", - req_id, - ), - ], + metadata=expected_metadata, retry=retry, timeout=timeout, ) @@ -1159,8 +1169,13 @@ def _execute_read( columns=tuple(COLUMNS), x_goog_spanner_request_id=req_id, ) - if request_options and request_options.request_tag: - expected_attributes["request.tag"] = request_options.request_tag + request_tag = ( + request_options.get("request_tag") + if isinstance(request_options, dict) + else getattr(request_options, "request_tag", None) + ) + if request_tag: + expected_attributes["request.tag"] = request_tag self.assertSpanAttributes( "CloudSpanner._Derived.read", attributes=expected_attributes ) @@ -1184,7 +1199,11 @@ def test_read_wo_multi_use(self, mock_region): ) def test_read_w_request_tag_success(self, mock_region): request_options = {"request_tag": "tag-1"} - self._execute_read(multi_use=False, request_options=request_options) + self._execute_read( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(request_tag="tag-1"), + ) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", @@ -1192,7 +1211,11 @@ def test_read_w_request_tag_success(self, mock_region): ) def test_read_w_transaction_tag_success(self, mock_region): request_options = {"transaction_tag": "tag-1-1"} - self._execute_read(multi_use=False, request_options=request_options) + self._execute_read( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(), + ) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", @@ -1200,7 +1223,11 @@ def test_read_w_transaction_tag_success(self, mock_region): ) 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) + self._execute_read( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(request_tag="tag-1"), + ) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", @@ -1208,7 +1235,41 @@ def test_read_w_request_and_transaction_tag_success(self, mock_region): ) def test_read_w_request_and_transaction_tag_dictionary_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) + self._execute_read( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(request_tag="tag-1"), + ) + + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + def test_read_in_read_write_transaction_with_tag(self, mock_region): + self._execute_read( + multi_use=False, + read_only=False, + transaction_tag="tx-tag", + expected_request_options=RequestOptions(transaction_tag="tx-tag"), + ) + + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + def test_read_in_read_write_transaction_caller_tag_preserved(self, mock_region): + self._execute_read( + multi_use=False, + read_only=False, + transaction_tag=None, + request_options={ + "request_tag": "r-tag", + "transaction_tag": "caller-tx-tag", + }, + expected_request_options=RequestOptions( + request_tag="r-tag", transaction_tag="caller-tx-tag" + ), + ) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", @@ -1369,6 +1430,9 @@ def _execute_sql_helper( directed_read_options=None, directed_read_options_at_client_level=None, use_multiplexed=False, + read_only=True, + transaction_tag=None, + expected_request_options=_DEFAULT_EXPECTED_REQUEST_OPTIONS, ): """Helper for testing _SnapshotBase.execute_sql(). Executes method and verifies transaction state, begin transaction API call, and span attributes and events. @@ -1437,16 +1501,29 @@ def _execute_sql_helper( api = database.spanner_api = build_spanner_api() api.execute_streaming_sql.return_value = iterator session = _Session(database) - derived = _build_snapshot_derived(session, multi_use=multi_use) + derived = _build_snapshot_derived( + session, multi_use=multi_use, read_only=read_only + ) + derived.transaction_tag = transaction_tag derived._read_request_count = count derived._execute_sql_request_count = sql_count if not first: derived._transaction_id = TXN_ID - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) + if expected_request_options is not _DEFAULT_EXPECTED_REQUEST_OPTIONS: + actual_expected_request_options = expected_request_options + elif request_options is not None: + actual_expected_request_options = RequestOptions(request_options) + if read_only: + actual_expected_request_options.transaction_tag = None + elif transaction_tag is not None: + actual_expected_request_options.transaction_tag = transaction_tag + elif not read_only and transaction_tag is not None: + actual_expected_request_options = RequestOptions( + transaction_tag=transaction_tag + ) + else: + actual_expected_request_options = None transaction_selector_pb = derived._build_transaction_selector_pb() @@ -1479,12 +1556,6 @@ def _execute_sql_helper( expected_query_options, query_options ) - expected_request_options = RequestOptions(request_options) - if derived.transaction_tag: - expected_request_options.transaction_tag = derived.transaction_tag - if not derived._read_only and request_options.request_tag: - expected_request_options.request_tag = request_options.request_tag - expected_directed_read_options = ( directed_read_options if directed_read_options is not None @@ -1499,21 +1570,21 @@ def _execute_sql_helper( param_types=PARAM_TYPES, query_mode=MODE, query_options=expected_query_options, - request_options=expected_request_options, + request_options=actual_expected_request_options, partition_token=partition, seqno=sql_count, directed_read_options=expected_directed_read_options, ) req_id = f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1" + expected_metadata = [ + ("google-cloud-resource-prefix", database.name), + ] + if not read_only: + expected_metadata.append(("x-goog-spanner-route-to-leader", "true")) + expected_metadata.append(("x-goog-spanner-request-id", req_id)) api.execute_streaming_sql.assert_called_once_with( request=expected_request, - metadata=[ - ("google-cloud-resource-prefix", database.name), - ( - "x-goog-spanner-request-id", - req_id, - ), - ], + metadata=expected_metadata, timeout=timeout, retry=retry, ) @@ -1527,8 +1598,13 @@ def _execute_sql_helper( "x_goog_spanner_request_id": req_id, }, ) - if request_options and request_options.request_tag: - expected_attributes["request.tag"] = request_options.request_tag + request_tag = ( + request_options.get("request_tag") + if isinstance(request_options, dict) + else getattr(request_options, "request_tag", None) + ) + if request_tag: + expected_attributes["request.tag"] = request_tag self.assertSpanAttributes( "CloudSpanner._Derived.execute_sql", @@ -1629,7 +1705,11 @@ def test_execute_sql_w_request_options(self, mock_region): ) def test_execute_sql_w_request_tag_success(self, mock_region): request_options = {"request_tag": "tag-1"} - self._execute_sql_helper(multi_use=False, request_options=request_options) + self._execute_sql_helper( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(request_tag="tag-1"), + ) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", @@ -1637,7 +1717,11 @@ def test_execute_sql_w_request_tag_success(self, mock_region): ) def test_execute_sql_w_transaction_tag_success(self, mock_region): request_options = {"transaction_tag": "tag-1-1"} - self._execute_sql_helper(multi_use=False, request_options=request_options) + self._execute_sql_helper( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(), + ) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", @@ -1645,7 +1729,11 @@ def test_execute_sql_w_transaction_tag_success(self, mock_region): ) 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) + self._execute_sql_helper( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(request_tag="tag-1"), + ) @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", @@ -1655,7 +1743,43 @@ def test_execute_sql_w_request_and_transaction_tag_dictionary_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) + self._execute_sql_helper( + multi_use=False, + request_options=request_options, + expected_request_options=RequestOptions(request_tag="tag-1"), + ) + + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + def test_execute_sql_in_read_write_transaction_with_tag(self, mock_region): + self._execute_sql_helper( + multi_use=False, + read_only=False, + transaction_tag="tx-tag", + expected_request_options=RequestOptions(transaction_tag="tx-tag"), + ) + + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + def test_execute_sql_in_read_write_transaction_caller_tag_preserved( + self, mock_region + ): + self._execute_sql_helper( + multi_use=False, + read_only=False, + transaction_tag=None, + request_options={ + "request_tag": "r-tag", + "transaction_tag": "caller-tx-tag", + }, + expected_request_options=RequestOptions( + request_tag="r-tag", transaction_tag="caller-tx-tag" + ), + ) def test_execute_sql_w_incorrect_tag_dictionary_error(self): request_options = {"incorrect_tag": "tag-1-1"} diff --git a/packages/google-cloud-spanner/tests/unit/test_transaction.py b/packages/google-cloud-spanner/tests/unit/test_transaction.py index a444873c2474..b5698b2750c5 100644 --- a/packages/google-cloud-spanner/tests/unit/test_transaction.py +++ b/packages/google-cloud-spanner/tests/unit/test_transaction.py @@ -756,6 +756,7 @@ def _execute_update_helper( timeout=gapic_v1.method.DEFAULT, begin=True, use_multiplexed=False, + transaction_tag=TRANSACTION_TAG, ): from google.protobuf.struct_pb2 import Struct @@ -789,17 +790,12 @@ def _execute_update_helper( session = _Session(database) transaction = self._make_one(session) - transaction.transaction_tag = TRANSACTION_TAG + transaction.transaction_tag = transaction_tag transaction._execute_sql_request_count = count if begin: transaction._transaction_id = TRANSACTION_ID - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - row_count = transaction.execute_update( DML_QUERY_WITH_PARAM, PARAMS, @@ -830,9 +826,13 @@ def _execute_update_helper( expected_query_options = _merge_query_options( expected_query_options, query_options ) - expected_request_options = RequestOptions(request_options) - if request_options.request_tag: - expected_request_options.request_tag = request_options.request_tag + if request_options is not None: + expected_request_options = RequestOptions(request_options) + expected_request_options.transaction_tag = transaction_tag + elif transaction_tag is not None: + expected_request_options = RequestOptions(transaction_tag=transaction_tag) + else: + expected_request_options = None expected_request = ExecuteSqlRequest( session=self.SESSION_NAME, @@ -862,8 +862,13 @@ def _execute_update_helper( expected_attributes = self._build_span_attributes( database, **{"db.statement": DML_QUERY_WITH_PARAM} ) - if request_options.request_tag: - expected_attributes["request.tag"] = request_options.request_tag + request_tag = ( + request_options.get("request_tag") + if isinstance(request_options, dict) + else getattr(request_options, "request_tag", None) + ) + if request_tag: + expected_attributes["request.tag"] = request_tag self.assertSpanAttributes( "CloudSpanner.Transaction.execute_update", attributes=expected_attributes ) @@ -986,6 +991,13 @@ def test_execute_update_w_query_options(self, mock_region): def test_execute_update_wo_begin(self, mock_region): self._execute_update_helper(begin=False) + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + def test_execute_update_wo_request_options_and_transaction_tag(self, mock_region): + self._execute_update_helper(transaction_tag=None) + @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", return_value="global", @@ -1046,6 +1058,7 @@ def _batch_update_helper( timeout=gapic_v1.method.DEFAULT, begin=True, use_multiplexed=False, + transaction_tag=TRANSACTION_TAG, ): from google.protobuf.struct_pb2 import Struct from google.rpc.status_pb2 import Status @@ -1110,17 +1123,12 @@ def _batch_update_helper( session = _Session(database) transaction = self._make_one(session) - transaction.transaction_tag = TRANSACTION_TAG + transaction.transaction_tag = transaction_tag transaction._execute_sql_request_count = count if begin: transaction._transaction_id = TRANSACTION_ID - if request_options is None: - request_options = RequestOptions() - elif type(request_options) is dict: - request_options = RequestOptions(request_options) - status, row_counts = transaction.batch_update( dml_statements, request_options=request_options, @@ -1155,8 +1163,13 @@ def _batch_update_helper( ExecuteBatchDmlRequest.Statement(sql=update_dml), ExecuteBatchDmlRequest.Statement(sql=delete_dml), ] - expected_request_options = request_options - expected_request_options.transaction_tag = TRANSACTION_TAG + if request_options is not None: + expected_request_options = RequestOptions(request_options) + expected_request_options.transaction_tag = transaction_tag + elif transaction_tag is not None: + expected_request_options = RequestOptions(transaction_tag=transaction_tag) + else: + expected_request_options = None expected_request = ExecuteBatchDmlRequest( session=self.SESSION_NAME, @@ -1192,6 +1205,13 @@ def _batch_update_helper( def test_batch_update_wo_begin(self, mock_region): self._batch_update_helper(begin=False) + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + def test_batch_update_wo_request_options_and_transaction_tag(self, mock_region): + self._batch_update_helper(transaction_tag=None) + @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", return_value="global",