diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py index 2972fa4f835e..1a2df8147706 100644 --- a/gcloud/datastore/__init__.py +++ b/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 diff --git a/gcloud/datastore/_implicit_environ.py b/gcloud/datastore/_implicit_environ.py index 7e6f5357dc8e..dab44fe782e2 100644 --- a/gcloud/datastore/_implicit_environ.py +++ b/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 diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py index 9d484c85f63b..754a3399a7a6 100644 --- a/gcloud/datastore/test___init__.py +++ b/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