@electric-sql/client 1.0.10 → 1.0.11

package/src/client.ts

@@ -5,9 +5,11 @@ import {
   Row,
   MaybePromise,
   GetExtensions,
+  ChangeMessage,
+  SnapshotMetadata,
 } from './types'
 import { MessageParser, Parser, TransformFunction } from './parser'
-import { getOffset, isUpToDateMessage } from './helpers'
+import { getOffset, isUpToDateMessage, isChangeMessage } from './helpers'
 import {
   FetchError,
   FetchBackoffAbortError,
@@ -43,12 +45,19 @@ import {
   PAUSE_STREAM,
   EXPERIMENTAL_LIVE_SSE_QUERY_PARAM,
   ELECTRIC_PROTOCOL_QUERY_PARAMS,
+  LOG_MODE_QUERY_PARAM,
+  SUBSET_PARAM_WHERE,
+  SUBSET_PARAM_WHERE_PARAMS,
+  SUBSET_PARAM_LIMIT,
+  SUBSET_PARAM_OFFSET,
+  SUBSET_PARAM_ORDER_BY,
 } from './constants'
 import {
   EventSourceMessage,
   fetchEventSource,
 } from '@microsoft/fetch-event-source'
 import { expiredShapesCache } from './expired-shapes-cache'
+import { SnapshotTracker } from './snapshot-tracker'
 
 const RESERVED_PARAMS: Set<ReservedParamKeys> = new Set([
   LIVE_CACHE_BUSTER_QUERY_PARAM,
@@ -58,6 +67,7 @@ const RESERVED_PARAMS: Set<ReservedParamKeys> = new Set([
 ])
 
 type Replica = `full` | `default`
+export type LogMode = `changes_only` | `full`
 
 /**
  * PostgreSQL-specific shape parameters that can be provided externally
@@ -114,11 +124,20 @@ export type ExternalParamsRecord<T extends Row<unknown> = Row> = {
   [K in string]: ParamValue | undefined
 } & Partial<PostgresParams<T>> & { [K in ReservedParamKeys]?: never }
 
+export type SubsetParams = {
+  where?: string
+  params?: Record<string, string>
+  limit?: number
+  offset?: number
+  orderBy?: string
+}
+
 type ReservedParamKeys =
   | typeof LIVE_CACHE_BUSTER_QUERY_PARAM
   | typeof SHAPE_HANDLE_QUERY_PARAM
   | typeof LIVE_QUERY_PARAM
   | typeof OFFSET_QUERY_PARAM
+  | `subset__${string}`
 
 /**
  * External headers type - what users provide.
@@ -258,6 +277,11 @@ export interface ShapeStreamOptions<T = never> {
    */
   experimentalLiveSse?: boolean
 
+  /**
+   * Initial data loading mode
+   */
+  mode?: LogMode
+
   signal?: AbortSignal
   fetchClient?: typeof fetch
   backoffOptions?: BackoffOptions
@@ -293,8 +317,20 @@ export interface ShapeStreamInterface<T extends Row<unknown> = Row> {
   lastOffset: Offset
   shapeHandle?: string
   error?: unknown
+  mode: LogMode
 
   forceDisconnectAndRefresh(): Promise<void>
+
+  requestSnapshot(params: {
+    where?: string
+    params?: Record<string, string>
+    limit: number
+    offset?: number
+    orderBy: string
+  }): Promise<{
+    metadata: SnapshotMetadata
+    data: Array<Message<T>>
+  }>
 }
 
 /**
@@ -383,8 +419,10 @@ export class ShapeStream<T extends Row<unknown> = Row>
   #liveCacheBuster: string // Seconds since our Electric Epoch 😎
   #lastSyncedAt?: number // unix time
   #isUpToDate: boolean = false
+  #isMidStream: boolean = true
   #connected: boolean = false
   #shapeHandle?: string
+  #mode: LogMode
   #schema?: Schema
   #onError?: ShapeStreamErrorHandler
   #requestAbortController?: AbortController
@@ -393,6 +431,10 @@ export class ShapeStream<T extends Row<unknown> = Row>
   #tickPromiseResolver?: () => void
   #tickPromiseRejecter?: (reason?: unknown) => void
   #messageChain = Promise.resolve<void[]>([]) // promise chain for incoming messages
+  #snapshotTracker = new SnapshotTracker()
+  #activeSnapshotRequests = 0 // counter for concurrent snapshot requests
+  #midStreamPromise?: Promise<void>
+  #midStreamPromiseResolver?: () => void
 
   constructor(options: ShapeStreamOptions<GetExtensions<T>>) {
     this.options = { subscribe: true, ...options }
@@ -405,6 +447,7 @@ export class ShapeStream<T extends Row<unknown> = Row>
       options.transformer
     )
     this.#onError = this.options.onError
+    this.#mode = this.options.mode ?? `full`
 
     const baseFetchClient =
       options.fetchClient ??
@@ -447,6 +490,10 @@ export class ShapeStream<T extends Row<unknown> = Row>
     return this.#lastOffset
   }
 
+  get mode() {
+    return this.#mode
+  }
+
   async #start(): Promise<void> {
     this.#started = true
 
@@ -573,7 +620,11 @@ export class ShapeStream<T extends Row<unknown> = Row>
     return this.#requestShape()
   }
 
-  async #constructUrl(url: string, resumingFromPause: boolean) {
+  async #constructUrl(
+    url: string,
+    resumingFromPause: boolean,
+    subsetParams?: SubsetParams
+  ) {
     // Resolve headers and params in parallel
     const [requestHeaders, params] = await Promise.all([
       resolveHeaders(this.options.headers),
@@ -583,9 +634,7 @@ export class ShapeStream<T extends Row<unknown> = Row>
     ])
 
     // Validate params after resolution
-    if (params) {
-      validateParams(params)
-    }
+    if (params) validateParams(params)
 
     const fetchUrl = new URL(url)
 
@@ -612,8 +661,22 @@ export class ShapeStream<T extends Row<unknown> = Row>
       }
     }
 
+    if (subsetParams) {
+      if (subsetParams.where)
+        setQueryParam(fetchUrl, SUBSET_PARAM_WHERE, subsetParams.where)
+      if (subsetParams.params)
+        setQueryParam(fetchUrl, SUBSET_PARAM_WHERE_PARAMS, subsetParams.params)
+      if (subsetParams.limit)
+        setQueryParam(fetchUrl, SUBSET_PARAM_LIMIT, subsetParams.limit)
+      if (subsetParams.offset)
+        setQueryParam(fetchUrl, SUBSET_PARAM_OFFSET, subsetParams.offset)
+      if (subsetParams.orderBy)
+        setQueryParam(fetchUrl, SUBSET_PARAM_ORDER_BY, subsetParams.orderBy)
+    }
+
     // Add Electric's internal parameters
     fetchUrl.searchParams.set(OFFSET_QUERY_PARAM, this.#lastOffset)
+    fetchUrl.searchParams.set(LOG_MODE_QUERY_PARAM, this.#mode)
 
     if (this.#isUpToDate) {
       // If we are resuming from a paused state, we don't want to perform a live request
@@ -705,6 +768,9 @@ export class ShapeStream<T extends Row<unknown> = Row>
   async #onMessages(batch: Array<Message<T>>, isSseMessage = false) {
     // Update isUpToDate
     if (batch.length > 0) {
+      // Set isMidStream to true when we receive any data
+      this.#isMidStream = true
+
       const lastMessage = batch[batch.length - 1]
       if (isUpToDateMessage(lastMessage)) {
         if (isSseMessage) {
@@ -718,9 +784,21 @@ export class ShapeStream<T extends Row<unknown> = Row>
         }
         this.#lastSyncedAt = Date.now()
         this.#isUpToDate = true
+        // Set isMidStream to false when we see an up-to-date message
+        this.#isMidStream = false
+        // Resolve the promise waiting for mid-stream to end
+        this.#midStreamPromiseResolver?.()
       }
 
-      await this.#publish(batch)
+      // Filter messages using snapshot tracker
+      const messagesToProcess = batch.filter((message) => {
+        if (isChangeMessage(message)) {
+          return !this.#snapshotTracker.shouldRejectMessage(message)
+        }
+        return true // Always process control messages
+      })
+
+      await this.#publish(messagesToProcess)
     }
   }
 
@@ -904,6 +982,24 @@ export class ShapeStream<T extends Row<unknown> = Row>
     return this.#tickPromise
   }
 
+  /** Await until we're not in the middle of a stream (i.e., until we see an up-to-date message) */
+  async #waitForStreamEnd() {
+    if (!this.#isMidStream) {
+      return
+    }
+    if (this.#midStreamPromise) {
+      return this.#midStreamPromise
+    }
+    this.#midStreamPromise = new Promise((resolve) => {
+      this.#midStreamPromiseResolver = resolve
+    })
+    this.#midStreamPromise.finally(() => {
+      this.#midStreamPromise = undefined
+      this.#midStreamPromiseResolver = undefined
+    })
+    return this.#midStreamPromise
+  }
+
   /**
    * Refreshes the shape stream.
    * This preemptively aborts any ongoing long poll and reconnects without
@@ -976,8 +1072,108 @@ export class ShapeStream<T extends Row<unknown> = Row>
     this.#liveCacheBuster = ``
     this.#shapeHandle = handle
     this.#isUpToDate = false
+    this.#isMidStream = true
     this.#connected = false
     this.#schema = undefined
+    this.#activeSnapshotRequests = 0
+  }
+
+  /**
+   * Request a snapshot for subset of data.
+   *
+   * Only available when mode is `changes_only`.
+   * Returns the insertion point & the data, but more importantly injects the data
+   * into the subscribed data stream. Returned value is unlikely to be useful for the caller,
+   * unless the caller has complicated additional logic.
+   *
+   * Data will be injected in a way that's also tracking further incoming changes, and it'll
+   * skip the ones that are already in the snapshot.
+   *
+   * @param opts - The options for the snapshot request.
+   * @returns The metadata and the data for the snapshot.
+   */
+  async requestSnapshot(opts: SubsetParams): Promise<{
+    metadata: SnapshotMetadata
+    data: Array<ChangeMessage<T>>
+  }> {
+    if (this.#mode === `full`) {
+      throw new Error(
+        `Snapshot requests are not supported in ${this.#mode} mode, as the consumer is guaranteed to observe all data`
+      )
+    }
+    // We shouldn't be getting a snapshot on a shape that's not started
+    if (!this.#started) await this.#start()
+
+    // Wait until we're not mid-stream before pausing
+    // This ensures we don't pause in the middle of a transaction
+    await this.#waitForStreamEnd()
+
+    // Pause the stream if this is the first snapshot request
+    this.#activeSnapshotRequests++
+
+    try {
+      if (this.#activeSnapshotRequests === 1) {
+        // Currently this cannot throw, but in case it can later it's in this try block to not have a stuck counter
+        this.#pause()
+      }
+
+      const { fetchUrl, requestHeaders } = await this.#constructUrl(
+        this.options.url,
+        true,
+        opts
+      )
+
+      const { metadata, data } = await this.#fetchSnapshot(
+        fetchUrl,
+        requestHeaders
+      )
+
+      const dataWithEndBoundary = (data as Array<Message<T>>).concat([
+        { headers: { control: `snapshot-end`, ...metadata } },
+      ])
+
+      this.#snapshotTracker.addSnapshot(
+        metadata,
+        new Set(data.map((message) => message.key))
+      )
+      this.#onMessages(dataWithEndBoundary, false)
+
+      return {
+        metadata,
+        data,
+      }
+    } finally {
+      // Resume the stream if this was the last snapshot request
+      this.#activeSnapshotRequests--
+      if (this.#activeSnapshotRequests === 0) {
+        this.#resume()
+      }
+    }
+  }
+
+  async #fetchSnapshot(url: URL, headers: Record<string, string>) {
+    const response = await this.#fetchClient(url.toString(), { headers })
+
+    if (!response.ok) {
+      throw new FetchError(
+        response.status,
+        undefined,
+        undefined,
+        Object.fromEntries([...response.headers.entries()]),
+        url.toString()
+      )
+    }
+
+    const { metadata, data } = await response.json()
+    const batch = this.#messageParser.parse<Array<ChangeMessage<T>>>(
+      JSON.stringify(data),
+      this.#schema!
+    )
+
+    return {
+      metadata,
+      data: batch,
+    }
   }
 }
 
