Skip to content
This repository

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Browse code

RunSQL migration operation and alpha SeparateDatabaseAndState op'n.

  • Loading branch information...
commit bacbbb481da8a582642a265459f2144db813dcee 1 parent 9079436
Andrew Godwin authored September 07, 2013
1  django/db/migrations/operations/__init__.py
... ...
@@ -1,2 +1,3 @@
1 1
 from .models import CreateModel, DeleteModel, AlterModelTable, AlterUniqueTogether, AlterIndexTogether
2 2
 from .fields import AddField, RemoveField, AlterField, RenameField
  3
+from .special import SeparateDatabaseAndState, RunSQL
95  django/db/migrations/operations/special.py
... ...
@@ -0,0 +1,95 @@
  1
+import re
  2
+from .base import Operation
  3
+from django.db import models, router
  4
+from django.db.migrations.state import ModelState
  5
+
  6
+
  7
+class SeparateDatabaseAndState(Operation):
  8
+    """
  9
+    Takes two lists of operations - ones that will be used for the database,
  10
+    and ones that will be used for the state change. This allows operations
  11
+    that don't support state change to have it applied, or have operations
  12
+    that affect the state or not the database, or so on.
  13
+    """
  14
+
  15
+    def __init__(self, database_operations=None, state_operations=None):
  16
+        self.database_operations = database_operations or []
  17
+        self.state_operations = state_operations or []
  18
+
  19
+    def state_forwards(self, app_label, state):
  20
+        for state_operation in self.state_operations:
  21
+            state_operation.state_forwards(app_label, state)
  22
+
  23
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
  24
+        # We calculate state separately in here since our state functions aren't useful
  25
+        for database_operation in self.database_operations:
  26
+            to_state = from_state.clone()
  27
+            database_operation.state_forwards(app_label, to_state)
  28
+            database_operation.database_forwards(self, app_label, schema_editor, from_state, to_state)
  29
+            from_state = to_state
  30
+
  31
+    def database_backwards(self, app_label, schema_editor, from_state, to_state):
  32
+        # We calculate state separately in here since our state functions aren't useful
  33
+        base_state = to_state
  34
+        for pos, database_operation in enumerate(reversed(self.database_operations)):
  35
+            to_state = base_state.clone()
  36
+            for dbop in self.database_operations[:-(pos+1)]:
  37
+                dbop.state_forwards(app_label, to_state)
  38
+            from_state = base_state.clone()
  39
+            database_operation.state_forwards(app_label, from_state)
  40
+            database_operation.database_backwards(self, app_label, schema_editor, from_state, to_state)
  41
+
  42
+    def describe(self):
  43
+        return "Custom state/database change combination"
  44
+
  45
+
  46
+class RunSQL(Operation):
  47
+    """
  48
+    Runs some raw SQL - a single statement by default, but it will attempt
  49
+    to parse and split it into multiple statements if multiple=True.
  50
+
  51
+    A reverse SQL statement may be provided.
  52
+
  53
+    Also accepts a list of operations that represent the state change effected
  54
+    by this SQL change, in case it's custom column/table creation/deletion.
  55
+    """
  56
+
  57
+    def __init__(self, sql, reverse_sql=None, state_operations=None, multiple=False):
  58
+        self.sql = sql
  59
+        self.reverse_sql = reverse_sql
  60
+        self.state_operations = state_operations or []
  61
+        self.multiple = multiple
  62
+
  63
+    def state_forwards(self, app_label, state):
  64
+        for state_operation in self.state_operations:
  65
+            state_operation.state_forwards(app_label, state)
  66
+
  67
+    def _split_sql(self, sql):
  68
+        regex = r"(?mx) ([^';]* (?:'[^']*'[^';]*)*)"
  69
+        comment_regex = r"(?mx) (?:^\s*$)|(?:--.*$)"
  70
+        # First, strip comments
  71
+        sql = "\n".join([x.strip().replace("%", "%%") for x in re.split(comment_regex, sql) if x.strip()])
  72
+        # Now get each statement
  73
+        for st in re.split(regex, sql)[1:][::2]:
  74
+            yield st
  75
+
  76
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
  77
+        if self.multiple:
  78
+            statements = self._split_sql(self.sql)
  79
+        else:
  80
+            statements = [self.sql]
  81
+        for statement in statements:
  82
+            schema_editor.execute(statement)
  83
+
  84
+    def database_backwards(self, app_label, schema_editor, from_state, to_state):
  85
+        if self.reverse_sql is None:
  86
+            raise NotImplementedError("You cannot reverse this operation")
  87
+        if self.multiple:
  88
+            statements = self._split_sql(self.reverse_sql)
  89
+        else:
  90
+            statements = [self.reverse_sql]
  91
+        for statement in statements:
  92
+            schema_editor.execute(statement)
  93
+
  94
+    def describe(self):
  95
+        return "Raw SQL operation"
26  tests/migrations/test_operations.py
@@ -280,6 +280,32 @@ def test_alter_index_together(self):
280 280
             operation.database_backwards("test_alinto", editor, new_state, project_state)
281 281
         self.assertIndexNotExists("test_alinto_pony", ["pink", "weight"])
282 282
 
  283
+    def test_run_sql(self):
  284
+        """
  285
+        Tests the AlterIndexTogether operation.
  286
+        """
  287
+        project_state = self.set_up_test_model("test_runsql")
  288
+        # Create the operation
  289
+        operation = migrations.RunSQL(
  290
+            "CREATE TABLE i_love_ponies (id int, special_thing int)",
  291
+            "DROP TABLE i_love_ponies",
  292
+            state_operations = [migrations.CreateModel("SomethingElse", [("id", models.AutoField(primary_key=True))])],
  293
+        )
  294
+        # Test the state alteration
  295
+        new_state = project_state.clone()
  296
+        operation.state_forwards("test_runsql", new_state)
  297
+        self.assertEqual(len(new_state.models["test_runsql", "somethingelse"].fields), 1)
  298
+        # Make sure there's no table
  299
+        self.assertTableNotExists("i_love_ponies")
  300
+        # Test the database alteration
  301
+        with connection.schema_editor() as editor:
  302
+            operation.database_forwards("test_runsql", editor, project_state, new_state)
  303
+        self.assertTableExists("i_love_ponies")
  304
+        # And test reversal
  305
+        with connection.schema_editor() as editor:
  306
+            operation.database_backwards("test_runsql", editor, new_state, project_state)
  307
+        self.assertTableNotExists("i_love_ponies")
  308
+
283 309
 
284 310
 class MigrateNothingRouter(object):
285 311
     """

0 notes on commit bacbbb4

Please sign in to comment.
Something went wrong with that request. Please try again.