Make PostgresModelDataMigrator work for scheduled rebuilds of a table

.github/workflows/release.yml

@@ -12,7 +12,7 @@ on:
         required: true
 
 concurrency:
-  group: release-${{ github.workflow }}-${{ github.ref }}-${{ inputs.version || '' }
+  group: release-${{ github.workflow }}-${{ github.ref }}-${{ inputs.version || '' }}
   cancel-in-progress: true
 
 jobs:
@@ -22,7 +22,7 @@ jobs:
     runs-on: ubuntu-latest
 
     container:
-      image: public.ecr.aws/docker/library/python:3.11-bullseye
+      image: public.ecr.aws/docker/library/python:3.11-bookworm
 
     steps:
       - name: Checkout
@@ -36,11 +36,8 @@ jobs:
 
       - name: Set version number
         env:
-          RELEASE_VERSION: |-
-            ${{ case(
-              github.event_name == 'workflow_dispatch', inputs.version,
-              github.ref_name,
-            ) }}
+          RELEASE_VERSION: ${{ github.event_name == 'workflow_dispatch' && inputs.version || github.ref_name }}
+        shell: bash
         run: |
           if [[ "$RELEASE_VERSION" =~ ^v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
             echo "__version__ = \"${RELEASE_VERSION:1}\"" > psqlextra/_version.py

.github/workflows/test.yml

@@ -17,7 +17,7 @@ jobs:
     runs-on: ubuntu-latest
 
     container:
-      image: public.ecr.aws/docker/library/python:3.11-bullseye
+      image: public.ecr.aws/docker/library/python:3.11-bookworm
 
     steps:
       - name: Checkout
@@ -43,25 +43,25 @@ jobs:
         include:
           - python-version: "3.7"
             postgres-version: "13"
-            debian-dist: "bullseye"
+            debian-dist: "bookworm"
           - python-version: "3.8"
             postgres-version: "13"
-            debian-dist: "bullseye"
+            debian-dist: "bookworm"
           - python-version: "3.9"
             postgres-version: "13"
-            debian-dist: "bullseye"
+            debian-dist: "bookworm"
           - python-version: "3.10"
             postgres-version: "16"
-            debian-dist: "bullseye"
+            debian-dist: "bookworm"
           - python-version: "3.11"
             postgres-version: "16"
-            debian-dist: "bullseye"
+            debian-dist: "bookworm"
           - python-version: "3.12"
             postgres-version: "16"
-            debian-dist: "bullseye"
+            debian-dist: "bookworm"
           - python-version: "3.13"
             postgres-version: "16"
-            debian-dist: "bullseye"
+            debian-dist: "bookworm"
 
     container:
       image: public.ecr.aws/docker/library/python:${{ matrix.python-version }}-${{ matrix.debian-dist }}

psqlextra/backend/schema.py

@@ -286,8 +286,9 @@ def clone_model_constraints_and_indexes_to_schema(
                 # a AccessExclusiveLock on the referenced table. We want to avoid
                 # that and created the FK as NOT VALID. We can run VALIDATE in
                 # a separate transaction later to validate the entries without
-                # acquiring a AccessExclusiveLock.
-                if field.remote_field:
+                # acquiring a AccessExclusiveLock. A field with `db_constraint=False`
+                # has no foreign key constraint to clone.
+                if field.remote_field and field.db_constraint:
                     with postgres_reset_local_search_path(
                         using=self.connection.alias
                     ):

psqlextra/contrib/model_data_migrator.py

@@ -99,16 +99,26 @@ class PostgresModelDataMigrator:
     The process is very similiar to how pg_repack rewrites
     an entire table without long-running locks on the table.
 
+    If any of the steps fails, the work and backup schemas are
+    deleted and the original table is left as it was.
+
     Attributes:
         model: The model to migrate.
         using: Optional name of the database connection to use.
         operation_timeout: Maximum amount of time a single statement
                            can take.
+        keep_backup_schema: Whether to keep the backup schema with
+                            the original table in it after a
+                            successful migration. Turn off for
+                            migrations that run on a schedule, which
+                            would otherwise leave a copy of the table
+                            behind on every run.
     """
 
     model: Type[models.Model]
     using: str = DEFAULT_DB_ALIAS
     operation_timeout: timedelta
+    keep_backup_schema: bool = True
 
     def __init__(self, logger) -> None:
         self.logger = logger
@@ -184,19 +194,32 @@ def migrate(self) -> PostgresModelDataMigratorState:
             (self._migrate_phase_4, "swapping"),
         ]
 
-        for index, (phase, description) in enumerate(phases):
-            phase_start_time = time.time()
-            logger.info(
-                f"Starting phase #{index + 1} of migrating {self.table_name}: {description}"
-            )
-            phase(state)
+        try:
+            for index, (phase, description) in enumerate(phases):
+                phase_start_time = time.time()
+                logger.info(
+                    f"Starting phase #{index + 1} of migrating {self.table_name}: {description}"
+                )
+                phase(state)
+                logger.info(
+                    f"Finished phase #{index + 1} of migrating {self.table_name}: {description}",
+                    task_time=time.time() - phase_start_time,
+                )
+        except Exception:
+            # The swap is the last phase and runs in a single transaction,
+            # so the original table never made it into the backup schema.
             logger.info(
-                f"Finished phase #{index + 1} of migrating {self.table_name}: {description}",
-                task_time=time.time() - phase_start_time,
+                f"Failed migrating {self.table_name}, deleting the work and backup schemas"
             )
+            state.work_schema.delete(cascade=True, using=self.using)
+            state.backup_schema.delete(cascade=True, using=self.using)
+            raise
 
         state.work_schema.delete(cascade=True, using=self.using)
 
+        if not self.keep_backup_schema:
+            state.backup_schema.delete(cascade=True, using=self.using)
+
         logger.info(
             f"Finished migrating {self.table_name}",
             task_time=time.time() - start_time,

tests/test_model_data_migrator.py

@@ -0,0 +1,110 @@
+from datetime import timedelta
+from unittest.mock import MagicMock
+
+import django
+import pytest
+
+from django.db import connection, models, transaction
+
+from psqlextra.contrib.model_data_migrator import PostgresModelDataMigrator
+from psqlextra.schema import PostgresSchema
+from psqlextra.settings import postgres_prepend_local_search_path
+
+from .fake_model import delete_fake_model, get_fake_model
+
+pytestmark = pytest.mark.skipif(
+    django.VERSION < (3, 2),
+    reason="The migrator clones models into a separate schema and uses durable transactions, which need Django >= 3.2",
+)
+
+
+@pytest.fixture
+def fake_model():
+    model = get_fake_model({"name": models.TextField()})
+
+    yield model
+
+    delete_fake_model(model)
+
+
+def _create_migrator(model, *, keep_backup_schema=True, fail=False):
+    class Migrator(PostgresModelDataMigrator):
+        operation_timeout = timedelta(minutes=1)
+
+        def fill_cloned_table_lockless(self, work_schema, default_schema):
+            with self.atomic():
+                with postgres_prepend_local_search_path([work_schema.name]):
+                    self.model.objects.create(name="new")
+
+            if fail:
+                raise RuntimeError("fill failed")
+
+        def clean_cloned_table(self, work_schema, default_schema):
+            pass
+
+        def fill_cloned_table_locked(self, work_schema, default_schema):
+            pass
+
+    Migrator.model = model
+    Migrator.keep_backup_schema = keep_backup_schema
+
+    return Migrator(MagicMock())
+
+
+def _list_schemas_of(model):
+    with connection.cursor() as cursor:
+        schema_names = connection.introspection.get_schema_list(cursor)
+
+    return [
+        schema_name
+        for schema_name in schema_names
+        if model._meta.db_table in schema_name
+    ]
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_swaps_in_the_filled_table(fake_model):
+    fake_model.objects.create(name="old")
+
+    _create_migrator(fake_model).migrate()
+
+    assert list(fake_model.objects.values_list("name", flat=True)) == ["new"]
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_keeps_the_backup_schema(fake_model):
+    fake_model.objects.create(name="old")
+
+    state = _create_migrator(fake_model).migrate()
+
+    assert _list_schemas_of(fake_model) == [state.backup_schema.name]
+
+    with transaction.atomic():
+        with postgres_prepend_local_search_path([state.backup_schema.name]):
+            with connection.cursor() as cursor:
+                cursor.execute(
+                    f"SELECT name FROM {connection.ops.quote_name(fake_model._meta.db_table)}"
+                )
+                assert cursor.fetchall() == [("old",)]
+
+    PostgresSchema(state.backup_schema.name).delete(cascade=True)
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_deletes_the_backup_schema_when_asked(
+    fake_model,
+):
+    _create_migrator(fake_model, keep_backup_schema=False).migrate()
+
+    assert _list_schemas_of(fake_model) == []
+
+
+@pytest.mark.django_db(transaction=True)
+def test_model_data_migrator_deletes_its_schemas_when_it_fails(fake_model):
+    fake_model.objects.create(name="old")
+
+    with pytest.raises(RuntimeError):
+        _create_migrator(fake_model, fail=True).migrate()
+
+    assert _list_schemas_of(fake_model) == []
+    assert list(fake_model.objects.values_list("name", flat=True)) == ["old"]

tests/test_schema_editor_clone_model_to_schema.py

@@ -197,6 +197,13 @@ def fake_model(fake_model_fk_target_1, fake_model_fk_target_2):
             "alternative_family": models.ForeignKey(
                 fake_model_fk_target_2, null=True, on_delete=models.SET_NULL
             ),
+            "unconstrained_family": models.ForeignKey(
+                fake_model_fk_target_2,
+                null=True,
+                on_delete=models.DO_NOTHING,
+                db_constraint=False,
+                related_name="+",
+            ),
         },
         meta_options=meta_options,
     )