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

This commit is contained in:
Anthony Sottile 2025-05-08 16:37:11 -04:00 committed by Mariusz Felisiak
parent 14fc2e9703
commit 5488530a27
5 changed files with 166 additions and 28 deletions

View file

@ -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")