Skip to content

Commit 063f804

Browse files
KSDaemonKSDaemon
committed
feat(snowflake-driver): Add support for export Buckets with paths
1 parent 9a9c3d4 commit 063f804

10 files changed

+49345
-13
lines changed

‎.github/workflows/drivers-tests.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,9 +228,12 @@ jobs:
228228
snowflake
229229
snowflake-encrypted-pk
230230
snowflake-export-bucket-s3
231+
snowflake-export-bucket-s3-prefix
231232
snowflake-export-bucket-azure
233+
snowflake-export-bucket-azure-prefix
232234
snowflake-export-bucket-azure-via-storage-integration
233235
snowflake-export-bucket-gcs
236+
snowflake-export-bucket-gcs-prefix
234237
# As per docs:
235238
# Secrets cannot be directly referenced in if: conditionals. Instead, consider setting
236239
# secrets as job-level environment variables, then referencing the environment variables
@@ -259,9 +262,12 @@ jobs:
259262
- snowflake
260263
- snowflake-encrypted-pk
261264
- snowflake-export-bucket-s3
265+
- snowflake-export-bucket-s3-prefix
262266
- snowflake-export-bucket-azure
267+
- snowflake-export-bucket-azure-prefix
263268
- snowflake-export-bucket-azure-via-storage-integration
264269
- snowflake-export-bucket-gcs
270+
- snowflake-export-bucket-gcs-prefix
265271
fail-fast: false
266272

267273
steps:

‎packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts‎

