fetch_azure_incremental_WI_Relationships.js

/**
 * fetch_work_item_parents.incremental.js
 *
 * Purpose: Incremental-only importer of ADO parent-child links into DX.
 *          Uses a per-project watermark based on ADO System.ChangedDate to avoid gaps and duplicates.
 *
 * Writes: batched multi-row upserts into custom.ado_work_item_links
 * Watermarks: custom.ado_wi_links_watermarks (organization_name, project_name, last_changed_at)
 *
 * Schema columns used in links table:
 *   child_work_item_source_id, parent_work_item_source_id, relation_url
 *
 * Network: ADO HTTPS via http://DX_PROXY_USER:DX_PROXY_PASS@proxy.getdx.net:80
 *
 * Env vars (required):
 *   DX_PG_URL       PostgreSQL connection string
 *   ADO_PAT         Azure DevOps Personal Access Token
 *   DX_PROXY_USER   Proxy username
 *   DX_PROXY_PASS   Proxy password
 *
 * Env vars (optional):
 *   DRY_RUN=1                   If set, do not write to DB, emit SQL file instead
 *   OUT_FILE=ado_links_inserts.sql
 *   BATCH_SIZE=200              workitemsbatch chunk size
 *   FIRST_RUN_LOOKBACK_DAYS=3650
 *   MIN_ID=11339                minimum source id to accept for both child and parent
 *   WRITE_BUFFER_SIZE=200       rows per upsert statement
 *   WRITE_FLUSH_MS=250          timer-based flush throttle
 *   SLEEP_BETWEEN_BATCHES_MS=300
 *   SLEEP_BETWEEN_PROJECTS_MS=500
 */

const { Client } = require("pg");
const fs = require("fs");
const path = require("path");
const axios = require("axios");
const HttpsProxyAgent = require("https-proxy-agent");

// Required env
const { DX_PG_URL, ADO_PAT, DX_PROXY_USER, DX_PROXY_PASS } = process.env;
if (!DX_PG_URL || !ADO_PAT) {
  console.error("Set DX_PG_URL and ADO_PAT");
  process.exit(1);
}
if (!DX_PROXY_USER || !DX_PROXY_PASS) {
  console.error("Set DX_PROXY_USER and DX_PROXY_PASS");
  process.exit(1);
}

// Tunables
const DRY_RUN = process.env.DRY_RUN === "1";
const OUT_FILE = process.env.OUT_FILE || "ado_links_inserts.sql";
const BATCH_SIZE = parseInt(process.env.BATCH_SIZE || "200", 10);
const FIRST_RUN_LOOKBACK_DAYS = parseInt(process.env.FIRST_RUN_LOOKBACK_DAYS || "3650", 10);
const MIN_ID = parseInt(process.env.MIN_ID || "11339", 10);
const WRITE_BUFFER_SIZE = parseInt(process.env.WRITE_BUFFER_SIZE || "200", 10);
const WRITE_FLUSH_MS = parseInt(process.env.WRITE_FLUSH_MS || "250", 10);
const SLEEP_BETWEEN_BATCHES_MS = parseInt(process.env.SLEEP_BETWEEN_BATCHES_MS || "300", 10);
const SLEEP_BETWEEN_PROJECTS_MS = parseInt(process.env.SLEEP_BETWEEN_PROJECTS_MS || "500", 10);

// Normalize DX connection string
const normalizePostgresURL = (url) =>
  url.startsWith("postgres://") ? url.replace("postgres://", "postgresql://") : url;

// Proxy agent for ADO HTTPS calls
const proxyUrl = `http://${encodeURIComponent(DX_PROXY_USER)}:${encodeURIComponent(DX_PROXY_PASS)}@proxy.getdx.net:80`;
const httpsAgent = new HttpsProxyAgent(proxyUrl);

// Axios wrapper with retry
let httpCalls = 0;
async function httpJson(url, { method = "GET", headers = {}, body = undefined }, attempt = 1) {
  httpCalls += 1;
  try {
    const res = await axios({
      url,
      method,
      headers,
      data: body,
      httpsAgent,
      proxy: false,
      timeout: 30000,
      validateStatus: () => true,
    });
    if (res.status >= 200 && res.status < 300) return res.data;

    if ((res.status === 429 || (res.status >= 500 && res.status < 600)) && attempt < 6) {
      const ra = parseInt(res.headers["retry-after"] || "0", 10);
      const ms = ra > 0 ? ra * 1000 : attempt * 1000;
      await sleep(ms);
      return httpJson(url, { method, headers, body }, attempt + 1);
    }
    throw new Error(
      `HTTP ${res.status}: ${res.statusText} | ${typeof res.data === "string" ? res.data : JSON.stringify(res.data)}`
    );
  } catch (err) {
    if (attempt < 6) {
      await sleep(attempt * 1000);
      return httpJson(url, { method, headers, body }, attempt + 1);
    }
    throw err;
  }
}

