Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-models/plain/models/migrations/loader.py: 39%

180 statements  

« 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 

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 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 ) 

139 

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] 

143 

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]] 

166 

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]}") 

194 

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) 

204 

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) 

217 

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() 

300 

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 ) 

325 

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 } 

342 

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. 

347 

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 ) 

353 

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