@electric-sql/client 0.5.1 → 0.6.1

package/src/client.ts

@@ -1,12 +1,19 @@
-import { Message, Offset, Schema, Row, MaybePromise } from './types'
+import {
+  Message,
+  Offset,
+  Schema,
+  Row,
+  MaybePromise,
+  GetExtensions,
+} from './types'
 import { MessageParser, Parser } from './parser'
 import { isUpToDateMessage } from './helpers'
-import { MessageProcessor, MessageProcessorInterface } from './queue'
 import { FetchError, FetchBackoffAbortError } from './error'
 import {
   BackoffDefaults,
   BackoffOptions,
   createFetchWithBackoff,
+  createFetchWithChunkBuffer,
 } from './fetch'
 import {
   CHUNK_LAST_OFFSET_HEADER,
@@ -21,7 +28,7 @@ import {
 /**
  * Options for constructing a ShapeStream.
  */
-export interface ShapeStreamOptions {
+export interface ShapeStreamOptions<T = never> {
   /**
    * The full URL to where the Shape is hosted. This can either be the Electric server
    * directly or a proxy. E.g. for a local Electric instance, you might set `http://localhost:3000/v1/shape/foo`
@@ -46,6 +53,13 @@ export interface ShapeStreamOptions {
    */
   shapeId?: string
   backoffOptions?: BackoffOptions
+
+  /**
+   * HTTP headers to attach to requests made by the client.
+   * Can be used for adding authentication headers.
+   */
+  headers?: Record<string, string>
+
   /**
    * Automatically fetch updates to the Shape. If you just want to sync the current
    * shape and stop, pass false.
@@ -53,10 +67,10 @@ export interface ShapeStreamOptions {
   subscribe?: boolean
   signal?: AbortSignal
   fetchClient?: typeof fetch
-  parser?: Parser
+  parser?: Parser<T>
 }
 
-export interface ShapeStreamInterface<T extends Row = Row> {
+export interface ShapeStreamInterface<T extends Row<unknown> = Row> {
   subscribe(
     callback: (messages: Message<T>[]) => MaybePromise<void>,
     onError?: (error: FetchError | Error) => void
@@ -108,10 +122,10 @@ export interface ShapeStreamInterface<T extends Row = Row> {
  * ```
  */
 
-export class ShapeStream<T extends Row = Row>
+export class ShapeStream<T extends Row<unknown> = Row>
   implements ShapeStreamInterface<T>
 {
-  readonly options: ShapeStreamOptions
+  readonly options: ShapeStreamOptions<GetExtensions<T>>
 
   readonly #fetchClient: typeof fetch
   readonly #messageParser: MessageParser<T>
@@ -119,7 +133,7 @@ export class ShapeStream<T extends Row = Row>
   readonly #subscribers = new Map<
     number,
     [
-      MessageProcessorInterface<Message<T>[]>,
+      (messages: Message<T>[]) => MaybePromise<void>,
       ((error: Error) => void) | undefined,
     ]
   >()
@@ -135,24 +149,26 @@ export class ShapeStream<T extends Row = Row>
   #shapeId?: string
   #schema?: Schema
 
-  constructor(options: ShapeStreamOptions) {
+  constructor(options: ShapeStreamOptions<GetExtensions<T>>) {
     validateOptions(options)
     this.options = { subscribe: true, ...options }
     this.#lastOffset = this.options.offset ?? `-1`
     this.#shapeId = this.options.shapeId
     this.#messageParser = new MessageParser<T>(options.parser)
 
-    this.#fetchClient = createFetchWithBackoff(
+    const baseFetchClient =
       options.fetchClient ??
-        ((...args: Parameters<typeof fetch>) => fetch(...args)),
-      {
-        ...(options.backoffOptions ?? BackoffDefaults),
-        onFailedAttempt: () => {
-          this.#connected = false
-          options.backoffOptions?.onFailedAttempt?.()
-        },
-      }
-    )
+      ((...args: Parameters<typeof fetch>) => fetch(...args))
+
+    const fetchWithBackoffClient = createFetchWithBackoff(baseFetchClient, {
+      ...(options.backoffOptions ?? BackoffDefaults),
+      onFailedAttempt: () => {
+        this.#connected = false
+        options.backoffOptions?.onFailedAttempt?.()
+      },
+    })
+
+    this.#fetchClient = createFetchWithChunkBuffer(fetchWithBackoffClient)
 
     this.start()
   }
@@ -190,7 +206,10 @@ export class ShapeStream<T extends Row = Row>
 
         let response!: Response
         try {
-          response = await this.#fetchClient(fetchUrl.toString(), { signal })
+          response = await this.#fetchClient(fetchUrl.toString(), {
+            signal,
+            headers: this.options.headers,
+          })
           this.#connected = true
         } catch (e) {
           if (e instanceof FetchBackoffAbortError) break // interrupted
@@ -199,14 +218,14 @@ export class ShapeStream<T extends Row = Row>
             // The request is invalid, most likely because the shape has been deleted.
             // We should start from scratch, this will force the shape to be recreated.
             this.#reset()
-            this.#publish(e.json as Message<T>[])
+            await this.#publish(e.json as Message<T>[])
             continue
           } else if (e.status == 409) {
             // Upon receiving a 409, we should start from scratch
             // with the newly provided shape ID
             const newShapeId = e.headers[SHAPE_ID_HEADER]
             this.#reset(newShapeId)
-            this.#publish(e.json as Message<T>[])
+            await this.#publish(e.json as Message<T>[])
             continue
           } else if (e.status >= 400 && e.status < 500) {
             // Notify subscribers
@@ -246,16 +265,17 @@ export class ShapeStream<T extends Row = Row>
 
         // Update isUpToDate
         if (batch.length > 0) {
+          const prevUpToDate = this.#isUpToDate
           const lastMessage = batch[batch.length - 1]
           if (isUpToDateMessage(lastMessage)) {
             this.#lastSyncedAt = Date.now()
-            if (!this.#isUpToDate) {
-              this.#isUpToDate = true
-              this.#notifyUpToDateSubscribers()
-            }
+            this.#isUpToDate = true
           }
 
-          this.#publish(batch)
+          await this.#publish(batch)
+          if (!prevUpToDate && this.#isUpToDate) {
+            this.#notifyUpToDateSubscribers()
+          }
         }
       }
     } finally {
@@ -268,9 +288,8 @@ export class ShapeStream<T extends Row = Row>
     onError?: (error: FetchError | Error) => void
   ) {
     const subscriptionId = Math.random()
-    const subscriber = new MessageProcessor(callback)
 
-    this.#subscribers.set(subscriptionId, [subscriber, onError])
+    this.#subscribers.set(subscriptionId, [callback, onError])
 
     return () => {
       this.#subscribers.delete(subscriptionId)
@@ -319,10 +338,18 @@ export class ShapeStream<T extends Row = Row>
     return !this.isUpToDate
   }
 
-  #publish(messages: Message<T>[]) {
-    this.#subscribers.forEach(([subscriber, _]) => {
-      subscriber.process(messages)
-    })
+  async #publish(messages: Message<T>[]): Promise<void> {
+    await Promise.all(
+      Array.from(this.#subscribers.values()).map(async ([callback, __]) => {
+        try {
+          await callback(messages)
+        } catch (err) {
+          queueMicrotask(() => {
+            throw err
+          })
+        }
+      })
+    )
   }
 
   #sendErrorToSubscribers(error: Error) {
