feat: support aborted transactions internal retry

google/cloud/spanner_dbapi/checksum.py

@@ -0,0 +1,73 @@
+# Copyright 2020 Google LLC
+#
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file or at
+# https://developers.google.com/open-source/licenses/bsd
+
+"""API to calculate checksums of SQL statements results."""
+
+import hashlib
+import pickle
+
+from google.api_core.exceptions import Aborted
+
+
+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 has
+    different length, or are not equal.
+
+    :type original: :class:`~google.cloud.spanner_v1.transaction.ResultsChecksum`
+    :param original: results checksum of the original transaction.
+
+    :type retried: :class:`~google.cloud.spanner_v1.transaction.ResultsChecksum`
+    :param retried: results checksum of the retried transaction.
+
+    :raises: :exc:`google.api_core.exceptions.Aborted` in case if checksums are not equal.
+    """
+    if len(retried) != len(original) or retried != original:
+        raise Aborted(
+            "The transaction was aborted and could not be retried due to a concurrent modification."
+        )

google/cloud/spanner_dbapi/connection.py

@@ -11,6 +11,8 @@
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.cloud import spanner_v1 as spanner
 
+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
@@ -40,6 +42,9 @@ 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
@@ -106,6 +111,39 @@ 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 transaction
+        will be re-executed. Results checksums of the original
+        statements and the retried ones will be compared.
+
+        :raises: :class:`google.api_core.exceptions.Aborted`
+            If results checksum of the retried statement is
+            not equal to the checksum of the original one.
+        """
+        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
+                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.
 
@@ -162,6 +200,7 @@ def commit(self):
         elif self._transaction:
             self._transaction.commit()
             self._release_session()
+            self._statements = []
 
     def rollback(self):
         """Rolls back any pending transaction.
@@ -174,6 +213,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."""
@@ -190,6 +230,33 @@ 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()
+        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
 
@@ -243,11 +310,11 @@ def connect(
     """
 
     client_info = ClientInfo(
-        user_agent=user_agent or DEFAULT_USER_AGENT, python_version=PY_VERSION,
+        user_agent=user_agent or DEFAULT_USER_AGENT, python_version=PY_VERSION
     )
 
     client = spanner.Client(
-        project=project, credentials=credentials, client_info=client_info,
+        project=project, credentials=credentials, client_info=client_info
     )
 
     instance = client.instance(instance_id)

google/cloud/spanner_dbapi/cursor.py

@@ -6,6 +6,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
@@ -14,7 +15,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
@@ -26,6 +27,7 @@
 
 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
@@ -46,6 +48,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
@@ -158,16 +162,18 @@ 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
-                )
-
-                self._result_set = transaction.execute_sql(
-                    sql, params, param_types=get_param_types(params)
+                sql, params = sql_pyformat_args_to_spanner(sql, args)
+
+                statement = {
+                    "sql": sql,
+                    "params": params,
+                    "param_types": get_param_types(params),
+                    "checksum": ResultsChecksum(),
+                }
+                self._res, self._checksum = self.connection.run_statement(
+                    statement
                 )
-                self._itr = PeekIterator(self._result_set)
+                self._itr = PeekIterator(self._res)
                 return
 
             if classification == parse_utils.STMT_NON_UPDATING:
@@ -207,9 +213,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.__iter__():
+                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
@@ -230,20 +258,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()

tests/spanner_dbapi/test_checksum.py

@@ -0,0 +1,54 @@
+# Copyright 2020 Google LLC
+#
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file or at
+# https://developers.google.com/open-source/licenses/bsd
+
+import unittest
+
+from google.api_core.exceptions import Aborted
+from google.cloud.spanner_dbapi.checksum import (
+    _compare_checksums,
+    ResultsChecksum,
+)
+
+
+class Test_compare_checksums(unittest.TestCase):
+    def test_equal(self):
+        original = ResultsChecksum()
+        original.consume_result(5)
+
+        retried = ResultsChecksum()
+        retried.consume_result(5)
+
+        self.assertIsNone(_compare_checksums(original, retried))
+
+    def test_less_results(self):
+        original = ResultsChecksum()
+        original.consume_result(5)
+
+        retried = ResultsChecksum()
+
+        with self.assertRaises(Aborted):
+            _compare_checksums(original, retried)
+
+    def test_more_results(self):
+        original = ResultsChecksum()
+        original.consume_result(5)
+
+        retried = ResultsChecksum()
+        retried.consume_result(5)
+        retried.consume_result(2)
+
+        with self.assertRaises(Aborted):
+            _compare_checksums(original, retried)
+
+    def test_mismatch(self):
+        original = ResultsChecksum()
+        original.consume_result(5)
+
+        retried = ResultsChecksum()
+        retried.consume_result(2)
+
+        with self.assertRaises(Aborted):
+            _compare_checksums(original, retried)