feat: use DML batches in executemany() method

google/cloud/spanner_dbapi/connection.py

@@ -32,6 +32,9 @@
 from google.cloud.spanner_dbapi.version import DEFAULT_USER_AGENT
 from google.cloud.spanner_dbapi.version import PY_VERSION
 
+from google.cloud.spanner_dbapi.exceptions import OperationalError
+from google.rpc.code_pb2 import ABORTED, OK
+
 
 AUTOCOMMIT_MODE_WARNING = "This method is non-operational in autocommit mode"
 MAX_INTERNAL_RETRIES = 50
@@ -175,43 +178,58 @@ def _rerun_previous_statements(self):
         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
+            if isinstance(statement, list):
+                statements, checksum = statement
+
+                transaction = self.transaction_checkout()
+                status, res = transaction.batch_update(statements)
+
+                if status.code == ABORTED:
+                    self.connection._transaction = None
+                    raise Aborted(status.details)
+                elif status.code != OK:
+                    raise OperationalError(status.details)
+
+                retried_checksum = ResultsChecksum()
+                retried_checksum.consume_result(res)
+                retried_checksum.consume_result(status.code)
+
+                _compare_checksums(checksum, retried_checksum)
             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))
+                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)
-                    except StopIteration:
-                        break
 
-                _compare_checksums(statement.checksum, retried_checksum)
+                    _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.
 
         Begin a new transaction, if there is no transaction in
         this connection yet. Return the begun one otherwise.
 
-        The method is non operational in autocommit mode.
-
         :rtype: :class:`google.cloud.spanner_v1.transaction.Transaction`
         :returns: A Cloud Spanner transaction object, ready to use.
         """
-        if not self.autocommit:
-            if not self.inside_transaction:
-                self._transaction = self._session_checkout().transaction()
-                self._transaction.begin()
+        if not self.inside_transaction:
+            self._transaction = self._session_checkout().transaction()
+            self._transaction.begin()
 
-            return self._transaction
+        return self._transaction
 
     def _raise_if_closed(self):
         """Helper to check the connection state before running a query.

google/cloud/spanner_dbapi/cursor.py

@@ -41,6 +41,8 @@
 from google.cloud.spanner_dbapi.utils import PeekIterator
 from google.cloud.spanner_dbapi.utils import StreamedManyResultSets
 
+from google.rpc.code_pb2 import ABORTED, OK
+
 _UNSET_COUNT = -1
 
 ColumnDetails = namedtuple("column_details", ["null_ok", "spanner_type"])
@@ -258,9 +260,57 @@ def executemany(self, operation, seq_of_params):
 
         many_result_set = StreamedManyResultSets()
 
-        for params in seq_of_params:
-            self.execute(operation, params)
-            many_result_set.add_iter(self._itr)
+        if classification in (parse_utils.STMT_INSERT, parse_utils.STMT_UPDATING):
+            statements = []
+
+            for params in seq_of_params:
+                sql, params = parse_utils.sql_pyformat_args_to_spanner(
+                    operation, params
+                )
+                statements.append((sql, params, get_param_types(params)))
+
+            if self.connection.autocommit:
+                transaction = self.connection.transaction_checkout()
+                status, res = transaction.batch_update(statements)
+                many_result_set.add_iter(res)
+
+                if status.code != OK:
+                    self.connection._transaction.rollback()
+                    self.connection._transaction = None
+                    raise OperationalError(status.details)
+
+                transaction.commit()
+            else:
+                retried = False
+                while True:
+                    try:
+                        transaction = self.connection.transaction_checkout()
+
+                        res_checksum = ResultsChecksum()
+                        if not retried:
+                            self.connection._statements.append(
+                                (statements, res_checksum)
+                            )
+
+                        status, res = transaction.batch_update(statements)
+                        many_result_set.add_iter(res)
+                        res_checksum.consume_result(res)
+                        res_checksum.consume_result(status.code)
+
+                        if status.code == ABORTED:
+                            self.connection._transaction = None
+                            raise Aborted(status.details)
+                        elif status.code != OK:
+                            raise OperationalError(status.details)
+                        break
+                    except Aborted:
+                        self.connection.retry_transaction()
+                        retried = True
+
+        else:
+            for params in seq_of_params:
+                self.execute(operation, params)
+                many_result_set.add_iter(self._itr)
 
         self._result_set = many_result_set
         self._itr = many_result_set

tests/unit/spanner_dbapi/test_connection.py

@@ -157,9 +157,6 @@ def test_transaction_checkout(self):
         mock_transaction.committed = mock_transaction.rolled_back = False
         self.assertEqual(connection.transaction_checkout(), mock_transaction)
 
-        connection._autocommit = True
-        self.assertIsNone(connection.transaction_checkout())
-
     def test_close(self):
         from google.cloud.spanner_dbapi import connect, InterfaceError
 

tests/unit/spanner_dbapi/test_cursor.py

@@ -337,6 +337,193 @@ def test_executemany(self):
             (mock.call(operation, (1,)), mock.call(operation, (2,)))
         )
 
