Skip to content

Commit 924b10b

Browse files
tseavertseaver
authored
fix: pass transaction's options to API in 'begin' (#143)
Closes #135.
1 parent 4f90d04 commit 924b10b

File tree

2 files changed

+96
-47
lines changed

2 files changed

+96
-47
lines changed

‎google/cloud/datastore/transaction.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,10 +176,12 @@ def Entity(*args, **kwargs):
176176
def __init__(self, client, read_only=False):
177177
super(Transaction, self).__init__(client)
178178
self._id = None
179+
179180
if read_only:
180181
options = TransactionOptions(read_only=TransactionOptions.ReadOnly())
181182
else:
182183
options = TransactionOptions()
184+
183185
self._options = options
184186

185187
@property
@@ -231,9 +233,13 @@ def begin(self, retry=None, timeout=None):
231233

232234
kwargs = _make_retry_timeout_kwargs(retry, timeout)
233235

236+
request = {
237+
"project_id": self.project,
238+
"transaction_options": self._options,
239+
}
234240
try:
235241
response_pb = self._client._datastore_api.begin_transaction(
236-
request={"project_id": self.project}, **kwargs
242+
request=request, **kwargs
237243
)
238244
self._id = response_pb.transaction
239245
except: # noqa: E722 do not use bare except, specify exception instead

‎tests/unit/test_transaction.py‎

Lines changed: 89 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -24,28 +24,50 @@ def _get_target_class():
2424

2525
return Transaction
2626

27-
def _get_options_class(self, **kw):
27+
def _make_one(self, client, **kw):
28+
return self._get_target_class()(client, **kw)
29+
30+
def _make_options(self, read_only=False, previous_transaction=None):
2831
from google.cloud.datastore_v1.types import TransactionOptions
2932

30-
return TransactionOptions
33+
kw = {}
3134

32-
def _make_one(self, client, **kw):
33-
return self._get_target_class()(client, **kw)
35+
if read_only:
36+
kw["read_only"] = TransactionOptions.ReadOnly()
3437

35-
def _make_options(self, **kw):
36-
return self._get_options_class()(**kw)
38+
return TransactionOptions(**kw)
3739

3840
def test_ctor_defaults(self):
3941
project = "PROJECT"
4042
client = _Client(project)
43+
4144
xact = self._make_one(client)
45+
4246
self.assertEqual(xact.project, project)
4347
self.assertIs(xact._client, client)
4448
self.assertIsNone(xact.id)
4549
self.assertEqual(xact._status, self._get_target_class()._INITIAL)
4650
self.assertEqual(xact._mutations, [])
4751
self.assertEqual(len(xact._partial_key_entities), 0)
4852

53+
def test_constructor_read_only(self):
54+
project = "PROJECT"
55+
id_ = 850302
56+
ds_api = _make_datastore_api(xact=id_)
57+
client = _Client(project, datastore_api=ds_api)
58+
options = self._make_options(read_only=True)
59+
60+
xact = self._make_one(client, read_only=True)
61+
62+
self.assertEqual(xact._options, options)
63+
64+
def _make_begin_request(self, project, read_only=False):
65+
expected_options = self._make_options(read_only=read_only)
66+
return {
67+
"project_id": project,
68+
"transaction_options": expected_options,
69+
}
70+
4971
def test_current(self):
5072
from google.cloud.datastore_v1.types import datastore as datastore_pb2
5173

@@ -57,24 +79,34 @@ def test_current(self):
5779
xact2 = self._make_one(client)
5880
self.assertIsNone(xact1.current())
5981
self.assertIsNone(xact2.current())
82+
6083
with xact1:
6184
self.assertIs(xact1.current(), xact1)
6285
self.assertIs(xact2.current(), xact1)
86+
6387
with _NoCommitBatch(client):
6488
self.assertIsNone(xact1.current())
6589
self.assertIsNone(xact2.current())
90+
6691
with xact2:
6792
self.assertIs(xact1.current(), xact2)
6893
self.assertIs(xact2.current(), xact2)
94+
6995
with _NoCommitBatch(client):
7096
self.assertIsNone(xact1.current())
7197
self.assertIsNone(xact2.current())
98+
7299
self.assertIs(xact1.current(), xact1)
73100
self.assertIs(xact2.current(), xact1)
101+
74102
self.assertIsNone(xact1.current())
75103
self.assertIsNone(xact2.current())
76104

77-
ds_api.rollback.assert_not_called()
105+
begin_txn = ds_api.begin_transaction
106+
self.assertEqual(begin_txn.call_count, 2)
107+
expected_request = self._make_begin_request(project)
108+
begin_txn.assert_called_with(request=expected_request)
109+
78110
commit_method = ds_api.commit
79111
self.assertEqual(commit_method.call_count, 2)
80112
mode = datastore_pb2.CommitRequest.Mode.TRANSACTIONAL
@@ -87,21 +119,35 @@ def test_current(self):
87119
}
88120
)
89121

90-
begin_txn = ds_api.begin_transaction
91-
self.assertEqual(begin_txn.call_count, 2)
92-
begin_txn.assert_called_with(request={"project_id": project})
122+
ds_api.rollback.assert_not_called()
93123

