Documents

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

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

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

const headers = {
  Authorization: `Bearer ${KEY}`,
  "Content-Type": "application/json",
};

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

// One bad record fails its whole
// batch later, in the task, so wait
// for each task.
async function waitFor({
  taskUid,
  status,
}) {
  let task = { status };
  const path =
    `/v1/tasks/${taskUid}` +
    `?index=${INDEX}`;
  while (
    task.status === "enqueued" ||
    task.status === "processing"
  ) {
    await sleep(500);
    task = (await call("GET", path)).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);
  }
  const n = documents.length;
  console.log(`indexed ${n}`);
}

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 json = JSON.stringify(record);
  const size =
    Buffer.byteLength(json) + 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);