petrolyst-infra/petrolyst-insights/dashboards/atoma/import_bundle.py

"""Import the ATOMA dashboards bundle into Superset, overwriting what it defines.

Runs inside the Superset container, with the image's Python:

    docker cp import_bundle.py <container>:/tmp/ && docker cp atoma-dashboards.zip <container>:/tmp/
    docker exec <container> /app/.venv/bin/python /tmp/import_bundle.py --show-database
    docker exec <container> /app/.venv/bin/python /tmp/import_bundle.py /tmp/atoma-dashboards.zip --user <username>

`superset import-dashboards` imports a dashboard's datasets and charts with
overwrite off, so an existing chart or dataset keeps its old definition. This
script runs the dataset, chart and dashboard importers in turn, each with
overwrite on (the bundle's metadata.yaml type is set to what each importer
expects). A dataset import replaces the column list, so columns missing from
the bundle are dropped; afterwards each dataset's columns are compared with the
live view and any drift is reported (exit status 2).

The database is never created or changed: the bundle's database uuid must
already exist under the name "Fuel", otherwise nothing is imported.

The dashboards (atoma-overview, atoma-site) get no role: admins see them in Superset, and
the Petrolyst app shows them with guest passes (PD-12). With --embedded-out FILE
--embed-domain URL, each of them is made embeddable from that origin (its embedding id stays
the same across imports), and FILE receives {slug: embedding id} for the Petrolyst server:

    docker exec <container> /app/.venv/bin/python /tmp/import_bundle.py /tmp/atoma-dashboards.zip \
        --user <username> --embedded-out /tmp/embedded-dashboards.json --embed-domain http://127.0.0.1:8098
"""

from __future__ import annotations

import argparse
import json
import sys
import zipfile

import yaml

DATABASE_NAME = "Fuel"
EMBEDDED_PREFIX = "atoma-"
STEPS = [("datasets", "SqlaTable"), ("charts", "Slice"), ("dashboards", "Dashboard")]


def main() -> int:
    parser = argparse.ArgumentParser(description="Import the ATOMA dashboards bundle.")
    parser.add_argument("bundle", nargs="?")
    parser.add_argument("--user", help="existing Superset user who will own new objects")
    parser.add_argument("--show-database", action="store_true",
                        help=f'print the uuid of "{DATABASE_NAME}" and stop')
    parser.add_argument("--embedded-out", help="write {slug: embedding id} of the app's dashboards")
    parser.add_argument("--embed-domain", help="the app's origin, e.g. http://127.0.0.1:8098")
    args = parser.parse_args()
    if not args.show_database and not (args.bundle and args.user):
        parser.error("give a bundle and --user, or --show-database")
    if bool(args.embedded_out) != bool(args.embed_domain):
        parser.error("--embedded-out and --embed-domain go together")

    from superset.app import create_app

    app = create_app()
    with app.app_context():
        from superset import db, security_manager
        from superset.commands.chart.importers.v1 import ImportChartsCommand
        from superset.commands.dashboard.importers.v1 import ImportDashboardsCommand
        from superset.commands.dataset.importers.v1 import ImportDatasetsCommand
        from superset.commands.importers.v1.utils import get_contents_from_bundle
        from superset.connectors.sqla.models import SqlaTable
        from superset.models.core import Database
        from superset.models.dashboard import Dashboard
        from superset.utils.core import override_user

        if args.show_database:
            found = db.session.query(Database).filter_by(database_name=DATABASE_NAME).all()
            for database in found:
                print(database.uuid)
            return 0 if found else 1

        user = security_manager.find_user(username=args.user)
        if user is None:
            print(f"no Superset user {args.user!r}", file=sys.stderr)
            return 1
        with zipfile.ZipFile(args.bundle) as bundle:
            contents = get_contents_from_bundle(bundle)
        databases = [yaml.safe_load(t) for p, t in contents.items() if p.startswith("databases/")]
        for config in databases:
            database = db.session.query(Database).filter_by(uuid=config["uuid"]).one_or_none()
            if database is None or database.database_name != DATABASE_NAME:
                print(f'database {config["uuid"]} is not "{DATABASE_NAME}" here; '
                      "rebuild the bundle with the uuid from --show-database", file=sys.stderr)
                return 1

        metadata = yaml.safe_load(contents["metadata.yaml"])
        commands = {"datasets": ImportDatasetsCommand, "charts": ImportChartsCommand,
                    "dashboards": ImportDashboardsCommand}
        with override_user(user):
            for step, model in STEPS:
                step_contents = dict(contents)
                step_contents["metadata.yaml"] = yaml.safe_dump({**metadata, "type": model})
                commands[step](step_contents, overwrite=True).run()
                count = sum(p.startswith(step + "/") for p in contents)
                print(f"{step}: {count} imported (overwrite)")

        embedded: dict[str, str] = {}
        for path, text in contents.items():
            if path.startswith("dashboards/"):
                dashboard = db.session.query(Dashboard).filter_by(
                    uuid=yaml.safe_load(text)["uuid"]).one()
                if dashboard.slug.startswith(EMBEDDED_PREFIX) and args.embedded_out:
                    from superset.daos.dashboard import EmbeddedDashboardDAO

                    config = EmbeddedDashboardDAO.upsert(dashboard, [args.embed_domain])
                    db.session.flush()
                    embedded[dashboard.slug] = str(config.uuid)
        db.session.commit()
        if args.embedded_out:
            with open(args.embedded_out, "w") as out:
                json.dump(dict(sorted(embedded.items())), out, indent=2)
            print(f"embedded: {len(embedded)} dashboard(s) for {args.embed_domain}, "
                  f"ids in {args.embedded_out}")

        drift = 0
        dataset_uuids = [yaml.safe_load(t)["uuid"] for p, t in contents.items()
                         if p.startswith("datasets/")]
        for dataset_uuid in dataset_uuids:
            dataset = db.session.query(SqlaTable).filter_by(uuid=dataset_uuid).one()
            physical = {c.column_name for c in dataset.columns if not c.expression}
            live = {c["column_name"] for c in dataset.external_metadata()}
            if physical != live:
                drift += 1
                print(f"drift {dataset.table_name}: only in bundle {sorted(physical - live)}, "
                      f"only in view {sorted(live - physical)}")
        print(f"column drift: {drift} dataset(s)")
        return 2 if drift else 0


if __name__ == "__main__":
    sys.exit(main())