Fixed #31473 -- Made sql_flush() use RESTART IDENTITY to reset sequences on PostgreSQL.

The sql_flush() positional argument sequences is replaced by the boolean
keyword-only argument reset_sequences. This ensures that the old
function signature can't be used by mistake when upgrading Django. When
the new argument is True, the sequences of the truncated tables will
reset. Using a single boolean value, rather than a list, allows making a
binary yes/no choice as to whether to reset all sequences rather than a
working on a completely different set.
This commit is contained in:
Jon Dufresne 2020-04-15 02:20:46 -07:00 committed by Mariusz Felisiak
parent 8005829bb9
commit 75410228df
13 changed files with 92 additions and 176 deletions

View File

@ -13,8 +13,12 @@ def sql_flush(style, connection, only_django=False, reset_sequences=True, allow_
tables = connection.introspection.django_table_names(only_existing=True, include_views=False) tables = connection.introspection.django_table_names(only_existing=True, include_views=False)
else: else:
tables = connection.introspection.table_names(include_views=False) tables = connection.introspection.table_names(include_views=False)
seqs = connection.introspection.sequence_list() if reset_sequences else () return connection.ops.sql_flush(
return connection.ops.sql_flush(style, tables, seqs, allow_cascade) style,
tables,
reset_sequences=reset_sequences,
allow_cascade=allow_cascade,
)
def emit_pre_migrate_signal(verbosity, interactive, db, **kwargs): def emit_pre_migrate_signal(verbosity, interactive, db, **kwargs):

View File

@ -382,16 +382,18 @@ class BaseDatabaseOperations:
""" """
return '' return ''
def sql_flush(self, style, tables, sequences, allow_cascade=False): def sql_flush(self, style, tables, *, reset_sequences=False, allow_cascade=False):
""" """
Return a list of SQL statements required to remove all data from Return a list of SQL statements required to remove all data from
the given database tables (without actually removing the tables the given database tables (without actually removing the tables
themselves) and the SQL statements required to reset the sequences themselves).
passed in `sequences`.
The `style` argument is a Style object as returned by either The `style` argument is a Style object as returned by either
color_style() or no_style() in django.core.management.color. color_style() or no_style() in django.core.management.color.
If `reset_sequences` is True, the list includes SQL statements required
to reset the sequences.
The `allow_cascade` argument determines whether truncation may cascade The `allow_cascade` argument determines whether truncation may cascade
to tables with foreign keys pointing the tables being truncated. to tables with foreign keys pointing the tables being truncated.
PostgreSQL requires a cascade even if these tables are empty. PostgreSQL requires a cascade even if these tables are empty.

View File

@ -193,29 +193,30 @@ class DatabaseOperations(BaseDatabaseOperations):
] ]
return 'RETURNING %s' % ', '.join(columns), () return 'RETURNING %s' % ', '.join(columns), ()
def sql_flush(self, style, tables, sequences, allow_cascade=False): def sql_flush(self, style, tables, *, reset_sequences=False, allow_cascade=False):
if not tables: if not tables:
return [] return []
sql = ['SET FOREIGN_KEY_CHECKS = 0;'] sql = ['SET FOREIGN_KEY_CHECKS = 0;']
tables = set(tables) if reset_sequences:
with_sequences = set(s['table'] for s in sequences) # It's faster to TRUNCATE tables that require a sequence reset
# It's faster to TRUNCATE tables that require a sequence reset since # since ALTER TABLE AUTO_INCREMENT is slower than TRUNCATE.
# ALTER TABLE AUTO_INCREMENT is slower than TRUNCATE. sql.extend(
sql.extend( '%s %s;' % (
'%s %s;' % ( style.SQL_KEYWORD('TRUNCATE'),
style.SQL_KEYWORD('TRUNCATE'), style.SQL_FIELD(self.quote_name(table_name)),
style.SQL_FIELD(self.quote_name(table_name)), ) for table_name in tables
) for table_name in tables.intersection(with_sequences) )
) else:
# Otherwise issue a simple DELETE since it's faster than TRUNCATE # Otherwise issue a simple DELETE since it's faster than TRUNCATE
# and preserves sequences. # and preserves sequences.
sql.extend( sql.extend(
'%s %s %s;' % ( '%s %s %s;' % (
style.SQL_KEYWORD('DELETE'), style.SQL_KEYWORD('DELETE'),
style.SQL_KEYWORD('FROM'), style.SQL_KEYWORD('FROM'),
style.SQL_FIELD(self.quote_name(table_name)), style.SQL_FIELD(self.quote_name(table_name)),
) for table_name in tables.difference(with_sequences) ) for table_name in tables
) )
sql.append('SET FOREIGN_KEY_CHECKS = 1;') sql.append('SET FOREIGN_KEY_CHECKS = 1;')
return sql return sql

