"""Management command to visualize migrations and dependencies for apps in the project."""
import re
from django.apps import apps as django_apps
from django.core.management.base import BaseCommand, CommandError
from django.db.migrations.loader import MigrationLoader
[docs]
class Command(BaseCommand):
"""A management command to Visualize migrations and dependencies for applications in the project."""
help = "Visualize migrations and dependencies for apps in the project, limiting to the specified apps, if any."
[docs]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.nodes = set()
self.edges = set()
self.app_labels = set()
self.loader = None
self.reverse_dependencies = {}
[docs]
def add_arguments(self, parser):
"""Add optional argument to specify apps to show migrations for."""
parser.add_argument(
"app_labels",
nargs="*",
help="Optional list of application labels.",
)
[docs]
def handle(self, *args, **options):
"""Handle the command."""
app_labels = options["app_labels"]
self.loader = MigrationLoader(None, ignore_no_migrations=True)
if not app_labels:
app_labels = [app_config.label for app_config in django_apps.get_app_configs()]
else:
# Validate that provided app labels are installed
for app_label in app_labels:
try:
django_apps.get_app_config(app_label)
except LookupError as e:
raise CommandError(f"{e}. Are you sure your INSTALLED_APPS setting is correct?") from e
self.app_labels = set(app_labels)
self._build_reverse_dependencies()
for idx, app in enumerate(app_labels):
self._print_success(f"[{app}]")
self._print_app_migrationgraph(app)
if idx != len(app_labels) - 1:
self.stdout.write("\n")
# Collect nodes from edges and print MermaidJS flowchart
self._collect_nodes_from_edges()
self._print_mermaidjs_flowchart()
def _build_reverse_dependencies(self):
"""Build a mapping of reverse dependencies for all migrations."""
self.reverse_dependencies = {}
for (app_label, migration_name), node in self.loader.graph.nodes.items():
for dep_app, dep_name in node.dependencies:
# Resolve __first__ dependencies
if dep_name == "__first__":
root_nodes = self.loader.graph.root_nodes(dep_app)
for root_node in root_nodes:
self.reverse_dependencies.setdefault(root_node, set()).add((app_label, migration_name))
else:
dep_key = (dep_app, dep_name)
self.reverse_dependencies.setdefault(dep_key, set()).add((app_label, migration_name))
def _print_migration_info(self, app_label, migration_name, node):
"""Print detailed information about a migration."""
self._print_label(f"{app_label}/{migration_name}")
# Print dependencies
dependencies = [(dep_app, dep_name) for dep_app, dep_name in node.dependencies if dep_app != app_label]
if dependencies:
self._print_title("\tDepends on:")
for dep_app, dep_name in dependencies:
# Resolve __first__ to actual migration name
if dep_name == "__first__":
root_nodes = self.loader.graph.root_nodes(dep_app)
for root_app, root_name in root_nodes:
self._print_notice(f"\t\t{root_app}/{root_name}")
else:
self._print_notice(f"\t\t{dep_app}/{dep_name}")
# Print reverse dependencies
migration_key = (app_label, migration_name)
reverse_deps = self.reverse_dependencies.get(migration_key, set())
reverse_deps = [
(rev_app, rev_name) for rev_app, rev_name in reverse_deps if rev_app in self.app_labels
] # Filter to requested apps
if reverse_deps:
self._print_title("\tDepended upon by:")
for rev_app, rev_name in sorted(reverse_deps):
self._print_notice(f"\t\t{rev_app}/{rev_name}")
def _get_node_key(self, node):
"""Get the app_label and migration_name for a node."""
# Handle both Node objects and tuples
if hasattr(node, "key"):
return node.key
return node
def _collect_nodes_from_edges(self):
"""Ensure all nodes in edges are included in nodes."""
for edge in self.edges:
from_node_key, to_node_key = edge
from_node = self.loader.graph.node_map.get(from_node_key)
to_node = self.loader.graph.node_map.get(to_node_key)
if from_node:
self.nodes.add(from_node)
if to_node:
self.nodes.add(to_node)
def _print_styled(self, style, text):
"""Print the text with the specified style."""
self.stdout.write(style(text) + "\n")
def _print_label(self, text):
"""Print the label."""
self._print_styled(self.style.MIGRATE_LABEL, text)
def _print_warn(self, text):
"""Print the warning."""
self._print_styled(self.style.WARNING, text)
def _print_title(self, text):
"""Print the title."""
self._print_styled(self.style.MIGRATE_HEADING, text)
def _print_notice(self, text):
"""Print the notice."""
self._print_styled(self.style.NOTICE, text)
def _print_error(self, text):
"""Print the error."""
self._print_styled(self.style.ERROR, text)
def _print_success(self, text):
"""Print the success."""
try:
style = self.style.SUCCESS
except AttributeError:
style = self.style.MIGRATE_SUCCESS
self._print_styled(style, text)
def _get_node_id(self, node_key):
"""Get the node ID."""
app_label, migration_name = node_key
node_id = f"{app_label}_{migration_name}"
# Replace any non-alphanumeric characters with underscores
node_id = re.sub(r"\W|^(?=\d)", "_", node_id)
return node_id
def _get_node_label(self, node_key):
"""Get the node label."""
app_label, migration_name = node_key
match = re.match(r"^(\d+)_", migration_name)
if match:
prefix = match.group(1)
else:
prefix = migration_name[:9]
label = f"{app_label}/{prefix}"
return label
def _get_node_sort_key(self, node_key):
"""Get the sort key for a node to ensure deterministic order.
Args:
node_key: The node key tuple (app_label, migration_name).
Returns:
str: The node ID used for sorting.
"""
return self._get_node_id(node_key)
def _get_edge_sort_key(self, edge):
"""Get the sort key for an edge to ensure deterministic order.
Args:
edge: The edge tuple (from_key, to_key).
Returns:
tuple: A tuple of (from_node_id, to_node_id) used for sorting.
"""
from_key, to_key = edge
return (self._get_node_id(from_key), self._get_node_id(to_key))
def _print_app_migrationgraph(self, app):
"""Print the migrations graph for the specified app."""
try:
migrations = []
# Collect all migrations for this app
for (app_label, migration_name), node in self.loader.graph.nodes.items():
if app_label == app:
migrations.append((migration_name, node))
# Add this node to our nodes set
self.nodes.add((app_label, migration_name))
# Collect edges from this node's dependencies
for dep_app, dep_name in node.dependencies:
if dep_name == "__first__":
root_nodes = self.loader.graph.root_nodes(dep_app)
for root_node in root_nodes:
self.edges.add((root_node, (app_label, migration_name)))
else:
self.edges.add(((dep_app, dep_name), (app_label, migration_name)))
# Sort migrations by name to ensure consistent ordering
migrations.sort()
for migration_name, node in migrations:
self._print_migration_info(app, migration_name, node)
except KeyError:
self._print_error(f"Migrations for `{app}` application were not found")
return
def _print_mermaidjs_flowchart(self):
"""Print the MermaidJS flowchart."""
self.stdout.write("\n_____________________")
self.stdout.write("Migrations Flowchart:\n")
self.stdout.write("\n```mermaid\n")
self.stdout.write("graph TD\n")
# First, define nodes
# Sort nodes to ensure deterministic output
for node_key in sorted(self.nodes, key=self._get_node_sort_key):
node_id = self._get_node_id(node_key)
node_label = self._get_node_label(node_key)
self.stdout.write(f' {node_id}["{node_label}"]\n')
# Then, define edges
# Sort edges to ensure deterministic output
for from_key, to_key in sorted(self.edges, key=self._get_edge_sort_key):
from_id = self._get_node_id(from_key)
to_id = self._get_node_id(to_key)
self.stdout.write(f" {from_id} --> {to_id}\n")
self.stdout.write("```\n")