From 8f4877c89d0f28e289399fbb46357623a49e7eb0 Mon Sep 17 00:00:00 2001 From: Loic Bistuer Date: Fri, 9 Jan 2015 00:10:10 +0700 Subject: [PATCH] Fixed #22583 -- Allowed RunPython and RunSQL to provide hints to the db router. Thanks Markus Holtermann and Tim Graham for the review. --- django/db/migrations/operations/base.py | 16 +- django/db/migrations/operations/special.py | 30 ++-- django/db/utils.py | 4 +- docs/ref/migration-operations.txt | 20 ++- docs/releases/1.8.txt | 14 ++ docs/topics/db/multi-db.txt | 6 +- docs/topics/migrations.txt | 20 ++- tests/migrations/test_multidb.py | 174 +++++++++++++++++++++ tests/migrations/test_operations.py | 38 ----- 9 files changed, 247 insertions(+), 75 deletions(-) create mode 100644 tests/migrations/test_multidb.py diff --git a/django/db/migrations/operations/base.py b/django/db/migrations/operations/base.py index 655f03a4c6..04089938c2 100644 --- a/django/db/migrations/operations/base.py +++ b/django/db/migrations/operations/base.py @@ -98,17 +98,15 @@ class Operation(object): """ return self.references_model(model_name, app_label) - def allowed_to_migrate(self, connection_alias, model): + def allowed_to_migrate(self, connection_alias, model, hints=None): """ - Returns if we're allowed to migrate the model. Checks the router, - if it's a proxy, if it's managed, and if it's swapped out. + Returns if we're allowed to migrate the model. """ - return ( - not model._meta.proxy and - not model._meta.swapped and - model._meta.managed and - router.allow_migrate(connection_alias, model) - ) + # Always skip if proxy, swapped out, or unmanaged. + if model and (model._meta.proxy or model._meta.swapped or not model._meta.managed): + return False + + return router.allow_migrate(connection_alias, model, **(hints or {})) def __repr__(self): return "<%s %s%s>" % ( diff --git a/django/db/migrations/operations/special.py b/django/db/migrations/operations/special.py index 00caa08684..91261f724f 100644 --- a/django/db/migrations/operations/special.py +++ b/django/db/migrations/operations/special.py @@ -63,10 +63,11 @@ class RunSQL(Operation): """ noop = '' - def __init__(self, sql, reverse_sql=None, state_operations=None): + def __init__(self, sql, reverse_sql=None, state_operations=None, hints=None): self.sql = sql self.reverse_sql = reverse_sql self.state_operations = state_operations or [] + self.hints = hints or {} def deconstruct(self): kwargs = { @@ -76,6 +77,8 @@ class RunSQL(Operation): kwargs['reverse_sql'] = self.reverse_sql if self.state_operations: kwargs['state_operations'] = self.state_operations + if self.hints: + kwargs['hints'] = self.hints return ( self.__class__.__name__, [], @@ -91,12 +94,14 @@ class RunSQL(Operation): state_operation.state_forwards(app_label, state) def database_forwards(self, app_label, schema_editor, from_state, to_state): - self._run_sql(schema_editor, self.sql) + if self.allowed_to_migrate(schema_editor.connection.alias, None, hints=self.hints): + self._run_sql(schema_editor, self.sql) def database_backwards(self, app_label, schema_editor, from_state, to_state): if self.reverse_sql is None: raise NotImplementedError("You cannot reverse this operation") - self._run_sql(schema_editor, self.reverse_sql) + if self.allowed_to_migrate(schema_editor.connection.alias, None, hints=self.hints): + self._run_sql(schema_editor, self.reverse_sql) def describe(self): return "Raw SQL operation" @@ -125,7 +130,7 @@ class RunPython(Operation): reduces_to_sql = False - def __init__(self, code, reverse_code=None, atomic=True): + def __init__(self, code, reverse_code=None, atomic=True, hints=None): self.atomic = atomic # Forwards code if not callable(code): @@ -138,6 +143,7 @@ class RunPython(Operation): if not callable(reverse_code): raise ValueError("RunPython must be supplied with callable arguments") self.reverse_code = reverse_code + self.hints = hints or {} def deconstruct(self): kwargs = { @@ -147,6 +153,8 @@ class RunPython(Operation): kwargs['reverse_code'] = self.reverse_code if self.atomic is not True: kwargs['atomic'] = self.atomic + if self.hints: + kwargs['hints'] = self.hints return ( self.__class__.__name__, [], @@ -163,16 +171,18 @@ class RunPython(Operation): pass def database_forwards(self, app_label, schema_editor, from_state, to_state): - # We now execute the Python code in a context that contains a 'models' - # object, representing the versioned models as an app registry. - # We could try to override the global cache, but then people will still - # use direct imports, so we go with a documentation approach instead. - self.code(from_state.apps, schema_editor) + if self.allowed_to_migrate(schema_editor.connection.alias, None, hints=self.hints): + # We now execute the Python code in a context that contains a 'models' + # object, representing the versioned models as an app registry. + # We could try to override the global cache, but then people will still + # use direct imports, so we go with a documentation approach instead. + self.code(from_state.apps, schema_editor) def database_backwards(self, app_label, schema_editor, from_state, to_state): if self.reverse_code is None: raise NotImplementedError("You cannot reverse this operation") - self.reverse_code(from_state.apps, schema_editor) + if self.allowed_to_migrate(schema_editor.connection.alias, None, hints=self.hints): + self.reverse_code(from_state.apps, schema_editor) def describe(self): return "Raw Python operation" diff --git a/django/db/utils.py b/django/db/utils.py index a355bbe466..fc2203ff47 100644 --- a/django/db/utils.py +++ b/django/db/utils.py @@ -316,7 +316,7 @@ class ConnectionRouter(object): return allow return obj1._state.db == obj2._state.db - def allow_migrate(self, db, model): + def allow_migrate(self, db, model, **hints): for router in self.routers: try: try: @@ -331,7 +331,7 @@ class ConnectionRouter(object): # If the router doesn't have a method, skip to the next one. pass else: - allow = method(db, model) + allow = method(db, model, **hints) if allow is not None: return allow return True diff --git a/docs/ref/migration-operations.txt b/docs/ref/migration-operations.txt index dcac7f26a6..8e13974cfe 100644 --- a/docs/ref/migration-operations.txt +++ b/docs/ref/migration-operations.txt @@ -206,7 +206,7 @@ Special Operations RunSQL ------ -.. class:: RunSQL(sql, reverse_sql=None, state_operations=None) +.. class:: RunSQL(sql, reverse_sql=None, state_operations=None, hints=None) Allows running of arbitrary SQL on the database - useful for more advanced features of database backends that Django doesn't support directly, like @@ -235,6 +235,11 @@ operation here so that the autodetector still has an up-to-date state of the model (otherwise, when you next run ``makemigrations``, it won't see any operation that adds that field and so will try to run it again). +The optional ``hints`` argument will be passed as ``**hints`` to the +:meth:`allow_migrate` method of database routers to assist them in making +routing decisions. See :ref:`topics-db-multi-db-hints` for more details on +database hints. + .. versionchanged:: 1.7.1 If you want to include literal percent signs in a query without parameters @@ -245,6 +250,8 @@ operation that adds that field and so will try to run it again). The ability to pass parameters to the ``sql`` and ``reverse_sql`` queries was added. + The ``hints`` argument was added. + .. attribute:: RunSQL.noop .. versionadded:: 1.8 @@ -258,7 +265,7 @@ operation that adds that field and so will try to run it again). RunPython --------- -.. class:: RunPython(code, reverse_code=None, atomic=True) +.. class:: RunPython(code, reverse_code=None, atomic=True, hints=None) Runs custom Python code in a historical context. ``code`` (and ``reverse_code`` if supplied) should be callable objects that accept two arguments; the first is @@ -267,6 +274,15 @@ match the operation's place in the project history, and the second is an instance of :class:`SchemaEditor `. +The optional ``hints`` argument will be passed as ``**hints`` to the +:meth:`allow_migrate` method of database routers to assist them in making a +routing decision. See :ref:`topics-db-multi-db-hints` for more details on +database hints. + +.. versionadded:: 1.8 + + The ``hints`` argument was added. + You are advised to write the code as a separate function above the ``Migration`` class in the migration file, and just pass it to ``RunPython``. Here's an example of using ``RunPython`` to create some initial objects on a ``Country`` diff --git a/docs/releases/1.8.txt b/docs/releases/1.8.txt index 427ec5d493..64765fcea8 100644 --- a/docs/releases/1.8.txt +++ b/docs/releases/1.8.txt @@ -462,6 +462,12 @@ Migrations attribute/method were added to ease in making ``RunPython`` and ``RunSQL`` operations reversible. +* The :class:`~django.db.migrations.operations.RunPython` and + :class:`~django.db.migrations.operations.RunSQL` operations now accept a + ``hints`` parameter that will be passed to :meth:`allow_migrate`. To take + advantage of this feature you must ensure that the ``allow_migrate()`` method + of all your routers accept ``**hints``. + Models ^^^^^^ @@ -1029,6 +1035,14 @@ Miscellaneous * :func:`django.utils.translation.get_language()` now returns ``None`` instead of :setting:`LANGUAGE_CODE` when translations are temporarily deactivated. +* The migration operations :class:`~django.db.migrations.operations.RunPython` + and :class:`~django.db.migrations.operations.RunSQL` now call the + :meth:`allow_migrate` method of database routers. In these cases the + ``model`` argument of ``allow_migrate()`` is set to ``None``, so the router + must properly handle this value. This is most useful when used together with + the newly introduced ``hints`` parameter for these operations, but it can + also be used to disable migrations from running on a particular database. + .. _deprecated-features-1.8: Features deprecated in 1.8 diff --git a/docs/topics/db/multi-db.txt b/docs/topics/db/multi-db.txt index 7d4f2465c9..76df8610e3 100644 --- a/docs/topics/db/multi-db.txt +++ b/docs/topics/db/multi-db.txt @@ -150,7 +150,7 @@ A database Router is a class that provides up to four methods: used by foreign key and many to many operations to determine if a relation should be allowed between two objects. -.. method:: allow_migrate(db, model) +.. method:: allow_migrate(db, model, **hints) Determine if the ``model`` should have tables/indexes created in the database with alias ``db``. Return True if the model should be @@ -293,7 +293,7 @@ send queries for the ``auth`` app to ``auth_db``:: return True return None - def allow_migrate(self, db, model): + def allow_migrate(self, db, model, **hints): """ Make sure the auth app only appears in the 'auth_db' database. @@ -333,7 +333,7 @@ from:: return True return None - def allow_migrate(self, db, model): + def allow_migrate(self, db, model, **hints): """ All non-auth models end up in this pool. """ diff --git a/docs/topics/migrations.txt b/docs/topics/migrations.txt index 48021c724e..1220b4b41f 100644 --- a/docs/topics/migrations.txt +++ b/docs/topics/migrations.txt @@ -545,28 +545,26 @@ attribute:: migrations.RunPython(forwards), ] -You can also use your database router's ``allow_migrate()`` method, but keep in -mind that the imported router needs to stay around as long as it is referenced -inside a migration: +.. versionadded:: 1.8 + +You can also provide hints that will be passed to the :meth:`allow_migrate()` +method of database routers as ``**hints``: .. snippet:: :filename: myapp/dbrouters.py class MyRouter(object): - def allow_migrate(self, db, model): - return db == 'default' + def allow_migrate(self, db, model, **hints): + if 'target_db' in hints: + return db == hints['target_db'] + return True Then, to leverage this in your migrations, do the following:: from django.db import migrations - from myappname.dbrouters import MyRouter - def forwards(apps, schema_editor): - MyModel = apps.get_model("myappname", "MyModel") - if not MyRouter().allow_migrate(schema_editor.connection.alias, MyModel): - return # Your migration code goes here class Migration(migrations.Migration): @@ -576,7 +574,7 @@ Then, to leverage this in your migrations, do the following:: ] operations = [ - migrations.RunPython(forwards), + migrations.RunPython(forwards, hints={'target_db': 'default'}), ] More advanced migrations diff --git a/tests/migrations/test_multidb.py b/tests/migrations/test_multidb.py new file mode 100644 index 0000000000..dec1d25dd0 --- /dev/null +++ b/tests/migrations/test_multidb.py @@ -0,0 +1,174 @@ +import unittest + +try: + import sqlparse +except ImportError: + sqlparse = None + +from django.db import migrations, models, connection +from django.db.migrations.state import ProjectState +from django.test import override_settings + +from .test_operations import OperationTestBase + + +class AgnosticRouter(object): + """ + A router that doesn't have an opinion regarding migrating. + """ + def allow_migrate(self, db, model, **hints): + return None + + +class MigrateNothingRouter(object): + """ + A router that doesn't allow migrating. + """ + def allow_migrate(self, db, model, **hints): + return False + + +class MigrateEverythingRouter(object): + """ + A router that always allows migrating. + """ + def allow_migrate(self, db, model, **hints): + return True + + +class MigrateWhenFooRouter(object): + """ + A router that allows migrating depending on a hint. + """ + def allow_migrate(self, db, model, **hints): + return hints.get('foo', False) + + +class MultiDBOperationTests(OperationTestBase): + multi_db = True + + def _test_create_model(self, app_label, should_run): + """ + Tests that CreateModel honours multi-db settings. + """ + operation = migrations.CreateModel( + "Pony", + [("id", models.AutoField(primary_key=True))], + ) + # Test the state alteration + project_state = ProjectState() + new_state = project_state.clone() + operation.state_forwards(app_label, new_state) + # Test the database alteration + self.assertTableNotExists("%s_pony" % app_label) + with connection.schema_editor() as editor: + operation.database_forwards(app_label, editor, project_state, new_state) + if should_run: + self.assertTableExists("%s_pony" % app_label) + else: + self.assertTableNotExists("%s_pony" % app_label) + # And test reversal + with connection.schema_editor() as editor: + operation.database_backwards(app_label, editor, new_state, project_state) + self.assertTableNotExists("%s_pony" % app_label) + + @override_settings(DATABASE_ROUTERS=[AgnosticRouter()]) + def test_create_model(self): + """ + Test when router doesn't have an opinion (i.e. CreateModel should run). + """ + self._test_create_model("test_mltdb_crmo", should_run=True) + + @override_settings(DATABASE_ROUTERS=[MigrateNothingRouter()]) + def test_create_model2(self): + """ + Test when router returns False (i.e. CreateModel shouldn't run). + """ + self._test_create_model("test_mltdb_crmo2", should_run=False) + + @override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()]) + def test_create_model3(self): + """ + Test when router returns True (i.e. CreateModel should run). + """ + self._test_create_model("test_mltdb_crmo3", should_run=True) + + def test_create_model4(self): + """ + Test multiple routers. + """ + with override_settings(DATABASE_ROUTERS=[AgnosticRouter(), AgnosticRouter()]): + self._test_create_model("test_mltdb_crmo4", should_run=True) + with override_settings(DATABASE_ROUTERS=[MigrateNothingRouter(), MigrateEverythingRouter()]): + self._test_create_model("test_mltdb_crmo4", should_run=False) + with override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter(), MigrateNothingRouter()]): + self._test_create_model("test_mltdb_crmo4", should_run=True) + + def _test_run_sql(self, app_label, should_run, hints=None): + with override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()]): + project_state = self.set_up_test_model(app_label) + + sql = """ + INSERT INTO {0}_pony (pink, weight) VALUES (1, 3.55); + INSERT INTO {0}_pony (pink, weight) VALUES (3, 5.0); + """.format(app_label) + + operation = migrations.RunSQL(sql, hints=hints or {}) + # Test the state alteration does nothing + new_state = project_state.clone() + operation.state_forwards(app_label, new_state) + self.assertEqual(new_state, project_state) + # Test the database alteration + self.assertEqual(project_state.apps.get_model(app_label, "Pony").objects.count(), 0) + with connection.schema_editor() as editor: + operation.database_forwards(app_label, editor, project_state, new_state) + Pony = project_state.apps.get_model(app_label, "Pony") + if should_run: + self.assertEqual(Pony.objects.count(), 2) + else: + self.assertEqual(Pony.objects.count(), 0) + + @unittest.skipIf(sqlparse is None and connection.features.requires_sqlparse_for_splitting, "Missing sqlparse") + @override_settings(DATABASE_ROUTERS=[MigrateNothingRouter()]) + def test_run_sql(self): + self._test_run_sql("test_mltdb_runsql", should_run=False) + + @unittest.skipIf(sqlparse is None and connection.features.requires_sqlparse_for_splitting, "Missing sqlparse") + @override_settings(DATABASE_ROUTERS=[MigrateWhenFooRouter()]) + def test_run_sql2(self): + self._test_run_sql("test_mltdb_runsql2", should_run=False) + self._test_run_sql("test_mltdb_runsql2", should_run=True, hints={'foo': True}) + + def _test_run_python(self, app_label, should_run, hints=None): + with override_settings(DATABASE_ROUTERS=[MigrateEverythingRouter()]): + project_state = self.set_up_test_model(app_label) + + # Create the operation + def inner_method(models, schema_editor): + Pony = models.get_model(app_label, "Pony") + Pony.objects.create(pink=1, weight=3.55) + Pony.objects.create(weight=5) + + operation = migrations.RunPython(inner_method, hints=hints or {}) + # Test the state alteration does nothing + new_state = project_state.clone() + operation.state_forwards(app_label, new_state) + self.assertEqual(new_state, project_state) + # Test the database alteration + self.assertEqual(project_state.apps.get_model(app_label, "Pony").objects.count(), 0) + with connection.schema_editor() as editor: + operation.database_forwards(app_label, editor, project_state, new_state) + Pony = project_state.apps.get_model(app_label, "Pony") + if should_run: + self.assertEqual(Pony.objects.count(), 2) + else: + self.assertEqual(Pony.objects.count(), 0) + + @override_settings(DATABASE_ROUTERS=[MigrateNothingRouter()]) + def test_run_python(self): + self._test_run_python("test_mltdb_runpython", should_run=False) + + @override_settings(DATABASE_ROUTERS=[MigrateWhenFooRouter()]) + def test_run_python2(self): + self._test_run_python("test_mltdb_runpython2", should_run=False) + self._test_run_python("test_mltdb_runpython2", should_run=True, hints={'foo': True}) diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py index f3e33b0f7e..4f46d1a015 100644 --- a/tests/migrations/test_operations.py +++ b/tests/migrations/test_operations.py @@ -1679,44 +1679,6 @@ class OperationTests(OperationTestBase): self.assertEqual(sorted(definition[2]), ["database_operations", "state_operations"]) -class MigrateNothingRouter(object): - """ - A router that doesn't allow storing any model in any database. - """ - def allow_migrate(self, db, model): - return False - - -@override_settings(DATABASE_ROUTERS=[MigrateNothingRouter()]) -class MultiDBOperationTests(MigrationTestBase): - multi_db = True - - def test_create_model(self): - """ - Tests that CreateModel honours multi-db settings. - """ - operation = migrations.CreateModel( - "Pony", - [ - ("id", models.AutoField(primary_key=True)), - ("pink", models.IntegerField(default=1)), - ], - ) - # Test the state alteration - project_state = ProjectState() - new_state = project_state.clone() - operation.state_forwards("test_crmo", new_state) - # Test the database alteration - self.assertTableNotExists("test_crmo_pony") - with connection.schema_editor() as editor: - operation.database_forwards("test_crmo", editor, project_state, new_state) - self.assertTableNotExists("test_crmo_pony") - # And test reversal - with connection.schema_editor() as editor: - operation.database_backwards("test_crmo", editor, new_state, project_state) - self.assertTableNotExists("test_crmo_pony") - - class SwappableOperationTests(OperationTestBase): """ Tests that key operations ignore swappable models