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