Adding a dependency graph class and tests

django/db/migrations/graph.py

+from django.utils.datastructures import SortedSet
+
+
+class MigrationsGraph(object):
+    """
+    Represents the digraph of all migrations in a project.
+
+    Each migration is a node, and each dependency is an edge. There are
+    no implicit dependencies between numbered migrations - the numbering is
+    merely a convention to aid file listing. Every new numbered migration
+    has a declared dependency to the previous number, meaning that VCS
+    branch merges can be detected and resolved.
+
+    Migrations files can be marked as replacing another set of migrations -
+    this is to support the "squash" feature. The graph handler isn't resposible
+    for these; instead, the code to load them in here should examine the
+    migration files and if the replaced migrations are all either unapplied
+    or not present, it should ignore the replaced ones, load in just the
+    replacing migration, and repoint any dependencies that pointed to the
+    replaced migrations to point to the replacing one.
+
+    A node should be a tuple: (applabel, migration_name) - but the code
+    here doesn't really care.
+    """
+
+    def __init__(self):
+        self.nodes = {}
+        self.dependencies = {}
+        self.dependents = {}
+
+    def add_node(self, node, implementation):
+        self.nodes[node] = implementation
+
+    def add_dependency(self, child, parent):
+        self.nodes[child] = None
+        self.nodes[parent] = None
+        self.dependencies.setdefault(child, set()).add(parent)
+        self.dependents.setdefault(parent, set()).add(child)
+
+    def forwards_plan(self, node):
+        """
+        Given a node, returns a list of which previous nodes (dependencies)
+        must be applied, ending with the node itself.
+        This is the list you would follow if applying the migrations to
+        a database.
+        """
+        if node not in self.nodes:
+            raise ValueError("Node %r not a valid node" % node)
+        return self.dfs(node, lambda x: self.dependencies.get(x, set()))
+
+    def backwards_plan(self, node):
+        """
+        Given a node, returns a list of which dependent nodes (dependencies)
+        must be unapplied, ending with the node itself.
+        This is the list you would follow if removing the migrations from
+        a database.
+        """
+        if node not in self.nodes:
+            raise ValueError("Node %r not a valid node" % node)
+        return self.dfs(node, lambda x: self.dependents.get(x, set()))
+
+    def dfs(self, start, get_children):
+        """
+        Dynamic programming based depth first search, for finding dependencies.
+        """
+        cache = {}
+        def _dfs(start, get_children, path):
+            # If we already computed this, use that (dynamic programming)
+            if (start, get_children) in cache:
+                return cache[(start, get_children)]
+            # If we've traversed here before, that's a circular dep
+            if start in path:
+                raise CircularDependencyException(path[path.index(start):] + [start])
+            # Build our own results list, starting with us
+            results = []
+            results.append(start)
+            # We need to add to results all the migrations this one depends on
+            children = sorted(get_children(start))
+            path.append(start)
+            for n in children:
+                results = _dfs(n, get_children, path) + results
+            path.pop()
+            # Use SortedSet to ensure only one instance of each result
+            results = list(SortedSet(results))
+            # Populate DP cache
+            cache[(start, get_children)] = results
+            # Done!
+            return results
+        return _dfs(start, get_children, [])
+
+
+class CircularDependencyException(Exception):
+    """
+    Raised when there's an impossible-to-resolve circular dependency.
+    """
+    pass

django/utils/datastructures.py

@@ -252,6 +252,36 @@ class SortedDict(dict):
         super(SortedDict, self).clear()
         self.keyOrder = []
 
+class SortedSet(object):
+    """
+    A set which keeps the ordering of the inserted items.
+    Currently backs onto SortedDict.
+    """
+
+    def __init__(self, iterable=None):
+        self.dict = SortedDict(((x, None) for x in iterable) if iterable else [])
+
+    def add(self, item):
+        self.dict[item] = None
+
+    def remove(self, item):
+        del self.dict[item]
+
+    def discard(self, item):
+        try:
+            self.remove(item)
+        except KeyError:
+            pass
+
+    def __iter__(self):
+        return iter(self.dict.keys())
+
+    def __contains__(self, item):
+        return item in self.dict
+
+    def __nonzero__(self):
+        return bool(self.dict)
+
 class MultiValueDictKeyError(KeyError):
     pass
 

tests/migrations/tests.py

