Skip to content

Commit 3c76dbf

Browse files
tseavertseaver
committed
Merge pull request #978 from tseaver/944-batch_transaction_query_hold_client
Update batch/transaction/query to hold client.
2 parents 681105d + d556b63 commit 3c76dbf

File tree

8 files changed

+337
-419
lines changed

8 files changed

+337
-419
lines changed

‎gcloud/datastore/batch.py‎

Lines changed: 22 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,11 @@
2121
https://cloud.google.com/datastore/docs/concepts/entities#Datastore_Batch_operations
2222
"""
2323

24-
from gcloud._helpers import _LocalStack
25-
from gcloud.datastore import _implicit_environ
2624
from gcloud.datastore import helpers
2725
from gcloud.datastore.key import _dataset_ids_equal
2826
from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
2927

3028

31-
_BATCHES = _LocalStack()
32-
33-
3429
class Batch(object):
3530
"""An abstraction representing a collected group of updates / deletes.
3631
@@ -62,34 +57,19 @@ class Batch(object):
6257
... do_some_work(batch)
6358
... raise Exception() # rolls back
6459
65-
:type dataset_id: :class:`str`.
66-
:param dataset_id: The ID of the dataset.
67-
68-
:type connection: :class:`gcloud.datastore.connection.Connection`
69-
:param connection: The connection used to connect to datastore.
70-
71-
:raises: :class:`ValueError` if either a connection or dataset ID
72-
are not set.
60+
:type client: :class:`gcloud.datastore.client.Client`
61+
:param client: The client used to connect to datastore.
7362
"""
7463
_id = None # "protected" attribute, always None for non-transactions
7564

76-
def __init__(self, dataset_id=None, connection=None):
77-
self._connection = (connection or
78-
_implicit_environ.get_default_connection())
79-
self._dataset_id = (dataset_id or
80-
_implicit_environ.get_default_dataset_id())
81-
82-
if self._connection is None or self._dataset_id is None:
83-
raise ValueError('A batch must have a connection and '
84-
'a dataset ID set.')
85-
65+
def __init__(self, client):
66+
self._client = client
8667
self._mutation = datastore_pb.Mutation()
8768
self._auto_id_entities = []
8869

89-
@staticmethod
90-
def current():
70+
def current(self):
9171
"""Return the topmost batch / transaction, or None."""
92-
return _BATCHES.top
72+
return self._client.current_batch
9373

9474
@property
9575
def dataset_id(self):
@@ -98,7 +78,16 @@ def dataset_id(self):
9878
:rtype: :class:`str`
9979
:returns: The dataset ID in which the batch will run.
10080
"""
101-
return self._dataset_id
81+
return self._client.dataset_id
82+
83+
@property
84+
def namespace(self):
85+
"""Getter for namespace in which the batch will run.
86+
87+
:rtype: :class:`str`
88+
:returns: The namespace in which the batch will run.
89+
"""
90+
return self._client.namespace
10291

10392
@property
10493
def connection(self):
@@ -107,7 +96,7 @@ def connection(self):
10796
:rtype: :class:`gcloud.datastore.connection.Connection`
10897
:returns: The connection over which the batch will run.
10998
"""
110-
return self._connection
99+
return self._client.connection
111100

112101
@property
113102
def mutation(self):
@@ -172,7 +161,7 @@ def put(self, entity):
172161
if entity.key is None:
173162
raise ValueError("Entity must have a key")
174163

175-
if not _dataset_ids_equal(self._dataset_id, entity.key.dataset_id):
164+
if not _dataset_ids_equal(self.dataset_id, entity.key.dataset_id):
176165
raise ValueError("Key must be from same dataset as batch")
177166

178167
_assign_entity_to_mutation(
@@ -190,7 +179,7 @@ def delete(self, key):
190179
if key.is_partial:
191180
raise ValueError("Key must be complete")
192181

193-
if not _dataset_ids_equal(self._dataset_id, key.dataset_id):
182+
if not _dataset_ids_equal(self.dataset_id, key.dataset_id):
194183
raise ValueError("Key must be from same dataset as batch")
195184

196185
key_pb = helpers._prepare_key_for_request(key.to_protobuf())
@@ -211,7 +200,7 @@ def commit(self):
211200
context manager.
212201
"""
213202
response = self.connection.commit(
214-
self._dataset_id, self.mutation, self._id)
203+
self.dataset_id, self.mutation, self._id)
215204
# If the back-end returns without error, we are guaranteed that
216205
# the response's 'insert_auto_id_key' will match (length and order)
217206
# the request's 'insert_auto_id` entities, which are derived from
@@ -229,7 +218,7 @@ def rollback(self):
229218
pass
230219

