Add support for serialized/transactional read in 'Client.get{,_multi}'.

gcloud/datastore/client.py

@@ -223,7 +223,7 @@ def current_transaction(self):
         if isinstance(transaction, Transaction):
             return transaction
 
-    def get(self, key, missing=None, deferred=None):
+    def get(self, key, missing=None, deferred=None, transaction=None):
         """Retrieve an entity from a single key (if it exists).
 
         .. note::
@@ -244,15 +244,19 @@ def get(self, key, missing=None, deferred=None):
         :param deferred: (Optional) If a list is passed, the keys returned
                          by the backend as "deferred" will be copied into it.
 
+        :type transaction: :class:`gcloud.datastore.transaction.Transaction`
+        :param transaction: (Optional) Transaction to use for read consistency.
+                            If not passed, uses current transaction, if set.
+
         :rtype: :class:`gcloud.datastore.entity.Entity` or ``NoneType``
         :returns: The requested entity if it exists.
         """
         entities = self.get_multi(keys=[key], missing=missing,
-                                  deferred=deferred)
+                                  deferred=deferred, transaction=transaction)
         if entities:
             return entities[0]
 
-    def get_multi(self, keys, missing=None, deferred=None):
+    def get_multi(self, keys, missing=None, deferred=None, transaction=None):
         """Retrieve entities, along with their attributes.
 
         :type keys: list of :class:`gcloud.datastore.key.Key`
@@ -268,6 +272,10 @@ def get_multi(self, keys, missing=None, deferred=None):
                          by the backend as "deferred" will be copied into it.
                          If the list is not empty, an error will occur.
 
+        :type transaction: :class:`gcloud.datastore.transaction.Transaction`
+        :param transaction: (Optional) Transaction to use for read consistency.
+                            If not passed, uses current transaction, if set.
+
         :rtype: list of :class:`gcloud.datastore.entity.Entity`
         :returns: The requested entities.
         :raises: :class:`ValueError` if one or more of ``keys`` has a project
@@ -281,7 +289,8 @@ def get_multi(self, keys, missing=None, deferred=None):
             if current_id != self.project:
                 raise ValueError('Keys do not match project')
 
-        transaction = self.current_transaction
+        if transaction is None:
+            transaction = self.current_transaction
 
         entity_pbs = _extended_lookup(
             connection=self.connection,

gcloud/datastore/test_client.py

@@ -231,8 +231,10 @@ def _get_multi(*args, **kw):
         self.assertEqual(_called_with[0][1]['keys'], [key])
         self.assertTrue(_called_with[0][1]['missing'] is None)
         self.assertTrue(_called_with[0][1]['deferred'] is None)
+        self.assertTrue(_called_with[0][1]['transaction'] is None)
 
     def test_get_hit(self):
+        TXN_ID = '123'
         _called_with = []
         _entity = object()
 
@@ -246,12 +248,13 @@ def _get_multi(*args, **kw):
 
         key, missing, deferred = object(), [], []
 
-        self.assertTrue(client.get(key, missing, deferred) is _entity)
+        self.assertTrue(client.get(key, missing, deferred, TXN_ID) is _entity)
 
         self.assertEqual(_called_with[0][0], ())
         self.assertEqual(_called_with[0][1]['keys'], [key])
         self.assertTrue(_called_with[0][1]['missing'] is missing)
         self.assertTrue(_called_with[0][1]['deferred'] is deferred)
+        self.assertEqual(_called_with[0][1]['transaction'], TXN_ID)
 
     def test_get_multi_no_keys(self):
         creds = object()
@@ -412,6 +415,40 @@ def test_get_multi_hit(self):
         self.assertEqual(list(result), ['foo'])
         self.assertEqual(result['foo'], 'Foo')
 
+    def test_get_multi_hit_w_transaction(self):
+        from gcloud.datastore.key import Key
+
+        TXN_ID = '123'
+        KIND = 'Kind'
+        ID = 1234
+        PATH = [{'kind': KIND, 'id': ID}]
+
+        # Make a found entity pb to be returned from mock backend.
+        entity_pb = _make_entity_pb(self.PROJECT, KIND, ID, 'foo', 'Foo')
+
+        # Make a connection to return the entity pb.
+        creds = object()
+        client = self._makeOne(credentials=creds)
+        client.connection._add_lookup_result([entity_pb])
+
+        key = Key(KIND, ID, project=self.PROJECT)
+        txn = client.transaction()
+        txn._id = TXN_ID
+        result, = client.get_multi([key], transaction=txn)
+        new_key = result.key
+
+        # Check the returned value is as expected.
+        self.assertFalse(new_key is key)
+        self.assertEqual(new_key.project, self.PROJECT)
+        self.assertEqual(new_key.path, PATH)
+        self.assertEqual(list(result), ['foo'])
+        self.assertEqual(result['foo'], 'Foo')
+
+        cw = client.connection._lookup_cw
+        self.assertEqual(len(cw), 1)
+        _, _, _, transaction_id = cw[0]
+        self.assertEqual(transaction_id, TXN_ID)
+
     def test_get_multi_hit_multiple_keys_same_project(self):
         from gcloud.datastore.key import Key
 

system_tests/datastore.py

@@ -417,7 +417,7 @@ def test_query_distinct_on(self):
 
 class TestDatastoreTransaction(TestDatastore):
 
-    def test_transaction(self):
+    def test_transaction_via_with_statement(self):
         entity = datastore.Entity(key=Config.CLIENT.key('Company', 'Google'))
         entity['url'] = u'www.google.com'
 
@@ -432,6 +432,43 @@ def test_transaction(self):
         self.case_entities_to_delete.append(retrieved_entity)
         self.assertEqual(retrieved_entity, entity)
 
+    def test_transaction_via_explicit_begin_get_commit(self):
+        # See https://github.com/GoogleCloudPlatform/gcloud-python/issues/1859
+        # Note that this example lacks the threading which provokes the race
+        # condition in that issue:  we are basically just exercising the
+        # "explict" path for using transactions.
+        BEFORE_1 = 100
+        BEFORE_2 = 0
+        TRANSFER_AMOUNT = 40
+        key1 = Config.CLIENT.key('account', '123')
+        account1 = datastore.Entity(key=key1)
+        account1['balance'] = BEFORE_1
+        key2 = Config.CLIENT.key('account', '234')
+        account2 = datastore.Entity(key=key2)
+        account2['balance'] = BEFORE_2
+        Config.CLIENT.put_multi([account1, account2])
+        self.case_entities_to_delete.append(account1)
+        self.case_entities_to_delete.append(account2)
+
+        def transfer_funds(client, from_key, to_key, amount):
+            xact = client.transaction()
+            xact.begin()
+            from_account = client.get(from_key, transaction=xact)
+            to_account = client.get(to_key, transaction=xact)
+            from_account['balance'] -= amount
+            to_account['balance'] += amount
+
+            xact.put(from_account)
+            xact.put(to_account)
+            xact.commit()
+
+        transfer_funds(Config.CLIENT, key1, key2, TRANSFER_AMOUNT)
+
+        after1 = Config.CLIENT.get(key1)
+        after2 = Config.CLIENT.get(key2)
+        self.assertEqual(after1['balance'], BEFORE_1 - TRANSFER_AMOUNT)
+        self.assertEqual(after2['balance'], BEFORE_2 + TRANSFER_AMOUNT)
+
     def test_failure_with_contention(self):
         contention_prop_name = 'baz'
         local_client = clone_client(Config.CLIENT)