1import pkgutil
2import sys
3from importlib import import_module, reload
4
5from plain.models.migrations.graph import MigrationGraph
6from plain.models.migrations.recorder import MigrationRecorder
7from plain.packages import packages
8
9from .exceptions import (
10 AmbiguityError,
11 BadMigrationError,
12 InconsistentMigrationHistory,
13 NodeNotFoundError,
14)
15
16MIGRATIONS_MODULE_NAME = "migrations"
17
18
19class MigrationLoader:
20 """
21 Load migration files from disk and their status from the database.
22
23 Migration files are expected to live in the "migrations" directory of
24 an app. Their names are entirely unimportant from a code perspective,
25 but will probably follow the 1234_name.py convention.
26
27 On initialization, this class will scan those directories, and open and
28 read the Python files, looking for a class called Migration, which should
29 inherit from plain.models.migrations.Migration. See
30 plain.models.migrations.migration for what that looks like.
31
32 Some migrations will be marked as "replacing" another set of migrations.
33 These are loaded into a separate set of migrations away from the main ones.
34 If all the migrations they replace are either unapplied or missing from
35 disk, then they are injected into the main set, replacing the named migrations.
36 Any dependency pointers to the replaced migrations are re-pointed to the
37 new migration.
38
39 This does mean that this class MUST also talk to the database as well as
40 to disk, but this is probably fine. We're already not just operating
41 in memory.
42 """
43
44 def __init__(
45 self,
46 connection,
47 load=True,
48 ignore_no_migrations=False,
49 replace_migrations=True,
50 ):
51 self.connection = connection
52 self.disk_migrations = None
53 self.applied_migrations = None
54 self.ignore_no_migrations = ignore_no_migrations
55 self.replace_migrations = replace_migrations
56 if load:
57 self.build_graph()
58
59 @classmethod
60 def migrations_module(cls, package_label):
61 """
62 Return the path to the migrations module for the specified package_label
63 and a boolean indicating if the module is specified in
64 settings.MIGRATION_MODULE.
65 """
66
67 app = packages.get_package_config(package_label)
68 if app.migrations_module is None:
69 return None, True
70 explicit = app.migrations_module != MIGRATIONS_MODULE_NAME
71 return f"{app.name}.{app.migrations_module}", explicit
72
73 def load_disk(self):
74 """Load the migrations from all INSTALLED_PACKAGES from disk."""
75 self.disk_migrations = {}
76 self.unmigrated_packages = set()
77 self.migrated_packages = set()
78 for package_config in packages.get_package_configs():
79 # Get the migrations module directory
80 module_name, explicit = self.migrations_module(package_config.label)
81 if module_name is None:
82 self.unmigrated_packages.add(package_config.label)
83 continue
84 was_loaded = module_name in sys.modules
85 try:
86 module = import_module(module_name)
87 except ModuleNotFoundError as e:
88 if (explicit and self.ignore_no_migrations) or (
89 not explicit and MIGRATIONS_MODULE_NAME in e.name.split(".")
90 ):
91 self.unmigrated_packages.add(package_config.label)
92 continue
93 raise
94 else:
95 # Module is not a package (e.g. migrations.py).
96 if not hasattr(module, "__path__"):
97 self.unmigrated_packages.add(package_config.label)
98 continue
99 # Empty directories are namespaces. Namespace packages have no
100 # __file__ and don't use a list for __path__. See
101 # https://docs.python.org/3/reference/import.html#namespace-packages
102 if getattr(module, "__file__", None) is None and not isinstance(
103 module.__path__, list
104 ):
105 self.unmigrated_packages.add(package_config.label)
106 continue
107 # Force a reload if it's already loaded (tests need this)
108 if was_loaded:
109 reload(module)
110 self.migrated_packages.add(package_config.label)
111 migration_names = {
112 name
113 for _, name, is_pkg in pkgutil.iter_modules(module.__path__)
114 if not is_pkg and name[0] not in "_~"
115 }
116 # Load migrations
117 for migration_name in migration_names:
118 migration_path = f"{module_name}.{migration_name}"
119 try:
120 migration_module = import_module(migration_path)
121 except ImportError as e:
122 if "bad magic number" in str(e):
123 raise ImportError(
124 "Couldn't import %r as it appears to be a stale "
125 ".pyc file." % migration_path
126 ) from e
127 else:
128 raise
129 if not hasattr(migration_module, "Migration"):
130 raise BadMigrationError(
131 "Migration {} in app {} has no Migration class".format(
132 migration_name, package_config.label
133 )
134 )
135 self.disk_migrations[
136 package_config.label, migration_name
137 ] = migration_module.Migration(
138 migration_name,
139 package_config.label,
140 )
141
142 def get_migration(self, package_label, name_prefix):
143 """Return the named migration or raise NodeNotFoundError."""
144 return self.graph.nodes[package_label, name_prefix]
145
146 def get_migration_by_prefix(self, package_label, name_prefix):
147 """
148 Return the migration(s) which match the given app label and name_prefix.
149 """
150 # Do the search
151 results = []
152 for migration_package_label, migration_name in self.disk_migrations:
153 if migration_package_label == package_label and migration_name.startswith(
154 name_prefix
155 ):
156 results.append((migration_package_label, migration_name))
157 if len(results) > 1:
158 raise AmbiguityError(
159 "There is more than one migration for '{}' with the prefix '{}'".format(
160 package_label, name_prefix
161 )
162 )
163 elif not results:
164 raise KeyError(
165 f"There is no migration for '{package_label}' with the prefix "
166 f"'{name_prefix}'"
167 )
168 else:
169 return self.disk_migrations[results[0]]
170
171 def check_key(self, key, current_package):
172 if (key[1] != "__first__" and key[1] != "__latest__") or key in self.graph:
173 return key
174 # Special-case __first__, which means "the first migration" for
175 # migrated packages, and is ignored for unmigrated packages. It allows
176 # makemigrations to declare dependencies on packages before they even have
177 # migrations.
178 if key[0] == current_package:
179 # Ignore __first__ references to the same app (#22325)
180 return
181 if key[0] in self.unmigrated_packages:
182 # This app isn't migrated, but something depends on it.
183 # The models will get auto-added into the state, though
184 # so we're fine.
185 return
186 if key[0] in self.migrated_packages:
187 try:
188 if key[1] == "__first__":
189 return self.graph.root_nodes(key[0])[0]
190 else: # "__latest__"
191 return self.graph.leaf_nodes(key[0])[0]
192 except IndexError:
193 if self.ignore_no_migrations:
194 return None
195 else:
196 raise ValueError(
197 "Dependency on app with no migrations: %s" % key[0]
198 )
199 raise ValueError("Dependency on unknown app: %s" % key[0])
200
201 def add_internal_dependencies(self, key, migration):
202 """
203 Internal dependencies need to be added first to ensure `__first__`
204 dependencies find the correct root node.
205 """
206 for parent in migration.dependencies:
207 # Ignore __first__ references to the same app.
208 if parent[0] == key[0] and parent[1] != "__first__":
209 self.graph.add_dependency(migration, key, parent, skip_validation=True)
210
211 def add_external_dependencies(self, key, migration):
212 for parent in migration.dependencies:
213 # Skip internal dependencies
214 if key[0] == parent[0]:
215 continue
216 parent = self.check_key(parent, key[0])
217 if parent is not None:
218 self.graph.add_dependency(migration, key, parent, skip_validation=True)
219 for child in migration.run_before:
220 child = self.check_key(child, key[0])
221 if child is not None:
222 self.graph.add_dependency(migration, child, key, skip_validation=True)
223
224 def build_graph(self):
225 """
226 Build a migration dependency graph using both the disk and database.
227 You'll need to rebuild the graph if you apply migrations. This isn't
228 usually a problem as generally migration stuff runs in a one-shot process.
229 """
230 # Load disk data
231 self.load_disk()
232 # Load database data
233 if self.connection is None:
234 self.applied_migrations = {}
235 else:
236 recorder = MigrationRecorder(self.connection)
237 self.applied_migrations = recorder.applied_migrations()
238 # To start, populate the migration graph with nodes for ALL migrations
239 # and their dependencies. Also make note of replacing migrations at this step.
240 self.graph = MigrationGraph()
241 self.replacements = {}
242 for key, migration in self.disk_migrations.items():
243 self.graph.add_node(key, migration)
244 # Replacing migrations.
245 if migration.replaces:
246 self.replacements[key] = migration
247 for key, migration in self.disk_migrations.items():
248 # Internal (same app) dependencies.
249 self.add_internal_dependencies(key, migration)
250 # Add external dependencies now that the internal ones have been resolved.
251 for key, migration in self.disk_migrations.items():
252 self.add_external_dependencies(key, migration)
253 # Carry out replacements where possible and if enabled.
254 if self.replace_migrations:
255 for key, migration in self.replacements.items():
256 # Get applied status of each of this migration's replacement
257 # targets.
258 applied_statuses = [
259 (target in self.applied_migrations) for target in migration.replaces
260 ]
261 # The replacing migration is only marked as applied if all of
262 # its replacement targets are.
263 if all(applied_statuses):
264 self.applied_migrations[key] = migration
265 else:
266 self.applied_migrations.pop(key, None)
267 # A replacing migration can be used if either all or none of
268 # its replacement targets have been applied.
269 if all(applied_statuses) or (not any(applied_statuses)):
270 self.graph.remove_replaced_nodes(key, migration.replaces)
271 else:
272 # This replacing migration cannot be used because it is
273 # partially applied. Remove it from the graph and remap
274 # dependencies to it (#25945).
275 self.graph.remove_replacement_node(key, migration.replaces)
276 # Ensure the graph is consistent.
277 try:
278 self.graph.validate_consistency()
279 except NodeNotFoundError as exc:
280 # Check if the missing node could have been replaced by any squash
281 # migration but wasn't because the squash migration was partially
282 # applied before. In that case raise a more understandable exception
283 # (#23556).
284 # Get reverse replacements.
285 reverse_replacements = {}
286 for key, migration in self.replacements.items():
287 for replaced in migration.replaces:
288 reverse_replacements.setdefault(replaced, set()).add(key)
289 # Try to reraise exception with more detail.
290 if exc.node in reverse_replacements:
291 candidates = reverse_replacements.get(exc.node, set())
292 is_replaced = any(
293 candidate in self.graph.nodes for candidate in candidates
294 )
295 if not is_replaced:
296 tries = ", ".join("{}.{}".format(*c) for c in candidates)
297 raise NodeNotFoundError(
298 "Migration {0} depends on nonexistent node ('{1}', '{2}'). "
299 "Plain tried to replace migration {1}.{2} with any of [{3}] "
300 "but wasn't able to because some of the replaced migrations "
301 "are already applied.".format(
302 exc.origin, exc.node[0], exc.node[1], tries
303 ),
304 exc.node,
305 ) from exc
306 raise
307 self.graph.ensure_not_cyclic()
308
309 def check_consistent_history(self, connection):
310 """
311 Raise InconsistentMigrationHistory if any applied migrations have
312 unapplied dependencies.
313 """
314 recorder = MigrationRecorder(connection)
315 applied = recorder.applied_migrations()
316 for migration in applied:
317 # If the migration is unknown, skip it.
318 if migration not in self.graph.nodes:
319 continue
320 for parent in self.graph.node_map[migration].parents:
321 if parent not in applied:
322 # Skip unapplied squashed migrations that have all of their
323 # `replaces` applied.
324 if parent in self.replacements:
325 if all(
326 m in applied for m in self.replacements[parent].replaces
327 ):
328 continue
329 raise InconsistentMigrationHistory(
330 "Migration {}.{} is applied before its dependency "
331 "{}.{} on database '{}'.".format(
332 migration[0],
333 migration[1],
334 parent[0],
335 parent[1],
336 connection.alias,
337 )
338 )
339
340 def detect_conflicts(self):
341 """
342 Look through the loaded graph and detect any conflicts - packages
343 with more than one leaf migration. Return a dict of the app labels
344 that conflict with the migration names that conflict.
345 """
346 seen_packages = {}
347 conflicting_packages = set()
348 for package_label, migration_name in self.graph.leaf_nodes():
349 if package_label in seen_packages:
350 conflicting_packages.add(package_label)
351 seen_packages.setdefault(package_label, set()).add(migration_name)
352 return {
353 package_label: sorted(seen_packages[package_label])
354 for package_label in conflicting_packages
355 }
356
357 def project_state(self, nodes=None, at_end=True):
358 """
359 Return a ProjectState object representing the most recent state
360 that the loaded migrations represent.
361
362 See graph.make_state() for the meaning of "nodes" and "at_end".
363 """
364 return self.graph.make_state(
365 nodes=nodes, at_end=at_end, real_packages=self.unmigrated_packages
366 )
367
368 def collect_sql(self, plan):
369 """
370 Take a migration plan and return a list of collected SQL statements
371 that represent the best-efforts version of that plan.
372 """
373 statements = []
374 state = None
375 for migration, backwards in plan:
376 with self.connection.schema_editor(
377 collect_sql=True, atomic=migration.atomic
378 ) as schema_editor:
379 if state is None:
380 state = self.project_state(
381 (migration.package_label, migration.name), at_end=False
382 )
383 if not backwards:
384 state = migration.apply(state, schema_editor, collect_sql=True)
385 else:
386 state = migration.unapply(state, schema_editor, collect_sql=True)
387 statements.extend(schema_editor.collected_sql)
388 return statements