Closes googleapis#4278 Datastore: Transaction Options (googleapis#4357)

chemelnucfin · web-flow · commit f4bdb330b4e0 · 2017-12-11T09:51:16.000-08:00
diff --git a/datastore/google/cloud/datastore/client.py b/datastore/google/cloud/datastore/client.py
@@ -511,9 +511,13 @@ def batch(self):
         """Proxy to :class:`google.cloud.datastore.batch.Batch`."""
         return Batch(self)
 
-    def transaction(self):
-        """Proxy to :class:`google.cloud.datastore.transaction.Transaction`."""
-        return Transaction(self)
+    def transaction(self, **kwargs):
+        """Proxy to :class:`google.cloud.datastore.transaction.Transaction`.
+
+        :type kwargs: dict
+        :param kwargs: Keyword arguments to be passed in.
+        """
+        return Transaction(self, **kwargs)
 
     def query(self, **kwargs):
         """Proxy to :class:`google.cloud.datastore.query.Query`.
diff --git a/datastore/google/cloud/datastore/transaction.py b/datastore/google/cloud/datastore/transaction.py
@@ -15,6 +15,7 @@
 """Create / interact with Google Cloud Datastore transactions."""
 
 from google.cloud.datastore.batch import Batch
+from google.cloud.datastore_v1.types import TransactionOptions
 
 
 class Transaction(Batch):
@@ -152,13 +153,22 @@ def Entity(*args, **kwargs):
 
     :type client: :class:`google.cloud.datastore.client.Client`
     :param client: the client used to connect to datastore.
+
+    :type read_only: bool
+    :param read_only: indicates the transaction is read only.
     """
 
     _status = None
 
-    def __init__(self, client):
+    def __init__(self, client, read_only=False):
         super(Transaction, self).__init__(client)
         self._id = None
+        if read_only:
+            options = TransactionOptions(
+                read_only=TransactionOptions.ReadOnly())
+        else:
+            options = TransactionOptions()
+        self._options = options
 
     @property
     def id(self):
@@ -234,3 +244,21 @@ def commit(self):
         finally:
             # Clear our own ID in case this gets accidentally reused.
             self._id = None
+
+    def put(self, entity):
+        """Adds an entity to be committed.
+
+        Ensures the transaction is not marked readonly.
+        Please see documentation at
+        :meth:`~google.cloud.datastore.batch.Batch.put`
+
+        :type entity: :class:`~google.cloud.datastore.entity.Entity`
+        :param entity: the entity to be saved.
+
+        :raises: :class:`RuntimeError` if the transaction
+                 is marked ReadOnly
+        """
+        if self._options.HasField('read_only'):
+            raise RuntimeError("Transaction is read only")
+        else:
+            super(Transaction, self).put(entity)
diff --git a/datastore/tests/system/test_system.py b/datastore/tests/system/test_system.py
@@ -22,6 +22,7 @@
 from google.cloud._helpers import UTC
 from google.cloud import datastore
 from google.cloud.datastore.helpers import GeoPoint
+from google.cloud.datastore_v1 import types
 from google.cloud.environment_vars import GCD_DATASET
 from google.cloud.exceptions import Conflict
 
diff --git a/datastore/tests/unit/test_client.py b/datastore/tests/unit/test_client.py
@@ -934,6 +934,22 @@ def test_transaction_defaults(self):
             self.assertIs(xact, mock_klass.return_value)
             mock_klass.assert_called_once_with(client)
 
+    def test_read_only_transaction_defaults(self):
+        from google.cloud.datastore.transaction import Transaction
+        from google.cloud.datastore_v1.types import TransactionOptions
+        creds = _make_credentials()
+        client = self._make_one(credentials=creds)
+        xact = client.transaction(read_only=True)
+        self.assertEqual(xact._options,
+                         TransactionOptions(
+                             read_only=TransactionOptions.ReadOnly()
+                         )
+        )
+        self.assertFalse(xact._options.HasField("read_write"))
+        self.assertTrue(xact._options.HasField("read_only"))
+        self.assertEqual(xact._options.read_only,
+                         TransactionOptions.ReadOnly())
+
     def test_query_w_client(self):
         KIND = 'KIND'
 
diff --git a/datastore/tests/unit/test_transaction.py b/datastore/tests/unit/test_transaction.py
@@ -22,12 +22,18 @@ class TestTransaction(unittest.TestCase):
     @staticmethod
     def _get_target_class():
         from google.cloud.datastore.transaction import Transaction
-
         return Transaction
 
+    def _get_options_class(self, **kw):
+        from google.cloud.datastore_v1.types import TransactionOptions
+        return TransactionOptions
+
     def _make_one(self, client, **kw):
         return self._get_target_class()(client, **kw)
 
+    def _make_options(self, **kw):
+        return self._get_options_class()(**kw)
+
     def test_ctor_defaults(self):
         project = 'PROJECT'
         client = _Client(project)
@@ -212,6 +218,27 @@ class Foo(Exception):
         self.assertIsNone(xact.id)
         self.assertEqual(ds_api.begin_transaction.call_count, 1)
 
+    def test_constructor_read_only(self):
+        project = 'PROJECT'
+        id_ = 850302
+        ds_api = _make_datastore_api(xact=id_)
+        client = _Client(project, datastore_api=ds_api)
+        read_only = self._get_options_class().ReadOnly()
+        options = self._make_options(read_only=read_only)
+        xact = self._make_one(client, read_only=True)
+        self.assertEqual(xact._options, options)
+
+    def test_put_read_only(self):
+        project = 'PROJECT'
+        id_ = 943243
+        ds_api = _make_datastore_api(xact_id=id_)
+        client = _Client(project, datastore_api=ds_api)
+        entity = _Entity()
+        xact = self._make_one(client, read_only=True)
+        xact.begin()
+        with self.assertRaises(RuntimeError):
+            xact.put(entity)
+
 
 def _make_key(kind, id_, project):
     from google.cloud.datastore_v1.proto import entity_pb2
