Add beforeQuery hook to query method (#39)

* Add beforeQuery hook

* Fix lints

* Bump version

* Adding onDuckdbShutdown to be called

Author: itsTalwar

.github/workflows/publish-package.yml

@@ -23,21 +23,25 @@ jobs:
           cd ./dist/meerkat-dbm
           echo "//npm.pkg.github.com/:_authToken=${{secrets.NPM_PACKAGE_PUBLISH_TOKEN}}" >> .npmrc
           npm publish
+        continue-on-error: true
 
       - run: |
           cd ./dist/meerkat-core
           echo "//npm.pkg.github.com/:_authToken=${{secrets.NPM_PACKAGE_PUBLISH_TOKEN}}" >> .npmrc
           npm publish
+        continue-on-error: true
 
       - run: |
           cd ./dist/meerkat-node
           echo "//npm.pkg.github.com/:_authToken=${{secrets.NPM_PACKAGE_PUBLISH_TOKEN}}" >> .npmrc
           npm publish
+        continue-on-error: true
           
       - run: |
           cd ./dist/meerkat-browser
           echo "//npm.pkg.github.com/:_authToken=${{secrets.NPM_PACKAGE_PUBLISH_TOKEN}}" >> .npmrc
           npm publish
-          
+        continue-on-error: true
+
         env:
           NODE_AUTH_TOKEN: ${{secrets.NPM_PACKAGE_PUBLISH_TOKEN}}

meerkat-dbm/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@devrev/meerkat-dbm",
-  "version": "0.0.132",
+  "version": "0.0.133",
   "dependencies": {
     "tslib": "^2.3.0",
     "@duckdb/duckdb-wasm": "^1.28.0",

meerkat-dbm/src/dbm/dbm.spec.ts

@@ -71,6 +71,35 @@ export class MockFileManager implements FileManagerType {
       delete this.fileBufferStore[fileName];
     }
   }
+
+  async getFilesNameForTables(tableNames: string[]): Promise<
+    {
+      tableName: string;
+      files: string[];
+    }[]
+  > {
+    const data: {
+      tableName: string;
+      files: string[];
+    }[] = [];
+
+    for (const tableName of tableNames) {
+      const files: string[] = [];
+
+      for (const key in this.fileBufferStore) {
+        if (this.fileBufferStore[key].tableName === tableName) {
+          files.push(key);
+        }
+      }
+
+      data.push({
+        tableName,
+        files,
+      });
+    }
+
+    return data;
+  }
 }
 
 const mockDB = {
@@ -135,6 +164,33 @@ describe('DBM', () => {
   });
 
   describe('queryWithTableNames', () => {
+    it('should call the beforeQuery hook', async () => {
+      const beforeQuery = jest.fn();
+
+      await fileManager.registerFileBuffer({
+        fileName: 'file1',
+        tableName: 'table1',
+        buffer: new Uint8Array(),
+      });
+
+      const result = await dbm.queryWithTableNames(
+        'SELECT * FROM table1',
+        ['table1'],
+        {
+          beforeQuery,
+        }
+      );
+
+      expect(beforeQuery).toBeCalledTimes(1);
+
+      expect(beforeQuery).toBeCalledWith([
+        {
+          tableName: 'table1',
+          files: ['file1'],
+        },
+      ]);
+    });
+
     it('should execute a query with table names', async () => {
       const result = await dbm.queryWithTableNames('SELECT * FROM table1', [
         'table1',
@@ -195,6 +251,8 @@ describe('DBM', () => {
       // If instanceManager.terminateDB is a method
       jest.spyOn(instanceManager, 'terminateDB');
 
+      const onDuckdbShutdown = jest.fn();
+
       // If instanceManager.terminateDB is a function
       instanceManager.terminateDB = jest.fn();
       const options: DBMConstructorOptions = {
@@ -204,6 +262,7 @@ describe('DBM', () => {
         onEvent: (event) => {
           console.log(event);
         },
+        onDuckdbShutdown: onDuckdbShutdown,
         options: {
           shutdownInactiveTime: 100,
         },
@@ -239,6 +298,12 @@ describe('DBM', () => {
        * wait for 200ms
        */
       await new Promise((resolve) => setTimeout(resolve, 200));
+
+      /**
+       * Expect onDuckdbShutdown to be called
+       */
+      expect(onDuckdbShutdown).toBeCalled();
+
       /**
        * Expect instanceManager.terminateDB to be called
        */

meerkat-dbm/src/dbm/dbm.ts

@@ -8,6 +8,7 @@ export interface DBMConstructorOptions {
   fileManager: FileManagerType;
   logger: DBMLogger;
   onEvent?: (event: DBMEvent) => void;
+  onDuckdbShutdown?: () => void;
   instanceManager: InstanceManagerType;
   options?: {
     /**
@@ -16,6 +17,18 @@ export interface DBMConstructorOptions {
     shutdownInactiveTime?: number;
   };
 }
+
+type BeforeQueryHook = (
+  data: {
+    tableName: string;
+    files: string[];
+  }[]
+) => Promise<void>;
+
+type QueryOptions = {
+  beforeQuery: BeforeQueryHook;
+};
+
 export class DBM {
   private fileManager: FileManagerType;
   private instanceManager: InstanceManagerType;
@@ -31,25 +44,39 @@ export class DBM {
      * Timestamp when the query was added to the queue
      */
     timestamp: number;
+    options?: { beforeQuery: BeforeQueryHook };
   }[] = [];
+  private beforeQuery?: ({
+    query,
+    tableByFiles,
+  }: {
+    query: string;
+    tableByFiles: {
+      tableName: string;
+      files: string[];
+    }[];
+  }) => Promise<void>;
   private queryQueueRunning = false;
   private logger: DBMLogger;
   private onEvent?: (event: DBMEvent) => void;
   private options: DBMConstructorOptions['options'];
   private terminateDBTimeout: NodeJS.Timeout | null = null;
+  private onDuckdbShutdown?: () => void;
 
   constructor({
     fileManager,
     logger,
     onEvent,
     options,
     instanceManager,
+    onDuckdbShutdown,
   }: DBMConstructorOptions) {
     this.fileManager = fileManager;
     this.logger = logger;
     this.onEvent = onEvent;
     this.options = options;
     this.instanceManager = instanceManager;
+    this.onDuckdbShutdown = onDuckdbShutdown;
   }
 
   private async _shutdown() {
@@ -58,6 +85,9 @@ export class DBM {
       this.connection = null;
     }
     this.logger.debug('Shutting down the DB');
+    if (this.onDuckdbShutdown) {
+      this.onDuckdbShutdown();
+    }
     await this.instanceManager.terminateDB();
   }
 
@@ -97,7 +127,11 @@ export class DBM {
     return this.connection;
   }
 
-  private async _queryWithTableNames(query: string, tableNames: string[]) {
+  private async _queryWithTableNames(
+    query: string,
+    tableNames: string[],
+    options?: QueryOptions
+  ) {
     /**
      * Load all the files into the database
      */
@@ -117,6 +151,17 @@ export class DBM {
       duration: endMountTime - startMountTime,
     });
 
+    const tablesFileData = await this.fileManager.getFilesNameForTables(
+      tableNames
+    );
+
+    /**
+     * Execute the beforeQuery hook
+     */
+    if (options?.beforeQuery) {
+      await options.beforeQuery(tablesFileData);
+    }
+
     /**
      * Execute the query
      */
@@ -189,11 +234,11 @@ export class DBM {
     /**
      * Get the first query
      */
-    const query = this.queriesQueue.shift();
+    const queueElement = this.queriesQueue.shift();
     /**
      * If there is no query, stop the queue
      */
-    if (!query) {
+    if (!queueElement) {
       await this._stopQueryQueue();
       return;
     }
@@ -202,41 +247,42 @@ export class DBM {
       const startTime = Date.now();
       this.logger.debug(
         'Time since query was added to the queue:',
-        startTime - query.timestamp,
+        startTime - queueElement.timestamp,
         'ms',
-        query.query
+        queueElement.query
       );
 
       this._emitEvent({
         event_name: 'query_queue_duration',
-        duration: startTime - query.timestamp,
+        duration: startTime - queueElement.timestamp,
       });
 
       /**
        * Execute the query
        */
       const result = await this._queryWithTableNames(
-        query.query,
-        query.tableNames
+        queueElement.query,
+        queueElement.tableNames,
+        queueElement.options
       );
       const endTime = Date.now();
 
       this.logger.debug(
         'Total time spent along with queue time',
-        endTime - query.timestamp,
+        endTime - queueElement.timestamp,
         'ms',
-        query.query
+        queueElement.query
       );
       /**
        * Resolve the promise
        */
-      query.promise.resolve(result);
+      queueElement.promise.resolve(result);
     } catch (error) {
-      this.logger.warn('Error while executing query:', query.query);
+      this.logger.warn('Error while executing query:', queueElement.query);
       /**
        * Reject the promise, so the caller can catch the error
        */
-      query.promise.reject(error);
+      queueElement.promise.reject(error);
     }
 
     /**
@@ -266,7 +312,11 @@ export class DBM {
     return this.queryQueueRunning;
   }
 
-  public async queryWithTableNames(query: string, tableNames: string[]) {
+  public async queryWithTableNames(
+    query: string,
+    tableNames: string[],
+    options?: { beforeQuery: BeforeQueryHook }
+  ) {
     const promise = new Promise((resolve, reject) => {
       this.queriesQueue.push({
         query,
@@ -276,6 +326,7 @@ export class DBM {
           reject,
         },
         timestamp: Date.now(),
+        options,
       });
     });
     this._startQueryQueue();

meerkat-dbm/src/file-manager/file-manager-type.ts

@@ -17,6 +17,12 @@ export interface FileManagerType {
   unmountFileBufferByTableNames: (tableName: string[]) => Promise<void>;
   getFilesByTableName(tableName: string): Promise<FileData[]>;
   dropFilesByTableName(tableName: string, fileNames: string[]): Promise<void>;
+  getFilesNameForTables(tableNames: string[]): Promise<
+    {
+      tableName: string;
+      files: string[];
+    }[]
+  >;
 }
 
 export interface FileManagerConstructorOptions {

meerkat-dbm/src/file-manager/indexed-db/indexed-db-file-manager.ts

@@ -128,6 +128,20 @@ export class IndexedDBFileManager implements FileManagerType {
     return fileData?.buffer;
   }
 
+  async getFilesNameForTables(tableNames: string[]): Promise<
+    {
+      tableName: string;
+      files: string[];
+    }[]
+  > {
+    const tableData = await this.indexedDB.tablesKey.bulkGet(tableNames);
+
+    return tableData.map((table) => ({
+      tableName: table?.tableName ?? '',
+      files: (table?.files ?? []).map((file) => file.fileName),
+    }));
+  }
+
   async mountFileBufferByTableNames(tableNames: string[]): Promise<void> {
     const tableData = await this.indexedDB.tablesKey.bulkGet(tableNames);
 

meerkat-dbm/src/file-manager/memory-file-manager.ts

@@ -46,6 +46,16 @@ export class MemoryDBFileManager implements FileManagerType {
     // not needed for memory file manager
   }
 
+  async getFilesNameForTables(tableNames: string[]): Promise<
+    {
+      tableName: string;
+      files: string[];
+    }[]
+  > {
+    // not needed for memory file manager
+    return [];
+  }
+
   async getFilesByTableName(tableName: string): Promise<FileData[]> {
     // not needed for memory file manager
     return [];