From 72ef10cf645c91deb9d8bf58fa2f20dff126754d Mon Sep 17 00:00:00 2001 From: Tim Graham Date: Sat, 28 Dec 2024 21:17:31 -0500 Subject: [PATCH] add RunMQL operation --- django_mongodb/migrations.py | 59 +++++++++++++++++++++ docs/source/migrations.rst | 72 ++++++++++++++++++++++++++ tests/migrations_/__init__.py | 0 tests/migrations_/test_operations.py | 76 ++++++++++++++++++++++++++++ 4 files changed, 207 insertions(+) create mode 100644 django_mongodb/migrations.py create mode 100644 docs/source/migrations.rst create mode 100644 tests/migrations_/__init__.py create mode 100644 tests/migrations_/test_operations.py diff --git a/django_mongodb/migrations.py b/django_mongodb/migrations.py new file mode 100644 index 00000000..a0144fe0 --- /dev/null +++ b/django_mongodb/migrations.py @@ -0,0 +1,59 @@ +from django.db import router +from django.db.migrations.operations.base import Operation + + +class RunMQL(Operation): + """ + Run some raw MQL. A reverse MQL statement may be provided. + + Also accept a list of operations that represent the state change effected + by this MQL change, in case it's custom column/table creation/deletion. + """ + + noop = "" + + def __init__(self, code, reverse_code=None, state_operations=None, hints=None, elidable=False): + # Forwards code + if not callable(code): + raise ValueError("RunMQL must be supplied with a callable") + self.code = code + # Reverse code + self.reverse_code = reverse_code + if reverse_code is not None and not callable(reverse_code): + raise ValueError("RunMQL must be supplied with callable arguments") + self.state_operations = state_operations or [] + self.hints = hints or {} + self.elidable = elidable + + def deconstruct(self): + kwargs = { + "code": self.code, + } + if self.reverse_code is not None: + kwargs["reverse_code"] = self.reverse_code + if self.state_operations: + kwargs["state_operations"] = self.state_operations + if self.hints: + kwargs["hints"] = self.hints + return (self.__class__.__qualname__, [], kwargs) + + @property + def reversible(self): + return self.reverse_code is not None + + def state_forwards(self, app_label, state): + for state_operation in self.state_operations: + state_operation.state_forwards(app_label, state) + + def database_forwards(self, app_label, schema_editor, from_state, to_state): + if router.allow_migrate(schema_editor.connection.alias, app_label, **self.hints): + self.code(schema_editor, schema_editor.get_database()) + + 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") + if router.allow_migrate(schema_editor.connection.alias, app_label, **self.hints): + self.reverse_code(schema_editor, schema_editor.get_database()) + + def describe(self): + return "Raw MQL operation" diff --git a/docs/source/migrations.rst b/docs/source/migrations.rst new file mode 100644 index 00000000..f98cc71f --- /dev/null +++ b/docs/source/migrations.rst @@ -0,0 +1,72 @@ +Migrations API reference +======================== + +You can use PyMongo operations in your migrations. In lieu of the ``RunMQL``, +use ``RunMQL``. + + + +``RunMQL`` +---------- + +.. class:: RunMQL(code, reverse_code=None, state_operations=None, hints=None, elidable=False) + +Allows running of arbitrary MQL on the database - useful for more advanced +features of database backends that Django doesn't support directly. + +``sql``, and ``reverse_sql`` if provided, should be strings of MQL to run on +the database. On most database backends (all but PostgreMQL), Django will +split the MQL into individual statements prior to executing them. + +If you want to include literal percent signs in the query, you have to double +them if you are passing parameters. + +The ``reverse_sql`` queries are executed when the migration is unapplied. They +should undo what is done by the ``sql`` queries. For example, to undo the above +insertion with a deletion:: + + migrations.RunMQL( + forward_func=[("INSERT INTO musician (name) VALUES (%s);", ["Reinhardt"])], + reverse_func=[("DELETE FROM musician where name=%s;", ["Reinhardt"])], + ) + +If ``reverse_sql`` is ``None`` (the default), the ``RunMQL`` operation is +irreversible. + +The ``state_operations`` argument allows you to supply operations that are +equivalent to the MQL in terms of project state. For example, if you are +manually creating a column, you should pass in a list containing an ``AddField`` +operation here so that the autodetector still has an up-to-date state of the +model. If you don't, when you next run ``makemigrations``, it won't see any +operation that adds that field and so will try to run it again. For example:: + + migrations.RunMQL( + "ALTER TABLE musician ADD COLUMN name varchar(255) NOT NULL;", + state_operations=[ + migrations.AddField( + "musician", + "name", + models.CharField(max_length=255), + ), + ], + ) + +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. + +The optional ``elidable`` argument determines whether or not the operation will +be removed (elided) when :ref:`squashing migrations `. + +.. attribute:: RunMQL.noop + + Pass the ``RunMQL.noop`` attribute to ``sql`` or ``reverse_sql`` when you + want the operation not to do anything in the given direction. This is + especially useful in making the operation reversible. + + + + + +def forwards_func(apps, schema_editor, database): diff --git a/tests/migrations_/__init__.py b/tests/migrations_/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/migrations_/test_operations.py b/tests/migrations_/test_operations.py new file mode 100644 index 00000000..1158e45d --- /dev/null +++ b/tests/migrations_/test_operations.py @@ -0,0 +1,76 @@ +from django.db import connection +from django.db.migrations.state import ProjectState +from migrations.test_base import OperationTestBase + +from django_mongodb.migrations import RunMQL + + +def create_collection(schema_editor, database): # noqa: ARG001 + database.create_collection("test_runmql") + + +def drop_collection(schema_editor, database): # noqa: ARG001 + database.drop_collection("test_runmql") + + +class RunMQLTests(OperationTestBase): + available_apps = ["migrations_"] + + def test_basic(self): + project_state = ProjectState() + operation = RunMQL(create_collection, reverse_code=drop_collection) + self.assertEqual(operation.describe(), "Raw MQL operation") + # Test the state alteration does nothing + new_state = project_state.clone() + operation.state_forwards("test_runmql", new_state) + self.assertEqual(new_state, project_state) + # Test the database alteration + self.assertTableNotExists("test_runmql") + with connection.schema_editor() as editor: + operation.database_forwards("test_runmql", editor, project_state, new_state) + self.assertTableExists("test_runmql") + # Now test reversal + self.assertTrue(operation.reversible) + with connection.schema_editor() as editor: + operation.database_backwards("test_runmql", editor, project_state, new_state) + self.assertTableNotExists("test_runmql") + # Test deconstruction + definition = operation.deconstruct() + self.assertEqual(definition[0], "RunMQL") + self.assertEqual(definition[1], []) + self.assertEqual(sorted(definition[2]), ["code", "reverse_code"]) + # Also test reversal fails, with an operation identical to above but + # without reverse_code set. + no_reverse_operation = RunMQL(create_collection) + self.assertFalse(no_reverse_operation.reversible) + with connection.schema_editor() as editor: + no_reverse_operation.database_forwards("test_runmql", editor, project_state, new_state) + with self.assertRaises(NotImplementedError): + no_reverse_operation.database_backwards( + "test_runmql", editor, new_state, project_state + ) + self.assertTableExists("test_runmql") + + def test_run_msql_no_reverse(self): + project_state = ProjectState() + new_state = project_state.clone() + operation = RunMQL(create_collection) + self.assertTableNotExists("test_runmql") + with connection.schema_editor() as editor: + operation.database_forwards("test_runmql", editor, project_state, new_state) + self.assertTableExists("test_runmql") + # And deconstruction + definition = operation.deconstruct() + self.assertEqual(definition[0], "RunMQL") + self.assertEqual(definition[1], []) + self.assertEqual(sorted(definition[2]), ["code"]) + + def test_elidable(self): + operation = RunMQL(create_collection) + self.assertIs(operation.reduce(operation, []), False) + elidable_operation = RunMQL(create_collection, elidable=True) + self.assertEqual(elidable_operation.reduce(operation, []), [operation]) + + def test_run_mql_invalid_code(self): + with self.assertRaisesMessage(ValueError, "RunMQL must be supplied with a callable"): + RunMQL("print 'ahahaha'")