Skip to content

Commit f1e538a

Browse files
feat(duckdb-driver): Add support for installing and loading DuckDB Community Extensions
1 parent 57bcbc4 commit f1e538a

File tree

4 files changed

+61
-33
lines changed

4 files changed

+61
-33
lines changed

docs/pages/product/configuration/data-sources/duckdb.mdx

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ deployment][ref-demo-deployment] in Cube Cloud.
7474
| `CUBEJS_DB_DUCKDB_S3_URL_STYLE` | To choose the S3 URL style(vhost or path) | 'vhost' or 'path' |||
7575
| `CUBEJS_DB_DUCKDB_S3_SESSION_TOKEN` | The token for the S3 session | A valid Session Token |||
7676
| `CUBEJS_DB_DUCKDB_EXTENSIONS` | A comma-separated list of DuckDB extensions to install and load | A comma-separated list of DuckDB extensions |||
77-
77+
| `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 |||
7878
## Pre-Aggregation Feature Support
7979

8080
### count_distinct_approx
@@ -131,4 +131,4 @@ connections are made over HTTPS.
131131
/product/caching/using-pre-aggregations#pre-aggregation-build-strategies
132132
[ref-schema-ref-types-formats-countdistinctapprox]: /reference/data-model/types-and-formats#count_distinct_approx
133133
[self-preaggs-batching]: #batching
134-
[ref-demo-deployment]: /product/deployment/cloud/deployments#demo-deployments
134+
[ref-demo-deployment]: /product/deployment/cloud/deployments#demo-deployments

docs/pages/reference/configuration/environment-variables.mdx

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,14 @@ A comma-separated list of DuckDB extensions to install and load.
332332
| ------------------------------------------- | ---------------------- | --------------------- |
333333
| A comma-separated list of DuckDB extensions | N/A | N/A |
334334

335+
## `CUBEJS_DB_DUCKDB_COMMUNITY_EXTENSIONS`
336+
337+
A comma-separated list of DuckDB community extensions to install and load.
338+
339+
| Possible Values | Default in Development | Default in Production |
340+
| ----------------------------------------------------- | ---------------------- | --------------------- |
341+
| A comma-separated list of DuckDB community extensions | N/A | N/A |
342+
335343
## `CUBEJS_DB_ELASTIC_APIKEY_ID`
336344

337345
The [ID of the API key from elastic.co][elastic-docs-api-keys]. Required when

packages/cubejs-backend-shared/src/env.ts

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1713,6 +1713,19 @@ const variables: Record<string, (...args: any) => any> = {
17131713
}
17141714
return [];
17151715
},
1716+
duckdbCommunityExtensions: ({
1717+
dataSource
1718+
}: {
1719+
dataSource: string,
1720+
}) => {
1721+
const extensions = process.env[
1722+
keyByDataSource('CUBEJS_DB_DUCKDB_COMMUNITY_EXTENSIONS', dataSource)
1723+
];
1724+
if (extensions) {
1725+
return extensions.split(',').map(e => e.trim());
1726+
}
1727+
return [];
1728+
},
17161729
/** ***************************************************************
17171730
* Presto Driver *
17181731
**************************************************************** */

packages/cubejs-duckdb-driver/src/DuckDBDriver.ts

Lines changed: 38 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,39 @@ export class DuckDBDriver extends BaseDriver implements DriverInterface {
5555
return super.toGenericType(columnType.toLowerCase());
5656
}
5757

58+
private async installExtensions(extensions: string[], execAsync: (sql: string, ...params: any[]) => Promise<void>, repository: string = ''): Promise<void> {
59+
repository = repository ? ` FROM ${repository}` : '';
60+
for (const extension of extensions) {
61+
try {
62+
await execAsync(`INSTALL ${extension}${repository}`);
63+
} catch (e) {
64+
if (this.logger) {
65+
console.error(`DuckDB - error on installing ${extension}`, { e });
66+
}
67+
// DuckDB will lose connection_ref on connection on error, this will lead to broken connection object
68+
throw e;
69+
}
70+
}
71+
}
72+
73+
private async loadExtensions(extensions: string[], execAsync: (sql: string, ...params: any[]) => Promise<void>): Promise<void> {
74+
for (const extension of extensions) {
75+
try {
76+
await execAsync(`LOAD ${extension}`);
77+
} catch (e) {
78+
if (this.logger) {
79+
console.error(`DuckDB - error on loading ${extension}`, { e });
80+
}
81+
// DuckDB will lose connection_ref on connection on error, this will lead to broken connection object
82+
throw e;
83+
}
84+
}
85+
}
86+
5887
protected async init(): Promise<InitPromise> {
5988
const token = getEnv('duckdbMotherDuckToken', this.config);
6089
const dbPath = getEnv('duckdbDatabasePath', this.config);
61-
90+
6291
// Determine the database URL based on the provided db_path or token
6392
let dbUrl: string;
6493
if (dbPath) {
@@ -119,7 +148,7 @@ export class DuckDBDriver extends BaseDriver implements DriverInterface {
119148
value: getEnv('duckdbS3SessionToken', this.config),
120149
}
121150
];
122-
151+
123152
for (const { key, value } of configuration) {
124153
if (value) {
125154
try {
@@ -135,34 +164,12 @@ export class DuckDBDriver extends BaseDriver implements DriverInterface {
135164
}
136165

137166
// Install & load extensions if configured in env variable.
138-
const extensions = getEnv('duckdbExtensions', this.config);
139-
for (const extension of extensions) {
140-
try {
141-
await execAsync(`INSTALL ${extension}`);
142-
} catch (e) {
143-
if (this.logger) {
144-
console.error(`DuckDB - error on installing ${extension}`, {
145-
e
146-
});
147-
}
148-
149-
// DuckDB will lose connection_ref on connection on error, this will lead to broken connection object
150-
throw e;
151-
}
152-
153-
try {
154-
await execAsync(`LOAD ${extension}`);
155-
} catch (e) {
156-
if (this.logger) {
157-
console.error(`DuckDB - error on loading ${extension}`, {
158-
e
159-
});
160-
}
161-
162-
// DuckDB will lose connection_ref on connection on error, this will lead to broken connection object
163-
throw e;
164-
}
165-
}
167+
const officialExtensions = getEnv('duckdbExtensions', this.config);
168+
await this.installExtensions(officialExtensions, execAsync);
169+
await this.loadExtensions(officialExtensions, execAsync);
170+
const communityExtensions = getEnv('duckdbCommunityExtensions', this.config);
171+
await this.installExtensions(communityExtensions, execAsync, 'community');
172+
await this.loadExtensions(communityExtensions, execAsync);
166173

167174
if (this.config.initSql) {
168175
try {
@@ -175,7 +182,7 @@ export class DuckDBDriver extends BaseDriver implements DriverInterface {
175182
}
176183
}
177184
}
178-
185+
179186
return {
180187
defaultConnection,
181188
db

0 commit comments

Comments
 (0)