Fixed #31653 -- Added AddConstraintNotValid()/ValidateConstraint() operations for PostgreSQL.

django/contrib/postgres/operations.py

@@ -2,8 +2,9 @@ from django.contrib.postgres.signals import (
     get_citext_oids, get_hstore_oids, register_type_handlers,
 )
 from django.db import NotSupportedError, router
-from django.db.migrations import AddIndex, RemoveIndex
+from django.db.migrations import AddConstraint, AddIndex, RemoveIndex
 from django.db.migrations.operations.base import Operation
+from django.db.models.constraints import CheckConstraint
 
 
 class CreateExtension(Operation):
@@ -256,3 +257,73 @@ class RemoveCollation(CollationOperation):
     @property
     def migration_name_fragment(self):
         return 'remove_collation_%s' % self.name.lower()
+
+
+class AddConstraintNotValid(AddConstraint):
+    """
+    Add a table constraint without enforcing validation, using PostgreSQL's
+    NOT VALID syntax.
+    """
+
+    def __init__(self, model_name, constraint):
+        if not isinstance(constraint, CheckConstraint):
+            raise TypeError(
+                'AddConstraintNotValid.constraint must be a check constraint.'
+            )
+        super().__init__(model_name, constraint)
+
+    def describe(self):
+        return 'Create not valid constraint %s on model %s' % (
+            self.constraint.name,
+            self.model_name,
+        )
+
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+        model = from_state.apps.get_model(app_label, self.model_name)
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            constraint_sql = self.constraint.create_sql(model, schema_editor)
+            if constraint_sql:
+                # Constraint.create_sql returns interpolated SQL which makes
+                # params=None a necessity to avoid escaping attempts on
+                # execution.
+                schema_editor.execute(str(constraint_sql) + ' NOT VALID', params=None)
+
+    @property
+    def migration_name_fragment(self):
+        return super().migration_name_fragment + '_not_valid'
+
+
+class ValidateConstraint(Operation):
+    """Validate a table NOT VALID constraint."""
+
+    def __init__(self, model_name, name):
+        self.model_name = model_name
+        self.name = name
+
+    def describe(self):
+        return 'Validate constraint %s on model %s' % (self.name, self.model_name)
+
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+        model = from_state.apps.get_model(app_label, self.model_name)
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            schema_editor.execute('ALTER TABLE %s VALIDATE CONSTRAINT %s' % (
+                schema_editor.quote_name(model._meta.db_table),
+                schema_editor.quote_name(self.name),
+            ))
+
+    def database_backwards(self, app_label, schema_editor, from_state, to_state):
+        # PostgreSQL does not provide a way to make a constraint invalid.
+        pass
+
+    def state_forwards(self, app_label, state):
+        pass
+
+    @property
+    def migration_name_fragment(self):
+        return '%s_validate_%s' % (self.model_name.lower(), self.name.lower())
+
+    def deconstruct(self):
+        return self.__class__.__name__, [], {
+            'model_name': self.model_name,
+            'name': self.name,
+        }

docs/ref/contrib/postgres/operations.txt

@@ -188,3 +188,39 @@ database.
 
     The ``CONCURRENTLY`` option is not supported inside a transaction (see
     :ref:`non-atomic migration <non-atomic-migrations>`).