View File

@ -404,7 +404,7 @@ END;
# Django's test suite. # Django's test suite.
return lru_cache(maxsize=512)(self.__foreign_key_constraints) return lru_cache(maxsize=512)(self.__foreign_key_constraints)
def sql_flush(self, style, tables, sequences, allow_cascade=False): def sql_flush(self, style, tables, *, reset_sequences=False, allow_cascade=False):
if not tables: if not tables:
return [] return []
@ -446,9 +446,15 @@ END;
style.SQL_FIELD(self.quote_name(constraint)), style.SQL_FIELD(self.quote_name(constraint)),
) for table, constraint in constraints ) for table, constraint in constraints
] ]
# Since we've just deleted all the rows, running our sequence ALTER if reset_sequences:
# code will reset the sequence to 0. sequences = [
sql.extend(self.sequence_reset_by_name_sql(style, sequences)) sequence
for sequence in self.connection.introspection.sequence_list()
if sequence['table'].upper() in truncated_tables
]
# Since we've just deleted all the rows, running our sequence ALTER
# code will reset the sequence to 0.
sql.extend(self.sequence_reset_by_name_sql(style, sequences))
return sql return sql
def sequence_reset_by_name_sql(self, style, sequences): def sequence_reset_by_name_sql(self, style, sequences):

View File

@ -117,28 +117,21 @@ class DatabaseOperations(BaseDatabaseOperations):
def set_time_zone_sql(self): def set_time_zone_sql(self):
return "SET TIME ZONE %s" return "SET TIME ZONE %s"
def sql_flush(self, style, tables, sequences, allow_cascade=False): def sql_flush(self, style, tables, *, reset_sequences=False, allow_cascade=False):
if not tables: if not tables:
return [] return []
# Perform a single SQL 'TRUNCATE x, y, z...;' statement. It allows us # Perform a single SQL 'TRUNCATE x, y, z...;' statement. It allows us
# to truncate tables referenced by a foreign key in any other table. # to truncate tables referenced by a foreign key in any other table.
tables_sql = ', '.join( sql_parts = [
style.SQL_FIELD(self.quote_name(table)) for table in tables style.SQL_KEYWORD('TRUNCATE'),
) ', '.join(style.SQL_FIELD(self.quote_name(table)) for table in tables),
]
if reset_sequences:
sql_parts.append(style.SQL_KEYWORD('RESTART IDENTITY'))
if allow_cascade: if allow_cascade:
sql = ['%s %s %s;' % ( sql_parts.append(style.SQL_KEYWORD('CASCADE'))
style.SQL_KEYWORD('TRUNCATE'), return ['%s;' % ' '.join(sql_parts)]
tables_sql,
style.SQL_KEYWORD('CASCADE'),
)]
else:
sql = ['%s %s;' % (
style.SQL_KEYWORD('TRUNCATE'),
tables_sql,
)]
sql.extend(self.sequence_reset_by_name_sql(style, sequences))
return sql
def sequence_reset_by_name_sql(self, style, sequences): def sequence_reset_by_name_sql(self, style, sequences):
# 'ALTER SEQUENCE sequence_name RESTART WITH 1;'... style SQL statements # 'ALTER SEQUENCE sequence_name RESTART WITH 1;'... style SQL statements

View File

@ -196,7 +196,7 @@ class DatabaseOperations(BaseDatabaseOperations):
# Django's test suite. # Django's test suite.
return lru_cache(maxsize=512)(self.__references_graph) return lru_cache(maxsize=512)(self.__references_graph)
def sql_flush(self, style, tables, sequences, allow_cascade=False): def sql_flush(self, style, tables, *, reset_sequences=False, allow_cascade=False):
if tables and allow_cascade: if tables and allow_cascade:
# Simulate TRUNCATE CASCADE by recursively collecting the tables # Simulate TRUNCATE CASCADE by recursively collecting the tables
# referencing the tables to be flushed. # referencing the tables to be flushed.