94124
def test_begin(self):
95125
project = "PROJECT"
96126
id_ = 889
97127
ds_api = _make_datastore_api(xact_id=id_)
98128
client = _Client(project, datastore_api=ds_api)
99129
xact = self._make_one(client)
130+
100131
xact.begin()
132+
101133
self.assertEqual(xact.id, id_)
102-
ds_api.begin_transaction.assert_called_once_with(
103-
request={"project_id": project}
104-
)
134+
135+
expected_request = self._make_begin_request(project)
136+
ds_api.begin_transaction.assert_called_once_with(request=expected_request)
137+
138+
def test_begin_w_readonly(self):
139+
project = "PROJECT"
140+
id_ = 889
141+
ds_api = _make_datastore_api(xact_id=id_)
142+
client = _Client(project, datastore_api=ds_api)
143+
xact = self._make_one(client, read_only=True)
144+
145+
xact.begin()
146+
147+
self.assertEqual(xact.id, id_)
148+
149+
expected_request = self._make_begin_request(project, read_only=True)
150+
ds_api.begin_transaction.assert_called_once_with(request=expected_request)
105151

106152
def test_begin_w_retry_w_timeout(self):
107153
project = "PROJECT"
@@ -116,8 +162,10 @@ def test_begin_w_retry_w_timeout(self):
116162
xact.begin(retry=retry, timeout=timeout)
117163

118164
self.assertEqual(xact.id, id_)
165+
166+
expected_request = self._make_begin_request(project)
119167
ds_api.begin_transaction.assert_called_once_with(
120-
request={"project_id": project}, retry=retry, timeout=timeout
168+
request=expected_request, retry=retry, timeout=timeout,
121169
)
122170

123171
def test_begin_tombstoned(self):
@@ -126,19 +174,23 @@ def test_begin_tombstoned(self):
126174
ds_api = _make_datastore_api(xact_id=id_)
127175
client = _Client(project, datastore_api=ds_api)
128176
xact = self._make_one(client)
177+
129178
xact.begin()
179+
130180
self.assertEqual(xact.id, id_)
131-
ds_api.begin_transaction.assert_called_once_with(
132-
request={"project_id": project}
133-
)
181+
182+
expected_request = self._make_begin_request(project)
183+
ds_api.begin_transaction.assert_called_once_with(request=expected_request)
134184

135185
xact.rollback()
186+
136187
client._datastore_api.rollback.assert_called_once_with(
137188
request={"project_id": project, "transaction": id_}
138189
)
139190
self.assertIsNone(xact.id)
140191

141-
self.assertRaises(ValueError, xact.begin)
192+
with self.assertRaises(ValueError):
193+
xact.begin()
142194

143195
def test_begin_w_begin_transaction_failure(self):
144196
project = "PROJECT"
@@ -152,9 +204,9 @@ def test_begin_w_begin_transaction_failure(self):
152204
xact.begin()
153205

154206
self.assertIsNone(xact.id)
155-
ds_api.begin_transaction.assert_called_once_with(
156-
request={"project_id": project}
157-
)
207+
208+
expected_request = self._make_begin_request(project)
209+
ds_api.begin_transaction.assert_called_once_with(request=expected_request)
158210

159211
def test_rollback(self):
160212
project = "PROJECT"
@@ -256,11 +308,14 @@ def test_context_manager_no_raise(self):
256308
ds_api = _make_datastore_api(xact_id=id_)
257309
client = _Client(project, datastore_api=ds_api)
258310
xact = self._make_one(client)
311+
259312
with xact:
260-
self.assertEqual(xact.id, id_)
261-
ds_api.begin_transaction.assert_called_once_with(
262-
request={"project_id": project}
263-
)
313+
self.assertEqual(xact.id, id_) # only set between begin / commit
314+
315+
self.assertIsNone(xact.id)
316+
317+
expected_request = self._make_begin_request(project)
318+
ds_api.begin_transaction.assert_called_once_with(request=expected_request)
264319

265320
mode = datastore_pb2.CommitRequest.Mode.TRANSACTIONAL
266321
client._datastore_api.commit.assert_called_once_with(
@@ -272,9 +327,6 @@ def test_context_manager_no_raise(self):
272327
},
273328
)
274329

275-
self.assertIsNone(xact.id)
276-
self.assertEqual(ds_api.begin_transaction.call_count, 1)
277-
278330
def test_context_manager_w_raise(self):
279331
class Foo(Exception):
280332
pass
@@ -288,29 +340,20 @@ class Foo(Exception):
288340
try:
289341
with xact:
290342
self.assertEqual(xact.id, id_)
291-
ds_api.begin_transaction.assert_called_once_with(
292-
request={"project_id": project}
293-
)
294343
raise Foo()
295344
except Foo:
296-
self.assertIsNone(xact.id)
297-
client._datastore_api.rollback.assert_called_once_with(
298-
request={"project_id": project, "transaction": id_}
299-
)
345+
pass
300346

301-
client._datastore_api.commit.assert_not_called()
302347
self.assertIsNone(xact.id)
303-
self.assertEqual(ds_api.begin_transaction.call_count, 1)
304348

305-
def test_constructor_read_only(self):
306-
project = "PROJECT"
307-
id_ = 850302
308-
ds_api = _make_datastore_api(xact=id_)
309-
client = _Client(project, datastore_api=ds_api)
310-
read_only = self._get_options_class().ReadOnly()
311-
options = self._make_options(read_only=read_only)
312-
xact = self._make_one(client, read_only=True)
313-
self.assertEqual(xact._options, options)
349+
expected_request = self._make_begin_request(project)
350+
ds_api.begin_transaction.assert_called_once_with(request=expected_request)
351+
352+
client._datastore_api.commit.assert_not_called()
353+
354+
client._datastore_api.rollback.assert_called_once_with(
355+
request={"project_id": project, "transaction": id_}
356+
)
314357

315358
def test_put_read_only(self):
316359
project = "PROJECT"

0 commit comments

Comments
 (0)