From 6762b807c8e31c76a3f6eeac0a6e2461cb0d5d74 Mon Sep 17 00:00:00 2001
From: SandroMaglione <lass.maglio@gmail.com>
Date: Fri, 10 Jan 2025 10:36:46 +0100
Subject: [PATCH 1/2] live query return `Promise` with results

---
 packages/pglite-react/src/hooks.ts | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/packages/pglite-react/src/hooks.ts b/packages/pglite-react/src/hooks.ts
index 73d9ced6..7e3b0380 100644
--- a/packages/pglite-react/src/hooks.ts
+++ b/packages/pglite-react/src/hooks.ts
@@ -18,7 +18,7 @@ function paramsEqual(
 }
 
 function useLiveQueryImpl<T = { [key: string]: unknown }>(
-  query: string | LiveQuery<T> | Promise<LiveQuery<T>>,
+  query: string | LiveQuery<T> | Promise<LiveQuery<T>> | (() => Promise<T[]>),
   params: unknown[] | undefined | null,
   key?: string,
 ): Omit<LiveQueryResults<T>, 'affectedRows'> | undefined {
@@ -27,7 +27,11 @@ function useLiveQueryImpl<T = { [key: string]: unknown }>(
   const liveQueryRef = useRef<LiveQuery<T> | undefined>()
   let liveQuery: LiveQuery<T> | undefined
   let liveQueryChanged = false
-  if (!(typeof query === 'string') && !(query instanceof Promise)) {
+  if (
+    !(typeof query === 'string') &&
+    !(typeof query === 'function') &&
+    !(query instanceof Promise)
+  ) {
     liveQuery = query
     liveQueryChanged = liveQueryRef.current !== liveQuery
     liveQueryRef.current = liveQuery
@@ -59,6 +63,17 @@ function useLiveQueryImpl<T = { [key: string]: unknown }>(
         cancelled = true
         ret.then(({ unsubscribe }) => unsubscribe())
       }
+    } else if (typeof query === 'function') {
+      query.then((liveQuery) => {
+        if (cancelled) return
+        liveQueryRef.current = liveQuery
+        setResults(liveQuery.initialResults)
+        liveQuery.subscribe(cb)
+      })
+      return () => {
+        cancelled = true
+        liveQueryRef.current?.unsubscribe(cb)
+      }
     } else if (query instanceof Promise) {
       query.then((liveQuery) => {
         if (cancelled) return
@@ -112,7 +127,11 @@ export function useLiveQuery<T = { [key: string]: unknown }>(
 ): LiveQueryResults<T> | undefined
 
 export function useLiveQuery<T = { [key: string]: unknown }>(
-  query: string | LiveQuery<T> | Promise<LiveQuery<T>>,
+  liveQueryPromise: () => Promise<T[]>,
+): LiveQueryResults<T> | undefined
+
+export function useLiveQuery<T = { [key: string]: unknown }>(
+  query: string | LiveQuery<T> | Promise<LiveQuery<T>> | (() => Promise<T[]>),
   params?: unknown[] | null,
 ): LiveQueryResults<T> | undefined {
   return useLiveQueryImpl<T>(query, params)

From 28ea4596b7d89f07f6ddd9caaabfb23971eb7c74 Mon Sep 17 00:00:00 2001
From: SandroMaglione <lass.maglio@gmail.com>
Date: Thu, 16 Jan 2025 16:01:20 +0100
Subject: [PATCH 2/2] `() => Promise<string | LiveQuery<T>>` live query
 implementation

---
 packages/pglite-react/src/hooks.ts | 35 +++++++++++++++++++++++++-----
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/packages/pglite-react/src/hooks.ts b/packages/pglite-react/src/hooks.ts
index 7e3b0380..92f29855 100644
--- a/packages/pglite-react/src/hooks.ts
+++ b/packages/pglite-react/src/hooks.ts
@@ -18,7 +18,11 @@ function paramsEqual(
 }
 
 function useLiveQueryImpl<T = { [key: string]: unknown }>(
-  query: string | LiveQuery<T> | Promise<LiveQuery<T>> | (() => Promise<T[]>),
+  query:
+    | string
+    | LiveQuery<T>
+    | Promise<LiveQuery<T>>
+    | (() => Promise<string | LiveQuery<T>>),
   params: unknown[] | undefined | null,
   key?: string,
 ): Omit<LiveQueryResults<T>, 'affectedRows'> | undefined {
@@ -53,6 +57,7 @@ function useLiveQueryImpl<T = { [key: string]: unknown }>(
       if (cancelled) return
       setResults(results)
     }
+
     if (typeof query === 'string') {
       const ret =
         key !== undefined
@@ -63,24 +68,37 @@ function useLiveQueryImpl<T = { [key: string]: unknown }>(
         cancelled = true
         ret.then(({ unsubscribe }) => unsubscribe())
       }
-    } else if (typeof query === 'function') {
+    } else if (query instanceof Promise) {
       query.then((liveQuery) => {
         if (cancelled) return
         liveQueryRef.current = liveQuery
         setResults(liveQuery.initialResults)
         liveQuery.subscribe(cb)
       })
+
       return () => {
         cancelled = true
         liveQueryRef.current?.unsubscribe(cb)
       }
-    } else if (query instanceof Promise) {
-      query.then((liveQuery) => {
+    } else if (typeof query === 'function') {
+      query().then((liveQueryOrString) => {
         if (cancelled) return
+        const liveQuery =
+          typeof liveQueryOrString === 'string'
+            ? key !== undefined
+              ? db.live.incrementalQuery<T>(
+                  liveQueryOrString,
+                  currentParams,
+                  key,
+                  cb,
+                )
+              : db.live.query<T>(liveQueryOrString, currentParams, cb)
+            : liveQueryOrString
         liveQueryRef.current = liveQuery
         setResults(liveQuery.initialResults)
         liveQuery.subscribe(cb)
       })
+
       return () => {
         cancelled = true
         liveQueryRef.current?.unsubscribe(cb)
@@ -88,6 +106,7 @@ function useLiveQueryImpl<T = { [key: string]: unknown }>(
     } else if (liveQuery) {
       setResults(liveQuery.initialResults)
       liveQuery.subscribe(cb)
+
       return () => {
         cancelled = true
         liveQuery.unsubscribe(cb)
@@ -127,11 +146,15 @@ export function useLiveQuery<T = { [key: string]: unknown }>(
 ): LiveQueryResults<T> | undefined
 
 export function useLiveQuery<T = { [key: string]: unknown }>(
-  liveQueryPromise: () => Promise<T[]>,
+  liveQueryPromise: () => Promise<string | LiveQuery<T>>,
 ): LiveQueryResults<T> | undefined
 
 export function useLiveQuery<T = { [key: string]: unknown }>(
-  query: string | LiveQuery<T> | Promise<LiveQuery<T>> | (() => Promise<T[]>),
+  query:
+    | string
+    | LiveQuery<T>
+    | Promise<LiveQuery<T>>
+    | (() => Promise<string | LiveQuery<T>>),
   params?: unknown[] | null,
 ): LiveQueryResults<T> | undefined {
   return useLiveQueryImpl<T>(query, params)