chore: add getFilesByTableName and dropFilesByTableName methods in indexeddb filemanager (#33)

* add: cache time logic

* fix: test

* update: logic

* revert: tests

* more changes

* add: test

* some cleanup

* bump: version

* update: name

* remove: filesForTable

* update: file manager type

* fix: build

* fix: mock manager

* update: name

Author: vpbs2

meerkat-dbm/package.json

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

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

@@ -2,6 +2,7 @@ import { AsyncDuckDB } from '@duckdb/duckdb-wasm';
 import log from 'loglevel';
 import {
   FileBufferStore,
+  FileData,
   FileManagerType,
 } from '../file-manager/file-manager-type';
 import { DBM, DBMConstructorOptions } from './dbm';
@@ -49,6 +50,27 @@ export class MockFileManager implements FileManagerType {
       }
     }
   }
+
+  async getFilesByTableName(tableName: string): Promise<FileData[]> {
+    const files: FileData[] = [];
+
+    for (const key in this.fileBufferStore) {
+      if (this.fileBufferStore[key].tableName === tableName) {
+        files.push({ fileName: key });
+      }
+    }
+
+    return files;
+  }
+
+  async dropFilesByTableName(
+    tableName: string,
+    fileNames: string[]
+  ): Promise<void> {
+    for (const fileName of fileNames) {
+      delete this.fileBufferStore[fileName];
+    }
+  }
 }
 
 const mockDB = {

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

@@ -15,6 +15,8 @@ export interface FileManagerType {
   getFileBuffer: (name: string) => Promise<Uint8Array | undefined>;
   mountFileBufferByTableNames: (tableName: string[]) => Promise<void>;
   unmountFileBufferByTableNames: (tableName: string[]) => Promise<void>;
+  getFilesByTableName(tableName: string): Promise<FileData[]>;
+  dropFilesByTableName(tableName: string, fileNames: string[]): Promise<void>;
 }
 
 export interface FileManagerConstructorOptions {

meerkat-dbm/src/file-manager/indexed-db/duckdb-database.ts renamed to meerkat-dbm/src/file-manager/indexed-db/duckdb-files-database.ts

@@ -6,12 +6,12 @@ import { File, Table } from '../file-manager-type';
  * https://dexie.org/docs/Version/Version.stores()#warning
  */
 
-export class DuckDBDatabase extends Dexie {
+export class DuckDBFilesDatabase extends Dexie {
   tablesKey: Dexie.Table<Table, string>;
   files: Dexie.Table<File, string>;
 
   constructor() {
-    super('DuckDBDatabase');
+    super('DuckDBFilesDatabase');
 
     this.version(1).stores({
       tablesKey: '&tableName',

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

@@ -2,7 +2,7 @@ import { AsyncDuckDB } from '@duckdb/duckdb-wasm';
 import 'fake-indexeddb/auto';
 import { InstanceManagerType } from '../../dbm/instance-manager';
 import { FILE_TYPES } from '../file-manager-type';
-import { DuckDBDatabase } from './duckdb-database';
+import { DuckDBFilesDatabase } from './duckdb-files-database';
 import { IndexedDBFileManager } from './indexed-db-file-manager';
 
 const mockDB = {
@@ -25,7 +25,7 @@ const mockDB = {
 describe('IndexedDBFileManager', () => {
   let fileManager: IndexedDBFileManager;
   let db: AsyncDuckDB;
-  let indexedDB: DuckDBDatabase;
+  let indexedDB: DuckDBFilesDatabase;
   let instanceManager: InstanceManagerType;
 
   const fileBuffer = {
@@ -65,7 +65,7 @@ describe('IndexedDBFileManager', () => {
   });
 
   beforeEach(async () => {
-    indexedDB = new DuckDBDatabase();
+    indexedDB = new DuckDBFilesDatabase();
     fileManager = new IndexedDBFileManager({
       fetchTableFileBuffers: async () => {
         return [];
@@ -124,29 +124,20 @@ describe('IndexedDBFileManager', () => {
     ]);
   });
 
-  it('should flush the database when initialized', async () => {
-    // Fetch the stored data in the indexedDB
+  it('should persist the data when the IndexedDB is reinitialized', async () => {
     const tableData = await indexedDB.tablesKey.toArray();
-    const fileBufferData = await indexedDB.files.toArray();
+    const fileData = await indexedDB.files.toArray();
 
-    /**
-     * There should be no tablesKey and no file buffers
-     */
-    expect(tableData.length).toBe(0);
-    expect(fileBufferData.length).toBe(0);
+    expect(tableData.length).toBe(2);
+    expect(fileData.length).toBe(3);
   });
 
   it('should override the file buffer when registering the same file again', async () => {
-    // Register single file buffer
-    await fileManager.registerFileBuffer(fileBuffer);
-
-    const tableData = await indexedDB.tablesKey.toArray();
     const fileBufferData1 = await indexedDB.files.toArray();
 
     /**
-     * There should be one file buffer, the buffer should be initial file buffer
+     * The buffer value should be initial file buffer
      */
-    expect(fileBufferData1.length).toBe(1);
     expect(fileBufferData1[0].buffer).toEqual(fileBuffer.buffer);
 
     // Register the same file with a different buffer
@@ -158,15 +149,56 @@ describe('IndexedDBFileManager', () => {
     const fileBufferData2 = await indexedDB.files.toArray();
 
     /**
-     * There should be one table with one file buffer
+     * The buffer value should be updated to the new buffer
      */
-    expect(tableData.length).toBe(1);
-    expect(tableData[0].files.length).toBe(1);
-
-    /**
-     * There should be one file buffer, the buffer should be updated to the new buffer
-     */
-    expect(fileBufferData2.length).toBe(1);
     expect(fileBufferData2[0].buffer).toEqual(new Uint8Array([1]));
   });
+
+  it('should return the files for a table stored', async () => {
+    const fileData = await fileManager.getFilesByTableName('taxi1');
+
+    expect(fileData).toEqual([
+      { fileName: 'taxi1.parquet', fileType: 'parquet' },
+      { fileName: 'taxi2.parquet', fileType: 'parquet' },
+    ]);
+  });
+
+  it('should drop the file buffers for a table', async () => {
+    // Drop a file when there are multiple files
+    await fileManager.dropFilesByTableName('taxi1', ['taxi1.parquet']);
+
+    const tableData1 = await indexedDB.tablesKey.toArray();
+    const fileBufferData1 = await indexedDB.files.toArray();
+
+    // Verify that the file is dropped
+    expect(tableData1[0]).toEqual({
+      tableName: 'taxi1',
+      files: [
+        {
+          fileName: 'taxi2.parquet',
+          fileType: 'parquet',
+        },
+      ],
+    });
+
+    expect(
+      fileBufferData1.some((file) => file.fileName !== 'taxi1.parquet')
+    ).toBe(true);
+
+    // Drop a file when there is only one file
+    await fileManager.dropFilesByTableName('taxi1', ['taxi2.parquet']);
+
+    const tableData2 = await indexedDB.tablesKey.toArray();
+    const fileBufferData2 = await indexedDB.files.toArray();
+
+    // Verify that the file is dropped
+    expect(tableData2[0]).toEqual({
+      tableName: 'taxi1',
+      files: [],
+    });
+
+    expect(
+      fileBufferData2.some((file) => file.fileName !== 'taxi2.parquet')
+    ).toBe(true);
+  });
 });

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

@@ -2,14 +2,15 @@ import { InstanceManagerType } from '../../dbm/instance-manager';
 import { mergeFileBufferStoreIntoTable } from '../../utils/merge-file-buffer-store-into-table';
 import {
   FileBufferStore,
+  FileData,
   FileManagerConstructorOptions,
   FileManagerType,
-  Table,
 } from '../file-manager-type';
-import { DuckDBDatabase } from './duckdb-database';
+import { DuckDBFilesDatabase } from './duckdb-files-database';
+
 export class IndexedDBFileManager implements FileManagerType {
-  private indexedDB: DuckDBDatabase;
-  private tables: Map<string, Table>;
+  // IndexedDB instance
+  private indexedDB: DuckDBFilesDatabase;
 
   fetchTableFileBuffers: (tableName: string) => Promise<FileBufferStore[]>;
   instanceManager: InstanceManagerType;
@@ -20,8 +21,7 @@ export class IndexedDBFileManager implements FileManagerType {
   }: FileManagerConstructorOptions) {
     this.fetchTableFileBuffers = fetchTableFileBuffers;
     this.instanceManager = instanceManager;
-    this.indexedDB = new DuckDBDatabase();
-    this.tables = new Map();
+    this.indexedDB = new DuckDBFilesDatabase();
   }
 
   /**
@@ -33,29 +33,30 @@ export class IndexedDBFileManager implements FileManagerType {
   }
 
   async initializeDB(): Promise<void> {
-    // Clear the database when initialized
-    await this._flushDB();
+    return;
   }
 
   async bulkRegisterFileBuffer(fileBuffers: FileBufferStore[]): Promise<void> {
     const tableNames = Array.from(
       new Set(fileBuffers.map((fileBuffer) => fileBuffer.tableName))
     );
 
+    const currentTableData = await this.indexedDB.tablesKey.toArray();
+
     const updatedTableMap = mergeFileBufferStoreIntoTable(
       fileBuffers,
-      this.tables
+      currentTableData
     );
 
     /**
      * Extracts the tables and files data from the tablesMap and fileBuffers
      * in format that can be stored in IndexedDB
      */
-    const tableData = tableNames.map((tableName) => {
+    const updatedTableData = tableNames.map((tableName) => {
       return { tableName, files: updatedTableMap.get(tableName)?.files ?? [] };
     });
 
-    const fileData = fileBuffers.map((fileBuffer) => {
+    const newFilesData = fileBuffers.map((fileBuffer) => {
       return { buffer: fileBuffer.buffer, fileName: fileBuffer.fileName };
     });
 
@@ -66,11 +67,9 @@ export class IndexedDBFileManager implements FileManagerType {
         this.indexedDB.tablesKey,
         this.indexedDB.files,
         async () => {
-          await this.indexedDB.tablesKey.bulkPut(tableData);
-
-          await this.indexedDB.files.bulkPut(fileData);
+          await this.indexedDB.tablesKey.bulkPut(updatedTableData);
 
-          this.tables = updatedTableMap;
+          await this.indexedDB.files.bulkPut(newFilesData);
         }
       )
       .catch((error) => {
@@ -90,9 +89,11 @@ export class IndexedDBFileManager implements FileManagerType {
   async registerFileBuffer(fileBuffer: FileBufferStore): Promise<void> {
     const { buffer, fileName, tableName } = fileBuffer;
 
+    const currentTableData = await this.indexedDB.tablesKey.toArray();
+
     const updatedTableMap = mergeFileBufferStoreIntoTable(
       [fileBuffer],
-      this.tables
+      currentTableData
     );
 
     // Update the tables and files table in IndexedDB
@@ -108,14 +109,14 @@ export class IndexedDBFileManager implements FileManagerType {
           });
 
           await this.indexedDB.files.put({ fileName, buffer });
-
-          this.tables = updatedTableMap;
         }
       )
       .catch((error) => {
         console.error(error);
       });
+
     const db = await this.instanceManager.getDB();
+
     // Register the file buffer in the DuckDB instance
     return db.registerFileBuffer(fileBuffer.fileName, fileBuffer.buffer);
   }
@@ -128,9 +129,7 @@ export class IndexedDBFileManager implements FileManagerType {
   }
 
   async mountFileBufferByTableNames(tableNames: string[]): Promise<void> {
-    const tableData = tableNames.map((tableName) => {
-      return this.tables.get(tableName);
-    });
+    const tableData = await this.indexedDB.tablesKey.bulkGet(tableNames);
 
     const promises = tableData.map(async (table) => {
       // Retrieve file names for the specified table
@@ -145,6 +144,7 @@ export class IndexedDBFileManager implements FileManagerType {
         filesData.map(async (file) => {
           if (file) {
             const db = await this.instanceManager.getDB();
+
             await db.registerFileBuffer(file.fileName, file.buffer);
           }
         })
@@ -155,9 +155,7 @@ export class IndexedDBFileManager implements FileManagerType {
   }
 
   async unmountFileBufferByTableNames(tableNames: string[]): Promise<void> {
-    const tableData = tableNames.map((tableName) => {
-      return this.tables.get(tableName);
-    });
+    const tableData = await this.indexedDB.tablesKey.bulkGet(tableNames);
 
     const promises = tableData.map(async (table) => {
       // Unregister file buffers from DuckDB for each table
@@ -171,4 +169,38 @@ export class IndexedDBFileManager implements FileManagerType {
 
     await Promise.all(promises);
   }
+
+  /**
+   * Get the list of files for the specified table name
+   */
+  async getFilesByTableName(tableName: string): Promise<FileData[]> {
+    const tableData = await this.indexedDB.tablesKey.get(tableName);
+
+    return tableData?.files ?? [];
+  }
+
+  /**
+   * Drop the specified files by tableName from the IndexedDB
+   */
+  async dropFilesByTableName(
+    tableName: string,
+    fileNames: string[]
+  ): Promise<void> {
+    const tableData = await this.indexedDB.tablesKey.get(tableName);
+
+    if (tableData) {
+      // Retrieve the files that are not dropped
+      const updatedFiles = tableData.files.filter(
+        (file) => !fileNames.includes(file.fileName)
+      );
+
+      await this.indexedDB.tablesKey.put({
+        tableName,
+        files: updatedFiles,
+      });
+    }
+
+    // Remove the files from the IndexedDB
+    await this.indexedDB.files.bulkDelete(fileNames);
+  }
 }

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

@@ -1,6 +1,7 @@
 import { InstanceManagerType } from '../dbm/instance-manager';
 import {
   FileBufferStore,
+  FileData,
   FileManagerConstructorOptions,
   FileManagerType,
 } from './file-manager-type';
@@ -44,4 +45,15 @@ export class MemoryDBFileManager implements FileManagerType {
   async unmountFileBufferByTableNames(tableNames: string[]): Promise<void> {
     // not needed for memory file manager
   }
+
+  async getFilesByTableName(tableName: string): Promise<FileData[]> {
+    // not needed for memory file manager
+    return [];
+  }
+  async dropFilesByTableName(
+    tableName: string,
+    fileNames: string[]
+  ): Promise<void> {
+    // not needed for memory file manager
+  }
 }