From 96d7f1d024bbbe94a70f09823ff922068293c8d7 Mon Sep 17 00:00:00 2001
From: Niva Vaani Sivakumar <niva.sivakumar@mongodb.com>
Date: Thu, 26 Jun 2025 14:39:26 -0400
Subject: [PATCH] first draft

---
 .../src/compass-query-storage.ts              | 130 ++++++++++++------
 1 file changed, 90 insertions(+), 40 deletions(-)

diff --git a/packages/my-queries-storage/src/compass-query-storage.ts b/packages/my-queries-storage/src/compass-query-storage.ts
index e9b6b3b6169..4b64166a83f 100644
--- a/packages/my-queries-storage/src/compass-query-storage.ts
+++ b/packages/my-queries-storage/src/compass-query-storage.ts
@@ -12,69 +12,119 @@ export type QueryStorageOptions = {
   basepath?: string;
 };
 
-export interface QueryStorageBackend<TData> {
-  loadAll(namespace?: string): Promise<TData[]>;
+export interface QueryDocumentBase {
+  _id: string;
+  _lastExecuted: Date;
+  _ns?: string;
+}
+
+// this will be either UserDataBackend (for desktop filesystem storage) or
+// ApiBackend (for Compass-Web cloud storage)
+export interface QueryStorageBackend<TData extends QueryDocumentBase> {
+  loadAll(): Promise<TData[]>;
+  readOne(id: string): Promise<TData | undefined>;
   updateAttributes(id: string, data: Partial<TData>): Promise<TData>;
   delete(id: string): Promise<boolean>;
-  saveQuery(data: Omit<TData, '_id' | '_lastExecuted'>): Promise<void>;
+  saveQuery(data: TData): Promise<void>;
+}
+
+export abstract class CompassQueryStorage<TData extends QueryDocumentBase>
+  implements QueryStorageBackend<TData>
+{
+  protected backend: QueryStorageBackend<TData>;
+
+  constructor(backend: QueryStorageBackend<TData>) {
+    this.backend = backend;
+  }
+
+  async loadAll(namespace?: string): Promise<TData[]> {
+    const items = await this.backend.loadAll();
+    return items
+      .sort((a, b) => b._lastExecuted.getTime() - a._lastExecuted.getTime())
+      .filter((x) => !namespace || x._ns === namespace);
+  }
+
+  async readOne(id: string): Promise<TData | undefined> {
+    return this.backend.readOne(id);
+  }
+
+  async updateAttributes(id: string, patch: Partial<TData>): Promise<TData> {
+    return this.backend.updateAttributes(id, patch);
+  }
+
+  async delete(id: string): Promise<boolean> {
+    return this.backend.delete(id);
+  }
+
+  abstract saveQuery(data: Omit<TData, '_id' | '_lastExecuted'>): Promise<void>;
 }
 
-export abstract class CompassQueryStorage<
+// UserDataBackend preserves all major functionality and adapts UserData to the QueryStorageBackend interface
+export class UserDataBackend<
   TSchema extends z.Schema,
-  TData extends z.output<TSchema> = z.output<TSchema>
+  TData extends z.output<TSchema> & QueryDocumentBase
 > implements QueryStorageBackend<TData>
 {
-  protected readonly userData: UserData<TSchema>;
-  constructor(
-    schemaValidator: TSchema,
-    protected readonly folder: string,
-    protected readonly options: QueryStorageOptions
-  ) {
-    this.userData = new UserData(schemaValidator, {
-      subdir: folder,
-      basePath: options.basepath,
-      serialize: (content) => EJSON.stringify(content, undefined, 2),
-      deserialize: (content) => EJSON.parse(content),
-    });
+  private userData: UserData<TSchema>;
+
+  constructor(userData: UserData<TSchema>) {
+    this.userData = userData;
   }
 
-  async loadAll(namespace?: string): Promise<TData[]> {
+  async loadAll(): Promise<TData[]> {
     try {
       const { data } = await this.userData.readAll();
-      const sortedData = data
-        .sort((a, b) => {
-          return b._lastExecuted.getTime() - a._lastExecuted.getTime();
-        })
-        .filter((x) => !namespace || x._ns === namespace);
-      return sortedData;
+      return data as TData[];
     } catch {
       return [];
     }
   }
 
-  async updateAttributes(id: string, data: Partial<TData>): Promise<TData> {
-    await this.userData.write(id, {
-      ...((await this.userData.readOne(id)) ?? {}),
-      ...data,
-    });
-    return await this.userData.readOne(id);
+  async readOne(id: string): Promise<TData | undefined> {
+    try {
+      const result = await this.userData.readOne(id);
+      return result as TData | undefined;
+    } catch {
+      // not sure if this is what I should be returning here
+      return undefined;
+    }
   }
 
-  async delete(id: string) {
+  async updateAttributes(id: string, patch: Partial<TData>): Promise<TData> {
+    // this is kind of a different approach than the original UserData
+    // just more error handling
+    const current = await this.readOne(id);
+    if (!current) {
+      // will throwing Errors here be bubbled up properly?
+      // or should I return undefined or null?
+      throw new Error(`Record not found for id ${id}`);
+    }
+    const updated = { ...current, ...patch };
+    await this.userData.write(id, updated as z.input<TSchema>);
+    const after = await this.readOne(id);
+    if (!after) {
+      throw new Error(`Failed to update record for id ${id}`);
+    }
+    return after;
+  }
+
+  async delete(id: string): Promise<boolean> {
     return await this.userData.delete(id);
   }
 
-  abstract saveQuery(data: any): Promise<void>;
+  async saveQuery(data: TData): Promise<void> {
+    await this.userData.write(data._id, data as z.input<TSchema>);
+  }
 }
 
 export class CompassRecentQueryStorage
-  extends CompassQueryStorage<typeof RecentQuerySchema, RecentQuery>
+  extends CompassQueryStorage<RecentQuery>
   implements RecentQueryStorage
 {
   private readonly maxAllowedQueries = 30;
 
-  constructor(options: QueryStorageOptions = {}) {
-    super(RecentQuerySchema, 'RecentQueries', options);
+  constructor(backend: QueryStorageBackend<RecentQuery>) {
+    super(backend);
   }
 
   async saveQuery(
@@ -93,16 +143,16 @@ export class CompassRecentQueryStorage
       _id,
       _lastExecuted: new Date(),
     };
-    await this.userData.write(_id, recentQuery);
+    await this.backend.saveQuery(recentQuery);
   }
 }
 
 export class CompassFavoriteQueryStorage
-  extends CompassQueryStorage<typeof FavoriteQuerySchema, FavoriteQuery>
+  extends CompassQueryStorage<FavoriteQuery>
   implements FavoriteQueryStorage
 {
-  constructor(options: QueryStorageOptions = {}) {
-    super(FavoriteQuerySchema, 'FavoriteQueries', options);
+  constructor(backend: QueryStorageBackend<FavoriteQuery>) {
+    super(backend);
   }
 
   async saveQuery(
@@ -118,6 +168,6 @@ export class CompassFavoriteQueryStorage
       _lastExecuted: new Date(),
       _dateSaved: new Date(),
     };
-    await this.userData.write(_id, favoriteQuery);
+    await this.backend.saveQuery(favoriteQuery);
   }
 }