View File

@ -519,6 +519,13 @@ backends.
* ``DatabaseClient.runshell()`` now requires an additional ``parameters`` * ``DatabaseClient.runshell()`` now requires an additional ``parameters``
argument as a list of extra arguments to pass on to the command-line client. argument as a list of extra arguments to pass on to the command-line client.
* The ``sequences`` positional argument of ``DatabaseOperations.sql_flush()``
is replaced by the boolean keyword-only argument ``reset_sequences``. If
``True``, the sequences of the truncated tables will be reset.
* The ``allow_cascade`` argument of ``DatabaseOperations.sql_flush()`` is now a
keyword-only argument.
Dropped support for MariaDB 10.1 Dropped support for MariaDB 10.1
-------------------------------- --------------------------------

View File

@ -43,7 +43,7 @@ class SimpleDatabaseOperationTests(SimpleTestCase):
def test_sql_flush(self): def test_sql_flush(self):
msg = 'subclasses of BaseDatabaseOperations must provide a sql_flush() method' msg = 'subclasses of BaseDatabaseOperations must provide a sql_flush() method'
with self.assertRaisesMessage(NotImplementedError, msg): with self.assertRaisesMessage(NotImplementedError, msg):
self.ops.sql_flush(None, None, None) self.ops.sql_flush(None, None)
def test_pk_default_value(self): def test_pk_default_value(self):
self.assertEqual(self.ops.pk_default_value(), 'DEFAULT') self.assertEqual(self.ops.pk_default_value(), 'DEFAULT')
@ -154,7 +154,7 @@ class SqlFlushTests(TransactionTestCase):
available_apps = ['backends'] available_apps = ['backends']
def test_sql_flush_no_tables(self): def test_sql_flush_no_tables(self):
self.assertEqual(connection.ops.sql_flush(no_style(), [], []), []) self.assertEqual(connection.ops.sql_flush(no_style(), []), [])
def test_execute_sql_flush_statements(self): def test_execute_sql_flush_statements(self):
with transaction.atomic(): with transaction.atomic():
@ -169,12 +169,7 @@ class SqlFlushTests(TransactionTestCase):
sql_list = connection.ops.sql_flush( sql_list = connection.ops.sql_flush(
no_style(), no_style(),
[Author._meta.db_table, Book._meta.db_table], [Author._meta.db_table, Book._meta.db_table],
[ reset_sequences=True,
{
'table': Author._meta.db_table,
'column': Author._meta.pk.db_column,
},
],
allow_cascade=True, allow_cascade=True,
) )
connection.ops.execute_sql_flush(connection.alias, sql_list) connection.ops.execute_sql_flush(connection.alias, sql_list)
@ -185,3 +180,5 @@ class SqlFlushTests(TransactionTestCase):
if connection.features.supports_sequence_reset: if connection.features.supports_sequence_reset:
author = Author.objects.create(name='F. Scott Fitzgerald') author = Author.objects.create(name='F. Scott Fitzgerald')
self.assertEqual(author.pk, 1) self.assertEqual(author.pk, 1)
book = Book.objects.create(author=author)
self.assertEqual(book.pk, 1)

View File

