superset2/scripts/seed_junction_load.py

#!/usr/bin/env python3
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
# ----------------------------------------------------------------------
# Stress-test data generator for the composite-PK migration (sc-105349).
#
# Bulk-inserts synthetic parent rows and many-to-many junction rows for
# the eight association tables that the composite-PK migration touches.
# Useful for measuring migration runtime at varying scales — run this at
# 100K / 1M / 5M / 10M rows and time the migration at each scale to
# verify the O(N log N) extrapolation.
#
# Idempotent: rerunning with the same target is a no-op; rerunning with
# a higher target adds rows up to the new total. Batched bulk INSERTs
# (10K rows per statement) make it fast on Postgres, MySQL, and SQLite.
#
# Usage (inside the Superset container):
#
#     docker exec superset-superset-1 \\
#         /app/.venv/bin/python /app/scripts/seed_junction_load.py \\
#         --dashboard-slices 1000000 \\
#         --slice-user 100000 \\
#         --dashboard-user 100000
#
# Run with no flags for the defaults shown below. Use ``--dry-run`` to
# print the planned inserts without writing anything.
#
# The script connects via Superset's standard ``DATABASE_*`` env vars
# (or ``SUPERSET__SQLALCHEMY_DATABASE_URI`` if set), so it works
# automatically inside the Superset container regardless of which
# metadata DB backend is in use.

from __future__ import annotations

import argparse
import logging
import os
import sys
import time
from contextlib import contextmanager
from typing import Iterator
from uuid import uuid4

import sqlalchemy as sa
from sqlalchemy.engine import Connection, Engine

logger = logging.getLogger("seed_junction_load")

# Bulk INSERT batch size. Larger values = fewer statements but more memory.
BATCH = 10_000

# Default per-junction-table target row counts. Tuned to mimic the shape
# of a large multi-team Superset install. Override via CLI flags.
DEFAULTS: dict[str, int] = {
    "dashboard_slices": 1_000_000,
    "slice_user": 100_000,
    "dashboard_user": 100_000,
    "dashboard_roles": 10_000,
}

# (junction_table, fk1_col, fk2_col, parent1_table, parent2_table)
# parents reference id columns; we generate (fk1, fk2) pairs by sampling
# from the parents' existing IDs.
JUNCTIONS: list[tuple[str, str, str, str, str]] = [
    ("dashboard_slices", "dashboard_id", "slice_id", "dashboards", "slices"),
    ("slice_user", "user_id", "slice_id", "ab_user", "slices"),
    ("dashboard_user", "user_id", "dashboard_id", "ab_user", "dashboards"),
    ("dashboard_roles", "dashboard_id", "role_id", "dashboards", "ab_role"),
]


# ----------------------------------------------------------------------
# Connection setup
# ----------------------------------------------------------------------


def build_engine() -> Engine:
    """Build a SQLAlchemy engine from Superset env vars."""
    if uri := os.environ.get("SUPERSET__SQLALCHEMY_DATABASE_URI"):
        logger.info("Using SUPERSET__SQLALCHEMY_DATABASE_URI from env")
        return sa.create_engine(uri)

    try:
        dialect = os.environ["DATABASE_DIALECT"]
        user = os.environ["DATABASE_USER"]
        password = os.environ["DATABASE_PASSWORD"]
        host = os.environ["DATABASE_HOST"]
        port = os.environ["DATABASE_PORT"]
        db = os.environ["DATABASE_DB"]
    except KeyError as exc:
        sys.exit(
            f"Missing env var {exc}; either set DATABASE_DIALECT/USER/PASSWORD/"
            f"HOST/PORT/DB or SUPERSET__SQLALCHEMY_DATABASE_URI before running."
        )

    uri = f"{dialect}://{user}:{password}@{host}:{port}/{db}"
    logger.info(
        "Built URI from DATABASE_* env vars (dialect=%s, host=%s)", dialect, host
    )
    return sa.create_engine(uri)


# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------


def uuid_value(dialect_name: str) -> bytes | str:
    """Return a UUID in the form the active dialect expects.

    MySQL stores UUIDs as ``BINARY(16)`` (16 raw bytes); Postgres has a
    native ``UUID`` type that accepts strings; SQLite stores them as
    BLOB/TEXT and accepts either. Branching here keeps the seed script
    backend-agnostic without depending on Superset's custom column types.
    """
    if dialect_name.startswith("mysql"):
        return uuid4().bytes
    return str(uuid4())


@contextmanager
def time_phase(name: str) -> Iterator[None]:
    """Log elapsed wall time for a named phase."""
    start = time.monotonic()
    logger.info("[%s] starting", name)
    try:
        yield
    finally:
        elapsed = time.monotonic() - start
        logger.info("[%s] done in %.2fs", name, elapsed)


def count_rows(conn: Connection, table: str) -> int:
    return conn.scalar(sa.text(f"SELECT COUNT(*) FROM {table}")) or 0  # noqa: S608


def existing_ids(conn: Connection, table: str, limit: int | None = None) -> list[int]:
    sql = f"SELECT id FROM {table} ORDER BY id"  # noqa: S608
    if limit is not None:
        sql += f" LIMIT {limit}"
    return [row[0] for row in conn.execute(sa.text(sql))]


# ----------------------------------------------------------------------
# Parent seeders
#
# Each function ensures the named parent table has at least ``target``
# rows by inserting synthetic ones with minimal-but-valid columns.
# Returns nothing; subsequent code reads back IDs via ``existing_ids``.
# ----------------------------------------------------------------------


def seed_dashboards(conn: Connection, target: int, dry_run: bool) -> None:
    current = count_rows(conn, "dashboards")
    if current >= target:
        logger.info(
            "dashboards: %d rows (target %d) — no insert needed", current, target
        )
        return
    needed = target - current
    logger.info("dashboards: %d → %d (+%d)", current, target, needed)
    if dry_run:
        return

    dialect = conn.engine.dialect.name
    sql = sa.text(
        "INSERT INTO dashboards (uuid, dashboard_title, slug, published) "
        "VALUES (:uuid, :title, :slug, :published)"
    )
    for batch_start in range(0, needed, BATCH):
        rows = [
            {
                "uuid": uuid_value(dialect),
                "title": f"seed_dashboard_{current + i}",
                "slug": f"seed-dashboard-{current + i}-{uuid4().hex[:8]}",
                "published": False,
            }
            for i in range(batch_start, min(batch_start + BATCH, needed))
        ]
        conn.execute(sql, rows)
        logger.info("  dashboards: inserted %d / %d", batch_start + len(rows), needed)


def seed_dbs(conn: Connection, dry_run: bool) -> int:
    """Ensure at least one row exists in ``dbs`` (parent of ``tables``).
    Returns the id to use as ``database_id`` when seeding ``tables``."""
    ids = existing_ids(conn, "dbs", limit=1)
    if ids:
        return ids[0]
    if dry_run:
        return -1  # placeholder
    dialect = conn.engine.dialect.name
    logger.info("dbs: inserting one synthetic database (no rows present)")
    conn.execute(
        sa.text(
            "INSERT INTO dbs (uuid, database_name, sqlalchemy_uri, expose_in_sqllab) "
            "VALUES (:uuid, :name, :uri, :expose)"
        ),
        {
            "uuid": uuid_value(dialect),
            "name": f"seed_db_{uuid4().hex[:8]}",
            "uri": "sqlite:///seed.db",
            "expose": False,
        },
    )
    return existing_ids(conn, "dbs", limit=1)[0]


