#825: Allow passing explicit connection to dataset API methods.

gcloud/datastore/dataset.py

@@ -27,40 +27,34 @@ class Dataset(object):
 
     :type dataset_id: string
     :param dataset_id: (required) dataset ID to pass to proxied API methods.
-
-    :type connection: :class:`gcloud.datastore.connection.Connection`, or None
-    :param connection: (optional) connection to pass to proxied API methods
     """
 
-    def __init__(self, dataset_id, connection=None):
+    def __init__(self, dataset_id):
         if dataset_id is None:
             raise ValueError('dataset_id required')
         self.dataset_id = dataset_id
-        self.connection = connection
 
-    def get(self, keys, missing=None, deferred=None):
+    def get(self, keys, missing=None, deferred=None, connection=None):
         """Proxy to :func:`gcloud.datastore.api.get`.
 
         Passes our ``dataset_id``.
         """
         return get(keys, missing=missing, deferred=deferred,
-                   connection=self.connection, dataset_id=self.dataset_id)
+                   connection=connection, dataset_id=self.dataset_id)
 
-    def put(self, entities):
+    def put(self, entities, connection=None):
         """Proxy to :func:`gcloud.datastore.api.put`.
 
         Passes our ``dataset_id``.
         """
-        return put(entities, connection=self.connection,
-                   dataset_id=self.dataset_id)
+        return put(entities, connection=connection, dataset_id=self.dataset_id)
 
-    def delete(self, keys):
+    def delete(self, keys, connection=None):
         """Proxy to :func:`gcloud.datastore.api.delete`.
 
         Passes our ``dataset_id``.
         """
-        return delete(keys, connection=self.connection,
-                      dataset_id=self.dataset_id)
+        return delete(keys, connection=connection, dataset_id=self.dataset_id)
 
     def key(self, *path_args, **kwargs):
         """Proxy to :class:`gcloud.datastore.key.Key`.
@@ -72,20 +66,19 @@ def key(self, *path_args, **kwargs):
         kwargs['dataset_id'] = self.dataset_id
         return Key(*path_args, **kwargs)
 
-    def batch(self):
+    def batch(self, connection=None):
         """Proxy to :class:`gcloud.datastore.batch.Batch`.
 
         Passes our ``dataset_id``.
         """
-        return Batch(dataset_id=self.dataset_id, connection=self.connection)
+        return Batch(dataset_id=self.dataset_id, connection=connection)
 
-    def transaction(self):
+    def transaction(self, connection=None):
         """Proxy to :class:`gcloud.datastore.transaction.Transaction`.
 
         Passes our ``dataset_id``.
         """