@ -4,7 +4,7 @@ from django.core.management.color import no_style
from django.db import connection from django.db import connection
from django.test import SimpleTestCase from django.test import SimpleTestCase
from ..models import Person, Square, Tag from ..models import Person, Tag
@unittest.skipUnless(connection.vendor == 'mysql', 'MySQL tests.') @unittest.skipUnless(connection.vendor == 'mysql', 'MySQL tests.')
@ -13,50 +13,35 @@ class MySQLOperationsTests(SimpleTestCase):
# allow_cascade doesn't change statements on MySQL. # allow_cascade doesn't change statements on MySQL.
for allow_cascade in [False, True]: for allow_cascade in [False, True]:
with self.subTest(allow_cascade=allow_cascade): with self.subTest(allow_cascade=allow_cascade):
statements = connection.ops.sql_flush(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
[],
allow_cascade=allow_cascade,
)
self.assertEqual(statements[0], 'SET FOREIGN_KEY_CHECKS = 0;')
# The tables are processed in an unordered set.
self.assertEqual( self.assertEqual(
sorted(statements[1:-1]), connection.ops.sql_flush(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
allow_cascade=allow_cascade,
),
[ [
'SET FOREIGN_KEY_CHECKS = 0;',
'DELETE FROM `backends_person`;', 'DELETE FROM `backends_person`;',
'DELETE FROM `backends_tag`;', 'DELETE FROM `backends_tag`;',
'SET FOREIGN_KEY_CHECKS = 1;',
], ],
) )
self.assertEqual(statements[-1], 'SET FOREIGN_KEY_CHECKS = 1;')
def test_sql_flush_sequences(self): def test_sql_flush_sequences(self):
# allow_cascade doesn't change statements on MySQL. # allow_cascade doesn't change statements on MySQL.
for allow_cascade in [False, True]: for allow_cascade in [False, True]:
with self.subTest(allow_cascade=allow_cascade): with self.subTest(allow_cascade=allow_cascade):
statements = connection.ops.sql_flush(
no_style(),
[Person._meta.db_table, Square._meta.db_table, Tag._meta.db_table],
[
{
'table': Person._meta.db_table,
'column': Person._meta.pk.db_column,
},
{
'table': Tag._meta.db_table,
'column': Tag._meta.pk.db_column,
},
],
allow_cascade=allow_cascade,
)
self.assertEqual(statements[0], 'SET FOREIGN_KEY_CHECKS = 0;')
# The tables are processed in an unordered set.
self.assertEqual( self.assertEqual(
sorted(statements[1:-1]), connection.ops.sql_flush(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
reset_sequences=True,
allow_cascade=allow_cascade,
),
[ [
'DELETE FROM `backends_square`;', 'SET FOREIGN_KEY_CHECKS = 0;',
'TRUNCATE `backends_person`;', 'TRUNCATE `backends_person`;',
'TRUNCATE `backends_tag`;', 'TRUNCATE `backends_tag`;',
'SET FOREIGN_KEY_CHECKS = 1;',
], ],
) )
self.assertEqual(statements[-1], 'SET FOREIGN_KEY_CHECKS = 1;')

View File

@ -31,7 +31,6 @@ class OperationsTests(unittest.TestCase):
statements = connection.ops.sql_flush( statements = connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[],
) )
# The tables and constraints are processed in an unordered set. # The tables and constraints are processed in an unordered set.
self.assertEqual( self.assertEqual(
@ -56,7 +55,6 @@ class OperationsTests(unittest.TestCase):
statements = connection.ops.sql_flush( statements = connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[],
allow_cascade=True, allow_cascade=True,
) )
# The tables and constraints are processed in an unordered set. # The tables and constraints are processed in an unordered set.
@ -83,16 +81,7 @@ class OperationsTests(unittest.TestCase):
statements = connection.ops.sql_flush( statements = connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[ reset_sequences=True,
{
'table': Person._meta.db_table,
'column': Person._meta.pk.db_column,
},
{
'table': Tag._meta.db_table,
'column': Tag._meta.pk.db_column,
},
],
) )
# The tables and constraints are processed in an unordered set. # The tables and constraints are processed in an unordered set.
self.assertEqual( self.assertEqual(
@ -121,16 +110,7 @@ class OperationsTests(unittest.TestCase):
statements = connection.ops.sql_flush( statements = connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[ reset_sequences=True,
{
'table': Person._meta.db_table,
'column': Person._meta.pk.db_column,
},
{
'table': Tag._meta.db_table,
'column': Tag._meta.pk.db_column,
},
],
allow_cascade=True, allow_cascade=True,
) )
# The tables and constraints are processed in an unordered set. # The tables and constraints are processed in an unordered set.
@ -153,6 +133,7 @@ class OperationsTests(unittest.TestCase):
'"BACKENDS__PERSON_ID_1DD5E829_F";', '"BACKENDS__PERSON_ID_1DD5E829_F";',
) )
# Sequences. # Sequences.
self.assertEqual(len(statements[5:]), 2) self.assertEqual(len(statements[5:]), 3)
self.assertIn('BACKENDS_PERSON_SQ', statements[5]) self.assertIn('BACKENDS_PERSON_SQ', statements[5])
self.assertIn('BACKENDS_TAG_SQ', statements[6]) self.assertIn('BACKENDS_VERYLONGMODELN7BE2_SQ', statements[6])
self.assertIn('BACKENDS_TAG_SQ', statements[7])

View File

