Skip to content

Commit f8e40d7

Browse files
tseavertseaver
committed
Move threadlocal stacks of batches to relevant connection module.
1 parent 6cf0852 commit f8e40d7

File tree

10 files changed

+71
-55
lines changed

10 files changed

+71
-55
lines changed

‎gcloud/datastore/api.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
from gcloud.datastore import _implicit_environ
2222
from gcloud.datastore.batch import Batch
23+
from gcloud.datastore.connection import _CONNECTIONS
2324
from gcloud.datastore.entity import Entity
2425
from gcloud.datastore.transaction import Transaction
2526
from gcloud.datastore import helpers
@@ -53,7 +54,7 @@ def _require_dataset_id(dataset_id=None, first_key=None):
5354
"""
5455
if dataset_id is not None:
5556
return dataset_id
56-
top = Batch.current()
57+
top = _CONNECTIONS.top
5758
if top is not None:
5859
return top.dataset_id
5960
if first_key is not None:
@@ -77,7 +78,7 @@ def _require_connection(connection=None):
7778
cannot be inferred from the environment.
7879
"""
7980
if connection is None:
80-
top = Batch.current()
81+
top = _CONNECTIONS.top
8182
if top is not None:
8283
connection = top.connection
8384
else:
@@ -262,7 +263,7 @@ def put(entities, connection=None, dataset_id=None):
262263
connection = _require_connection(connection)
263264
dataset_id = _require_dataset_id(dataset_id, entities[0].key)
264265

265-
current = Batch.current()
266+
current = _CONNECTIONS.top
266267
in_batch = current is not None
267268
if not in_batch:
268269
current = Batch(dataset_id=dataset_id, connection=connection)
@@ -298,7 +299,7 @@ def delete(keys, connection=None, dataset_id=None):
298299
dataset_id = _require_dataset_id(dataset_id, keys[0])
299300

300301
# We allow partial keys to attempt a delete, the backend will fail.
301-
current = Batch.current()
302+
current = _CONNECTIONS.top
302303
in_batch = current is not None
303304
if not in_batch:
304305
current = Batch(dataset_id=dataset_id, connection=connection)

‎gcloud/datastore/batch.py‎

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

24-
from gcloud._helpers import _LocalStack
2524
from gcloud.datastore import _implicit_environ
2625
from gcloud.datastore import helpers
26+
from gcloud.datastore.connection import _CONNECTIONS
2727
from gcloud.datastore.key import _dataset_ids_equal
2828
from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
2929

3030