+    def test_executemany_insert_batch_non_autocommit(self):
+        from google.cloud.spanner_v1.param_types import INT64
+        from google.cloud.spanner_dbapi import connect
+        from google.rpc.code_pb2 import OK
+
+        sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)"""
+
+        with mock.patch(
+            "google.cloud.spanner_v1.instance.Instance.exists", return_value=True
+        ):
+            with mock.patch(
+                "google.cloud.spanner_v1.database.Database.exists", return_value=True,
+            ):
+                connection = connect("test-instance", "test-database")
+
+        connection._transaction = mock.Mock(committed=False, rolled_back=False)
+        connection._transaction.batch_update = mock.Mock(
+            return_value=[mock.Mock(code=OK), []]
+        )
+
+        cursor = connection.cursor()
+        cursor.executemany(sql, [(1, 2, 3, 4), (5, 6, 7, 8)])
+
+        connection._transaction.batch_update.assert_called_once_with(
+            [
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 1, "a1": 2, "a2": 3, "a3": 4},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 5, "a1": 6, "a2": 7, "a3": 8},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+            ]
+        )
+
+    def test_executemany_insert_batch_autocommit(self):
+        from google.cloud.spanner_v1.param_types import INT64
+        from google.cloud.spanner_dbapi import connect
+        from google.rpc.code_pb2 import OK
+
+        sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)"""
+
+        with mock.patch(
+            "google.cloud.spanner_v1.instance.Instance.exists", return_value=True
+        ):
+            with mock.patch(
+                "google.cloud.spanner_v1.database.Database.exists", return_value=True,
+            ):
+                connection = connect("test-instance", "test-database")
+
+        connection.autocommit = True
+
+        connection._transaction = mock.Mock(committed=False, rolled_back=False)
+        connection._transaction.batch_update = mock.Mock(
+            return_value=[mock.Mock(code=OK), []]
+        )
+        connection._transaction.commit = mock.Mock()
+
+        cursor = connection.cursor()
+        cursor.executemany(sql, [(1, 2, 3, 4), (5, 6, 7, 8)])
+
+        connection._transaction.batch_update.assert_called_once_with(
+            [
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 1, "a1": 2, "a2": 3, "a3": 4},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 5, "a1": 6, "a2": 7, "a3": 8},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+            ]
+        )
+        connection._transaction.commit.assert_called_once()
+
+    def test_executemany_insert_batch_failed(self):
+        from google.cloud.spanner_dbapi import connect
+        from google.cloud.spanner_dbapi.exceptions import OperationalError
+        from google.rpc.code_pb2 import UNKNOWN
+
+        sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)"""
+        err_details = "Details here"
+
+        with mock.patch(
+            "google.cloud.spanner_v1.instance.Instance.exists", return_value=True
+        ):
+            with mock.patch(
+                "google.cloud.spanner_v1.database.Database.exists", return_value=True,
+            ):
+                connection = connect("test-instance", "test-database")
+
+        connection.autocommit = True
+        cursor = connection.cursor()
+
+        connection._transaction = mock.Mock(committed=False, rolled_back=False)
+        connection._transaction.batch_update = mock.Mock(
+            return_value=(mock.Mock(code=UNKNOWN, details=err_details), [])
+        )
+
+        with self.assertRaisesRegex(OperationalError, err_details):
+            cursor.executemany(sql, [(1, 2, 3, 4), (5, 6, 7, 8)])
+
+    def test_executemany_insert_batch_aborted(self):
+        from google.cloud.spanner_dbapi import connect
+        from google.cloud.spanner_dbapi.checksum import ResultsChecksum
+        from google.cloud.spanner_v1.param_types import INT64
+        from google.rpc.code_pb2 import ABORTED, OK
+
+        sql = """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (%s, %s, %s, %s)"""
+        err_details = "Aborted details here"
+
+        with mock.patch(
+            "google.cloud.spanner_v1.instance.Instance.exists", return_value=True
+        ):
+            with mock.patch(
+                "google.cloud.spanner_v1.database.Database.exists", return_value=True,
+            ):
+                connection = connect("test-instance", "test-database")
+
+        transaction1 = mock.Mock(committed=False, rolled_back=False)
+        transaction1.batch_update = mock.Mock(
+            side_effect=[(mock.Mock(code=ABORTED, details=err_details), [])]
+        )
+
+        transaction2 = mock.Mock(committed=False, rolled_back=False)
+        transaction2.batch_update = mock.Mock(side_effect=[(mock.Mock(code=OK), [])])
+
+        connection.transaction_checkout = mock.Mock(
+            side_effect=[transaction1, transaction2]
+        )
+        connection.retry_transaction = mock.Mock()
+
+        cursor = connection.cursor()
+        cursor.executemany(sql, [(1, 2, 3, 4), (5, 6, 7, 8)])
+
+        transaction1.batch_update.assert_called_with(
+            [
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 1, "a1": 2, "a2": 3, "a3": 4},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 5, "a1": 6, "a2": 7, "a3": 8},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+            ]
+        )
+        transaction2.batch_update.assert_called_with(
+            [
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 1, "a1": 2, "a2": 3, "a3": 4},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 5, "a1": 6, "a2": 7, "a3": 8},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+            ]
+        )
+        connection.retry_transaction.assert_called_once()
+
+        self.assertEqual(
+            connection._statements[0][0],
+            [
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 1, "a1": 2, "a2": 3, "a3": 4},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+                (
+                    """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""",
+                    {"a0": 5, "a1": 6, "a2": 7, "a3": 8},
+                    {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
+                ),
+            ],
+        )
+        self.assertIsInstance(connection._statements[0][1], ResultsChecksum)
+
     @unittest.skipIf(
         sys.version_info[0] < 3, "Python 2 has an outdated iterator definition"
     )