231220
def __enter__(self):
232-
_BATCHES.push(self)
221+
self._client._push_batch(self)
233222
self.begin()
234223
return self
235224

@@ -240,7 +229,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
240229
else:
241230
self.rollback()
242231
finally:
243-
_BATCHES.pop()
232+
self._client._pop_batch()
244233

245234

246235
def _assign_entity_to_mutation(mutation_pb, entity, auto_id_entities):

‎gcloud/datastore/client.py‎

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,19 @@ def current_batch(self):
164164
"""
165165
return self._batch_stack.top
166166

167+
@property
168+
def current_transaction(self):
169+
"""Currently-active transaction.
170+
171+
:rtype: :class:`gcloud.datastore.transaction.Transaction`, or an object
172+
implementing its API, or ``NoneType`` (if no transaction is
173+
active).
174+
:returns: The transaction at the toop of the batch stack.
175+
"""
176+
transaction = self.current_batch
177+
if isinstance(transaction, Transaction):
178+
return transaction
179+
167180
def get(self, key, missing=None, deferred=None):
168181
"""Retrieve an entity from a single key (if it exists).
169182
@@ -222,7 +235,7 @@ def get_multi(self, keys, missing=None, deferred=None):
222235
if ids != [self.dataset_id]:
223236
raise ValueError('Keys do not match dataset ID')
224237

225-
transaction = Transaction.current()
238+
transaction = self.current_transaction
226239

