HiveGrid-Wireframes/hivegrid-cloud/apps/worker/src/jobs/export.ts

import type { Writable } from "node:stream";

import type { Upload } from "@aws-sdk/lib-storage";
import { S3Client } from "@aws-sdk/client-s3";
import type { EnvRecord } from "@toolctrl/config";
import { serializeExportTableToCsvStream } from "@toolctrl/exports/csv";
import { serializeExportTableToJsonStream } from "@toolctrl/exports/json";
import { ExportRowCapExceededError, assertExportRowCap, resolveExportRowCap, type ExportFormat, type TenantPlanTier } from "@toolctrl/exports/limits";
import { serializeExportTableToPdfStream } from "@toolctrl/exports/pdf";
import { createS3MultipartUploadSink } from "@toolctrl/exports/s3-multipart";
import type { ExportRecord } from "@toolctrl/exports/types";
import { serializeExportTableToXlsxStream } from "@toolctrl/exports/xlsx";
import type { SharedStore } from "@toolctrl/shared-store";
import type { ExportJob } from "@toolctrl/queue";

import { workerLogger } from "../logger.js";

export interface WorkerExportRunResult {
  artifactBucket: string;
  artifactKey: string;
  rowCount: number;
  expiresAt: Date;
}

/**
 * Async row-source contract — implementations stream the rows of the
 * target export in chunks (typically a Prisma cursor) without buffering
 * the full dataset. The job pipeline orchestrates serialisation,
 * progress-tracking and S3 multipart upload around this iterator.
 */
export interface ExportRowSourceContext {
  payload: ExportJob;
  plan: TenantPlanTier;
  /** Effective row cap (min of plan and format hard cap). */
  effectiveRowCap: number;
}

export interface ExportRowSource {
  /** Best-effort row-count estimate, used for XLSX strategy selection. */
  estimateRowCount(context: ExportRowSourceContext): Promise<number | undefined>;
  /** Column names in export order. */
  resolveColumns(context: ExportRowSourceContext): Promise<readonly string[]>;
  /** Async iterator over the row stream. */
  iterateRows(context: ExportRowSourceContext): AsyncIterable<ExportRecord>;
}

export interface ExportTenantPlanResolver {
  (tenantId: string): Promise<TenantPlanTier>;
}

export interface ExportArtifactKeyBuilder {
  (payload: ExportJob): { bucket: string; key: string; contentType: string; expiresAt: Date };
}

export interface ExportJobDependencies {
  s3Client: S3Client;
  sharedStore?: SharedStore;
  rowSource: ExportRowSource;
  resolveTenantPlan: ExportTenantPlanResolver;
  buildArtifactKey: ExportArtifactKeyBuilder;
  /** Test seam — overrides the per-format streamer registry. */
  streamerRegistry?: ExportStreamerRegistry;
  /** Progress-publish frequency in rows; defaults to 5_000. */
  progressFlushInterval?: number;
  /** TTL applied to the progress key in seconds; defaults to 3600 (1h). */
  progressTtlSeconds?: number;
  /** Test seam — overrides the @aws-sdk/lib-storage Upload class. */
  uploadCtor?: typeof Upload;
}

export type ExportFormatNarrow = ExportFormat;

export interface ExportStreamerRegistry {
  csv: (writeStream: Writable, columns: readonly string[], rows: AsyncIterable<ExportRecord>, estimatedRowCount: number | undefined) => Promise<{ rowCount: number }>;
  json: (writeStream: Writable, columns: readonly string[], rows: AsyncIterable<ExportRecord>, estimatedRowCount: number | undefined) => Promise<{ rowCount: number }>;
  xlsx: (writeStream: Writable, columns: readonly string[], rows: AsyncIterable<ExportRecord>, estimatedRowCount: number | undefined) => Promise<{ rowCount: number }>;
  pdf: (writeStream: Writable, columns: readonly string[], rows: AsyncIterable<ExportRecord>, estimatedRowCount: number | undefined) => Promise<{ rowCount: number }>;
}

const DEFAULT_REGISTRY: ExportStreamerRegistry = {
  csv: async (writeStream, columns, rows) =>
    serializeExportTableToCsvStream(writeStream, { columns: [...columns], rows }),
  json: async (writeStream, columns, rows) =>
    serializeExportTableToJsonStream(writeStream, { columns: [...columns], rows }),
  xlsx: async (writeStream, columns, rows, estimatedRowCount) =>
    serializeExportTableToXlsxStream(writeStream, {
      columns: [...columns],
      rows,
      ...(estimatedRowCount !== undefined ? { estimatedRowCount } : {}),
    }),
  pdf: async (writeStream, columns, rows) => {
    // PDF needs a buffered table because react-pdf renders the document
    // tree synchronously; we cap below the format hard limit so heap
    // usage stays bounded.
    const buffered: ExportRecord[] = [];
    for await (const row of rows) {
      buffered.push(row);
    }
    return serializeExportTableToPdfStream(writeStream, { columns: [...columns], rows: buffered });
  },
};

