Skip to content

Commit

Permalink
Adding implicit dataset ID support for Compute Engine.
Browse files Browse the repository at this point in the history
  • Loading branch information
dhermes committed Jan 27, 2015
1 parent d9fd60e commit 6267857
Show file tree
Hide file tree
Showing 3 changed files with 149 additions and 7 deletions.
13 changes: 8 additions & 5 deletions gcloud/datastore/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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

Expand Down
38 changes: 37 additions & 1 deletion gcloud/datastore/_implicit_environ.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,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:
Expand All @@ -18,7 +21,7 @@


def app_engine_id():
"""Gets the App Engine application ID if it can be found.
"""Gets the App Engine application ID if it can be implied.
:rtype: string or ``NoneType``
:returns: App Engine application ID if running in App Engine,
Expand All @@ -28,3 +31,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 implied.
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
105 changes: 104 additions & 1 deletion gcloud/datastore/test___init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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

0 comments on commit 6267857

Please sign in to comment.