backfill_rds_snapshot_copy_tasks.py

"""Incrementally backfills COPY_FILES_TASK / COPY_CHANGES_TASK / COPY_SENT_MESSAGES_TASK's
RDS_SNAPSHOTS_STAGE backlog one snapshot_date at a time (SNOW-580).

These tasks time out trying to COPY INTO their entire backlog (150K+ files,
1TB+) in a single statement, even on COMPUTE_MEDIUM. This script drains the
backlog day by day instead: each COPY INTO is scoped to one snapshot_date via
the PATTERN clause, which keeps each run's file count/volume small enough to
finish well inside the warehouse statement timeout. It's safe to re-run --
COPY INTO skips files it has already loaded -- so a partial or interrupted run
can just be re-invoked to pick up where it left off.

Shells out to the `snow` CLI (see AGENTS.md) rather than holding credentials
directly.

Usage:
    uv run --group snowflake python analytics/backfill_rds_snapshot_copy_tasks.py \\
        --target files --warehouse compute_medium --max-dates 5
"""

import argparse
import json
import subprocess
import sys
import time

DATABASE = "SYNAPSE_DATA_WAREHOUSE"
SCHEMA = "RDS_LANDING"
STAGE = "RDS_SNAPSHOTS_STAGE"

TARGETS = {
    "files": {
        "table": "files",
        "stage_subpath": "rds-snapshot/FILES/",
        "select": """
            SELECT
                $1:ID::BIGINT                                 AS id,
                $1:ETAG::VARCHAR                              AS etag,
                $1:PREVIEW_ID::BIGINT                         AS preview_id,
                $1:CREATED_ON::TIMESTAMP_NTZ(9)               AS created_on,
                $1:CREATED_BY::BIGINT                         AS created_by,
                $1:UPDATED_ON::TIMESTAMP_NTZ(9)               AS updated_on,
                $1:METADATA_TYPE::VARCHAR                     AS metadata_type,
                $1:CONTENT_TYPE::VARCHAR                      AS content_type,
                $1:CONTENT_SIZE::BIGINT                       AS content_size,
                $1:CONTENT_MD5::VARCHAR                       AS content_md5,
                $1:BUCKET_NAME::VARCHAR                       AS bucket_name,
                $1:NAME::VARCHAR                              AS name,
                $1:KEY::VARCHAR                               AS key,
                $1:STORAGE_LOCATION_ID::BIGINT                AS storage_location_id,
                $1:ENDPOINT::VARCHAR                          AS endpoint,
                ($1:IS_PREVIEW::NUMBER != 0)                  AS is_preview,
                $1:STATUS::VARCHAR                            AS status,
                SPLIT_PART(METADATA$FILENAME, '/', 3)::BIGINT AS stack,
                SPLIT_PART(METADATA$FILENAME, '/', 4)::DATE   AS snapshot_date,
                METADATA$FILENAME                             AS filename
            FROM @{stage}/rds-snapshot/FILES/
        """,
        "file_format": None,
    },
    "changes": {
        "table": "changes",
        "stage_subpath": "rds-snapshot/CHANGES/",
        "select": """
            SELECT
                $1:CHANGE_NUM::BIGINT                         AS change_num,
                $1:TIME_STAMP::TIMESTAMP_NTZ(9)               AS time_stamp,
                $1:OBJECT_ID::BIGINT                          AS object_id,
                $1:OBJECT_VERSION::BIGINT                     AS object_version,
                $1:USER_ID::BIGINT                            AS user_id,
                $1:OBJECT_TYPE::VARCHAR                       AS object_type,
                $1:CHANGE_TYPE::VARCHAR                       AS change_type,
                SPLIT_PART(METADATA$FILENAME, '/', 3)::BIGINT AS stack,
                SPLIT_PART(METADATA$FILENAME, '/', 4)::DATE   AS snapshot_date,
                METADATA$FILENAME                             AS filename
            FROM @{stage}/rds-snapshot/CHANGES/
        """,
        "file_format": None,
    },
    "sent_messages": {
        "table": "sent_messages",
        "stage_subpath": "rds-snapshot/SENT_MESSAGES/",
        "select": """
            SELECT
                $1:CHANGE_NUM::BIGINT                         AS change_num,
                $1:TIME_STAMP::TIMESTAMP_NTZ(9)               AS time_stamp,
                $1:OBJECT_ID::BIGINT                          AS object_id,
                $1:OBJECT_VERSION::BIGINT                     AS object_version,
                $1:OBJECT_TYPE::VARCHAR                       AS object_type,
                SPLIT_PART(METADATA$FILENAME, '/', 3)::BIGINT AS stack,
                SPLIT_PART(METADATA$FILENAME, '/', 4)::DATE   AS snapshot_date,
                METADATA$FILENAME                             AS filename
            FROM @{stage}/rds-snapshot/SENT_MESSAGES/
        """,
        "file_format": None,
    },
}


