chore: import upstream snapshot with attribution

apps/sim/lib/data-drains/destinations/bigquery.ts

@@ -0,0 +1,333 @@
+import { createHash } from 'node:crypto'
+import { createLogger } from '@sim/logger'
+import { toError } from '@sim/utils/errors'
+import { backoffWithJitter, parseRetryAfter } from '@sim/utils/retry'
+import { JWT } from 'google-auth-library'
+import { z } from 'zod'
+import {
+  type ParsedServiceAccount,
+  parseNdjsonObjects,
+  parseServiceAccount,
+  refineServiceAccountJson,
+  sleepUntilAborted,
+} from '@/lib/data-drains/destinations/utils'
+import type { DeliveryMetadata, DrainDestination } from '@/lib/data-drains/types'
+
+const logger = createLogger('DataDrainBigQueryDestination')
+
+/**
+ * Uses the legacy `tabledata.insertAll` streaming endpoint. The Storage Write
+ * API offers exactly-once semantics and lower pricing but requires gRPC; we
+ * stay on insertAll for simplicity and direct HTTP support.
+ */
+
+/** `insertdata` for streaming inserts; `readonly` for the `tables.get` probe in `test()`. */
+const SCOPES = [
+  'https://www.googleapis.com/auth/bigquery.insertdata',
+  'https://www.googleapis.com/auth/bigquery.readonly',
+]
+
+/** Standard project IDs are 6-30 chars; the optional `domain.tld:` prefix supports legacy domain-scoped projects. */
+const PROJECT_ID_RE = /^([a-z][a-z0-9.-]{0,61}[a-z0-9]:)?[a-z][a-z0-9-]{4,28}[a-z0-9]$/
+const DATASET_RE = /^[A-Za-z0-9_]{1,1024}$/
+const TABLE_RE = /^[\p{L}\p{M}\p{N}\p{Pc}\p{Pd} ]{1,1024}$/u
+
+const USER_AGENT = 'sim-data-drain/1.0'
+/** Per-request streaming limits: 10 MB body, 50,000 rows, 1 MB per row. */
+const MAX_REQUEST_BYTES = 10 * 1024 * 1024
+const MAX_ROWS_PER_REQUEST = 50_000
+const MAX_ROW_BYTES = 1024 * 1024
+/** `insertId` is capped at 128 characters (encoded length). */
+const MAX_INSERT_ID_LENGTH = 128
+const PER_ATTEMPT_TIMEOUT_MS = 60_000
+
+const bigqueryConfigSchema = z.object({
+  projectId: z
+    .string()
+    .min(6, 'projectId is required')
+    .refine((v) => PROJECT_ID_RE.test(v), {
+      message: 'projectId must match Google Cloud project ID rules',
+    }),
+  datasetId: z
+    .string()
+    .min(1, 'datasetId is required')
+    .refine((v) => DATASET_RE.test(v), {
+      message: 'datasetId may only contain ASCII letters, digits, and underscores (max 1024 chars)',
+    }),
+  tableId: z
+    .string()
+    .min(1, 'tableId is required')
+    .refine((v) => TABLE_RE.test(v), {
+      message:
+        'tableId may only contain Unicode letters/marks/numbers, connectors, dashes, and spaces (max 1024 chars)',
+    })
+    .refine((v) => Buffer.byteLength(v, 'utf8') <= 1024, {
+      message: 'tableId must be at most 1024 bytes when UTF-8 encoded',
+    }),
+})
+
+const bigqueryCredentialsSchema = z
+  .object({
+    serviceAccountJson: z.string().min(1, 'serviceAccountJson is required'),
+  })
+  .superRefine(refineServiceAccountJson)
+
+export type BigQueryDestinationConfig = z.infer<typeof bigqueryConfigSchema>
+export type BigQueryDestinationCredentials = z.infer<typeof bigqueryCredentialsSchema>
+
+function buildJwt(account: ParsedServiceAccount): JWT {
+  return new JWT({ email: account.clientEmail, key: account.privateKey, scopes: SCOPES })
+}
+
+async function getAccessToken(jwt: JWT, forceRefresh = false): Promise<string> {
+  if (forceRefresh) {
+    /** Clearing `credentials` forces `getAccessToken` to mint a new token instead of returning the cached one. */
+    jwt.credentials = {}
+  }
+  const { token } = await jwt.getAccessToken()
+  if (!token) throw new Error('Failed to obtain BigQuery access token')
+  return token
+}
+
+interface InsertAllInput {
+  config: BigQueryDestinationConfig
+  rows: Record<string, unknown>[]
+  metadata: DeliveryMetadata
+  jwt: JWT
+  signal: AbortSignal
+}
+
+interface InsertAllError {
+  index: number
+  errors: Array<{ reason?: string; message?: string; location?: string }>
+}
+
+async function postInsertAll(
+  input: InsertAllInput,
+  url: string,
+  body: string,
+  forceRefresh = false
+): Promise<Response> {
+  const token = await getAccessToken(input.jwt, forceRefresh)
+  const perAttempt = AbortSignal.any([input.signal, AbortSignal.timeout(PER_ATTEMPT_TIMEOUT_MS)])
+  try {
+    return await fetch(url, {
+      method: 'POST',
+      headers: {
+        Authorization: `Bearer ${token}`,
+        'Content-Type': 'application/json',
+        'User-Agent': USER_AGENT,
+      },
+      body,
+      signal: perAttempt,
+    })
+  } catch (error) {
+    logger.warn('BigQuery request failed', {
+      table: `${input.config.projectId}.${input.config.datasetId}.${input.config.tableId}`,
+      error: toError(error).message,
+    })
+    throw error
+  }
+}
+
+/**
+ * Builds a stable `insertId` for best-effort dedup (~60s window). Prefixed
+ * with `drainId` so (runId, sequence) collisions across drains do not
+ * accidentally dedupe each other's rows. With UUID drain/run IDs the raw
+ * form fits well under 128 chars; if anything pushes it over (e.g. a future
+ * non-UUID id), hash the prefix and keep the row-distinguishing `index`
+ * suffix intact so BigQuery does not silently dedupe distinct rows.
+ */
+function buildInsertId(metadata: DeliveryMetadata, index: number): string {
+  const raw = `${metadata.drainId}-${metadata.runId}-${metadata.sequence}-${index}`
+  if (raw.length <= MAX_INSERT_ID_LENGTH) return raw
+  const prefixHash = createHash('sha1')
+    .update(`${metadata.drainId}-${metadata.runId}-${metadata.sequence}`)
+    .digest('hex')
+  return `${prefixHash}-${index}`
+}
+
+const RETRYABLE_STATUSES = new Set([408, 429, 500, 502, 503, 504])
+const MAX_RETRY_ATTEMPTS = 3
+const BASE_RETRY_DELAY_MS = 250
+
+/**
+ * Streams a chunk of rows to `tabledata.insertAll`.
+ *
+ * Partial-success caveat: BigQuery may return HTTP 200 with a non-empty
+ * `insertErrors` array. Rows not listed there are inserted and dedup-keyed by
+ * `insertId` for ~60s. We throw on any `insertErrors`; retries within the
+ * dedup window are safe, but retries after it may duplicate succeeded rows.
+ */
+async function insertAll(input: InsertAllInput): Promise<void> {
+  if (input.rows.length > MAX_ROWS_PER_REQUEST) {
+    throw new Error(
+      `BigQuery insertAll chunk has ${input.rows.length} rows, exceeds the ${MAX_ROWS_PER_REQUEST} per-request limit`
+    )
+  }
+  const url = `https://bigquery.googleapis.com/bigquery/v2/projects/${encodeURIComponent(input.config.projectId)}/datasets/${encodeURIComponent(input.config.datasetId)}/tables/${encodeURIComponent(input.config.tableId)}/insertAll`
+  /**
+   * `skipInvalidRows: false` and `ignoreUnknownValues: false` surface schema
+   * mismatches as `insertErrors` instead of silently dropping data — drains
+   * should fail loudly so operators notice the schema drift.
+   */
+  const payload = {
+    skipInvalidRows: false,
+    ignoreUnknownValues: false,
+    rows: input.rows.map((row, index) => {
+      const rowBytes = Buffer.byteLength(JSON.stringify(row), 'utf8')
+      if (rowBytes > MAX_ROW_BYTES) {
+        throw new Error(
+          `BigQuery row at index ${index} is ${rowBytes} bytes, exceeds the ${MAX_ROW_BYTES}-byte per-row limit`
+        )
+      }
+      return {
+        insertId: buildInsertId(input.metadata, index),
+        json: row,
+      }
+    }),
+  }
+  const body = JSON.stringify(payload)
+  const byteLength = Buffer.byteLength(body, 'utf8')
+  if (byteLength > MAX_REQUEST_BYTES) {
+    throw new Error(
+      `BigQuery insertAll body is ${byteLength} bytes, exceeds the ${MAX_REQUEST_BYTES}-byte per-request limit`
+    )
+  }
+  let attempt = 0
+  let response: Response | undefined
+  let refreshedOnce = false
+  while (true) {
+    attempt++
+    try {
+      response = await postInsertAll(input, url, body)
+      /** A 401 retry doesn't count against the 5xx/429 budget — token refresh is a one-shot recovery. */
+      if (response.status === 401 && !refreshedOnce) {
+        refreshedOnce = true
+        logger.debug('BigQuery returned 401; refreshing access token and retrying once')
+        /** Drain the 401 body before discarding so undici can return the socket to the keep-alive pool. */
+        await response.text().catch(() => '')
+        response = await postInsertAll(input, url, body, true)
+      }
+      if (!RETRYABLE_STATUSES.has(response.status)) break
+      if (attempt >= MAX_RETRY_ATTEMPTS) break
+      const retryAfterHeaderMs = parseRetryAfter(response.headers.get('retry-after'))
+      const retryAfterMs = backoffWithJitter(attempt, retryAfterHeaderMs, {
+        baseMs: BASE_RETRY_DELAY_MS,
+      })
+      logger.warn('BigQuery insertAll transient error; retrying', {
+        status: response.status,
+        attempt,
+        retryAfterMs,
+      })
+      /** Drain the body so the keep-alive connection can be reused. */
+      await response.text().catch(() => '')
+      await sleepUntilAborted(retryAfterMs, input.signal)
+      if (input.signal.aborted) throw input.signal.reason ?? new Error('Aborted')
+    } catch (error) {
+      /**
+       * Connection-level failures (DNS, socket reset, timeout) never produce
+       * a Response — treat them like 5xx and retry with backoff. Re-throw
+       * aborts unwrapped so callers see the cancellation reason.
+       */
+      if (input.signal.aborted) throw input.signal.reason ?? error
+      if (attempt >= MAX_RETRY_ATTEMPTS) throw error
+      const retryAfterMs = backoffWithJitter(attempt, null, { baseMs: BASE_RETRY_DELAY_MS })
+      logger.warn('BigQuery insertAll network error; retrying', {
+        attempt,
+        retryAfterMs,
+        error: toError(error).message,
+      })
+      await sleepUntilAborted(retryAfterMs, input.signal)
+      if (input.signal.aborted) throw input.signal.reason ?? new Error('Aborted')
+    }
+  }
+  if (!response) throw new Error('BigQuery insertAll failed: no response')
+  if (!response.ok) {
+    const text = await response.text().catch(() => '')
+    throw new Error(`BigQuery insertAll failed (HTTP ${response.status}): ${text}`)
+  }
+  const result = (await response.json().catch(() => ({}))) as {
+    insertErrors?: InsertAllError[]
+  }
+  if (result.insertErrors && result.insertErrors.length > 0) {
+    const failedIndices = result.insertErrors.map((e) => e.index)
+    const total = input.rows.length
+    const failed = result.insertErrors.length
+    const succeeded = total - failed
+    logger.warn('BigQuery insertAll returned partial failure', {
+      partialFailure: true,
+      table: `${input.config.projectId}.${input.config.datasetId}.${input.config.tableId}`,
+      succeededRows: succeeded,
+      failedRows: failed,
+      failedIndices: failedIndices.slice(0, 20),
+    })
+    const summary = result.insertErrors
+      .slice(0, 3)
+      .map(
+        (e) =>
+          `row ${e.index}: ${e.errors.map((er) => er.message ?? er.reason ?? 'unknown').join('; ')}`
+      )
+      .join(' | ')
+    throw new Error(
+      `BigQuery insertAll partial failure: ${failed} of ${total} rows failed (indices: ${failedIndices.slice(0, 20).join(',')}${failedIndices.length > 20 ? ',...' : ''}); ${succeeded} rows were inserted and are dedup-keyed by insertId for ~60s — retries within that window are safe, but retries after the window may duplicate the succeeded rows. First errors: ${summary}`
+    )
+  }
+}
+
+export const bigqueryDestination: DrainDestination<
+  BigQueryDestinationConfig,
+  BigQueryDestinationCredentials
+> = {
+  type: 'bigquery',
+  displayName: 'Google BigQuery',
+  configSchema: bigqueryConfigSchema,
+  credentialsSchema: bigqueryCredentialsSchema,
+
+  /**
+   * Probes table existence, IAM access, and credential validity in a single
+   * `tables.get` call. `fields=id` minimises response size — we only care
+   * whether the call succeeds, not the payload.
+   */
+  async test({ config, credentials, signal }) {
+    const account = parseServiceAccount(credentials.serviceAccountJson)
+    const jwt = buildJwt(account)
+    const token = await getAccessToken(jwt)
+    const url = `https://bigquery.googleapis.com/bigquery/v2/projects/${encodeURIComponent(config.projectId)}/datasets/${encodeURIComponent(config.datasetId)}/tables/${encodeURIComponent(config.tableId)}?fields=id`
+    const perAttempt = AbortSignal.any([signal, AbortSignal.timeout(PER_ATTEMPT_TIMEOUT_MS)])
+    const response = await fetch(url, {
+      headers: { Authorization: `Bearer ${token}` },
+      signal: perAttempt,
+    })
+    if (!response.ok) {
+      const text = await response.text().catch(() => '')
+      throw new Error(`BigQuery probe failed (HTTP ${response.status}): ${text}`)
+    }
+    /** Drain the success body so undici can return the socket to the keep-alive pool. */
+    await response.text().catch(() => '')
+  },
+
+  openSession({ config, credentials }) {
+    const account = parseServiceAccount(credentials.serviceAccountJson)
+    const jwt = buildJwt(account)
+    return {
+      async deliver({ body, metadata, signal }) {
+        const rows = parseNdjsonObjects(body, { requireObject: true }) as Record<string, unknown>[]
+        if (rows.length === 0) {
+          return {
+            locator: `bigquery://${config.projectId}/${config.datasetId}/${config.tableId}#${metadata.runId}-${metadata.sequence}`,
+          }
+        }
+        await insertAll({ config, rows, metadata, jwt, signal })
+        logger.debug('BigQuery chunk delivered', {
+          table: `${config.projectId}.${config.datasetId}.${config.tableId}`,
+          rows: rows.length,
+        })
+        return {
+          locator: `bigquery://${config.projectId}/${config.datasetId}/${config.tableId}#${metadata.runId}-${metadata.sequence}`,
+        }
+      },
+      async close() {},
+    }
+  },
+}