sentinel/bin/generate-model-dependency-fixtures

#!/usr/bin/env python
from __future__ import annotations

from sentry.runner import configure

configure()

import click

from sentry.backup.dependencies import (
    DependenciesJSONEncoder,
    dependencies,
    get_model_name,
    sorted_dependencies,
)
from sentry.testutils.factories import get_fixture_path  # noqa

encoder = DependenciesJSONEncoder(
    sort_keys=True,
    ensure_ascii=True,
    check_circular=True,
    allow_nan=True,
    indent=2,
    encoding="utf-8",
)


@click.command()
def main():
    """Used for generating fixtures for the model dependency map."""

    click.echo("\nThis script can take up to 30 seconds, please be patient...\n")

    # Do all of the calculation before any file writes, so that we don't end up with partial
    # overwrites.
    detailed = {str(k): v for k, v in dependencies().items()}
    flat = {k: v.flatten() for k, v in detailed.items()}
    sorted = sorted_dependencies()
    truncate = [dependencies()[get_model_name(m)].table_name for m in sorted_dependencies()]

    det_path = get_fixture_path("backup", "model_dependencies", "detailed.json")
    with open(det_path, "w+") as fixture:
        fixture.write(encoder.encode(detailed))

    flat_path = get_fixture_path("backup", "model_dependencies", "flat.json")
    with open(flat_path, "w+") as fixture:
        fixture.write(encoder.encode(flat))

    det_path = get_fixture_path("backup", "model_dependencies", "sorted.json")
    with open(det_path, "w+") as fixture:
        fixture.write(encoder.encode(sorted))

    # Print tables in the order one would need to list them for a TRUNCATE statement.
    det_path = get_fixture_path("backup", "model_dependencies", "truncate.json")
    with open(det_path, "w+") as fixture:
        fixture.write(encoder.encode(truncate))

    click.echo(
        f"\nSuccess! The dependency mapping fixtures at {[det_path, flat_path]} were updated.\n"
    )


if __name__ == "__main__":
    main()