You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

graph.py 13KB

5 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319
  1. from functools import total_ordering
  2. from django.db.migrations.state import ProjectState
  3. from .exceptions import CircularDependencyError, NodeNotFoundError
  4. @total_ordering
  5. class Node:
  6. """
  7. A single node in the migration graph. Contains direct links to adjacent
  8. nodes in either direction.
  9. """
  10. def __init__(self, key):
  11. self.key = key
  12. self.children = set()
  13. self.parents = set()
  14. def __eq__(self, other):
  15. return self.key == other
  16. def __lt__(self, other):
  17. return self.key < other
  18. def __hash__(self):
  19. return hash(self.key)
  20. def __getitem__(self, item):
  21. return self.key[item]
  22. def __str__(self):
  23. return str(self.key)
  24. def __repr__(self):
  25. return '<%s: (%r, %r)>' % (self.__class__.__name__, self.key[0], self.key[1])
  26. def add_child(self, child):
  27. self.children.add(child)
  28. def add_parent(self, parent):
  29. self.parents.add(parent)
  30. class DummyNode(Node):
  31. """
  32. A node that doesn't correspond to a migration file on disk.
  33. (A squashed migration that was removed, for example.)
  34. After the migration graph is processed, all dummy nodes should be removed.
  35. If there are any left, a nonexistent dependency error is raised.
  36. """
  37. def __init__(self, key, origin, error_message):
  38. super().__init__(key)
  39. self.origin = origin
  40. self.error_message = error_message
  41. def raise_error(self):
  42. raise NodeNotFoundError(self.error_message, self.key, origin=self.origin)
  43. class MigrationGraph:
  44. """
  45. Represent the digraph of all migrations in a project.
  46. Each migration is a node, and each dependency is an edge. There are
  47. no implicit dependencies between numbered migrations - the numbering is
  48. merely a convention to aid file listing. Every new numbered migration
  49. has a declared dependency to the previous number, meaning that VCS
  50. branch merges can be detected and resolved.
  51. Migrations files can be marked as replacing another set of migrations -
  52. this is to support the "squash" feature. The graph handler isn't responsible
  53. for these; instead, the code to load them in here should examine the
  54. migration files and if the replaced migrations are all either unapplied
  55. or not present, it should ignore the replaced ones, load in just the
  56. replacing migration, and repoint any dependencies that pointed to the
  57. replaced migrations to point to the replacing one.
  58. A node should be a tuple: (app_path, migration_name). The tree special-cases
  59. things within an app - namely, root nodes and leaf nodes ignore dependencies
  60. to other apps.
  61. """
  62. def __init__(self):
  63. self.node_map = {}
  64. self.nodes = {}
  65. def add_node(self, key, migration):
  66. assert key not in self.node_map
  67. node = Node(key)
  68. self.node_map[key] = node
  69. self.nodes[key] = migration
  70. def add_dummy_node(self, key, origin, error_message):
  71. node = DummyNode(key, origin, error_message)
  72. self.node_map[key] = node
  73. self.nodes[key] = None
  74. def add_dependency(self, migration, child, parent, skip_validation=False):
  75. """
  76. This may create dummy nodes if they don't yet exist. If
  77. `skip_validation=True`, validate_consistency() should be called
  78. afterwards.
  79. """
  80. if child not in self.nodes:
  81. error_message = (
  82. "Migration %s dependencies reference nonexistent"
  83. " child node %r" % (migration, child)
  84. )
  85. self.add_dummy_node(child, migration, error_message)
  86. if parent not in self.nodes:
  87. error_message = (
  88. "Migration %s dependencies reference nonexistent"
  89. " parent node %r" % (migration, parent)
  90. )
  91. self.add_dummy_node(parent, migration, error_message)
  92. self.node_map[child].add_parent(self.node_map[parent])
  93. self.node_map[parent].add_child(self.node_map[child])
  94. if not skip_validation:
  95. self.validate_consistency()
  96. def remove_replaced_nodes(self, replacement, replaced):
  97. """
  98. Remove each of the `replaced` nodes (when they exist). Any
  99. dependencies that were referencing them are changed to reference the
  100. `replacement` node instead.
  101. """
  102. # Cast list of replaced keys to set to speed up lookup later.
  103. replaced = set(replaced)
  104. try:
  105. replacement_node = self.node_map[replacement]
  106. except KeyError as err:
  107. raise NodeNotFoundError(
  108. "Unable to find replacement node %r. It was either never added"
  109. " to the migration graph, or has been removed." % (replacement,),
  110. replacement
  111. ) from err
  112. for replaced_key in replaced:
  113. self.nodes.pop(replaced_key, None)
  114. replaced_node = self.node_map.pop(replaced_key, None)
  115. if replaced_node:
  116. for child in replaced_node.children:
  117. child.parents.remove(replaced_node)
  118. # We don't want to create dependencies between the replaced
  119. # node and the replacement node as this would lead to
  120. # self-referencing on the replacement node at a later iteration.
  121. if child.key not in replaced:
  122. replacement_node.add_child(child)
  123. child.add_parent(replacement_node)
  124. for parent in replaced_node.parents:
  125. parent.children.remove(replaced_node)
  126. # Again, to avoid self-referencing.
  127. if parent.key not in replaced:
  128. replacement_node.add_parent(parent)
  129. parent.add_child(replacement_node)
  130. def remove_replacement_node(self, replacement, replaced):
  131. """
  132. The inverse operation to `remove_replaced_nodes`. Almost. Remove the
  133. replacement node `replacement` and remap its child nodes to `replaced`
  134. - the list of nodes it would have replaced. Don't remap its parent
  135. nodes as they are expected to be correct already.
  136. """
  137. self.nodes.pop(replacement, None)
  138. try:
  139. replacement_node = self.node_map.pop(replacement)
  140. except KeyError as err:
  141. raise NodeNotFoundError(
  142. "Unable to remove replacement node %r. It was either never added"
  143. " to the migration graph, or has been removed already." % (replacement,),
  144. replacement
  145. ) from err
  146. replaced_nodes = set()
  147. replaced_nodes_parents = set()
  148. for key in replaced:
  149. replaced_node = self.node_map.get(key)
  150. if replaced_node:
  151. replaced_nodes.add(replaced_node)
  152. replaced_nodes_parents |= replaced_node.parents
  153. # We're only interested in the latest replaced node, so filter out
  154. # replaced nodes that are parents of other replaced nodes.
  155. replaced_nodes -= replaced_nodes_parents
  156. for child in replacement_node.children:
  157. child.parents.remove(replacement_node)
  158. for replaced_node in replaced_nodes:
  159. replaced_node.add_child(child)
  160. child.add_parent(replaced_node)
  161. for parent in replacement_node.parents:
  162. parent.children.remove(replacement_node)
  163. # NOTE: There is no need to remap parent dependencies as we can
  164. # assume the replaced nodes already have the correct ancestry.
  165. def validate_consistency(self):
  166. """Ensure there are no dummy nodes remaining in the graph."""
  167. [n.raise_error() for n in self.node_map.values() if isinstance(n, DummyNode)]
  168. def forwards_plan(self, target):
  169. """
  170. Given a node, return a list of which previous nodes (dependencies) must
  171. be applied, ending with the node itself. This is the list you would
  172. follow if applying the migrations to a database.
  173. """
  174. if target not in self.nodes:
  175. raise NodeNotFoundError("Node %r not a valid node" % (target,), target)
  176. return self.iterative_dfs(self.node_map[target])
  177. def backwards_plan(self, target):
  178. """
  179. Given a node, return a list of which dependent nodes (dependencies)
  180. must be unapplied, ending with the node itself. This is the list you
  181. would follow if removing the migrations from a database.
  182. """
  183. if target not in self.nodes:
  184. raise NodeNotFoundError("Node %r not a valid node" % (target,), target)
  185. return self.iterative_dfs(self.node_map[target], forwards=False)
  186. def iterative_dfs(self, start, forwards=True):
  187. """Iterative depth-first search for finding dependencies."""
  188. visited = []
  189. visited_set = set()
  190. stack = [(start, False)]
  191. while stack:
  192. node, processed = stack.pop()
  193. if node in visited_set:
  194. pass
  195. elif processed:
  196. visited_set.add(node)
  197. visited.append(node.key)
  198. else:
  199. stack.append((node, True))
  200. stack += [(n, False) for n in sorted(node.parents if forwards else node.children)]
  201. return visited
  202. def root_nodes(self, app=None):
  203. """
  204. Return all root nodes - that is, nodes with no dependencies inside
  205. their app. These are the starting point for an app.
  206. """
  207. roots = set()
  208. for node in self.nodes:
  209. if all(key[0] != node[0] for key in self.node_map[node].parents) and (not app or app == node[0]):
  210. roots.add(node)
  211. return sorted(roots)
  212. def leaf_nodes(self, app=None):
  213. """
  214. Return all leaf nodes - that is, nodes with no dependents in their app.
  215. These are the "most current" version of an app's schema.
  216. Having more than one per app is technically an error, but one that
  217. gets handled further up, in the interactive command - it's usually the
  218. result of a VCS merge and needs some user input.
  219. """
  220. leaves = set()
  221. for node in self.nodes:
  222. if all(key[0] != node[0] for key in self.node_map[node].children) and (not app or app == node[0]):
  223. leaves.add(node)
  224. return sorted(leaves)
  225. def ensure_not_cyclic(self):
  226. # Algo from GvR:
  227. # https://neopythonic.blogspot.com/2009/01/detecting-cycles-in-directed-graph.html
  228. todo = set(self.nodes)
  229. while todo:
  230. node = todo.pop()
  231. stack = [node]
  232. while stack:
  233. top = stack[-1]
  234. for child in self.node_map[top].children:
  235. # Use child.key instead of child to speed up the frequent
  236. # hashing.
  237. node = child.key
  238. if node in stack:
  239. cycle = stack[stack.index(node):]
  240. raise CircularDependencyError(", ".join("%s.%s" % n for n in cycle))
  241. if node in todo:
  242. stack.append(node)
  243. todo.remove(node)
  244. break
  245. else:
  246. node = stack.pop()
  247. def __str__(self):
  248. return 'Graph: %s nodes, %s edges' % self._nodes_and_edges()
  249. def __repr__(self):
  250. nodes, edges = self._nodes_and_edges()
  251. return '<%s: nodes=%s, edges=%s>' % (self.__class__.__name__, nodes, edges)
  252. def _nodes_and_edges(self):
  253. return len(self.nodes), sum(len(node.parents) for node in self.node_map.values())
  254. def _generate_plan(self, nodes, at_end):
  255. plan = []
  256. for node in nodes:
  257. for migration in self.forwards_plan(node):
  258. if migration not in plan and (at_end or migration not in nodes):
  259. plan.append(migration)
  260. return plan
  261. def make_state(self, nodes=None, at_end=True, real_apps=None):
  262. """
  263. Given a migration node or nodes, return a complete ProjectState for it.
  264. If at_end is False, return the state before the migration has run.
  265. If nodes is not provided, return the overall most current project state.
  266. """
  267. if nodes is None:
  268. nodes = list(self.leaf_nodes())
  269. if not nodes:
  270. return ProjectState()
  271. if not isinstance(nodes[0], tuple):
  272. nodes = [nodes]
  273. plan = self._generate_plan(nodes, at_end)
  274. project_state = ProjectState(real_apps=real_apps)
  275. for node in plan:
  276. project_state = self.nodes[node].mutate_state(project_state, preserve=False)
  277. return project_state
  278. def __contains__(self, node):
  279. return node in self.nodes