@ -14,7 +14,6 @@ class PostgreSQLOperationsTests(SimpleTestCase):
connection.ops.sql_flush( connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[],
), ),
['TRUNCATE "backends_person", "backends_tag";'], ['TRUNCATE "backends_person", "backends_tag";'],
) )
@ -24,61 +23,28 @@ class PostgreSQLOperationsTests(SimpleTestCase):
connection.ops.sql_flush( connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[],
allow_cascade=True, allow_cascade=True,
), ),
['TRUNCATE "backends_person", "backends_tag" CASCADE;'], ['TRUNCATE "backends_person", "backends_tag" CASCADE;'],
) )
def test_sql_flush_sequences(self): def test_sql_flush_sequences(self):
sequence_reset_sql = (
"SELECT setval(pg_get_serial_sequence('%s','id'), 1, false);"
)
self.assertEqual( self.assertEqual(
connection.ops.sql_flush( connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[ reset_sequences=True,
{
'table': Person._meta.db_table,
'column': Person._meta.pk.db_column,
},
{
'table': Tag._meta.db_table,
'column': Tag._meta.pk.db_column,
},
],
), ),
[ ['TRUNCATE "backends_person", "backends_tag" RESTART IDENTITY;'],
'TRUNCATE "backends_person", "backends_tag";',
sequence_reset_sql % '"backends_person"',
sequence_reset_sql % '"backends_tag"',
],
) )
def test_sql_flush_sequences_allow_cascade(self): def test_sql_flush_sequences_allow_cascade(self):
sequence_reset_sql = (
"SELECT setval(pg_get_serial_sequence('%s','id'), 1, false);"
)
self.assertEqual( self.assertEqual(
connection.ops.sql_flush( connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[ reset_sequences=True,
{
'table': Person._meta.db_table,
'column': Person._meta.pk.db_column,
},
{
'table': Tag._meta.db_table,
'column': Tag._meta.pk.db_column,
},
],
allow_cascade=True, allow_cascade=True,
), ),
[ ['TRUNCATE "backends_person", "backends_tag" RESTART IDENTITY CASCADE;'],
'TRUNCATE "backends_person", "backends_tag" CASCADE;',
sequence_reset_sql % '"backends_person"',
sequence_reset_sql % '"backends_tag"',
],
) )

View File

@ -14,7 +14,6 @@ class SQLiteOperationsTests(TestCase):
connection.ops.sql_flush( connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[],
), ),
[ [
'DELETE FROM "backends_person";', 'DELETE FROM "backends_person";',
@ -26,7 +25,6 @@ class SQLiteOperationsTests(TestCase):
statements = connection.ops.sql_flush( statements = connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[],
allow_cascade=True, allow_cascade=True,
) )
self.assertEqual( self.assertEqual(
@ -47,16 +45,7 @@ class SQLiteOperationsTests(TestCase):
connection.ops.sql_flush( connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[ reset_sequences=True,
{
'table': Person._meta.db_table,
'column': Person._meta.pk.db_column,
},
{
'table': Tag._meta.db_table,
'column': Tag._meta.pk.db_column,
},
],
), ),
[ [
'DELETE FROM "backends_person";', 'DELETE FROM "backends_person";',
@ -69,16 +58,7 @@ class SQLiteOperationsTests(TestCase):
statements = connection.ops.sql_flush( statements = connection.ops.sql_flush(
no_style(), no_style(),
[Person._meta.db_table, Tag._meta.db_table], [Person._meta.db_table, Tag._meta.db_table],
[ reset_sequences=True,
{
'table': Person._meta.db_table,
'column': Person._meta.pk.db_column,
},
{
'table': Tag._meta.db_table,
'column': Tag._meta.pk.db_column,
},
],
allow_cascade=True, allow_cascade=True,
) )
self.assertEqual( self.assertEqual(

View File

@ -161,13 +161,7 @@ class LongNameTest(TransactionTestCase):
VLM._meta.db_table, VLM._meta.db_table,
VLM_m2m._meta.db_table, VLM_m2m._meta.db_table,
] ]
sequences = [ sql_list = connection.ops.sql_flush(no_style(), tables, reset_sequences=True)
{
'column': VLM._meta.pk.column,
'table': VLM._meta.db_table
},
]
sql_list = connection.ops.sql_flush(no_style(), tables, sequences)
with connection.cursor() as cursor: with connection.cursor() as cursor:
for statement in sql_list: for statement in sql_list:
cursor.execute(statement) cursor.execute(statement)