Fixed #36377 -- Added hints support to CreateExtension and subclasses.

1 files changed, 54 insertions, 0 deletions

diff --git a/tests/postgres_tests/test_operations.py b/tests/postgres_tests/test_operations.py
index 322f38148b..551898c80a 100644
--- a/tests/postgres_tests/test_operations.py
+++ b/tests/postgres_tests/test_operations.py
@@ -238,6 +238,11 @@ class NoMigrationRouter:
         return False
 
 
+class MigrateWhenHinted:
+    def allow_migrate(self, db, app_label, **hints):
+        return hints.get("a_hint", False)
+
+
 @unittest.skipUnless(connection.vendor == "postgresql", "PostgreSQL specific tests.")
 class CreateExtensionTests(PostgreSQLTestCase):
     app_label = "test_allow_create_extention"
@@ -289,6 +294,55 @@ class CreateExtensionTests(PostgreSQLTestCase):
         self.assertEqual(len(captured_queries), 2)
         self.assertIn("DROP EXTENSION IF EXISTS", captured_queries[1]["sql"])
 
+    @override_settings(DATABASE_ROUTERS=[MigrateWhenHinted()])
+    def test_allow_migrate_based_on_hints(self):
+        operation_no_hints = CreateExtension("tablefunc")
+        self.assertEqual(operation_no_hints.hints, {})
+
+        operation_hints = CreateExtension("tablefunc", hints={"a_hint": True})
+        self.assertEqual(operation_hints.hints, {"a_hint": True})
+
+        project_state = ProjectState()
+        new_state = project_state.clone()
+
+        with (
+            CaptureQueriesContext(connection) as captured_queries,
+            connection.schema_editor(atomic=False) as editor,
+        ):
+            operation_no_hints.database_forwards(
+                self.app_label, editor, project_state, new_state
+            )
+        self.assertEqual(len(captured_queries), 0)
+
+        with (
+            CaptureQueriesContext(connection) as captured_queries,
+            connection.schema_editor(atomic=False) as editor,
+        ):
+            operation_no_hints.database_backwards(
+                self.app_label, editor, project_state, new_state
+            )
+        self.assertEqual(len(captured_queries), 0)
+
+        with (
+            CaptureQueriesContext(connection) as captured_queries,
+            connection.schema_editor(atomic=False) as editor,
+        ):
+            operation_hints.database_forwards(
+                self.app_label, editor, project_state, new_state
+            )
+        self.assertEqual(len(captured_queries), 4)
+        self.assertIn("CREATE EXTENSION IF NOT EXISTS", captured_queries[1]["sql"])
+
+        with (
+            CaptureQueriesContext(connection) as captured_queries,
+            connection.schema_editor(atomic=False) as editor,
+        ):
+            operation_hints.database_backwards(
+                self.app_label, editor, project_state, new_state
+            )
+        self.assertEqual(len(captured_queries), 2)
+        self.assertIn("DROP EXTENSION IF EXISTS", captured_queries[1]["sql"])
+
     def test_create_existing_extension(self):
         operation = BloomExtension()
         self.assertEqual(operation.migration_name_fragment, "create_extension_bloom")