Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-models/plain/models/migrations/loader.py: 51%
180 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
1import pkgutil
2import sys
3from importlib import import_module, reload
5from plain.models.migrations.graph import MigrationGraph
6from plain.models.migrations.recorder import MigrationRecorder
7from plain.packages import packages
9from .exceptions import (
10 AmbiguityError,
11 BadMigrationError,
12 InconsistentMigrationHistory,
13 NodeNotFoundError,
14)
16MIGRATIONS_MODULE_NAME = "migrations"
19class MigrationLoader:
20 """
21 Load migration files from disk and their status from the database.
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.
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.
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.
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 """
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()
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 """
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
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 f"Couldn't import {migration_path!r} as it appears to be a stale "
125 ".pyc file."
126 ) from e
127 else:
128 raise
129 if not hasattr(migration_module, "Migration"):
130 raise BadMigrationError(
131 f"Migration {migration_name} in app {package_config.label} has no Migration class"
132 )
133 self.disk_migrations[package_config.label, migration_name] = (
134 migration_module.Migration(
135 migration_name,
136 package_config.label,
137 )
138 )
140 def get_migration(self, package_label, name_prefix):
141 """Return the named migration or raise NodeNotFoundError."""
142 return self.graph.nodes[package_label, name_prefix]
144 def get_migration_by_prefix(self, package_label, name_prefix):
145 """
146 Return the migration(s) which match the given app label and name_prefix.
147 """
148 # Do the search
149 results = []
150 for migration_package_label, migration_name in self.disk_migrations:
151 if migration_package_label == package_label and migration_name.startswith(
152 name_prefix
153 ):
154 results.append((migration_package_label, migration_name))
155 if len(results) > 1:
156 raise AmbiguityError(
157 f"There is more than one migration for '{package_label}' with the prefix '{name_prefix}'"
158 )
159 elif not results:
160 raise KeyError(
161 f"There is no migration for '{package_label}' with the prefix "
162 f"'{name_prefix}'"
163 )
164 else:
165 return self.disk_migrations[results[0]]
167 def check_key(self, key, current_package):
168 if (key[1] != "__first__" and key[1] != "__latest__") or key in self.graph:
169 return key
170 # Special-case __first__, which means "the first migration" for
171 # migrated packages, and is ignored for unmigrated packages. It allows
172 # makemigrations to declare dependencies on packages before they even have
173 # migrations.
174 if key[0] == current_package:
175 # Ignore __first__ references to the same app (#22325)
176 return
177 if key[0] in self.unmigrated_packages:
178 # This app isn't migrated, but something depends on it.
179 # The models will get auto-added into the state, though
180 # so we're fine.
181 return
182 if key[0] in self.migrated_packages:
183 try:
184 if key[1] == "__first__":
185 return self.graph.root_nodes(key[0])[0]
186 else: # "__latest__"
187 return self.graph.leaf_nodes(key[0])[0]
188 except IndexError:
189 if self.ignore_no_migrations:
190 return None
191 else:
192 raise ValueError(f"Dependency on app with no migrations: {key[0]}")
193 raise ValueError(f"Dependency on unknown app: {key[0]}")
195 def add_internal_dependencies(self, key, migration):
196 """
197 Internal dependencies need to be added first to ensure `__first__`
198 dependencies find the correct root node.
199 """
200 for parent in migration.dependencies:
201 # Ignore __first__ references to the same app.
202 if parent[0] == key[0] and parent[1] != "__first__":
203 self.graph.add_dependency(migration, key, parent, skip_validation=True)
205 def add_external_dependencies(self, key, migration):
206 for parent in migration.dependencies:
207 # Skip internal dependencies
208 if key[0] == parent[0]:
209 continue
210 parent = self.check_key(parent, key[0])
211 if parent is not None:
212 self.graph.add_dependency(migration, key, parent, skip_validation=True)
213 for child in migration.run_before:
214 child = self.check_key(child, key[0])
215 if child is not None:
216 self.graph.add_dependency(migration, child, key, skip_validation=True)
218 def build_graph(self):
219 """
220 Build a migration dependency graph using both the disk and database.
221 You'll need to rebuild the graph if you apply migrations. This isn't
222 usually a problem as generally migration stuff runs in a one-shot process.
223 """
224 # Load disk data
225 self.load_disk()
226 # Load database data
227 if self.connection is None:
228 self.applied_migrations = {}
229 else:
230 recorder = MigrationRecorder(self.connection)
231 self.applied_migrations = recorder.applied_migrations()
232 # To start, populate the migration graph with nodes for ALL migrations
233 # and their dependencies. Also make note of replacing migrations at this step.
234 self.graph = MigrationGraph()
235 self.replacements = {}
236 for key, migration in self.disk_migrations.items():
237 self.graph.add_node(key, migration)
238 # Replacing migrations.
239 if migration.replaces:
240 self.replacements[key] = migration
241 for key, migration in self.disk_migrations.items():
242 # Internal (same app) dependencies.
243 self.add_internal_dependencies(key, migration)
244 # Add external dependencies now that the internal ones have been resolved.
245 for key, migration in self.disk_migrations.items():
246 self.add_external_dependencies(key, migration)
247 # Carry out replacements where possible and if enabled.
248 if self.replace_migrations:
249 for key, migration in self.replacements.items():
250 # Get applied status of each of this migration's replacement
251 # targets.
252 applied_statuses = [
253 (target in self.applied_migrations) for target in migration.replaces
254 ]
255 # The replacing migration is only marked as applied if all of
256 # its replacement targets are.
257 if all(applied_statuses):
258 self.applied_migrations[key] = migration
259 else:
260 self.applied_migrations.pop(key, None)
261 # A replacing migration can be used if either all or none of
262 # its replacement targets have been applied.
263 if all(applied_statuses) or (not any(applied_statuses)):
264 self.graph.remove_replaced_nodes(key, migration.replaces)
265 else:
266 # This replacing migration cannot be used because it is
267 # partially applied. Remove it from the graph and remap
268 # dependencies to it (#25945).
269 self.graph.remove_replacement_node(key, migration.replaces)
270 # Ensure the graph is consistent.
271 try:
272 self.graph.validate_consistency()
273 except NodeNotFoundError as exc:
274 # Check if the missing node could have been replaced by any squash
275 # migration but wasn't because the squash migration was partially
276 # applied before. In that case raise a more understandable exception
277 # (#23556).
278 # Get reverse replacements.
279 reverse_replacements = {}
280 for key, migration in self.replacements.items():
281 for replaced in migration.replaces:
282 reverse_replacements.setdefault(replaced, set()).add(key)
283 # Try to reraise exception with more detail.
284 if exc.node in reverse_replacements:
285 candidates = reverse_replacements.get(exc.node, set())
286 is_replaced = any(
287 candidate in self.graph.nodes for candidate in candidates
288 )
289 if not is_replaced:
290 tries = ", ".join("{}.{}".format(*c) for c in candidates)
291 raise NodeNotFoundError(
292 f"Migration {exc.origin} depends on nonexistent node ('{exc.node[0]}', '{exc.node[1]}'). "
293 f"Plain tried to replace migration {exc.node[0]}.{exc.node[1]} with any of [{tries}] "
294 "but wasn't able to because some of the replaced migrations "
295 "are already applied.",
296 exc.node,
297 ) from exc
298 raise
299 self.graph.ensure_not_cyclic()
301 def check_consistent_history(self, connection):
302 """
303 Raise InconsistentMigrationHistory if any applied migrations have
304 unapplied dependencies.
305 """
306 recorder = MigrationRecorder(connection)
307 applied = recorder.applied_migrations()
308 for migration in applied:
309 # If the migration is unknown, skip it.
310 if migration not in self.graph.nodes:
311 continue
312 for parent in self.graph.node_map[migration].parents:
313 if parent not in applied:
314 # Skip unapplied squashed migrations that have all of their
315 # `replaces` applied.
316 if parent in self.replacements:
317 if all(
318 m in applied for m in self.replacements[parent].replaces
319 ):
320 continue
321 raise InconsistentMigrationHistory(
322 f"Migration {migration[0]}.{migration[1]} is applied before its dependency "
323 f"{parent[0]}.{parent[1]} on database '{connection.alias}'."
324 )
326 def detect_conflicts(self):
327 """
328 Look through the loaded graph and detect any conflicts - packages
329 with more than one leaf migration. Return a dict of the app labels
330 that conflict with the migration names that conflict.
331 """
332 seen_packages = {}
333 conflicting_packages = set()
334 for package_label, migration_name in self.graph.leaf_nodes():
335 if package_label in seen_packages:
336 conflicting_packages.add(package_label)
337 seen_packages.setdefault(package_label, set()).add(migration_name)
338 return {
339 package_label: sorted(seen_packages[package_label])
340 for package_label in conflicting_packages
341 }
343 def project_state(self, nodes=None, at_end=True):
344 """
345 Return a ProjectState object representing the most recent state
346 that the loaded migrations represent.
348 See graph.make_state() for the meaning of "nodes" and "at_end".
349 """
350 return self.graph.make_state(
351 nodes=nodes, at_end=at_end, real_packages=self.unmigrated_packages
352 )
354 def collect_sql(self, plan):
355 """
356 Take a migration plan and return a list of collected SQL statements
357 that represent the best-efforts version of that plan.
358 """
359 statements = []
360 state = None
361 for migration, backwards in plan:
362 with self.connection.schema_editor(
363 collect_sql=True, atomic=migration.atomic
364 ) as schema_editor:
365 if state is None:
366 state = self.project_state(
367 (migration.package_label, migration.name), at_end=False
368 )
369 if not backwards:
370 state = migration.apply(state, schema_editor, collect_sql=True)
371 else:
372 state = migration.unapply(state, schema_editor, collect_sql=True)
373 statements.extend(schema_editor.collected_sql)
374 return statements