Repository: googleapis/python-spanner-django

Issue:
dbapi: rewrite INSERT statement
While running `python3 manage.py migrate` with SQL
```SQL
sql: INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)
params: ['contenttypes', '0001_initial', '2019-11-13 20:59:05.166706+00:00']
```
We got
```shell
Traceback (most recent call last):
  File "manage.py", line 21, in <module>
    main()
  File "manage.py", line 17, in main
    execute_from_command_line(sys.argv)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/core/management/__init__.py", line 381, in execute_from_command_line
    utility.execute()
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/core/management/__init__.py", line 375, in execute
    self.fetch_command(subcommand).run_from_argv(self.argv)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/core/management/base.py", line 323, in run_from_argv
    self.execute(*args, **cmd_options)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/core/management/base.py", line 364, in execute
    output = self.handle(*args, **options)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/core/management/base.py", line 83, in wrapped
    res = handle_func(*args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/core/management/commands/migrate.py", line 234, in handle
    fake_initial=fake_initial,
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/migrations/executor.py", line 117, in migrate
    state = self._migrate_all_forwards(state, plan, full_plan, fake=fake, fake_initial=fake_initial)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/migrations/executor.py", line 147, in _migrate_all_forwards
    state = self.apply_migration(state, migration, fake=fake, fake_initial=fake_initial)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/migrations/executor.py", line 246, in apply_migration
    self.record_migration(migration)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/migrations/executor.py", line 261, in record_migration
    self.recorder.record_applied(migration.app_label, migration.name)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/migrations/recorder.py", line 83, in record_applied
    self.migration_qs.create(app=app, name=name)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/models/query.py", line 422, in create
    obj.save(force_insert=True, using=self.db)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/models/base.py", line 741, in save
    force_update=force_update, update_fields=update_fields)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/models/base.py", line 779, in save_base
    force_update, using, update_fields,
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/models/base.py", line 870, in _save_table
    result = self._do_insert(cls._base_manager, using, fields, update_pk, raw)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/models/base.py", line 908, in _do_insert
    using=using, raw=raw)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/models/manager.py", line 82, in manager_method
    return getattr(self.get_queryset(), name)(*args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/models/query.py", line 1186, in _insert
    return query.get_compiler(using=using).execute_sql(return_id)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/models/sql/compiler.py", line 1335, in execute_sql
    cursor.execute(sql, params)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/backends/utils.py", line 99, in execute
    return super().execute(sql, params)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/backends/utils.py", line 67, in execute
    return self._execute_with_wrappers(sql, params, many=False, executor=self._execute)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/backends/utils.py", line 76, in _execute_with_wrappers
    return executor(sql, params, many, context)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/django/db/backends/utils.py", line 84, in _execute
    return self.cursor.execute(sql, params)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/spanner/dbapi/cursor.py", line 90, in execute
    param_types=param_types,
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/google/cloud/spanner_v1/session.py", line 299, in run_in_transaction
    return_value = func(txn, *args, **kw)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/spanner/dbapi/cursor.py", line 103, in __do_execute_update
    res = transaction.execute_update(sql, *args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/google/cloud/spanner_v1/transaction.py", line 188, in execute_update
    params_pb = self._make_params_pb(params, param_types)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/google/cloud/spanner_v1/transaction.py", line 157, in _make_params_pb
    fields={key: _make_value_pb(value) for key, value in params.items()}
```
but the problem here is that for inserts into Cloud Spanner we need to invoke
https://googleapis.dev/python/spanner/latest/transaction-api.html#google.cloud.spanner_v1.transaction.Transaction.insert
which is
```python
 res = transaction.insert_or_update(
                table=table_name,
                columns=columns,  # List or iterable
                values=seq_of_params)
```


Discussion:

Resolving patch:
diff --git a/spanner/dbapi/cursor.py b/spanner/dbapi/cursor.py
--- a/spanner/dbapi/cursor.py
+++ b/spanner/dbapi/cursor.py
@@ -15,7 +15,9 @@
 import google.api_core.exceptions as grpc_exceptions
 
 from .exceptions import IntegrityError, OperationalError, ProgrammingError
-from .parse_utils import STMT_DDL, STMT_NON_UPDATING, classify_stmt
+from .parse_utils import (
+    STMT_DDL, STMT_INSERT, STMT_NON_UPDATING, classify_stmt, parse_insert,
+)
 
 _UNSET_COUNT = -1
 
@@ -82,6 +84,16 @@ def execute(self, sql, args=None):
                     sql, args or None,
                     param_types=param_types,
                 )
