Skip to content

Commit

Permalink
add RunMQL operation
Browse files Browse the repository at this point in the history
  • Loading branch information
timgraham committed Dec 29, 2024
1 parent c90a00c commit 72ef10c
Show file tree
Hide file tree
Showing 4 changed files with 207 additions and 0 deletions.
59 changes: 59 additions & 0 deletions django_mongodb/migrations.py
Original file line number Diff line number Diff line change
@@ -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"
72 changes: 72 additions & 0 deletions docs/source/migrations.rst
Original file line number Diff line number Diff line change
@@ -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):
Empty file added tests/migrations_/__init__.py
Empty file.
76 changes: 76 additions & 0 deletions tests/migrations_/test_operations.py
Original file line number Diff line number Diff line change
@@ -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'")

0 comments on commit 72ef10c

Please sign in to comment.