afterhours-shift-manager/scripts/cutover/retarget_holiday_schedules_to_sqs.py

#!/usr/bin/env python3
"""Retarget outstanding holiday-* Scheduler one-offs from Lambda to jobs SQS.

Lists holiday-activate-* / holiday-deactivate-* in the destination account and
updates Target to the jobs queue with Input
``{"event":"holiday","action":"activate|deactivate","date":"YYYY-MM-DD"}``.
Dry-run unless --execute. Does not create schedules that are already past.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from datetime import datetime, timezone
from zoneinfo import ZoneInfo

import boto3
from botocore.exceptions import ClientError

PROD_ACCOUNT = "011934824531"
DEV_ACCOUNT = "710827005802"
PREFIXES = ("holiday-activate-", "holiday-deactivate-")
_AT = re.compile(r"^at\((\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})\)$")
_DATE = re.compile(r"(\d{4}-\d{2}-\d{2})")


def _client(profile: str, region: str):
    return boto3.Session(profile_name=profile, region_name=region).client("scheduler")


def _account(profile: str) -> str:
    return (
        boto3.Session(profile_name=profile)
        .client("sts")
        .get_caller_identity()["Account"]
    )


def schedule_when(detail: dict) -> datetime | None:
    expr = (detail.get("ScheduleExpression") or "").strip()
    tzname = detail.get("ScheduleExpressionTimezone") or "America/New_York"
    match = _AT.match(expr)
    if match:
        naive = datetime.strptime(match.group(1), "%Y-%m-%dT%H:%M:%S")
        return naive.replace(tzinfo=ZoneInfo(tzname)).astimezone(timezone.utc)
    at = detail.get("EndDate") or detail.get("StartDate")
    if at is None:
        return None
    if at.tzinfo is None:
        return at.replace(tzinfo=timezone.utc)
    return at.astimezone(timezone.utc)


def action_and_date(name: str, existing_input: str) -> tuple[str, str]:
    action = "deactivate" if name.startswith("holiday-deactivate-") else "activate"
    date = ""
    if existing_input:
        try:
            parsed = json.loads(existing_input)
        except json.JSONDecodeError:
            parsed = {}
        if isinstance(parsed, dict):
            date = str(parsed.get("date") or "")
            if parsed.get("action") in {"activate", "deactivate"}:
                action = parsed["action"]
    if not date:
        compact = name.split("-")[-1]
        if len(compact) == 8 and compact.isdigit():
            date = f"{compact[0:4]}-{compact[4:6]}-{compact[6:8]}"
    if not date:
        match = _DATE.search(existing_input or "")
        if match:
            date = match.group(1)
    if not date:
        raise ValueError(f"cannot derive date from {name}")
    return action, date


def holiday_sqs_input(action: str, date: str) -> str:
    return json.dumps({"event": "holiday", "action": action, "date": date})


def _list_holiday(client):
    names = []
    token = None
    while True:
        kwargs = {"GroupName": "default"}
        if token:
            kwargs["NextToken"] = token
        resp = client.list_schedules(**kwargs)
        for item in resp.get("Schedules", []):
            name = item.get("Name", "")
            if name.startswith(PREFIXES):
                names.append(name)
        token = resp.get("NextToken")
        if not token:
            return names


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--profile", required=True)
    parser.add_argument("--region", default="us-east-1")
    parser.add_argument("--queue-arn", required=True)
    parser.add_argument(
        "--role-arn",
        default="arn:aws:iam::011934824531:role/tf-managed/afterhours-shift-manager-holiday-scheduler",
    )
    parser.add_argument("--execute", action="store_true")
    args = parser.parse_args()

    account = _account(args.profile)
    if account not in {PROD_ACCOUNT, DEV_ACCOUNT}:
        print("profile is not seahaven-prod or seahaven-dev", file=sys.stderr)
        return 2

    client = _client(args.profile, args.region)
    now = datetime.now(timezone.utc)
    updated = 0
    skipped = 0
    failed = 0

    for name in _list_holiday(client):
        detail = client.get_schedule(Name=name, GroupName="default")
        expr = detail.get("ScheduleExpression", "")
        when = schedule_when(detail)
        if when is not None and when < now:
            print(f"skip past {name} expr={expr}")
            skipped += 1
            continue
        try:
            action, date = action_and_date(
                name, detail.get("Target", {}).get("Input", "")
            )
        except ValueError as exc:
            print(f"skip {exc}", file=sys.stderr)
            skipped += 1
            continue
        payload = holiday_sqs_input(action, date)
        print(f"would retarget {name} action={action} date={date}")
        if not args.execute:
            continue
        try:
            client.update_schedule(
                Name=name,
                GroupName="default",
                ScheduleExpression=expr,
                ScheduleExpressionTimezone=detail.get(
                    "ScheduleExpressionTimezone", "America/New_York"
                ),
                FlexibleTimeWindow={"Mode": "OFF"},
                ActionAfterCompletion=detail.get("ActionAfterCompletion", "DELETE"),
                Target={
                    "Arn": args.queue_arn,
                    "RoleArn": args.role_arn,
                    "Input": payload,
                },
            )
            updated += 1
        except ClientError as exc:
            code = exc.response["Error"]["Code"]
            print(f"failed {name}: {code}", file=sys.stderr)
            failed += 1

    print(f"updated={updated} skipped={skipped} failed={failed} execute={args.execute}")
    if not args.execute:
        print("dry-run; pass --execute to UpdateSchedule")
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())