+            elif classification == STMT_INSERT:
+                self.__session.run_in_transaction(
+                    self.__do_execute_insert,
+                    sql,
+                    # Insert statements' params are only passed as tuples or lists,
+                    # yet for do_execute_update, we've got to pass in list of list.
+                    # https://googleapis.dev/python/spanner/latest/transaction-api.html\
+                    #           #google.cloud.spanner_v1.transaction.Transaction.insert
+                    [args] if args else None,
+                )
             else:
                 self.__session.run_in_transaction(
                     self.__do_execute_update,
@@ -106,6 +118,14 @@ def __do_execute_update(self, transaction, sql, *args, **kwargs):
 
         return res
 
+    def __do_execute_insert(self, transaction, sql, params):
+        parts = parse_insert(sql)
+        return transaction.insert_or_update(
+            table=parts.get('table'),
+            columns=parts.get('columns'),
+            values=params,
+        )
+
     def __do_execute_non_update(self, sql, *args, **kwargs):
         # Reference
         #  https://googleapis.dev/python/spanner/latest/session-api.html#google.cloud.spanner_v1.session.Session.execute_sql
diff --git a/spanner/dbapi/parse_utils.py b/spanner/dbapi/parse_utils.py
--- a/spanner/dbapi/parse_utils.py
+++ b/spanner/dbapi/parse_utils.py
@@ -213,6 +213,7 @@ def validate_instance_config(instance_config):
 STMT_DDL = 'DDL'
 STMT_NON_UPDATING = 'NON_UPDATING'
 STMT_UPDATING = 'UPDATING'
+STMT_INSERT = 'INSERT'
 
 # Heuristic for identifying statements that don't need to be run as updates.
 re_NON_UPDATE = re.compile(r'^\s*(SELECT|ANALYZE|AUDIT|EXPLAIN|SHOW)', re.UNICODE | re.IGNORECASE)
@@ -223,11 +224,38 @@ def validate_instance_config(instance_config):
     re.UNICODE | re.IGNORECASE | re.DOTALL,
 )
 
+re_IS_INSERT = re.compile(
+    r'^\s*(INSERT)',
+    re.UNICODE | re.IGNORECASE | re.DOTALL,
+)
+
 
 def classify_stmt(sql):
     if re_DDL.match(sql):
         return STMT_DDL
+    elif re_IS_INSERT.match(sql):
+        return STMT_INSERT
     elif re_NON_UPDATE.match(sql):
         return STMT_NON_UPDATING
     else:
         return STMT_UPDATING
+
+
+re_INSERT = re.compile(
+    # Only match the `INSERT INTO <table_name> (columns...)
+    # otherwise the rest of the statement could be a complex
+    # operation.
+    r'^\s*INSERT INTO (?P<table_name>[^\s\(\)]+)\s+\((?P<columns>[^\(\)]+)\)',
+    re.IGNORECASE | re.UNICODE | re.DOTALL,
+)
+
+
+def parse_insert(insert_sql):
+    match = re_INSERT.search(insert_sql)
+    if match:
+        return {
+            'table': match.group('table_name'),
+            'columns': [mi.strip() for mi in match.group('columns').split(',')],
+        }
+    else:
+        return None


Test patch:
diff --git a/tests/spanner/dbapi/test_parse_utils.py b/tests/spanner/dbapi/test_parse_utils.py
--- a/tests/spanner/dbapi/test_parse_utils.py
+++ b/tests/spanner/dbapi/test_parse_utils.py
@@ -17,7 +17,8 @@
 from spanner.dbapi.exceptions import Error
 from spanner.dbapi.parse_utils import (
     STMT_DDL, STMT_NON_UPDATING, classify_stmt, extract_connection_params,
-    parse_spanner_url, reINSTANCE_CONFIG, validate_instance_config,
+    parse_insert, parse_spanner_url, reINSTANCE_CONFIG,
+    validate_instance_config,
 )
 
 
@@ -184,3 +185,27 @@ def test_classify_stmt(self):
             sql, want_classification = tt
             got_classification = classify_stmt(sql)
             self.assertEqual(got_classification, want_classification, 'Classification mismatch')
+
+    def test_parse_insert(self):
+        cases = [
+            (
+                'INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)',
+                {
+                    'table': 'django_migrations',
+                    'columns': ['app', 'name', 'applied'],
+                },
+            ),
+            (
+                'INSERT INTO sales.addresses (street, city, state, zip_code)'
+                'SELECT street, city, state, zip_code FROM sales.customers'
+                'ORDER BY first_name, last_name',
+                {
+                    'table': 'sales.addresses',
+                    'columns': ['street', 'city', 'state', 'zip_code'],
+                },
+            ),
+        ]
+        for sql, want in cases:
+            with self.subTest(sql=sql):
+                got = parse_insert(sql)
+                self.assertEqual(got, want, 'Mismatch with parse_insert of `%s`' % sql)