feat(dbapi): add aborted transactions retry support (#168)

Fixes #34. See googleapis/python-spanner-django#544.

Author: Ilya Gurov

google/cloud/spanner_dbapi/checksum.py

@@ -0,0 +1,80 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""API to calculate checksums of SQL statements results."""
+
+import hashlib
+import pickle
+
+from google.cloud.spanner_dbapi.exceptions import RetryAborted
+
+
+class ResultsChecksum:
+ """Cumulative checksum.
+
+ Used to calculate a total checksum of all the results
+ returned by operations executed within transaction.
+ Includes methods for checksums comparison.
+ These checksums are used while retrying an aborted
+ transaction to check if the results of a retried transaction
+ are equal to the results of the original transaction.
+ """
+
+ def __init__(self):
+ self.checksum = hashlib.sha256()
+ self.count = 0 # counter of consumed results
+
+ def __len__(self):
+ """Return the number of consumed results.
+
+ :rtype: :class:`int`
+ :returns: The number of results.
+ """
+ return self.count
+
+ def __eq__(self, other):
+ """Check if checksums are equal.
+
+ :type other: :class:`google.cloud.spanner_dbapi.checksum.ResultsChecksum`
+ :param other: Another checksum to compare with this one.
+ """
+ return self.checksum.digest() == other.checksum.digest()
+
+ def consume_result(self, result):
+ """Add the given result into the checksum.
+
+ :type result: Union[int, list]
+ :param result: Streamed row or row count from an UPDATE operation.
+ """
+ self.checksum.update(pickle.dumps(result))
+ self.count += 1
+
+
+def _compare_checksums(original, retried):
+ """Compare the given checksums.
+
+ Raise an error if the given checksums are not equal.
+
+ :type original: :class:`~google.cloud.spanner_dbapi.checksum.ResultsChecksum`
+ :param original: results checksum of the original transaction.
+
+ :type retried: :class:`~google.cloud.spanner_dbapi.checksum.ResultsChecksum`
+ :param retried: results checksum of the retried transaction.
+
+ :raises: :exc:`google.cloud.spanner_dbapi.exceptions.RetryAborted` in case if checksums are not equal.
+ """
+ if retried != original:
+ raise RetryAborted(
+ "The transaction was aborted and could not be retried due to a concurrent modification."
+ )

google/cloud/spanner_dbapi/connection.py

@@ -14,18 +14,24 @@
 
 """DB-API Connection for the Google Cloud Spanner."""
 
+import time
 import warnings
 
+from google.api_core.exceptions import Aborted
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.cloud import spanner_v1 as spanner
+from google.cloud.spanner_v1.session import _get_retry_delay
 
+from google.cloud.spanner_dbapi.checksum import _compare_checksums
+from google.cloud.spanner_dbapi.checksum import ResultsChecksum
 from google.cloud.spanner_dbapi.cursor import Cursor
 from google.cloud.spanner_dbapi.exceptions import InterfaceError
 from google.cloud.spanner_dbapi.version import DEFAULT_USER_AGENT
 from google.cloud.spanner_dbapi.version import PY_VERSION
 
 
 AUTOCOMMIT_MODE_WARNING = "This method is non-operational in autocommit mode"
+MAX_INTERNAL_RETRIES = 50
 
 
 class Connection:
@@ -48,9 +54,16 @@ def __init__(self, instance, database):
 
  self._transaction = None
  self._session = None
+ # SQL statements, which were executed
+ # within the current transaction
+ self._statements = []
 
  self.is_closed = False
  self._autocommit = False
+ # indicator to know if the session pool used by
+ # this connection should be cleared on the
+ # connection close
+ self._own_pool = True
 
  @property
  def autocommit(self):
@@ -114,6 +127,58 @@ def _release_session(self):
  self.database._pool.put(self._session)
  self._session = None
 
+ def retry_transaction(self):
+ """Retry the aborted transaction.
+
+ All the statements executed in the original transaction
+ will be re-executed in new one. Results checksums of the
+ original statements and the retried ones will be compared.
+
+ :raises: :class:`google.cloud.spanner_dbapi.exceptions.RetryAborted`
+ If results checksum of the retried statement is
+ not equal to the checksum of the original one.
+ """
+ attempt = 0
+ while True:
+ self._transaction = None
+ attempt += 1
+ if attempt > MAX_INTERNAL_RETRIES:
+ raise
+
+ try:
+ self._rerun_previous_statements()
+ break
+ except Aborted as exc:
+ delay = _get_retry_delay(exc.errors[0], attempt)
+ if delay:
+ time.sleep(delay)
+
+ def _rerun_previous_statements(self):
+ """
+ Helper to run all the remembered statements
+ from the last transaction.
+ """
+ for statement in self._statements:
+ res_iter, retried_checksum = self.run_statement(statement, retried=True)
+ # executing all the completed statements
+ if statement != self._statements[-1]:
+ for res in res_iter:
+ retried_checksum.consume_result(res)
+
+ _compare_checksums(statement.checksum, retried_checksum)
+ # executing the failed statement
+ else:
+ # streaming up to the failed result or
+ # to the end of the streaming iterator
+ while len(retried_checksum) < len(statement.checksum):
+ try:
+ res = next(iter(res_iter))
+ retried_checksum.consume_result(res)
+ except StopIteration:
+ break
+
+ _compare_checksums(statement.checksum, retried_checksum)
+
  def transaction_checkout(self):
  """Get a Cloud Spanner transaction.
 
@@ -158,6 +223,9 @@ def close(self):
  ):
  self._transaction.rollback()
 
