diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py index 3c097a01ed..de359e8005 100644 --- a/google/cloud/spanner_v1/database.py +++ b/google/cloud/spanner_v1/database.py @@ -188,6 +188,7 @@ def __init__( self._reconciling = False self._directed_read_options = self._instance._client.directed_read_options self._proto_descriptors = proto_descriptors + self._channel_id = 0 # It'll be created when _spanner_api is created. if pool is None: pool = BurstyPool(database_role=database_role) @@ -446,22 +447,16 @@ def spanner_api(self): client_info=client_info, client_options=client_options, ) - return self._spanner_api - @property - def _channel_id(self): - """ - Helper to retrieve the associated channelID for the spanner_api. - This property is paramount to x-goog-spanner-request-id. - """ - with self.__transport_lock: - api = self.spanner_api - channel_id = self.__transports_to_channel_id.get(api._transport, None) - if channel_id is None: - channel_id = len(self.__transports_to_channel_id) + 1 - self.__transports_to_channel_id[api._transport] = channel_id - - return channel_id + with self.__transport_lock: + transport = self._spanner_api._transport + channel_id = self.__transports_to_channel_id.get(transport, None) + if channel_id is None: + channel_id = len(self.__transports_to_channel_id) + 1 + self.__transports_to_channel_id[transport] = channel_id + self._channel_id = channel_id + + return self._spanner_api def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]): return _metadata_with_request_id( diff --git a/tests/mockserver_tests/mock_server_test_base.py b/tests/mockserver_tests/mock_server_test_base.py index 771e414fd7..0fd68848b6 100644 --- a/tests/mockserver_tests/mock_server_test_base.py +++ b/tests/mockserver_tests/mock_server_test_base.py @@ -133,11 +133,8 @@ def setup_class(cls): def teardown_class(cls): if MockServerTestBase.server is not None: MockServerTestBase.server.stop(grace=None) + Client.NTH_CLIENT.reset() MockServerTestBase.server = None - self.reset() - - def reset(self): - Client.NTH_CLIENT.reset() def setup_method(self, *args, **kwargs): self._client = None diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index af6d500dad..a6ab1f176a 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -554,6 +554,10 @@ def test_create_grpc_error(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -583,6 +587,10 @@ def test_create_already_exists(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -611,6 +619,10 @@ def test_create_instance_not_found(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -649,6 +661,10 @@ def test_create_success(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -688,6 +704,10 @@ def test_create_success_w_encryption_config_dict(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -725,6 +745,10 @@ def test_create_success_w_proto_descriptors(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -745,6 +769,10 @@ def test_exists_grpc_error(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -764,6 +792,10 @@ def test_exists_not_found(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -785,6 +817,10 @@ def test_exists_success(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -805,6 +841,10 @@ def test_reload_grpc_error(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -825,6 +865,10 @@ def test_reload_not_found(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -886,12 +930,20 @@ def test_reload_success(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) api.get_database.assert_called_once_with( name=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.2.1", + ), ], ) @@ -920,6 +972,10 @@ def test_update_ddl_grpc_error(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -948,6 +1004,10 @@ def test_update_ddl_not_found(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -977,6 +1037,10 @@ def test_update_ddl(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1006,6 +1070,10 @@ def test_update_ddl_w_operation_id(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1034,6 +1102,10 @@ def test_update_success(self): update_mask=field_mask, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1064,6 +1136,10 @@ def test_update_ddl_w_proto_descriptors(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1084,6 +1160,10 @@ def test_drop_grpc_error(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1104,6 +1184,10 @@ def test_drop_not_found(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1123,6 +1207,10 @@ def test_drop_success(self): database=self.DATABASE_NAME, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1196,22 +1284,34 @@ def _execute_partitioned_dml_helper( exclude_txn_from_change_streams=exclude_txn_from_change_streams, ) - api.begin_transaction.assert_called_with( - session=session.name, - options=txn_options, - metadata=[ - ("google-cloud-resource-prefix", database.name), - ("x-goog-spanner-route-to-leader", "true"), - ( - "x-goog-spanner-request-id", - f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", - ), - ], - ) if retried: self.assertEqual(api.begin_transaction.call_count, 2) + api.begin_transaction.assert_called_with( + session=session.name, + options=txn_options, + metadata=[ + ("google-cloud-resource-prefix", database.name), + ("x-goog-spanner-route-to-leader", "true"), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.2", + ), + ], + ) else: self.assertEqual(api.begin_transaction.call_count, 1) + api.begin_transaction.assert_called_with( + session=session.name, + options=txn_options, + metadata=[ + ("google-cloud-resource-prefix", database.name), + ("x-goog-spanner-route-to-leader", "true"), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), + ], + ) if params: expected_params = Struct( @@ -1249,7 +1349,7 @@ def _execute_partitioned_dml_helper( ("x-goog-spanner-route-to-leader", "true"), ( "x-goog-spanner-request-id", - f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1", + f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.{database._channel_id}.2.1", ), ], ) @@ -1271,6 +1371,10 @@ def _execute_partitioned_dml_helper( metadata=[ ("google-cloud-resource-prefix", database.name), ("x-goog-spanner-route-to-leader", "true"), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.2.2", + ), ], ) self.assertEqual(api.execute_streaming_sql.call_count, 2) @@ -1547,6 +1651,10 @@ def test_restore_grpc_error(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1575,6 +1683,10 @@ def test_restore_not_found(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1614,6 +1726,10 @@ def test_restore_success(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1657,6 +1773,10 @@ def test_restore_success_w_encryption_config_dict(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1806,6 +1926,10 @@ def test_list_database_roles_grpc_error(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) @@ -1829,6 +1953,10 @@ def test_list_database_roles_defaults(self): request=expected_request, metadata=[ ("google-cloud-resource-prefix", database.name), + ( + "x-goog-spanner-request-id", + f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1", + ), ], ) self.assertIsNotNone(resp)