Skip to content

Commit 9130c03

Browse files
tseavertseaver
committed
Allow passing explicit connection to dataset API methods.
See #825.
1 parent 8b5c5b9 commit 9130c03

File tree

2 files changed

+116
-18
lines changed

2 files changed

+116
-18
lines changed

‎gcloud/datastore/dataset.py‎

Lines changed: 24 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,29 +38,35 @@ def __init__(self, dataset_id, connection=None):
3838
self.dataset_id = dataset_id
3939
self.connection = connection
4040

41-
def get(self, keys, missing=None, deferred=None):
41+
def get(self, keys, missing=None, deferred=None, connection=None):
4242
"""Proxy to :func:`gcloud.datastore.api.get`.
4343
4444
Passes our ``dataset_id``.
4545
"""
46+
if connection is None:
47+
connection = self.connection
4648
return get(keys, missing=missing, deferred=deferred,
47-
connection=self.connection, dataset_id=self.dataset_id)
49+
connection=connection, dataset_id=self.dataset_id)
4850

49-
def put(self, entities):
51+
def put(self, entities, connection=None):
5052
"""Proxy to :func:`gcloud.datastore.api.put`.
5153
5254
Passes our ``dataset_id``.
5355
"""
54-
return put(entities, connection=self.connection,
55-
dataset_id=self.dataset_id)
56+
if connection is None:
57+
connection = self.connection
5658

57-
def delete(self, keys):
59+
return put(entities, connection=connection, dataset_id=self.dataset_id)
60+
61+
def delete(self, keys, connection=None):
5862
"""Proxy to :func:`gcloud.datastore.api.delete`.
5963
6064
Passes our ``dataset_id``.
6165
"""
62-
return delete(keys, connection=self.connection,
63-
dataset_id=self.dataset_id)
66+
if connection is None:
67+
connection = self.connection
68+
69+
return delete(keys, connection=connection, dataset_id=self.dataset_id)
6470

6571
def key(self, *path_args, **kwargs):
6672
"""Proxy to :class:`gcloud.datastore.key.Key`.
@@ -72,20 +78,25 @@ def key(self, *path_args, **kwargs):
7278
kwargs['dataset_id'] = self.dataset_id
7379
return Key(*path_args, **kwargs)
7480

75-
def batch(self):
81+
def batch(self, connection=None):
7682
"""Proxy to :class:`gcloud.datastore.batch.Batch`.
7783
7884
Passes our ``dataset_id``.
7985
"""
80-
return Batch(dataset_id=self.dataset_id, connection=self.connection)
86+
if connection is None:
87+
connection = self.connection
8188

82-
def transaction(self):
89+
return Batch(dataset_id=self.dataset_id, connection=connection)
90+
91+
def transaction(self, connection=None):
8392
"""Proxy to :class:`gcloud.datastore.transaction.Transaction`.
8493
8594
Passes our ``dataset_id``.
8695
"""
87-
return Transaction(dataset_id=self.dataset_id,
88-
connection=self.connection)
96+
if connection is None:
97+
connection = self.connection
98+
99+
return Transaction(dataset_id=self.dataset_id, connection=connection)
89100

90101
def query(self, **kwargs):
91102
"""Proxy to :class:`gcloud.datastore.query.Query`.

‎gcloud/datastore/test_dataset.py‎

Lines changed: 92 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ def _get(*args, **kw):
6060
self.assertTrue(_called_with[0][1]['connection'] is None)
6161
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
6262

63-
def test_get_explicit(self):
63+
def test_get_explicit_w_connection_attr(self):
6464
from gcloud.datastore import dataset as MUT
6565
from gcloud._testing import _Monkey
6666

@@ -82,6 +82,28 @@ def _get(*args, **kw):
8282
self.assertTrue(_called_with[0][1]['connection'] is conn)
8383
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
8484

85+
def test_get_explicit_w_passed_connection(self):
86+
from gcloud.datastore import dataset as MUT
87+
from gcloud._testing import _Monkey
88+
89+
_called_with = []
90+
91+
def _get(*args, **kw):
92+
_called_with.append((args, kw))
93+
94+
conn = object()
95+
dataset = self._makeOne(connection=None)
96+
key, missing, deferred = object(), [], []
97+
98+
with _Monkey(MUT, get=_get):
99+
dataset.get([key], missing, deferred, connection=conn)
100+
101+
self.assertEqual(_called_with[0][0], ([key],))
102+
self.assertTrue(_called_with[0][1]['missing'] is missing)
103+
self.assertTrue(_called_with[0][1]['deferred'] is deferred)
104+
self.assertTrue(_called_with[0][1]['connection'] is conn)
105+
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
106+
85107
def test_put_wo_connection(self):
86108
from gcloud.datastore import dataset as MUT
87109
from gcloud._testing import _Monkey
@@ -101,7 +123,7 @@ def _put(*args, **kw):
101123
self.assertTrue(_called_with[0][1]['connection'] is None)
102124
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
103125

104-
def test_put_w_connection(self):
126+
def test_put_w_connection_attr(self):
105127
from gcloud.datastore import dataset as MUT
106128
from gcloud._testing import _Monkey
107129

