Add registerOPFSFileName to API, available in the shell via

carlopi · Jan 16, 2025 · f848df5 · f848df5
1 parent 72e03b4
commit f848df5
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 0 deletions.
diff --git a/packages/duckdb-wasm-shell/crate/src/duckdb/async_duckdb.rs b/packages/duckdb-wasm-shell/crate/src/duckdb/async_duckdb.rs
@@ -63,6 +63,11 @@ extern "C" {
         conn: ConnectionID,
     ) -> Result<JsValue, JsValue>;
 
+    #[wasm_bindgen(catch, method, js_name = "registerOPFSFileName")]
+    async fn register_opfs_file_name(
+        this: &JsAsyncDuckDB,
+        text: &str,
+    ) -> Result<JsValue, JsValue>;
     #[wasm_bindgen(catch, method, js_name = "collectFileStatistics")]
     async fn collect_file_statistics(
         this: &JsAsyncDuckDB,
@@ -158,6 +163,14 @@ impl AsyncDuckDB {
         Ok(())
     }
 
+    pub async fn register_opfs_file_name(
+        &self,
+        file: &str,
+    ) -> Result<(), js_sys::Error> {
+        self.bindings.register_opfs_file_name(file).await?;
+        Ok(())
+    }
+
     /// Enable file statistics
     pub async fn export_file_statistics(
         &self,

diff --git a/packages/duckdb-wasm-shell/crate/src/shell.rs b/packages/duckdb-wasm-shell/crate/src/shell.rs
@@ -377,6 +377,18 @@ impl Shell {
         Shell::with(|s| s.write(&buffer));
     }
 
+    pub fn register_opfs_file_name(name: &str) {
+        let db_ptr = Shell::with(|s| s.db.clone());
+        let name_copy = name.to_string();
+        spawn_local(async move {
+            let db = match db_ptr {
+                Some(ref db) => db.read().unwrap(),
+                None => return,
+            };
+            db.register_opfs_file_name(&name_copy).await.unwrap();
+        });
+    }
+
     pub fn collect_file_statistics(name: &str, enable: bool) {
         let db_ptr = Shell::with(|s| s.db.clone());
         let name_copy = name.to_string();
@@ -532,6 +544,13 @@ impl Shell {
                     }
                 }
             }
+            "register" => {
+                let filename = args[subcmd.len()..].trim();
+                db.register_opfs_file_name(filename).await.unwrap();
+                Shell::with_mut(|s| {
+                    s.writeln(&format!("Registering OPFS file handle for: {}", filename))
+                });
+            }
             "track" => {
                 let filename = args[subcmd.len()..].trim();
                 db.collect_file_statistics(filename, true).await.unwrap();
@@ -666,6 +685,7 @@ impl Shell {
                         ".files drop            Drop all files.\r\n",
                         ".files drop $FILE      Drop a single file.\r\n",
                         ".files track $FILE     Collect file statistics.\r\n",
+                        ".files register $FILE  Register OPFS file handle.\r\n",
                         ".files paging $FILE    Show file paging.\r\n",
                         ".files reads $FILE     Show file reads.\r\n",
                         ".open $FILE            Open database file.\r\n",

diff --git a/packages/duckdb-wasm-shell/crate/src/shell_api.rs b/packages/duckdb-wasm-shell/crate/src/shell_api.rs
@@ -65,6 +65,11 @@ pub fn writeln(text: &str) {
     Shell::with(|s| s.writeln(text));
 }
 
+#[wasm_bindgen(js_name = "registerOPFSFileName")]
+pub fn register_opfs_file_name(name: &str) {
+    Shell::register_opfs_file_name(name);
+}
+
 #[wasm_bindgen(js_name = "collectFileStatistics")]
 pub fn collect_file_statistics(name: &str, enable: bool) {
     Shell::collect_file_statistics(name, enable);

diff --git a/packages/duckdb-wasm/src/bindings/bindings_interface.ts b/packages/duckdb-wasm/src/bindings/bindings_interface.ts
@@ -62,6 +62,7 @@ export interface DuckDBBindings {
     flushFiles(): void;
     copyFileToPath(name: string, path: string): void;
     copyFileToBuffer(name: string): Uint8Array;
+    registerOPFSFileName(file: string): void;
     collectFileStatistics(file: string, enable: boolean): void;
     exportFileStatistics(file: string): FileStatistics;
 }
diff --git a/packages/duckdb-wasm/src/parallel/async_bindings.ts b/packages/duckdb-wasm/src/parallel/async_bindings.ts
@@ -160,6 +160,7 @@ export class AsyncDuckDB implements AsyncDuckDBBindings {
         switch (task.type) {
             case WorkerRequestType.CLOSE_PREPARED:
             case WorkerRequestType.COLLECT_FILE_STATISTICS:
+            case WorkerRequestType.REGISTER_OPFS_FILE_NAME:
             case WorkerRequestType.COPY_FILE_TO_PATH:
             case WorkerRequestType.DISCONNECT:
             case WorkerRequestType.DROP_FILE:
@@ -546,6 +547,15 @@ export class AsyncDuckDB implements AsyncDuckDBBindings {
         await this.postTask(task, []);
     }
 
+    /** Enable file statistics */
+    public async registerOPFSFileName(name: string): Promise<void> {
+        const task = new WorkerTask<WorkerRequestType.REGISTER_OPFS_FILE_NAME, [string], null>(
+            WorkerRequestType.REGISTER_OPFS_FILE_NAME,
+            [name],
+        );
+        await this.postTask(task, []);
+    }
+
     /** Enable file statistics */
     public async collectFileStatistics(name: string, enable: boolean): Promise<void> {
         const task = new WorkerTask<WorkerRequestType.COLLECT_FILE_STATISTICS, [string, boolean], null>(

diff --git a/packages/duckdb-wasm/src/parallel/worker_dispatcher.ts b/packages/duckdb-wasm/src/parallel/worker_dispatcher.ts
@@ -360,6 +360,11 @@ export abstract class AsyncDuckDBDispatcher implements Logger {
                     this.sendOK(request);
                     break;
 
+                case WorkerRequestType.REGISTER_OPFS_FILE_NAME:
+                    this._bindings.registerOPFSFileName(request.data[0]);
+                    this.sendOK(request);
+                    break;
+
                 case WorkerRequestType.EXPORT_FILE_STATISTICS: {
                     this.postMessage(
                         {

diff --git a/packages/duckdb-wasm/src/parallel/worker_request.ts b/packages/duckdb-wasm/src/parallel/worker_request.ts
@@ -14,6 +14,7 @@ export enum WorkerRequestType {
     CANCEL_PENDING_QUERY = 'CANCEL_PENDING_QUERY',
     CLOSE_PREPARED = 'CLOSE_PREPARED',
     COLLECT_FILE_STATISTICS = 'COLLECT_FILE_STATISTICS',
+    REGISTER_OPFS_FILE_NAME = 'REGISTER_OPFS_FILE_NAME',
     CONNECT = 'CONNECT',
     COPY_FILE_TO_BUFFER = 'COPY_FILE_TO_BUFFER',
     COPY_FILE_TO_PATH = 'COPY_FILE_TO_PATH',
@@ -108,6 +109,7 @@ export type WorkerRequestVariant =
     | WorkerRequest<WorkerRequestType.CLOSE_PREPARED, [ConnectionID, StatementID]>
     | WorkerRequest<WorkerRequestType.CANCEL_PENDING_QUERY, number>
     | WorkerRequest<WorkerRequestType.COLLECT_FILE_STATISTICS, [string, boolean]>
+    | WorkerRequest<WorkerRequestType.REGISTER_OPFS_FILE_NAME, [string]>
     | WorkerRequest<WorkerRequestType.CONNECT, null>
     | WorkerRequest<WorkerRequestType.COPY_FILE_TO_BUFFER, string>
     | WorkerRequest<WorkerRequestType.COPY_FILE_TO_PATH, [string, string]>
@@ -166,6 +168,7 @@ export type WorkerResponseVariant =
 
 export type WorkerTaskVariant =
     | WorkerTask<WorkerRequestType.COLLECT_FILE_STATISTICS, [string, boolean], null>
+    | WorkerTask<WorkerRequestType.REGISTER_OPFS_FILE_NAME, [string], null>
     | WorkerTask<WorkerRequestType.CLOSE_PREPARED, [number, number], null>
     | WorkerTask<WorkerRequestType.CONNECT, null, ConnectionID>
     | WorkerTask<WorkerRequestType.COPY_FILE_TO_BUFFER, string, Uint8Array>