@@ -356,7 +383,7 @@ export class ShapeStream<T extends Row = Row>
   }
 }
 
-function validateOptions(options: Partial<ShapeStreamOptions>): void {
+function validateOptions<T>(options: Partial<ShapeStreamOptions<T>>): void {
   if (!options.url) {
     throw new Error(`Invalid shape option. It must provide the url`)
   }

package/src/constants.ts

@@ -1,6 +1,7 @@
-export const SHAPE_ID_HEADER = `x-electric-shape-id`
-export const CHUNK_LAST_OFFSET_HEADER = `x-electric-chunk-last-offset`
-export const SHAPE_SCHEMA_HEADER = `x-electric-schema`
+export const SHAPE_ID_HEADER = `electric-shape-id`
+export const CHUNK_LAST_OFFSET_HEADER = `electric-chunk-last-offset`
+export const CHUNK_UP_TO_DATE_HEADER = `electric-chunk-up-to-date`
+export const SHAPE_SCHEMA_HEADER = `electric-schema`
 export const SHAPE_ID_QUERY_PARAM = `shape_id`
 export const OFFSET_QUERY_PARAM = `offset`
 export const WHERE_QUERY_PARAM = `where`

package/src/fetch.ts

@@ -1,3 +1,11 @@
+import {
+  CHUNK_LAST_OFFSET_HEADER,
+  CHUNK_UP_TO_DATE_HEADER,
+  LIVE_QUERY_PARAM,
+  OFFSET_QUERY_PARAM,
+  SHAPE_ID_HEADER,
+  SHAPE_ID_QUERY_PARAM,
+} from './constants'
 import { FetchError, FetchBackoffAbortError } from './error'
 
 export interface BackoffOptions {
@@ -77,3 +85,188 @@ export function createFetchWithBackoff(
     }
   }
 }
+
+interface ChunkPrefetchOptions {
+  maxChunksToPrefetch: number
+}
+
+const ChunkPrefetchDefaults = {
+  maxChunksToPrefetch: 2,
+}
+
+/**
+ * Creates a fetch client that prefetches subsequent log chunks for
+ * consumption by the shape stream without waiting for the chunk bodies
+ * themselves to be loaded.
+ *
+ * @param fetchClient the client to wrap
+ * @param prefetchOptions options to configure prefetching
+ * @returns wrapped client with prefetch capabilities
+ */
+export function createFetchWithChunkBuffer(
+  fetchClient: typeof fetch,
+  prefetchOptions: ChunkPrefetchOptions = ChunkPrefetchDefaults
+): typeof fetch {
+  const { maxChunksToPrefetch } = prefetchOptions
+
+  let prefetchQueue: PrefetchQueue
+
+  const prefetchClient = async (...args: Parameters<typeof fetchClient>) => {
+    const url = args[0].toString()
+
+    // try to consume from the prefetch queue first, and if request is
+    // not present abort the prefetch queue as it must no longer be valid
+    const prefetchedRequest = prefetchQueue?.consume(...args)
+    if (prefetchedRequest) {
+      return prefetchedRequest
+    }
+
+    prefetchQueue?.abort()
+
+    // perform request and fire off prefetch queue if request is eligible
+    const response = await fetchClient(...args)
+    const nextUrl = getNextChunkUrl(url, response)
+    if (nextUrl) {
+      prefetchQueue = new PrefetchQueue({
+        fetchClient,
+        maxPrefetchedRequests: maxChunksToPrefetch,
+        url: nextUrl,
+        requestInit: args[1],
+      })
+    }
+
+    return response
+  }
+
+  return prefetchClient
+}
+
+class PrefetchQueue {
+  readonly #fetchClient: typeof fetch
+  readonly #maxPrefetchedRequests: number
+  readonly #prefetchQueue = new Map<
+    string,
+    [Promise<Response>, AbortController]
+  >()
+  #queueHeadUrl: string | void
+  #queueTailUrl: string
+
+  constructor(options: {
+    url: Parameters<typeof fetch>[0]
+    requestInit: Parameters<typeof fetch>[1]
+    maxPrefetchedRequests: number
+    fetchClient?: typeof fetch
+  }) {
+    this.#fetchClient =
+      options.fetchClient ??
+      ((...args: Parameters<typeof fetch>) => fetch(...args))
+    this.#maxPrefetchedRequests = options.maxPrefetchedRequests
+    this.#queueHeadUrl = options.url.toString()
+    this.#queueTailUrl = this.#queueHeadUrl
+    this.#prefetch(options.url, options.requestInit)
+  }
+
+  abort(): void {
+    this.#prefetchQueue.forEach(([_, aborter]) => aborter.abort())
+  }
+
+  consume(...args: Parameters<typeof fetch>): Promise<Response> | void {
+    const url = args[0].toString()
+
+    const request = this.#prefetchQueue.get(url)?.[0]
+    // only consume if request is in queue and is the queue "head"
+    // if request is in the queue but not the head, the queue is being
+    // consumed out of order and should be restarted
+    if (!request || url !== this.#queueHeadUrl) return
+    this.#prefetchQueue.delete(url)
+
+    // fire off new prefetch since request has been consumed
+    request
+      .then((response) => {
+        const nextUrl = getNextChunkUrl(url, response)
+        this.#queueHeadUrl = nextUrl
+        if (!this.#prefetchQueue.has(this.#queueTailUrl)) {
+          this.#prefetch(this.#queueTailUrl, args[1])
+        }
+      })
+      .catch(() => {})
+
+    return request
+  }
+
+  #prefetch(...args: Parameters<typeof fetch>): void {
+    const url = args[0].toString()
+
+    // only prefetch when queue is not full
+    if (this.#prefetchQueue.size >= this.#maxPrefetchedRequests) return
+
+    // initialize aborter per request, to avoid aborting consumed requests that
+    // are still streaming their bodies to the consumer
+    const aborter = new AbortController()
+
+    try {
+      const request = this.#fetchClient(url, {
+        ...(args[1] ?? {}),
+        signal: chainAborter(aborter, args[1]?.signal),
+      })
+      this.#prefetchQueue.set(url, [request, aborter])
+      request
+        .then((response) => {
+          // only keep prefetching if response chain is uninterrupted
+          if (!response.ok || aborter.signal.aborted) return
+
+          const nextUrl = getNextChunkUrl(url, response)
+
+          // only prefetch when there is a next URL
+          if (!nextUrl || nextUrl === url) return
+
+          this.#queueTailUrl = nextUrl
+          return this.#prefetch(nextUrl, args[1])
+        })
+        .catch(() => {})
+    } catch (_) {
+      // ignore prefetch errors
+    }
+  }
+}
+
+/**
+ * Generate the next chunk's URL if the url and response are valid
+ */
+function getNextChunkUrl(url: string, res: Response): string | void {
+  const shapeId = res.headers.get(SHAPE_ID_HEADER)
+  const lastOffset = res.headers.get(CHUNK_LAST_OFFSET_HEADER)
+  const isUpToDate = res.headers.has(CHUNK_UP_TO_DATE_HEADER)
+
+  // only prefetch if shape ID and offset for next chunk are available, and
+  // response is not already up-to-date
+  if (!shapeId || !lastOffset || isUpToDate) return
+
+  const nextUrl = new URL(url)
+
+  // don't prefetch live requests, rushing them will only
+  // potentially miss more recent data
+  if (nextUrl.searchParams.has(LIVE_QUERY_PARAM)) return
+
+  nextUrl.searchParams.set(SHAPE_ID_QUERY_PARAM, shapeId)
+  nextUrl.searchParams.set(OFFSET_QUERY_PARAM, lastOffset)
+  return nextUrl.toString()
+}
+
+/**
+ * Chains an abort controller on an optional source signal's
+ * aborted state - if the source signal is aborted, the provided abort
+ * controller will also abort
+ */
+function chainAborter(
+  aborter: AbortController,
+  sourceSignal?: AbortSignal
+): AbortSignal {
+  if (!sourceSignal) return aborter.signal
+  if (sourceSignal.aborted) aborter.abort()
+  else
+    sourceSignal.addEventListener(`abort`, () => aborter.abort(), {
+      once: true,
+    })
+  return aborter.signal
+}

