Stream an export and clean up failure

mkdir -p examples
cp -R node_modules/@db3.ai/app/src/storage/examples/. examples/

npx tsx examples/runStreamExport.ts

import { randomUUID } from 'node:crypto';
import { Readable } from 'node:stream';
import type { StorageDisk } from '@db3.ai/app/storage';

/** Options for a bounded-memory CSV export and deliberate failure in the lab. */
export interface CsvExportOptions {
	/** Number of synthetic rows to generate without retaining all rows in memory. */
	rows: number;
	/** Lab-only failure point; omit for a normal export. */
	failAt?: number;
}

/**
 * Streams a generated CSV to a staging path, publishing only after success.
 *
 * The staging path is owned by this operation. A failed stream removes that
 * object and leaves the previous destination untouched. Cross-provider moves
 * are not assumed atomic; concurrent publication needs an application policy.
 *
 * @param disk - Application-owned disk selected before the operation.
 * @param destination - Authorized relative output path, not raw request input.
 * @param options - Row count and optional controlled failure point.
 */
export async function writeCsvExport(disk: StorageDisk, destination: string, options: CsvExportOptions): Promise<void> {
	if (!Number.isInteger(options.rows) || options.rows < 0 || options.rows > 1_000_000) throw new Error('rows must be an integer between 0 and 1000000');
	const staging = `pending/${randomUUID()}.csv`;
	const stream = Readable.from(csvRows(options));
	try {
		await disk.writeStream(staging, stream, { mimeType: 'text/csv', visibility: 'private' });
		await disk.move(staging, destination);
	} catch (error) {
		stream.destroy();
		if (await disk.exists(staging)) await disk.delete(staging);
		throw error;
	}
}

/** Produces one safe synthetic CSV row at a time; real exports must escape data. */
function* csvRows(options: CsvExportOptions): Generator<string> {
	yield 'id,value\n';
	for (let index = 0; index < options.rows; index++) {
		if (index === options.failAt) throw new Error('Controlled export interruption');
		yield `${index},note-${index}\n`;
	}
}

import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { Writable } from 'node:stream';
import { pipeline } from 'node:stream/promises';
import { pathToFileURL } from 'node:url';
import { Storage } from '@db3.ai/app/storage';
import { writeCsvExport } from './writeCsvExport';

/**
 * Exercises streamed file publication, interrupted-write cleanup and retry.
 * @returns Counts and cleanup outcomes without buffering the complete export.
 */
export async function runStreamExport() {
	const root = await mkdtemp(join(tmpdir(), 'db3-stream-guide-'));
	const disk = new Storage({ disks: { local: { driver: 'local', root } } }).disk('local');
	try {
		await writeCsvExport(disk, 'exports/latest.csv', { rows: 50_000 });
		const sizeBefore = await disk.size('exports/latest.csv');
		let interruptionRejected = false;
		try { await writeCsvExport(disk, 'exports/latest.csv', { rows: 10, failAt: 2 }); } catch (error) {
			// Drivers may wrap failures while preserving the original stream error.
			const cause = error instanceof Error && error.cause instanceof Error ? error.cause : error;
			if (!(cause instanceof Error) || cause.message !== 'Controlled export interruption') throw error;
			interruptionRejected = true;
		}
		const previousPreserved = await disk.size('exports/latest.csv') === sizeBefore;
		const pendingFiles = (await disk.list('pending').toArray()).filter(entry => entry.isFile).length;
		await writeCsvExport(disk, 'exports/repaired.csv', { rows: 3 });
		let lines = 0;
		let bytes = 0;
		const sink = new Writable({
			/** Counts each chunk without retaining the complete file. */
			write(chunk: Buffer, _encoding, callback) { bytes += chunk.length; for (const byte of chunk) if (byte === 10) lines++; callback(); },
		});
		await pipeline(await disk.readStream('exports/latest.csv'), sink);
		return { lines, complete: bytes === sizeBefore, interruptionRejected, previousPreserved, pendingFiles, repaired: await disk.exists('exports/repaired.csv') };
	} finally { await rm(root, { recursive: true, force: true }); }
}

if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) console.log(JSON.stringify(await runStreamExport(), null, 2));

import { expect, it } from 'vitest';
import { runStreamExport } from '../../examples/runStreamExport';

it('streams a large export, preserves the previous file on interruption and cleans partial bytes', async () => {
	expect(await runStreamExport()).toEqual({ lines: 50_001, complete: true, interruptionRejected: true, previousPreserved: true, pendingFiles: 0, repaired: true });
});

npx vitest run tests/storage/runStreamExport.test.ts
npx tsc --noEmit --target ES2022 --module ESNext --moduleResolution Bundler --types node --skipLibCheck examples/*.ts