Python Projects for Data Engineers

import json
import os
import time
from datetime import datetime, timedelta
from pathlib import Path

import duckdb
import requests

STATE = Path("state.json")
OVERLAP = timedelta(minutes=5)


def get(session, url, params=None, attempts=6):
    for attempt in range(attempts):
        r = session.get(url, params=params, timeout=30)
        if r.status_code == 200:
            return r
        limited = r.status_code in (403, 429)
        if limited and "retry-after" in r.headers:
            time.sleep(int(r.headers["retry-after"]))
        elif limited and r.headers.get("x-ratelimit-remaining") == "0":
            time.sleep(max(int(r.headers["x-ratelimit-reset"]) - time.time(), 0) + 1)
        elif r.status_code == 429 or r.status_code >= 500:
            time.sleep(min(2 ** attempt, 60))
        else:
            r.raise_for_status()  # a 401 or 404 will not fix itself
    raise RuntimeError(f"gave up on {url} after {attempts} attempts")


def load(repo, db="github.duckdb"):
    state = json.loads(STATE.read_text()) if STATE.exists() else {}
    cursor = datetime.fromisoformat(state.get(repo, "1970-01-01T00:00:00+00:00"))
    params = {
        "state": "all",
        "sort": "updated",
        "direction": "desc",  # an update mid-crawl repeats a row, never skips one
        "since": (cursor - OVERLAP).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "per_page": 100,
    }
    url, rows = f"https://api.github.com/repos/{repo}/issues", []
    with requests.Session() as s:
        s.headers["Authorization"] = f"Bearer {os.environ['GITHUB_TOKEN']}"
        while url:
            r = get(s, url, params)
            rows += [(i["id"], i["number"], i["updated_at"], json.dumps(i)) for i in r.json()]
            url, params = r.links.get("next", {}).get("url"), None  # the link keeps the query

    con = duckdb.connect(db)
    con.execute("""CREATE TABLE IF NOT EXISTS issues (
        id BIGINT PRIMARY KEY, number INTEGER, updated_at TIMESTAMPTZ, payload JSON)""")
    con.executemany("""INSERT INTO issues VALUES (?, ?, ?, ?)
        ON CONFLICT DO UPDATE SET number = EXCLUDED.number,
            updated_at = EXCLUDED.updated_at, payload = EXCLUDED.payload""", rows)
    con.close()

    if rows:  # advance the cursor only after the rows are stored
        state[repo] = max(row[2] for row in rows).replace("Z", "+00:00")
        STATE.write_text(json.dumps(state))