31-
_BATCHES = _LocalStack()
32-
33-
3431
class Batch(object):
3532
"""An abstraction representing a collected group of updates / deletes.
3633
@@ -90,7 +87,7 @@ def __init__(self, dataset_id=None, connection=None):
9087
@staticmethod
9188
def current():
9289
"""Return the topmost batch / transaction, or None."""
93-
return _BATCHES.top
90+
return _CONNECTIONS.top
9491

9592
@property
9693
def dataset_id(self):
@@ -229,7 +226,7 @@ def rollback(self):
229226
pass
230227

231228
def __enter__(self):
232-
_BATCHES.push(self)
229+
_CONNECTIONS.push(self)
233230
self.begin()
234231
return self
235232

@@ -240,7 +237,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
240237
else:
241238
self.rollback()
242239
finally:
243-
_BATCHES.pop()
240+
_CONNECTIONS.pop()
244241

245242

246243
def _assign_entity_to_mutation(mutation_pb, entity, auto_id_entities):

‎gcloud/datastore/connection.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,14 @@
1717
import os
1818

1919
from gcloud import connection
20+
from gcloud._helpers import _LocalStack
2021
from gcloud.exceptions import make_exception
2122
from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
2223

2324

25+
_CONNECTIONS = _LocalStack()
26+
27+
2428
SCOPE = ('https://www.googleapis.com/auth/datastore',
2529
'https://www.googleapis.com/auth/userinfo.email')
2630
"""The scopes required for authenticating as a Cloud Datastore consumer."""

‎gcloud/datastore/test_api.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -925,13 +925,13 @@ def __init__(self, dataset_id, connection):
925925
self._batch = Batch(dataset_id, connection)
926926

927927
def __enter__(self):
928-
from gcloud.datastore.batch import _BATCHES
929-
_BATCHES.push(self._batch)
928+
from gcloud.datastore.connection import _CONNECTIONS
929+
_CONNECTIONS.push(self._batch)
930930
return self._batch
931931

932932
def __exit__(self, *args):
933-
from gcloud.datastore.batch import _BATCHES
934-
_BATCHES.pop()
933+
from gcloud.datastore.connection import _CONNECTIONS
934+
_CONNECTIONS.pop()
935935

936936

937937
class _NoCommitTransaction(object):
@@ -942,13 +942,13 @@ def __init__(self, dataset_id, connection, transaction_id='TRANSACTION'):
942942
xact._id = transaction_id
943943

944944
def __enter__(self):
945-
from gcloud.datastore.batch import _BATCHES
946-
_BATCHES.push(self._transaction)
945+
from gcloud.datastore.connection import _CONNECTIONS
946+
_CONNECTIONS.push(self._transaction)
947947
return self._transaction
948948

949949
def __exit__(self, *args):
950-
from gcloud.datastore.batch import _BATCHES
951-
_BATCHES.pop()
950+
from gcloud.datastore.connection import _CONNECTIONS
951+
_CONNECTIONS.pop()
952952

953953

954954
class _HttpMultiple(object):

‎gcloud/datastore/test_batch.py‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -279,21 +279,21 @@ def test_commit_w_auto_id_entities(self):
279279
self.assertEqual(key._id, _NEW_ID)
280280

281281
def test_as_context_mgr_wo_error(self):
282-
from gcloud.datastore.batch import _BATCHES
282+
from gcloud.datastore.connection import _CONNECTIONS
283283
_DATASET = 'DATASET'
284284
_PROPERTIES = {'foo': 'bar'}
285285
connection = _Connection()
286286
entity = _Entity(_PROPERTIES)
287287
key = entity.key = _Key(_DATASET)
288288

289-
self.assertEqual(list(_BATCHES), [])
289+
self.assertEqual(list(_CONNECTIONS), [])
290290

291291
with self._makeOne(dataset_id=_DATASET,
292292
connection=connection) as batch:
293-
self.assertEqual(list(_BATCHES), [batch])
293+
self.assertEqual(list(_CONNECTIONS), [batch])
294294
batch.put(entity)
295295

296-
self.assertEqual(list(_BATCHES), [])
296+
self.assertEqual(list(_CONNECTIONS), [])
297297

298298
insert_auto_ids = list(batch.mutation.insert_auto_id)
299299
self.assertEqual(len(insert_auto_ids), 0)
@@ -305,7 +305,7 @@ def test_as_context_mgr_wo_error(self):
305305
self.assertEqual(connection._committed, [(_DATASET, batch.mutation)])
306306

307307
def test_as_context_mgr_nested(self):
308-
from gcloud.datastore.batch import _BATCHES
308+
from gcloud.datastore.connection import _CONNECTIONS
309309
_DATASET = 'DATASET'
310310
_PROPERTIES = {'foo': 'bar'}
311311
connection = _Connection()
@@ -314,20 +314,20 @@ def test_as_context_mgr_nested(self):
314314
entity2 = _Entity(_PROPERTIES)
315315
key2 = entity2.key = _Key(_DATASET)
316316

317-
self.assertEqual(list(_BATCHES), [])
317+
self.assertEqual(list(_CONNECTIONS), [])
318318

319319
with self._makeOne(dataset_id=_DATASET,
320320
connection=connection) as batch1:
321-
self.assertEqual(list(_BATCHES), [batch1])
321+
self.assertEqual(list(_CONNECTIONS), [batch1])
322322
batch1.put(entity1)
323323
with self._makeOne(dataset_id=_DATASET,
324324
connection=connection) as batch2:
325-
self.assertEqual(list(_BATCHES), [batch2, batch1])
325+
self.assertEqual(list(_CONNECTIONS), [batch2, batch1])
326326
batch2.put(entity2)
327327

328-
self.assertEqual(list(_BATCHES), [batch1])
328+
self.assertEqual(list(_CONNECTIONS), [batch1])
329329

330-
self.assertEqual(list(_BATCHES), [])
330+
self.assertEqual(list(_CONNECTIONS), [])
331331

332332
insert_auto_ids = list(batch1.mutation.insert_auto_id)
333333
self.assertEqual(len(insert_auto_ids), 0)
@@ -350,25 +350,25 @@ def test_as_context_mgr_nested(self):
350350
(_DATASET, batch1.mutation)])
351351

352352
def test_as_context_mgr_w_error(self):
353-
from gcloud.datastore.batch import _BATCHES
353+
from gcloud.datastore.connection import _CONNECTIONS
354354
_DATASET = 'DATASET'
355355
_PROPERTIES = {'foo': 'bar'}
356356
connection = _Connection()
357357
entity = _Entity(_PROPERTIES)
358358
key = entity.key = _Key(_DATASET)
359359

360-
self.assertEqual(list(_BATCHES), [])
360+
self.assertEqual(list(_CONNECTIONS), [])
361361

362362
try:
363363
with self._makeOne(dataset_id=_DATASET,
364364
connection=connection) as batch:
365-
self.assertEqual(list(_BATCHES), [batch])
365+
self.assertEqual(list(_CONNECTIONS), [batch])
366366
batch.put(entity)
367367
raise ValueError("testing")
368368
except ValueError:
369369
pass
370370

371-
self.assertEqual(list(_BATCHES), [])
371+
self.assertEqual(list(_CONNECTIONS), [])
372372

373373
insert_auto_ids = list(batch.mutation.insert_auto_id)
374374
self.assertEqual(len(insert_auto_ids), 0)

‎gcloud/storage/_helpers.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
import base64
2222

2323
from gcloud.storage._implicit_environ import get_default_connection
24-
from gcloud.storage.batch import Batch
24+
from gcloud.storage.connection import _CONNECTIONS
2525

2626

2727
class _PropertyMixin(object):
@@ -124,9 +124,8 @@ def _require_connection(connection=None):
124124
:raises: :class:`EnvironmentError` if ``connection`` is ``None``, and
125125
cannot be inferred from the environment.
126126
"""
127-
# NOTE: We use current Batch directly since it inherits from Connection.
128127
if connection is None:
129-
connection = Batch.current()
128+
connection = _CONNECTIONS.top
130129

