feat: add metrics (part 3) (#1305)

This PR adds `x-goog-api-client` header to
- access token and id token refresh requests, for compute engine credentials / user credentials / service account credentials / impersonated credentials
- reauth start and continue requests
- metadata server ping requests

Previous PRs:
Part 1: #1298 
Part 2: #1303

Author: arithmetic1728

google/auth/compute_engine/_metadata.py

@@ -29,6 +29,7 @@
 from google.auth import _helpers
 from google.auth import environment_vars
 from google.auth import exceptions
+from google.auth import metrics
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -121,13 +122,13 @@ def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3):
     #       the metadata resolution was particularly slow. The latter case is
     #       "unlikely".
     retries = 0
+    headers = _METADATA_HEADERS.copy()
+    headers[metrics.API_CLIENT_HEADER] = metrics.mds_ping()
+
     while retries < retry_count:
         try:
             response = request(
-                url=_METADATA_IP_ROOT,
-                method="GET",
-                headers=_METADATA_HEADERS,
-                timeout=timeout,
+                url=_METADATA_IP_ROOT, method="GET", headers=headers, timeout=timeout
             )
 
             metadata_flavor = response.headers.get(_METADATA_FLAVOR_HEADER)
@@ -150,7 +151,13 @@ def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3):
 
 
 def get(
-    request, path, root=_METADATA_ROOT, params=None, recursive=False, retry_count=5
+    request,
+    path,
+    root=_METADATA_ROOT,
+    params=None,
+    recursive=False,
+    retry_count=5,
+    headers=None,
 ):
     """Fetch a resource from the metadata server.
 
@@ -167,6 +174,7 @@ def get(
             details.
         retry_count (int): How many times to attempt connecting to metadata
             server using above timeout.
+        headers (Optional[Mapping[str, str]]): Headers for the request.
 
     Returns:
         Union[Mapping, str]: If the metadata server returns JSON, a mapping of
@@ -180,6 +188,10 @@ def get(
     base_url = urlparse.urljoin(root, path)
     query_params = {} if params is None else params
 
+    headers_to_use = _METADATA_HEADERS.copy()
+    if headers:
+        headers_to_use.update(headers)
+
     if recursive:
         query_params["recursive"] = "true"
 
@@ -188,7 +200,7 @@ def get(
     retries = 0
     while retries < retry_count:
         try:
-            response = request(url=url, method="GET", headers=_METADATA_HEADERS)
+            response = request(url=url, method="GET", headers=headers_to_use)
             break
 
         except exceptions.TransportError as e:
@@ -300,8 +312,12 @@ def get_service_account_token(request, service_account="default", scopes=None):
     else:
         params = None
 
+    metrics_header = {
+        metrics.API_CLIENT_HEADER: metrics.token_request_access_token_mds()
+    }
+
     path = "instance/service-accounts/{0}/token".format(service_account)
-    token_json = get(request, path, params=params)
+    token_json = get(request, path, params=params, headers=metrics_header)
     token_expiry = _helpers.utcnow() + datetime.timedelta(
         seconds=token_json["expires_in"]
     )

google/auth/compute_engine/credentials.py

@@ -378,7 +378,12 @@ def _call_metadata_identity_endpoint(self, request):
         try:
             path = "instance/service-accounts/default/identity"
             params = {"audience": self._target_audience, "format": "full"}
-            id_token = _metadata.get(request, path, params=params)
+            metrics_header = {
+                metrics.API_CLIENT_HEADER: metrics.token_request_id_token_mds()
+            }
+            id_token = _metadata.get(
+                request, path, params=params, headers=metrics_header
+            )
         except exceptions.TransportError as caught_exc:
             new_exc = exceptions.RefreshError(caught_exc)
             six.raise_from(new_exc, caught_exc)

google/auth/impersonated_credentials.py

@@ -265,7 +265,10 @@ def _update_token(self, request):
             "lifetime": str(self._lifetime) + "s",
         }
 
-        headers = {"Content-Type": "application/json"}
+        headers = {
+            "Content-Type": "application/json",
+            metrics.API_CLIENT_HEADER: metrics.token_request_access_token_impersonate(),
+        }
 
         # Apply the source credentials authentication info.
         self._source_credentials.apply(headers)
@@ -426,7 +429,10 @@ def refresh(self, request):
             "includeEmail": self._include_email,
         }
 
-        headers = {"Content-Type": "application/json"}
+        headers = {
+            "Content-Type": "application/json",
+            metrics.API_CLIENT_HEADER: metrics.token_request_id_token_impersonate(),
+        }
 
         authed_session = AuthorizedSession(
             self._target_credentials._source_credentials, auth_request=request

google/oauth2/_client.py

@@ -34,6 +34,7 @@
 from google.auth import _helpers
 from google.auth import exceptions
 from google.auth import jwt
+from google.auth import metrics
 from google.auth import transport
 
 _URLENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
@@ -146,6 +147,7 @@ def _token_endpoint_request_no_throw(
     access_token=None,
     use_json=False,
     can_retry=True,
+    headers=None,
     **kwargs
 ):
     """Makes a request to the OAuth 2.0 authorization server's token endpoint.
