Enrich a CSV of domains with company data

"""Enrich a CSV of domains with company data from HappyFetch.

Usage: python enrich.py domains.csv enriched.csv
Set MAX_AGE_DAYS to accept older cached answers, which cost less.
"""
import csv
import os
import sys
import time
import uuid
from concurrent.futures import ThreadPoolExecutor

import requests

API = "https://api.happyfetch.net"
HEADERS = {"Authorization": f"Bearer {os.environ['HAPPYFETCH_API_KEY']}"}
MAX_AGE = int(os.environ["MAX_AGE_DAYS"]) * 86400 if os.environ.get("MAX_AGE_DAYS") else None
COLUMNS = ["domain", "name", "website", "linkedin", "industry", "employees", "size",
           "founded", "city", "country", "cached", "credits", "error"]


def lookup(domain: str) -> dict:
    key = str(uuid.uuid4())  # one key per domain, reused on every retry: a retry never pays twice
    body = {"domain": domain, **({"max_age": MAX_AGE} if MAX_AGE else {})}
    for attempt in range(6):
        try:
            res = requests.post(f"{API}/v1/domain", headers={**HEADERS, "Idempotency-Key": key}, json=body, timeout=90)
            answer = res.json()
        except (requests.RequestException, ValueError):
            time.sleep(2**attempt)  # a timeout or a dropped connection
            continue
        if res.ok:
            return to_row(domain, answer)
        error = answer["error"]
        if not error["retryable"]:
            return {"domain": domain, "credits": answer["meta"]["credits_charged"], "error": error["code"]}
        time.sleep(error.get("retry_after_seconds") or 2**attempt)
    return {"domain": domain, "error": "gave_up"}


def to_row(domain: str, answer: dict) -> dict:
    company, meta = answer["data"], answer["meta"]
    hq = company["headquarter"] or {}
    size = company["employeeCountRange"] or {}
    return {
        "domain": domain,
        "name": company["name"],
        "website": company["websiteUrl"],
        "linkedin": company["linkedInUrl"],
        "industry": company["industry"],
        "employees": company["employeeCount"],
        "size": f"{size['start']}-{size['end'] or ''}" if size else None,
        "founded": (company["foundedOn"] or {}).get("year"),
        "city": hq.get("city"),
        "country": hq.get("country"),
        "cached": "yes" if meta["cache"]["status"] == "hit" else "no",
        "credits": meta["credits_charged"],
    }


def main(src: str, dst: str) -> None:
    with open(src, newline="") as f:
        domains = list(dict.fromkeys(row["domain"].strip() for row in csv.DictReader(f) if row["domain"].strip()))
    rows = []
    with ThreadPoolExecutor(max_workers=4) as pool:
        for done, row in enumerate(pool.map(lookup, domains), 1):
            rows.append(row)
            print(f"\r{done:,}/{len(domains):,}", end="", flush=True)
    print()
    with open(dst, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=COLUMNS)
        writer.writeheader()
        writer.writerows(rows)
    spent = sum(row.get("credits") or 0 for row in rows)
    errors = sum(1 for row in rows if row.get("error"))
    print(f"Wrote {len(rows):,} rows to {dst}: {spent:,} credits, {errors:,} with an error code.")


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])

// Enrich a CSV of domains with company data from HappyFetch.
// Usage: node enrich.mjs domains.csv enriched.csv   (npm install csv-parse csv-stringify)
// Set MAX_AGE_DAYS to accept older cached answers, which cost less.
import { randomUUID } from "node:crypto";
import { readFileSync, writeFileSync } from "node:fs";
import { parse } from "csv-parse/sync";
import { stringify } from "csv-stringify/sync";

const API = "https://api.happyfetch.net";
const KEY = process.env.HAPPYFETCH_API_KEY;
const MAX_AGE = process.env.MAX_AGE_DAYS ? Number(process.env.MAX_AGE_DAYS) * 86400 : undefined;
const COLUMNS = ["domain", "name", "website", "linkedin", "industry", "employees", "size",
  "founded", "city", "country", "cached", "credits", "error"];
const sleep = (seconds) => new Promise((resolve) => setTimeout(resolve, seconds * 1000));

async function lookup(domain) {
  const idempotencyKey = randomUUID(); // one key per domain, reused on every retry: a retry never pays twice
  for (let attempt = 0; attempt < 6; attempt++) {
    let res, answer;
    try {
      res = await fetch(`${API}/v1/domain`, {
        method: "POST",
        headers: { Authorization: `Bearer ${KEY}`, "Content-Type": "application/json", "Idempotency-Key": idempotencyKey },
        body: JSON.stringify({ domain, max_age: MAX_AGE }),
        signal: AbortSignal.timeout(90_000),
      });
      answer = await res.json();
    } catch {
      await sleep(2 ** attempt); // a timeout or a dropped connection
      continue;
    }
    if (res.ok) return toRow(domain, answer);
    const { error, meta } = answer;
    if (!error.retryable) return { domain, credits: meta.credits_charged, error: error.code };
    await sleep(error.retry_after_seconds ?? 2 ** attempt);
  }
  return { domain, error: "gave_up" };
}

function toRow(domain, { data: company, meta }) {
  const size = company.employeeCountRange;
  return {
    domain,
    name: company.name,
    website: company.websiteUrl,
    linkedin: company.linkedInUrl,
    industry: company.industry,
    employees: company.employeeCount,
    size: size ? `${size.start}-${size.end ?? ""}` : null,
    founded: company.foundedOn?.year ?? null,
    city: company.headquarter?.city ?? null,
    country: company.headquarter?.country ?? null,
    cached: meta.cache.status === "hit" ? "yes" : "no",
    credits: meta.credits_charged,
  };
}

const [src, dst] = process.argv.slice(2);
const domains = [...new Set(parse(readFileSync(src), { columns: true, skip_empty_lines: true }).map((r) => r.domain?.trim()).filter(Boolean))];
const rows = new Array(domains.length);
let next = 0;
let done = 0;
async function worker() {
  while (next < domains.length) {
    const i = next++;
    rows[i] = await lookup(domains[i]);
    process.stdout.write(`\r${(++done).toLocaleString()}/${domains.length.toLocaleString()}`);
  }
}
await Promise.all(Array.from({ length: 4 }, worker)); // four lookups at a time
writeFileSync(dst, stringify(rows, { header: true, columns: COLUMNS }));
const spent = rows.reduce((sum, r) => sum + (r.credits ?? 0), 0);
const errors = rows.filter((r) => r.error).length;
console.log(`\nWrote ${rows.length.toLocaleString()} rows to ${dst}: ${spent.toLocaleString()} credits, ${errors.toLocaleString()} with an error code.`);