+ if self._own_pool:
+ self.database._pool.clear()
+
  self.is_closed = True
 
  def commit(self):
@@ -168,8 +236,13 @@ def commit(self):
  if self._autocommit:
  warnings.warn(AUTOCOMMIT_MODE_WARNING, UserWarning, stacklevel=2)
  elif self._transaction:
- self._transaction.commit()
- self._release_session()
+ try:
+ self._transaction.commit()
+ self._release_session()
+ self._statements = []
+ except Aborted:
+ self.retry_transaction()
+ self.commit()
 
  def rollback(self):
  """Rolls back any pending transaction.
@@ -182,6 +255,7 @@ def rollback(self):
  elif self._transaction:
  self._transaction.rollback()
  self._release_session()
+ self._statements = []
 
  def cursor(self):
  """Factory to create a DB-API Cursor."""
@@ -198,6 +272,32 @@ def run_prior_DDL_statements(self):
 
  return self.database.update_ddl(ddl_statements).result()
 
+ def run_statement(self, statement, retried=False):
+ """Run single SQL statement in begun transaction.
+
+ This method is never used in autocommit mode. In
+ !autocommit mode however it remembers every executed
+ SQL statement with its parameters.
+
+ :type statement: :class:`dict`
+ :param statement: SQL statement to execute.
+
+ :rtype: :class:`google.cloud.spanner_v1.streamed.StreamedResultSet`,
+ :class:`google.cloud.spanner_dbapi.checksum.ResultsChecksum`
+ :returns: Streamed result set of the statement and a
+ checksum of this statement results.
+ """
+ transaction = self.transaction_checkout()
+ if not retried:
+ self._statements.append(statement)
+
+ return (
+ transaction.execute_sql(
+ statement.sql, statement.params, param_types=statement.param_types,
+ ),
+ ResultsChecksum() if retried else statement.checksum,
+ )
+
  def __enter__(self):
  return self
 
@@ -207,7 +307,12 @@ def __exit__(self, etype, value, traceback):
 
 
 def connect(
- instance_id, database_id, project=None, credentials=None, pool=None, user_agent=None
+ instance_id,
+ database_id,
+ project=None,
+ credentials=None,
+ pool=None,
+ user_agent=None,
 ):
  """Creates a connection to a Google Cloud Spanner database.
 
