Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 72 additions & 11 deletions packages/google-api-core/google/api_core/gapic_v1/method.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

import enum
import functools
from typing import List, Tuple

from google.api_core import grpc_helpers
from google.api_core.gapic_v1 import client_info
Expand Down Expand Up @@ -57,6 +58,52 @@ def _apply_decorators(func, decorators):
return func


def _deduplicate_metadata_tokens(*headers: str) -> str:
"""
Given one or more metadata payload strings, create a combined
string with deduplicated tokens, while preserving token order.

Inputs are expected contain a set of metadata tokens separated by spaces
Example: `gl-python/3.14.0 grpc/1.76.0 gax/2.29.0 gapic/3.8.0 pb/6.33.4`

Args:
*headers: one or more metadata payload strings

Returns:
a single combined payload string
"""
# Split all non-empty headers into individual tokens
token_list = " ".join(filter(None, headers)).split()
# Deduplicate while preserving order
return " ".join(dict.fromkeys(token_list))


def _extract_metrics_header(metadata) -> Tuple[str, List[Tuple[str, str]]]:
"""Extract x-google-api-client header from metadata list.

Args:
metadata (Sequence[Tuple[str, str]]): The metadata to extract from.

Returns:
A tuple containing:
- A sequence of remaining metadata tuples.
- a string representing the header value.
"""
if not metadata:
return "", []

key_to_find = client_info.METRICS_METADATA_KEY

metric_str = _deduplicate_metadata_tokens(
" ".join([v for k, v in metadata if k == key_to_find])
)
if not metric_str:
return "", list(metadata)

arbitrary_metadata = [item for item in metadata if item[0] != key_to_find]
return metric_str, arbitrary_metadata


class _GapicCallable(object):
"""Callable that applies retry, timeout, and metadata logic.

Expand Down Expand Up @@ -90,7 +137,16 @@ def __init__(
self._retry = retry
self._timeout = timeout
self._compression = compression
self._metadata = metadata
# Pre-extract the x-goog-api-client header from the initialized metadata.
self._x_goog_api_client, remaining = _extract_metrics_header(metadata)
self._static_metadata = tuple(remaining)
if self._x_goog_api_client:
self._default_metadata = (
(client_info.METRICS_METADATA_KEY, self._x_goog_api_client),
*self._static_metadata,
)
else:
self._default_metadata = self._static_metadata

def __call__(
self, *args, timeout=DEFAULT, retry=DEFAULT, compression=DEFAULT, **kwargs
Expand All @@ -112,16 +168,21 @@ def __call__(
# Apply all applicable decorators.
wrapped_func = _apply_decorators(self._target, [retry, timeout])

# Add the user agent metadata to the call.
if self._metadata is not None:
metadata = kwargs.get("metadata", [])
# Due to the nature of invocation, None should be treated the same
# as not specified.
if metadata is None:
metadata = []
metadata = list(metadata)
metadata.extend(self._metadata)
kwargs["metadata"] = metadata
if user_metadata := kwargs.get("metadata"):
# Add the user agent metadata to the call.
final_metadata = list(self._static_metadata)
user_x_goog, remaining = _extract_metrics_header(user_metadata)

merged_header = _deduplicate_metadata_tokens(
self._x_goog_api_client, user_x_goog
)
if merged_header:
final_metadata.append((client_info.METRICS_METADATA_KEY, merged_header))
final_metadata.extend(remaining)
kwargs["metadata"] = final_metadata
elif self._default_metadata:
kwargs["metadata"] = self._default_metadata

if self._compression is not None:
kwargs["compression"] = compression

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,8 @@
if they are not already set.
"""

from typing import Union
import uuid
from typing import Union

import google.protobuf.message

Expand Down
19 changes: 14 additions & 5 deletions packages/google-api-core/google/api_core/grpc_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,11 +254,20 @@ def _create_composite_credentials(
request = google.auth.transport.requests.Request()

# Create the metadata plugin for inserting the authorization header.
metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin(
credentials,
request,
default_host=default_host,
)
try:
metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin(
credentials,
request,
default_host=default_host,
suppress_metrics_header=True,
)
except TypeError:
# Support older versions of google-auth that do not accept suppress_metrics_header
metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin(
credentials,
request,
default_host=default_host,
)

# Create a set of grpc.CallCredentials using the metadata plugin.
google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin)
Expand Down
84 changes: 84 additions & 0 deletions packages/google-api-core/tests/unit/gapic/test_method.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,69 @@ def test_invoke_wrapped_method_with_metadata_as_none():
assert len(metadata) == 1


def test_extract_metrics_header_duplicate_tokens():
metadata = [
("x-goog-api-client", "token1 token2"),
("x-goog-api-client", "token2 token3 token1"),
("other-header", "value"),
("x-goog-api-client", "token4 token2"),
]

metric_str, arbitrary_metadata = (
google.api_core.gapic_v1.method._extract_metrics_header(metadata)
)

# Should maintain order of first appearance and eliminate duplicates
assert metric_str == "token1 token2 token3 token4"
assert arbitrary_metadata == [("other-header", "value")]


def test_invoke_wrapped_method_with_duplicate_x_goog_api_client_metadata():
method = mock.Mock(spec=["__call__"])

# Create a custom ClientInfo with defined properties so we know exactly what is returned
client_info = google.api_core.gapic_v1.client_info.ClientInfo(
user_agent="custom-user-agent/1.0",
python_version="3.14.0",
grpc_version="1.76.0",
api_core_version="2.29.0",
)