def run_snow_sql(sql, connection, role):
    cmd = ["snow", "sql", "-q", sql, "--format=JSON"]
    if connection:
        cmd += ["--connection", connection]
    if role:
        cmd += ["--role", role]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"snow sql failed:\n{result.stderr or result.stdout}")
    return json.loads(result.stdout)


def list_backlog_dates(stage_subpath, connection, role):
    list_sql = f"LIST @{DATABASE}.{SCHEMA}.{STAGE}/{stage_subpath};"
    dates_sql = f"""
        {list_sql}
        SELECT DISTINCT SPLIT_PART("name", '/', 7) AS snapshot_date
        FROM TABLE(RESULT_SCAN(LAST_QUERY_ID()))
        ORDER BY snapshot_date;
    """
    result = run_snow_sql(dates_sql, connection, role)
    dates_result = result[-1]
    return [row["SNAPSHOT_DATE"] for row in dates_result]


def copy_one_date(target_name, target, snapshot_date, warehouse, connection, role):
    select_sql = target["select"].format(stage=f"{DATABASE}.{SCHEMA}.{STAGE}")
    file_format_clause = f"FILE_FORMAT = ({target['file_format']})" if target["file_format"] else ""
    copy_sql = f"""
        USE WAREHOUSE {warehouse};
        USE SCHEMA {DATABASE}.{SCHEMA};
        COPY INTO {target['table']} FROM (
            {select_sql}
        )
        {file_format_clause}
        PATTERN = '.*\\/[0-9]+\\/{snapshot_date}\\/.*\\.gz\\.parquet';
    """
    result = run_snow_sql(copy_sql, connection, role)
    copy_result = result[-1]
    rows_loaded = sum(row.get("rows_loaded", 0) or 0 for row in copy_result)
    failed = [row for row in copy_result if row.get("status") not in ("LOADED", "LOAD_SKIPPED")]
    return rows_loaded, failed


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--target", choices=sorted(TARGETS), required=True)
    parser.add_argument("--warehouse", default="compute_medium")
    parser.add_argument("--max-dates", type=int, default=10, help="Safety cap per invocation; re-run to continue.")
    parser.add_argument("--start-date", help="Skip dates before this (inclusive), e.g. 2026-06-15.")
    parser.add_argument("--connection", help="Named `snow` CLI connection to use.")
    parser.add_argument("--role", help="Snowflake role to run as (needs INSERT on the target table).")
    parser.add_argument("--dry-run", action="store_true", help="List the dates that would be processed and exit.")
    args = parser.parse_args()

    target = TARGETS[args.target]
    dates = list_backlog_dates(target["stage_subpath"], args.connection, args.role)
    if args.start_date:
        dates = [d for d in dates if d >= args.start_date]
    dates = dates[: args.max_dates]

    if not dates:
        print("No backlog dates found for this target/start-date.")
        return

    if args.dry_run:
        print(f"Would process {len(dates)} date(s): {dates[0]} .. {dates[-1]}")
        return

    print(f"Processing {len(dates)} date(s) for '{args.target}' on {args.warehouse}: {dates[0]} .. {dates[-1]}")
    for snapshot_date in dates:
        start = time.monotonic()
        try:
            rows_loaded, failed = copy_one_date(
                args.target, target, snapshot_date, args.warehouse, args.connection, args.role
            )
        except RuntimeError as e:
            print(f"  {snapshot_date}: ERROR ({time.monotonic() - start:.0f}s) -- {e}", file=sys.stderr)
            continue
        elapsed = time.monotonic() - start
        status = "OK" if not failed else f"{len(failed)} FILE(S) FAILED"
        print(f"  {snapshot_date}: {status}, {rows_loaded} row(s) loaded ({elapsed:.0f}s)")
        for row in failed:
            print(f"    {row.get('file')}: {row.get('status')} -- {row.get('first_error')}", file=sys.stderr)

    print("Done. Re-run with --start-date past the last processed date (or a higher --max-dates) to continue.")


if __name__ == "__main__":
    main()