def seed_tables(conn: Connection, target: int, dry_run: bool) -> None:
    current = count_rows(conn, "tables")
    if current >= target:
        logger.info("tables: %d rows (target %d) — no insert needed", current, target)
        return
    needed = target - current
    logger.info("tables: %d → %d (+%d)", current, target, needed)
    if dry_run:
        return

    database_id = seed_dbs(conn, dry_run=False)
    dialect = conn.engine.dialect.name
    sql = sa.text(
        "INSERT INTO tables (uuid, table_name, database_id) "
        "VALUES (:uuid, :name, :db_id)"
    )
    for batch_start in range(0, needed, BATCH):
        rows = [
            {
                "uuid": uuid_value(dialect),
                "name": f"seed_table_{current + i}",
                "db_id": database_id,
            }
            for i in range(batch_start, min(batch_start + BATCH, needed))
        ]
        conn.execute(sql, rows)
        logger.info("  tables: inserted %d / %d", batch_start + len(rows), needed)


def seed_slices(conn: Connection, target: int, dry_run: bool) -> None:
    current = count_rows(conn, "slices")
    if current >= target:
        logger.info("slices: %d rows (target %d) — no insert needed", current, target)
        return
    needed = target - current
    logger.info("slices: %d → %d (+%d)", current, target, needed)
    if dry_run:
        return

    # Slices reference tables.id; ensure at least one ``tables`` row exists
    # so the FK is satisfiable (datasource_id is nullable but we set it for
    # realism). The migration test doesn't care, but a real Superset that
    # re-renders these slices does.
    seed_tables(conn, target=1, dry_run=False)
    table_id = existing_ids(conn, "tables", limit=1)[0]
    dialect = conn.engine.dialect.name
    sql = sa.text(
        "INSERT INTO slices "
        "(uuid, slice_name, datasource_id, datasource_type, viz_type) "
        "VALUES (:uuid, :name, :ds_id, :ds_type, :viz)"
    )
    for batch_start in range(0, needed, BATCH):
        rows = [
            {
                "uuid": uuid_value(dialect),
                "name": f"seed_slice_{current + i}",
                "ds_id": table_id,
                "ds_type": "table",
                "viz": "table",
            }
            for i in range(batch_start, min(batch_start + BATCH, needed))
        ]
        conn.execute(sql, rows)
        logger.info("  slices: inserted %d / %d", batch_start + len(rows), needed)


def seed_users(conn: Connection, target: int, dry_run: bool) -> None:
    current = count_rows(conn, "ab_user")
    if current >= target:
        logger.info("ab_user: %d rows (target %d) — no insert needed", current, target)
        return
    needed = target - current
    logger.info("ab_user: %d → %d (+%d)", current, target, needed)
    if dry_run:
        return

    sql = sa.text(
        "INSERT INTO ab_user (first_name, last_name, username, email, active) "
        "VALUES (:first, :last, :username, :email, :active)"
    )
    for batch_start in range(0, needed, BATCH):
        rows = [
            {
                "first": "seed",
                "last": f"user_{current + i}",
                "username": f"seed_user_{current + i}_{uuid4().hex[:8]}",
                "email": f"seed_user_{current + i}_{uuid4().hex[:8]}@example.invalid",
                "active": True,
            }
            for i in range(batch_start, min(batch_start + BATCH, needed))
        ]
        conn.execute(sql, rows)
        logger.info("  ab_user: inserted %d / %d", batch_start + len(rows), needed)


def seed_roles(conn: Connection, target: int, dry_run: bool) -> None:
    current = count_rows(conn, "ab_role")
    if current >= target:
        logger.info("ab_role: %d rows (target %d) — no insert needed", current, target)
        return
    needed = target - current
    logger.info("ab_role: %d → %d (+%d)", current, target, needed)
    if dry_run:
        return

    sql = sa.text("INSERT INTO ab_role (name) VALUES (:name)")
    for batch_start in range(0, needed, BATCH):
        rows = [
            {"name": f"seed_role_{current + i}_{uuid4().hex[:8]}"}
            for i in range(batch_start, min(batch_start + BATCH, needed))
        ]
        conn.execute(sql, rows)
        logger.info("  ab_role: inserted %d / %d", batch_start + len(rows), needed)