+
+Adding constraints without enforcing validation
+===============================================
+
+.. versionadded:: 4.0
+
+PostgreSQL supports the ``NOT VALID`` option with the ``ADD CONSTRAINT``
+statement to add check constraints without enforcing validation on existing
+rows. This option is useful if you want to skip the potentially lengthy scan of
+the table to verify that all existing rows satisfy the constraint.
+
+To validate check constraints created with the ``NOT VALID`` option at a later
+point of time, use the
+:class:`~django.contrib.postgres.operations.ValidateConstraint` operation.
+
+See `the PostgreSQL documentation <https://www.postgresql.org/docs/current/
+sql-altertable.html#SQL-ALTERTABLE-NOTES>`__ for more details.
+
+.. class:: AddConstraintNotValid(model_name, constraint)
+
+    Like :class:`~django.db.migrations.operations.AddConstraint`, but avoids
+    validating the constraint on existing rows.
+
+.. class:: ValidateConstraint(model_name, name)
+
+    Scans through the table and validates the given check constraint on
+    existing rows.
+
+.. note::
+
+    ``AddConstraintNotValid`` and ``ValidateConstraint`` operations should be
+    performed in two separate migrations. Performing both operations in the
+    same atomic migration has the same effect as
+    :class:`~django.db.migrations.operations.AddConstraint`, whereas performing
+    them in a single non-atomic migration, may leave your database in an
+    inconsistent state if the ``ValidateConstraint`` operation fails.

docs/releases/4.0.txt

@@ -122,6 +122,15 @@ Minor features
 * The PostgreSQL backend now supports connecting by a service name. See
   :ref:`postgresql-connection-settings` for more details.
 
+* The new :class:`~django.contrib.postgres.operations.AddConstraintNotValid`
+  operation allows creating check constraints on PostgreSQL without verifying
+  that all existing rows satisfy the new constraint.
+
+* The new :class:`~django.contrib.postgres.operations.ValidateConstraint`
+  operation allows validating check constraints which were created using
+  :class:`~django.contrib.postgres.operations.AddConstraintNotValid` on
+  PostgreSQL.
+
 :mod:`django.contrib.redirects`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

tests/postgres_tests/test_operations.py

@@ -3,9 +3,11 @@ from unittest import mock
 
 from migrations.test_base import OperationTestBase
 
-from django.db import NotSupportedError, connection
+from django.db import (
+    IntegrityError, NotSupportedError, connection, transaction,
+)
 from django.db.migrations.state import ProjectState
-from django.db.models import Index
+from django.db.models import CheckConstraint, Index, Q, UniqueConstraint
 from django.db.utils import ProgrammingError
 from django.test import modify_settings, override_settings, skipUnlessDBFeature
 from django.test.utils import CaptureQueriesContext
@@ -15,8 +17,9 @@ from . import PostgreSQLTestCase
 try:
     from django.contrib.postgres.indexes import BrinIndex, BTreeIndex
     from django.contrib.postgres.operations import (
-        AddIndexConcurrently, BloomExtension, CreateCollation, CreateExtension,
-        RemoveCollation, RemoveIndexConcurrently,
+        AddConstraintNotValid, AddIndexConcurrently, BloomExtension,
+        CreateCollation, CreateExtension, RemoveCollation,
+        RemoveIndexConcurrently, ValidateConstraint,
     )
 except ImportError:
     pass
@@ -392,3 +395,102 @@ class RemoveCollationTests(PostgreSQLTestCase):
         self.assertEqual(name, 'RemoveCollation')
         self.assertEqual(args, [])
         self.assertEqual(kwargs, {'name': 'C_test', 'locale': 'C'})
