diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py index 965cf8730e92..cdb739e0b78b 100644 --- a/gcloud/datastore/__init__.py +++ b/gcloud/datastore/__init__.py @@ -59,12 +59,10 @@ def set_default_dataset(dataset_id=None): - """Determines auth settings from local enviroment. + """Set default dataset ID either explicitly or implicitly as fall-back. - Sets a default dataset either explicitly or via the local - enviroment. Currently only supports enviroment variable but will - implicitly support App Engine, Compute Engine and other environments - in the future. + 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 @@ -80,6 +78,16 @@ def set_default_dataset(dataset_id=None): _implicit_environ.DATASET = get_dataset(dataset_id) +def set_default_connection(connection=None): + """Set default connection either explicitly or implicitly as fall-back. + + :type connection: :class:`gcloud.datastore.connection.Connection` + :param connection: A connection provided to be the default. + """ + connection = connection or get_connection() + _implicit_environ.CONNECTION = connection + + def get_connection(): """Shortcut method to establish a connection to the Cloud Datastore. diff --git a/gcloud/datastore/_implicit_environ.py b/gcloud/datastore/_implicit_environ.py index 088067831820..7fcb14dd570a 100644 --- a/gcloud/datastore/_implicit_environ.py +++ b/gcloud/datastore/_implicit_environ.py @@ -1,7 +1,7 @@ """Module to provide implicit behavior based on enviroment. Acts as a mutable namespace to allow the datastore package to -imply the current dataset from the enviroment. +imply the current dataset and connection from the enviroment. Also provides a base class for classes in the `datastore` package which could utilize the implicit enviroment. @@ -11,6 +11,9 @@ DATASET = None """Module global to allow persistent implied dataset from enviroment.""" +CONNECTION = None +"""Module global to allow persistent implied connection from enviroment.""" + class _DatastoreBase(object): """Base for all classes in the datastore package. diff --git a/gcloud/datastore/test___init__.py b/gcloud/datastore/test___init__.py index 603eee1543c2..4f5782f0e67b 100644 --- a/gcloud/datastore/test___init__.py +++ b/gcloud/datastore/test___init__.py @@ -89,6 +89,43 @@ def test_set_explicit(self): self._test_with_environ({}, DATASET_ID, dataset_id=DATASET_ID) +class Test_set_default_connection(unittest2.TestCase): + + def setUp(self): + from gcloud.datastore import _implicit_environ + self._replaced_connection = _implicit_environ.CONNECTION + _implicit_environ.CONNECTION = None + + def tearDown(self): + from gcloud.datastore import _implicit_environ + _implicit_environ.CONNECTION = self._replaced_connection + + def _callFUT(self, connection=None): + from gcloud.datastore import set_default_connection + return set_default_connection(connection=connection) + + def test_set_explicit(self): + from gcloud.datastore import _implicit_environ + + self.assertEqual(_implicit_environ.CONNECTION, None) + fake_cnxn = object() + self._callFUT(connection=fake_cnxn) + self.assertEqual(_implicit_environ.CONNECTION, fake_cnxn) + + def test_set_implicit(self): + from gcloud._testing import _Monkey + from gcloud import datastore + from gcloud.datastore import _implicit_environ + + self.assertEqual(_implicit_environ.CONNECTION, None) + + fake_cnxn = object() + with _Monkey(datastore, get_connection=lambda: fake_cnxn): + self._callFUT() + + self.assertEqual(_implicit_environ.CONNECTION, fake_cnxn) + + class Test_get_dataset(unittest2.TestCase): def _callFUT(self, dataset_id):