# ----------------------------------------------------------------------
# Junction seeder
# ----------------------------------------------------------------------


def _load_existing_pairs(
    conn: Connection, junction: str, fk1_col: str, fk2_col: str
) -> set[tuple[int, int]]:
    """Load existing ``(fk1, fk2)`` pairs from a junction table into a set.

    Used so the seeder can skip them when generating new pairs (junction
    tables enforce uniqueness on the FK pair). Memory is ~32 bytes/tuple
    on CPython, so 10M existing pairs is ~320MB — acceptable for a dev
    machine. The junction / column names come from ``JUNCTIONS``, not
    user input, so the f-string interpolation is safe.
    """
    sql_text = f"SELECT {fk1_col}, {fk2_col} FROM {junction}"  # noqa: S608
    return {(row[0], row[1]) for row in conn.execute(sa.text(sql_text))}


def _generate_new_pairs(
    p1_ids: list[int],
    p2_ids: list[int],
    existing_pairs: set[tuple[int, int]],
) -> Iterator[tuple[int, int]]:
    """Yield ``(fk1, fk2)`` pairs from the parent1 × parent2 cross-product
    that are not already in ``existing_pairs``."""
    for fk1 in p1_ids:
        for fk2 in p2_ids:
            if (fk1, fk2) not in existing_pairs:
                yield (fk1, fk2)


def seed_junction(
    conn: Connection,
    junction: str,
    fk1_col: str,
    fk2_col: str,
    parent1: str,
    parent2: str,
    target: int,
    dry_run: bool,
) -> None:
    """Bulk-insert junction rows up to ``target`` rows total.

    Generates ``(fk1, fk2)`` pairs by walking the cross-product of
    parent1 IDs × parent2 IDs in row-major order, skipping pairs that
    already exist. Walking the cross-product deterministically keeps
    the script replayable: re-running with the same target is a no-op,
    and re-running with a higher target appends new pairs in a stable
    order regardless of how many runs preceded.
    """
    current = count_rows(conn, junction)
    if current >= target:
        logger.info(
            "%s: %d rows (target %d) — no insert needed", junction, current, target
        )
        return
    needed = target - current
    logger.info("%s: %d → %d (+%d)", junction, current, target, needed)
    if dry_run:
        return

    p1_ids = existing_ids(conn, parent1)
    p2_ids = existing_ids(conn, parent2)
    max_pairs = len(p1_ids) * len(p2_ids)
    if max_pairs < target:
        sys.exit(
            f"Cannot reach {target} rows in {junction}: "
            f"only {max_pairs} unique pairs available "
            f"({len(p1_ids)} × {len(p2_ids)}). "
            f"Increase parent targets and rerun."
        )

    existing_pairs: set[tuple[int, int]] = (
        _load_existing_pairs(conn, junction, fk1_col, fk2_col) if current > 0 else set()
    )
    if existing_pairs:
        logger.info(
            "  %s: loaded %d existing pairs into avoidance set",
            junction,
            len(existing_pairs),
        )

    insert_sql = sa.text(
        f"INSERT INTO {junction} ({fk1_col}, {fk2_col}) "  # noqa: S608
        f"VALUES (:fk1, :fk2)"
    )

    inserted = 0
    batch: list[dict[str, int]] = []
    for fk1, fk2 in _generate_new_pairs(p1_ids, p2_ids, existing_pairs):
        batch.append({"fk1": fk1, "fk2": fk2})
        inserted += 1
        if len(batch) == BATCH or inserted == needed:
            conn.execute(insert_sql, batch)
            logger.info("  %s: inserted %d / %d", junction, inserted, needed)
            batch = []
            if inserted == needed:
                return
    if inserted < needed:
        sys.exit(
            f"Ran out of unique pairs at {inserted}/{needed} for {junction} "
            f"(parents have {len(p1_ids)} × {len(p2_ids)} = {max_pairs} pairs, "
            f"{len(existing_pairs)} already present)"
        )


