From 7d710d8543d9afeff6f96013d6ca50afe9084195 Mon Sep 17 00:00:00 2001
From: Swen Kooij <swen@sectorlabs.ro>
Date: Mon, 19 Oct 2020 18:14:24 +0300
Subject: [PATCH] Make AddListPartition and AddRangePartition reversible

Fixes #101
---
 .../operations/add_list_partition.py          |  7 +++
 .../operations/add_range_partition.py         |  7 +++
 tests/test_migration_operations.py            | 47 +++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/psqlextra/backend/migrations/operations/add_list_partition.py b/psqlextra/backend/migrations/operations/add_list_partition.py
index 787727ee..0c718bbe 100644
--- a/psqlextra/backend/migrations/operations/add_list_partition.py
+++ b/psqlextra/backend/migrations/operations/add_list_partition.py
@@ -43,6 +43,13 @@ def database_forwards(self, app_label, schema_editor, from_state, to_state):
         if self.allow_migrate_model(schema_editor.connection.alias, model):
             schema_editor.add_list_partition(model, self.name, self.values)
 
+    def database_backwards(
+        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.delete_partition(model, self.name)
+
     def deconstruct(self):
         name, args, kwargs = super().deconstruct()
         kwargs["values"] = self.values
diff --git a/psqlextra/backend/migrations/operations/add_range_partition.py b/psqlextra/backend/migrations/operations/add_range_partition.py
index 74b63e38..41a94ea5 100644
--- a/psqlextra/backend/migrations/operations/add_range_partition.py
+++ b/psqlextra/backend/migrations/operations/add_range_partition.py
@@ -53,6 +53,13 @@ def database_forwards(self, app_label, schema_editor, from_state, to_state):
                 model, self.name, self.from_values, self.to_values
             )
 
+    def database_backwards(
+        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.delete_partition(model, self.name)
+
     def deconstruct(self):
         name, args, kwargs = super().deconstruct()
 
diff --git a/tests/test_migration_operations.py b/tests/test_migration_operations.py
index 48690128..f34aef9e 100644
--- a/tests/test_migration_operations.py
+++ b/tests/test_migration_operations.py
@@ -126,6 +126,53 @@ def test_migration_operations_delete_partitioned_table(method, create_model):
     assert _partitioned_table_exists(create_operation)
 
 
+@pytest.mark.parametrize(
+    "method,add_partition_operation",
+    [
+        (
+            PostgresPartitioningMethod.LIST,
+            operations.PostgresAddDefaultPartition(
+                model_name="test", name="pt1"
+            ),
+        ),
+        (
+            PostgresPartitioningMethod.RANGE,
+            operations.PostgresAddRangePartition(
+                model_name="test",
+                name="pt1",
+                from_values="2019-01-01",
+                to_values="2019-02-01",
+            ),
+        ),
+        (
+            PostgresPartitioningMethod.LIST,
+            operations.PostgresAddListPartition(
+                model_name="test", name="pt1", values=["car", "boat"]
+            ),
+        ),
+    ],
+)
+def test_migration_operations_add_partition(
+    method, add_partition_operation, create_model
+):
+    """Tests whether adding partitions and then rolling them back works as
+    expected."""
+
+    create_operation = create_model(method)
+    state = migrations.state.ProjectState.from_apps(apps)
+
+    # migrate forwards
+    apply_migration([create_operation, add_partition_operation], state)
+    assert _partition_exists(create_operation, add_partition_operation)
+
+    # rollback
+    apply_migration(
+        [create_operation, add_partition_operation], state, backwards=True
+    )
+
+    assert not _partition_exists(create_operation, add_partition_operation)
+
+
 @pytest.mark.parametrize(
     "method,add_partition_operation,delete_partition_operation",
     [