131130
if connection is None:
132131
connection = get_default_connection()

‎gcloud/storage/batch.py‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,10 @@
2626

2727
import six
2828

29-
from gcloud._helpers import _LocalStack
3029
from gcloud.exceptions import make_exception
3130
from gcloud.storage import _implicit_environ
3231
from gcloud.storage.connection import Connection
33-
34-
35-
_BATCHES = _LocalStack()
32+
from gcloud.storage.connection import _CONNECTIONS
3633

3734

3835
class MIMEApplicationHTTP(MIMEApplication):
@@ -245,18 +242,18 @@ def finish(self):
245242
@staticmethod
246243
def current():
247244
"""Return the topmost batch, or None."""
248-
return _BATCHES.top
245+
return _CONNECTIONS.top
249246

250247
def __enter__(self):
251-
_BATCHES.push(self)
248+
_CONNECTIONS.push(self)
252249
return self
253250

254251
def __exit__(self, exc_type, exc_val, exc_tb):
255252
try:
256253
if exc_type is None:
257254
self.finish()
258255
finally:
259-
_BATCHES.pop()
256+
_CONNECTIONS.pop()
260257

261258

262259
def _generate_faux_mime_message(parser, response, content):

‎gcloud/storage/connection.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@
1515
"""Create / interact with gcloud storage connections."""
1616

1717
from gcloud import connection as base_connection
18+
from gcloud._helpers import _LocalStack
19+
20+
21+
_CONNECTIONS = _LocalStack()
1822

1923