@@ -120,6 +142,25 @@ def _put(*args, **kw):
120142
self.assertTrue(_called_with[0][1]['connection'] is conn)
121143
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
122144

145+
def test_put_w_passed_connection(self):
146+
from gcloud.datastore import dataset as MUT
147+
from gcloud._testing import _Monkey
148+
149+
_called_with = []
150+
151+
def _put(*args, **kw):
152+
_called_with.append((args, kw))
153+
154+
entity, conn = object(), object()
155+
dataset = self._makeOne()
156+
157+
with _Monkey(MUT, put=_put):
158+
dataset.put([entity], connection=conn)
159+
160+
self.assertEqual(_called_with[0][0], ([entity],))
161+
self.assertTrue(_called_with[0][1]['connection'] is conn)
162+
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
163+
123164
def test_delete_wo_connection(self):
124165
from gcloud.datastore import dataset as MUT
125166
from gcloud._testing import _Monkey
@@ -139,7 +180,7 @@ def _delete(*args, **kw):
139180
self.assertTrue(_called_with[0][1]['connection'] is None)
140181
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
141182

142-
def test_delete_w_connection(self):
183+
def test_delete_w_connection_attr(self):
143184
from gcloud.datastore import dataset as MUT
144185
from gcloud._testing import _Monkey
145186

@@ -157,6 +198,24 @@ def _delete(*args, **kw):
157198
self.assertTrue(_called_with[0][1]['connection'] is conn)
158199
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
159200

201+
def test_delete_w_passed_connection(self):
202+
from gcloud.datastore import dataset as MUT
203+
from gcloud._testing import _Monkey
204+
205+
_called_with = []
206+
207+
def _delete(*args, **kw):
208+
_called_with.append((args, kw))
209+
210+
key, conn = object(), object()
211+
dataset = self._makeOne()
212+
with _Monkey(MUT, delete=_delete):
213+
dataset.delete([key], connection=conn)
214+
215+
self.assertEqual(_called_with[0][0], ([key],))
216+
self.assertTrue(_called_with[0][1]['connection'] is conn)
217+
self.assertEqual(_called_with[0][1]['dataset_id'], self.DATASET_ID)
218+
160219
def test_key_w_dataset_id(self):
161220
KIND = 'KIND'
162221
ID = 1234
@@ -191,7 +250,7 @@ def test_batch_wo_connection(self):
191250
self.assertEqual(batch.kwargs,
192251
{'dataset_id': self.DATASET_ID, 'connection': None})
193252

194-
def test_batch_w_connection(self):
253+
def test_batch_w_connection_attr(self):
195254
from gcloud.datastore import dataset as MUT
196255
from gcloud._testing import _Monkey
197256
conn = object()
@@ -205,6 +264,20 @@ def test_batch_w_connection(self):
205264
self.assertEqual(batch.kwargs,
206265
{'dataset_id': self.DATASET_ID, 'connection': conn})
207266

267+
def test_batch_w_passed_connection(self):
268+
from gcloud.datastore import dataset as MUT
269+
from gcloud._testing import _Monkey
270+
conn = object()
271+
dataset = self._makeOne()
272+
273+
with _Monkey(MUT, Batch=_Dummy):
274+
batch = dataset.batch(connection=conn)
275+
276+
self.assertTrue(isinstance(batch, _Dummy))
277+
self.assertEqual(batch.args, ())
278+
self.assertEqual(batch.kwargs,
279+
{'dataset_id': self.DATASET_ID, 'connection': conn})
280+
208281
def test_transaction_wo_connection(self):
209282
from gcloud.datastore import dataset as MUT
210283
from gcloud._testing import _Monkey
@@ -218,7 +291,7 @@ def test_transaction_wo_connection(self):
218291
self.assertEqual(xact.kwargs,
219292
{'dataset_id': self.DATASET_ID, 'connection': None})
220293

221-
def test_transaction_w_connection(self):
294+
def test_transaction_w_connection_attr(self):
222295
from gcloud.datastore import dataset as MUT
223296
from gcloud._testing import _Monkey
224297
conn = object()
@@ -232,6 +305,20 @@ def test_transaction_w_connection(self):
232305
self.assertEqual(xact.kwargs,
233306
{'dataset_id': self.DATASET_ID, 'connection': conn})
234307

308+
def test_transaction_w_passed_connection(self):
309+
from gcloud.datastore import dataset as MUT
310+
from gcloud._testing import _Monkey
311+
conn = object()
312+
dataset = self._makeOne()
313+
314+
with _Monkey(MUT, Transaction=_Dummy):
315+
xact = dataset.transaction(connection=conn)
316+
317+
self.assertTrue(isinstance(xact, _Dummy))
318+
self.assertEqual(xact.args, ())
319+
self.assertEqual(xact.kwargs,
320+
{'dataset_id': self.DATASET_ID, 'connection': conn})
321+
235322
def test_query_w_dataset_id(self):
236323
KIND = 'KIND'
237324
dataset = self._makeOne()

0 commit comments

Comments
 (0)