-        return Transaction(dataset_id=self.dataset_id,
-                           connection=self.connection)
+        return Transaction(dataset_id=self.dataset_id, connection=connection)
 
     def query(self, **kwargs):
         """Proxy to :class:`gcloud.datastore.query.Query`.

gcloud/datastore/test_dataset.py

@@ -23,44 +23,41 @@ def _getTargetClass(self):
         from gcloud.datastore.dataset import Dataset
         return Dataset
 
-    def _makeOne(self, dataset_id=DATASET_ID, connection=None):
-        return self._getTargetClass()(dataset_id, connection)
+    def _makeOne(self, dataset_id=DATASET_ID):
+        return self._getTargetClass()(dataset_id)
 
     def test_ctor_w_dataset_id_None(self):
         self.assertRaises(ValueError, self._makeOne, None)
 
-    def test_ctor_w_dataset_id_no_connection(self):
+    def test_ctor_w_dataset_id(self):
         dataset = self._makeOne()
         self.assertEqual(dataset.dataset_id, self.DATASET_ID)
 
-    def test_ctor_w_dataset_id_w_connection(self):
-        conn = object()
-        dataset = self._makeOne(connection=conn)
-        self.assertEqual(dataset.dataset_id, self.DATASET_ID)
-        self.assertTrue(dataset.connection is conn)
-
-    def test_get_defaults(self):
+    def test_get_defaults_w_implicit_connection(self):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
+        from gcloud.datastore._testing import _monkey_defaults
 
         _called_with = []
 
         def _get(*args, **kw):
             _called_with.append((args, kw))
 
+        conn = object()
         dataset = self._makeOne()
         key = object()
 
         with _Monkey(MUT, get=_get):
-            dataset.get([key])
+            with _monkey_defaults(connection=conn):
+                dataset.get([key])
 
         self.assertEqual(_called_with[0][0], ([key],))
         self.assertTrue(_called_with[0][1]['missing'] is None)
         self.assertTrue(_called_with[0][1]['deferred'] is None)
         self.assertTrue(_called_with[0][1]['connection'] is None)
         self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
 
-    def test_get_explicit(self):
+    def test_get_explicit_w_explicit_connection(self):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
 
@@ -70,38 +67,40 @@ def _get(*args, **kw):
             _called_with.append((args, kw))
 
         conn = object()
-        dataset = self._makeOne(connection=conn)
+        dataset = self._makeOne()
         key, missing, deferred = object(), [], []
 
         with _Monkey(MUT, get=_get):
-            dataset.get([key], missing, deferred)
+            dataset.get([key], missing, deferred, connection=conn)
 
         self.assertEqual(_called_with[0][0], ([key],))
         self.assertTrue(_called_with[0][1]['missing'] is missing)
         self.assertTrue(_called_with[0][1]['deferred'] is deferred)
         self.assertTrue(_called_with[0][1]['connection'] is conn)
         self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
 
-    def test_put_wo_connection(self):
-        from gcloud.datastore import dataset as MUT
+    def test_put_w_implicit_connection(self):
         from gcloud._testing import _Monkey
+        from gcloud.datastore import dataset as MUT
+        from gcloud.datastore._testing import _monkey_defaults
 
         _called_with = []
 
         def _put(*args, **kw):
             _called_with.append((args, kw))
 
+        entity, conn = object(), object()
         dataset = self._makeOne()
-        entity = object()
 
         with _Monkey(MUT, put=_put):
-            dataset.put([entity])
+            with _monkey_defaults(connection=conn):
+                dataset.put([entity])
 
         self.assertEqual(_called_with[0][0], ([entity],))
         self.assertTrue(_called_with[0][1]['connection'] is None)
         self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
 
-    def test_put_w_connection(self):
+    def test_put_w_explicit_connection(self):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
 
@@ -111,35 +110,36 @@ def _put(*args, **kw):
             _called_with.append((args, kw))
 
         entity, conn = object(), object()
-        dataset = self._makeOne(connection=conn)
+        dataset = self._makeOne()
 
         with _Monkey(MUT, put=_put):
-            dataset.put([entity])
+            dataset.put([entity], connection=conn)
 
         self.assertEqual(_called_with[0][0], ([entity],))
         self.assertTrue(_called_with[0][1]['connection'] is conn)
         self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
 
-    def test_delete_wo_connection(self):
-        from gcloud.datastore import dataset as MUT
+    def test_delete_w_implicit_connection(self):
         from gcloud._testing import _Monkey
+        from gcloud.datastore import dataset as MUT
+        from gcloud.datastore._testing import _monkey_defaults
 
         _called_with = []
 
         def _delete(*args, **kw):
             _called_with.append((args, kw))
 
+        key, conn = object(), object()
         dataset = self._makeOne()
-        key = object()
-
         with _Monkey(MUT, delete=_delete):
-            dataset.delete([key])
+            with _monkey_defaults(connection=conn):
+                dataset.delete([key])
 
         self.assertEqual(_called_with[0][0], ([key],))
         self.assertTrue(_called_with[0][1]['connection'] is None)
         self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
 
-    def test_delete_w_connection(self):
+    def test_delete_w_explicit_connection(self):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
 
@@ -149,9 +149,9 @@ def _delete(*args, **kw):
             _called_with.append((args, kw))
 
         key, conn = object(), object()
-        dataset = self._makeOne(connection=conn)
+        dataset = self._makeOne()
         with _Monkey(MUT, delete=_delete):
-            dataset.delete([key])
+            dataset.delete([key], connection=conn)
 
         self.assertEqual(_called_with[0][0], ([key],))
         self.assertTrue(_called_with[0][1]['connection'] is conn)
@@ -178,54 +178,60 @@ def test_key_wo_dataset_id(self):
         self.assertEqual(key.args, (KIND, ID))
         self.assertEqual(key.kwargs, {'dataset_id': self.DATASET_ID})
 
-    def test_batch_wo_connection(self):
-        from gcloud.datastore import dataset as MUT
+    def test_batch_w_implicit_connection(self):
         from gcloud._testing import _Monkey
+        from gcloud.datastore import dataset as MUT
+        from gcloud.datastore._testing import _monkey_defaults
+        conn = object()
         dataset = self._makeOne()
 
         with _Monkey(MUT, Batch=_Dummy):
-            batch = dataset.batch()
+            with _monkey_defaults(connection=conn):
+                batch = dataset.batch()
 
         self.assertTrue(isinstance(batch, _Dummy))
         self.assertEqual(batch.args, ())
         self.assertEqual(batch.kwargs,
                          {'dataset_id': self.DATASET_ID, 'connection': None})
 
-    def test_batch_w_connection(self):
+    def test_batch_w_explicit_connection(self):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
         conn = object()
-        dataset = self._makeOne(connection=conn)
+        dataset = self._makeOne()
 
         with _Monkey(MUT, Batch=_Dummy):
-            batch = dataset.batch()
+            batch = dataset.batch(connection=conn)
 
         self.assertTrue(isinstance(batch, _Dummy))
         self.assertEqual(batch.args, ())
         self.assertEqual(batch.kwargs,
                          {'dataset_id': self.DATASET_ID, 'connection': conn})
 
-    def test_transaction_wo_connection(self):
-        from gcloud.datastore import dataset as MUT
+    def test_transaction_w_implicit_connection(self):
         from gcloud._testing import _Monkey
+        from gcloud.datastore import dataset as MUT
+        from gcloud.datastore._testing import _monkey_defaults
+        conn = object()
         dataset = self._makeOne()
 
         with _Monkey(MUT, Transaction=_Dummy):
-            xact = dataset.transaction()
+            with _monkey_defaults(connection=conn):
+                xact = dataset.transaction()
 
         self.assertTrue(isinstance(xact, _Dummy))
         self.assertEqual(xact.args, ())
         self.assertEqual(xact.kwargs,
                          {'dataset_id': self.DATASET_ID, 'connection': None})
 
-    def test_transaction_w_connection(self):
+    def test_transaction_w_explicit_connection(self):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
         conn = object()
-        dataset = self._makeOne(connection=conn)
+        dataset = self._makeOne()
 
         with _Monkey(MUT, Transaction=_Dummy):
-            xact = dataset.transaction()
+            xact = dataset.transaction(connection=conn)
 
         self.assertTrue(isinstance(xact, _Dummy))
         self.assertEqual(xact.args, ())