fix: fix execute insert for homogeneous statement (#233)

googleapis · Feb 18, 2021 · 36b12a7 · 36b12a7
1 parent 539f145
commit 36b12a7
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 41 deletions.
diff --git a/google/cloud/spanner_dbapi/parse_utils.py b/google/cloud/spanner_dbapi/parse_utils.py
@@ -306,19 +306,15 @@ def parse_insert(insert_sql, params):
         # Case c)
 
         columns = [mi.strip(" `") for mi in match.group("columns").split(",")]
-        sql_params_list = []
-        insert_sql_preamble = "INSERT INTO %s (%s) VALUES %s" % (
-            match.group("table_name"),
-            match.group("columns"),
-            values.argv[0],
-        )
         values_pyformat = [str(arg) for arg in values.argv]
         rows_list = rows_for_insert_or_update(columns, params, values_pyformat)
-        insert_sql_preamble = sanitize_literals_for_upload(insert_sql_preamble)
-        for row in rows_list:
-            sql_params_list.append((insert_sql_preamble, row))
 
-        return {"sql_params_list": sql_params_list}
+        return {
+            "homogenous": True,
+            "table": match.group("table_name"),
+            "columns": columns,
+            "values": rows_list,
+        }
 
     # Case d)
     # insert_sql is of the form:

diff --git a/tests/unit/spanner_dbapi/test_connection.py b/tests/unit/spanner_dbapi/test_connection.py
@@ -379,6 +379,25 @@ def test_run_statement_dont_remember_retried_statements(self):
 
         self.assertEqual(len(connection._statements), 0)
 
+    def test_run_statement_w_homogeneous_insert_statements(self):
+        """Check that Connection executed homogeneous insert statements."""
+        from google.cloud.spanner_dbapi.checksum import ResultsChecksum
+        from google.cloud.spanner_dbapi.cursor import Statement
+
+        sql = "INSERT INTO T (f1, f2) VALUES (%s, %s), (%s, %s)"
+        params = ["a", "b", "c", "d"]
+        param_types = {"f1": str, "f2": str}
+
+        connection = self._make_connection()
+
+        statement = Statement(sql, params, param_types, ResultsChecksum(), True)
+        with mock.patch(
+            "google.cloud.spanner_dbapi.connection.Connection.transaction_checkout"
+        ):
+            connection.run_statement(statement, retried=True)
+
+        self.assertEqual(len(connection._statements), 0)
+
     def test_clear_statements_on_commit(self):
         """
         Check that all the saved statements are

diff --git a/tests/unit/spanner_dbapi/test_parse_utils.py b/tests/unit/spanner_dbapi/test_parse_utils.py
@@ -72,32 +72,20 @@ def test_parse_insert(self):
                 "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)",
                 [1, 2, 3, 4, 5, 6],
                 {
-                    "sql_params_list": [
-                        (
-                            "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)",
-                            (1, 2, 3),
-                        ),
-                        (
-                            "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)",
-                            (4, 5, 6),
-                        ),
-                    ]
+                    "homogenous": True,
+                    "table": "django_migrations",
+                    "columns": ["app", "name", "applied"],
+                    "values": [(1, 2, 3), (4, 5, 6)],
                 },
             ),
             (
                 "INSERT INTO django_migrations(app, name, applied) VALUES (%s, %s, %s)",
                 [1, 2, 3, 4, 5, 6],
                 {
-                    "sql_params_list": [
-                        (
-                            "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)",
-                            (1, 2, 3),
-                        ),
-                        (
-                            "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)",
-                            (4, 5, 6),
-                        ),
-                    ]
+                    "homogenous": True,
+                    "table": "django_migrations",
+                    "columns": ["app", "name", "applied"],
+                    "values": [(1, 2, 3), (4, 5, 6)],
                 },
             ),
             (
@@ -118,25 +106,23 @@ def test_parse_insert(self):
             ),
             (
                 "INSERT INTO ap (n, ct, cn) "
-                "VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s),(%s,      %s, %s)",
+                "VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s),(%s,%s, %s)",
                 (1, 2, 3, 4, 5, 6, 7, 8, 9),
                 {
-                    "sql_params_list": [
-                        ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (1, 2, 3)),
-                        ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (4, 5, 6)),
-                        ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (7, 8, 9)),
-                    ]
+                    "homogenous": True,
+                    "table": "ap",
+                    "columns": ["n", "ct", "cn"],
+                    "values": [(1, 2, 3), (4, 5, 6), (7, 8, 9)],
                 },
             ),
             (
                 "INSERT INTO `no` (`yes`) VALUES (%s)",
                 (1, 4, 5),
                 {
-                    "sql_params_list": [
-                        ("INSERT INTO `no` (`yes`) VALUES (%s)", (1,)),
-                        ("INSERT INTO `no` (`yes`) VALUES (%s)", (4,)),
-                        ("INSERT INTO `no` (`yes`) VALUES (%s)", (5,)),
-                    ]
+                    "homogenous": True,
+                    "table": "`no`",
+                    "columns": ["yes"],
+                    "values": [(1,), (4,), (5,)],
                 },
             ),
             (