# ----------------------------------------------------------------------
# Orchestration
# ----------------------------------------------------------------------


def required_parent_count(target_pairs: int, other_parent: int) -> int:
    """How many rows we need in this parent so that
    (this_parent × other_parent) ≥ target_pairs."""
    if other_parent == 0:
        # Bootstrapping: assume we'll create at least 1
        other_parent = 1
    return -(-target_pairs // other_parent)  # ceil(target_pairs / other_parent)


def _compute_parent_requirements(targets: dict[str, int]) -> dict[str, int]:
    """For each parent table, return the minimum row count needed so that
    parent1 × parent2 ≥ target for every junction it participates in.

    Allocates ceil(sqrt(target)) rows per parent, balanced across the two
    parents of each junction. The actual junction seeder will then walk
    the cross-product to produce the target number of unique pairs.
    """
    parent_req: dict[str, int] = {}
    for junction, _, _, p1, p2 in JUNCTIONS:
        target = targets.get(junction, 0)
        if target == 0:
            continue
        sqrt_n = int(target**0.5) + 1
        parent_req[p1] = max(parent_req.get(p1, 0), sqrt_n)
        parent_req[p2] = max(parent_req.get(p2, 0), sqrt_n)
    return parent_req


def _seed_parents(conn: Connection, parent_req: dict[str, int], dry_run: bool) -> None:
    """Seed parent tables in dependency order:
    independent parents (ab_user, ab_role) first, then dashboards / slices /
    tables (which transitively depend on dbs, seeded inside seed_tables)."""
    if "ab_user" in parent_req:
        seed_users(conn, parent_req["ab_user"], dry_run)
    if "ab_role" in parent_req:
        seed_roles(conn, parent_req["ab_role"], dry_run)
    if "dashboards" in parent_req:
        seed_dashboards(conn, parent_req["dashboards"], dry_run)
    if "slices" in parent_req:
        seed_slices(conn, parent_req["slices"], dry_run)
    if "tables" in parent_req:
        seed_tables(conn, parent_req["tables"], dry_run)


def _seed_all_junctions(
    conn: Connection, targets: dict[str, int], dry_run: bool
) -> None:
    for junction, fk1, fk2, p1, p2 in JUNCTIONS:
        target = targets.get(junction, 0)
        if target == 0:
            continue
        with time_phase(f"junction:{junction}"):
            seed_junction(conn, junction, fk1, fk2, p1, p2, target, dry_run)


def run(targets: dict[str, int], dry_run: bool) -> None:
    engine = build_engine()
    with engine.begin() as conn:
        parent_req = _compute_parent_requirements(targets)
        logger.info("Required parent row counts: %s", parent_req)

        with time_phase("parents"):
            _seed_parents(conn, parent_req, dry_run)

        with time_phase("junctions"):
            _seed_all_junctions(conn, targets, dry_run)


# ----------------------------------------------------------------------
# CLI
# ----------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    for table, default in DEFAULTS.items():
        parser.add_argument(
            f"--{table.replace('_', '-')}",
            type=int,
            default=default,
            help=f"target row count for {table} (default: {default:,})",
        )
    parser.add_argument(
        "--dry-run",
        "-n",
        action="store_true",
        help="print planned inserts without writing to the DB",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="increase log verbosity",
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%H:%M:%S",
    )

    targets = {table: getattr(args, table) for table in DEFAULTS}

    logger.info("Targets: %s", targets)
    logger.info("Dry run: %s", args.dry_run)

    with time_phase("total"):
        run(targets, dry_run=args.dry_run)


if __name__ == "__main__":
    main()