Added support for installing and loading DuckDB Extensions.

cube-js · Sep 24, 2024 · 53f5f22 · 53f5f22
1 parent 73d9314
commit 53f5f22
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 0 deletions.
diff --git a/docs/pages/product/configuration/data-sources/duckdb.mdx b/docs/pages/product/configuration/data-sources/duckdb.mdx
@@ -73,6 +73,7 @@ deployment][ref-demo-deployment] in Cube Cloud.
 | `CUBEJS_DB_DUCKDB_S3_USE_SSL`           | Use SSL for connection                                                                                                | A boolean                                                     |    ❌    |                                    ❌                                     |
 | `CUBEJS_DB_DUCKDB_S3_URL_STYLE`         | To choose the S3 URL style(vhost or path)                                                                             | 'vhost' or 'path'                                             |    ❌    |                                    ❌                                     |
 | `CUBEJS_DB_DUCKDB_S3_SESSION_TOKEN`     | The token for the S3 session                                                                                          | A valid Session Token                                         |    ❌    |                                    ✅                                     |
+| `CUBEJS_DB_DUCKDB_EXTENSIONS`           | A comma-separated list of DuckDB extensions	to install and load                                                       | A comma-separated list of DuckDB extensions                   |    ❌    |                                    ✅                                     |
 
 ## Pre-Aggregation Feature Support
 

diff --git a/docs/pages/reference/configuration/environment-variables.mdx b/docs/pages/reference/configuration/environment-variables.mdx
@@ -315,6 +315,14 @@ The S3 session token.
 | -------------------------------------- | ---------------------- | --------------------- |
 | A valid S3 session token               | N/A                    | N/A                   |
 
+## `CUBEJS_DB_DUCKDB_EXTENSIONS`
+
+A comma-separated list of DuckDB extensions	to install and load.
+
+| Possible Values                             | Default in Development | Default in Production |
+| ------------------------------------------- | ---------------------- | --------------------- |
+| A comma-separated list of DuckDB extensions | N/A                    | N/A                   |
+
 ## `CUBEJS_DB_ELASTIC_APIKEY_ID`
 
 The [ID of the API key from elastic.co][elastic-docs-api-keys]. Required when

diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts
@@ -1535,6 +1535,20 @@ const variables: Record<string, (...args: any) => any> = {
     ]
   ),
 
+  duckdbExtensions: ({
+    dataSource
+  }: {
+    dataSource: string,
+  }) => {
+    const extensions = process.env[
+      keyByDataSource('CUBEJS_DB_DUCKDB_EXTENSIONS', dataSource)
+    ]
+    if (extensions) {
+      return extensions.split(',').map(e => e.trim());
+    }
+    return [];
+  },
+
   /**
    * Presto catalog.
    */

diff --git a/packages/cubejs-duckdb-driver/src/DuckDBDriver.ts b/packages/cubejs-duckdb-driver/src/DuckDBDriver.ts
@@ -160,6 +160,36 @@ export class DuckDBDriver extends BaseDriver implements DriverInterface {
       }
     }
 
+    // Install & load extensions if configured in env variable.
+    const extensions = getEnv('duckdbExtensions', this.config);
+    for (const extension of extensions) {
+      try {
+        await execAsync('INSTALL ' + extension);
+      } catch (e) {
+        if (this.logger) {
+          console.error('DuckDB - error on installing ' + extension, {
+            e
+          });
+        }
+
+        // DuckDB will lose connection_ref on connection on error, this will lead to broken connection object
+        throw e;
+      }
+
+      try {
+        await execAsync('LOAD ' + extension);
+      } catch (e) {
+        if (this.logger) {
+          console.error('DuckDB - error on loading ' + extension, {
+            e
+          });
+        }
+
+        // DuckDB will lose connection_ref on connection on error, this will lead to broken connection object
+        throw e;
+      }
+    }
+
     if (this.config.initSql) {
       try {
         await execAsync(this.config.initSql);