Skip to content

Commit 317ddfd

Browse files
tseavertseaver
committed
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).
1 parent c368db3 commit 317ddfd

2 files changed

Lines changed: 35 additions & 26 deletions

File tree

‎gcloud/datastore/dataset.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,31 +25,32 @@
2525
class Dataset(object):
2626
"""Convenience wrapper for invoking APIs/factories w/ a dataset ID."""
2727

28-
def __init__(self, dataset_id):
28+
def __init__(self, dataset_id, connection=None):
2929
if dataset_id is None:
3030
raise ValueError('dataset_id required')
3131
self.dataset_id = dataset_id
32+
self.connection = connection
3233

33-
def get(self, keys, missing=None, deferred=None, connection=None):
34+
def get(self, keys, missing=None, deferred=None):
3435
"""Proxy to :func:`gcloud.datastore.api.get`.
3536
3637
Passes our ``dataset_id``.
3738
"""
38-
return get(keys, missing, deferred, connection, self.dataset_id)
39+
return get(keys, missing, deferred, self.connection, self.dataset_id)
3940

40-
def put(self, entities, connection=None):
41+
def put(self, entities):
4142
"""Proxy to :func:`gcloud.datastore.api.put`.
4243
4344
Passes our ``dataset_id``.
4445
"""
45-
return put(entities, connection, dataset_id=self.dataset_id)
46+
return put(entities, self.connection, dataset_id=self.dataset_id)
4647

47-
def delete(self, keys, connection=None):
48+
def delete(self, keys):
4849
"""Proxy to :func:`gcloud.datastore.api.delete`.
4950
5051
Passes our ``dataset_id``.
5152
"""
52-
return delete(keys, connection, dataset_id=self.dataset_id)
53+
return delete(keys, self.connection, dataset_id=self.dataset_id)
5354

5455
def key(self, *path_args, **kwargs):
5556
"""Proxy to :func:`gcloud.datastore.key.Key`.
@@ -62,19 +63,20 @@ def key(self, *path_args, **kwargs):
6263
kwargs['dataset_id'] = self.dataset_id
6364
return Key(*path_args, **kwargs)
6465

65-
def batch(self, connection=None):
66+
def batch(self):
6667
"""Proxy to :func:`gcloud.datastore.batch.Batch`.
6768
6869
Passes our ``dataset_id``.
6970
"""
70-
return Batch(dataset_id=self.dataset_id, connection=connection)
71+
return Batch(dataset_id=self.dataset_id, connection=self.connection)
7172

72-
def transaction(self, connection=None):
73+
def transaction(self):
7374
"""Proxy to :func:`gcloud.datastore.transaction.Transaction`.
7475
7576
Passes our ``dataset_id``.
7677
"""
77-
return Transaction(dataset_id=self.dataset_id, connection=connection)
78+
return Transaction(dataset_id=self.dataset_id,
79+
connection=self.connection)
7880

7981
def query(self,
8082
kind=None,

‎gcloud/datastore/test_dataset.py‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -23,16 +23,22 @@ def _getTargetClass(self):
2323
from gcloud.datastore.dataset import Dataset
2424
return Dataset
2525

26-
def _makeOne(self, dataset_id=DATASET_ID):
27-
return self._getTargetClass()(dataset_id=dataset_id)
26+
def _makeOne(self, dataset_id=DATASET_ID, connection=None):
27+
return self._getTargetClass()(dataset_id, connection)
2828

29-
def test_ctor_w_None(self):
29+
def test_ctor_w_dataset_id_None(self):
3030
self.assertRaises(ValueError, self._makeOne, None)
3131

32-
def test_ctor_w_dataset_id(self):
32+
def test_ctor_w_dataset_id_no_connection(self):
3333
dataset = self._makeOne()
3434
self.assertEqual(dataset.dataset_id, self.DATASET_ID)
3535

36+
def test_ctor_w_dataset_id_w_connection(self):
37+
conn = object()
38+
dataset = self._makeOne(connection=conn)
39+
self.assertEqual(dataset.dataset_id, self.DATASET_ID)
40+
self.assertTrue(dataset.connection is conn)
41+
3642
def test_get_defaults(self):
3743
from gcloud.datastore import dataset as MUT
3844
from gcloud._testing import _Monkey
@@ -60,11 +66,12 @@ def test_get_explicit(self):
6066
def _get(*args, **kw):
6167
_called_with.append((args, kw))
6268

63-
dataset = self._makeOne()
64-
key, missing, deferred, conn = object(), [], [], object()
69+
conn = object()
70+
dataset = self._makeOne(connection=conn)
71+
key, missing, deferred = object(), [], []
6572

6673
with _Monkey(MUT, get=_get):
67-
dataset.get([key], missing, deferred, conn)
74+
dataset.get([key], missing, deferred)
6875

6976
args = ([key], missing, deferred, conn, self.DATASET_ID)
7077
self.assertEqual(_called_with, [(args, {})])
@@ -96,11 +103,11 @@ def test_put_w_connection(self):
96103
def _put(*args, **kw):
97104
_called_with.append((args, kw))
98105

99-
dataset = self._makeOne()
100106
entity, conn = object(), object()
107+
dataset = self._makeOne(connection=conn)
101108

102109
with _Monkey(MUT, put=_put):
103-
dataset.put([entity], conn)
110+
dataset.put([entity])
104111

105112
self.assertEqual(_called_with,
106113
[(([entity], conn), {'dataset_id': self.DATASET_ID})])
@@ -132,10 +139,10 @@ def test_delete_w_connection(self):
132139
def _delete(*args, **kw):
133140
_called_with.append((args, kw))
134141

135-
dataset = self._makeOne()
136142
key, conn = object(), object()
143+
dataset = self._makeOne(connection=conn)
137144
with _Monkey(MUT, delete=_delete):
138-
dataset.delete([key], conn)
145+
dataset.delete([key])
139146

140147
self.assertEqual(_called_with,
141148
[(([key], conn), {'dataset_id': self.DATASET_ID})])
@@ -191,11 +198,11 @@ def test_batch_wo_connection(self):
191198
def test_batch_w_connection(self):
192199
from gcloud.datastore import dataset as MUT
193200
from gcloud._testing import _Monkey
194-
dataset = self._makeOne()
195201
conn = object()
202+
dataset = self._makeOne(connection=conn)
196203

197204
with _Monkey(MUT, Batch=_Dummy):
198-
batch = dataset.batch(conn)
205+
batch = dataset.batch()
199206

200207
self.assertTrue(isinstance(batch, _Dummy))
201208
self.assertEqual(batch.args, ())
@@ -218,11 +225,11 @@ def test_transaction_wo_connection(self):
218225
def test_transaction_w_connection(self):
219226
from gcloud.datastore import dataset as MUT
220227
from gcloud._testing import _Monkey
221-
dataset = self._makeOne()
222228
conn = object()
229+
dataset = self._makeOne(connection=conn)
223230

224231
with _Monkey(MUT, Transaction=_Dummy):
225-
xact = dataset.transaction(conn)
232+
xact = dataset.transaction()
226233

227234
self.assertTrue(isinstance(xact, _Dummy))
228235
self.assertEqual(xact.args, ())

0 commit comments

Comments
 (0)