Export results in bulk

"""Backfill one own brand's mention rows into a newline-delimited JSON file.

Usage: python backfill_mentions.py <org_id> <brand_id> <from> <to>
Dates are YYYY-MM-DD (UTC), inclusive.
"""
import json
import os
import sys
import time

import requests

API = "https://api.surfais.com/v1"
PAGE_SIZE = 1000        # the largest page /mentions serves
MAX_RETRY_AFTER = 300   # never sleep longer than this on a 429 or 503
MAX_WALKS = 5           # give up after this many restarted walks
SETTLE_SECONDS = 60     # how long the change token must hold still

session = requests.Session()
session.headers["Authorization"] = f"Bearer {os.environ['SURFAIS_API_KEY']}"


class RestartWalk(Exception):
    """The brand's data changed under the walk. Start again from the first page."""


def get(path, params=None):
    """GET one JSON body. Waits out 429 and 503, and raises RestartWalk when the walk is void."""
    while True:
        resp = session.get(f"{API}{path}", params=params, timeout=60)
        if resp.status_code == 200:
            pause_when_minute_is_spent(resp)
            return resp.json()
        try:
            body = resp.json()
            error, request_id = body["error"], body["request_id"]
        except (ValueError, KeyError):
            error, request_id = {"code": "unknown", "message": resp.text[:200]}, None
        wait = int(resp.headers.get("Retry-After", "0"))
        if resp.status_code in (429, 503) and 0 < wait <= MAX_RETRY_AFTER:
            time.sleep(wait)
            continue
        if resp.status_code == 409 and error["code"] == "scan_in_progress":
            raise RestartWalk("scan_in_progress")  # the settle wait outlasts its Retry-After
        details = error.get("details")
        stale = isinstance(details, dict) and details.get("reason") == "stale"
        if resp.status_code == 400 and error["code"] == "invalid_cursor" and stale:
            raise RestartWalk("stale cursor")
        sys.exit(f"{resp.status_code} {error['code']}: {error['message']} (request {request_id})")


def pause_when_minute_is_spent(resp):
    """Sleep to the end of the minute window rather than run into a 429."""
    if resp.headers.get("X-RateLimit-Remaining") == "0":
        reset = int(resp.headers.get("X-RateLimit-Reset", "0"))
        time.sleep(max(0.0, reset - time.time()) + 1)


def walk_mentions(org_id, brand_id, date_from, date_to):
    """Yield every mention row in the window, newest date first."""
    params = {"from": date_from, "to": date_to, "limit": PAGE_SIZE}
    while True:
        page = get(f"/orgs/{org_id}/brands/{brand_id}/mentions", params)
        yield from page["data"]
        if not page["pagination"]["has_more"]:
            return
        # The cursor carries the window and the filters. Send it alone.
        params = {"cursor": page["pagination"]["next_cursor"], "limit": PAGE_SIZE}


def wait_until_settled(org_id, brand_id):
    """Block until the brand's change token stops moving."""
    path = f"/orgs/{org_id}/brands/{brand_id}"
    token = get(path)["data"]["last_scan_changed_at"]
    while True:
        time.sleep(SETTLE_SECONDS)
        latest = get(path)["data"]["last_scan_changed_at"]
        if latest == token:
            return
        token = latest


def main():
    org_id, brand_id, date_from, date_to = sys.argv[1:5]
    out_path = f"mentions-{brand_id}.ndjson"
    for walk in range(1, MAX_WALKS + 1):
        try:
            rows = 0
            # Write to a temporary file, so a restarted walk never leaves half a result.
            with open(f"{out_path}.part", "w", encoding="utf-8") as out:
                for row in walk_mentions(org_id, brand_id, date_from, date_to):
                    out.write(json.dumps(row) + "\n")
                    rows += 1
            os.replace(f"{out_path}.part", out_path)
            print(f"wrote {rows} rows to {out_path}")
            return
        except RestartWalk as reason:
            print(f"walk {walk} abandoned ({reason}); waiting for the brand to settle", file=sys.stderr)
            wait_until_settled(org_id, brand_id)
    sys.exit("the brand's data kept changing; run the backfill again once its scan has finished")


if __name__ == "__main__":
    main()

// backfill-mentions.mjs — Node.js 18 or later
// Usage: node backfill-mentions.mjs <orgId> <brandId> <from> <to>
// Dates are YYYY-MM-DD (UTC), inclusive.
import { once } from "node:events";
import { createWriteStream } from "node:fs";
import { rename } from "node:fs/promises";
import { setTimeout as sleep } from "node:timers/promises";

