From 5488530a272b863794484ee2b027294ff2ec86d2 Mon Sep 17 00:00:00 2001 From: Anthony Sottile Date: Thu, 8 May 2025 16:37:11 -0400 Subject: [PATCH] Fixed #36377 -- Added hints support to CreateExtension and subclasses. --- AUTHORS | 1 + django/contrib/postgres/operations.py | 41 ++++++----- docs/ref/contrib/postgres/operations.txt | 90 +++++++++++++++++++++--- docs/releases/6.0.txt | 8 +++ tests/postgres_tests/test_operations.py | 54 ++++++++++++++ 5 files changed, 166 insertions(+), 28 deletions(-) diff --git a/AUTHORS b/AUTHORS index f492b36357..7beafefd6f 100644 --- a/AUTHORS +++ b/AUTHORS @@ -100,6 +100,7 @@ answer newbie questions, and generally made Django that much better: Anssi Kääriäinen ant9000@netwise.it Anthony Briggs + anthony sottile Anthony Wright Antoine Chéneau Anton Samarchyan diff --git a/django/contrib/postgres/operations.py b/django/contrib/postgres/operations.py index 84360febf9..c09d7874c1 100644 --- a/django/contrib/postgres/operations.py +++ b/django/contrib/postgres/operations.py @@ -13,15 +13,16 @@ class CreateExtension(Operation): reversible = True category = OperationCategory.ADDITION - def __init__(self, name): + def __init__(self, name, hints=None): self.name = name + self.hints = hints or {} def state_forwards(self, app_label, state): pass def database_forwards(self, app_label, schema_editor, from_state, to_state): if schema_editor.connection.vendor != "postgresql" or not router.allow_migrate( - schema_editor.connection.alias, app_label + schema_editor.connection.alias, app_label, **self.hints ): return if not self.extension_exists(schema_editor, self.name): @@ -42,7 +43,9 @@ class CreateExtension(Operation): ) def database_backwards(self, app_label, schema_editor, from_state, to_state): - if not router.allow_migrate(schema_editor.connection.alias, app_label): + if not router.allow_migrate( + schema_editor.connection.alias, app_label, **self.hints + ): return if self.extension_exists(schema_editor, self.name): schema_editor.execute( @@ -69,43 +72,43 @@ class CreateExtension(Operation): class BloomExtension(CreateExtension): - def __init__(self): - self.name = "bloom" + def __init__(self, hints=None): + super().__init__("bloom", hints=hints) class BtreeGinExtension(CreateExtension): - def __init__(self): - self.name = "btree_gin" + def __init__(self, hints=None): + super().__init__("btree_gin", hints=hints) class BtreeGistExtension(CreateExtension): - def __init__(self): - self.name = "btree_gist" + def __init__(self, hints=None): + super().__init__("btree_gist", hints=hints) class CITextExtension(CreateExtension): - def __init__(self): - self.name = "citext" + def __init__(self, hints=None): + super().__init__("citext", hints=hints) class CryptoExtension(CreateExtension): - def __init__(self): - self.name = "pgcrypto" + def __init__(self, hints=None): + super().__init__("pgcrypto", hints=hints) class HStoreExtension(CreateExtension): - def __init__(self): - self.name = "hstore" + def __init__(self, hints=None): + super().__init__("hstore", hints=hints) class TrigramExtension(CreateExtension): - def __init__(self): - self.name = "pg_trgm" + def __init__(self, hints=None): + super().__init__("pg_trgm", hints=hints) class UnaccentExtension(CreateExtension): - def __init__(self): - self.name = "unaccent" + def __init__(self, hints=None): + super().__init__("unaccent", hints=hints) class NotInTransactionMixin: diff --git a/docs/ref/contrib/postgres/operations.txt b/docs/ref/contrib/postgres/operations.txt index 1c4cd562d1..e2b9af4f60 100644 --- a/docs/ref/contrib/postgres/operations.txt +++ b/docs/ref/contrib/postgres/operations.txt @@ -41,7 +41,7 @@ them. In that case, connect to your Django database and run the query ``CreateExtension`` =================== -.. class:: CreateExtension(name) +.. class:: CreateExtension(name, hints=None) An ``Operation`` subclass which installs a PostgreSQL extension. For common extensions, use one of the more specific subclasses below. @@ -50,63 +50,135 @@ them. In that case, connect to your Django database and run the query This is a required argument. The name of the extension to be installed. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + ``BloomExtension`` ================== -.. class:: BloomExtension() +.. class:: BloomExtension(hints=None) Installs the ``bloom`` extension. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + ``BtreeGinExtension`` ===================== -.. class:: BtreeGinExtension() +.. class:: BtreeGinExtension(hints=None) Installs the ``btree_gin`` extension. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + ``BtreeGistExtension`` ====================== -.. class:: BtreeGistExtension() +.. class:: BtreeGistExtension(hints=None) Installs the ``btree_gist`` extension. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + ``CITextExtension`` =================== -.. class:: CITextExtension() +.. class:: CITextExtension(hints=None) Installs the ``citext`` extension. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + ``CryptoExtension`` =================== -.. class:: CryptoExtension() +.. class:: CryptoExtension(hints=None) Installs the ``pgcrypto`` extension. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + ``HStoreExtension`` =================== -.. class:: HStoreExtension() +.. class:: HStoreExtension(hints=None) Installs the ``hstore`` extension and also sets up the connection to interpret hstore data for possible use in subsequent migrations. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + ``TrigramExtension`` ==================== -.. class:: TrigramExtension() +.. class:: TrigramExtension(hints=None) Installs the ``pg_trgm`` extension. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + ``UnaccentExtension`` ===================== -.. class:: UnaccentExtension() +.. class:: UnaccentExtension(hints=None) Installs the ``unaccent`` extension. + .. attribute:: hints + + .. versionadded:: 6.0 + + The optional ``hints`` argument will be passed as ``**hints`` to the + :meth:`allow_migrate` method of database routers to assist them in + :ref:`making routing decisions `. + .. _manage-postgresql-collations: Managing collations using migrations diff --git a/docs/releases/6.0.txt b/docs/releases/6.0.txt index c5fa1bac63..9d8120f5c2 100644 --- a/docs/releases/6.0.txt +++ b/docs/releases/6.0.txt @@ -135,6 +135,14 @@ Minor features now include system checks to verify that ``django.contrib.postgres`` is an installed app. +* The :class:`.CreateExtension`, :class:`.BloomExtension`, + :class:`.BtreeGinExtension`, :class:`.BtreeGistExtension`, + :class:`.CITextExtension`, :class:`.CryptoExtension`, + :class:`.HStoreExtension`, :class:`.TrigramExtension`, and + :class:`.UnaccentExtension` operations now support the optional ``hints`` + parameter. This allows providing database hints to database routers to assist + them in :ref:`making routing decisions `. + :mod:`django.contrib.redirects` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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")