@@ -1007,6 +1203,7 @@ function validateOptions<T>(options: Partial<ShapeStreamOptions<T>>): void {
   if (
     options.offset !== undefined &&
     options.offset !== `-1` &&
+    options.offset !== `now` &&
     !options.handle
   ) {
     throw new MissingShapeHandleError()

package/src/constants.ts

@@ -16,6 +16,12 @@ export const WHERE_PARAMS_PARAM = `params`
 export const EXPERIMENTAL_LIVE_SSE_QUERY_PARAM = `experimental_live_sse`
 export const FORCE_DISCONNECT_AND_REFRESH = `force-disconnect-and-refresh`
 export const PAUSE_STREAM = `pause-stream`
+export const LOG_MODE_QUERY_PARAM = `log`
+export const SUBSET_PARAM_WHERE = `subset__where`
+export const SUBSET_PARAM_LIMIT = `subset__limit`
+export const SUBSET_PARAM_OFFSET = `subset__offset`
+export const SUBSET_PARAM_ORDER_BY = `subset__order_by`
+export const SUBSET_PARAM_WHERE_PARAMS = `subset__params`
 
 // Query parameters that should be passed through when proxying Electric requests
 export const ELECTRIC_PROTOCOL_QUERY_PARAMS: Array<string> = [
@@ -24,4 +30,10 @@ export const ELECTRIC_PROTOCOL_QUERY_PARAMS: Array<string> = [
   OFFSET_QUERY_PARAM,
   LIVE_CACHE_BUSTER_QUERY_PARAM,
   EXPIRED_HANDLE_QUERY_PARAM,
+  LOG_MODE_QUERY_PARAM,
+  SUBSET_PARAM_WHERE,
+  SUBSET_PARAM_LIMIT,
+  SUBSET_PARAM_OFFSET,
+  SUBSET_PARAM_ORDER_BY,
+  SUBSET_PARAM_WHERE_PARAMS,
 ]

package/src/fetch.ts

@@ -5,6 +5,11 @@ import {
   OFFSET_QUERY_PARAM,
   SHAPE_HANDLE_HEADER,
   SHAPE_HANDLE_QUERY_PARAM,
+  SUBSET_PARAM_LIMIT,
+  SUBSET_PARAM_OFFSET,
+  SUBSET_PARAM_ORDER_BY,
+  SUBSET_PARAM_WHERE,
+  SUBSET_PARAM_WHERE_PARAMS,
 } from './constants'
 import {
   FetchError,
@@ -206,11 +211,24 @@ export function createFetchWithResponseHeadersCheck(
 
       const addMissingHeaders = (requiredHeaders: Array<string>) =>
         missingHeaders.push(...requiredHeaders.filter((h) => !headers.has(h)))
-      addMissingHeaders(requiredElectricResponseHeaders)
 
       const input = args[0]
       const urlString = input.toString()
       const url = new URL(urlString)
+
+      // Snapshot responses (subset params) return a JSON object and do not include Electric chunk headers
+      const isSnapshotRequest = [
+        SUBSET_PARAM_WHERE,
+        SUBSET_PARAM_WHERE_PARAMS,
+        SUBSET_PARAM_LIMIT,
+        SUBSET_PARAM_OFFSET,
+        SUBSET_PARAM_ORDER_BY,
+      ].some((p) => url.searchParams.has(p))
+      if (isSnapshotRequest) {
+        return response
+      }
+
+      addMissingHeaders(requiredElectricResponseHeaders)
       if (url.searchParams.get(LIVE_QUERY_PARAM) === `true`) {
         addMissingHeaders(requiredLiveResponseHeaders)
       }

package/src/helpers.ts

@@ -1,4 +1,12 @@
-import { ChangeMessage, ControlMessage, Message, Offset, Row } from './types'
+import {
+  ChangeMessage,
+  ControlMessage,
+  Message,
+  NormalizedPgSnapshot,
+  Offset,
+  PostgresSnapshot,
+  Row,
+} from './types'
 
 /**
  * Type guard for checking {@link Message} is {@link ChangeMessage}.
@@ -64,3 +72,28 @@ export function getOffset(message: ControlMessage): Offset | undefined {
   }
   return `${lsn}_0` as Offset
 }
+
+/**
+ * Checks if a transaction is visible in a snapshot.
+ *
+ * @param txid - the transaction id to check
+ * @param snapshot - the information about the snapshot
+ * @returns true if the transaction is visible in the snapshot
+ */
+export function isVisibleInSnapshot(
+  txid: number | bigint | `${bigint}`,
+  snapshot: PostgresSnapshot | NormalizedPgSnapshot
+): boolean {
+  const xid = BigInt(txid)
+  const xmin = BigInt(snapshot.xmin)
+  const xmax = BigInt(snapshot.xmax)
+  const xip = snapshot.xip_list.map(BigInt)
+
+  // If the transaction id is less than the minimum transaction id, it is visible in the snapshot.
+  // If the transaction id is less than the maximum transaction id and not in the list of active
+  //   transactions at the time of the snapshot, it has been committed before the snapshot was taken
+  //   and is therefore visible in the snapshot.
+  // Otherwise, it is not visible in the snapshot.
+
+  return xid < xmin || (xid < xmax && !xip.includes(xid))
+}

package/src/index.ts

@@ -1,7 +1,11 @@
 export * from './client'
 export * from './shape'
 export * from './types'
-export { isChangeMessage, isControlMessage } from './helpers'
+export {
+  isChangeMessage,
+  isControlMessage,
+  isVisibleInSnapshot,
+} from './helpers'
 export { FetchError } from './error'
 export { type BackoffOptions, BackoffDefaults } from './fetch'
 export { ELECTRIC_PROTOCOL_QUERY_PARAMS } from './constants'

package/src/shape.ts

@@ -1,7 +1,7 @@
 import { Message, Offset, Row } from './types'
 import { isChangeMessage, isControlMessage } from './helpers'
 import { FetchError } from './error'
-import { ShapeStreamInterface } from './client'
+import { LogMode, ShapeStreamInterface } from './client'
 
 export type ShapeData<T extends Row<unknown> = Row> = Map<string, T>
 export type ShapeChangedCallback<T extends Row<unknown> = Row> = (data: {
@@ -52,6 +52,9 @@ export class Shape<T extends Row<unknown> = Row> {
 
   readonly #data: ShapeData<T> = new Map()
   readonly #subscribers = new Map<number, ShapeChangedCallback<T>>()
+  readonly #insertedKeys = new Set<string>()
+  readonly #requestedSubSnapshots = new Set<string>()
+  #reexecuteSnapshotsPending = false
   #status: ShapeStatus = `syncing`
   #error: FetchError | false = false
 
@@ -125,6 +128,26 @@ export class Shape<T extends Row<unknown> = Row> {
     return this.stream.isConnected()
   }
 
+  /** Current log mode of the underlying stream */
+  get mode(): LogMode {
+    return this.stream.mode
+  }
+
+  /**
+   * Request a snapshot for subset of data. Only available when mode is changes_only.
+   * Returns void; data will be emitted via the stream and processed by this Shape.
+   */
+  async requestSnapshot(
+    params: Parameters<ShapeStreamInterface<T>[`requestSnapshot`]>[0]
+  ): Promise<void> {
+    // Track this snapshot request for future re-execution on shape rotation
+    const key = JSON.stringify(params)
+    this.#requestedSubSnapshots.add(key)
+    // Ensure the stream is up-to-date so schema is available for parsing
+    await this.#awaitUpToDate()
+    await this.stream.requestSnapshot(params)
+  }
+
   subscribe(callback: ShapeChangedCallback<T>): () => void {
     const subscriptionId = Math.random()
 
@@ -149,19 +172,43 @@ export class Shape<T extends Row<unknown> = Row> {
     messages.forEach((message) => {
       if (isChangeMessage(message)) {
         shouldNotify = this.#updateShapeStatus(`syncing`)
-        switch (message.headers.operation) {
-          case `insert`:
-            this.#data.set(message.key, message.value)
-            break
-          case `update`:
-            this.#data.set(message.key, {
-              ...this.#data.get(message.key)!,
-              ...message.value,
-            })
-            break
-          case `delete`:
-            this.#data.delete(message.key)
-            break
+        if (this.mode === `full`) {
+          switch (message.headers.operation) {
+            case `insert`:
+              this.#data.set(message.key, message.value)
+              break
+            case `update`:
+              this.#data.set(message.key, {
+                ...this.#data.get(message.key)!,
+                ...message.value,
+              })
+              break
+            case `delete`:
+              this.#data.delete(message.key)
+              break
+          }
+        } else {
+          // changes_only: only apply updates/deletes for keys for which we observed an insert
+          switch (message.headers.operation) {
+            case `insert`:
+              this.#insertedKeys.add(message.key)
+              this.#data.set(message.key, message.value)
+              break
+            case `update`:
+              if (this.#insertedKeys.has(message.key)) {
+                this.#data.set(message.key, {
+                  ...this.#data.get(message.key)!,
+                  ...message.value,
+                })
+              }
+              break
+            case `delete`:
+              if (this.#insertedKeys.has(message.key)) {
+                this.#data.delete(message.key)
+                this.#insertedKeys.delete(message.key)
+              }
+              break
+          }
         }
       }
 
@@ -169,11 +216,18 @@ export class Shape<T extends Row<unknown> = Row> {
         switch (message.headers.control) {
           case `up-to-date`:
             shouldNotify = this.#updateShapeStatus(`up-to-date`)
+            if (this.#reexecuteSnapshotsPending) {
+              this.#reexecuteSnapshotsPending = false
+              void this.#reexecuteSnapshots()
+            }
             break
           case `must-refetch`:
             this.#data.clear()
+            this.#insertedKeys.clear()
             this.#error = false
             shouldNotify = this.#updateShapeStatus(`syncing`)
+            // Flag to re-execute sub-snapshots once the new shape is up-to-date
+            this.#reexecuteSnapshotsPending = true
             break
         }
       }
@@ -182,6 +236,42 @@ export class Shape<T extends Row<unknown> = Row> {
     if (shouldNotify) this.#notify()
   }
 
+  async #reexecuteSnapshots(): Promise<void> {
+    // Wait until stream is up-to-date again (ensures schema is available)
+    await this.#awaitUpToDate()
+
+    // Re-execute all snapshots concurrently
+    await Promise.all(
+      Array.from(this.#requestedSubSnapshots).map(async (jsonParams) => {
+        try {
+          const snapshot = JSON.parse(jsonParams)
+          await this.stream.requestSnapshot(snapshot)
+        } catch (_) {
+          // Ignore and continue; errors will be surfaced via stream onError
+        }
+      })
+    )
+  }
+
+  async #awaitUpToDate(): Promise<void> {
+    if (this.stream.isUpToDate) return
+    await new Promise<void>((resolve) => {
+      const check = () => {
+        if (this.stream.isUpToDate) {
+          clearInterval(interval)
+          unsub()
+          resolve()
+        }
+      }
+      const interval = setInterval(check, 10)
+      const unsub = this.stream.subscribe(
+        () => check(),
+        () => check()
+      )
+      check()
+    })
+  }
+
   #updateShapeStatus(status: ShapeStatus): boolean {
     const stateChanged = this.#status !== status
     this.#status = status