export function createExportJobRunner(deps: ExportJobDependencies): (payload: ExportJob) => Promise<WorkerExportRunResult> {
  const flushInterval = deps.progressFlushInterval ?? 5_000;
  const progressTtl = deps.progressTtlSeconds ?? 3_600;
  const registry = deps.streamerRegistry ?? DEFAULT_REGISTRY;

  return async function runExport(payload: ExportJob): Promise<WorkerExportRunResult> {
    const format = normaliseFormat(payload.format);
    const plan = await deps.resolveTenantPlan(payload.tenantId);
    const effectiveRowCap = resolveExportRowCap(plan, format);
    const context: ExportRowSourceContext = { payload, plan, effectiveRowCap };

    const estimate = await deps.rowSource.estimateRowCount(context);
    if (estimate !== undefined) {
      assertExportRowCap(estimate, plan, format);
    }

    const columns = await deps.rowSource.resolveColumns(context);
    const artifact = deps.buildArtifactKey(payload);
    const { writeStream, completionPromise } = createS3MultipartUploadSink({
      client: deps.s3Client,
      bucket: artifact.bucket,
      key: artifact.key,
      contentType: artifact.contentType,
      metadata: { tenantId: payload.tenantId, exportJobId: payload.exportJobId, exportType: payload.exportType },
      ...(deps.uploadCtor ? { uploadCtor: deps.uploadCtor } : {}),
    });

    let processedRows = 0;
    const guardedRows = (async function* () {
      for await (const row of deps.rowSource.iterateRows(context)) {
        if (processedRows >= effectiveRowCap) {
          throw new ExportRowCapExceededError(processedRows + 1, effectiveRowCap, plan, format);
        }
        processedRows += 1;
        if (deps.sharedStore && processedRows % flushInterval === 0) {
          // Best-effort progress publish — failure here must not abort
          // the export so we swallow + log only.
          await publishProgress(deps.sharedStore, payload.exportJobId, processedRows, estimate, progressTtl);
        }
        yield row;
      }
    })();

    try {
      const result = await registry[format](writeStream, columns, guardedRows, estimate);
      writeStream.end();
      const completion = await completionPromise;

      if (deps.sharedStore) {
        await publishProgress(deps.sharedStore, payload.exportJobId, result.rowCount, estimate, progressTtl);
      }

      return {
        artifactBucket: completion.bucket,
        artifactKey: completion.key,
        rowCount: result.rowCount,
        expiresAt: artifact.expiresAt,
      };
    } catch (error) {
      writeStream.destroy(error instanceof Error ? error : new Error(String(error)));
      throw error;
    }
  };
}

async function publishProgress(
  store: SharedStore,
  exportJobId: string,
  processedRows: number,
  totalEstimate: number | undefined,
  ttlSeconds: number,
): Promise<void> {
  try {
    await store.set(
      `export:progress:${exportJobId}`,
      JSON.stringify({ processedRows, totalEstimate: totalEstimate ?? null, updatedAt: new Date().toISOString() }),
      { ttlSeconds },
    );
  } catch (error) {
    workerLogger.warn({ error, exportJobId }, "Failed to publish export progress to shared store.");
  }
}

function normaliseFormat(format: ExportJob["format"]): ExportFormat {
  switch (format) {
    case "csv":
    case "json":
      return format;
    default: {
      // The queue payload only declares csv/json today, but the worker is
      // ready for xlsx/pdf as soon as the producer side widens the union.
      const widened = format as unknown as ExportFormat;
      if (widened === "xlsx" || widened === "pdf") return widened;
      throw new Error(`Unsupported export format: ${String(format)}`);
    }
  }
}

/**
 * Resolve the S3Client used by the worker for export multipart uploads.
 * Mirrors the scanService env shape so operators configure a single set of
 * STORAGE_* variables for both the scanner and the exporter.
 */
export function createWorkerExportS3Client(env: EnvRecord = process.env as EnvRecord): S3Client | undefined {
  if (!env.STORAGE_REGION || !env.STORAGE_ACCESS_KEY_ID || !env.STORAGE_SECRET_ACCESS_KEY) {
    return undefined;
  }
  return new S3Client({
    endpoint: env.STORAGE_ENDPOINT,
    region: env.STORAGE_REGION,
    forcePathStyle: env.STORAGE_FORCE_PATH_STYLE === "true",
    credentials: {
      accessKeyId: env.STORAGE_ACCESS_KEY_ID,
      secretAccessKey: env.STORAGE_SECRET_ACCESS_KEY,
      ...(env.STORAGE_SESSION_TOKEN ? { sessionToken: env.STORAGE_SESSION_TOKEN } : {}),
    },
  });
}