227240
entity_pbs = _extended_lookup(
228241
connection=self.connection,
@@ -274,12 +287,12 @@ def put_multi(self, entities):
274287
if not entities:
275288
return
276289

277-
current = Batch.current()
290+
current = self.current_batch
278291
in_batch = current is not None
279292

280293
if not in_batch:
281-
current = Batch(dataset_id=self.dataset_id,
282-
connection=self.connection)
294+
current = self.batch()
295+
283296
for entity in entities:
284297
current.put(entity)
285298

@@ -310,12 +323,12 @@ def delete_multi(self, keys):
310323
return
311324

312325
# We allow partial keys to attempt a delete, the backend will fail.
313-
current = Batch.current()
326+
current = self.current_batch
314327
in_batch = current is not None
315328

316329
if not in_batch:
317-
current = Batch(dataset_id=self.dataset_id,
318-
connection=self.connection)
330+
current = self.batch()
331+
319332
for key in keys:
320333
current.delete(key)
321334

@@ -368,15 +381,14 @@ def batch(self):
368381
369382
Passes our ``dataset_id``.
370383
"""
371-
return Batch(dataset_id=self.dataset_id, connection=self.connection)
384+
return Batch(self)
372385

373386
def transaction(self):
374387
"""Proxy to :class:`gcloud.datastore.transaction.Transaction`.
375388
376389
Passes our ``dataset_id``.
377390
"""
378-
return Transaction(dataset_id=self.dataset_id,
379-
connection=self.connection)
391+
return Transaction(self)
380392

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

‎gcloud/datastore/query.py‎

Lines changed: 21 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,9 @@
1717
import base64
1818

1919
from gcloud._helpers import _ensure_tuple_or_list
20-
from gcloud.datastore import _implicit_environ
2120
from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
2221
from gcloud.datastore import helpers
2322
from gcloud.datastore.key import Key
24-
from gcloud.datastore.transaction import Transaction
2523

2624

2725
class Query(object):
@@ -30,15 +28,19 @@ class Query(object):
3028
This class serves as an abstraction for creating a query over data
3129
stored in the Cloud Datastore.
3230
31+
:type client: :class:`gcloud.datastore.client.Client`
32+
:param client: The client used to connect to datastore.
33+
3334
:type kind: string
3435
:param kind: The kind to query.
3536
3637
:type dataset_id: string
3738
:param dataset_id: The ID of the dataset to query. If not passed,
38-
uses the implicit default.
39+
uses the client's value.
3940
4041
:type namespace: string or None
41-
:param namespace: The namespace to which to restrict results.
42+
:param namespace: The namespace to which to restrict results. If not
43+
passed, uses the client's value.
4244
4345
:type ancestor: :class:`gcloud.datastore.key.Key` or None
4446
:param ancestor: key of the ancestor to which this query's results are
@@ -71,6 +73,7 @@ class Query(object):
7173
"""Mapping of operator strings and their protobuf equivalents."""
7274

7375
def __init__(self,
76+
client,
7477
kind=None,
7578
dataset_id=None,
7679
namespace=None,
@@ -80,15 +83,10 @@ def __init__(self,
8083
order=(),
8184
group_by=()):
8285

83-
if dataset_id is None:
84-
dataset_id = _implicit_environ.get_default_dataset_id()
85-
86-
if dataset_id is None:
87-
raise ValueError("No dataset ID supplied, and no default set.")
88-
89-
self._dataset_id = dataset_id
86+
self._client = client
9087
self._kind = kind
91-
self._namespace = namespace
88+
self._dataset_id = dataset_id or client.dataset_id
89+
self._namespace = namespace or client.namespace
9290
self._ancestor = ancestor
9391
self._filters = []
9492
# Verify filters passed in.
@@ -294,7 +292,7 @@ def group_by(self, value):
294292
self._group_by[:] = value
295293

296294
def fetch(self, limit=None, offset=0, start_cursor=None, end_cursor=None,
297-
connection=None):
295+
client=None):
298296
"""Execute the Query; return an iterator for the matching entities.
299297
300298
For example::
@@ -319,22 +317,19 @@ def fetch(self, limit=None, offset=0, start_cursor=None, end_cursor=None,
319317
:type end_cursor: bytes
320318
:param end_cursor: An optional cursor passed through to the iterator.
321319
322-
:type connection: :class:`gcloud.datastore.connection.Connection`
323-
:param connection: An optional cursor passed through to the iterator.
324-
If not supplied, uses the implicit default.
320+
:type client: :class:`gcloud.datastore.client.Client`
321+
:param client: client used to connect to datastore.
322+
If not supplied, uses the query's value.
325323
326324
:rtype: :class:`Iterator`
327325
:raises: ValueError if ``connection`` is not passed and no implicit
328326
default has been set.
329327
"""
330-
if connection is None:
331-
connection = _implicit_environ.get_default_connection()
332-
333-
if connection is None:
334-
raise ValueError("No connection passed, and no default set")
328+
if client is None:
329+
client = self._client
335330

336331
return Iterator(
337-
self, connection, limit, offset, start_cursor, end_cursor)
332+
self, client, limit, offset, start_cursor, end_cursor)
338333

339334

340335
class Iterator(object):
@@ -347,10 +342,10 @@ class Iterator(object):
347342
datastore_pb.QueryResultBatch.MORE_RESULTS_AFTER_LIMIT,
348343
)
349344

350-
def __init__(self, query, connection, limit=None, offset=0,
345+
def __init__(self, query, client, limit=None, offset=0,
351346
start_cursor=None, end_cursor=None):
352347
self._query = query
353-
self._connection = connection
348+
self._client = client
354349
self._limit = limit
355350
self._offset = offset
356351
self._start_cursor = start_cursor
@@ -380,9 +375,9 @@ def next_page(self):
380375

381376
pb.offset = self._offset
382377

383-
transaction = Transaction.current()
378+
transaction = self._client.current_transaction
384379

385-
query_results = self._connection.run_query(
380+
query_results = self._client.connection.run_query(
386381
query_pb=pb,
387382
dataset_id=self._query.dataset_id,
388383
namespace=self._query.namespace,

0 commit comments

Comments
 (0)