@@ -1185,7 +1185,15 @@ def get_update_sql(self, rows):
11851185
11861186 return self .sql , sql_args
11871187
1188- def get_upsert_sql (self , rows , unique_fields , update_fields , auto_field_name = None , only_insert = False ):
1188+ def get_upsert_sql (
1189+ self ,
1190+ rows ,
1191+ unique_fields ,
1192+ update_fields ,
1193+ auto_field_name = None ,
1194+ only_insert = False ,
1195+ return_rows = True
1196+ ):
11891197 """
11901198 Generates the postgres specific sql necessary to perform an upsert (ON CONFLICT)
11911199
@@ -1235,22 +1243,22 @@ def get_upsert_sql(self, rows, unique_fields, update_fields, auto_field_name=Non
12351243 row_values_sql = ', ' .join (row_values )
12361244
12371245 if update_fields :
1238- self .sql = 'INSERT INTO {0} ({1}) VALUES {2} ON CONFLICT ({3}) DO UPDATE SET {4} RETURNING {5}' .format (
1246+ self .sql = 'INSERT INTO {0} ({1}) VALUES {2} ON CONFLICT ({3}) DO UPDATE SET {4} {5}' .format (
12391247 self .tables [0 ].get_identifier (),
12401248 all_field_names_sql ,
12411249 row_values_sql ,
12421250 unique_field_names_sql ,
12431251 update_fields_sql ,
1244- '* '
1252+ 'RETURNING *' if return_rows else ' '
12451253 )
12461254 else :
1247- self .sql = 'INSERT INTO {0} ({1}) VALUES {2} ON CONFLICT ({3}) {4} RETURNING {5}' .format (
1255+ self .sql = 'INSERT INTO {0} ({1}) VALUES {2} ON CONFLICT ({3}) {4} {5}' .format (
12481256 self .tables [0 ].get_identifier (),
12491257 all_field_names_sql ,
12501258 row_values_sql ,
12511259 unique_field_names_sql ,
12521260 'DO UPDATE SET {0}=EXCLUDED.{0}' .format (unique_fields [0 ].column ),
1253- '* '
1261+ 'RETURNING *' if return_rows else ' '
12541262 )
12551263
12561264 return self .sql , sql_args
@@ -1737,7 +1745,13 @@ def upsert(self, rows, unique_fields, update_fields, return_rows=False, return_m
17371745 return_value = []
17381746
17391747 if rows :
1740- sql , sql_args = self .get_upsert_sql (rows , unique_fields , update_fields , auto_field_name = auto_field_name )
1748+ sql , sql_args = self .get_upsert_sql (
1749+ rows ,
1750+ unique_fields ,
1751+ update_fields ,
1752+ auto_field_name = auto_field_name ,
1753+ return_rows = return_rows or return_models
1754+ )
17411755
17421756 # get the cursor to execute the query
17431757 cursor = self .get_cursor ()
@@ -1755,6 +1769,7 @@ def upsert(self, rows, unique_fields, update_fields, return_rows=False, return_m
17551769 update_fields ,
17561770 auto_field_name = auto_field_name ,
17571771 only_insert = True ,
1772+ return_rows = return_rows or return_models
17581773 )
17591774
17601775 # get the cursor to execute the query
0 commit comments