fix: increment seqno before execute calls to prevent InvalidArgument … (#19)

* fix: increment seqno before execute calls to prevent InvalidArgument errors after a previous error

* make assignments atomic

* add and update tests

* revert snapshot.py change

* formatting

Co-authored-by: larkee <[email protected]>

Author: larkee

google/cloud/spanner_v1/transaction.py

@@ -201,6 +201,11 @@ def execute_update(
         transaction = self._make_txn_selector()
         api = database.spanner_api
 
+        seqno, self._execute_sql_count = (
+            self._execute_sql_count,
+            self._execute_sql_count + 1,
+        )
+
         # Query-level options have higher precedence than client-level and
         # environment-level options
         default_query_options = database._instance._client._query_options
@@ -214,11 +219,9 @@ def execute_update(
             param_types=param_types,
             query_mode=query_mode,
             query_options=query_options,
-            seqno=self._execute_sql_count,
+            seqno=seqno,
             metadata=metadata,
         )
-
-        self._execute_sql_count += 1
         return response.stats.row_count_exact
 
     def batch_update(self, statements):
@@ -259,15 +262,18 @@ def batch_update(self, statements):
         transaction = self._make_txn_selector()
         api = database.spanner_api
 
+        seqno, self._execute_sql_count = (
+            self._execute_sql_count,
+            self._execute_sql_count + 1,
+        )
+
         response = api.execute_batch_dml(
             session=self._session.name,
             transaction=transaction,
             statements=parsed,
-            seqno=self._execute_sql_count,
+            seqno=seqno,
             metadata=metadata,
         )
-
-        self._execute_sql_count += 1
         row_counts = [
             result_set.stats.row_count_exact for result_set in response.result_sets
         ]

tests/unit/test_snapshot.py

@@ -311,6 +311,8 @@ def test_execute_sql_other_error(self):
         with self.assertRaises(RuntimeError):
             list(derived.execute_sql(SQL_QUERY))
 
+        self.assertEqual(derived._execute_sql_count, 1)
+
     def test_execute_sql_w_params_wo_param_types(self):
         database = _Database()
         session = _Session(database)

tests/unit/test_transaction.py

@@ -413,6 +413,19 @@ def test_execute_update_new_transaction(self):
     def test_execute_update_w_count(self):
         self._execute_update_helper(count=1)
 
+    def test_execute_update_error(self):
+        database = _Database()
+        database.spanner_api = self._make_spanner_api()
+        database.spanner_api.execute_sql.side_effect = RuntimeError()
+        session = _Session(database)
+        transaction = self._make_one(session)
+        transaction._transaction_id = self.TRANSACTION_ID
+
+        with self.assertRaises(RuntimeError):
+            transaction.execute_update(DML_QUERY)
+
+        self.assertEqual(transaction._execute_sql_count, 1)
+
     def test_execute_update_w_query_options(self):
         from google.cloud.spanner_v1.proto.spanner_pb2 import ExecuteSqlRequest
 
@@ -513,6 +526,31 @@ def test_batch_update_wo_errors(self):
     def test_batch_update_w_errors(self):
         self._batch_update_helper(error_after=2, count=1)
 
+    def test_batch_update_error(self):
+        database = _Database()
+        api = database.spanner_api = self._make_spanner_api()
+        api.execute_batch_dml.side_effect = RuntimeError()
+        session = _Session(database)
+        transaction = self._make_one(session)
+        transaction._transaction_id = self.TRANSACTION_ID
+
+        insert_dml = "INSERT INTO table(pkey, desc) VALUES (%pkey, %desc)"
+        insert_params = {"pkey": 12345, "desc": "DESCRIPTION"}
+        insert_param_types = {"pkey": "INT64", "desc": "STRING"}
+        update_dml = 'UPDATE table SET desc = desc + "-amended"'
+        delete_dml = "DELETE FROM table WHERE desc IS NULL"
+
+        dml_statements = [
+            (insert_dml, insert_params, insert_param_types),
+            update_dml,
+            delete_dml,
+        ]
+
+        with self.assertRaises(RuntimeError):
+            transaction.batch_update(dml_statements)
+
+        self.assertEqual(transaction._execute_sql_count, 1)
+
     def test_context_mgr_success(self):
         import datetime
         from google.cloud.spanner_v1.proto.spanner_pb2 import CommitResponse