From ccd1a0c41dcb96a5f6a3d7a0434ed22f3a91d236 Mon Sep 17 00:00:00 2001
From: Kevin De Porre <kevin@electric-sql.com>
Date: Thu, 3 Oct 2024 15:11:46 +0200
Subject: [PATCH 1/5] Make Value, Parser, etc. generic such that they can be
 parameterized with additional types supported by custom parsers.

---
 packages/typescript-client/src/client.ts      | 23 ++++++----
 packages/typescript-client/src/helpers.ts     |  6 +--
 packages/typescript-client/src/parser.ts      | 45 ++++++++++++-------
 packages/typescript-client/src/shape.ts       |  6 +--
 packages/typescript-client/src/types.ts       | 24 +++++++---
 .../typescript-client/test/client.test-d.ts   | 11 +++++
 .../test/support/test-helpers.ts              |  2 +-
 7 files changed, 79 insertions(+), 38 deletions(-)

diff --git a/packages/typescript-client/src/client.ts b/packages/typescript-client/src/client.ts
index 254b95be2a..f86eeb4d29 100644
--- a/packages/typescript-client/src/client.ts
+++ b/packages/typescript-client/src/client.ts
@@ -1,4 +1,11 @@
-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 { FetchError, FetchBackoffAbortError } from './error'
@@ -21,7 +28,7 @@ import {
 /**
  * Options for constructing a ShapeStream.
  */
-export interface ShapeStreamOptions {
+export interface ShapeStreamOptions<T> {
   /**
    * 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`
@@ -53,10 +60,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 +115,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>
@@ -135,7 +142,7 @@ 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`
@@ -366,7 +373,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`)
   }
diff --git a/packages/typescript-client/src/helpers.ts b/packages/typescript-client/src/helpers.ts
index c3bb5ab299..c30ed129f0 100644
--- a/packages/typescript-client/src/helpers.ts
+++ b/packages/typescript-client/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`
diff --git a/packages/typescript-client/src/parser.ts b/packages/typescript-client/src/parser.ts
index a519551407..814c8a9ae5 100644
--- a/packages/typescript-client/src/parser.ts
+++ b/packages/typescript-client/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
diff --git a/packages/typescript-client/src/shape.ts b/packages/typescript-client/src/shape.ts
index 43c397fcc4..6012850fee 100644
--- a/packages/typescript-client/src/shape.ts
+++ b/packages/typescript-client/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()
diff --git a/packages/typescript-client/src/types.ts b/packages/typescript-client/src/types.ts
index 677b2fce73..2332a4e706 100644
--- a/packages/typescript-client/src/types.ts
+++ b/packages/typescript-client/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
 }
diff --git a/packages/typescript-client/test/client.test-d.ts b/packages/typescript-client/test/client.test-d.ts
index e6c3f395d5..1a81cc69dd 100644
--- a/packages/typescript-client/test/client.test-d.ts
+++ b/packages/typescript-client/test/client.test-d.ts
@@ -12,6 +12,7 @@ type CustomRow = {
   foo: number
   bar: boolean
   baz: string
+  ts: Date
 }
 
 describe(`client`, () => {
@@ -30,6 +31,11 @@ describe(`client`, () => {
     it(`should infer correct return type when provided`, () => {
       const shapeStream = new ShapeStream<CustomRow>({
         url: ``,
+        parser: {
+          timestampz: (date: string) => {
+            return new Date(date)
+          },
+        },
       })
 
       shapeStream.subscribe((msgs) => {
@@ -61,6 +67,11 @@ describe(`client`, () => {
     it(`should infer correct return type when provided`, async () => {
       const shapeStream = new ShapeStream<CustomRow>({
         url: ``,
+        parser: {
+          timestampz: (date: string) => {
+            return new Date(date)
+          },
+        },
       })
       const shape = new Shape(shapeStream)
       expectTypeOf(shape).toEqualTypeOf<Shape<CustomRow>>()
diff --git a/packages/typescript-client/test/support/test-helpers.ts b/packages/typescript-client/test/support/test-helpers.ts
index 6b6c52fdec..edeed4eb1f 100644
--- a/packages/typescript-client/test/support/test-helpers.ts
+++ b/packages/typescript-client/test/support/test-helpers.ts
@@ -14,7 +14,7 @@ export function makePgClient(overrides: ClientConfig = {}) {
   })
 }
 
-export function forEachMessage<T extends Row>(
+export function forEachMessage<T extends Row<unknown>>(
   stream: ShapeStreamInterface<T>,
   controller: AbortController,
   handler: (

From 3a2a265df0039ec2926ca1080f5ec94c1ffc296c Mon Sep 17 00:00:00 2001
From: Kevin De Porre <kevin@electric-sql.com>
Date: Thu, 3 Oct 2024 15:13:03 +0200
Subject: [PATCH 2/5] changeset

---
 .changeset/pink-rivers-clean.md | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 .changeset/pink-rivers-clean.md

diff --git a/.changeset/pink-rivers-clean.md b/.changeset/pink-rivers-clean.md
new file mode 100644
index 0000000000..d4c2fad298
--- /dev/null
+++ b/.changeset/pink-rivers-clean.md
@@ -0,0 +1,5 @@
+---
+"@electric-sql/client": patch
+---
+
+Make parser generic such that it can be parameterized with additional types supported by custom parsers.

From 63028e5bb536d311a6b13498e853c73183c63d8b Mon Sep 17 00:00:00 2001
From: Kevin De Porre <kevin@electric-sql.com>
Date: Thu, 3 Oct 2024 15:30:42 +0200
Subject: [PATCH 3/5] Update react-hooks with generics.

---
 packages/react-hooks/src/react-hooks.tsx      | 43 ++++++++++++-------
 .../react-hooks/test/react-hooks.test-d.ts    |  3 ++
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/packages/react-hooks/src/react-hooks.tsx b/packages/react-hooks/src/react-hooks.tsx
index 2640bacc8e..adcdb77fb6 100644
--- a/packages/react-hooks/src/react-hooks.tsx
+++ b/packages/react-hooks/src/react-hooks.tsx
@@ -3,15 +3,19 @@ import {
   ShapeStream,
   ShapeStreamOptions,
   Row,
+  GetExtensions,
 } from '@electric-sql/client'
 import React from 'react'
 import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/with-selector.js'
 
-const streamCache = new Map<string, ShapeStream>()
-const shapeCache = new Map<ShapeStream, Shape>()
+type UnknownShape = Shape<Row<unknown>>
+type UnknownShapeStream = ShapeStream<Row<unknown>>
 
-export async function preloadShape<T extends Row = Row>(
-  options: ShapeStreamOptions
+const streamCache = new Map<string, UnknownShapeStream>()
+const shapeCache = new Map<UnknownShapeStream, UnknownShape>()
+
+export async function preloadShape<T extends Row<unknown> = Row>(
+  options: ShapeStreamOptions<GetExtensions<T>>
 ): Promise<Shape<T>> {
   const shapeStream = getShapeStream<T>(options)
   const shape = getShape<T>(shapeStream)
@@ -19,12 +23,12 @@ export async function preloadShape<T extends Row = Row>(
   return shape
 }
 
-export function sortedOptionsHash(options: ShapeStreamOptions): string {
+export function sortedOptionsHash<T>(options: ShapeStreamOptions<T>): string {
   return JSON.stringify(options, Object.keys(options).sort())
 }
 
-export function getShapeStream<T extends Row = Row>(
-  options: ShapeStreamOptions
+export function getShapeStream<T extends Row<unknown>>(
+  options: ShapeStreamOptions<GetExtensions<T>>
 ): ShapeStream<T> {
   const shapeHash = sortedOptionsHash(options)
 
@@ -42,7 +46,9 @@ export function getShapeStream<T extends Row = Row>(
   }
 }
 
-export function getShape<T extends Row>(shapeStream: ShapeStream<T>): Shape<T> {
+export function getShape<T extends Row<unknown>>(
+  shapeStream: ShapeStream<T>
+): Shape<T> {
   // If the stream is already cached, return
   if (shapeCache.has(shapeStream)) {
     // Return the ShapeStream
@@ -57,7 +63,7 @@ export function getShape<T extends Row>(shapeStream: ShapeStream<T>): Shape<T> {
   }
 }
 
-export interface UseShapeResult<T extends Row = Row> {
+export interface UseShapeResult<T extends Row<unknown> = Row> {
   /**
    * The array of rows that make up the Shape.
    * @type {T[]}
@@ -76,14 +82,19 @@ export interface UseShapeResult<T extends Row = Row> {
   isError: boolean
 }
 
-function shapeSubscribe<T extends Row>(shape: Shape<T>, callback: () => void) {
+function shapeSubscribe<T extends Row<unknown>>(
+  shape: Shape<T>,
+  callback: () => void
+) {
   const unsubscribe = shape.subscribe(callback)
   return () => {
     unsubscribe()
   }
 }
 
-function parseShapeData<T extends Row>(shape: Shape<T>): UseShapeResult<T> {
+function parseShapeData<T extends Row<unknown>>(
+  shape: Shape<T>
+): UseShapeResult<T> {
   return {
     data: [...shape.valueSync.values()],
     isLoading: shape.isLoading(),
@@ -98,19 +109,21 @@ function identity<T>(arg: T): T {
   return arg
 }
 
-interface UseShapeOptions<SourceData extends Row, Selection>
-  extends ShapeStreamOptions {
+interface UseShapeOptions<SourceData extends Row<unknown>, Selection>
+  extends ShapeStreamOptions<GetExtensions<SourceData>> {
   selector?: (value: UseShapeResult<SourceData>) => Selection
 }
 
 export function useShape<
-  SourceData extends Row = Row,
+  SourceData extends Row<unknown> = Row,
   Selection = UseShapeResult<SourceData>,
 >({
   selector = identity as (arg: UseShapeResult<SourceData>) => Selection,
   ...options
 }: UseShapeOptions<SourceData, Selection>): Selection {
-  const shapeStream = getShapeStream<SourceData>(options as ShapeStreamOptions)
+  const shapeStream = getShapeStream<SourceData>(
+    options as ShapeStreamOptions<GetExtensions<SourceData>>
+  )
   const shape = getShape<SourceData>(shapeStream)
 
   const useShapeData = React.useMemo(() => {
diff --git a/packages/react-hooks/test/react-hooks.test-d.ts b/packages/react-hooks/test/react-hooks.test-d.ts
index 4df3e371dd..6f7c934db8 100644
--- a/packages/react-hooks/test/react-hooks.test-d.ts
+++ b/packages/react-hooks/test/react-hooks.test-d.ts
@@ -15,6 +15,7 @@ describe(`useShape`, () => {
     foo: number
     bar: boolean
     baz: string
+    ts: Date
   }
 
   it(`should infer correct return type when a selector is provided`, () => {
@@ -25,6 +26,7 @@ describe(`useShape`, () => {
           foo: 5,
           bar: true,
           baz: `str`,
+          ts: new Date(),
         }
       },
     })
@@ -41,6 +43,7 @@ describe(`useShape`, () => {
           foo: 5,
           bar: true,
           baz: `str`,
+          ts: new Date(),
         }
       },
     })

From f0c486f853487b61beb677440899670ddb91553a Mon Sep 17 00:00:00 2001
From: Kevin De Porre <kevin@electric-sql.com>
Date: Thu, 3 Oct 2024 16:22:48 +0200
Subject: [PATCH 4/5] Provide default type for ShapeStreamOptions' type
 parameter

---
 packages/typescript-client/src/client.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/typescript-client/src/client.ts b/packages/typescript-client/src/client.ts
index f86eeb4d29..3366a8ebd7 100644
--- a/packages/typescript-client/src/client.ts
+++ b/packages/typescript-client/src/client.ts
@@ -28,7 +28,7 @@ import {
 /**
  * Options for constructing a ShapeStream.
  */
-export interface ShapeStreamOptions<T> {
+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`

From b5f2321bf04963fb41ab32f5430355c00e56caeb Mon Sep 17 00:00:00 2001
From: Kevin De Porre <kevin@electric-sql.com>
Date: Thu, 3 Oct 2024 16:25:18 +0200
Subject: [PATCH 5/5] Update tanstack example

---
 examples/tanstack-example/src/Example.tsx     | 2 +-
 examples/tanstack-example/src/match-stream.ts | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/tanstack-example/src/Example.tsx b/examples/tanstack-example/src/Example.tsx
index cc81aeb09e..3ae8f9156d 100644
--- a/examples/tanstack-example/src/Example.tsx
+++ b/examples/tanstack-example/src/Example.tsx
@@ -61,7 +61,7 @@ export const Example = () => {
   const { data: items } = useShape<Item>(itemShape())
   const submissions: Item[] = useMutationState({
     filters: { status: `pending` },
-    select: (mutation) => mutation.state.context as Item | undefined,
+    select: (mutation) => mutation.state.context as Item,
   }).filter((item) => item !== undefined)
 
   const { mutateAsync: addItemMut } = useMutation({
diff --git a/examples/tanstack-example/src/match-stream.ts b/examples/tanstack-example/src/match-stream.ts
index 6c5372bf9b..195f2a6080 100644
--- a/examples/tanstack-example/src/match-stream.ts
+++ b/examples/tanstack-example/src/match-stream.ts
@@ -5,7 +5,7 @@ import {
   isChangeMessage,
 } from "@electric-sql/client"
 
-export async function matchStream<T extends Row>({
+export async function matchStream<T extends Row<unknown>>({
   stream,
   operations,
   matchFn,