krz/hutch-stats

  1from __future__ import annotations
  2
  3import argparse
  4from datetime import UTC, datetime, timedelta
  5from pathlib import Path
  6import re
  7
  8from sqlalchemy import select
  9
 10from srht_contrib.config import Settings
 11from srht_contrib.db import make_session_factory
 12from srht_contrib.models import TrackedActor
 13
 14
 15USERNAME_RE = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._-]{0,61}[A-Za-z0-9])?$")
 16
 17
 18def _normalize_actor(raw_username: str) -> str | None:
 19    username = raw_username.strip()
 20    if not username or username.startswith("#"):
 21        return None
 22    if username.startswith("~"):
 23        username = username[1:]
 24    if not USERNAME_RE.fullmatch(username):
 25        return None
 26    return f"~{username}"
 27
 28
 29def _iter_usernames(path: Path) -> list[str]:
 30    usernames: list[str] = []
 31    seen: set[str] = set()
 32    for raw_line in path.read_text(encoding="utf-8").splitlines():
 33        actor = _normalize_actor(raw_line)
 34        if actor is None:
 35            continue
 36        if actor in seen:
 37            continue
 38        seen.add(actor)
 39        usernames.append(actor)
 40    return usernames
 41
 42
 43def enqueue_actors(username_file: Path, *, stagger_seconds: int = 60, start_at: datetime | None = None) -> int:
 44    settings = Settings()
 45    session_factory = make_session_factory(settings)
 46    usernames = _iter_usernames(username_file)
 47    queued_at = start_at or datetime.now(tz=UTC)
 48    inserted = 0
 49
 50    batch_size = 250
 51    queued_in_batch = 0
 52
 53    with session_factory() as db:
 54        for index, actor in enumerate(usernames):
 55            next_poll_after = queued_at + timedelta(seconds=index * stagger_seconds)
 56            tracked_actor = db.scalar(select(TrackedActor).where(TrackedActor.actor == actor))
 57            if tracked_actor is None:
 58                tracked_actor = TrackedActor(
 59                    actor=actor,
 60                    is_active=True,
 61                    discovery_state="queued",
 62                    queued_for_discovery_at=queued_at,
 63                    next_poll_after=next_poll_after,
 64                    recent_backfill_status="pending",
 65                )
 66                db.add(tracked_actor)
 67                inserted += 1
 68            else:
 69                tracked_actor.is_active = True
 70                if tracked_actor.queued_for_discovery_at is None:
 71                    tracked_actor.queued_for_discovery_at = queued_at
 72                if tracked_actor.last_polled_at is None and tracked_actor.discovery_state != "indexed":
 73                    tracked_actor.discovery_state = "queued"
 74                    tracked_actor.next_poll_after = next_poll_after
 75
 76            queued_in_batch += 1
 77            if queued_in_batch >= batch_size:
 78                db.commit()
 79                queued_in_batch = 0
 80
 81        if queued_in_batch:
 82            db.commit()
 83
 84    return inserted
 85
 86
 87def main() -> None:
 88    parser = argparse.ArgumentParser(description="Durably enqueue SourceHut actors without polling them immediately.")
 89    parser.add_argument(
 90        "username_file",
 91        nargs="?",
 92        default="srht_usernames.txt",
 93        help="Path to a newline-delimited SourceHut username file.",
 94    )
 95    parser.add_argument(
 96        "--stagger-seconds",
 97        type=int,
 98        default=60,
 99        help="Seconds to space out each actor's first eligible poll time.",
100    )
101    args = parser.parse_args()
102
103    inserted = enqueue_actors(Path(args.username_file), stagger_seconds=args.stagger_seconds)
104    print(f"Enqueued {inserted} new actors from {args.username_file}.")
105
106
107if __name__ == "__main__":
108    main()