+
+
+@unittest.skipUnless(connection.vendor == 'postgresql', 'PostgreSQL specific tests.')
+@modify_settings(INSTALLED_APPS={'append': 'migrations'})
+class AddConstraintNotValidTests(OperationTestBase):
+    app_label = 'test_add_constraint_not_valid'
+
+    def test_non_check_constraint_not_supported(self):
+        constraint = UniqueConstraint(fields=['pink'], name='pony_pink_uniq')
+        msg = 'AddConstraintNotValid.constraint must be a check constraint.'
+        with self.assertRaisesMessage(TypeError, msg):
+            AddConstraintNotValid(model_name='pony', constraint=constraint)
+
+    def test_add(self):
+        table_name = f'{self.app_label}_pony'
+        constraint_name = 'pony_pink_gte_check'
+        constraint = CheckConstraint(check=Q(pink__gte=4), name=constraint_name)
+        operation = AddConstraintNotValid('Pony', constraint=constraint)
+        project_state, new_state = self.make_test_state(self.app_label, operation)
+        self.assertEqual(
+            operation.describe(),
+            f'Create not valid constraint {constraint_name} on model Pony',
+        )
+        self.assertEqual(
+            operation.migration_name_fragment,
+            f'pony_{constraint_name}_not_valid',
+        )
+        self.assertEqual(
+            len(new_state.models[self.app_label, 'pony'].options['constraints']),
+            1,
+        )
+        self.assertConstraintNotExists(table_name, constraint_name)
+        Pony = new_state.apps.get_model(self.app_label, 'Pony')
+        self.assertEqual(len(Pony._meta.constraints), 1)
+        Pony.objects.create(pink=2, weight=1.0)
+        # Add constraint.
+        with connection.schema_editor(atomic=True) as editor:
+            operation.database_forwards(self.app_label, editor, project_state, new_state)
+        msg = f'check constraint "{constraint_name}"'
+        with self.assertRaisesMessage(IntegrityError, msg), transaction.atomic():
+            Pony.objects.create(pink=3, weight=1.0)
+        self.assertConstraintExists(table_name, constraint_name)
+        # Reversal.
+        with connection.schema_editor(atomic=True) as editor:
+            operation.database_backwards(self.app_label, editor, project_state, new_state)
+        self.assertConstraintNotExists(table_name, constraint_name)
+        Pony.objects.create(pink=3, weight=1.0)
+        # Deconstruction.
+        name, args, kwargs = operation.deconstruct()
+        self.assertEqual(name, 'AddConstraintNotValid')
+        self.assertEqual(args, [])
+        self.assertEqual(kwargs, {'model_name': 'Pony', 'constraint': constraint})
+
+
+@unittest.skipUnless(connection.vendor == 'postgresql', 'PostgreSQL specific tests.')
+@modify_settings(INSTALLED_APPS={'append': 'migrations'})
+class ValidateConstraintTests(OperationTestBase):
+    app_label = 'test_validate_constraint'
+
+    def test_validate(self):
+        constraint_name = 'pony_pink_gte_check'
+        constraint = CheckConstraint(check=Q(pink__gte=4), name=constraint_name)
+        operation = AddConstraintNotValid('Pony', constraint=constraint)
+        project_state, new_state = self.make_test_state(self.app_label, operation)
+        Pony = new_state.apps.get_model(self.app_label, 'Pony')
+        obj = Pony.objects.create(pink=2, weight=1.0)
+        # Add constraint.
+        with connection.schema_editor(atomic=True) as editor:
+            operation.database_forwards(self.app_label, editor, project_state, new_state)
+        project_state = new_state
+        new_state = new_state.clone()
+        operation = ValidateConstraint('Pony', name=constraint_name)
+        operation.state_forwards(self.app_label, new_state)
+        self.assertEqual(
+            operation.describe(),
+            f'Validate constraint {constraint_name} on model Pony',
+        )
+        self.assertEqual(
+            operation.migration_name_fragment,
+            f'pony_validate_{constraint_name}',
+        )
+        # Validate constraint.
+        with connection.schema_editor(atomic=True) as editor:
+            msg = f'check constraint "{constraint_name}"'
+            with self.assertRaisesMessage(IntegrityError, msg):
+                operation.database_forwards(self.app_label, editor, project_state, new_state)
+        obj.pink = 5
+        obj.save()
+        with connection.schema_editor(atomic=True) as editor:
+            operation.database_forwards(self.app_label, editor, project_state, new_state)
+        # Reversal is a noop.
+        with connection.schema_editor() as editor:
+            with self.assertNumQueries(0):
+                operation.database_backwards(self.app_label, editor, new_state, project_state)
+        # Deconstruction.
+        name, args, kwargs = operation.deconstruct()
+        self.assertEqual(name, 'ValidateConstraint')
+        self.assertEqual(args, [])
+        self.assertEqual(kwargs, {'model_name': 'Pony', 'name': constraint_name})