Coverage for /Users/davegaeddert/Development/dropseed/plain/plain-models/plain/models/migrations/loader.py: 53%

180 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-16 22:04 -0500

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