@@ -161,6 +163,7 @@ def _token_endpoint_request_no_throw(
         use_json (Optional(bool)): Use urlencoded format or json format for the
             content type. The default value is False.
         can_retry (bool): Enable or disable request retry behavior.
+        headers (Optional[Mapping[str, str]]): The headers for the request.
         kwargs: Additional arguments passed on to the request method. The
             kwargs will be passed to `requests.request` method, see:
             https://docs.python-requests.org/en/latest/api/#requests.request.
@@ -176,18 +179,21 @@ def _token_endpoint_request_no_throw(
           is retryable.
     """
     if use_json:
-        headers = {"Content-Type": _JSON_CONTENT_TYPE}
+        headers_to_use = {"Content-Type": _JSON_CONTENT_TYPE}
         body = json.dumps(body).encode("utf-8")
     else:
-        headers = {"Content-Type": _URLENCODED_CONTENT_TYPE}
+        headers_to_use = {"Content-Type": _URLENCODED_CONTENT_TYPE}
         body = urllib.parse.urlencode(body).encode("utf-8")
 
     if access_token:
-        headers["Authorization"] = "Bearer {}".format(access_token)
+        headers_to_use["Authorization"] = "Bearer {}".format(access_token)
+
+    if headers:
+        headers_to_use.update(headers)
 
     def _perform_request():
         response = request(
-            method="POST", url=token_uri, headers=headers, body=body, **kwargs
+            method="POST", url=token_uri, headers=headers_to_use, body=body, **kwargs
         )
         response_body = (
             response.data.decode("utf-8")
@@ -231,6 +237,7 @@ def _token_endpoint_request(
     access_token=None,
     use_json=False,
     can_retry=True,
+    headers=None,
     **kwargs
 ):
     """Makes a request to the OAuth 2.0 authorization server's token endpoint.
@@ -245,6 +252,7 @@ def _token_endpoint_request(
         use_json (Optional(bool)): Use urlencoded format or json format for the
             content type. The default value is False.
         can_retry (bool): Enable or disable request retry behavior.
+        headers (Optional[Mapping[str, str]]): The headers for the request.
         kwargs: Additional arguments passed on to the request method. The
             kwargs will be passed to `requests.request` method, see:
             https://docs.python-requests.org/en/latest/api/#requests.request.
@@ -268,6 +276,7 @@ def _token_endpoint_request(
         access_token=access_token,
         use_json=use_json,
         can_retry=can_retry,
+        headers=headers,
         **kwargs
     )
     if not response_status_ok:
@@ -301,7 +310,13 @@ def jwt_grant(request, token_uri, assertion, can_retry=True):
     body = {"assertion": assertion, "grant_type": _JWT_GRANT_TYPE}
 
     response_data = _token_endpoint_request(
-        request, token_uri, body, can_retry=can_retry
+        request,
+        token_uri,
+        body,
+        can_retry=can_retry,
+        headers={
+            metrics.API_CLIENT_HEADER: metrics.token_request_access_token_sa_assertion()
+        },
     )
 
     try:
@@ -384,7 +399,13 @@ def id_token_jwt_grant(request, token_uri, assertion, can_retry=True):
     body = {"assertion": assertion, "grant_type": _JWT_GRANT_TYPE}
 
     response_data = _token_endpoint_request(
-        request, token_uri, body, can_retry=can_retry
+        request,
+        token_uri,
+        body,
+        can_retry=can_retry,
+        headers={
+            metrics.API_CLIENT_HEADER: metrics.token_request_id_token_sa_assertion()
+        },
     )
 
     try:

google/oauth2/reauth.py

@@ -37,6 +37,7 @@
 from six.moves import range
 
 from google.auth import exceptions
+from google.auth import metrics
 from google.oauth2 import _client
 from google.oauth2 import challenges
 
@@ -94,9 +95,15 @@ def _get_challenges(
     body = {"supportedChallengeTypes": supported_challenge_types}
     if requested_scopes:
         body["oauthScopesForDomainPolicyLookup"] = requested_scopes
+    metrics_header = {metrics.API_CLIENT_HEADER: metrics.reauth_start()}
 
     return _client._token_endpoint_request(
-        request, _REAUTH_API + ":start", body, access_token=access_token, use_json=True
+        request,
+        _REAUTH_API + ":start",
+        body,
+        access_token=access_token,
+        use_json=True,
+        headers=metrics_header,
     )
 
 
@@ -123,13 +130,15 @@ def _send_challenge_result(
         "action": "RESPOND",
         "proposalResponse": client_input,
     }
+    metrics_header = {metrics.API_CLIENT_HEADER: metrics.reauth_continue()}
 
     return _client._token_endpoint_request(
         request,
         _REAUTH_API + "/{}:continue".format(session_id),
         body,
         access_token=access_token,
         use_json=True,
+        headers=metrics_header,
     )
 
 
@@ -320,9 +329,10 @@ def refresh_grant(
         body["scope"] = " ".join(scopes)
     if rapt_token:
         body["rapt"] = rapt_token
+    metrics_header = {metrics.API_CLIENT_HEADER: metrics.token_request_user()}
 
     response_status_ok, response_data, retryable_error = _client._token_endpoint_request_no_throw(
-        request, token_uri, body
+        request, token_uri, body, headers=metrics_header
     )
     if (
         not response_status_ok
@@ -345,7 +355,9 @@ def refresh_grant(
             response_status_ok,
             response_data,
             retryable_error,
-        ) = _client._token_endpoint_request_no_throw(request, token_uri, body)
+        ) = _client._token_endpoint_request_no_throw(
+            request, token_uri, body, headers=metrics_header
+        )
 
     if not response_status_ok:
         _client._handle_error_response(response_data, retryable_error)