code/fill_repo_cap.py

"""Targeted fill pass: fetch the content of bundle files that the crawler skipped with skipped='repo_cap'
(the 100 MB-per-repo memory guard). Only a handful of huge skill-collection repos are affected.

1. Collect targets from the bucket (processed/github_delta/<run>/files where skipped='repo_cap'): the exact
   git blob SHA of every skipped file is already recorded, so we fetch exactly those blobs.
2. Per repo: blobless shallow clone, then batched `git fetch --stdin <blob oids>` + `cat-file --batch`,
   written straight to parquet in small batches (bounded RAM regardless of repo size).
3. Rows use the original skill_uid / rel_path, so they supersede the repo_cap placeholder rows:
   processed/github_delta/<run>/files/fill-repo_cap-<n>.parquet   (consumers: prefer the row WITH content)

Usage: python code/fill_repo_cap.py [run ...]     (default: all three runs). Resumable via work/fill/done.txt.
"""
import os, sys, json, shutil, subprocess, hashlib, time
from pathlib import Path
import pyarrow as pa, pyarrow.parquet as pq

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "code"))
from bucket_query import urls, connect
from build_github_delta import FI_SCHEMA
from huggingface_hub import HfApi

RUNS = sys.argv[1:] or ["pass1_sourcegraph", "pass2_new", "backfill_gitskills"]
WORK = ROOT / "work/fill"; WORK.mkdir(parents=True, exist_ok=True)
DONE = WORK / "done.txt"; done = set(DONE.read_text().split()) if DONE.exists() else set()
ENV = {**os.environ, "GIT_TERMINAL_PROMPT": "0", "GIT_LFS_SKIP_SMUDGE": "1"}
NO_LAZY = {**ENV, "GIT_NO_LAZY_FETCH": "1"}
BATCH_OIDS, BATCH_BYTES, MAX_BIN = 1500, 48 << 20, 256 << 10
api = HfApi()


def git(*a, cwd=None, inp=None, env=ENV, timeout=900):
    return subprocess.run(["git", *a], cwd=cwd, input=inp, capture_output=True, env=env, timeout=timeout)


RAW_DIR = {"pass1_sourcegraph": "github_delta", "pass2_new": "github_delta_p2", "backfill_gitskills": "github_backfill"}


def clone(repo):
    d = WORK / hashlib.md5(repo.encode()).hexdigest(); shutil.rmtree(d, ignore_errors=True)
    r = git("clone", "-q", "--depth", "1", "--filter=blob:none", "--no-checkout", "--single-branch",
            f"https://github.com/{repo}.git", str(d), timeout=600)
    if r.returncode:
        print(f"  {repo}: clone failed ({r.stderr.decode(errors='replace')[:120].strip()}) — will retry next run", flush=True)
        return None
    return d


def fetch_and_write(run, repo, rows, d, tag):
    """rows carry blob_sha; fetch those blobs in batches and upload fill rows. Returns (rows_out, blobs_got, blobs)."""
    by_sha = {}
    for x in rows: by_sha.setdefault(x["blob_sha"], []).append(x)
    oids = list(by_sha); n_out = n_got = 0
    batches, cur, cur_b = [], [], 0   # batch by bytes (~48 MB) as well as count: capped files can be up to 1 MB each
    for o in oids:
        sz = by_sha[o][0]["size"] or 0
        if cur and (len(cur) >= BATCH_OIDS or cur_b + sz > BATCH_BYTES): batches.append(cur); cur, cur_b = [], 0
        cur.append(o); cur_b += sz
    if cur: batches.append(cur)
    for part, chunk in enumerate(batches):
        git("-c", "fetch.negotiationAlgorithm=noop", "fetch", "-q", "--no-tags", "--no-write-fetch-head",
            "--recurse-submodules=no", "--filter=blob:none", "origin", "--stdin", cwd=d, inp=("\n".join(chunk) + "\n").encode())
        buf = git("cat-file", "--batch", cwd=d, env=NO_LAZY, inp=("\n".join(chunk) + "\n").encode()).stdout
        got, i = {}, 0
        while i < len(buf):
            nl = buf.index(b"\n", i); hdr = buf[i:nl].decode().split()
            if len(hdr) < 3 or hdr[1] == "missing": i = nl + 1; continue
            n = int(hdr[2]); got[hdr[0]] = buf[nl + 1: nl + 1 + n]; i = nl + 2 + n
        del buf
        out = []
        for sha in chunk:
            c_ = got.get(sha)
            for x in by_sha[sha]:
                text = c_ is not None and b"\0" not in c_[:8192]
                keep_bin = c_ is not None and not text and len(c_) <= MAX_BIN
                out.append({**{k2: x[k2] for k2 in ("skill_uid", "source", "repo", "skill_dir", "rel_path", "size", "blob_sha")},
                            "content": c_.decode("utf-8", "replace") if text else None,
                            "content_bin": c_ if keep_bin else None, "has_content": text or keep_bin,
                            "skipped": None if (text or keep_bin) else ("binary_large" if c_ is not None else "fetch_failed")})
        n_got += len(got); n_out += len(out)
        name = f"fill-repo_cap-{tag}{hashlib.md5(repo.encode()).hexdigest()[:10]}-{part:03d}.parquet"
        p = WORK / name; pq.write_table(pa.Table.from_pylist(out, FI_SCHEMA), p, compression="zstd")
        for attempt in range(6):
            try:
                api.batch_bucket_files("Mercity/SkillsStorage", add=[(str(p), f"processed/github_delta/{run}/files/{name}")]); break
            except Exception:
                time.sleep(15 * (attempt + 1))
        else:
            raise RuntimeError(f"upload failed {name}")
        p.unlink(); del got, out
    return n_out, n_got, len(oids)


