From 782b70ecc8b376974f0e3b963d1b9a5463e23aa9 Mon Sep 17 00:00:00 2001
From: Christiaan Landman <chriz.ek@gmail.com>
Date: Thu, 17 Oct 2024 10:40:30 +0200
Subject: [PATCH] WatchedQuery now extends BaseObserver.

---
 packages/react/src/WatchedQuery.ts           | 23 ++++++++++----------
 packages/react/src/hooks/useSuspenseQuery.ts | 10 +++++----
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/packages/react/src/WatchedQuery.ts b/packages/react/src/WatchedQuery.ts
index 79af9738..41770f6e 100644
--- a/packages/react/src/WatchedQuery.ts
+++ b/packages/react/src/WatchedQuery.ts
@@ -1,4 +1,4 @@
-import { AbstractPowerSyncDatabase, CompilableQuery } from '@powersync/common';
+import { AbstractPowerSyncDatabase, BaseListener, BaseObserver, CompilableQuery } from '@powersync/common';
 import { AdditionalOptions } from './hooks/useQuery';
 
 export class Query<T> {
@@ -7,9 +7,11 @@ export class Query<T> {
   queryParameters: any[];
 }
 
-export class WatchedQuery {
-  listeners = new Set<() => void>();
+export interface WatchedQueryListener extends BaseListener {
+  onUpdate: () => void;
+}
 
+export class WatchedQuery extends BaseObserver<WatchedQueryListener> {
   readyPromise: Promise<void>;
   isReady: boolean = false;
   currentData: any[] | undefined;
@@ -27,6 +29,7 @@ export class WatchedQuery {
   private disposer: () => void;
 
   constructor(db: AbstractPowerSyncDatabase, query: Query<unknown>, options: AdditionalOptions, disposer: () => void) {
+    super();
     this.db = db;
     this.query = query;
     this.options = options;
@@ -65,12 +68,12 @@ export class WatchedQuery {
     return release;
   }
 
-  addListener(l: () => void) {
-    this.listeners.add(l);
+  registerListener(listener: Partial<WatchedQueryListener>): () => void {
+    const disposer = super.registerListener(listener);
 
     this.maybeListen();
     return () => {
-      this.listeners.delete(l);
+      disposer();
       this.maybeDispose();
     };
   }
@@ -142,9 +145,7 @@ export class WatchedQuery {
     this.currentError = undefined;
     this.resolveReady?.();
 
-    for (let listener of this.listeners) {
-      listener();
-    }
+    this.iterateListeners((l) => l?.onUpdate());
   }
 
   private setError(error: any) {
@@ -153,9 +154,7 @@ export class WatchedQuery {
     this.currentError = error;
     this.resolveReady?.();
 
-    for (let listener of this.listeners) {
-      listener();
-    }
+    this.iterateListeners((l) => l?.onUpdate());
   }
 
   private maybeDispose() {
diff --git a/packages/react/src/hooks/useSuspenseQuery.ts b/packages/react/src/hooks/useSuspenseQuery.ts
index e43bdd23..7cd43bfc 100644
--- a/packages/react/src/hooks/useSuspenseQuery.ts
+++ b/packages/react/src/hooks/useSuspenseQuery.ts
@@ -48,10 +48,12 @@ export const useSuspenseQuery = <T = any>(
   const [_counter, setUpdateCounter] = React.useState(0);
 
   React.useEffect(() => {
-    const dispose = q.addListener(() => {
-      setUpdateCounter((counter) => {
-        return counter + 1;
-      });
+    const dispose = q.registerListener({
+      onUpdate: () => {
+        setUpdateCounter((counter) => {
+          return counter + 1;
+        });
+      }
     });
 
     releaseTemporaryHold.current?.();