Skip to content

Commit

Permalink
feat(duckdb-driver): Add support for installing and loading DuckDB Co…
Browse files Browse the repository at this point in the history
…mmunity Extensions
  • Loading branch information
benyaminsalimi committed Feb 1, 2025
1 parent 57bcbc4 commit f1e538a
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 33 deletions.
4 changes: 2 additions & 2 deletions docs/pages/product/configuration/data-sources/duckdb.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ deployment][ref-demo-deployment] in Cube Cloud.
| `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 |||

| `CUBEJS_DB_DUCKDB_COMMUNITY_EXTENSIONS` | A comma-separated list of DuckDB community extensions to install and load | A comma-separated list of DuckDB community extensions |||
## Pre-Aggregation Feature Support

### count_distinct_approx
Expand Down Expand Up @@ -131,4 +131,4 @@ connections are made over HTTPS.
/product/caching/using-pre-aggregations#pre-aggregation-build-strategies
[ref-schema-ref-types-formats-countdistinctapprox]: /reference/data-model/types-and-formats#count_distinct_approx
[self-preaggs-batching]: #batching
[ref-demo-deployment]: /product/deployment/cloud/deployments#demo-deployments
[ref-demo-deployment]: /product/deployment/cloud/deployments#demo-deployments
8 changes: 8 additions & 0 deletions docs/pages/reference/configuration/environment-variables.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -332,6 +332,14 @@ A comma-separated list of DuckDB extensions to install and load.
| ------------------------------------------- | ---------------------- | --------------------- |
| A comma-separated list of DuckDB extensions | N/A | N/A |

## `CUBEJS_DB_DUCKDB_COMMUNITY_EXTENSIONS`

A comma-separated list of DuckDB community extensions to install and load.

| Possible Values | Default in Development | Default in Production |
| ----------------------------------------------------- | ---------------------- | --------------------- |
| A comma-separated list of DuckDB community 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
Expand Down
13 changes: 13 additions & 0 deletions packages/cubejs-backend-shared/src/env.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1713,6 +1713,19 @@ const variables: Record<string, (...args: any) => any> = {
}
return [];
},
duckdbCommunityExtensions: ({
dataSource
}: {
dataSource: string,
}) => {
const extensions = process.env[
keyByDataSource('CUBEJS_DB_DUCKDB_COMMUNITY_EXTENSIONS', dataSource)
];
if (extensions) {
return extensions.split(',').map(e => e.trim());
}
return [];
},
/** ***************************************************************
* Presto Driver *
**************************************************************** */
Expand Down
69 changes: 38 additions & 31 deletions packages/cubejs-duckdb-driver/src/DuckDBDriver.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,10 +55,39 @@ export class DuckDBDriver extends BaseDriver implements DriverInterface {
return super.toGenericType(columnType.toLowerCase());
}

private async installExtensions(extensions: string[], execAsync: (sql: string, ...params: any[]) => Promise<void>, repository: string = ''): Promise<void> {
repository = repository ? ` FROM ${repository}` : '';
for (const extension of extensions) {
try {
await execAsync(`INSTALL ${extension}${repository}`);
} 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;
}
}
}

private async loadExtensions(extensions: string[], execAsync: (sql: string, ...params: any[]) => Promise<void>): Promise<void> {
for (const extension of extensions) {
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;
}
}
}

protected async init(): Promise<InitPromise> {
const token = getEnv('duckdbMotherDuckToken', this.config);
const dbPath = getEnv('duckdbDatabasePath', this.config);

// Determine the database URL based on the provided db_path or token
let dbUrl: string;
if (dbPath) {
Expand Down Expand Up @@ -119,7 +148,7 @@ export class DuckDBDriver extends BaseDriver implements DriverInterface {
value: getEnv('duckdbS3SessionToken', this.config),
}
];

for (const { key, value } of configuration) {
if (value) {
try {
Expand All @@ -135,34 +164,12 @@ 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;
}
}
const officialExtensions = getEnv('duckdbExtensions', this.config);
await this.installExtensions(officialExtensions, execAsync);
await this.loadExtensions(officialExtensions, execAsync);
const communityExtensions = getEnv('duckdbCommunityExtensions', this.config);
await this.installExtensions(communityExtensions, execAsync, 'community');
await this.loadExtensions(communityExtensions, execAsync);

if (this.config.initSql) {
try {
Expand All @@ -175,7 +182,7 @@ export class DuckDBDriver extends BaseDriver implements DriverInterface {
}
}
}

return {
defaultConnection,
db
Expand Down

0 comments on commit f1e538a

Please sign in to comment.