Documents

// node import.mjs records.json
import { readFileSync } from "node:fs";

const HEADBAND = process.env.HEADBAND_URL ?? "https://headband.dev";
const KEY = process.env.HEADBAND_ADMIN_KEY; // hb_adm_..., on a server only
const INDEX = process.env.HEADBAND_INDEX ?? "products";
const PRIMARY_KEY = process.env.HEADBAND_PRIMARY_KEY ?? "id"; // "objectID" for an Algolia export
const MAX_BYTES = 3.5 * 1024 * 1024; // under the 4 MiB body limit, with room to spare

const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));

// Waits out 429s for the time Retry-After names, then answers status and body.
async function call(method, path, body) {
  for (;;) {
    const res = await fetch(`${HEADBAND}${path}`, {
      method,
      headers: { Authorization: `Bearer ${KEY}`, "Content-Type": "application/json" },
      body: body === undefined ? undefined : JSON.stringify(body),
    });
    if (res.status === 429) {
      await sleep(Number(res.headers.get("retry-after") ?? 1) * 1000);
      continue;
    }
    return { status: res.status, data: await res.json().catch(() => null) };
  }
}

// One bad record fails its whole batch later, in the task, so wait for each task.
async function waitFor({ taskUid, status }) {
  let task = { status };
  while (task.status === "enqueued" || task.status === "processing") {
    await sleep(500);
    task = (await call("GET", `/v1/tasks/${taskUid}?index=${INDEX}`)).data;
  }
  if (task.status !== "succeeded") {
    throw new Error(`task ${taskUid} ${task.status}: ${JSON.stringify(task.error)}`);
  }
}

async function send(documents) {
  const res = await call("POST", "/v1/bulk", { index: INDEX, primaryKey: PRIMARY_KEY, documents });
  if (res.status >= 300) throw new Error(`bulk ${res.status}: ${JSON.stringify(res.data)}`);
  for (const task of res.data.tasks) await waitFor(task);
  console.log(`indexed ${documents.length}`);
}

const records = JSON.parse(readFileSync(process.argv[2], "utf8"));

// 409 means the index exists already.
const created = await call("POST", "/v1/indexes", { uid: INDEX, primaryKey: PRIMARY_KEY });
if (created.status >= 400 && created.status !== 409) throw new Error(JSON.stringify(created.data));

let batch = [];
let bytes = 0;
for (const record of records) {
  const size = Buffer.byteLength(JSON.stringify(record)) + 1;
  if (bytes + size > MAX_BYTES && batch.length > 0) {
    await send(batch);
    batch = [];
    bytes = 0;
  }
  batch.push(record);
  bytes += size;
}
if (batch.length > 0) await send(batch);