def targets(run, by_path):
    tf = WORK / f"targets-{run}{'-bypath' if by_path else ''}.parquet"
    if not tf.exists():
        c = connect("300MB")
        c.execute(f"""copy (select skill_uid, source, repo, skill_dir, rel_path, size, blob_sha from read_parquet({urls(f'processed/github_delta/{run}/files/')})
                      where skipped = 'repo_cap' and blob_sha is {'' if by_path else 'not '}null) to '{tf}' (format parquet)""")
        c.close()
    t = pq.read_table(tf); by_repo = {}
    for r in t.to_pylist(): by_repo.setdefault(r["repo"], []).append(r)
    return by_repo


for by_path in (False, True):
    # by_path=True: files skipped on the codeload path have no blob SHA recorded -> resolve path -> SHA with
    # `git ls-tree` at the exact commit that was crawled (full commit id from raw/<run>/repos.jsonl), then fetch.
    for run in RUNS:
        by_repo = targets(run, by_path)
        print(f"{run}{' (by path)' if by_path else ''}: {sum(map(len, by_repo.values()))} files in {len(by_repo)} repos", flush=True)
        commits = {}
        if by_path and by_repo:
            for l in open(ROOT / "data/raw" / RAW_DIR[run] / "repos.jsonl"):
                m = json.loads(l)
                if m.get("commit"): commits[m["repo"]] = m["commit"]
        for repo in sorted(by_repo):
            key = f"{run}|{repo}{'|bypath' if by_path else ''}"
            if key in done: continue
            rows = by_repo[repo]
            d = clone(repo)
            if d is None: continue
            if by_path:
                ref = commits.get(repo) or "HEAD"
                if ref != "HEAD" and git("rev-parse", "HEAD", cwd=d).stdout.decode().strip() != ref:
                    git("fetch", "-q", "--depth", "1", "--filter=blob:none", "origin", ref, cwd=d, timeout=600)
                full = lambda x: x["rel_path"] if x["skill_dir"] == "." else f"{x['skill_dir']}/{x['rel_path']}"
                dirs = sorted({x["skill_dir"] for x in rows}); sha_of = {}
                for k in range(0, len(dirs), 200):
                    out = git("--literal-pathspecs", "ls-tree", "-r", "-z", ref, "--", *dirs[k:k + 200], cwd=d).stdout
                    for rec in out.decode(errors="replace").split("\0"):
                        if "\t" not in rec: continue
                        info, path = rec.split("\t", 1); parts = info.split()
                        if len(parts) == 3 and parts[1] == "blob": sha_of[path] = parts[2]
                rows = [{**x, "blob_sha": sha_of[full(x)]} for x in rows if full(x) in sha_of]
            n_out, n_got, n_oids = fetch_and_write(run, repo, rows, d, "bypath-" if by_path else "")
            shutil.rmtree(d, ignore_errors=True)
            with open(DONE, "a") as f: f.write(key + "\n")
            print(f"  {repo}: {n_out} rows written ({n_got}/{n_oids} blobs fetched)" + (f", {len(by_repo[repo]) - len(rows)} paths not found at commit" if by_path else ""), flush=True)
print("fill done", flush=True)