Rework OPFS, add prepareFileHandle and prepareFileHandles

duckdb · Jan 16, 2025 · 8445e12 · 8445e12
1 parent 0f29e56
commit 8445e12
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 5 deletions.
diff --git a/packages/duckdb-wasm/src/bindings/bindings_base.ts b/packages/duckdb-wasm/src/bindings/bindings_base.ts
@@ -469,6 +469,19 @@ export abstract class DuckDBBindingsBase implements DuckDBBindings {
         }
         dropResponseBuffers(this.mod);
     }
+    public async prepareFileHandle(fileName: string, protocol: DuckDBDataProtocol): Promise<void> {
+        if (protocol === DuckDBDataProtocol.BROWSER_FSACCESS && this._runtime.prepareFileHandles) {
+            const list = await this._runtime.prepareFileHandles([fileName], DuckDBDataProtocol.BROWSER_FSACCESS);
+            for (const item of list) {
+                const { handle, path: filePath, fromCached } = item;
+                if (!fromCached && handle.getSize()) {
+                    await this.registerFileHandleAsync(filePath, handle, DuckDBDataProtocol.BROWSER_FSACCESS, true);
+                }
+            }
+            return;
+        }
+        throw new Error(`prepareFileHandle: unsupported protocol ${protocol}`);
+    }
     /** Prepare a file handle that could only be acquired aschronously */
     public async prepareDBFileHandle(path: string, protocol: DuckDBDataProtocol): Promise<void> {
         if (protocol === DuckDBDataProtocol.BROWSER_FSACCESS && this._runtime.prepareDBFileHandle) {
@@ -601,8 +614,14 @@ export abstract class DuckDBBindingsBase implements DuckDBBindings {
         dropResponseBuffers(this.mod);
         return copy;
     }
-
     /** Enable tracking of file statistics */
+    public registerOPFSFileName(file: string): Promise<void> {
+	if (file.startsWith("opfs://")) {
+		return this.prepareFileHandle(file, DuckDBDataProtocol.BROWSER_FSACCESS);
+	} else {
+                throw new Error("Not an OPFS file name: " + file);
+        }
+    }
     public collectFileStatistics(file: string, enable: boolean): void {
         const [s, d, n] = callSRet(this.mod, 'duckdb_web_collect_file_stats', ['string', 'boolean'], [file, enable]);
         if (s !== StatusCode.SUCCESS) {

diff --git a/packages/duckdb-wasm/src/bindings/bindings_interface.ts b/packages/duckdb-wasm/src/bindings/bindings_interface.ts
@@ -54,6 +54,7 @@ export interface DuckDBBindings {
         protocol: DuckDBDataProtocol,
         directIO: boolean,
     ): Promise<HandleType>;
+    prepareFileHandle(path: string, protocol: DuckDBDataProtocol): Promise<void>;
     prepareDBFileHandle(path: string, protocol: DuckDBDataProtocol): Promise<void>;
     globFiles(path: string): WebFile[];
     dropFile(name: string): void;

diff --git a/packages/duckdb-wasm/src/bindings/runtime.ts b/packages/duckdb-wasm/src/bindings/runtime.ts
@@ -157,6 +157,8 @@ export interface DuckDBRuntime {
     removeFile(mod: DuckDBModule, pathPtr: number, pathLen: number): void;
 
     // Prepare a file handle that could only be acquired aschronously
+    prepareFileHandle?: (path: string, protocol: DuckDBDataProtocol) => Promise<PreparedDBFileHandle[]>;
+    prepareFileHandles?: (path: string[], protocol: DuckDBDataProtocol) => Promise<PreparedDBFileHandle[]>;
     prepareDBFileHandle?: (path: string, protocol: DuckDBDataProtocol) => Promise<PreparedDBFileHandle[]>;
 
     // Call a scalar UDF function

diff --git a/packages/duckdb-wasm/src/bindings/runtime_browser.ts b/packages/duckdb-wasm/src/bindings/runtime_browser.ts
@@ -19,20 +19,24 @@ import * as udf from './udf_runtime';
 const OPFS_PREFIX_LEN = 'opfs://'.length;
 const PATH_SEP_REGEX = /\/|\\/;
 
+
 export const BROWSER_RUNTIME: DuckDBRuntime & {
     _files: Map<string, any>;
     _fileInfoCache: Map<number, DuckDBFileInfo>;
     _globalFileInfo: DuckDBGlobalFileInfo | null;
     _preparedHandles: Record<string, FileSystemSyncAccessHandle>;
+    _opfsRoot: FileSystemDirectoryHandle | null;
 
     getFileInfo(mod: DuckDBModule, fileId: number): DuckDBFileInfo | null;
     getGlobalFileInfo(mod: DuckDBModule): DuckDBGlobalFileInfo | null;
+    assignOPFSRoot(): Promise<void>;
 } = {
     _files: new Map<string, any>(),
     _fileInfoCache: new Map<number, DuckDBFileInfo>(),
     _udfFunctions: new Map(),
     _globalFileInfo: null,
     _preparedHandles: {} as any,
+    _opfsRoot: null,
 
     getFileInfo(mod: DuckDBModule, fileId: number): DuckDBFileInfo | null {
         try {
@@ -101,11 +105,15 @@ export const BROWSER_RUNTIME: DuckDBRuntime & {
             return null;
         }
     },
-
+    async assignOPFSRoot(): Promise<void> {
+        if (!BROWSER_RUNTIME._opfsRoot) {
+            BROWSER_RUNTIME._opfsRoot = await navigator.storage.getDirectory();
+	}
+    },
     /** Prepare a file handle that could only be acquired aschronously */
-    async prepareDBFileHandle(dbPath: string, protocol: DuckDBDataProtocol): Promise<PreparedDBFileHandle[]> {
+    async prepareFileHandles(filePaths: string[], protocol: DuckDBDataProtocol): Promise<PreparedDBFileHandle[]> {
         if (protocol === DuckDBDataProtocol.BROWSER_FSACCESS) {
-            const filePaths = [dbPath, `${dbPath}.wal`];
+            await BROWSER_RUNTIME.assignOPFSRoot();
             const prepare = async (path: string): Promise<PreparedDBFileHandle> => {
                 if (BROWSER_RUNTIME._files.has(path)) {
                     return {
@@ -114,7 +122,7 @@ export const BROWSER_RUNTIME: DuckDBRuntime & {
                         fromCached: true,
                     };
                 }
-                const opfsRoot = await navigator.storage.getDirectory();
+                const opfsRoot = BROWSER_RUNTIME._opfsRoot!;
                 let dirHandle: FileSystemDirectoryHandle = opfsRoot;
                 // check if mkdir -p is needed
                 const opfsPath = path.slice(OPFS_PREFIX_LEN);
@@ -156,6 +164,14 @@ export const BROWSER_RUNTIME: DuckDBRuntime & {
             }
             return result;
         }
+        throw new Error(`Unsupported protocol ${protocol} for paths ${filePaths} with protocol ${protocol}`);
+    },
+    /** Prepare a file handle that could only be acquired aschronously */
+    async prepareDBFileHandle(dbPath: string, protocol: DuckDBDataProtocol): Promise<PreparedDBFileHandle[]> {
+        if (protocol === DuckDBDataProtocol.BROWSER_FSACCESS && this.prepareFileHandles) {
+            const filePaths = [dbPath, `${dbPath}.wal`];
+            return this.prepareFileHandles(filePaths, protocol);
+        }
         throw new Error(`Unsupported protocol ${protocol} for path ${dbPath} with protocol ${protocol}`);
     },