Adding implicit dataset ID support for Compute Engine.

Fixes #475.

Author: dhermes

gcloud/datastore/__init__.py

@@ -72,11 +72,11 @@
 def set_default_dataset_id(dataset_id=None):
     """Set default dataset ID either explicitly or implicitly as fall-back.
 
-    In implicit case, currently only supports enviroment variable but will
-    support App Engine, Compute Engine and other environments in the future.
-
-    Local environment variable used is:
-    - GCLOUD_DATASET_ID
+    In implicit case, supports three cases. In order of precedence, the
+    implicit cases are:
+    - GCLOUD_DATASET_ID environment variable
+    - Google App Engine application ID
+    - Google Compute Engine project ID (from metadata server)
 
     :type dataset_id: string
     :param dataset_id: Optional. The dataset ID to use as default.
@@ -87,6 +87,9 @@ def set_default_dataset_id(dataset_id=None):
     if dataset_id is None:
         dataset_id = _implicit_environ.app_engine_id()
 
+    if dataset_id is None:
+        dataset_id = _implicit_environ.compute_engine_id()
+
     if dataset_id is not None:
         _implicit_environ.DATASET_ID = dataset_id
 

gcloud/datastore/_implicit_environ.py

@@ -18,6 +18,9 @@
 imply the current dataset ID and connection from the enviroment.
 """
 
+import httplib2
+import socket
+
 try:
     from google.appengine.api import app_identity
 except ImportError:
@@ -42,3 +45,36 @@ def app_engine_id():
         return None
 
     return app_identity.get_application_id()
+
+
+def compute_engine_id():
+    """Gets the Compute Engine project ID if it can be inferred.
+
+    Uses 169.254.169.254 for the metadata server to avoid request
+    latency from DNS lookup.
+
+    See https://cloud.google.com/compute/docs/metadata#metadataserver
+    for information about this IP address. (This IP is also used for
+    Amazon EC2 instances, so the metadata flavor is crucial.)
+
+    See https://github.com/google/oauth2client/issues/93 for context about
+    DNS latency.
+
+    :rtype: string or ``NoneType``
+    :returns: Compute Engine project ID if the metadata service is available,
+              else ``None``.
+    """
+    http = httplib2.Http(timeout=0.1)
+    uri = 'http://169.254.169.254/computeMetadata/v1/project/project-id'
+    headers = {'Metadata-Flavor': 'Google'}
+
+    response = content = None
+    try:
+        response, content = http.request(uri, method='GET', headers=headers)
+    except socket.timeout:
+        pass
+
+    if response is None or response['status'] != '200':
+        return None
+
+    return content

gcloud/datastore/test___init__.py

@@ -96,14 +96,85 @@ def test_set_implicit_both_env_and_appengine(self):
         from gcloud.datastore import _implicit_environ
 
         IMPLICIT_DATASET_ID = 'IMPLICIT'
+        APP_IDENTITY = _AppIdentity('GAE')
+
+        with self._monkey(IMPLICIT_DATASET_ID):
+            with _Monkey(_implicit_environ, app_identity=APP_IDENTITY):
+                self._callFUT()
+
+        self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID)
+
+    def _implicit_compute_engine_helper(self, status):
+        from gcloud._testing import _Monkey
+        from gcloud.datastore import _implicit_environ
+
+        COMPUTE_ENGINE_ID = 'GCE'
+        HTTPLIB2 = _Httplib2(COMPUTE_ENGINE_ID, status=status)
+        if status == '200':
+            EXPECTED_ID = COMPUTE_ENGINE_ID
+        else:
+            EXPECTED_ID = None
+
+        with self._monkey(None):
+            with _Monkey(_implicit_environ, httplib2=HTTPLIB2):
+                self._callFUT()
+
+        self.assertEqual(_implicit_environ.DATASET_ID, EXPECTED_ID)
+        self.assertEqual(len(HTTPLIB2._http_instances), 1)
+
+        (timeout, http_instance), = HTTPLIB2._http_instances
+        self.assertEqual(timeout, 0.1)
+        self.assertEqual(
+            http_instance._called_uris,
+            ['http://169.254.169.254/computeMetadata/v1/project/project-id'])
+        expected_kwargs = {
+            'method': 'GET',
+            'headers': {
+                'Metadata-Flavor': 'Google',
+            },
+        }
+        self.assertEqual(http_instance._called_kwargs, [expected_kwargs])
+
+    def test_set_implicit_from_compute_engine(self):
+        self._implicit_compute_engine_helper('200')
+
+    def test_set_implicit_from_compute_engine_bad_status(self):
+        self._implicit_compute_engine_helper('404')
+
+    def test_set_implicit_from_compute_engine_raise_timeout(self):
+        self._implicit_compute_engine_helper('RAISE')
+
+    def test_set_implicit_both_appengine_and_compute(self):
+        from gcloud._testing import _Monkey
+        from gcloud.datastore import _implicit_environ
+
         APP_ENGINE_ID = 'GAE'
         APP_IDENTITY = _AppIdentity(APP_ENGINE_ID)
+        HTTPLIB2 = _Httplib2('GCE')
+
+        with self._monkey(None):
+            with _Monkey(_implicit_environ, app_identity=APP_IDENTITY,
+                         httplib2=HTTPLIB2):
+                self._callFUT()
+
+        self.assertEqual(_implicit_environ.DATASET_ID, APP_ENGINE_ID)
+        self.assertEqual(len(HTTPLIB2._http_instances), 0)
+
+    def test_set_implicit_three_env_appengine_and_compute(self):
+        from gcloud._testing import _Monkey
+        from gcloud.datastore import _implicit_environ
+
+        IMPLICIT_DATASET_ID = 'IMPLICIT'
+        APP_IDENTITY = _AppIdentity('GAE')
+        HTTPLIB2 = _Httplib2('GCE')
 
         with self._monkey(IMPLICIT_DATASET_ID):
-            with _Monkey(_implicit_environ, app_identity=APP_IDENTITY):
+            with _Monkey(_implicit_environ, app_identity=APP_IDENTITY,
+                         httplib2=HTTPLIB2):
                 self._callFUT()
 
         self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID)
+        self.assertEqual(len(HTTPLIB2._http_instances), 0)
 
 
 class Test_set_default_connection(unittest2.TestCase):
@@ -201,3 +272,35 @@ def __init__(self, app_id):
 
     def get_application_id(self):
         return self.app_id
+
+
+class _Http(object):
+
+    def __init__(self, parent):
+        self.parent = parent
+        self._called_uris = []
+        self._called_kwargs = []
+
+    def request(self, uri, **kwargs):
+        import socket
+
+        self._called_uris.append(uri)
+        self._called_kwargs.append(kwargs)
+
+        if self.parent.status == 'RAISE':
+            raise socket.timeout('timed out')
+        else:
+            return {'status': self.parent.status}, self.parent.project_id
+
+
+class _Httplib2(object):
+
+    def __init__(self, project_id, status='200'):
+        self.project_id = project_id
+        self.status = status
+        self._http_instances = []
+
+    def Http(self, timeout=None):
+        result = _Http(self)
+        self._http_instances.append((timeout, result))
+        return result