feat: Inline Begin transction for RW transactions

Author: surbhigarg92

google/cloud/spanner_v1/pool.py

@@ -515,7 +515,6 @@ def begin_pending_transactions(self):
  """Begin all transactions for sessions added to the pool."""
  while not self._pending_sessions.empty():
  session = self._pending_sessions.get()
- session._transaction.begin()
  super(TransactionPingingPool, self).put(session)
 
 

google/cloud/spanner_v1/session.py

@@ -352,9 +352,7 @@ def run_in_transaction(self, func, *args, **kw):
  txn.transaction_tag = transaction_tag
  else:
  txn = self._transaction
- if txn._transaction_id is None:
- txn.begin()
-
+ 
  try:
  attempts += 1
  return_value = func(txn, *args, **kw)

google/cloud/spanner_v1/transaction.py

@@ -61,9 +61,7 @@ def _check_state(self):
  :raises: :exc:`ValueError` if the object's state is invalid for making
  API requests.
  """
- if self._transaction_id is None:
- raise ValueError("Transaction is not begun")
-
+ 
  if self.committed is not None:
  raise ValueError("Transaction is already committed")
 
@@ -78,7 +76,11 @@ def _make_txn_selector(self):
  :returns: a selector configured for read-write transaction semantics.
  """
  self._check_state()
- return TransactionSelector(id=self._transaction_id)
+ 
+ if self._transaction_id is None:
+ return TransactionSelector(begin=TransactionOptions(read_write=TransactionOptions.ReadWrite()))
+ else:
+ return TransactionSelector(id=self._transaction_id)
 
  def begin(self):
  """Begin a transaction on the database.
@@ -111,15 +113,17 @@ def begin(self):
  def rollback(self):
  """Roll back a transaction on the database."""
  self._check_state()
- database = self._session._database
- api = database.spanner_api
- metadata = _metadata_with_prefix(database.name)
- with trace_call("CloudSpanner.Rollback", self._session):
- api.rollback(
- session=self._session.name,
- transaction_id=self._transaction_id,
- metadata=metadata,
- )
+
+ if self._transaction_id is not None:
+ database = self._session._database
+ api = database.spanner_api
+ metadata = _metadata_with_prefix(database.name)
+ with trace_call("CloudSpanner.Rollback", self._session):
+ api.rollback(
+ session=self._session.name,
+ transaction_id=self._transaction_id,
+ metadata=metadata,
+ )
  self.rolled_back = True
  del self._session._transaction
 
@@ -142,6 +146,8 @@ def commit(self, return_commit_stats=False, request_options=None):
  :raises ValueError: if there are no mutations to commit.
  """
  self._check_state()
+ if self._transaction_id is None:
+ self.begin()
 
  database = self._session._database
  api = database.spanner_api
@@ -302,6 +308,10 @@ def execute_update(
  response = api.execute_sql(
  request=request, metadata=metadata, retry=retry, timeout=timeout
  )
+
+ if self._transaction_id is None and response.metadata.transaction is not None:
+ self._transaction_id = response.metadata.transaction.id
+
  return response.stats.row_count_exact
 
  def batch_update(self, statements, request_options=None):
@@ -378,11 +388,15 @@ def batch_update(self, statements, request_options=None):
  row_counts = [
  result_set.stats.row_count_exact for result_set in response.result_sets
  ]
+
+ for result_set in response.result_sets:
+ if self._transaction_id is None and result_set.metadata.transaction is not None:
+ self._transaction_id = result_set.metadata.transaction.id
+ 
  return response.status, row_counts
 
  def __enter__(self):
  """Begin ``with`` block."""
- self.begin()
  return self
 
  def __exit__(self, exc_type, exc_val, exc_tb):

tests/unit/test_pool.py

@@ -656,7 +656,7 @@ def test_bind(self):
  for session in SESSIONS:
  session.create.assert_not_called()
  txn = session._transaction
- txn.begin.assert_called_once_with()
+ txn.begin.assert_not_called()
 
  self.assertTrue(pool._pending_sessions.empty())
 
@@ -685,7 +685,7 @@ def test_bind_w_timestamp_race(self):
  for session in SESSIONS:
  session.create.assert_not_called()
  txn = session._transaction
- txn.begin.assert_called_once_with()
+ txn.begin.assert_not_called()
 
  self.assertTrue(pool._pending_sessions.empty())
 
@@ -771,7 +771,7 @@ def test_begin_pending_transactions_non_empty(self):
  pool.begin_pending_transactions() # no raise
 
  for txn in TRANSACTIONS:
- txn.begin.assert_called_once_with()
+ txn.begin.assert_not_called()
 
  self.assertTrue(pending.empty())
 

tests/unit/test_session.py

@@ -725,17 +725,6 @@ def unit_of_work(txn, *args, **kw):
  self.assertEqual(args, ())
  self.assertEqual(kw, {})
 
- expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
- gax_api.begin_transaction.assert_called_once_with(
- session=self.SESSION_NAME,
- options=expected_options,
- metadata=[("google-cloud-resource-prefix", database.name)],
- )
- gax_api.rollback.assert_called_once_with(
- session=self.SESSION_NAME,
- transaction_id=TRANSACTION_ID,
- metadata=[("google-cloud-resource-prefix", database.name)],
- )
 
  def test_run_in_transaction_callback_raises_non_abort_rpc_error(self):
  from google.api_core.exceptions import Cancelled
@@ -780,12 +769,6 @@ def unit_of_work(txn, *args, **kw):
  self.assertEqual(args, ())
  self.assertEqual(kw, {})
 
- expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
- gax_api.begin_transaction.assert_called_once_with(
- session=self.SESSION_NAME,
- options=expected_options,
- metadata=[("google-cloud-resource-prefix", database.name)],
- )
  gax_api.rollback.assert_not_called()
 
  def test_run_in_transaction_w_args_w_kwargs_wo_abort(self):
@@ -1141,16 +1124,10 @@ def unit_of_work(txn, *args, **kw):
  self.assertEqual(kw, {})
 
  expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
- self.assertEqual(
- gax_api.begin_transaction.call_args_list,
- [
- mock.call(
- session=self.SESSION_NAME,
- options=expected_options,
- metadata=[("google-cloud-resource-prefix", database.name)],
- )
- ]
- * 2,
+ gax_api.begin_transaction.assert_called_once_with(
+ session=self.SESSION_NAME,
+ options=expected_options,
+ metadata=[("google-cloud-resource-prefix", database.name)],
  )
  request = CommitRequest(
  session=self.SESSION_NAME,