Returning ids/rows after bulk insert (#107)

* Implement can_return_rows_from_bulk_insert feature, returning ids or rows after bulk inserts.

* Since mssql-django supports Django 2.2, we also need the pre-Django 3.0 version of feature flag can_return_rows_from_bulk_insert (namely, can_return_ids_from_bulk_insert) (cf. https://docs.djangoproject.com/en/4.0/releases/3.0/#database-backend-api)

* My alternative changes on SQLInsertCompiler.as_sql.

Maybe a bit ambitious, as we completely forsake the SCOPE_IDENTITY strategy (dead code path - we keep the code here, but we could decide not to, really) in favor of OUTPUT strategy.

* Don't try to use the OUTPUT clause when inserting without fields

* Actually we don't really have to offer the feature for Django 2.2, so let's only set can_return_rows_from_bulk_insert to True and not can_return_ids_from_bulk_insert

* Tentative fix: when there are returning fields, but no fields (which means default values insertion - for n objects of course!), we must still fulfill our contract, and return the appropriate rows.

This means we won't use INSERT INTO (...) DEFAULT VALUES n times, but a single INSERT INTO (...) VALUES (DEFAULT, (...), DEFAULT), (...), (DEFAULT, (...), DEFAULT)
Also: be more thorough re the infamous feature flag rename from Django 3.0

* Using MERGE INTO to support Bulk Insertion of multiple rows into a table with only an IDENTITY column.

* Add a link to a reference web page.

* Attempt to make Django 2.2 tests pass

* Get back to a lighter diff of as_sql function vs. original

* Use a query to generate sequence of numbers instead of using the master....spt_values table.

* Update mssql/operations.py

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

* Simplification & refactoring

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

Author: jean-frenette-optel

mssql/compiler.py

@@ -426,6 +426,16 @@ def get_returned_fields(self):
             return self.returning_fields
         return self.return_id
 
+    def can_return_columns_from_insert(self):
+        if django.VERSION >= (3, 0, 0):
+            return self.connection.features.can_return_columns_from_insert
+        return self.connection.features.can_return_id_from_insert
+
+    def can_return_rows_from_bulk_insert(self):
+        if django.VERSION >= (3, 0, 0):
+            return self.connection.features.can_return_rows_from_bulk_insert
+        return self.connection.features.can_return_ids_from_bulk_insert
+
     def fix_auto(self, sql, opts, fields, qn):
         if opts.auto_field is not None:
             # db_column is None if not explicitly specified by model field
@@ -441,15 +451,39 @@ def fix_auto(self, sql, opts, fields, qn):
 
         return sql
 
+    def bulk_insert_default_values_sql(self, table):
+        seed_rows_number = 8
+        cross_join_power = 4  # 8^4 = 4096 > maximum allowed batch size for the backend = 1000
+
+        def generate_seed_rows(n):
+            return " UNION ALL ".join("SELECT 1 AS x" for _ in range(n))
+
+        def cross_join(p):
+            return ", ".join("SEED_ROWS AS _%s" % i for i in range(p))
+
+        return """
+        WITH SEED_ROWS AS (%s)
+            MERGE INTO %s
+            USING (
+                SELECT TOP %s * FROM (SELECT 1 as x FROM %s) FAKE_ROWS
+            ) FAKE_DATA
+            ON 1 = 0
+            WHEN NOT MATCHED THEN
+            INSERT DEFAULT VALUES
+        """ % (generate_seed_rows(seed_rows_number),
+               table,
+               len(self.query.objs),
+               cross_join(cross_join_power))
+
     def as_sql(self):
         # We don't need quote_name_unless_alias() here, since these are all
         # going to be column names (so we can avoid the extra overhead).
         qn = self.connection.ops.quote_name
         opts = self.query.get_meta()
         result = ['INSERT INTO %s' % qn(opts.db_table)]
-        fields = self.query.fields or [opts.pk]
 
         if self.query.fields:
+            fields = self.query.fields
             result.append('(%s)' % ', '.join(qn(f.column) for f in fields))
             values_format = 'VALUES (%s)'
             value_rows = [
@@ -470,11 +504,31 @@ def as_sql(self):
 
         placeholder_rows, param_rows = self.assemble_as_sql(fields, value_rows)
 
-        if self.get_returned_fields() and self.connection.features.can_return_id_from_insert:
-            result.insert(0, 'SET NOCOUNT ON')
-            result.append((values_format + ';') % ', '.join(placeholder_rows[0]))
-            params = [param_rows[0]]
-            result.append('SELECT CAST(SCOPE_IDENTITY() AS bigint)')
+        if self.get_returned_fields() and self.can_return_columns_from_insert():
+            if self.can_return_rows_from_bulk_insert():
+                if not(self.query.fields):
+                    # There isn't really a single statement to bulk multiple DEFAULT VALUES insertions,
+                    # so we have to use a workaround:
+                    # https://dba.stackexchange.com/questions/254771/insert-multiple-rows-into-a-table-with-only-an-identity-column
+                    result = [self.bulk_insert_default_values_sql(qn(opts.db_table))]
+                    r_sql, self.returning_params = self.connection.ops.return_insert_columns(self.get_returned_fields())
+                    if r_sql:
+                        result.append(r_sql)
+                    sql = " ".join(result) + ";"
+                    return [(sql, None)]
+                # Regular bulk insert
+                params = []
+                r_sql, self.returning_params = self.connection.ops.return_insert_columns(self.get_returned_fields())
+                if r_sql:
+                    result.append(r_sql)
+                    params += [self.returning_params]
+                params += param_rows
+                result.append(self.connection.ops.bulk_insert_sql(fields, placeholder_rows))
+            else:
+                result.insert(0, 'SET NOCOUNT ON')
+                result.append((values_format + ';') % ', '.join(placeholder_rows[0]))
+                params = [param_rows[0]]
+                result.append('SELECT CAST(SCOPE_IDENTITY() AS bigint)')
             sql = [(" ".join(result), tuple(chain.from_iterable(params)))]
         else:
             if can_bulk:

mssql/features.py

@@ -12,6 +12,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     can_introspect_small_integer_field = True
     can_return_columns_from_insert = True
     can_return_id_from_insert = True
+    can_return_rows_from_bulk_insert = True
     can_rollback_ddl = True
     can_use_chunked_reads = False
     for_update_after_from = True

mssql/operations.py

@@ -188,6 +188,24 @@ def datetime_trunc_sql(self, lookup_type, field_name, tzname):
             sql = "CONVERT(datetime2, CONVERT(varchar, %s, 20))" % field_name
         return sql
 
+    def fetch_returned_insert_rows(self, cursor):
+        """
+        Given a cursor object that has just performed an INSERT...OUTPUT INSERTED
+        statement into a table, return the list of returned data.
+        """
+        return cursor.fetchall()
+
+    def return_insert_columns(self, fields):
+        if not fields:
+            return '', ()
+        columns = [
+            '%s.%s' % (
+                'INSERTED',
+                self.quote_name(field.column),
+            ) for field in fields
+        ]
+        return 'OUTPUT %s' % ', '.join(columns), ()
+
     def for_update_sql(self, nowait=False, skip_locked=False, of=()):
         if skip_locked:
             return 'WITH (ROWLOCK, UPDLOCK, READPAST)'