Hold an optional connection on the dataset.

Pass it through to the API methods, rather than having the caller pass it.

Don't try to resolve it if None is passed (the API methods do that for us).

gcloud/datastore/dataset.py

@@ -25,31 +25,32 @@
 class Dataset(object):
     """Convenience wrapper for invoking APIs/factories w/ a dataset ID."""
 
-    def __init__(self, dataset_id):
+    def __init__(self, dataset_id, connection=None):
         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, connection=None):
+    def get(self, keys, missing=None, deferred=None):
         """Proxy to :func:`gcloud.datastore.api.get`.
 
         Passes our ``dataset_id``.
         """
-        return get(keys, missing, deferred, connection, self.dataset_id)
+        return get(keys, missing, deferred, self.connection, self.dataset_id)
 
-    def put(self, entities, connection=None):
+    def put(self, entities):
         """Proxy to :func:`gcloud.datastore.api.put`.
 
         Passes our ``dataset_id``.
         """
-        return put(entities, connection, dataset_id=self.dataset_id)
+        return put(entities, self.connection, dataset_id=self.dataset_id)
 
-    def delete(self, keys, connection=None):
+    def delete(self, keys):
         """Proxy to :func:`gcloud.datastore.api.delete`.
 
         Passes our ``dataset_id``.
         """
-        return delete(keys, connection, dataset_id=self.dataset_id)
+        return delete(keys, self.connection, dataset_id=self.dataset_id)
 
     def key(self, *path_args, **kwargs):
         """Proxy to :func:`gcloud.datastore.key.Key`.
@@ -62,19 +63,20 @@ def key(self, *path_args, **kwargs):
         kwargs['dataset_id'] = self.dataset_id
         return Key(*path_args, **kwargs)
 
-    def batch(self, connection=None):
+    def batch(self):
         """Proxy to :func:`gcloud.datastore.batch.Batch`.
 
         Passes our ``dataset_id``.
         """
-        return Batch(dataset_id=self.dataset_id, connection=connection)
+        return Batch(dataset_id=self.dataset_id, connection=self.connection)
 
-    def transaction(self, connection=None):
+    def transaction(self):
         """Proxy to :func:`gcloud.datastore.transaction.Transaction`.
 
         Passes our ``dataset_id``.
         """
-        return Transaction(dataset_id=self.dataset_id, connection=connection)
+        return Transaction(dataset_id=self.dataset_id,
+                           connection=self.connection)
 
     def query(self,
               kind=None,

gcloud/datastore/test_dataset.py

@@ -23,16 +23,22 @@ def _getTargetClass(self):
         from gcloud.datastore.dataset import Dataset
         return Dataset
 
-    def _makeOne(self, dataset_id=DATASET_ID):
-        return self._getTargetClass()(dataset_id=dataset_id)
+    def _makeOne(self, dataset_id=DATASET_ID, connection=None):
+        return self._getTargetClass()(dataset_id, connection)
 
-    def test_ctor_w_None(self):
+    def test_ctor_w_dataset_id_None(self):
         self.assertRaises(ValueError, self._makeOne, None)
 
-    def test_ctor_w_dataset_id(self):
+    def test_ctor_w_dataset_id_no_connection(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):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
@@ -60,11 +66,12 @@ def test_get_explicit(self):
         def _get(*args, **kw):
             _called_with.append((args, kw))
 
-        dataset = self._makeOne()
-        key, missing, deferred, conn = object(), [], [], object()
+        conn = object()
+        dataset = self._makeOne(connection=conn)
+        key, missing, deferred = object(), [], []
 
         with _Monkey(MUT, get=_get):
-            dataset.get([key], missing, deferred, conn)
+            dataset.get([key], missing, deferred)
 
         args = ([key], missing, deferred, conn, self.DATASET_ID)
         self.assertEqual(_called_with, [(args, {})])
@@ -96,11 +103,11 @@ def test_put_w_connection(self):
         def _put(*args, **kw):
             _called_with.append((args, kw))
 
-        dataset = self._makeOne()
         entity, conn = object(), object()
+        dataset = self._makeOne(connection=conn)
 
         with _Monkey(MUT, put=_put):
-            dataset.put([entity], conn)
+            dataset.put([entity])
 
         self.assertEqual(_called_with,
                          [(([entity], conn), {'dataset_id': self.DATASET_ID})])
@@ -132,10 +139,10 @@ def test_delete_w_connection(self):
         def _delete(*args, **kw):
             _called_with.append((args, kw))
 
-        dataset = self._makeOne()
         key, conn = object(), object()
+        dataset = self._makeOne(connection=conn)
         with _Monkey(MUT, delete=_delete):
-            dataset.delete([key], conn)
+            dataset.delete([key])
 
         self.assertEqual(_called_with,
                          [(([key], conn), {'dataset_id': self.DATASET_ID})])
@@ -191,11 +198,11 @@ def test_batch_wo_connection(self):
     def test_batch_w_connection(self):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
-        dataset = self._makeOne()
         conn = object()
+        dataset = self._makeOne(connection=conn)
 
         with _Monkey(MUT, Batch=_Dummy):
-            batch = dataset.batch(conn)
+            batch = dataset.batch()
 
         self.assertTrue(isinstance(batch, _Dummy))
         self.assertEqual(batch.args, ())
@@ -218,11 +225,11 @@ def test_transaction_wo_connection(self):
     def test_transaction_w_connection(self):
         from gcloud.datastore import dataset as MUT
         from gcloud._testing import _Monkey
-        dataset = self._makeOne()
         conn = object()
+        dataset = self._makeOne(connection=conn)
 
         with _Monkey(MUT, Transaction=_Dummy):
-            xact = dataset.transaction(conn)
+            xact = dataset.transaction()
 
         self.assertTrue(isinstance(xact, _Dummy))
         self.assertEqual(xact.args, ())