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)
else:
tables = connection.introspection.table_names(include_views=False)
seqs = connection.introspection.sequence_list() if reset_sequences else ()
return connection.ops.sql_flush(style, tables, seqs, allow_cascade)
return connection.ops.sql_flush(
style,
tables,
reset_sequences=reset_sequences,
allow_cascade=allow_cascade,
)
def emit_pre_migrate_signal(verbosity, interactive, db, **kwargs):

View File

@ -382,16 +382,18 @@ class BaseDatabaseOperations:
"""
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
the given database tables (without actually removing the tables
themselves) and the SQL statements required to reset the sequences
passed in `sequences`.
themselves).
The `style` argument is a Style object as returned by either
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
to tables with foreign keys pointing the tables being truncated.
PostgreSQL requires a cascade even if these tables are empty.

View File

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

View File

@ -404,7 +404,7 @@ END;
# Django's test suite.
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:
return []
@ -446,9 +446,15 @@ END;
style.SQL_FIELD(self.quote_name(constraint)),
) for table, constraint in constraints
]
# 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))
if reset_sequences:
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
def sequence_reset_by_name_sql(self, style, sequences):

View File

@ -117,28 +117,21 @@ class DatabaseOperations(BaseDatabaseOperations):
def set_time_zone_sql(self):
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:
return []
# Perform a single SQL 'TRUNCATE x, y, z...;' statement. It allows us
# to truncate tables referenced by a foreign key in any other table.
tables_sql = ', '.join(
style.SQL_FIELD(self.quote_name(table)) for table in tables
)
sql_parts = [
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:
sql = ['%s %s %s;' % (
style.SQL_KEYWORD('TRUNCATE'),
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
sql_parts.append(style.SQL_KEYWORD('CASCADE'))
return ['%s;' % ' '.join(sql_parts)]
def sequence_reset_by_name_sql(self, style, sequences):
# 'ALTER SEQUENCE sequence_name RESTART WITH 1;'... style SQL statements

View File

@ -196,7 +196,7 @@ class DatabaseOperations(BaseDatabaseOperations):
# Django's test suite.
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:
# Simulate TRUNCATE CASCADE by recursively collecting the tables
# referencing the tables to be flushed.

View File

@ -519,6 +519,13 @@ backends.
* ``DatabaseClient.runshell()`` now requires an additional ``parameters``
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
--------------------------------

View File

@ -43,7 +43,7 @@ class SimpleDatabaseOperationTests(SimpleTestCase):
def test_sql_flush(self):
msg = 'subclasses of BaseDatabaseOperations must provide a sql_flush() method'
with self.assertRaisesMessage(NotImplementedError, msg):
self.ops.sql_flush(None, None, None)
self.ops.sql_flush(None, None)
def test_pk_default_value(self):
self.assertEqual(self.ops.pk_default_value(), 'DEFAULT')
@ -154,7 +154,7 @@ class SqlFlushTests(TransactionTestCase):
available_apps = ['backends']
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):
with transaction.atomic():
@ -169,12 +169,7 @@ class SqlFlushTests(TransactionTestCase):
sql_list = connection.ops.sql_flush(
no_style(),
[Author._meta.db_table, Book._meta.db_table],
[
{
'table': Author._meta.db_table,
'column': Author._meta.pk.db_column,
},
],
reset_sequences=True,
allow_cascade=True,
)
connection.ops.execute_sql_flush(connection.alias, sql_list)
@ -185,3 +180,5 @@ class SqlFlushTests(TransactionTestCase):
if connection.features.supports_sequence_reset:
author = Author.objects.create(name='F. Scott Fitzgerald')
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.test import SimpleTestCase
from ..models import Person, Square, Tag
from ..models import Person, Tag
@unittest.skipUnless(connection.vendor == 'mysql', 'MySQL tests.')
@ -13,50 +13,35 @@ class MySQLOperationsTests(SimpleTestCase):
# allow_cascade doesn't change statements on MySQL.
for allow_cascade in [False, True]:
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(
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_tag`;',
'SET FOREIGN_KEY_CHECKS = 1;',
],
)
self.assertEqual(statements[-1], 'SET FOREIGN_KEY_CHECKS = 1;')
def test_sql_flush_sequences(self):
# allow_cascade doesn't change statements on MySQL.
for allow_cascade in [False, True]:
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(
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_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(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
[],
)
# The tables and constraints are processed in an unordered set.
self.assertEqual(
@ -56,7 +55,6 @@ class OperationsTests(unittest.TestCase):
statements = connection.ops.sql_flush(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
[],
allow_cascade=True,
)
# The tables and constraints are processed in an unordered set.
@ -83,16 +81,7 @@ class OperationsTests(unittest.TestCase):
statements = connection.ops.sql_flush(
no_style(),
[Person._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,
},
],
reset_sequences=True,
)
# The tables and constraints are processed in an unordered set.
self.assertEqual(
@ -121,16 +110,7 @@ class OperationsTests(unittest.TestCase):
statements = connection.ops.sql_flush(
no_style(),
[Person._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,
},
],
reset_sequences=True,
allow_cascade=True,
)
# The tables and constraints are processed in an unordered set.
@ -153,6 +133,7 @@ class OperationsTests(unittest.TestCase):
'"BACKENDS__PERSON_ID_1DD5E829_F";',
)
# Sequences.
self.assertEqual(len(statements[5:]), 2)
self.assertEqual(len(statements[5:]), 3)
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(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
[],
),
['TRUNCATE "backends_person", "backends_tag";'],
)
@ -24,61 +23,28 @@ class PostgreSQLOperationsTests(SimpleTestCase):
connection.ops.sql_flush(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
[],
allow_cascade=True,
),
['TRUNCATE "backends_person", "backends_tag" CASCADE;'],
)
def test_sql_flush_sequences(self):
sequence_reset_sql = (
"SELECT setval(pg_get_serial_sequence('%s','id'), 1, false);"
)
self.assertEqual(
connection.ops.sql_flush(
no_style(),
[Person._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,
},
],
reset_sequences=True,
),
[
'TRUNCATE "backends_person", "backends_tag";',
sequence_reset_sql % '"backends_person"',
sequence_reset_sql % '"backends_tag"',
],
['TRUNCATE "backends_person", "backends_tag" RESTART IDENTITY;'],
)
def test_sql_flush_sequences_allow_cascade(self):
sequence_reset_sql = (
"SELECT setval(pg_get_serial_sequence('%s','id'), 1, false);"
)
self.assertEqual(
connection.ops.sql_flush(
no_style(),
[Person._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,
},
],
reset_sequences=True,
allow_cascade=True,
),
[
'TRUNCATE "backends_person", "backends_tag" CASCADE;',
sequence_reset_sql % '"backends_person"',
sequence_reset_sql % '"backends_tag"',
],
['TRUNCATE "backends_person", "backends_tag" RESTART IDENTITY CASCADE;'],
)

View File

@ -14,7 +14,6 @@ class SQLiteOperationsTests(TestCase):
connection.ops.sql_flush(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
[],
),
[
'DELETE FROM "backends_person";',
@ -26,7 +25,6 @@ class SQLiteOperationsTests(TestCase):
statements = connection.ops.sql_flush(
no_style(),
[Person._meta.db_table, Tag._meta.db_table],
[],
allow_cascade=True,
)
self.assertEqual(
@ -47,16 +45,7 @@ class SQLiteOperationsTests(TestCase):
connection.ops.sql_flush(
no_style(),
[Person._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,
},
],
reset_sequences=True,
),
[
'DELETE FROM "backends_person";',
@ -69,16 +58,7 @@ class SQLiteOperationsTests(TestCase):
statements = connection.ops.sql_flush(
no_style(),
[Person._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,
},
],
reset_sequences=True,
allow_cascade=True,
)
self.assertEqual(

View File

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