const AUTH_HEADER = { Authorization: "Basic " + Buffer.from(":" + ADO_PAT).toString("base64") };

// Small helpers
function escapeSql(s) { return s.replace(/'/g, "''"); }
function parentIdFromUrl(url) { const n = parseInt(url.split("/").pop(), 10); return Number.isFinite(n) ? n : null; }
function chunk(arr, n) { const out = []; for (let i = 0; i < arr.length; i += n) out.push(arr.slice(i, i + n)); return out; }
async function sleep(ms) { return new Promise((r) => setTimeout(r, ms)); }
function toIso(d) { return new Date(d).toISOString(); }

// SQL constants
const ENSURE_WATERMARK_SQL = `
CREATE TABLE IF NOT EXISTS custom.ado_wi_links_watermarks (
  organization_name text NOT NULL,
  project_name      text NOT NULL,
  last_changed_at   timestamptz NOT NULL,
  PRIMARY KEY (organization_name, project_name)
);
`;

const LIST_SCOPES_SQL = `
SELECT
  ao.name AS organization_name,
  ap.name AS project_name,
  w.last_changed_at AS last_changed_at
FROM ado_projects ap
JOIN ado_organizations ao
  ON ap.organization_id::bigint = ao.id
LEFT JOIN custom.ado_wi_links_watermarks w
  ON w.organization_name = ao.name
 AND w.project_name      = ap.name
WHERE ap.api_accessible = TRUE
  AND ap.allowlisted    = TRUE
  AND ao.api_accessible = TRUE
ORDER BY ao.name, ap.name;
`;

const UPSERT_WATERMARK_SQL = `
INSERT INTO custom.ado_wi_links_watermarks (organization_name, project_name, last_changed_at)
VALUES ($1, $2, $3)
ON CONFLICT (organization_name, project_name)
DO UPDATE SET last_changed_at = EXCLUDED.last_changed_at;
`;

// Build one row insert for dry run output file
function insertSql(childId, parentId, relUrl) {
  const urlSql = relUrl ? `'${escapeSql(relUrl)}'` : "NULL";
  return (
    "INSERT INTO custom.ado_work_item_links " +
    "(child_work_item_source_id, parent_work_item_source_id, relation_url) VALUES " +
    `(${childId}, ${parentId}, ${urlSql}) ` +
    "ON CONFLICT (child_work_item_source_id, parent_work_item_source_id) " +
    "DO UPDATE SET relation_url = EXCLUDED.relation_url;"
  );
}

// Buffered writer for live mode with multi row upserts
function buildWriter(DRY_RUN, client, sqlLines) {
  if (DRY_RUN) {
    const writer = async (childId, parentId, relUrl) => {
      sqlLines.push(insertSql(childId, parentId, relUrl));
    };
    writer.flush = async () => {};
    return writer;
  }

  const buf = [];
  let flushTimer = null;

  async function flush() {
    if (buf.length === 0) return;
    const rows = buf.splice(0, buf.length);

    const values = [];
    const params = [];
    let p = 1;
    for (const r of rows) {
      values.push(`($${p++}, $${p++}, $${p++})`);
      params.push(r.childId, r.parentId, r.relUrl);
    }
    const sql = `
      INSERT INTO custom.ado_work_item_links
        (child_work_item_source_id, parent_work_item_source_id, relation_url)
      VALUES ${values.join(",")}
      ON CONFLICT (child_work_item_source_id, parent_work_item_source_id)
      DO UPDATE SET relation_url = EXCLUDED.relation_url
    `;

    // one retry on transient error
    for (let attempt = 1; ; attempt++) {
      try {
        await client.query(sql, params);
        break;
      } catch (e) {
        if (attempt >= 2) throw e;
        await sleep(500);
      }
    }

    await sleep(100);
  }

  async function scheduleFlush() {
    if (flushTimer) return;
    flushTimer = setTimeout(async () => {
      flushTimer = null;
      try { await flush(); } catch (e) { console.error(e.message || String(e)); }
    }, WRITE_FLUSH_MS);
  }

  const writer = async (childId, parentId, relUrl) => {
    buf.push({ childId, parentId, relUrl });
    if (buf.length >= WRITE_BUFFER_SIZE) {
      await flush();
    } else {
      await scheduleFlush();
    }
  };
  writer.flush = flush;
  return writer;
}

// Incremental fetch using WIQL by ChangedDate > sinceISO and workitemsbatch expand=Relations
async function incrementalProject(org, project, sinceISO, writeFn) {
  const base = `https://dev.azure.com/${encodeURIComponent(org)}`;
  const wiql = `
    SELECT [System.Id]
    FROM WorkItems
    WHERE [System.TeamProject] = '${escapeSql(project)}'
      AND [System.ChangedDate] > '${sinceISO}'
    ORDER BY [System.ChangedDate] ASC
  `;

  const idsRes = await httpJson(
    `${base}/${encodeURIComponent(project)}/_apis/wit/wiql?api-version=7.1-preview.2`,
    { method: "POST", headers: { ...AUTH_HEADER, "Content-Type": "application/json" }, body: JSON.stringify({ query: wiql }) }
  );
  const ids = (idsRes.workItems || []).map((w) => w.id).filter(Number.isFinite);
  if (ids.length === 0) return { scanned: 0, links: 0, maxChangedAt: null };

  let links = 0;
  let maxChangedAt = null;

  for (const slice of chunk(ids, BATCH_SIZE)) {
    const batch = await httpJson(
      `${base}/_apis/wit/workitemsbatch?api-version=7.1`,
      {
        method: "POST",
        headers: { ...AUTH_HEADER, "Content-Type": "application/json" },
        body: JSON.stringify({ ids: slice, fields: ["System.Id", "System.ChangedDate"], expand: "Relations" }),
      }
    );

    for (const wi of batch.value || []) {
      const childId = wi.id;
      const changed = wi.fields && wi.fields["System.ChangedDate"] ? new Date(wi.fields["System.ChangedDate"]) : null;
      if (changed && (!maxChangedAt || changed > maxChangedAt)) maxChangedAt = changed;

      for (const rel of wi.relations || []) {
        if (rel.rel !== "System.LinkTypes.Hierarchy-Reverse" || !rel.url) continue;
        const parentId = parentIdFromUrl(rel.url);
        if (!parentId) continue;
        if (childId < MIN_ID || parentId < MIN_ID) continue;
        const relUrl = `${base}/_apis/wit/workItems/${parentId}`;
        await writeFn(childId, parentId, relUrl);
        links += 1;
      }
    }

    await sleep(SLEEP_BETWEEN_BATCHES_MS);
  }

  return { scanned: ids.length, links, maxChangedAt };
}

async function main() {
  const started = new Date();
  const client = new Client({ connectionString: normalizePostgresURL(DX_PG_URL), ssl: { rejectUnauthorized: false } });
  await client.connect();

  // Ensure watermark table exists
  await client.query(ENSURE_WATERMARK_SQL);

  // Read scopes and existing watermarks
  const { rows: scopes } = await client.query(LIST_SCOPES_SQL);
  if (!scopes.length) {
    console.log("No API accessible and allowlisted projects found");
    await client.end();
    return;
  }

  const sqlLines = [];
  const writer = buildWriter(DRY_RUN, client, sqlLines);

  let totalLinksProcessed = 0;

  try {
    for (const { organization_name: org, project_name: project, last_changed_at } of scopes) {
      const lockKey = `${org}:${project}`;

      if (!DRY_RUN) {
        await client.query("BEGIN");
        await client.query("SET LOCAL statement_timeout = '30s'");
        await client.query("SET LOCAL lock_timeout = '5s'");
        await client.query("SET LOCAL idle_in_transaction_session_timeout = '60s'");
        // advisory lock to avoid concurrent runs per project
        await client.query("SELECT pg_advisory_lock(hashtext($1))", [lockKey]);
      }

      try {
        const since = last_changed_at
          ? new Date(last_changed_at)
          : new Date(Date.now() - FIRST_RUN_LOOKBACK_DAYS * 86400 * 1000);

        const { scanned, links, maxChangedAt } = await incrementalProject(org, project, toIso(since), writer);
        totalLinksProcessed += links;

        if (!DRY_RUN && writer.flush) await writer.flush();

        if (!DRY_RUN) {
          // advance watermark only if we actually saw newer changes
          if (maxChangedAt) {
            await client.query(UPSERT_WATERMARK_SQL, [org, project, toIso(maxChangedAt)]);
          }
          await client.query("COMMIT");
        }

        console.log(`[incremental] ${org} / ${project}: since=${toIso(since)} scanned=${scanned} upserted=${links}`);
      } catch (e) {
        if (!DRY_RUN) {
          try { await client.query("ROLLBACK"); } catch {}
        }
        console.error(`[error] ${org} / ${project}: ${e.message || String(e)}`);
        throw e;
      } finally {
        if (!DRY_RUN) {
          try { await client.query("SELECT pg_advisory_unlock(hashtext($1))", [lockKey]); } catch {}
        }
      }

      await sleep(SLEEP_BETWEEN_PROJECTS_MS);
    }

    if (DRY_RUN) {
      const header = ["-- Dry run output", "BEGIN;"].join("\n");
      const footer = ["-- Cursor not advanced in dry run", "COMMIT;"].join("\n");
      const file = path.resolve(process.cwd(), OUT_FILE);
      fs.writeFileSync(file, header + "\n" + sqlLines.join("\n") + "\n" + footer + "\n", "utf8");
      console.log(`Dry run file: ${file}`);
    }
  } catch (err) {
    process.exitCode = 1;
  } finally {
    await client.end();
    const ended = new Date();
    console.log(`HTTP calls: ${httpCalls}`);
    console.log(`Run time: ${Math.round((ended - started) / 1000)}s`);
    console.log(`Total links processed: ${totalLinksProcessed}.`);
  }
}

main();