package/src/helpers.ts

@@ -17,7 +17,7 @@ import { ChangeMessage, ControlMessage, Message, Row } from './types'
  * }
  * ```
  */
-export function isChangeMessage<T extends Row = Row>(
+export function isChangeMessage<T extends Row<unknown> = Row>(
   message: Message<T>
 ): message is ChangeMessage<T> {
   return `key` in message
@@ -40,13 +40,13 @@ export function isChangeMessage<T extends Row = Row>(
  * }
  * ```
  */
-export function isControlMessage<T extends Row = Row>(
+export function isControlMessage<T extends Row<unknown> = Row>(
   message: Message<T>
 ): message is ControlMessage {
   return !isChangeMessage(message)
 }
 
-export function isUpToDateMessage<T extends Row = Row>(
+export function isUpToDateMessage<T extends Row<unknown> = Row>(
   message: Message<T>
 ): message is ControlMessage & { up_to_date: true } {
   return isControlMessage(message) && message.headers.control === `up-to-date`

package/src/parser.ts

@@ -1,17 +1,23 @@
-import { ColumnInfo, Message, Row, Schema, Value } from './types'
+import { ColumnInfo, GetExtensions, Message, Row, Schema, Value } from './types'
 
 type NullToken = null | `NULL`
 type Token = Exclude<string, NullToken>
 type NullableToken = Token | NullToken
-export type ParseFunction = (
+export type ParseFunction<Extensions = never> = (
   value: Token,
   additionalInfo?: Omit<ColumnInfo, `type` | `dims`>
-) => Value
-type NullableParseFunction = (
+) => Value<Extensions>
+type NullableParseFunction<Extensions = never> = (
   value: NullableToken,
   additionalInfo?: Omit<ColumnInfo, `type` | `dims`>
-) => Value
-export type Parser = { [key: string]: ParseFunction }
+) => Value<Extensions>
+/**
+ * @typeParam Extensions - Additional types that can be parsed by this parser beyond the standard SQL types.
+ *                         Defaults to no additional types.
+ */
+export type Parser<Extensions = never> = {
+  [key: string]: ParseFunction<Extensions>
+}
 
 const parseNumber = (value: string) => Number(value)
 const parseBool = (value: string) => value === `true` || value === `t`
@@ -31,7 +37,10 @@ export const defaultParser: Parser = {
 }
 
 // Taken from: https://github.com/electric-sql/pglite/blob/main/packages/pglite/src/types.ts#L233-L279
-export function pgArrayParser(value: Token, parser?: ParseFunction): Value {
+export function pgArrayParser<Extensions>(
+  value: Token,
+  parser?: ParseFunction<Extensions>
+): Value<Extensions> {
   let i = 0
   let char = null
   let str = ``
@@ -39,7 +48,7 @@ export function pgArrayParser(value: Token, parser?: ParseFunction): Value {
   let last = 0
   let p: string | undefined = undefined
 
-  function loop(x: string): Value[] {
+  function loop(x: string): Array<Value<Extensions>> {
     const xs = []
     for (; i < x.length; i++) {
       char = x[i]
@@ -79,9 +88,9 @@ export function pgArrayParser(value: Token, parser?: ParseFunction): Value {
   return loop(value)[0]
 }
 
-export class MessageParser<T extends Row> {
-  private parser: Parser
-  constructor(parser?: Parser) {
+export class MessageParser<T extends Row<unknown>> {
+  private parser: Parser<GetExtensions<T>>
+  constructor(parser?: Parser<GetExtensions<T>>) {
     // Merge the provided parser with the default parser
     // to use the provided parser whenever defined
     // and otherwise fall back to the default parser
@@ -96,7 +105,7 @@ export class MessageParser<T extends Row> {
       // But `typeof null === 'object'` so we need to make an explicit check.
       if (key === `value` && typeof value === `object` && value !== null) {
         // Parse the row values
-        const row = value as Record<string, Value>
+        const row = value as Record<string, Value<GetExtensions<T>>>
         Object.keys(row).forEach((key) => {
           row[key] = this.parseRow(key, row[key] as NullableToken, schema)
         })
@@ -106,7 +115,11 @@ export class MessageParser<T extends Row> {
   }
 
   // Parses the message values using the provided parser based on the schema information
-  private parseRow(key: string, value: NullableToken, schema: Schema): Value {
+  private parseRow(
+    key: string,
+    value: NullableToken,
+    schema: Schema
+  ): Value<GetExtensions<T>> {
     const columnInfo = schema[key]
     if (!columnInfo) {
       // We don't have information about the value
@@ -137,11 +150,11 @@ export class MessageParser<T extends Row> {
   }
 }
 
-function makeNullableParser(
-  parser: ParseFunction,
+function makeNullableParser<Extensions>(
+  parser: ParseFunction<Extensions>,
   columnInfo: ColumnInfo,
   columnName?: string
-): NullableParseFunction {
+): NullableParseFunction<Extensions> {
   const isNullable = !(columnInfo.not_null ?? false)
   // The sync service contains `null` value for a column whose value is NULL
   // but if the column value is an array that contains a NULL value

package/src/shape.ts

@@ -3,8 +3,8 @@ import { isChangeMessage, isControlMessage } from './helpers'
 import { FetchError } from './error'
 import { ShapeStreamInterface } from './client'
 
-export type ShapeData<T extends Row = Row> = Map<string, T>
-export type ShapeChangedCallback<T extends Row = Row> = (
+export type ShapeData<T extends Row<unknown> = Row> = Map<string, T>
+export type ShapeChangedCallback<T extends Row<unknown> = Row> = (
   value: ShapeData<T>
 ) => void
 
@@ -39,7 +39,7 @@ export type ShapeChangedCallback<T extends Row = Row> = (
  *       console.log(shapeData)
  *     })
  */
-export class Shape<T extends Row = Row> {
+export class Shape<T extends Row<unknown> = Row> {
   readonly #stream: ShapeStreamInterface<T>
 
   readonly #data: ShapeData<T> = new Map()

package/src/types.ts

@@ -1,13 +1,21 @@
-export type Value =
+/**
+ * Default types for SQL but can be extended with additional types when using a custom parser.
+ * @typeParam Extensions - Additional value types.
+ */
+export type Value<Extensions = never> =
   | string
   | number
   | boolean
   | bigint
   | null
-  | Value[]
-  | { [key: string]: Value }
+  | Extensions
+  | Value<Extensions>[]
+  | { [key: string]: Value<Extensions> }
+
+export type Row<Extensions = never> = Record<string, Value<Extensions>>
 
-export type Row = { [key: string]: Value }
+export type GetExtensions<T extends Row<unknown>> =
+  T extends Row<infer Extensions> ? Extensions : never
 
 export type Offset = `-1` | `${number}_${number}`
 
@@ -19,7 +27,7 @@ export type ControlMessage = {
   headers: Header & { control: `up-to-date` | `must-refetch` }
 }
 
-export type ChangeMessage<T extends Row = Row> = {
+export type ChangeMessage<T extends Row<unknown> = Row> = {
   key: string
   value: T
   headers: Header & { operation: `insert` | `update` | `delete` }
@@ -27,7 +35,9 @@ export type ChangeMessage<T extends Row = Row> = {
 }
 
 // Define the type for a record
-export type Message<T extends Row = Row> = ControlMessage | ChangeMessage<T>
+export type Message<T extends Row<unknown> = Row> =
+  | ControlMessage
+  | ChangeMessage<T>
 
 /**
  * Common properties for all columns.
@@ -104,7 +114,7 @@ export type ColumnInfo =
 
 export type Schema = { [key: string]: ColumnInfo }
 
-export type TypedMessages<T extends Row = Row> = {
+export type TypedMessages<T extends Row<unknown> = Row> = {
   messages: Array<Message<T>>
   schema: ColumnInfo
 }

package/src/queue.ts

@@ -1,62 +0,0 @@
-import { MaybePromise } from './types'
-
-function isThenable(value: MaybePromise<void>): value is Promise<void> {
-  return (
-    !!value &&
-    typeof value === `object` &&
-    `then` in value &&
-    typeof value.then === `function`
-  )
-}
-
-/**
- * Processes messages asynchronously in order.
- */
-export class AsyncProcessingQueue {
-  #processingChain: MaybePromise<void> = undefined
-
-  public process(callback: () => MaybePromise<void>): MaybePromise<void> {
-    this.#processingChain = isThenable(this.#processingChain)
-      ? this.#processingChain.then(callback)
-      : callback()
-    return this.#processingChain
-  }
-
-  public async waitForProcessing(): Promise<void> {
-    let currentChain: MaybePromise<void>
-    do {
-      currentChain = this.#processingChain
-      await currentChain
-    } while (this.#processingChain !== currentChain)
-  }
-}
-
-export interface MessageProcessorInterface<T> {
-  process(messages: T): MaybePromise<void>
-  waitForProcessing(): Promise<void>
-}
-
-/**
- * Receives messages, puts them on a queue and processes
- * them synchronously or asynchronously by passing to a
- * registered callback function.
- *
- * @constructor
- * @param {(message: T) => MaybePromise<void>} callback function
- */
-export class MessageProcessor<T> implements MessageProcessorInterface<T> {
-  readonly #queue = new AsyncProcessingQueue()
-  readonly #callback: (messages: T) => MaybePromise<void>
-
-  constructor(callback: (messages: T) => MaybePromise<void>) {
-    this.#callback = callback
-  }
-
-  public process(messages: T): void {
-    this.#queue.process(() => this.#callback(messages))
-  }
-
-  public async waitForProcessing(): Promise<void> {
-    await this.#queue.waitForProcessing()
-  }
-}