@dotdo/postgres 0.1.0

package/src/streaming/response-streaming.ts

@@ -0,0 +1,557 @@
+/**
+ * Response Streaming - HTTP Response streaming utilities
+ *
+ * Provides functions for converting async iterables to Web Streams API
+ * ReadableStream for use in HTTP responses.
+ *
+ * ## Supported Formats
+ *
+ * - **ReadableStream<T>**: Raw objects for custom processing
+ * - **NDJSON**: Newline-delimited JSON for efficient parsing
+ * - **SSE**: Server-Sent Events for real-time web clients
+ * - **JSON Array**: Traditional JSON array with streaming output
+ *
+ * ## Usage with Cloudflare Workers
+ *
+ * ```typescript
+ * export default {
+ *   async fetch(request: Request, env: Env): Promise<Response> {
+ *     const rows = streamer.stream('SELECT * FROM users')
+ *     return new Response(createNDJSONStream(rows), {
+ *       headers: {
+ *         'Content-Type': 'application/x-ndjson',
+ *         'Transfer-Encoding': 'chunked',
+ *       },
+ *     })
+ *   },
+ * }
+ * ```
+ *
+ * @module streaming/response-streaming
+ */
+
+import type { Row } from './query-streamer'
+
+// ============================================================================
+// Type Definitions
+// ============================================================================
+
+/**
+ * Options for creating readable streams.
+ */
+export interface ReadableStreamOptions {
+  /** AbortSignal for cancellation support */
+  signal?: AbortSignal
+  /** High water mark for internal buffer (default: 1) */
+  highWaterMark?: number
+}
+
+/**
+ * Options for NDJSON stream.
+ */
+export interface NDJSONStreamOptions extends ReadableStreamOptions {
+  /** Include metadata as first line (default: false) */
+  includeMetadata?: boolean
+  /** Metadata object to include as first line */
+  metadata?: Record<string, unknown>
+}
+
+// ============================================================================
+// Shared Utilities
+// ============================================================================
+
+/** Shared TextEncoder instance for all stream functions */
+const textEncoder = new TextEncoder()
+
+/**
+ * Create an abort-aware wrapper around an async iterator.
+ * @internal
+ */
+function createAbortAwareIterator<T>(
+  iterator: AsyncIterator<T>,
+  signal?: AbortSignal
+): { aborted: boolean; next: () => Promise<IteratorResult<T>> } {
+  let aborted = false
+
+  if (signal) {
+    signal.addEventListener('abort', () => {
+      aborted = true
+    }, { once: true })
+  }
+
+  return {
+    get aborted() { return aborted },
+    async next() {
+      if (aborted) {
+        return { value: undefined, done: true } as IteratorResult<T>
+      }
+      return iterator.next()
+    },
+  }
+}
+
+/**
+ * Safely close an async iterator if it has a return method.
+ * @internal
+ */
+function closeIterator(iterator: AsyncIterator<unknown>): void {
+  if (typeof iterator.return === 'function') {
+    iterator.return()
+  }
+}
+
+// ============================================================================
+// Stream Creation Functions
+// ============================================================================
+
+/**
+ * Create a ReadableStream from an async iterable of rows.
+ *
+ * This enables zero-buffering HTTP responses - rows are streamed
+ * directly to the response without buffering the entire result.
+ *
+ * @param source - Async iterable of rows
+ * @param options - Stream options
+ * @returns Web Streams API ReadableStream
+ *
+ * @example
+ * ```typescript
+ * const stream = createReadableStream(queryStream)
+ * return new Response(stream, {
+ *   headers: { 'Content-Type': 'application/json' }
+ * })
+ * ```
+ */
+export function createReadableStream<T extends Row = Row>(
+  source: AsyncIterable<T>,
+  options: ReadableStreamOptions = {}
+): ReadableStream<T> {
+  const rawIterator = source[Symbol.asyncIterator]()
+  const { aborted, next } = createAbortAwareIterator(rawIterator, options.signal)
+
+  return new ReadableStream<T>({
+    async pull(controller) {
+      if (aborted) {
+        controller.close()
+        return
+      }
+
+      try {
+        const result = await next()
+
+        if (result.done) {
+          controller.close()
+          return
+        }
+
+        controller.enqueue(result.value)
+      } catch (error) {
+        controller.error(error)
+      }
+    },
+
+    cancel() {
+      closeIterator(rawIterator)
+    },
+  }, {
+    highWaterMark: options.highWaterMark ?? 1,
+  })
+}
+
+/**
+ * Create a ReadableStream of Uint8Array chunks in NDJSON format.
+ *
+ * NDJSON (Newline Delimited JSON) is ideal for streaming responses
+ * as each line is a complete, parseable JSON object. This format is
+ * efficient for both parsing and memory usage.
+ *
+ * @param source - Async iterable of rows
+ * @param options - NDJSON stream options
+ * @returns Web Streams API ReadableStream of Uint8Array
+ *
+ * @example
+ * ```typescript
+ * const stream = createNDJSONStream(queryStream)
+ * return new Response(stream, {
+ *   headers: {
+ *     'Content-Type': 'application/x-ndjson',
+ *     'Transfer-Encoding': 'chunked'
+ *   }
+ * })
+ * ```
+ */
+export function createNDJSONStream<T extends Row = Row>(
+  source: AsyncIterable<T>,
+  options: NDJSONStreamOptions = {}
+): ReadableStream<Uint8Array> {
+  const rawIterator = source[Symbol.asyncIterator]()
+  const { aborted, next } = createAbortAwareIterator(rawIterator, options.signal)
+  let metadataSent = false
+
+  return new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      if (aborted) {
+        controller.close()
+        return
+      }
+
+      // Send metadata first if configured
+      if (options.includeMetadata && !metadataSent) {
+        metadataSent = true
+        const metaLine = JSON.stringify({ _meta: true, ...options.metadata }) + '\n'
+        controller.enqueue(textEncoder.encode(metaLine))
+        return
+      }
+
+      try {
+        const result = await next()
+
+        if (result.done) {
+          controller.close()
+          return
+        }
+
+        const line = JSON.stringify(result.value) + '\n'
+        controller.enqueue(textEncoder.encode(line))
+      } catch (error) {
+        controller.error(error)
+      }
+    },
+
+    cancel() {
+      closeIterator(rawIterator)
+    },
+  }, {
+    highWaterMark: options.highWaterMark ?? 1,
+  })
+}
+
+/**
+ * Create a Server-Sent Events (SSE) stream from an async iterable.
+ *
+ * SSE format allows real-time streaming to web clients with automatic
+ * reconnection and event typing support. Each row is sent as a "row" event,
+ * with a "done" event sent when the stream completes.
+ *
+ * Event types:
+ * - `row`: Contains a single row as JSON data
+ * - `done`: Signals stream completion
+ * - `error`: Contains error information if streaming fails
+ *
+ * @param source - Async iterable of rows
+ * @param options - Stream options
+ * @returns Web Streams API ReadableStream of Uint8Array
+ *
+ * @example
+ * ```typescript
+ * const stream = createSSEStream(queryStream)
+ * return new Response(stream, {
+ *   headers: {
+ *     'Content-Type': 'text/event-stream',
+ *     'Cache-Control': 'no-cache',
+ *     'Connection': 'keep-alive'
+ *   }
+ * })
+ * ```
+ */
+export function createSSEStream<T extends Row = Row>(
+  source: AsyncIterable<T>,
+  options: ReadableStreamOptions = {}
+): ReadableStream<Uint8Array> {
+  const rawIterator = source[Symbol.asyncIterator]()
+  const { aborted, next } = createAbortAwareIterator(rawIterator, options.signal)
+  let eventId = 0
+
+  return new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      if (aborted) {
+        // Send done event before closing on abort
+        controller.enqueue(textEncoder.encode(`event: done\ndata: {}\n\n`))
+        controller.close()
+        return
+      }
+
+      try {
+        const result = await next()
+
+        if (result.done) {
+          controller.enqueue(textEncoder.encode(`event: done\ndata: {}\n\n`))
+          controller.close()
+          return
+        }
+
+        eventId++
+        const sseMessage = `id: ${eventId}\nevent: row\ndata: ${JSON.stringify(result.value)}\n\n`
+        controller.enqueue(textEncoder.encode(sseMessage))
+      } catch (error) {
+        const errorMessage = error instanceof Error ? error.message : String(error)
+        const errorEvent = `event: error\ndata: ${JSON.stringify({ error: errorMessage })}\n\n`
+        controller.enqueue(textEncoder.encode(errorEvent))
+        controller.error(error)
+      }
+    },
+
+    cancel() {
+      closeIterator(rawIterator)
+    },
+  }, {
+    highWaterMark: options.highWaterMark ?? 1,
+  })
+}
+
+/**
+ * Create a JSON array stream that wraps rows in array syntax.
+ *
+ * Useful when clients expect a JSON array but you want to stream.
+ * Outputs valid JSON: `[{row1},{row2},...]`
+ *
+ * Note: While this produces valid JSON, clients cannot parse it
+ * incrementally. For better streaming support, prefer NDJSON.
+ *
+ * @param source - Async iterable of rows
+ * @param options - Stream options
+ * @returns Web Streams API ReadableStream of Uint8Array
+ *
+ * @example
+ * ```typescript
+ * const stream = createJSONArrayStream(queryStream)
+ * return new Response(stream, {
+ *   headers: { 'Content-Type': 'application/json' }
+ * })
+ * ```
+ */
+export function createJSONArrayStream<T extends Row = Row>(
+  source: AsyncIterable<T>,
+  options: ReadableStreamOptions = {}
+): ReadableStream<Uint8Array> {
+  const rawIterator = source[Symbol.asyncIterator]()
+  const { aborted, next } = createAbortAwareIterator(rawIterator, options.signal)
+  let firstItem = true
+  let started = false
+
+  return new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      // Start array on first pull
+      if (!started) {
+        started = true
+        controller.enqueue(textEncoder.encode('['))
+        return
+      }
+
+      if (aborted) {
+        controller.enqueue(textEncoder.encode(']'))
+        controller.close()
+        return
+      }
+
+      try {
+        const result = await next()
+
+        if (result.done) {
+          controller.enqueue(textEncoder.encode(']'))
+          controller.close()
+          return
+        }
+
+        const prefix = firstItem ? '' : ','
+        firstItem = false
+        controller.enqueue(textEncoder.encode(prefix + JSON.stringify(result.value)))
+      } catch (error) {
+        // Close the array on error to maintain valid JSON structure
+        controller.enqueue(textEncoder.encode(']'))
+        controller.error(error)
+      }
+    },
+
+    cancel() {
+      closeIterator(rawIterator)
+    },
+  }, {
+    highWaterMark: options.highWaterMark ?? 1,
+  })
+}
+
+/**
+ * StreamResult interface for fluent streaming API
+ */
+export interface StreamResult<T extends Row = Row> extends AsyncIterable<T> {
+  readonly sql: string
+  readonly params: unknown[]
+  batch(size: number): AsyncIterable<T[]>
+  cursor(): Promise<{ fetch: (count: number) => Promise<T[]>; close: () => Promise<void> }>
+  toReadableStream(): ReadableStream<T>
+  toNDJSONStream(): ReadableStream<Uint8Array>
+  toSSEStream(): ReadableStream<Uint8Array>
+  toJSONArrayStream(): ReadableStream<Uint8Array>
+}
+
+/**
+ * Create a StreamResult from an async generator
+ *
+ * Provides a fluent API for streaming query results to various formats.
+ *
+ * @param source - Async iterable of rows
+ * @param sql - Original SQL query
+ * @param params - Query parameters
+ * @returns StreamResult with conversion methods
+ */
+export function createStreamResult<T extends Row = Row>(
+  source: AsyncIterable<T>,
+  sql: string,
+  params: unknown[] = []
+): StreamResult<T> {
+  // Cache the source as an array for reuse
+  let cachedRows: T[] | null = null
+
+  async function* getSource(): AsyncIterable<T> {
+    if (cachedRows) {
+      for (const row of cachedRows) {
+        yield row
+      }
+      return
+    }
+
+    const rows: T[] = []
+    for await (const row of source) {
+      rows.push(row)
+      yield row
+    }
+    cachedRows = rows
+  }
+
+  return {
+    sql,
+    params,
+
+    async *[Symbol.asyncIterator]() {
+      yield* getSource()
+    },
+
+    batch(size: number): AsyncIterable<T[]> {
+      const src = getSource()
+      return {
+        async *[Symbol.asyncIterator]() {
+          let batch: T[] = []
+          for await (const row of src) {
+            batch.push(row)
+            if (batch.length >= size) {
+              yield batch
+              batch = []
+            }
+          }
+          if (batch.length > 0) {
+            yield batch
+          }
+        },
+      }
+    },
+
+    async cursor() {
+      const rows = cachedRows || []
+      if (!cachedRows) {
+        for await (const row of source) {
+          rows.push(row)
+        }
+        cachedRows = rows
+      }
+
+      let position = 0
+      return {
+        async fetch(count: number): Promise<T[]> {
+          const result = rows.slice(position, position + count)
+          position += result.length
+          return result
+        },
+        async close() {
+          position = rows.length
+        },
+      }
+    },
+
+    toReadableStream() {
+      return createReadableStream(getSource())
+    },
+
+    toNDJSONStream() {
+      return createNDJSONStream(getSource())
+    },
+
+    toSSEStream() {
+      return createSSEStream(getSource())
+    },
+
+    toJSONArrayStream() {
+      return createJSONArrayStream(getSource())
+    },
+  }
+}
+
+/**
+ * Pipe an async iterable through a transform
+ *
+ * @param source - Source async iterable
+ * @param transform - Transform function
+ * @returns Transformed async iterable
+ */
+export async function* pipeThrough<T, U>(
+  source: AsyncIterable<T>,
+  transform: (item: T) => U | Promise<U>
+): AsyncGenerator<U> {
+  for await (const item of source) {
+    yield await transform(item)
+  }
+}
+
+/**
+ * Filter an async iterable
+ *
+ * @param source - Source async iterable
+ * @param predicate - Filter predicate
+ * @returns Filtered async iterable
+ */
+export async function* filter<T>(
+  source: AsyncIterable<T>,
+  predicate: (item: T) => boolean | Promise<boolean>
+): AsyncGenerator<T> {
+  for await (const item of source) {
+    if (await predicate(item)) {
+      yield item
+    }
+  }
+}
+
+/**
+ * Take first N items from an async iterable
+ *
+ * @param source - Source async iterable
+ * @param count - Number of items to take
+ * @returns Limited async iterable
+ */
+export async function* take<T>(source: AsyncIterable<T>, count: number): AsyncGenerator<T> {
+  if (count <= 0) return
+
+  let taken = 0
+  for await (const item of source) {
+    yield item
+    taken++
+    if (taken >= count) break
+  }
+}
+
+/**
+ * Skip first N items from an async iterable
+ *
+ * @param source - Source async iterable
+ * @param count - Number of items to skip
+ * @returns Async iterable without first N items
+ */
+export async function* skip<T>(source: AsyncIterable<T>, count: number): AsyncGenerator<T> {
+  let skipped = 0
+  for await (const item of source) {
+    if (skipped < count) {
+      skipped++
+      continue
+    }
+    yield item
+  }
+}