2024
SCOPE = ('https://www.googleapis.com/auth/devstorage.full_control',

‎gcloud/storage/test__helpers.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -294,10 +294,10 @@ def __init__(self, connection):
294294
self._connection = connection
295295

296296
def __enter__(self):
297-
from gcloud.storage.batch import _BATCHES
298-
_BATCHES.push(self._connection)
297+
from gcloud.storage.connection import _CONNECTIONS
298+
_CONNECTIONS.push(self._connection)
299299
return self._connection
300300

301301
def __exit__(self, *args):
302-
from gcloud.storage.batch import _BATCHES
303-
_BATCHES.pop()
302+
from gcloud.storage.connection import _CONNECTIONS
303+
_CONNECTIONS.pop()

‎gcloud/storage/test_batch.py‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,20 @@ def test_finish_empty(self):
222222
self.assertRaises(ValueError, batch.finish)
223223
self.assertTrue(connection.http is http)
224224

225+
def test_current(self):
226+
from gcloud.storage.connection import _CONNECTIONS
227+
klass = self._getTargetClass()
228+
http = _HTTP() # no requests expected
229+
connection = _Connection(http=http)
230+
self.assertTrue(klass.current() is None)
231+
batch = self._makeOne(connection)
232+
_CONNECTIONS.push(batch)
233+
try:
234+
self.assertTrue(klass.current() is batch)
235+
finally:
236+
_CONNECTIONS.pop()
237+
self.assertTrue(klass.current() is None)
238+
225239
def _check_subrequest_no_payload(self, chunk, method, url):
226240
lines = chunk.splitlines()
227241
# blank + 2 headers + blank + request + blank + blank
@@ -370,27 +384,27 @@ def test_finish_nonempty_non_multipart_response(self):
370384
self.assertRaises(ValueError, batch.finish)
371385

372386
def test_as_context_mgr_wo_error(self):
373-
from gcloud.storage.batch import _BATCHES
387+
from gcloud.storage.connection import _CONNECTIONS
374388
URL = 'http://example.com/api'
375389
expected = _Response()
376390
expected['content-type'] = 'multipart/mixed; boundary="DEADBEEF="'
377391
http = _HTTP((expected, _THREE_PART_MIME_RESPONSE))
378392
connection = _Connection(http=http)
379393

380-
self.assertEqual(list(_BATCHES), [])
394+
self.assertEqual(list(_CONNECTIONS), [])
381395

382396
target1 = _MockObject()
383397
target2 = _MockObject()
384398
target3 = _MockObject()
385399
with self._makeOne(connection) as batch:
386-
self.assertEqual(list(_BATCHES), [batch])
400+
self.assertEqual(list(_CONNECTIONS), [batch])
387401
batch._make_request('POST', URL, {'foo': 1, 'bar': 2},
388402
target_object=target1)
389403
batch._make_request('PATCH', URL, {'bar': 3},
390404
target_object=target2)
391405
batch._make_request('DELETE', URL, target_object=target3)
392406

393-
self.assertEqual(list(_BATCHES), [])
407+
self.assertEqual(list(_CONNECTIONS), [])
394408
self.assertEqual(len(batch._requests), 3)
395409
self.assertEqual(batch._requests[0][0], 'POST')
396410
self.assertEqual(batch._requests[1][0], 'PATCH')
@@ -404,19 +418,19 @@ def test_as_context_mgr_wo_error(self):
404418

405419
def test_as_context_mgr_w_error(self):
406420
from gcloud.storage.batch import _FutureDict
407-
from gcloud.storage.batch import _BATCHES
421+
from gcloud.storage.connection import _CONNECTIONS
408422
URL = 'http://example.com/api'
409423
http = _HTTP()
410424
connection = _Connection(http=http)
411425

412-
self.assertEqual(list(_BATCHES), [])
426+
self.assertEqual(list(_CONNECTIONS), [])
413427

414428
target1 = _MockObject()
415429
target2 = _MockObject()
416430
target3 = _MockObject()
417431
try:
418432
with self._makeOne(connection) as batch:
419-
self.assertEqual(list(_BATCHES), [batch])
433+
self.assertEqual(list(_CONNECTIONS), [batch])
420434
batch._make_request('POST', URL, {'foo': 1, 'bar': 2},
421435
target_object=target1)
422436
batch._make_request('PATCH', URL, {'bar': 3},
@@ -426,7 +440,7 @@ def test_as_context_mgr_w_error(self):
426440
except ValueError:
427441
pass
428442

429-
self.assertEqual(list(_BATCHES), [])
443+
self.assertEqual(list(_CONNECTIONS), [])
430444
self.assertEqual(len(http._requests), 0)
431445
self.assertEqual(len(batch._requests), 3)
432446
self.assertEqual(batch._target_objects, [target1, target2, target3])

0 commit comments

Comments
 (0)