Fix useLiveQuery with typed PGlite providers

.changeset/typed-providers-query.md

@@ -0,0 +1,5 @@
+---
+'@electric-sql/pglite-react': patch
+---
+
+Make live query hooks use databases supplied by typed PGlite providers.

packages/pglite-react/src/provider.tsx

@@ -16,11 +16,12 @@ interface PGliteProviderSet<T extends PGliteWithLive> {
   usePGlite: UsePGlite<T>
 }
 
-/**
- * Create a typed set of {@link PGliteProvider} and {@link usePGlite}.
- */
-function makePGliteProvider<T extends PGliteWithLive>(): PGliteProviderSet<T> {
-  const ctx = createContext<T | undefined>(undefined)
+const defaultContext = createContext<PGliteWithLive | undefined>(undefined)
+
+function makePGliteProviderSet<T extends PGliteWithLive>(
+  ctx: React.Context<T | undefined>,
+  bridgeDefaultContext: boolean,
+): PGliteProviderSet<T> {
   return {
     usePGlite: ((db?: T) => {
       const dbProvided = useContext(ctx)
@@ -36,11 +37,28 @@ function makePGliteProvider<T extends PGliteWithLive>(): PGliteProviderSet<T> {
       return dbProvided
     }) as UsePGlite<T>,
     PGliteProvider: ({ children, db }: Props<T>) => {
-      return <ctx.Provider value={db}>{children}</ctx.Provider>
+      const provider = <ctx.Provider value={db}>{children}</ctx.Provider>
+
+      return bridgeDefaultContext ? (
+        <defaultContext.Provider value={db}>{provider}</defaultContext.Provider>
+      ) : (
+        provider
+      )
     },
   }
 }
 
-const { PGliteProvider, usePGlite } = makePGliteProvider<PGliteWithLive>()
+/**
+ * Create a typed set of {@link PGliteProvider} and {@link usePGlite}.
+ */
+function makePGliteProvider<T extends PGliteWithLive>(): PGliteProviderSet<T> {
+  const ctx = createContext<T | undefined>(undefined)
+  return makePGliteProviderSet(ctx, true)
+}
+
+const { PGliteProvider, usePGlite } = makePGliteProviderSet(
+  defaultContext,
+  false,
+)
 
 export { makePGliteProvider, PGliteProvider, usePGlite }

packages/pglite-react/test/provider.test.tsx

@@ -1,10 +1,15 @@
-import { describe, it, expect } from 'vitest'
+import { describe, it, expect, vi } from 'vitest'
 import { renderHook } from '@testing-library/react'
 import { waitFor } from '@testing-library/dom'
 import React from 'react'
 import { PGlite } from '@electric-sql/pglite'
 import { live, PGliteWithLive } from '@electric-sql/pglite/live'
-import { makePGliteProvider, PGliteProvider, usePGlite } from '../src'
+import {
+  makePGliteProvider,
+  PGliteProvider,
+  useLiveQuery,
+  usePGlite,
+} from '../src'
 
 describe('provider', () => {
   it('can receive PGlite', async () => {
@@ -40,4 +45,43 @@ describe('provider', () => {
 
     await waitFor(() => expect(result.current).toBe(db))
   })
+
+  it('makes useLiveQuery available under a typed provider', async () => {
+    const initialResults = {
+      rows: [{ value: 1 }],
+      fields: [{ name: 'value', dataTypeID: 23 }],
+    }
+    const query = vi.fn(
+      (
+        _query: string,
+        _params: unknown[] | undefined | null,
+        callback: (results: typeof initialResults) => void,
+      ) => {
+        callback(initialResults)
+        return Promise.resolve({
+          initialResults,
+          subscribe: vi.fn(),
+          unsubscribe: vi.fn(async () => undefined),
+          refresh: vi.fn(async () => undefined),
+        })
+      },
+    )
+    const db = { live: { query } } as unknown as PGliteWithLive
+    const { PGliteProvider: PGliteProviderTyped } =
+      makePGliteProvider<PGliteWithLive>()
+    const wrapper = ({ children }: { children: React.ReactNode }) => {
+      return <PGliteProviderTyped db={db}>{children}</PGliteProviderTyped>
+    }
+
+    const { result } = renderHook(() => useLiveQuery('SELECT 1 AS value'), {
+      wrapper,
+    })
+
+    await waitFor(() => expect(result.current?.rows).toEqual([{ value: 1 }]))
+    expect(query).toHaveBeenCalledWith(
+      'SELECT 1 AS value',
+      undefined,
+      expect.any(Function),
+    )
+  })
 })