Skip to content

Commit 38a5094

Browse files
RusovDmitriyFrank-TXS
authored andcommitted
feat(prestodb-driver): Support s3 export bucket (cube-js#9736)
1 parent ca1ebb0 commit 38a5094

File tree

1 file changed

+29
-6
lines changed

1 file changed

+29
-6
lines changed

packages/cubejs-prestodb-driver/src/PrestoDriver.ts

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,12 @@ const presto = require('presto-client');
3030

3131
export type PrestoDriverExportBucket = {
3232
exportBucket?: string,
33-
bucketType?: 'gcs',
33+
bucketType?: 'gcs' | 's3',
3434
credentials?: any,
35+
accessKeyId?: string,
36+
secretAccessKey?: string,
37+
exportBucketRegion?: string,
38+
exportBucketS3AdvancedFS?: boolean,
3539
exportBucketCsvEscapeSymbol?: string,
3640
};
3741

@@ -50,7 +54,7 @@ export type PrestoDriverConfiguration = PrestoDriverExportBucket & {
5054
queryTimeout?: number;
5155
};
5256

53-
const SUPPORTED_BUCKET_TYPES = ['gcs'];
57+
const SUPPORTED_BUCKET_TYPES = ['gcs', 's3'];
5458
/**
5559
* Presto driver class.
5660
*/
@@ -103,8 +107,11 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
103107
...(authToken ? { custom_auth: `Bearer ${authToken}` } : {}),
104108
...(dbPassword ? { basic_auth: { user: dbUser, password: dbPassword } } : {}),
105109
ssl: this.getSslOptions(dataSource),
106-
bucketType: getEnv('dbExportBucketType', { supported: ['gcs'], dataSource }),
110+
bucketType: getEnv('dbExportBucketType', { supported: SUPPORTED_BUCKET_TYPES, dataSource }),
107111
exportBucket: getEnv('dbExportBucket', { dataSource }),
112+
accessKeyId: getEnv('dbExportBucketAwsKey', { dataSource }),
113+
secretAccessKey: getEnv('dbExportBucketAwsSecret', { dataSource }),
114+
exportBucketRegion: getEnv('dbExportBucketAwsRegion', { dataSource }),
108115
credentials: getEnv('dbExportGCSCredentials', { dataSource }),
109116
...config
110117
};
@@ -312,7 +319,12 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
312319

313320
const { schema, tableName } = this.splitTableFullName(params.tableFullName);
314321
const tableWithCatalogAndSchema = `${this.config.catalog}.${schema}.${tableName}`;
315-
const protocol = bucketType === 'gcs' ? 'gs' : bucketType;
322+
323+
const protocol = {
324+
gcs: 'gs',
325+
s3: this.config.exportBucketS3AdvancedFS ? 's3a' : 's3'
326+
}[bucketType || 'gcs'];
327+
316328
const externalLocation = `${protocol}://${exportBucket}/${schema}/${tableName}`;
317329
const withParams = `( external_location = '${externalLocation}', format = 'CSV')`;
318330
const select = `SELECT ${this.generateTableColumnsForExport(types)} FROM (${params.fromSql})`;
@@ -345,12 +357,23 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
345357
if (!this.config.exportBucket) {
346358
throw new Error('Export bucket is not configured.');
347359
}
348-
const { bucketType, exportBucket, credentials } = this.config;
360+
const { bucketType, exportBucket } = this.config;
349361
const { schema, tableName } = this.splitTableFullName(tableFullName);
350362

351363
switch (bucketType) {
352364
case 'gcs':
353-
return this.extractFilesFromGCS({ credentials }, exportBucket, `${schema}/${tableName}`);
365+
return this.extractFilesFromGCS({ credentials: this.config.credentials }, exportBucket, `${schema}/${tableName}`);
366+
case 's3':
367+
return this.extractUnloadedFilesFromS3({
368+
credentials: this.config.accessKeyId && this.config.secretAccessKey
369+
? {
370+
accessKeyId: this.config.accessKeyId,
371+
secretAccessKey: this.config.secretAccessKey,
372+
}
373+
: undefined,
374+
region: this.config.exportBucketRegion,
375+
},
376+
exportBucket, `${schema}/${tableName}`);
354377
default:
355378
throw new Error(`Unsupported export bucket type: ${bucketType}`);
356379
}

0 commit comments

Comments
 (0)