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