diff --git a/logging/google/cloud/logging/_helpers.py b/logging/google/cloud/logging/_helpers.py index 8e17a9538e76..c7fab41bc4e8 100644 --- a/logging/google/cloud/logging/_helpers.py +++ b/logging/google/cloud/logging/_helpers.py @@ -14,11 +14,17 @@ """Common logging helpers.""" +import requests from google.cloud.logging.entries import ProtobufEntry from google.cloud.logging.entries import StructEntry from google.cloud.logging.entries import TextEntry +METADATA_URL = 'http://metadata/computeMetadata/v1/' +METADATA_HEADERS = { + 'Metadata-Flavor': 'Google' +} + def entry_from_resource(resource, client, loggers): """Detect correct entry type from resource and instantiate. @@ -46,3 +52,32 @@ def entry_from_resource(resource, client, loggers): return ProtobufEntry.from_api_repr(resource, client, loggers) raise ValueError('Cannot parse log entry resource.') + + +def retrieve_metadata_server(metadata_key): + """Retrieve the metadata key in the metadata server. + + See: https://cloud.google.com/compute/docs/storing-retrieving-metadata + + :type metadata_key: str + :param metadata_key: Key of the metadata which will form the url. You can + also supply query parameters after the metadata key. + e.g. "tags?alt=json" + + :rtype: str + :returns: The value of the metadata key returned by the metadata server. + """ + url = METADATA_URL + metadata_key + + try: + response = requests.get(url, headers=METADATA_HEADERS) + + if response.status_code == requests.codes.ok: + return response.text + + except requests.exceptions.RequestException: + # Ignore the exception, connection failed means the attribute does not + # exist in the metadata server. + pass + + return None diff --git a/logging/google/cloud/logging/client.py b/logging/google/cloud/logging/client.py index 23ec84ec67d0..ae20dd48fcf6 100644 --- a/logging/google/cloud/logging/client.py +++ b/logging/google/cloud/logging/client.py @@ -31,6 +31,7 @@ from google.cloud.client import ClientWithProject from google.cloud.environment_vars import DISABLE_GRPC +from google.cloud.logging._helpers import retrieve_metadata_server from google.cloud.logging._http import Connection from google.cloud.logging._http import _LoggingAPI as JSONLoggingAPI from google.cloud.logging._http import _MetricsAPI as JSONMetricsAPI @@ -55,8 +56,8 @@ _APPENGINE_FLEXIBLE_ENV_FLEX = 'GAE_INSTANCE' """Environment variable set in App Engine when env:flex is set.""" -_CONTAINER_ENGINE_ENV = 'KUBERNETES_SERVICE' -"""Environment variable set in a Google Container Engine environment.""" +_GKE_CLUSTER_NAME = 'instance/attributes/cluster-name' +"""Attribute in metadata server when in GKE environment.""" class Client(ClientWithProject): @@ -301,10 +302,12 @@ def get_default_handler(self): :rtype: :class:`logging.Handler` :returns: The default log handler based on the environment """ + gke_cluster_name = retrieve_metadata_server(_GKE_CLUSTER_NAME) + if (_APPENGINE_FLEXIBLE_ENV_VM in os.environ or _APPENGINE_FLEXIBLE_ENV_FLEX in os.environ): return AppEngineHandler(self) - elif _CONTAINER_ENGINE_ENV in os.environ: + elif gke_cluster_name is not None: return ContainerEngineHandler() else: return CloudLoggingHandler(self) diff --git a/logging/tests/unit/test__helpers.py b/logging/tests/unit/test__helpers.py index 7cc2d392514c..93532eed0c05 100644 --- a/logging/tests/unit/test__helpers.py +++ b/logging/tests/unit/test__helpers.py @@ -15,6 +15,8 @@ import unittest +import mock + class Test_entry_from_resource(unittest.TestCase): @@ -53,6 +55,69 @@ def test_proto_payload(self): self._payload_helper('protoPayload', 'ProtobufEntry') +class Test_retrieve_metadata_server(unittest.TestCase): + + @staticmethod + def _call_fut(metadata_key): + from google.cloud.logging._helpers import retrieve_metadata_server + + return retrieve_metadata_server(metadata_key) + + def test_metadata_exists(self): + status_code_ok = 200 + response_text = 'my-gke-cluster' + metadata_key = 'test_key' + + response_mock = ResponseMock(status_code=status_code_ok) + response_mock.text = response_text + + requests_mock = mock.Mock() + requests_mock.get.return_value = response_mock + requests_mock.codes.ok = status_code_ok + + patch = mock.patch( + 'google.cloud.logging._helpers.requests', + requests_mock) + + with patch: + metadata = self._call_fut(metadata_key) + + self.assertEqual(metadata, response_text) + + def test_metadata_does_not_exist(self): + status_code_ok = 200 + status_code_not_found = 404 + metadata_key = 'test_key' + + response_mock = ResponseMock(status_code=status_code_not_found) + + requests_mock = mock.Mock() + requests_mock.get.return_value = response_mock + requests_mock.codes.ok = status_code_ok + + patch = mock.patch( + 'google.cloud.logging._helpers.requests', + requests_mock) + + with patch: + metadata = self._call_fut(metadata_key) + + self.assertIsNone(metadata) + + def test_request_exception(self): + metadata_key = 'test_url_cannot_connect' + metadata_url = 'http://metadata.invalid/' + + patch = mock.patch( + 'google.cloud.logging._helpers.METADATA_URL', + new=metadata_url) + + with patch: + metadata = self._call_fut(metadata_key) + + self.assertIsNone(metadata) + + class EntryMock(object): def __init__(self): @@ -62,3 +127,10 @@ def __init__(self): def from_api_repr(self, resource, client, loggers): self.called = (resource, client, loggers) return self.sentinel + + +class ResponseMock(object): + + def __init__(self, status_code, text='test_response_text'): + self.status_code = status_code + self.text = text diff --git a/logging/tests/unit/test_client.py b/logging/tests/unit/test_client.py index 37bfc5c18214..bb16e85c7ae5 100644 --- a/logging/tests/unit/test_client.py +++ b/logging/tests/unit/test_client.py @@ -581,16 +581,18 @@ def test_get_default_handler_app_engine(self): self.assertIsInstance(handler, AppEngineHandler) def test_get_default_handler_container_engine(self): - import os - from google.cloud._testing import _Monkey - from google.cloud.logging.client import _CONTAINER_ENGINE_ENV from google.cloud.logging.handlers import ContainerEngineHandler - client = self._make_one(project=self.PROJECT, - credentials=_make_credentials(), - _use_grpc=False) + client = self._make_one( + project=self.PROJECT, + credentials=_make_credentials(), + _use_grpc=False) + + patch = mock.patch( + 'google.cloud.logging.client.retrieve_metadata_server', + return_value='test-gke-cluster') - with _Monkey(os, environ={_CONTAINER_ENGINE_ENV: 'True'}): + with patch: handler = client.get_default_handler() self.assertIsInstance(handler, ContainerEngineHandler)