Merge pull request #286 from Canner/chore/duckdb-config

Chore: add setting duckdb config via env
Canner · Aug 21, 2023 · 8c5d9c2 · 8c5d9c2
2 parents 56f1cdb + 50ca982
commit 8c5d9c2
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 14 deletions.
diff --git a/packages/extension-driver-duckdb/src/lib/duckdbDataSource.ts b/packages/extension-driver-duckdb/src/lib/duckdbDataSource.ts
@@ -106,6 +106,7 @@ export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
     this.logRequest(firstDataSQL, parameters, options);
     const connection = db.connect();
     await this.loadExtensions(connection, configurationParameters);
+    await this.setExecConfig(connection);
     if (restDataSQL) this.logRequest(restDataSQL, parameters, options);
     const [firstData, restDataStream] = await this.acquireData(
       firstDataSQL,
@@ -297,27 +298,77 @@ export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
 
   // The dafault duckdb thread is 16
   // Setting thread below your CPU core number may result in enhanced performance, according to our observations.
-  private async setThread(db: duckdb.Database) {
-    const thread = process.env['DUCKDB_THREADS'];
+  // private async setThread(db: duckdb.Database) {
+  //   const thread = process.env['DUCKDB_THREADS'];
 
-    if (!thread) return;
-    await new Promise((resolve, reject) => {
-      db.run(`SET threads=${Number(thread)}`, (err: any) => {
-        if (err) reject(err);
-        this.logger.debug(`Set thread to ${thread}`);
-        resolve(true);
-      });
-    });
-  }
+  //   if (!thread) return;
+  //   await new Promise((resolve, reject) => {
+  //     db.run(`SET threads=${Number(thread)}`, (err: any) => {
+  //       if (err) reject(err);
+  //       this.logger.debug(`Set thread to ${thread}`);
+  //       resolve(true);
+  //     });
+  //   });
+  // }
 
   private async initDatabase(dbPath: string) {
-    const db = new duckdb.Database(dbPath);
+    const readonlyMode = process.env['DUCKDB_READONLY_MODE'];
+    let db;
+    if (readonlyMode) {
+      db = new duckdb.Database(dbPath, duckdb.OPEN_READONLY);
+      this.logger.debug(`Open database in readonly mode`);
+    } else {
+      db = new duckdb.Database(dbPath);
+      this.logger.debug(`Open database in automatic mode`);
+    }
     const conn = db.connect();
-    await this.setThread(db);
+    // await this.setThread(db);
+    await this.setInitConfig(conn);
     await this.installExtensions(conn);
+    await this.getCurrentConfig(conn);
     return db;
   }
 
+  private async setConfigList(conn: duckdb.Connection, configs: string[] | []) {
+    if (!configs.length) return;
+    await Promise.all(
+      configs.map((config) => {
+        return new Promise((resolve, reject) => {
+          conn.run(`SET ${config}`, (err: any) => {
+            if (err) reject(err);
+            this.logger.debug(`Set config ${config}`);
+            resolve(true);
+          });
+        });
+      })
+    );
+  }
+
+  private async setInitConfig(conn: duckdb.Connection) {
+    // threads=1, dir='path',
+    const configListStr = process.env['DUCKDB_INIT_CONFIG'];
+    const configs = configListStr?.split(',') || [];
+    await this.setConfigList(conn, configs);
+  }
+  private async setExecConfig(conn: duckdb.Connection) {
+    // threads=1, dir='path',
+    const configListStr = process.env['DUCKDB_EXEC_CONFIG'];
+    const configs = configListStr?.split(',') || [];
+    await this.setConfigList(conn, configs);
+  }
+
+  private async getCurrentConfig(conn: duckdb.Connection) {
+    return await new Promise((resolve, reject) => {
+      conn.all('select * from duckdb_settings()', (err: any, res: any) => {
+        if (err) reject(err);
+        for (const config of res) {
+          this.logger.debug(`Duckdb config: ${config.name} = ${config.value}`);
+        }
+        resolve(true);
+      });
+    });
+  }
+
   private async installExtensions(
     connection: duckdb.Connection
   ): Promise<void> {

diff --git a/types/duckdb.d.ts b/types/duckdb.d.ts
@@ -47,6 +47,11 @@ declare module 'duckdb' {
     /**
      * @param path  a file name for a persistent DB or :memory: for in-memory database
      */
+    constructor(
+      path: string,
+      accessMode?: number | Record<string, string>,
+      callback?: Callback<any>
+    );
     constructor(path: string): Database;
 
     /** Create a new connection to execute query, each connection has own transaction context */
@@ -59,7 +64,7 @@ declare module 'duckdb' {
     /**
      * Wait all scheduled database tasks have completed.
      * @param callback trigger callback function when all scheduled database tasks have completed.
-     */
+     */ s;
     wait(callback?: Callback<void>): void;
     /**
      * Query and trigger callback function when query completed
@@ -121,4 +126,20 @@ declare module 'duckdb' {
     [Symbol.asyncIterator](): Iterator<any>;
     nextChunk(): Promise<any[][]>;
   }
+
+  export const ERROR: number;
+
+  export const OPEN_CREATE: number;
+
+  export const OPEN_FULLMUTEX: number;
+
+  export const OPEN_PRIVATECACHE: number;
+
+  export const OPEN_READONLY: number;
+
+  export const OPEN_READWRITE: number;
+
+  export const OPEN_SHAREDCACHE: number;
+
+  export const INTERRUPT: number;
 }