#!/usr/bin/env python from __future__ import annotations from typing import Iterable from sentry.backup.scopes import RelocationScope from sentry.runner import configure from sentry.silo.base import SiloMode configure() from enum import Enum, unique from string import Template import click from django.db import models from sentry.backup.dependencies import ForeignField, ForeignFieldKind, ModelRelations, dependencies digraph = Template( """ digraph Models { ranksep = 8; rankdir=LR node [style="rounded,filled",shape="rectangle"]; subgraph cluster_legend { label = "Legend"; fontsize="40" node [shape="plaintext",style="none"] key1 [label=<
HybridCloudForeignKey
HybridCloudForeignKey (nullable)
Explicit ForeignKey
Explicit ForeignKey (nullable)
Implicit ForeignKey
Implicit ForeignKey (nullable)
Control Silo Model
Control Silo Model (dangling)
Region Silo Model
Region Silo Model (dangling)
Unexported Model
>] key2 [label=<
 
 
 
 
 
 
 
 
 
 
 
>] key1:i1:e -> key2:i1:w [color="#008b00ff",style=solid] key1:i2:e -> key2:i2:w [color="#008b0066",style=dashed] key1:i3:e -> key2:i3:w [color="#0000eeff",style=solid] key1:i4:e -> key2:i4:w [color="#0000ee66",style=dashed] key1:i5:e -> key2:i5:w [color="#cd0000ff",style=solid] key1:i6:e -> key2:i6:w [color="#cd000066",style=dashed] } $clusters $edges } """ ) cluster = Template( """ subgraph cluster_$num { label="$name Relocation Scope" style="rounded,filled" shape="rectangle" fillcolor="$fill" fontsize="40" color="#c0c0c0" $nodes } """ ) @unique class ClusterColor(Enum): Purple = "#fff0f5" # lavenderblush Yellow = "#f0e68c" # khaki Green = "#f0fff0" # honeydew Blue = "#cae1ff" # lightsteelblue1 @unique class NodeColor(Enum): Red = "#ffb6c1" # lightpink Blue = "#add8e6" # lightblue @unique class EdgeColor(Enum): Hybrid = "#008b00" # green4 Explicit = "#0000ee" # blue2 Implicit = "#cd0000" # red3 def print_model_node(mr: ModelRelations, silo: SiloMode) -> str: id = mr.model.__name__ color = NodeColor.Red if silo == SiloMode.CONTROL else NodeColor.Blue opacity = "66" if mr.dangling else "ff" return f""""{id}" [fillcolor="{color.value}{opacity}",color="#000000{opacity}"];""" def print_rel_scope_subgraph( name: str, num: int, rels: Iterable[ModelRelations], color: ClusterColor ) -> str: return cluster.substitute( num=num, name=name, fill=color.value, nodes="\n ".join([print_model_node(mr, mr.silos[0]) for mr in rels]), ) def print_edges(mr: ModelRelations) -> str: if len(mr.foreign_keys) == 0: return "" src = mr.model return "\n ".join([print_edge(src, ff.model, ff) for ff in mr.foreign_keys.values()]) def print_edge(src: models.base.ModelBase, dest: models.base.ModelBase, field: ForeignField) -> str: color = EdgeColor.Explicit if field.kind == ForeignFieldKind.HybridCloudForeignKey: color = EdgeColor.Hybrid elif field.kind == ForeignFieldKind.ImplicitForeignKey: color = EdgeColor.Implicit style = "dashed" if field.nullable else "solid" return f""""{src.__name__}":e -> "{dest.__name__}":w [color="{color.value}",style={style}];""" def get_most_permissive_relocation_scope(mr: ModelRelations) -> RelocationScope: if isinstance(mr.relocation_scope, set): return sorted(list(mr.relocation_scope), key=lambda obj: obj.value * -1)[0] return mr.relocation_scope @click.command() @click.option("--show-excluded", default=False, is_flag=True, help="Show unexportable models too") def main(show_excluded: bool): """Generate a graphviz spec for the current model dependency graph.""" # Get all dependencies, filtering as necessary. deps = sorted(dependencies().values(), key=lambda mr: mr.model.__name__) if not show_excluded: deps = list(filter(lambda m: m.relocation_scope != RelocationScope.Excluded, deps)) # Group by most permissive region scope. user_scoped = filter( lambda m: get_most_permissive_relocation_scope(m) == RelocationScope.User, deps ) org_scoped = filter( lambda m: get_most_permissive_relocation_scope(m) == RelocationScope.Organization, deps ) config_scoped = filter( lambda m: get_most_permissive_relocation_scope(m) == RelocationScope.Config, deps ) global_scoped = filter( lambda m: get_most_permissive_relocation_scope(m) == RelocationScope.Global, deps ) # Print nodes. clusters = "".join( [ print_rel_scope_subgraph("User", 1, user_scoped, ClusterColor.Green), print_rel_scope_subgraph("Organization", 2, org_scoped, ClusterColor.Purple), print_rel_scope_subgraph("Config", 3, config_scoped, ClusterColor.Blue), print_rel_scope_subgraph("Global", 4, global_scoped, ClusterColor.Yellow), ] ) # Print edges. edges = "\n ".join(filter(lambda s: s, [print_edges(mr) for mr in deps])) click.echo(digraph.substitute(clusters=clusters, edges=edges)) if __name__ == "__main__": main()