feat: add metrics (part 1)

google/auth/credentials.py

@@ -22,6 +22,7 @@
 
 from google.auth import _helpers, environment_vars
 from google.auth import exceptions
+from google.auth import metrics
 
 
 @six.add_metaclass(abc.ABCMeta)
@@ -100,6 +101,21 @@ def refresh(self, request):
         # (pylint doesn't recognize that this is abstract)
         raise NotImplementedError("Refresh must be implemented")
 
+    def _metric_header_for_usage(self):
+        """The x-goog-api-client header for token usage metric.
+
+        This header will be added to the API service requests in before_request
+        method. For example, "cred-type/sa-jwt" means service account self
+        signed jwt access token is used in the API service request
+        authorization header. Children credentials classes need to override
+        this method to provide the header value, if the token usage metric is
+        needed.
+
+        Returns:
+            str: The x-goog-api-client header value.
+        """
+        return None
+
     def apply(self, headers, token=None):
         """Apply the token to the authentication header.
 
@@ -133,6 +149,7 @@ def before_request(self, request, method, url, headers):
         # the http request.)
         if not self.valid:
             self.refresh(request)
+        metrics.add_metric_header(headers, self._metric_header_for_usage())
         self.apply(headers)
 
 

google/auth/metrics.py

@@ -0,0 +1,165 @@
+# Copyright 2023 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.
+
+""" We use x-goog-api-client header to report metrics. This module provides
+the constants and helper methods to construct x-goog-api-client header.
+"""
+
+import platform
+
+from google.auth import version
+
+
+API_CLIENT_HEADER = "x-goog-api-client"
+
+# Auth request type
+REQUEST_TYPE_ACCESS_TOKEN = "auth-request-type/at"
+REQUEST_TYPE_ID_TOKEN = "auth-request-type/it"
+REQUEST_TYPE_MDS_PING = "auth-request-type/mds"
+REQUEST_TYPE_REAUTH_START = "auth-request-type/re-start"
+REQUEST_TYPE_REAUTH_CONTINUE = "auth-request-type/re-cont"
+
+# Credential type
+CRED_TYPE_USER = "cred-type/u"
+CRED_TYPE_SA_ASSERTION = "cred-type/sa"
+CRED_TYPE_SA_JWT = "cred-type/jwt"
+CRED_TYPE_SA_MDS = "cred-type/mds"
+CRED_TYPE_SA_IMPERSONATE = "cred-type/imp"
+
+
+# Versions
+def auth_lib_version():
+    return "auth/" + version.__version__
+
+
+def python_version():
+    return "gl-python/" + platform.python_version()
+
+
+# Token request metric header values
+
+# x-goog-api-client header value for access token request via metadata server.
+# Example: "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/mds"
+def token_request_access_token_mds():
+    return "{} {} {} {}".format(
+        python_version(),
+        auth_lib_version(),
+        REQUEST_TYPE_ACCESS_TOKEN,
+        CRED_TYPE_SA_MDS,
+    )
+
+
+# x-goog-api-client header value for ID token request via metadata server.
+# Example: "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/mds"
+def token_request_id_token_mds():
+    return "{} {} {} {}".format(
+        python_version(), auth_lib_version(), REQUEST_TYPE_ID_TOKEN, CRED_TYPE_SA_MDS
+    )
+
+
+# x-goog-api-client header value for impersonated credentials access token request.
+# Example: "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/imp"
+def token_request_access_token_impersonate():
+    return "{} {} {} {}".format(
+        python_version(),
+        auth_lib_version(),
+        REQUEST_TYPE_ACCESS_TOKEN,
+        CRED_TYPE_SA_IMPERSONATE,
+    )
+
+
+# x-goog-api-client header value for impersonated credentials ID token request.
+# Example: "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/imp"
+def token_request_id_token_impersonate():
+    return "{} {} {} {}".format(
+        python_version(),
+        auth_lib_version(),
+        REQUEST_TYPE_ID_TOKEN,
+        CRED_TYPE_SA_IMPERSONATE,
+    )
+
+
+# x-goog-api-client header value for service account credentials access token
+# request (assertion flow).
+# Example: "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/sa"
+def token_request_access_token_sa_assertion():
+    return "{} {} {} {}".format(
+        python_version(),
+        auth_lib_version(),
+        REQUEST_TYPE_ACCESS_TOKEN,
+        CRED_TYPE_SA_ASSERTION,
+    )
+
+
+# x-goog-api-client header value for service account credentials ID token
+# request (assertion flow).
+# Example: "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/sa"
+def token_request_id_token_sa_assertion():
+    return "{} {} {} {}".format(
+        python_version(),
+        auth_lib_version(),
+        REQUEST_TYPE_ID_TOKEN,
+        CRED_TYPE_SA_ASSERTION,
+    )
+
+
+# x-goog-api-client header value for user credentials token request.
+# Example: "gl-python/3.7 auth/1.1 cred-type/u"
+def token_request_user():
+    return "{} {} {}".format(python_version(), auth_lib_version(), CRED_TYPE_USER)
+
+
+# Miscellenous metrics
+
+# x-goog-api-client header value for metadata server ping.
+# Example: "gl-python/3.7 auth/1.1 auth-request-type/mds"
+def mds_ping():
+    return "{} {} {}".format(
+        python_version(), auth_lib_version(), REQUEST_TYPE_MDS_PING
+    )
+
+
+# x-goog-api-client header value for reauth start endpoint calls.
+# Example: "gl-python/3.7 auth/1.1 auth-request-type/re-start"
+def reauth_start():
+    return "{} {} {}".format(
+        python_version(), auth_lib_version(), REQUEST_TYPE_REAUTH_START
+    )
+
+
+# x-goog-api-client header value for reauth continue endpoint calls.
+# Example: "gl-python/3.7 auth/1.1 cred-type/re-cont"
+def reauth_continue():
+    return "{} {} {}".format(
+        python_version(), auth_lib_version(), REQUEST_TYPE_REAUTH_CONTINUE
+    )
+
+
+def add_metric_header(headers, metric_header_value):
+    """Add x-goog-api-client header with the given value.
+
+    Args:
+        headers (Mapping[str, str]): The headers to which we will add the
+            metric header.
+        metric_header_value (Optional[str]): If value is None, do nothing;
+            if headers already has a x-goog-api-client header, append the value
+            to the existing header; otherwise add a new x-goog-api-client
+            header with the given value.
+    """
+    if not metric_header_value:
+        return
+    if API_CLIENT_HEADER not in headers:
+        headers[API_CLIENT_HEADER] = metric_header_value
+    else:
+        headers[API_CLIENT_HEADER] += " " + metric_header_value

tests/test_credentials.py

@@ -28,6 +28,14 @@ def with_quota_project(self, quota_project_id):
         raise NotImplementedError()
 
 
+class CredentialsImplWithMetrics(credentials.Credentials):
+    def refresh(self, request):
+        self.token = request
+
+    def _metric_header_for_usage(self):
+        return "foo"
+
+
 def test_credentials_constructor():
     credentials = CredentialsImpl()
     assert not credentials.token
@@ -83,6 +91,15 @@ def test_before_request():
     assert headers["authorization"] == "Bearer token"
 
 
+def test_before_request_metrics():
+    credentials = CredentialsImplWithMetrics()
+    request = "token"
+    headers = {}
+
+    credentials.before_request(request, "http://example.com", "GET", headers)
+    assert headers["x-goog-api-client"] == "foo"
+
+
 def test_anonymous_credentials_ctor():
     anon = credentials.AnonymousCredentials()
     assert anon.token is None

tests/test_metrics.py

@@ -0,0 +1,74 @@
+# Copyright 2014 Google Inc.
+#
+# 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.
+
+import platform
+
+import mock
+
+from google.auth import metrics
+from google.auth import version
+
+
+def test_add_metric_header():
+    headers = {}
+    metrics.add_metric_header(headers, None)
+    assert headers == {}
+
+    headers = {"x-goog-api-client": "foo"}
+    metrics.add_metric_header(headers, "bar")
+    assert headers == {"x-goog-api-client": "foo bar"}
+
+    headers = {}
+    metrics.add_metric_header(headers, "bar")
+    assert headers == {"x-goog-api-client": "bar"}
+
+
+def test_versions():
+    assert metrics.auth_lib_version() == "auth/" + version.__version__
+    assert metrics.python_version() == "gl-python/" + platform.python_version()
+
+
+@mock.patch("google.auth.metrics.auth_lib_version", return_value="auth/1.1")
+@mock.patch("google.auth.metrics.python_version", return_value="gl-python/3.7")
+def test_metric_values(mock_python_version, mock_auth_lib_version):
+    assert (
+        metrics.token_request_access_token_mds()
+        == "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/mds"
+    )
+    assert (
+        metrics.token_request_id_token_mds()
+        == "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/mds"
+    )
+    assert (
+        metrics.token_request_access_token_impersonate()
+        == "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/imp"
+    )
+    assert (
+        metrics.token_request_id_token_impersonate()
+        == "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/imp"
+    )
+    assert (
+        metrics.token_request_access_token_sa_assertion()
+        == "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/sa"
+    )
+    assert (
+        metrics.token_request_id_token_sa_assertion()
+        == "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/sa"
+    )
+    assert metrics.token_request_user() == "gl-python/3.7 auth/1.1 cred-type/u"
+    assert metrics.mds_ping() == "gl-python/3.7 auth/1.1 auth-request-type/mds"
+    assert metrics.reauth_start() == "gl-python/3.7 auth/1.1 auth-request-type/re-start"
+    assert (
+        metrics.reauth_continue() == "gl-python/3.7 auth/1.1 auth-request-type/re-cont"
+    )