diff --git a/google/cloud/bigtable_v2/services/bigtable/async_client.py b/google/cloud/bigtable_v2/services/bigtable/async_client.py index 0ae20f3a2..84832ffd6 100644 --- a/google/cloud/bigtable_v2/services/bigtable/async_client.py +++ b/google/cloud/bigtable_v2/services/bigtable/async_client.py @@ -1584,7 +1584,9 @@ async def prepare_query( if regex_match and regex_match.group("name"): header_params["name"] = regex_match.group("name") - if request.app_profile_id: + if request.app_profile_id is not None: + # prepare_query currently requires app_profile_id header to be set + # even when the request param is unpopulated TODO: remove after support is added header_params["app_profile_id"] = request.app_profile_id if header_params: @@ -1704,7 +1706,9 @@ def execute_query( if regex_match and regex_match.group("name"): header_params["name"] = regex_match.group("name") - if request.app_profile_id: + if request.app_profile_id is not None: + # execute_query currently requires app_profile_id header to be set + # even when the request param is unpopulated TODO: remove after support is added header_params["app_profile_id"] = request.app_profile_id if header_params: diff --git a/google/cloud/bigtable_v2/services/bigtable/client.py b/google/cloud/bigtable_v2/services/bigtable/client.py index 2835a10cf..c5b14b54a 100644 --- a/google/cloud/bigtable_v2/services/bigtable/client.py +++ b/google/cloud/bigtable_v2/services/bigtable/client.py @@ -2033,7 +2033,9 @@ def prepare_query( if regex_match and regex_match.group("name"): header_params["name"] = regex_match.group("name") - if request.app_profile_id: + if request.app_profile_id is not None: + # prepare_query currently requires app_profile_id header to be set + # even when the request param is unpopulated TODO: remove after support is added header_params["app_profile_id"] = request.app_profile_id if header_params: @@ -2150,7 +2152,9 @@ def execute_query( if regex_match and regex_match.group("name"): header_params["name"] = regex_match.group("name") - if request.app_profile_id: + if request.app_profile_id is not None: + # execute_query currently requires app_profile_id header to be set + # even when the request param is unpopulated TODO: remove after support is added header_params["app_profile_id"] = request.app_profile_id if header_params: diff --git a/tests/unit/gapic/bigtable_v2/test_bigtable.py b/tests/unit/gapic/bigtable_v2/test_bigtable.py index 059e6a58c..1750be32b 100644 --- a/tests/unit/gapic/bigtable_v2/test_bigtable.py +++ b/tests/unit/gapic/bigtable_v2/test_bigtable.py @@ -7442,7 +7442,11 @@ def test_prepare_query_routing_parameters_request_1_grpc(): assert args[0] == request_msg - expected_headers = {"name": "projects/sample1/instances/sample2"} + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1072 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } assert ( gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] ) @@ -7494,7 +7498,11 @@ def test_execute_query_routing_parameters_request_1_grpc(): assert args[0] == request_msg - expected_headers = {"name": "projects/sample1/instances/sample2"} + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1072 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } assert ( gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] ) @@ -8548,7 +8556,11 @@ async def test_prepare_query_routing_parameters_request_1_grpc_asyncio(): assert args[0] == request_msg - expected_headers = {"name": "projects/sample1/instances/sample2"} + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1072 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } assert ( gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] ) @@ -8611,7 +8623,11 @@ async def test_execute_query_routing_parameters_request_1_grpc_asyncio(): assert args[0] == request_msg - expected_headers = {"name": "projects/sample1/instances/sample2"} + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1072 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } assert ( gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] ) @@ -10862,7 +10878,11 @@ def test_prepare_query_routing_parameters_request_1_rest(): assert args[0] == request_msg - expected_headers = {"name": "projects/sample1/instances/sample2"} + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1072 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } assert ( gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] ) @@ -10912,7 +10932,11 @@ def test_execute_query_routing_parameters_request_1_rest(): assert args[0] == request_msg - expected_headers = {"name": "projects/sample1/instances/sample2"} + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1072 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } assert ( gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] ) diff --git a/tests/unit/test_sql_routing_parameters.py b/tests/unit/test_sql_routing_parameters.py new file mode 100644 index 000000000..fa9316369 --- /dev/null +++ b/tests/unit/test_sql_routing_parameters.py @@ -0,0 +1,188 @@ +# -*- coding: utf-8 -*- +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock + from unittest.mock import AsyncMock # type: ignore # noqa: F401 +except ImportError: # pragma: NO COVER + import mock +import pytest + +from grpc.experimental import aio + +try: + from google.auth.aio import credentials as ga_credentials_async + + HAS_GOOGLE_AUTH_AIO = True +except ImportError: # pragma: NO COVER + HAS_GOOGLE_AUTH_AIO = False + +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers_async +from google.auth import credentials as ga_credentials +from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient +from google.cloud.bigtable_v2.services.bigtable.client import BigtableClient +from google.cloud.bigtable_v2.types import bigtable + +# This test file duplicates the gapic request header tests so that the temporary fix +# for SQL app_profile_id header handling can not be override by GAPIC. +# TODO: remove this once the fix is upstreamed + + +def async_anonymous_credentials(): + if HAS_GOOGLE_AUTH_AIO: + return ga_credentials_async.AnonymousCredentials() + return ga_credentials.AnonymousCredentials() + + +def test_prepare_query_routing_parameters_request_1_grpc(): + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object(type(client.transport.prepare_query), "__call__") as call: + call.return_value = bigtable.PrepareQueryResponse() + client.prepare_query( + request={"instance_name": "projects/sample1/instances/sample2"} + ) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, kw = call.mock_calls[0] + request_msg = bigtable.PrepareQueryRequest( + **{"instance_name": "projects/sample1/instances/sample2"} + ) + + assert args[0] == request_msg + + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1109 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } + assert ( + gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] + ) + + +@pytest.mark.asyncio +async def test_prepare_query_routing_parameters_request_1_grpc_asyncio(): + client = BigtableAsyncClient( + credentials=async_anonymous_credentials(), + transport="grpc_asyncio", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object(type(client.transport.prepare_query), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + bigtable.PrepareQueryResponse( + prepared_query=b"prepared_query_blob", + ) + ) + await client.prepare_query( + request={"instance_name": "projects/sample1/instances/sample2"} + ) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, kw = call.mock_calls[0] + request_msg = bigtable.PrepareQueryRequest( + **{"instance_name": "projects/sample1/instances/sample2"} + ) + + assert args[0] == request_msg + + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1109 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } + assert ( + gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] + ) + + +def test_execute_query_routing_parameters_request_1_grpc(): + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object(type(client.transport.execute_query), "__call__") as call: + call.return_value = iter([bigtable.ExecuteQueryResponse()]) + client.execute_query( + request={"instance_name": "projects/sample1/instances/sample2"} + ) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, kw = call.mock_calls[0] + request_msg = bigtable.ExecuteQueryRequest( + **{"instance_name": "projects/sample1/instances/sample2"} + ) + + assert args[0] == request_msg + + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1109 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } + assert ( + gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] + ) + + +@pytest.mark.asyncio +async def test_execute_query_routing_parameters_request_1_grpc_asyncio(): + client = BigtableAsyncClient( + credentials=async_anonymous_credentials(), + transport="grpc_asyncio", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object(type(client.transport.execute_query), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = mock.Mock(aio.UnaryStreamCall, autospec=True) + call.return_value.read = mock.AsyncMock( + side_effect=[bigtable.ExecuteQueryResponse()] + ) + await client.execute_query( + request={"instance_name": "projects/sample1/instances/sample2"} + ) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, kw = call.mock_calls[0] + request_msg = bigtable.ExecuteQueryRequest( + **{"instance_name": "projects/sample1/instances/sample2"} + ) + + assert args[0] == request_msg + + # expect app_profile_id while temporary patch is in place: https://github.com/googleapis/python-bigtable/pull/1109 + expected_headers = { + "name": "projects/sample1/instances/sample2", + "app_profile_id": "", + } + assert ( + gapic_v1.routing_header.to_grpc_metadata(expected_headers) in kw["metadata"] + )