Use metadata server to detect GKE environment

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

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)

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

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)