diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py index 0a9c515a5ed8b..2972fa4f835e3 100644 --- a/gcloud/datastore/__init__.py +++ b/gcloud/datastore/__init__.py @@ -84,6 +84,9 @@ def set_default_dataset_id(dataset_id=None): if dataset_id is None: dataset_id = os.getenv(_DATASET_ENV_VAR_NAME) + if dataset_id is None: + dataset_id = _implicit_environ.app_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 8410ab7cdece6..f210270ac71da 100644 --- a/gcloud/datastore/_implicit_environ.py +++ b/gcloud/datastore/_implicit_environ.py @@ -4,9 +4,27 @@ imply the current dataset ID and connection from the enviroment. """ +try: + from google.appengine.api import app_identity +except ImportError: + app_identity = None + DATASET_ID = None """Module global to allow persistent implied dataset ID from enviroment.""" CONNECTION = None """Module global to allow persistent implied connection from enviroment.""" + + +def app_engine_id(): + """Gets the App Engine application ID if it can be found. + + :rtype: string or ``NoneType`` + :returns: App Engine application ID if running in App Engine, + else ``None``. + """ + if app_identity is None: + return None + + return app_identity.get_application_id() diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py index 7c836b1d7c0ac..9d484c85f63bd 100644 --- a/gcloud/datastore/test___init__.py +++ b/gcloud/datastore/test___init__.py @@ -78,6 +78,33 @@ def test_set_explicit_None_w_env_var_set(self): self._callFUT(None) self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID) + def test_set_implicit_from_appengine(self): + from gcloud._testing import _Monkey + from gcloud.datastore import _implicit_environ + + APP_ENGINE_ID = 'GAE' + APP_IDENTITY = _AppIdentity(APP_ENGINE_ID) + + with self._monkey(None): + with _Monkey(_implicit_environ, app_identity=APP_IDENTITY): + self._callFUT() + + self.assertEqual(_implicit_environ.DATASET_ID, APP_ENGINE_ID) + + def test_set_implicit_both_env_and_appengine(self): + from gcloud._testing import _Monkey + from gcloud.datastore import _implicit_environ + + IMPLICIT_DATASET_ID = 'IMPLICIT' + APP_ENGINE_ID = 'GAE' + APP_IDENTITY = _AppIdentity(APP_ENGINE_ID) + + 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) + class Test_set_default_connection(unittest2.TestCase): @@ -165,3 +192,12 @@ def test_it(self): self.assertTrue(isinstance(found, Connection)) self.assertTrue(found._credentials is client._signed) self.assertTrue(client._get_app_default_called) + + +class _AppIdentity(object): + + def __init__(self, app_id): + self.app_id = app_id + + def get_application_id(self): + return self.app_id