diff --git a/scripts/sweep_stale_upstream.py b/scripts/sweep_stale_upstream.py index 2f88184..264e4f5 100644 --- a/scripts/sweep_stale_upstream.py +++ b/scripts/sweep_stale_upstream.py @@ -67,6 +67,12 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--apply", action="store_true", help="write corrections (default: dry run)") parser.add_argument("--limit", type=int, default=0, help="stop after N accounts (0 = all)") + parser.add_argument( + "--statuses", + default="deactivated", + help="comma-separated upstream_status values to sweep (default: deactivated — " + "the only state the migration bug produces; 'all' sweeps every non-active status)", + ) args = parser.parse_args() database_url = os.environ.get("DATABASE_URL") @@ -74,11 +80,19 @@ def main() -> None: sys.exit("DATABASE_URL is required") conn = psycopg.connect(database_url) - rows = conn.execute( - "SELECT uid, did, upstream_status FROM account" - " WHERE status = 'active' AND upstream_status != 'active' ORDER BY uid" - ).fetchall() - print(f"{len(rows)} candidate accounts (status=active, upstream_status!=active)") + if args.statuses == "all": + rows = conn.execute( + "SELECT uid, did, upstream_status FROM account" + " WHERE status = 'active' AND upstream_status != 'active' ORDER BY uid" + ).fetchall() + else: + statuses = [s.strip() for s in args.statuses.split(",") if s.strip()] + rows = conn.execute( + "SELECT uid, did, upstream_status FROM account" + " WHERE status = 'active' AND upstream_status = ANY(%s) ORDER BY uid", + (statuses,), + ).fetchall() + print(f"{len(rows)} candidate accounts (status=active, upstream_status in {args.statuses})") client = httpx.Client(timeout=10, headers={"user-agent": "zlay-sweep (atproto-relay)"}) checked = corrected = confirmed = failed = 0