add RunMQL operation

Author: timgraham

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"

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 <migration-squashing>`.
+
+.. 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):

tests/migrations_/__init__.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'")