Lines changed: 57 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -545,10 +545,25 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
545545
} else {
546546
const types = await this.queryColumnTypes(options.query.sql, options.query.params);
547547
const connection = await this.getConnection();
548-
const { bucketType, bucketName } =
548+
const { bucketType } =
549549
<SnowflakeDriverExportBucket> this.config.exportBucket;
550+
551+
let bucketName: string;
552+
let exportPrefix: string;
553+
let path: string;
554+
555+
if (bucketType === 'azure') {
556+
({ bucketName, path } = this.parseBucketUrl(this.config.exportBucket!.bucketName));
557+
const pathArr = path.split('/');
558+
bucketName = `${bucketName}/${pathArr[0]}`;
559+
exportPrefix = pathArr.length > 1 ? `${pathArr.slice(1).join('/')}/${tableName}` : tableName;
560+
} else {
561+
({ bucketName, path } = this.parseBucketUrl(this.config.exportBucket!.bucketName));
562+
exportPrefix = path ? `${path}/${tableName}` : tableName;
563+
}
564+
550565
const unloadSql = `
551-
COPY INTO '${bucketType}://${bucketName}/${tableName}/'
566+
COPY INTO '${bucketType}://${bucketName}/${exportPrefix}/'
552567
FROM (${options.query.sql})
553568
${this.exportOptionsClause(options)}`;
554569
const result = await this.execute<UnloadResponse[]>(
@@ -594,10 +609,25 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
594609
): Promise<TableStructure> {
595610
const types = await this.tableColumnTypes(tableName);
596611
const connection = await this.getConnection();
597-
const { bucketType, bucketName } =
612+
const { bucketType } =
598613
<SnowflakeDriverExportBucket> this.config.exportBucket;
614+
615+
let bucketName: string;
616+
let exportPrefix: string;
617+
let path: string;
618+
619+
if (bucketType === 'azure') {
620+
({ bucketName, path } = this.parseBucketUrl(this.config.exportBucket!.bucketName));
621+
const pathArr = path.split('/');
622+
bucketName = `${bucketName}/${pathArr[0]}`;
623+
exportPrefix = pathArr.length > 1 ? `${pathArr.slice(1).join('/')}/${tableName}` : tableName;
624+
} else {
625+
({ bucketName, path } = this.parseBucketUrl(this.config.exportBucket!.bucketName));
626+
exportPrefix = path ? `${path}/${tableName}` : tableName;
627+
}
628+
599629
const unloadSql = `
600-
COPY INTO '${bucketType}://${bucketName}/${tableName}/'
630+
COPY INTO '${bucketType}://${bucketName}/${exportPrefix}/'
601631
FROM ${tableName}
602632
${this.exportOptionsClause(options)}`;
603633
const result = await this.execute<UnloadResponse[]>(
@@ -695,36 +725,50 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
695725
* Returns an array of signed URLs of the unloaded csv files.
696726
*/
697727
private async getCsvFiles(tableName: string): Promise<string[]> {
698-
const { bucketType, bucketName } =
728+
const { bucketType } =
699729
<SnowflakeDriverExportBucket> this.config.exportBucket;
700730

701731
if (bucketType === 's3') {
702-
const cfg = <SnowflakeDriverExportAWS> this.config.exportBucket;
732+
const { keyId, secretKey, region } = <SnowflakeDriverExportAWS> this.config.exportBucket;
733+
734+
const { bucketName, path } = this.parseBucketUrl(this.config.exportBucket!.bucketName);
735+
const exportPrefix = path ? `${path}/${tableName}` : tableName;
703736

704737
return this.extractUnloadedFilesFromS3(
705738
{
706739
credentials: {
707-
accessKeyId: cfg.keyId,
708-
secretAccessKey: cfg.secretKey,
740+
accessKeyId: keyId,
741+
secretAccessKey: secretKey,
709742
},
710-
region: cfg.region,
743+
region,
711744
},
712745
bucketName,
713-
tableName,
746+
exportPrefix,
714747
);
715748
} else if (bucketType === 'gcs') {
716749
const { credentials } = (
717750
<SnowflakeDriverExportGCS> this.config.exportBucket
718751
);
719-
return this.extractFilesFromGCS({ credentials }, bucketName, tableName);
752+
753+
const { bucketName, path } = this.parseBucketUrl(this.config.exportBucket!.bucketName);
754+
const exportPrefix = path ? `${path}/${tableName}` : tableName;
755+
756+
return this.extractFilesFromGCS({ credentials }, bucketName, exportPrefix);
720757
} else if (bucketType === 'azure') {
721758
const { azureKey, sasToken, clientId, tenantId, tokenFilePath } = (
722759
<SnowflakeDriverExportAzure> this.config.exportBucket
723760
);
761+
762+
const { bucketName, path } = this.parseBucketUrl(this.config.exportBucket!.bucketName);
763+
const pathArr = path.split('/');
764+
const azureBucketPath = `${bucketName}/${pathArr[0]}`;
765+
766+
const exportPrefix = pathArr.length > 1 ? `${pathArr.slice(1).join('/')}/${tableName}` : tableName;
767+
724768
return this.extractFilesFromAzure(
725769
{ azureKey, sasToken, clientId, tenantId, tokenFilePath },
726-
bucketName,
727-
tableName,
770+
azureBucketPath,
771+
exportPrefix,
728772
);
729773
} else {
730774
throw new Error(`Unsupported export bucket type: ${bucketType}`);

‎packages/cubejs-testing-drivers/fixtures/snowflake.json‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,17 @@
1111
}
1212
}
1313
},
14+
"export-bucket-s3-prefix": {
15+
"cube": {
16+
"environment": {
17+
"CUBEJS_DB_EXPORT_BUCKET_TYPE": "s3",
18+
"CUBEJS_DB_EXPORT_BUCKET": "snowflake-drivers-tests-preaggs/testing_prefix/for_export_buckets",
19+
"CUBEJS_DB_EXPORT_BUCKET_AWS_KEY": "${DRIVERS_TESTS_CUBEJS_DB_EXPORT_BUCKET_AWS_KEY}",
20+
"CUBEJS_DB_EXPORT_BUCKET_AWS_SECRET": "${DRIVERS_TESTS_CUBEJS_DB_EXPORT_BUCKET_AWS_SECRET}",
21+
"CUBEJS_DB_EXPORT_BUCKET_AWS_REGION": "us-west-1"
22+
}
23+
}
24+
},
1425
"export-bucket-azure": {
1526
"cube": {
1627
"environment": {
@@ -21,6 +32,16 @@
2132
}
2233
}
2334
},
35+
"export-bucket-azure-prefix": {
36+
"cube": {
37+
"environment": {
38+
"CUBEJS_DB_EXPORT_BUCKET_TYPE": "azure",
39+
"CUBEJS_DB_EXPORT_BUCKET": "coreteamdevtest.blob.core.windows.net/snowflake-drivers-tests-preaggs/testing_prefix/for_export_buckets/",
40+
"CUBEJS_DB_EXPORT_BUCKET_AZURE_KEY": "${DRIVERS_TESTS_CUBEJS_DB_EXPORT_BUCKET_AZURE_KEY}",
41+
"CUBEJS_DB_EXPORT_BUCKET_AZURE_SAS_TOKEN": "${DRIVERS_TESTS_CUBEJS_DB_EXPORT_BUCKET_AZURE_SAS_TOKEN}"
42+
}
43+
}
44+
},
2445
"export-bucket-azure-via-storage-integration": {
2546
"cube": {
2647
"environment": {
@@ -41,6 +62,16 @@
4162
}
4263
}
4364
},
65+
"export-bucket-gcs-prefix": {
66+
"cube": {
67+
"environment": {
68+
"CUBEJS_DB_EXPORT_BUCKET_TYPE": "gcs",
69+
"CUBEJS_DB_EXPORT_BUCKET": "snowflake-drivers-tests-preaggs/testing_prefix/for_export_buckets",
70+
"CUBEJS_DB_EXPORT_INTEGRATION": "drivers_tests_preaggs_gcs",
71+
"CUBEJS_DB_EXPORT_GCS_CREDENTIALS": "${DRIVERS_TESTS_CUBEJS_DB_EXPORT_GCS_CREDENTIALS}"
72+
}
73+
}
74+
},
4475
"encrypted-pk": {
4576
"cube": {
4677
"environment": {

‎packages/cubejs-testing-drivers/package.json‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,9 +49,12 @@
4949
"snowflake-full": "yarn test-driver -i dist/test/snowflake-full.test.js",
5050
"snowflake-encrypted-pk-full": "yarn test-driver -i dist/test/snowflake-encrypted-pk-full.test.js",
5151
"snowflake-export-bucket-s3-full": "yarn test-driver -i dist/test/snowflake-export-bucket-s3-full.test.js",
52+
"snowflake-export-bucket-s3-prefix-full": "yarn test-driver -i dist/test/snowflake-export-bucket-s3-prefix-full.test.js",
5253
"snowflake-export-bucket-azure-full": "yarn test-driver -i dist/test/snowflake-export-bucket-azure-full.test.js",
54+
"snowflake-export-bucket-azure-prefix-full": "yarn test-driver -i dist/test/snowflake-export-bucket-azure-prefix-full.test.js",
5355
"snowflake-export-bucket-azure-via-storage-integration-full": "yarn test-driver -i dist/test/snowflake-export-bucket-azure-via-storage-integration-full.test.js",
5456
"snowflake-export-bucket-gcs-full": "yarn test-driver -i dist/test/snowflake-export-bucket-gcs-full.test.js",
57+
"snowflake-export-bucket-gcs-prefix-full": "yarn test-driver -i dist/test/snowflake-export-bucket-gcs-prefix-full.test.js",
5558
"redshift-driver": "yarn test-driver -i dist/test/redshift-driver.test.js",
5659
"redshift-core": "yarn test-driver -i dist/test/redshift-core.test.js",
5760
"redshift-full": "yarn test-driver -i dist/test/redshift-full.test.js",

0 commit comments

Comments
 (0)