wrapped_method = google.api_core.gapic_v1.method.wrap_method(
method, client_info=client_info
)

# Invoke the wrapped method with an explicit user-provided custom header that contains duplicates
# both within its own items and overlapping with the default client_info
wrapped_method(
mock.sentinel.request,
metadata=[
("x-goog-api-client", "override-client/2.0"),
(
"x-goog-api-client",
"override-client/2.0 grpc/1.76.0 custom-user-agent/1.0",
),
("other-header", "value"),
],
)

method.assert_called_once_with(mock.sentinel.request, metadata=mock.ANY)
metadata = method.call_args[1]["metadata"]

# There should only be one "x-goog-api-client" header, containing both values joined by space,
# plus the other-header.
assert len(metadata) == 2
metadata_dict = dict(metadata)
assert "other-header" in metadata_dict
assert metadata_dict["other-header"] == "value"
assert "x-goog-api-client" in metadata_dict
# Verify both the user-provided override value and the library system telemetry are merged explicitly
assert (
metadata_dict["x-goog-api-client"]
== "custom-user-agent/1.0 gl-python/3.14.0 grpc/1.76.0 gax/2.29.0 override-client/2.0"
)


@mock.patch("time.sleep")
def test_wrap_method_with_default_retry_and_timeout_and_compression(unused_sleep):
method = mock.Mock(
Expand Down Expand Up @@ -248,3 +311,24 @@ def test_wrap_method_with_call_not_supported():
with pytest.raises(ValueError) as exc_info:
google.api_core.gapic_v1.method.wrap_method(method, with_call=True)
assert "with_call=True is only supported for unary calls" in str(exc_info.value)


@pytest.mark.parametrize(
"headers,expected",
[
((), ""),
(("",), ""),
((None,), ""),
(("", None, ""), ""),
(("token1",), "token1"),
(("token1 token1",), "token1"),
(("token1", "token1"), "token1"),
(("token1 token2 token1",), "token1 token2"),
(("token1", "token2", "token1"), "token1 token2"),
(("token1 token2", "token2 token3"), "token1 token2 token3"),
(("token1", None, "token2", "", "token1"), "token1 token2"),
],
)
def test__deduplicate_metadata_tokens(headers, expected):
dedup = google.api_core.gapic_v1.method._deduplicate_metadata_tokens
assert dedup(*headers) == expected
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@

from google.api_core.gapic_v1.requests import setup_request_id


# --- Mock Request Helper Classes ---


Expand Down
10 changes: 9 additions & 1 deletion packages/google-auth/google/auth/transport/grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,16 +48,21 @@ class AuthMetadataPlugin(grpc.AuthMetadataPlugin):
default_host (Optional[str]): A host like "pubsub.googleapis.com".
This is used when a self-signed JWT is created from service
account credentials.
suppress_metrics_header (bool): When enabled, ``x-goog-api-client``
will be stripped from authorization headers.
"""

def __init__(self, credentials, request, default_host=None):
def __init__(
self, credentials, request, default_host=None, *, suppress_metrics_header=False
):
# pylint: disable=no-value-for-parameter
# pylint doesn't realize that the super method takes no arguments
# because this class is the same name as the superclass.
super(AuthMetadataPlugin, self).__init__()
self._credentials = credentials
self._request = request
self._default_host = default_host
self._suppress_metrics_header = suppress_metrics_header

def _get_authorization_headers(self, context):
"""Gets the authorization headers for a request.
Expand All @@ -81,6 +86,9 @@ def _get_authorization_headers(self, context):
self._request, context.method_name, context.service_url, headers
)

if self._suppress_metrics_header and "x-goog-api-client" in headers:
del headers["x-goog-api-client"]

return list(headers.items())

def __call__(self, context, callback):
Expand Down
29 changes: 29 additions & 0 deletions packages/google-auth/tests/transport/test_grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,35 @@ def test__get_authorization_headers_with_service_account_and_default_host(self):
"https://{}/".format(default_host)
)

def test_suppress_metrics_header(self):
credentials = mock.create_autospec(service_account.Credentials)

# Mock credentials before_request that adds metric and authorization
def mock_before_request(request, method, url, headers):
headers["x-goog-api-client"] = "foo"
headers["authorization"] = "Bearer token"

credentials.before_request.side_effect = mock_before_request
request = mock.create_autospec(transport.Request)

# By default, suppress_metrics_header=False
plugin = google.auth.transport.grpc.AuthMetadataPlugin(credentials, request)
context = mock.create_autospec(grpc.AuthMetadataContext, instance=True)
context.method_name = "methodName"
context.service_url = "https://pubsub.googleapis.com/methodName"

headers = dict(plugin._get_authorization_headers(context))
assert "x-goog-api-client" in headers
assert headers["x-goog-api-client"] == "foo"

# With suppress_metrics_header=True
plugin_suppressed = google.auth.transport.grpc.AuthMetadataPlugin(
credentials, request, suppress_metrics_header=True
)
headers_suppressed = dict(plugin_suppressed._get_authorization_headers(context))
assert "x-goog-api-client" not in headers_suppressed
assert headers_suppressed["authorization"] == "Bearer token"


@mock.patch(
"google.auth.transport._mtls_helper.get_client_ssl_credentials", autospec=True
Expand Down
Loading