@@ -261,4 +366,8 @@ def connect(
  if not database.exists():
  raise ValueError("database '%s' does not exist." % database_id)
 
- return Connection(instance, database)
+ conn = Connection(instance, database)
+ if pool is not None:
+ conn._own_pool = False
+
+ return conn

google/cloud/spanner_dbapi/cursor.py

@@ -14,6 +14,7 @@
 
 """Database cursor for Google Cloud Spanner DB-API."""
 
+from google.api_core.exceptions import Aborted
 from google.api_core.exceptions import AlreadyExists
 from google.api_core.exceptions import FailedPrecondition
 from google.api_core.exceptions import InternalServerError
@@ -22,7 +23,7 @@
 from collections import namedtuple
 
 from google.cloud import spanner_v1 as spanner
-
+from google.cloud.spanner_dbapi.checksum import ResultsChecksum
 from google.cloud.spanner_dbapi.exceptions import IntegrityError
 from google.cloud.spanner_dbapi.exceptions import InterfaceError
 from google.cloud.spanner_dbapi.exceptions import OperationalError
@@ -34,11 +35,13 @@
 
 from google.cloud.spanner_dbapi import parse_utils
 from google.cloud.spanner_dbapi.parse_utils import get_param_types
+from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner
 from google.cloud.spanner_dbapi.utils import PeekIterator
 
 _UNSET_COUNT = -1
 
 ColumnDetails = namedtuple("column_details", ["null_ok", "spanner_type"])
+Statement = namedtuple("Statement", "sql, params, param_types, checksum")
 
 
 class Cursor(object):
@@ -54,6 +57,8 @@ def __init__(self, connection):
  self._row_count = _UNSET_COUNT
  self.connection = connection
  self._is_closed = False
+ # the currently running SQL statement results checksum
+ self._checksum = None
 
  # the number of rows to fetch at a time with fetchmany()
  self.arraysize = 1
@@ -166,12 +171,13 @@ def execute(self, sql, args=None):
  self.connection.run_prior_DDL_statements()
 
  if not self.connection.autocommit:
- transaction = self.connection.transaction_checkout()
-
- sql, params = parse_utils.sql_pyformat_args_to_spanner(sql, args)
+ sql, params = sql_pyformat_args_to_spanner(sql, args)
 
- self._result_set = transaction.execute_sql(
- sql, params, param_types=get_param_types(params)
+ statement = Statement(
+ sql, params, get_param_types(params), ResultsChecksum(),
+ )
+ (self._result_set, self._checksum,) = self.connection.run_statement(
+ statement
  )
  self._itr = PeekIterator(self._result_set)
  return
@@ -213,9 +219,31 @@ def fetchone(self):
  self._raise_if_closed()
 
  try:
- return next(self)
+ res = next(self)
+ self._checksum.consume_result(res)
+ return res
  except StopIteration:
- return None
+ return
+ except Aborted:
+ self.connection.retry_transaction()
+ return self.fetchone()
+
+ def fetchall(self):
+ """Fetch all (remaining) rows of a query result, returning them as
+ a sequence of sequences.
+ """
+ self._raise_if_closed()
+
+ res = []
+ try:
+ for row in self:
+ self._checksum.consume_result(row)
+ res.append(row)
+ except Aborted:
+ self._connection.retry_transaction()
+ return self.fetchall()
+
+ return res
 
  def fetchmany(self, size=None):
  """Fetch the next set of rows of a query result, returning a sequence
@@ -236,20 +264,17 @@ def fetchmany(self, size=None):
  items = []
  for i in range(size):
  try:
- items.append(tuple(self.__next__()))
+ res = next(self)
+ self._checksum.consume_result(res)
+ items.append(res)
  except StopIteration:
  break
+ except Aborted:
+ self._connection.retry_transaction()
+ return self.fetchmany(size)
 
  return items
 
- def fetchall(self):
- """Fetch all (remaining) rows of a query result, returning them as
- a sequence of sequences.
- """
- self._raise_if_closed()
-
- return list(self.__iter__())
-
  def nextset(self):
  """A no-op, raising an error if the cursor or connection is closed."""
  self._raise_if_closed()