diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py index 6f77e05fba51..eb03a318cd82 100644 --- a/gcloud/datastore/__init__.py +++ b/gcloud/datastore/__init__.py @@ -67,6 +67,7 @@ """The scopes required for authenticating as a Cloud Datastore consumer.""" _DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID' +_GCD_DATASET_ENV_VAR_NAME = 'DATASTORE_DATASET' def set_default_dataset_id(dataset_id=None): @@ -86,6 +87,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 = os.getenv(_GCD_DATASET_ENV_VAR_NAME) + if dataset_id is None: dataset_id = _implicit_environ.app_engine_id() diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py index 916599590d83..9c259556239b 100644 --- a/gcloud/datastore/connection.py +++ b/gcloud/datastore/connection.py @@ -14,12 +14,17 @@ """Connections to gcloud datastore API servers.""" +import os + from gcloud import connection from gcloud.exceptions import make_exception from gcloud.datastore import _datastore_v1_pb2 as datastore_pb from gcloud.datastore import helpers +_GCD_HOST_ENV_VAR_NAME = 'DATASTORE_HOST' + + class Connection(connection.Connection): """A connection to the Google Cloud Datastore via the Protobuf API. @@ -30,6 +35,10 @@ class Connection(connection.Connection): :param credentials: The OAuth2 Credentials to use for this connection. """ + API_BASE_URL = os.getenv(_GCD_HOST_ENV_VAR_NAME, + connection.Connection.API_BASE_URL) + """The base of the API call URL.""" + API_VERSION = 'v1beta2' """The version of the API, used in building the API call's URL.""" diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py index e4907862a10f..140bbddc3a93 100644 --- a/gcloud/datastore/test___init__.py +++ b/gcloud/datastore/test___init__.py @@ -30,11 +30,11 @@ def _callFUT(self, dataset_id=None): from gcloud.datastore import set_default_dataset_id return set_default_dataset_id(dataset_id=dataset_id) - def _monkeyEnviron(self, implicit_dataset_id): + def _monkeyEnviron(self, implicit_dataset_id, environ=None): import os from gcloud._testing import _Monkey from gcloud.datastore import _DATASET_ENV_VAR_NAME - environ = {_DATASET_ENV_VAR_NAME: implicit_dataset_id} + environ = environ or {_DATASET_ENV_VAR_NAME: implicit_dataset_id} return _Monkey(os, getenv=environ.get) def _monkeyImplicit(self, connection=None, app_identity=None): @@ -112,6 +112,55 @@ def test_set_explicit_None_w_env_var_set(self): self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID) + def test_set_from_gcd_env_var(self): + from gcloud.datastore import _GCD_DATASET_ENV_VAR_NAME + from gcloud.datastore import _implicit_environ + + GCD_DATASET_ID = 'GCD-IMPLICIT' + ENVIRON = {_GCD_DATASET_ENV_VAR_NAME: GCD_DATASET_ID} + + with self._monkeyEnviron(None, environ=ENVIRON): + with self._monkeyImplicit(): + self._callFUT() + + self.assertEqual(_implicit_environ.DATASET_ID, GCD_DATASET_ID) + + def test_set_gcd_and_production_env_vars(self): + from gcloud.datastore import _DATASET_ENV_VAR_NAME + from gcloud.datastore import _GCD_DATASET_ENV_VAR_NAME + from gcloud.datastore import _implicit_environ + + IMPLICIT_DATASET_ID = 'IMPLICIT' + GCD_DATASET_ID = 'GCD-IMPLICIT' + ENVIRON = { + _DATASET_ENV_VAR_NAME: IMPLICIT_DATASET_ID, + _GCD_DATASET_ENV_VAR_NAME: GCD_DATASET_ID, + } + + with self._monkeyEnviron(None, environ=ENVIRON): + with self._monkeyImplicit(): + self._callFUT() + + self.assertNotEqual(_implicit_environ.DATASET_ID, GCD_DATASET_ID) + self.assertEqual(_implicit_environ.DATASET_ID, IMPLICIT_DATASET_ID) + + def test_set_gcd_env_vars_and_appengine(self): + from gcloud.datastore import _GCD_DATASET_ENV_VAR_NAME + from gcloud.datastore import _implicit_environ + + GCD_DATASET_ID = 'GCD-IMPLICIT' + ENVIRON = {_GCD_DATASET_ENV_VAR_NAME: GCD_DATASET_ID} + + APP_ENGINE_ID = 'GAE' + APP_IDENTITY = _AppIdentity(APP_ENGINE_ID) + + with self._monkeyEnviron(None, environ=ENVIRON): + with self._monkeyImplicit(app_identity=APP_IDENTITY): + self._callFUT() + + self.assertNotEqual(_implicit_environ.DATASET_ID, APP_ENGINE_ID) + self.assertEqual(_implicit_environ.DATASET_ID, GCD_DATASET_ID) + def test_set_implicit_from_appengine(self): from gcloud.datastore import _implicit_environ diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py index 487c7b7d6e24..30e202509699 100644 --- a/gcloud/datastore/test_connection.py +++ b/gcloud/datastore/test_connection.py @@ -46,6 +46,47 @@ def _verifyProtobufCall(self, called_with, URI, conn): self.assertEqual(called_with['headers']['User-Agent'], conn.USER_AGENT) + def test_default_url(self): + from gcloud.connection import Connection + + klass = self._getTargetClass() + self.assertEqual(klass.API_BASE_URL, Connection.API_BASE_URL) + + def test_custom_url(self): + import os + import sys + from gcloud._testing import _Monkey + from gcloud.connection import Connection as BaseConnection + from gcloud.datastore.connection import _GCD_HOST_ENV_VAR_NAME + + HOST = object() + fake_environ = {_GCD_HOST_ENV_VAR_NAME: HOST} + + def fake_getenv(key, default_val=None): + return fake_environ.get(key, default_val) + + # We want to temporarily all the gcloud.datastore modules except for + # the gcloud.datastore package itself. Then re-import. + gcloud_keys = [key for key in sys.modules + if 'gcloud.datastore.' in key] + modules_as_is = sys.modules.copy() + try: + for key in gcloud_keys: + sys.modules.pop(key) + + with _Monkey(os, getenv=fake_getenv): + from gcloud.datastore.connection import Connection + finally: + sys.modules.update(modules_as_is) + + self.assertNotEqual(Connection.API_BASE_URL, + BaseConnection.API_BASE_URL) + self.assertEqual(Connection.API_BASE_URL, HOST) + + # Make sure the restored import. + from gcloud.datastore.connection import Connection + self.assertEqual(Connection.API_BASE_URL, BaseConnection.API_BASE_URL) + def test_ctor_defaults(self): conn = self._makeOne() self.assertEqual(conn.credentials, None)