Function main

misc/python/materialize/cli/mz_workload_capture.py:199–640

()

def main() -> int:
    parser = argparse.ArgumentParser(
        prog="mz-workload-capture",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="Records a workload profile from a running Materialize instance (without actual data)",
    )

    parser.add_argument("mz_url", type=str)
    # Default cluster to use
    parser.add_argument("--cluster", type=str, default="quickstart")
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        help="Path to write the workload.yml, - for stdout",
        default=f"workload_{datetime.now(timezone.utc).strftime('%Y-%m-%dT%H-%M-%S')}.yml",
    )
    parser.add_argument(
        "--time",
        type=int,
        help="How long of a query/data ingestion history to capture, in seconds",
        default=360,
    )
    # Currently too slow to always enable, not used yet.
    parser.add_argument("--avg-column-size", action=argparse.BooleanOptionalAction)
    parser.add_argument("-v", "--verbose", action=argparse.BooleanOptionalAction)

    args = parser.parse_args()

    global VERBOSE
    VERBOSE = args.verbose

    conn = psycopg.connect(args.mz_url)
    conn.autocommit = True
    with conn.cursor() as cur:
        cur.execute(SQL("SET CLUSTER = {}").format(Identifier(args.cluster)))

    workload = {
        "databases": {},
        "clusters": {},
        "queries": [],
        "mz_workload_version": "1.0.0",
    }

    with timed("Fetching clusters"):
        for cluster, managed in query(
            conn, "SELECT name, managed FROM mz_clusters WHERE id LIKE 'u%'"
        ):
            workload["clusters"][cluster] = {"managed": managed}
            if managed:
                workload["clusters"][cluster]["create_sql"] = query(
                    conn,
                    SQL("SELECT create_sql FROM (SHOW CREATE CLUSTER {})").format(
                        Identifier(cluster)
                    ),
                )[0][0]

    with timed("Fetching databases"):