const API = "https://api.surfais.com/v1";
const PAGE_SIZE = 1000; // the largest page /mentions serves
const MAX_RETRY_AFTER = 300; // never sleep longer than this on a 429 or 503
const MAX_WALKS = 5; // give up after this many restarted walks
const SETTLE_SECONDS = 60; // how long the change token must hold still

const headers = { Authorization: `Bearer ${process.env.SURFAIS_API_KEY}` };

// The brand's data changed under the walk. Start again from the first page.
class RestartWalk extends Error {}

// GET one JSON body. Waits out 429 and 503, and throws RestartWalk when the walk is void.
async function get(path, params = {}) {
  const url = `${API}${path}?${new URLSearchParams(params)}`;
  for (;;) {
    const res = await fetch(url, { headers, signal: AbortSignal.timeout(60_000) });
    const text = await res.text();
    let body;
    try {
      body = JSON.parse(text);
    } catch {
      body = { error: { code: "unknown", message: text.slice(0, 200) }, request_id: null };
    }
    if (res.status === 200) {
      await pauseWhenMinuteIsSpent(res);
      return body;
    }
    const wait = Number(res.headers.get("retry-after") ?? 0);
    if ((res.status === 429 || res.status === 503) && wait > 0 && wait <= MAX_RETRY_AFTER) {
      await sleep(wait * 1000);
      continue;
    }
    if (res.status === 409 && body.error.code === "scan_in_progress") {
      throw new RestartWalk("scan_in_progress"); // the settle wait outlasts its Retry-After
    }
    const stale = body.error.details?.reason === "stale";
    if (res.status === 400 && body.error.code === "invalid_cursor" && stale) {
      throw new RestartWalk("stale cursor");
    }
    throw new Error(`${res.status} ${body.error.code}: ${body.error.message} (request ${body.request_id})`);
  }
}

// Sleep to the end of the minute window rather than run into a 429.
async function pauseWhenMinuteIsSpent(res) {
  if (res.headers.get("x-ratelimit-remaining") === "0") {
    const reset = Number(res.headers.get("x-ratelimit-reset") ?? 0);
    await sleep(Math.max(0, reset * 1000 - Date.now()) + 1000);
  }
}

// Yield every mention row in the window, newest date first.
async function* walkMentions(orgId, brandId, from, to) {
  let params = { from, to, limit: PAGE_SIZE };
  for (;;) {
    const page = await get(`/orgs/${orgId}/brands/${brandId}/mentions`, params);
    yield* page.data;
    if (!page.pagination.has_more) return;
    // The cursor carries the window and the filters. Send it alone.
    params = { cursor: page.pagination.next_cursor, limit: PAGE_SIZE };
  }
}

// Resolve once the brand's change token stops moving.
async function waitUntilSettled(orgId, brandId) {
  const path = `/orgs/${orgId}/brands/${brandId}`;
  let token = (await get(path)).data.last_scan_changed_at;
  for (;;) {
    await sleep(SETTLE_SECONDS * 1000);
    const latest = (await get(path)).data.last_scan_changed_at;
    if (latest === token) return;
    token = latest;
  }
}

async function main() {
  const [orgId, brandId, from, to] = process.argv.slice(2);
  const outPath = `mentions-${brandId}.ndjson`;
  for (let walk = 1; walk <= MAX_WALKS; walk++) {
    // Write to a temporary file, so a restarted walk never leaves half a result.
    const out = createWriteStream(`${outPath}.part`, { encoding: "utf8" });
    try {
      let rows = 0;
      for await (const row of walkMentions(orgId, brandId, from, to)) {
        if (!out.write(`${JSON.stringify(row)}\n`)) await once(out, "drain");
        rows++;
      }
      out.end();
      await once(out, "finish");
      await rename(`${outPath}.part`, outPath);
      console.log(`wrote ${rows} rows to ${outPath}`);
      return;
    } catch (err) {
      out.destroy();
      if (!(err instanceof RestartWalk)) throw err;
      console.error(`walk ${walk} abandoned (${err.message}); waiting for the brand to settle`);
      await waitUntilSettled(orgId, brandId);
    }
  }
  throw new Error("the brand's data kept changing; run the backfill again once its scan has finished");
}

main().catch((err) => {
  console.error(err.message);
  process.exit(1);
});