feat: add metrics (part 2) (#1303)

Author: arithmetic1728

google/auth/compute_engine/credentials.py

@@ -28,6 +28,7 @@
 from google.auth import exceptions
 from google.auth import iam
 from google.auth import jwt
+from google.auth import metrics
 from google.auth.compute_engine import _metadata
 from google.oauth2 import _client
 
@@ -94,6 +95,9 @@ def _retrieve_info(self, request):
         if self._scopes is None:
             self._scopes = info["scopes"]
 
+    def _metric_header_for_usage(self):
+        return metrics.CRED_TYPE_SA_MDS
+
     def refresh(self, request):
         """Refresh the access token and scopes.
 

google/auth/impersonated_credentials.py

@@ -37,6 +37,7 @@
 from google.auth import credentials
 from google.auth import exceptions
 from google.auth import jwt
+from google.auth import metrics
 
 _DEFAULT_TOKEN_LIFETIME_SECS = 3600  # 1 hour in seconds
 
@@ -238,6 +239,9 @@ def __init__(
         self._quota_project_id = quota_project_id
         self._iam_endpoint_override = iam_endpoint_override
 
+    def _metric_header_for_usage(self):
+        return metrics.CRED_TYPE_SA_IMPERSONATE
+
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
         self._update_token(request)

google/oauth2/credentials.py

@@ -42,6 +42,7 @@
 from google.auth import _helpers
 from google.auth import credentials
 from google.auth import exceptions
+from google.auth import metrics
 from google.oauth2 import reauth
 
 _LOGGER = logging.getLogger(__name__)
@@ -287,6 +288,9 @@ def with_token_uri(self, token_uri):
             enable_reauth_refresh=self._enable_reauth_refresh,
         )
 
+    def _metric_header_for_usage(self):
+        return metrics.CRED_TYPE_USER
+
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
         scopes = self._scopes if self._scopes is not None else self._default_scopes

google/oauth2/service_account.py

@@ -78,6 +78,7 @@
 from google.auth import credentials
 from google.auth import exceptions
 from google.auth import jwt
+from google.auth import metrics
 from google.oauth2 import _client
 
 _DEFAULT_TOKEN_LIFETIME_SECS = 3600  # 1 hour in seconds
@@ -400,6 +401,16 @@ def _make_authorization_grant_assertion(self):
 
         return token
 
+    def _use_self_signed_jwt(self):
+        # Since domain wide delegation doesn't work with self signed JWT. If
+        # subject exists, then we should not use self signed JWT.
+        return self._subject is None and self._jwt_credentials is not None
+
+    def _metric_header_for_usage(self):
+        if self._use_self_signed_jwt():
+            return metrics.CRED_TYPE_SA_JWT
+        return metrics.CRED_TYPE_SA_ASSERTION
+
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
         if (
@@ -414,9 +425,7 @@ def refresh(self, request):
                 "domain wide delegation is not supported for non-default universe domain"
             )
 
-        # Since domain wide delegation doesn't work with self signed JWT. If
-        # subject exists, then we should not use self signed JWT.
-        if self._subject is None and self._jwt_credentials is not None:
+        if self._use_self_signed_jwt():
             self._jwt_credentials.refresh(request)
             self.token = self._jwt_credentials.token.decode()
             self.expiry = self._jwt_credentials.expiry

tests/compute_engine/test_credentials.py

@@ -185,6 +185,15 @@ def test_with_scopes(self):
 
         assert self.credentials._scopes == scopes
 
+    def test_token_usage_metrics(self):
+        self.credentials.token = "token"
+        self.credentials.expiry = None
+
+        headers = {}
+        self.credentials.before_request(mock.Mock(), None, None, headers)
+        assert headers["authorization"] == "Bearer token"
+        assert headers["x-goog-api-client"] == "cred-type/mds"
+
 
 class TestIDTokenCredentials(object):
     credentials = None

tests/oauth2/test_credentials.py

@@ -69,6 +69,16 @@ def test_default_state(self):
         assert credentials.rapt_token == self.RAPT_TOKEN
         assert credentials.refresh_handler is None
 
+    def test_token_usage_metrics(self):
+        credentials = self.make_credentials()
+        credentials.token = "token"
+        credentials.expiry = None
+
+        headers = {}
+        credentials.before_request(mock.Mock(), None, None, headers)
+        assert headers["authorization"] == "Bearer token"
+        assert headers["x-goog-api-client"] == "cred-type/u"
+
     def test_refresh_handler_setter_and_getter(self):
         scopes = ["email", "profile"]
         original_refresh_handler = mock.Mock(return_value=("ACCESS_TOKEN_1", None))

tests/oauth2/test_service_account.py

@@ -408,6 +408,37 @@ def test__create_self_signed_jwt_always_use_jwt_access(self, jwt):
         credentials._create_self_signed_jwt(None)
         jwt.from_signing_credentials.assert_not_called()
 
+    def test_token_usage_metrics_assertion(self):
+        credentials = service_account.Credentials(
+            SIGNER,
+            self.SERVICE_ACCOUNT_EMAIL,
+            self.TOKEN_URI,
+            always_use_jwt_access=False,
+        )
+        credentials.token = "token"
+        credentials.expiry = None
+
+        headers = {}
+        credentials.before_request(mock.Mock(), None, None, headers)
+        assert headers["authorization"] == "Bearer token"
+        assert headers["x-goog-api-client"] == "cred-type/sa"
+
+    def test_token_usage_metrics_self_signed_jwt(self):
+        credentials = service_account.Credentials(
+            SIGNER,
+            self.SERVICE_ACCOUNT_EMAIL,
+            self.TOKEN_URI,
+            always_use_jwt_access=True,
+        )
+        credentials._create_self_signed_jwt("foo.googleapis.com")
+        credentials.token = "token"
+        credentials.expiry = None
+
+        headers = {}
+        credentials.before_request(mock.Mock(), None, None, headers)
+        assert headers["authorization"] == "Bearer token"
+        assert headers["x-goog-api-client"] == "cred-type/jwt"
+
     @mock.patch("google.oauth2._client.jwt_grant", autospec=True)
     def test_refresh_success(self, jwt_grant):
         credentials = self.make_credentials()

tests/test_impersonated_credentials.py

@@ -162,6 +162,16 @@ def make_request(
 
         return request
 
+    def test_token_usage_metrics(self):
+        credentials = self.make_credentials()
+        credentials.token = "token"
+        credentials.expiry = None
+
+        headers = {}
+        credentials.before_request(mock.Mock(), None, None, headers)
+        assert headers["authorization"] == "Bearer token"
+        assert headers["x-goog-api-client"] == "cred-type/imp"
+
     @pytest.mark.parametrize("use_data_bytes", [True, False])
     def test_refresh_success(self, use_data_bytes, mock_donor_credentials):
         credentials = self.make_credentials(lifetime=None)