Enrich a CSV of LinkedIn URLs

#!/usr/bin/env python3
"""Enrich a CSV of LinkedIn profile URLs with the Forager API.

usage: enrich_csv.py input.csv output.csv [--work-emails] [--yes]
"""
import argparse
import csv
import os
import re
import sys
import time
from urllib.parse import urlparse

import requests

BASE = "https://api-v2.forager.ai/api"
ACCOUNT = os.environ["FORAGER_ACCOUNT_ID"]
HEADERS = {"X-API-KEY": os.environ["FORAGER_API_KEY"]}
OUT_FIELDS = ["linkedin_url", "handle", "id", "full_name", "headline", "current_title",
              "current_company", "current_domain", "work_email", "work_email_result", "error"]
PERSON_LOOKUP_CREDITS = 1   # https://www.forager.ai/docs/api-overview/credit-pricing
WORK_EMAIL_CREDITS = 5      # per address found; nothing when not found


def handle_from_url(url):
    """'https://www.linkedin.com/in/jane-doe-example/?x=1' -> 'jane-doe-example'."""
    url = url.strip()
    if "://" not in url:
        url = "https://" + url
    match = re.search(r"/in/([^/?#]+)", urlparse(url).path + "?")
    return match.group(1) if match else None


def read_rows(path):
    with open(path, newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def done_handles(path):
    if not os.path.exists(path):
        return set()
    with open(path, newline="", encoding="utf-8") as f:
        return {row["handle"] for row in csv.DictReader(f) if row.get("id")}

class AuthError(Exception):
    pass


def forager(method, path, json=None, retries=3):
    wait = 2.0
    for attempt in range(retries + 1):
        resp = requests.request(method, f"{BASE}/{path}", headers=HEADERS, json=json, timeout=60)
        if resp.status_code == 403:
            raise AuthError("Forager rejected FORAGER_API_KEY / FORAGER_ACCOUNT_ID: " + resp.text)
        if resp.status_code == 429 and attempt < retries:
            time.sleep(float(resp.headers.get("Retry-After", wait)))
            wait *= 2
            continue
        return resp
    return resp


def lookup_person(handle):
    """Return (row fields, error). 1 credit."""
    resp = forager("POST", f"{ACCOUNT}/datastorage/person_detail_lookup/",
                   json={"linkedin_public_identifier": handle})
    if resp.status_code != 200:
        return {}, f"HTTP {resp.status_code}: {resp.text[:200]}"
    person = resp.json()
    current = next((r for r in person.get("roles", []) if r.get("is_current")), None)
    org = (current or {}).get("organization") or {}
    return {
        "id": person["id"],
        "full_name": person.get("full_name", ""),
        "headline": person.get("headline") or "",
        "current_title": (current or {}).get("role_title") or "",
        "current_company": org.get("name") or (current or {}).get("organization_name") or "",
        "current_domain": org.get("domain") or "",
    }, ""


def bulk_work_emails(handles, max_wait=900):
    """Submit one batch for every handle, poll until it completes, return {handle: (email, result)}.
    5 credits per address found. The batch is rejected, and nothing is charged, if the balance
    cannot cover it."""
    resp = forager("POST", f"v2/{ACCOUNT}/datastorage/person_contacts_lookup/work_emails/bulk/",
                   json={"email_lookups": [{"linkedin_public_identifier": h} for h in handles]})
    if resp.status_code != 202:
        raise RuntimeError(f"bulk submit failed: HTTP {resp.status_code}: {resp.text[:300]}")
    batch_id = resp.json()["batch_id"]
    deadline = time.monotonic() + max_wait
    while time.monotonic() < deadline:
        poll = forager("GET", f"v2/{ACCOUNT}/datastorage/person_contacts_lookup/work_emails/bulk/{batch_id}/")
        batch = poll.json()
        if batch["status"] in ("completed", "failed"):
            break
        time.sleep(float(poll.headers.get("Retry-After", 10)))   # polling is free
    else:
        raise TimeoutError(f"batch {batch_id} still {batch['status']} after {max_wait}s")
    out = {}
    for item in batch.get("results", []):                        # one per submitted row, in order
        handle = handles[item["submitted_index"]]
        if item["status"] == "failed":
            out[handle] = ("", f"failed:{item.get('error_code')}")
        else:
            emails = item.get("emails") or []
            out[handle] = (emails[0] if emails else "", item.get("result") or "")
    return out

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("input")
    ap.add_argument("output")
    ap.add_argument("--work-emails", action="store_true", help="also find work emails (bulk, 5 credits per address found)")
    ap.add_argument("--yes", action="store_true", help="skip the cost confirmation")
    args = ap.parse_args()

    rows = read_rows(args.input)
    skip = done_handles(args.output)
    todo = []
    for row in rows:
        handle = handle_from_url(row.get("linkedin_url", ""))
        if handle and handle not in skip:
            todo.append((row["linkedin_url"], handle))
    bad = sum(1 for row in rows if not handle_from_url(row.get("linkedin_url", "")))

    print(f"{len(rows)} rows, {len(skip)} already done, {bad} without a usable /in/ URL, {len(todo)} to enrich")
    estimate = len(todo) * PERSON_LOOKUP_CREDITS
    if args.work_emails:
        estimate += len(todo) * WORK_EMAIL_CREDITS
        print(f"up to {estimate} credits ({PERSON_LOOKUP_CREDITS} per lookup, plus {WORK_EMAIL_CREDITS} per work email found)")
    else:
        print(f"{estimate} credits ({PERSON_LOOKUP_CREDITS} per lookup)")
    if not todo:
        return
    if not args.yes and input("continue? [y/N] ").strip().lower() != "y":
        sys.exit(1)

    new_file = not os.path.exists(args.output)
    enriched = []
    with open(args.output, "a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=OUT_FIELDS)
        if new_file:
            writer.writeheader()
        for url, handle in todo:
            fields, error = lookup_person(handle)
            record = {"linkedin_url": url, "handle": handle, "error": error, **fields}
            enriched.append(record)
            writer.writerow({k: record.get(k, "") for k in OUT_FIELDS})
            f.flush()
            print(f"{handle}: {record.get('full_name') or error}")

    if args.work_emails:
        found = [r["handle"] for r in enriched if r.get("id")]
        if found:
            emails = bulk_work_emails(found)
            for record in enriched:
                record["work_email"], record["work_email_result"] = emails.get(record["handle"], ("", ""))
            # Rewrite the file so the email columns land on the rows written above.
            existing = [r for r in read_rows(args.output) if r["handle"] not in found]
            with open(args.output, "w", newline="", encoding="utf-8") as f:
                writer = csv.DictWriter(f, fieldnames=OUT_FIELDS)
                writer.writeheader()
                for record in existing + enriched:
                    writer.writerow({k: record.get(k, "") for k in OUT_FIELDS})
            print(f"work emails: {sum(1 for r in enriched if r.get('work_email'))} found")


if __name__ == "__main__":
    try:
        main()
    except AuthError as e:
        sys.exit(str(e))

3 rows, 0 already done, 0 without a usable /in/ URL, 3 to enrich
up to 18 credits (1 per lookup, plus 5 per work email found)
continue? [y/N] y
jane-doe-example: Jane Doe
john-roe-example: John Roe
alex-example-smith: Alex Smith
work emails: 2 found

linkedin_url,handle,id,full_name,headline,current_title,current_company,current_domain,work_email,work_email_result,error
https://www.linkedin.com/in/jane-doe-example/,jane-doe-example,48213,Jane Doe,Head of Data at Acme Example,Head of Data,Acme Example,acme.com,jane.doe@acme.com,found,
https://www.linkedin.com/in/john-roe-example,john-roe-example,48377,John Roe,VP Sales,VP Sales,Example Labs,example.com,john.roe@example.com,found,
linkedin.com/in/alex-example-smith/,alex-example-smith,48901,Alex Smith,Founder,Founder,Smith Example Co,,,not_found,