+from django.test import TransactionTestCase
+from django.db.migrations.graph import MigrationsGraph, CircularDependencyException
+
+
+class GraphTests(TransactionTestCase):
+    """
+    Tests the digraph structure.
+    """
+
+    def test_simple_graph(self):
+        """
+        Tests a basic dependency graph:
+
+        app_a:  0001 <-- 0002 <--- 0003 <-- 0004
+                                 /
+        app_b:  0001 <-- 0002 <-/
+        """
+        # Build graph
+        graph = MigrationsGraph()
+        graph.add_dependency(("app_a", "0004"), ("app_a", "0003"))
+        graph.add_dependency(("app_a", "0003"), ("app_a", "0002"))
+        graph.add_dependency(("app_a", "0002"), ("app_a", "0001"))
+        graph.add_dependency(("app_a", "0003"), ("app_b", "0002"))
+        graph.add_dependency(("app_b", "0002"), ("app_b", "0001"))
+        # Test root migration case
+        self.assertEqual(
+            graph.forwards_plan(("app_a", "0001")),
+            [('app_a', '0001')],
+        )
+        # Test branch B only
+        self.assertEqual(
+            graph.forwards_plan(("app_b", "0002")),
+            [("app_b", "0001"), ("app_b", "0002")],
+        )
+        # Test whole graph
+        self.assertEqual(
+            graph.forwards_plan(("app_a", "0004")),
+            [('app_b', '0001'), ('app_b', '0002'), ('app_a', '0001'), ('app_a', '0002'), ('app_a', '0003'), ('app_a', '0004')],
+        )
+        # Test reverse to b:0002
+        self.assertEqual(
+            graph.backwards_plan(("app_b", "0002")),
+            [('app_a', '0004'), ('app_a', '0003'), ('app_b', '0002')],
+        )
+
+    def test_complex_graph(self):
+        """
+        Tests a complex dependency graph:
+
+        app_a:  0001 <-- 0002 <--- 0003 <-- 0004
+                      \        \ /         /
+        app_b:  0001 <-\ 0002 <-X         /
+                      \          \       /
+        app_c:         \ 0001 <-- 0002 <-
+        """
+        # Build graph
+        graph = MigrationsGraph()
+        graph.add_dependency(("app_a", "0004"), ("app_a", "0003"))
+        graph.add_dependency(("app_a", "0003"), ("app_a", "0002"))
+        graph.add_dependency(("app_a", "0002"), ("app_a", "0001"))
+        graph.add_dependency(("app_a", "0003"), ("app_b", "0002"))
+        graph.add_dependency(("app_b", "0002"), ("app_b", "0001"))
+        graph.add_dependency(("app_a", "0004"), ("app_c", "0002"))
+        graph.add_dependency(("app_c", "0002"), ("app_c", "0001"))
+        graph.add_dependency(("app_c", "0001"), ("app_b", "0001"))
+        graph.add_dependency(("app_c", "0002"), ("app_a", "0002"))
+        # Test branch C only
+        self.assertEqual(
+            graph.forwards_plan(("app_c", "0002")),
+            [('app_b', '0001'), ('app_c', '0001'), ('app_a', '0001'), ('app_a', '0002'), ('app_c', '0002')],
+        )
+        # Test whole graph
+        self.assertEqual(
+            graph.forwards_plan(("app_a", "0004")),
+            [('app_b', '0001'), ('app_c', '0001'), ('app_a', '0001'), ('app_a', '0002'), ('app_c', '0002'), ('app_b', '0002'), ('app_a', '0003'), ('app_a', '0004')],
+        )
+        # Test reverse to b:0001
+        self.assertEqual(
+            graph.backwards_plan(("app_b", "0001")),
+            [('app_a', '0004'), ('app_c', '0002'), ('app_c', '0001'), ('app_a', '0003'), ('app_b', '0002'), ('app_b', '0001')],
+        )
+
+    def test_circular_graph(self):
+        """
+        Tests a circular dependency graph.
+        """
+        # Build graph
+        graph = MigrationsGraph()
+        graph.add_dependency(("app_a", "0003"), ("app_a", "0002"))
+        graph.add_dependency(("app_a", "0002"), ("app_a", "0001"))
+        graph.add_dependency(("app_a", "0001"), ("app_b", "0002"))
+        graph.add_dependency(("app_b", "0002"), ("app_b", "0001"))
+        graph.add_dependency(("app_b", "0001"), ("app_a", "0003"))
+        # Test whole graph
+        self.assertRaises(
+            CircularDependencyException,
+            graph.forwards_plan, ("app_a", "0003"),
+        )