Allow passing explicit connection to dataset API methods.

See #825.

gcloud/datastore/dataset.py

@@ -38,29 +38,35 @@ def __init__(self, dataset_id, connection=None):
  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``.
  """
+ if connection is None:
+ connection = self.connection
  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)
+ if connection is None:
+ connection = self.connection
 
- def delete(self, keys):
+ return put(entities, connection=connection, dataset_id=self.dataset_id)
+
+ 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)
+ if connection is None:
+ connection = self.connection
+
+ 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 +78,25 @@ 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)
+ if connection is None:
+ connection = self.connection
 
- def transaction(self):
+ return Batch(dataset_id=self.dataset_id, connection=connection)
+
+ 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)
+ if connection is None:
+ 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

@@ -60,7 +60,7 @@ def _get(*args, **kw):
  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_connection_attr(self):
  from gcloud.datastore import dataset as MUT
  from gcloud._testing import _Monkey
 
@@ -82,6 +82,28 @@ def _get(*args, **kw):
  self.assertTrue(_called_with[0][1]['connection'] is conn)
  self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
 
+ def test_get_explicit_w_passed_connection(self):
+ from gcloud.datastore import dataset as MUT
+ from gcloud._testing import _Monkey
+
+ _called_with = []
+
+ def _get(*args, **kw):
+ _called_with.append((args, kw))
+
+ conn = object()
+ dataset = self._makeOne(connection=None)
+ key, missing, deferred = object(), [], []
+
+ with _Monkey(MUT, get=_get):
+ 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
  from gcloud._testing import _Monkey
@@ -101,7 +123,7 @@ def _put(*args, **kw):
  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_connection_attr(self):
  from gcloud.datastore import dataset as MUT
  from gcloud._testing import _Monkey
 
@@ -120,6 +142,25 @@ def _put(*args, **kw):
  self.assertTrue(_called_with[0][1]['connection'] is conn)
  self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
 
+ def test_put_w_passed_connection(self):
+ from gcloud.datastore import dataset as MUT
+ from gcloud._testing import _Monkey
+
+ _called_with = []
+
+ def _put(*args, **kw):
+ _called_with.append((args, kw))
+
+ entity, conn = object(), object()
+ dataset = self._makeOne()
+
+ with _Monkey(MUT, put=_put):
+ 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
  from gcloud._testing import _Monkey
@@ -139,7 +180,7 @@ def _delete(*args, **kw):
  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_connection_attr(self):
  from gcloud.datastore import dataset as MUT
  from gcloud._testing import _Monkey
 
@@ -157,6 +198,24 @@ def _delete(*args, **kw):
  self.assertTrue(_called_with[0][1]['connection'] is conn)
  self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
 
+ def test_delete_w_passed_connection(self):
+ from gcloud.datastore import dataset as MUT
+ from gcloud._testing import _Monkey
+
+ _called_with = []
+
+ def _delete(*args, **kw):
+ _called_with.append((args, kw))
+
+ key, conn = object(), object()
+ dataset = self._makeOne()
+ with _Monkey(MUT, delete=_delete):
+ dataset.delete([key], connection=conn)
+
+ self.assertEqual(_called_with[0][0], ([key],))
+ self.assertTrue(_called_with[0][1]['connection'] is conn)
+ self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
+
  def test_key_w_dataset_id(self):
  KIND = 'KIND'
  ID = 1234
@@ -191,7 +250,7 @@ def test_batch_wo_connection(self):
  self.assertEqual(batch.kwargs,
  {'dataset_id': self.DATASET_ID, 'connection': None})
 
- def test_batch_w_connection(self):
+ def test_batch_w_connection_attr(self):
  from gcloud.datastore import dataset as MUT
  from gcloud._testing import _Monkey
  conn = object()
@@ -205,6 +264,20 @@ def test_batch_w_connection(self):
  self.assertEqual(batch.kwargs,
  {'dataset_id': self.DATASET_ID, 'connection': conn})
 
+ def test_batch_w_passed_connection(self):
+ from gcloud.datastore import dataset as MUT
+ from gcloud._testing import _Monkey
+ conn = object()
+ dataset = self._makeOne()
+
+ with _Monkey(MUT, Batch=_Dummy):
+ 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
  from gcloud._testing import _Monkey
@@ -218,7 +291,7 @@ def test_transaction_wo_connection(self):
  self.assertEqual(xact.kwargs,
  {'dataset_id': self.DATASET_ID, 'connection': None})
 
- def test_transaction_w_connection(self):
+ def test_transaction_w_connection_attr(self):
  from gcloud.datastore import dataset as MUT
  from gcloud._testing import _Monkey
  conn = object()
@@ -232,6 +305,20 @@ def test_transaction_w_connection(self):
  self.assertEqual(xact.kwargs,
  {'dataset_id': self.DATASET_ID, 'connection': conn})
 
+ def test_transaction_w_passed_connection(self):
+ from gcloud.datastore import dataset as MUT
+ from gcloud._testing import _Monkey
+ conn = object()
+ dataset = self._makeOne()
+
+ with _Monkey(MUT, Transaction=_Dummy):
+ xact = dataset.transaction(connection=conn)
+
+ self.assertTrue(isinstance(xact, _Dummy))
+ self.assertEqual(xact.args, ())
+ self.assertEqual(xact.kwargs,
+ {'dataset_id': self.DATASET_ID, 'connection': conn})
+
  def test_query_w_dataset_id(self):
  KIND = 'KIND'
  dataset = self._makeOne()