Skip to content

Commit

Permalink
feat: add metrics (part 3)
Browse files Browse the repository at this point in the history
  • Loading branch information
arithmetic1728 committed May 25, 2023
1 parent ebd5af7 commit d0435ac
Show file tree
Hide file tree
Showing 13 changed files with 355 additions and 56 deletions.
30 changes: 23 additions & 7 deletions google/auth/compute_engine/_metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down Expand Up @@ -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)
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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"

Expand All @@ -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:
Expand Down Expand Up @@ -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"]
)
Expand Down
7 changes: 6 additions & 1 deletion google/auth/compute_engine/credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
10 changes: 8 additions & 2 deletions google/auth/impersonated_credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
33 changes: 27 additions & 6 deletions google/oauth2/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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")
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
18 changes: 15 additions & 3 deletions google/oauth2/reauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
)


Expand All @@ -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,
)


Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down
Loading

0 comments on commit d0435ac

Please sign in to comment.