| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -156,6 +156,7 @@ def get( | |||
| 156 | 156 | recursive=False, | |
| 157 | 157 | retry_count=5, | |
| 158 | 158 | headers=None, | |
| 159 | + return_none_for_not_found_error=False, | ||
| 159 | 160 | ): | |
| 160 | 161 | """Fetch a resource from the metadata server. | |
| 161 | 162 | ||
@@ -173,6 +174,8 @@ def get( | |||
| 173 | 174 | retry_count (int): How many times to attempt connecting to metadata | |
| 174 | 175 | server using above timeout. | |
| 175 | 176 | headers (Optional[Mapping[str, str]]): Headers for the request. | |
| 177 | + return_none_for_not_found_error (Optional[bool]): If True, returns None | ||
| 178 | + for 404 error instead of throwing an exception. | ||
| 176 | 179 | ||
| 177 | 180 | Returns: | |
| 178 | 181 | Union[Mapping, str]: If the metadata server returns JSON, a mapping of | |
@@ -216,8 +219,17 @@ def get( | |||
| 216 | 219 | "metadata service. Compute Engine Metadata server unavailable".format(url) | |
| 217 | 220 | ) | |
| 218 | 221 | ||
| 222 | + content = _helpers.from_bytes(response.data) | ||
| 223 | + | ||
| 224 | + if response.status == http_client.NOT_FOUND and return_none_for_not_found_error: | ||
| 225 | + _LOGGER.info( | ||
| 226 | + "Compute Engine Metadata server call to %s returned 404, reason: %s", | ||
| 227 | + path, | ||
| 228 | + content, | ||
| 229 | + ) | ||
| 230 | + return None | ||
| 231 | + | ||
| 219 | 232 | if response.status == http_client.OK: | |
| 220 | - content = _helpers.from_bytes(response.data) | ||
| 221 | 233 | if ( | |
| 222 | 234 | _helpers.parse_content_type(response.headers["content-type"]) | |
| 223 | 235 | == "application/json" | |
@@ -232,14 +244,14 @@ def get( | |||
| 232 | 244 | raise new_exc from caught_exc | |
| 233 | 245 | else: | |
| 234 | 246 | return content | |
| 235 | - else: | ||
| 236 | - raise exceptions.TransportError( | ||
| 237 | - "Failed to retrieve {} from the Google Compute Engine " | ||
| 238 | - "metadata service. Status: {} Response:\n{}".format( | ||
| 239 | - url, response.status, response.data | ||
| 240 | - ), | ||
| 241 | - response, | ||
| 242 | - ) | ||
| 247 | + | ||
| 248 | + raise exceptions.TransportError( | ||
| 249 | + "Failed to retrieve {} from the Google Compute Engine " | ||
| 250 | + "metadata service. Status: {} Response:\n{}".format( | ||
| 251 | + url, response.status, response.data | ||
| 252 | + ), | ||
| 253 | + response, | ||
| 254 | + ) | ||
| 243 | 255 | ||
| 244 | 256 | ||
| 245 | 257 | def get_project_id(request): | |
@@ -259,6 +271,29 @@ def get_project_id(request): | |||
| 259 | 271 | return get(request, "project/project-id") | |
| 260 | 272 | ||
| 261 | 273 | ||
| 274 | + def get_universe_domain(request): | ||
| 275 | + """Get the universe domain value from the metadata server. | ||
| 276 | + | ||
| 277 | + Args: | ||
| 278 | + request (google.auth.transport.Request): A callable used to make | ||
| 279 | + HTTP requests. | ||
| 280 | + | ||
| 281 | + Returns: | ||
| 282 | + str: The universe domain value. If the universe domain endpoint is not | ||
| 283 | + not found, return the default value, which is googleapis.com | ||
| 284 | + | ||
| 285 | + Raises: | ||
| 286 | + google.auth.exceptions.TransportError: if an error other than | ||
| 287 | + 404 occurs while retrieving metadata. | ||
| 288 | + """ | ||
| 289 | + universe_domain = get( | ||
| 290 | + request, "universe/universe_domain", return_none_for_not_found_error=True | ||
| 291 | + ) | ||
| 292 | + if not universe_domain: | ||
| 293 | + return "googleapis.com" | ||
| 294 | + return universe_domain | ||
| 295 | + | ||
| 296 | + | ||
| 262 | 297 | def get_service_account_info(request, service_account="default"): | |
| 263 | 298 | """Get information about a service account from the metadata server. | |
| 264 | 299 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -73,6 +73,7 @@ def __init__( | |||
| 73 | 73 | self._quota_project_id = quota_project_id | |
| 74 | 74 | self._scopes = scopes | |
| 75 | 75 | self._default_scopes = default_scopes | |
| 76 | + self._universe_domain_cached = False | ||
| 76 | 77 | ||
| 77 | 78 | def _retrieve_info(self, request): | |
| 78 | 79 | """Retrieve information about the service account. | |
@@ -131,6 +132,14 @@ def service_account_email(self): | |||
| 131 | 132 | def requires_scopes(self): | |
| 132 | 133 | return not self._scopes | |
| 133 | 134 | ||
| 135 | + @property | ||
| 136 | + def universe_domain(self): | ||
| 137 | + if self._universe_domain_cached: | ||
| 138 | + return self._universe_domain | ||
| 139 | + self._universe_domain = _metadata.get_universe_domain() | ||
| 140 | + self._universe_domain_cached = True | ||
| 141 | + return self._universe_domain | ||
| 142 | + | ||
| 134 | 143 | @_helpers.copy_docstring(credentials.CredentialsWithQuotaProject) | |
| 135 | 144 | def with_quota_project(self, quota_project_id): | |
| 136 | 145 | return self.__class__( | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -325,6 +325,18 @@ def test_get_failure(): | |||
| 325 | 325 | ) | |
| 326 | 326 | ||
| 327 | 327 | ||
| 328 | + def test_get_return_none_for_not_found_error(): | ||
| 329 | + request = make_request("Metadata error", status=http_client.NOT_FOUND) | ||
| 330 | + | ||
| 331 | + assert _metadata.get(request, PATH, return_none_for_not_found_error=True) is None | ||
| 332 | + | ||
| 333 | + request.assert_called_once_with( | ||
| 334 | + method="GET", | ||
| 335 | + url=_metadata._METADATA_ROOT + PATH, | ||
| 336 | + headers=_metadata._METADATA_HEADERS, | ||
| 337 | + ) | ||
| 338 | + | ||
| 339 | + | ||
| 328 | 340 | def test_get_failure_connection_failed(): | |
| 329 | 341 | request = make_request("") | |
| 330 | 342 | request.side_effect = exceptions.TransportError() | |
@@ -371,6 +383,53 @@ def test_get_project_id(): | |||
| 371 | 383 | assert project_id == project | |
| 372 | 384 | ||
| 373 | 385 | ||
| 386 | + def test_get_universe_domain_success(): | ||
| 387 | + request = make_request( | ||
| 388 | + "fake_universe_domain", headers={"content-type": "text/plain"} | ||
| 389 | + ) | ||
| 390 | + | ||
| 391 | + universe_domain = _metadata.get_universe_domain(request) | ||
| 392 | + | ||
| 393 | + request.assert_called_once_with( | ||
| 394 | + method="GET", | ||
| 395 | + url=_metadata._METADATA_ROOT + "universe/universe_domain", | ||
| 396 | + headers=_metadata._METADATA_HEADERS, | ||
| 397 | + ) | ||
| 398 | + assert universe_domain == "fake_universe_domain" | ||
| 399 | + | ||
| 400 | + | ||
| 401 | + def test_get_universe_domain_not_found(): | ||
| 402 | + # Test that if the universe domain endpoint returns 404 error, we should | ||
| 403 | + # use googleapis.com as the universe domain | ||
| 404 | + request = make_request("not found", status=http_client.NOT_FOUND) | ||
| 405 | + | ||
| 406 | + universe_domain = _metadata.get_universe_domain(request) | ||
| 407 | + | ||
| 408 | + request.assert_called_once_with( | ||
| 409 | + method="GET", | ||
| 410 | + url=_metadata._METADATA_ROOT + "universe/universe_domain", | ||
| 411 | + headers=_metadata._METADATA_HEADERS, | ||
| 412 | + ) | ||
| 413 | + assert universe_domain == "googleapis.com" | ||
| 414 | + | ||
| 415 | + | ||
| 416 | + def test_get_universe_domain_other_error(): | ||
| 417 | + # Test that if the universe domain endpoint returns an error other than 404 | ||
| 418 | + # we should throw the error | ||
| 419 | + request = make_request("unauthorized", status=http_client.UNAUTHORIZED) | ||
| 420 | + | ||
| 421 | + with pytest.raises(exceptions.TransportError) as excinfo: | ||
| 422 | + _metadata.get_universe_domain(request) | ||
| 423 | + | ||
| 424 | + assert excinfo.match(r"unauthorized") | ||
| 425 | + | ||
| 426 | + request.assert_called_once_with( | ||
| 427 | + method="GET", | ||
| 428 | + url=_metadata._METADATA_ROOT + "universe/universe_domain", | ||
| 429 | + headers=_metadata._METADATA_HEADERS, | ||
| 430 | + ) | ||
| 431 | + | ||
| 432 | + | ||
| 374 | 433 | @mock.patch( | |
| 375 | 434 | "google.auth.metrics.token_request_access_token_mds", | |
| 376 | 435 | return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -208,6 +208,26 @@ def test_token_usage_metrics(self): | |||
| 208 | 208 | assert headers["authorization"] == "Bearer token" | |
| 209 | 209 | assert headers["x-goog-api-client"] == "cred-type/mds" | |
| 210 | 210 | ||
| 211 | + @mock.patch( | ||
| 212 | + "google.auth.compute_engine._metadata.get_universe_domain", | ||
| 213 | + return_value="fake_universe_domain", | ||
| 214 | + ) | ||
| 215 | + def test_universe_domain(self, get_universe_domain): | ||
| 216 | + self.credentials._universe_domain_cached = False | ||
| 217 | + self.credentials._universe_domain = "googleapis.com" | ||
| 218 | + | ||
| 219 | + # calling the universe_domain property should trigger a call to | ||
| 220 | + # get_universe_domain to fetch the value. The value should be cached. | ||
| 221 | + assert self.credentials.universe_domain == "fake_universe_domain" | ||
| 222 | + assert self.credentials._universe_domain == "fake_universe_domain" | ||
| 223 | + assert self.credentials._universe_domain_cached | ||
| 224 | + get_universe_domain.assert_called_once() | ||
| 225 | + | ||
| 226 | + # calling the universe_domain property the second time should use the | ||
| 227 | + # cached value instead of calling get_universe_domain | ||
| 228 | + assert self.credentials.universe_domain == "fake_universe_domain" | ||
| 229 | + get_universe_domain.assert_called_once() | ||
| 230 | + | ||
| 211 | 231 | ||
| 212 | 232 | class TestIDTokenCredentials(object): | |
| 213 | 233 | credentials = None | |
| Back | FazBrowse Home | New Git URL |
0 commit comments