11import types
22from itertools import chain
33
4+ import django
45from django .db .models .aggregates import Avg , Count , StdDev , Variance
56from django .db .models .expressions import Exists , OrderBy , Ref , Subquery , Value
67from django .db .models .functions import (
@@ -370,9 +371,9 @@ def as_sql(self, with_limits=True, with_col_aliases=False):
370371 # Finally do cleanup - get rid of the joins we created above.
371372 self .query .reset_refcounts (refcounts_before )
372373
373- def compile (self , node , select_format = False ):
374+ def compile (self , node , * args , ** kwargs ):
374375 node = self ._as_microsoft (node )
375- return super ().compile (node , select_format )
376+ return super ().compile (node , * args , ** kwargs )
376377
377378 def collapse_group_by (self , expressions , having ):
378379 expressions = super ().collapse_group_by (expressions , having )
@@ -421,6 +422,25 @@ def _as_microsoft(self, node):
421422
422423
423424class SQLInsertCompiler (compiler .SQLInsertCompiler , SQLCompiler ):
425+ def get_returned_fields (self ):
426+ if django .VERSION >= (3 , 0 , 0 ):
427+ return self .returning_fields
428+ return self .return_id
429+
430+ def fix_auto (self , sql , opts , fields , qn ):
431+ if opts .auto_field is not None :
432+ # db_column is None if not explicitly specified by model field
433+ auto_field_column = opts .auto_field .db_column or opts .auto_field .column
434+ columns = [f .column for f in fields ]
435+ if auto_field_column in columns :
436+ id_insert_sql = []
437+ table = qn (opts .db_table )
438+ sql_format = 'SET IDENTITY_INSERT %s ON; %s; SET IDENTITY_INSERT %s OFF'
439+ for q , p in sql :
440+ id_insert_sql .append ((sql_format % (table , q , table ), p ))
441+ sql = id_insert_sql
442+
443+ return sql
424444
425445 def as_sql (self ):
426446 # We don't need quote_name_unless_alias() here, since these are all
@@ -447,38 +467,28 @@ def as_sql(self):
447467 # queries and generate their own placeholders. Doing that isn't
448468 # necessary and it should be possible to use placeholders and
449469 # expressions in bulk inserts too.
450- can_bulk = (not self .return_id and self .connection .features .has_bulk_insert ) and self .query .fields
470+ can_bulk = (not self .get_returned_fields () and self .connection .features .has_bulk_insert ) and self .query .fields
451471
452472 placeholder_rows , param_rows = self .assemble_as_sql (fields , value_rows )
453473
454- if self .return_id and self .connection .features .can_return_id_from_insert :
474+ if self .get_returned_fields () and self .connection .features .can_return_id_from_insert :
455475 result .insert (0 , 'SET NOCOUNT ON' )
456476 result .append ((values_format + ';' ) % ', ' .join (placeholder_rows [0 ]))
457477 params = [param_rows [0 ]]
458478 result .append ('SELECT CAST(SCOPE_IDENTITY() AS bigint)' )
459- return [(" " .join (result ), tuple (chain .from_iterable (params )))]
460-
461- if can_bulk :
462- result .append (self .connection .ops .bulk_insert_sql (fields , placeholder_rows ))
463- sql = [(" " .join (result ), tuple (p for ps in param_rows for p in ps ))]
479+ sql = [(" " .join (result ), tuple (chain .from_iterable (params )))]
464480 else :
465- sql = [
466- (" " .join (result + [values_format % ", " .join (p )]), vals )
467- for p , vals in zip (placeholder_rows , param_rows )
468- ]
481+ if can_bulk :
482+ result .append (self .connection .ops .bulk_insert_sql (fields , placeholder_rows ))
483+ sql = [(" " .join (result ), tuple (p for ps in param_rows for p in ps ))]
484+ else :
485+ sql = [
486+ (" " .join (result + [values_format % ", " .join (p )]), vals )
487+ for p , vals in zip (placeholder_rows , param_rows )
488+ ]
469489
470490 if self .query .fields :
471- if opts .auto_field is not None :
472- # db_column is None if not explicitly specified by model field
473- auto_field_column = opts .auto_field .db_column or opts .auto_field .column
474- columns = [f .column for f in fields ]
475- if auto_field_column in columns :
476- id_insert_sql = []
477- table = qn (opts .db_table )
478- sql_format = 'SET IDENTITY_INSERT %s ON; %s; SET IDENTITY_INSERT %s OFF'
479- for q , p in sql :
480- id_insert_sql .append ((sql_format % (table , q , table ), p ))
481- sql = id_insert_sql
491+ sql = self .fix_auto (sql , opts , fields , qn )
482492
483493 return sql
484494
0 commit comments