Discourse AI Custom Tool: BigQuery Dataset Schema

bigquery_schema_tool.js

/* eslint-disable no-undef, no-unused-vars */

// ---- Base64 ----

const B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

function b64Dec(s) {
  s = s.replace(/[\s=]/g, "");
  const o = [];
  let buf = 0, bits = 0;
  for (let i = 0; i < s.length; i++) {
    buf = (buf << 6) | B64.indexOf(s[i]);
    bits += 6;
    if (bits >= 8) { bits -= 8; o.push((buf >> bits) & 0xff); }
  }
  return o;
}

function b64Enc(b) {
  let s = "";
  for (let i = 0; i < b.length; i += 3) {
    const x = b[i], y = b[i + 1] || 0, z = b[i + 2] || 0;
    s += B64[x >> 2] + B64[((x & 3) << 4) | (y >> 4)];
    if (i + 1 < b.length) s += B64[((y & 15) << 2) | (z >> 6)];
    if (i + 2 < b.length) s += B64[z & 63];
  }
  return s.replace(/\+/g, "-").replace(/\//g, "_");
}

function str2b(s) {
  const b = [];
  for (let i = 0; i < s.length; i++) b.push(s.charCodeAt(i));
  return b;
}

// ---- SHA-256 ----

const SK = [
  0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
  0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
  0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
  0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
  0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
  0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
  0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
  0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2
];

function rr(x, n) { return ((x >>> n) | (x << (32 - n))) >>> 0; }

function sha256(msg) {
  const p = msg.slice();
  const bl = msg.length * 8;
  p.push(0x80);
  while (p.length % 64 !== 56) p.push(0);
  p.push(0, 0, 0, 0, (bl >>> 24) & 0xff, (bl >>> 16) & 0xff, (bl >>> 8) & 0xff, bl & 0xff);

  let h0 = 0x6a09e667, h1 = 0xbb67ae85, h2 = 0x3c6ef372, h3 = 0xa54ff53a;
  let h4 = 0x510e527f, h5 = 0x9b05688c, h6 = 0x1f83d9ab, h7 = 0x5be0cd19;

  for (let off = 0; off < p.length; off += 64) {
    const w = Array(64);
    for (let i = 0; i < 16; i++)
      w[i] = ((p[off+i*4] << 24) | (p[off+i*4+1] << 16) |
              (p[off+i*4+2] << 8) | p[off+i*4+3]) >>> 0;
    for (let i = 16; i < 64; i++)
      w[i] = (w[i-16] + (rr(w[i-15],7) ^ rr(w[i-15],18) ^ (w[i-15]>>>3)) +
              w[i-7] + (rr(w[i-2],17) ^ rr(w[i-2],19) ^ (w[i-2]>>>10))) >>> 0;

    let a = h0, b = h1, c = h2, d = h3, e = h4, f = h5, g = h6, hh = h7;
    for (let i = 0; i < 64; i++) {
      const t1 = (hh + (rr(e,6) ^ rr(e,11) ^ rr(e,25)) +
                  ((e & f) ^ (~e & g)) + SK[i] + w[i]) >>> 0;
      const t2 = ((rr(a,2) ^ rr(a,13) ^ rr(a,22)) +
                  ((a & b) ^ (a & c) ^ (b & c))) >>> 0;
      hh = g; g = f; f = e; e = (d + t1) >>> 0;
      d = c; c = b; b = a; a = (t1 + t2) >>> 0;
    }
    h0 = (h0+a)>>>0; h1 = (h1+b)>>>0; h2 = (h2+c)>>>0; h3 = (h3+d)>>>0;
    h4 = (h4+e)>>>0; h5 = (h5+f)>>>0; h6 = (h6+g)>>>0; h7 = (h7+hh)>>>0;
  }

  const out = [];
  for (const v of [h0, h1, h2, h3, h4, h5, h6, h7])
    out.push((v >>> 24) & 0xff, (v >>> 16) & 0xff, (v >>> 8) & 0xff, v & 0xff);
  return out;
}

// ---- ASN.1 DER parser (PKCS8 → RSA key) ----

function parsePKCS8(pem) {
  const der = b64Dec(pem.replace(/-----[^-]+-----/g, "").replace(/\s/g, ""));
  let p = 0;

  function readLen() {
    var l = der[p++];
    if (l & 0x80) { var nb = l & 0x7f; l = 0; for (var i = 0; i < nb; i++) l = (l << 8) | der[p++]; }
    return l;
  }
  function readSeq()  { p++; readLen(); }
  function readInt()  { p++; var l = readLen(); var v = der.slice(p, p + l); p += l; return v; }
  function readSkip() { p++; var l = readLen(); p += l; }

  readSeq();
  readInt();
  readSkip();
  p++; readLen();
  readSeq();
  readInt();
  var n = readInt();
  readInt();
  var d = readInt();
  return { n: n, d: d };
}

// ---- RSA PKCS1v15-SHA256 sign ----

function b2bi(bytes) {
  if (bytes.length === 0) return 0n;
  let h = "0x";
  for (let i = 0; i < bytes.length; i++) h += bytes[i].toString(16).padStart(2, "0");
  return BigInt(h);
}

function bi2b(n, len) {
  let h = n.toString(16);
  if (h.length & 1) h = "0" + h;
  while (h.length < len * 2) h = "00" + h;
  const o = [];
  for (let i = 0; i < h.length; i += 2) o.push(parseInt(h.substr(i, 2), 16));
  return o;
}

function modpow(base, exp, mod) {
  let r = 1n;
  base = base % mod;
  while (exp > 0n) {
    if (exp & 1n) r = r * base % mod;
    exp >>= 1n;
    base = base * base % mod;
  }
  return r;
}

const DI_PREFIX = [
  0x30,0x31,0x30,0x0d,0x06,0x09,0x60,0x86,0x48,0x01,
  0x65,0x03,0x04,0x02,0x01,0x05,0x00,0x04,0x20
];

function rsaSign(hash, key) {
  const nBI = b2bi(key.n), dBI = b2bi(key.d);
  const kLen = Math.ceil(nBI.toString(16).length / 2);
  const di = DI_PREFIX.concat(hash);
  const em = [0x00, 0x01];
  for (let i = 0; i < kLen - di.length - 3; i++) em.push(0xff);
  em.push(0x00, ...di);
  return bi2b(modpow(b2bi(em), dBI, nBI), kLen);
}

// ---- JWT + OAuth2 ----

function signJWT(cred) {
  const now = Math.floor(Date.now() / 1000);
  const hdr = b64Enc(str2b(JSON.stringify({ alg: "RS256", typ: "JWT" })));
  const pay = b64Enc(str2b(JSON.stringify({
    iss: cred.client_email,
    scope: "https://www.googleapis.com/auth/bigquery.readonly",
    aud: "https://oauth2.googleapis.com/token",
    iat: now,
    exp: now + 3600
  })));
  const msg = hdr + "." + pay;
  const sig = rsaSign(sha256(str2b(msg)), parsePKCS8(cred.private_key));
  return msg + "." + b64Enc(sig);
}

function getToken(cred) {
  const jwt = signJWT(cred);
  const r = http.post("https://oauth2.googleapis.com/token", {
    headers: { "Content-Type": "application/x-www-form-urlencoded" },
    body: "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&assertion=" + jwt
  });
  if (r.status !== 200) throw new Error("OAuth2 failed (" + r.status + "): " + r.body);
  return JSON.parse(r.body).access_token;
}

// ---- BigQuery Schema ----

function flattenFields(fields, prefix) {
  var out = [];
  for (var i = 0; i < fields.length; i++) {
    var f = fields[i];
    var name = prefix ? prefix + "." + f.name : f.name;
    var entry = { name: name, type: f.type };
    if (f.mode && f.mode !== "NULLABLE") entry.mode = f.mode;
    if (f.description) entry.description = f.description;
    out.push(entry);
    if (f.fields && f.fields.length > 0) {
      out = out.concat(flattenFields(f.fields, name));
    }
  }
  return out;
}

var lastDetail = "";

function invoke(params) {
  var dataset = (params.dataset || "").trim();
  if (!dataset) return { error: "dataset parameter is required" };

  var table = (params.table || "").trim();

  var cred = JSON.parse(secrets.get("json_cred"));
  var token = getToken(cred);
  var project = cred.project_id;

  var base = "https://bigquery.googleapis.com/bigquery/v2/projects/" +
             encodeURIComponent(project) + "/datasets/" +
             encodeURIComponent(dataset);

  if (table) {
    // Mode 2: single table — full schema detail (3 HTTP calls: oauth + get table)
    lastDetail = dataset + "." + table;

    var tr = http.get(base + "/tables/" + encodeURIComponent(table), {
      headers: { Authorization: "Bearer " + token }
    });
    if (tr.status !== 200) throw new Error("BigQuery (" + tr.status + "): " + tr.body);

    var tableInfo = JSON.parse(tr.body);
    var fields = (tableInfo.schema && tableInfo.schema.fields) || [];
    var columns = flattenFields(fields, "");

    var result = {
      project: project,
      dataset: dataset,
      table: table,
      type: tableInfo.type || "TABLE",
      columns: columns
    };
    if (tableInfo.numRows) result.rows = tableInfo.numRows;
    if (tableInfo.numBytes) result.size_bytes = tableInfo.numBytes;
    if (tableInfo.description) result.description = tableInfo.description;

    // Bound response size
    var json = JSON.stringify(result);
    if (json.length > 50000) {
      // Strip descriptions and nested fields
      result.columns = result.columns.filter(function(c) {
        return c.name.indexOf(".") === -1;
      }).map(function(c) {
        return { name: c.name, type: c.type };
      });
      result.note = "Schema trimmed — nested fields and descriptions removed to fit size limit";
    }

    return result;
  }

  // Mode 1: dataset overview — list tables with metadata (2 HTTP calls: oauth + list)
  lastDetail = dataset;

  var lr = http.get(base + "/tables?maxResults=200", {
    headers: { Authorization: "Bearer " + token }
  });
  if (lr.status !== 200) throw new Error("BigQuery (" + lr.status + "): " + lr.body);

  var body = JSON.parse(lr.body);
  var tables = (body.tables || []).map(function(t) {
    var entry = {
      table: t.tableReference.tableId,
      type: t.type || "TABLE"
    };
    if (t.labels) entry.labels = t.labels;
    return entry;
  });

  return {
    project: project,
    dataset: dataset,
    table_count: tables.length,
    tables: tables,
    hint: "Call again with table parameter to get full column schema for a specific table"
  };
}

function details() {
  return lastDetail ? "<code>" + lastDetail + "</code>" : "";
}