Merge pull request #72 from kraken-tech/add-deferrable-to-safer-uniqu…

…e-constraint

Add deferrable to safer unique constraint

CHANGELOG.md

@@ -8,6 +8,11 @@ Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to
 
 ## [Unreleased]
 
+### Added
+
+- Enhanced `SaferAddUniqueConstraint` to support a `UniqueConstraint` with the
+  `deferrable` argument.
+
 ## [0.1.16] - 2025-01-08
 
 ### Added

docs/conf.py

@@ -13,7 +13,10 @@
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
-extensions = ["sphinx_rtd_theme"]
+extensions = [
+    "sphinx_rtd_theme",
+    "sphinx_design",
+]
 html_theme = "sphinx_rtd_theme"
 
 html_context = {

docs/usage/operations.rst

@@ -285,6 +285,19 @@ Class Definitions
       -- Perform the ALTER TABLE using the unique index just created.
       ALTER TABLE "myapp_mymodel" ADD CONSTRAINT "foo_unique" UNIQUE USING INDEX "foo_unique_idx";
 
+    .. dropdown:: Information about ``deferrable``
+        :color: info
+        :icon: info
+
+        The ``deferrable`` argument of ``UniqueConstraint`` is respected.
+
+        That is, if set to ``models.Deferrable.DEFERRED``, the ``ALTER TABLE``
+        command above will include the suffix ``DEFERRABLE INITIALLY
+        DEFERRED``.
+
+        The other value for ``models.Deferrable`` is ``IMMEDIATE``. No changes
+        are performed on the ``ALTER TABLE`` statement in this case as
+        ``IMMEDIATE`` is the default Postgres behaviour.
 
     How to use
     ----------

pyproject.toml

@@ -71,6 +71,7 @@ docs = [
   "sphinx>=7.4.7",
   "sphinx_rtd_theme>=2.0.0",
   "sphinx_lint>=0.9.1",
+  "sphinx_design>=0.6.1",
 ]
 dev = [
     # Testing

requirements/development.txt

@@ -40,7 +40,7 @@ docutils==0.21.2
     # via
     #   sphinx
     #   sphinx-rtd-theme
-environs==14.0.0
+environs==14.1.0
     # via django-pg-migration-tools (pyproject.toml)
 filelock==3.16.1
     # via virtualenv
@@ -56,7 +56,7 @@ jinja2==3.1.5
     # via sphinx
 markupsafe==3.0.2
     # via jinja2
-marshmallow==3.24.1
+marshmallow==3.25.1
     # via environs
 mypy==1.14.1
     # via django-pg-migration-tools (pyproject.toml)
@@ -103,15 +103,18 @@ regex==2024.11.6
     # via sphinx-lint
 requests==2.32.3
     # via sphinx
-ruff==0.8.6
+ruff==0.9.1
     # via django-pg-migration-tools (pyproject.toml)
 snowballstemmer==2.2.0
     # via sphinx
 sphinx==8.1.3
     # via
     #   django-pg-migration-tools (pyproject.toml)
+    #   sphinx-design
     #   sphinx-rtd-theme
     #   sphinxcontrib-jquery
+sphinx-design==0.6.1
+    # via django-pg-migration-tools (pyproject.toml)
 sphinx-lint==1.0.0
     # via django-pg-migration-tools (pyproject.toml)
 sphinx-rtd-theme==3.0.2

requirements/docs.txt

@@ -35,8 +35,11 @@ snowballstemmer==2.2.0
 sphinx==8.1.3
     # via
     #   django-pg-migration-tools (pyproject.toml)
+    #   sphinx-design
     #   sphinx-rtd-theme
     #   sphinxcontrib-jquery
+sphinx-design==0.6.1
+    # via django-pg-migration-tools (pyproject.toml)
 sphinx-lint==1.0.0
     # via django-pg-migration-tools (pyproject.toml)
 sphinx-rtd-theme==3.0.2

requirements/pytest-in-nox-psycopg2.txt

@@ -23,13 +23,13 @@ django-stubs==5.1.1
     # via django-pg-migration-tools (pyproject.toml)
 django-stubs-ext==5.1.1
     # via django-stubs
-environs==14.0.0
+environs==14.1.0
     # via django-pg-migration-tools (pyproject.toml)
 filelock==3.16.1
     # via virtualenv
 iniconfig==2.0.0
     # via pytest
-marshmallow==3.24.1
+marshmallow==3.25.1
     # via environs
 nox==2024.10.9
     # via django-pg-migration-tools (pyproject.toml)

requirements/pytest-in-nox-psycopg3.txt

@@ -23,13 +23,13 @@ django-stubs==5.1.1
     # via django-pg-migration-tools (pyproject.toml)
 django-stubs-ext==5.1.1
     # via django-stubs
-environs==14.0.0
+environs==14.1.0
     # via django-pg-migration-tools (pyproject.toml)
 filelock==3.16.1
     # via virtualenv
 iniconfig==2.0.0
     # via pytest
-marshmallow==3.24.1
+marshmallow==3.25.1
     # via environs
 nox==2024.10.9
     # via django-pg-migration-tools (pyproject.toml)

src/django_pg_migration_tools/operations.py

@@ -455,6 +455,19 @@ def create_unique_constraint(
         alter_table_sql = base_sql.split(" UNIQUE")[0]
         sql = f'{alter_table_sql} UNIQUE USING INDEX "{index.name}"'
 
+        if constraint.deferrable == models.Deferrable.DEFERRED:
+            sql += " DEFERRABLE INITIALLY DEFERRED"
+        else:
+            # Note that there are only two options for models.Deferrable:
+            # IMMEDIATE or DEFERRED. models.Deferrable.IMMEDIATE is the default
+            # and most common case. DEFERRED is also Postgres default when no
+            # deferrable setting was specified. This means that having an
+            # "else" here that would append "DEFERRABLE INITIALLY IMMEDIATE"
+            # to the query for the "IMMEDIATE" case is not necessary. This
+            # mimics what Django already does, which is to not include that
+            # suffix.
+            pass
+
         # Now we can execute the schema change. We have lock timeouts back in
         # place after creating the index that would prevent this operation from
         # running for too long if it's blocked by another query. Otherwise,

tests/django_pg_migration_tools/test_operations.py

@@ -945,6 +945,128 @@ def test_when_not_allowed_to_migrate_by_the_router(self):
         # the router.
         assert len(queries) == 0
 
+    @pytest.mark.django_db(transaction=True)
+    def test_when_deferred_set(self):
+        # Prove that:
+        #   - An invalid index doesn't exist.
+        #   - The constraint doesn't exist yet.
+        with connection.cursor() as cursor:
+            cursor.execute(
+                psycopg_sql.SQL(operations.IndexQueries.CHECK_INVALID_INDEX)
+                .format(index_name=psycopg_sql.Literal("unique_int_field"))
+                .as_string(cursor.connection)
+            )
+            assert not cursor.fetchone()
+            cursor.execute(
+                psycopg_sql.SQL(operations.ConstraintQueries.CHECK_EXISTING_CONSTRAINT)
+                .format(constraint_name=psycopg_sql.Literal("unique_int_field"))
+                .as_string(cursor.connection)
+            )
+            assert not cursor.fetchone()
+            # Also, set the lock_timeout to check it has been returned to
+            # its original value once the unique index creation is completed.
+            cursor.execute(_SET_LOCK_TIMEOUT)
+
+        project_state = ProjectState()
+        project_state.add_model(ModelState.from_model(IntModel))
+        new_state = project_state.clone()
+
+        operation = operations.SaferAddUniqueConstraint(
+            model_name="intmodel",
+            constraint=UniqueConstraint(
+                fields=("int_field",),
+                name="unique_int_field",
+                deferrable=models.Deferrable.DEFERRED,
+            ),
+        )
+
+        operation.state_forwards(self.app_label, new_state)
+        # Proceed to add the unique index followed by the constraint:
+        with connection.schema_editor(atomic=False, collect_sql=False) as editor:
+            with utils.CaptureQueriesContext(connection) as queries:
+                operation.database_forwards(
+                    self.app_label, editor, project_state, new_state
+                )
+
+        with connection.cursor() as cursor:
+            cursor.execute(
+                _CHECK_CONSTRAINT_EXISTS_QUERY,
+                {
+                    "table_name": "example_app_intmodel",
+                    "constraint_name": "unique_int_field",
+                },
+            )
+            assert cursor.fetchone()
+
+        # Assert on the sequence of expected SQL queries:
+        #
+        # 1. Check whether the constraint already exists.
+        assert queries[0]["sql"] == dedent("""
+            SELECT conname
+            FROM pg_catalog.pg_constraint
+            WHERE conname = 'unique_int_field';
+            """)
+        # 2. Check the original lock_timeout value to be able to restore it
+        # later.
+        assert queries[1]["sql"] == "SHOW lock_timeout;"
+        # 3. Remove the timeout.
+        assert queries[2]["sql"] == "SET lock_timeout = '0';"
+        # 4. Verify if the index is invalid.
+        assert queries[3]["sql"] == dedent("""
+            SELECT relname
+            FROM pg_class, pg_index
+            WHERE (
+                pg_index.indisvalid = false
+                AND pg_index.indexrelid = pg_class.oid
+                AND relname = 'unique_int_field'
+            );
+            """)
+        # 5. Finally create the index concurrently.
+        assert (
+            queries[4]["sql"]
+            == 'CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "unique_int_field" ON "example_app_intmodel" ("int_field")'
+        )
+        # 6. Set the timeout back to what it was originally.
+        assert queries[5]["sql"] == "SET lock_timeout = '1s';"
+
+        # 7. Add the table constraint with the DEFERRED option set.
+        assert (
+            queries[6]["sql"]
+            == 'ALTER TABLE "example_app_intmodel" ADD CONSTRAINT "unique_int_field" UNIQUE USING INDEX "unique_int_field" DEFERRABLE INITIALLY DEFERRED'
+        )
+
+        # Reverse the migration to drop the index and constraint, and verify
+        # that the lock_timeout queries are correct.
+        with connection.schema_editor(atomic=False, collect_sql=False) as editor:
+            with utils.CaptureQueriesContext(connection) as reverse_queries:
+                operation.database_backwards(
+                    self.app_label, editor, new_state, project_state
+                )
+
+        # 1. Check that the constraint is still there.
+        assert queries[0]["sql"] == dedent("""
+            SELECT conname
+            FROM pg_catalog.pg_constraint
+            WHERE conname = 'unique_int_field';
+            """)
+
+        # 2. perform the ALTER TABLE.
+        assert (
+            reverse_queries[1]["sql"]
+            == 'ALTER TABLE "example_app_intmodel" DROP CONSTRAINT "unique_int_field"'
+        )
+
+        # Verify the constraint doesn't exist any more.
+        with connection.cursor() as cursor:
+            cursor.execute(
+                _CHECK_CONSTRAINT_EXISTS_QUERY,
+                {
+                    "table_name": "example_app_intmodel",
+                    "constraint_name": "unique_int_field",
+                },
+            )
+            assert not cursor.fetchone()
+
     @pytest.mark.django_db(transaction=True)
     def test_raises_if_constraint_already_exists(self):
         project_state = ProjectState()