feat(snowflake-driver): support azure exports buckets (#8730)

* add azure export bucket options to Snowflake driver

* add env var for Azure SAS Token

* add STORAGE_INTEGRATION variant to export bucket flows

* Simplify isUnloadSupported()

* add extractFilesFromGCS() to BaseDriver Class

* implement extractFilesFromAzure() in SnowFlake Driver

* some refactoring in extractFilesFromAzure()

* debug and fix extractFilesFromAzure()

* remove error if no csv files were exported (that's ok)

* move extractFilesFromAzure to BaseDriver from Snowflake

* switch Athena Driver to use extractUnloadedFilesFromS3() from BaseDvier impl

* switch RedShift Driver to use extractUnloadedFilesFromS3() from BaseDvier impl

* switch Databricks Driver to use extractUnloadedFilesFromS3() and extractFilesFromAzure() from BaseDvier impl

* Remove unused/cleanup from databricks driver

* update extractFilesFromGCS() in BaseDriver to receive credentials instead of Storage instance

* remove dep on @google-cloud/storage in Databricks for export bucket flow

* fix Databricks driver export bucket flows after refactoring

* set up CI to test clickhouse with s3 export bucket

* Rename databricks export bucket test to reflect the s3 type

* fix clickhouse-export-bucket-s3-full script in package.json

* set up CI to test Snowflake with s3 export bucket

* rename databricks tests snapshot after renaming test job

* improve extractUnloadedFilesFromS3 flow to support different bucket naming schemes

* add snowflake test snapshots for export bucket test

* add clickhouse test snapshots for export bucket test

* update databricks with bucket tests snapshots

* rename athena test to reflect that it uses export bucket

* reformat CLOUD_DATABASES in drivers tests CI job + add clickhouse-export-bucket-s3

* add snowflake export bucket to azure driver test

* add databricks export bucket to azure driver test

* add snowflake export bucket to azure via storage integration driver test

* add snowflake export bucket to gcs via storage integration driver test

* improve/refactor exported files filtering in extractFilesFromAzure()

* set env secrets for drivers tests

* Fix createExportBucket in Snowflake (making sasToken optional)

* Fix CUBEJS_DB_EXPORT_BUCKET env var for tests

* Remove includeIncrementalSchemaSuite from every test suite besides one + add comment

* Fix CUBEJS_DB_EXPORT_BUCKET env var for tests in databricks

* fix databricks export to azure flow + align test config

* Fix blob object filtering during unloading in azure

* sync @google-cloud/storage package used across packages

* yarn lock sync

* remove not needed includeIncrementalSchemaSuite in some tests

* rename bigquery → bigquery-export-bucket-gcs to align with the rest tests

* Add comments for exporting urls from S3 in BigQueryDriver

Author: KSDaemon

.github/workflows/drivers-tests.yml

@@ -175,7 +175,18 @@ jobs:
     needs: [latest-tag-sha, build]
     if: (needs['latest-tag-sha'].outputs.sha != github.sha)
     env:
-      CLOUD_DATABASES: athena bigquery databricks-jdbc databricks-jdbc-export-bucket snowflake
+      CLOUD_DATABASES: >
+        athena-export-bucket-s3
+        bigquery-export-bucket-gcs
+        clickhouse-export-bucket-s3
+        databricks-jdbc
+        databricks-jdbc-export-bucket-s3
+        databricks-jdbc-export-bucket-azure
+        snowflake
+        snowflake-export-bucket-s3
+        snowflake-export-bucket-azure
+        snowflake-export-bucket-azure-via-storage-integration
+        snowflake-export-bucket-gcs
       # As per docs:
       # Secrets cannot be directly referenced in if: conditionals. Instead, consider setting
       # secrets as job-level environment variables, then referencing the environment variables
@@ -186,15 +197,21 @@ jobs:
         node:
           - 20.x
         database:
-          - athena
-          - bigquery
+          - athena-export-bucket-s3
+          - bigquery-export-bucket-gcs
           - clickhouse
+          - clickhouse-export-bucket-s3
           - databricks-jdbc
-          - databricks-jdbc-export-bucket
+          - databricks-jdbc-export-bucket-s3
+          - databricks-jdbc-export-bucket-azure
           - mssql
           - mysql
           - postgres
           - snowflake
+          - snowflake-export-bucket-s3
+          - snowflake-export-bucket-azure
+          - snowflake-export-bucket-azure-via-storage-integration
+          - snowflake-export-bucket-gcs
       fail-fast: false
 
     steps:
@@ -258,6 +275,13 @@ jobs:
           # BigQuery
           DRIVERS_TESTS_CUBEJS_DB_BQ_CREDENTIALS: ${{ secrets.DRIVERS_TESTS_CUBEJS_DB_BQ_CREDENTIALS }}
 
+          #GCS
+          DRIVERS_TESTS_CUBEJS_DB_EXPORT_GCS_CREDENTIALS: ${{ secrets.DRIVERS_TESTS_CUBEJS_DB_EXPORT_GCS_CREDENTIALS }}
+
+          # Azure
+          DRIVERS_TESTS_CUBEJS_DB_EXPORT_BUCKET_AZURE_KEY: ${{ secrets.DRIVERS_TESTS_CUBEJS_DB_EXPORT_BUCKET_AZURE_KEY }}
+          DRIVERS_TESTS_CUBEJS_DB_EXPORT_BUCKET_AZURE_SAS_TOKEN: ${{ secrets.DRIVERS_TESTS_CUBEJS_DB_EXPORT_BUCKET_AZURE_SAS_TOKEN }}
+
           # Databricks
           DRIVERS_TESTS_CUBEJS_DB_DATABRICKS_URL: ${{ secrets.DRIVERS_TESTS_CUBEJS_DB_DATABRICKS_URL }}
           DRIVERS_TESTS_CUBEJS_DB_DATABRICKS_TOKEN: ${{ secrets.DRIVERS_TESTS_CUBEJS_DB_DATABRICKS_TOKEN }}

packages/cubejs-athena-driver/package.json

@@ -29,8 +29,6 @@
   "types": "dist/src/index.d.ts",
   "dependencies": {
     "@aws-sdk/client-athena": "^3.22.0",
-    "@aws-sdk/client-s3": "^3.49.0",
-    "@aws-sdk/s3-request-presigner": "^3.49.0",
     "@cubejs-backend/base-driver": "^0.36.0",
     "@cubejs-backend/shared": "^0.36.0",
     "sqlstring": "^2.3.1"

packages/cubejs-athena-driver/src/AthenaDriver.ts

@@ -17,8 +17,6 @@ import {
   ColumnInfo,
   StartQueryExecutionCommandInput,
 } from '@aws-sdk/client-athena';
-import { S3, GetObjectCommand } from '@aws-sdk/client-s3';
-import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
 import * as stream from 'stream';
 import {
   BaseDriver,
@@ -126,7 +124,7 @@ export class AthenaDriver extends BaseDriver implements DriverInterface {
       getEnv('athenaAwsSecret', { dataSource });
 
     const { schema, ...restConfig } = config;
-    
+
     this.schema = schema ||
       getEnv('dbName', { dataSource }) ||
       getEnv('dbSchema', { dataSource });
@@ -438,31 +436,18 @@ export class AthenaDriver extends BaseDriver implements DriverInterface {
    * Returns an array of signed URLs of the unloaded csv files.
    */
   private async getCsvFiles(tableName: string): Promise<string[]> {
-    const client = new S3({
-      credentials: this.config.credentials,
-      region: this.config.region,
-    });
     const { bucket, prefix } = AthenaDriver.splitS3Path(
       `${this.config.exportBucket}/${tableName}`
     );
-    const list = await client.listObjectsV2({
-      Bucket: bucket,
-      Prefix: prefix.slice(1), // skip leading
-    });
-    if (!list.Contents) {
-      return [];
-    } else {
-      const files = await Promise.all(
-        list.Contents.map(async (file) => {
-          const command = new GetObjectCommand({
-            Bucket: bucket,
-            Key: file.Key,
-          });
-          return getSignedUrl(client, command, { expiresIn: 3600 });
-        })
-      );
-      return files;
-    }
+
+    return this.extractUnloadedFilesFromS3(
+      {
+        credentials: this.config.credentials,
+        region: this.config.region,
+      },
+      bucket,
+      prefix.slice(1),
+    );
   }
 
   public informationSchemaQuery() {

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

@@ -718,7 +718,7 @@ const variables: Record<string, (...args: any) => any> = {
   ),
 
   /**
-   * AWS Key for the AWS based export bucket srorage.
+   * AWS Key for the AWS based export bucket storage.
    */
   dbExportBucketAwsKey: ({
     dataSource,
@@ -731,7 +731,7 @@ const variables: Record<string, (...args: any) => any> = {
   ),
 
   /**
-   * AWS Secret for the AWS based export bucket srorage.
+   * AWS Secret for the AWS based export bucket storage.
    */
   dbExportBucketAwsSecret: ({
     dataSource,
@@ -744,7 +744,7 @@ const variables: Record<string, (...args: any) => any> = {
   ),
 
   /**
-   * AWS Region for the AWS based export bucket srorage.
+   * AWS Region for the AWS based export bucket storage.
    */
   dbExportBucketAwsRegion: ({
     dataSource,
@@ -757,7 +757,7 @@ const variables: Record<string, (...args: any) => any> = {
   ),
 
   /**
-   * Azure Key for the Azure based export bucket srorage.
+   * Azure Key for the Azure based export bucket storage.
    */
   dbExportBucketAzureKey: ({
     dataSource,
@@ -769,6 +769,19 @@ const variables: Record<string, (...args: any) => any> = {
     ]
   ),
 
+  /**
+   * Azure SAS Token for the Azure based export bucket storage.
+   */
+  dbExportAzureSasToken: ({
+    dataSource,
+  }: {
+    dataSource: string,
+  }) => (
+    process.env[
+      keyByDataSource('CUBEJS_DB_EXPORT_BUCKET_AZURE_SAS_TOKEN', dataSource)
+    ]
+  ),
+
   /**
    * Export bucket options for Integration based.
    */

packages/cubejs-base-driver/package.json

@@ -31,7 +31,9 @@
   "dependencies": {
     "@aws-sdk/client-s3": "^3.49.0",
     "@aws-sdk/s3-request-presigner": "^3.49.0",
+    "@azure/storage-blob": "^12.9.0",
     "@cubejs-backend/shared": "^0.36.0",
+    "@google-cloud/storage": "^7.13.0",
     "ramda": "^0.27.0"
   },
   "devDependencies": {

packages/cubejs-base-driver/src/BaseDriver.ts

@@ -18,6 +18,14 @@ import { reduce } from 'ramda';
 import fs from 'fs';
 import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
 import { S3, GetObjectCommand, S3ClientConfig } from '@aws-sdk/client-s3';
+import { Storage } from '@google-cloud/storage';
+import {
+  BlobServiceClient,
+  StorageSharedKeyCredential,
+  ContainerSASPermissions,
+  SASProtocol,
+  generateBlobSASQueryParameters,
+} from '@azure/storage-blob';
 
 import { cancelCombinator } from './utils';
 import {
@@ -44,6 +52,15 @@ import {
   ForeignKeysQueryResult,
 } from './driver.interface';
 
+export type AzureStorageClientConfig = {
+  azureKey: string,
+  sasToken?: string,
+};
+
+export type GoogleStorageClientConfig = {
+  credentials: any,
+};
+
 const sortByKeys = (unordered: any) => {
   const ordered: any = {};
 
@@ -467,7 +484,7 @@ export abstract class BaseDriver implements DriverInterface {
     const conditionString = conditions.join(' OR ');
 
     const query = this.getColumnsForSpecificTablesQuery(conditionString);
-    
+
     const [primaryKeys, foreignKeys] = await Promise.all([
       this.primaryKeys(conditionString, parameters),
       this.foreignKeys(conditionString, parameters)
@@ -648,6 +665,10 @@ export abstract class BaseDriver implements DriverInterface {
     prefix: string
   ): Promise<string[]> {
     const storage = new S3(clientOptions);
+    // It looks that different driver configurations use different formats
+    // for the bucket - some expect only names, some - full url-like names.
+    // So we unify this.
+    bucketName = bucketName.replace(/^[a-zA-Z]+:\/\//, '');
 
     const list = await storage.listObjectsV2({
       Bucket: bucketName,
@@ -672,4 +693,71 @@ export abstract class BaseDriver implements DriverInterface {
 
     throw new Error('Unable to retrieve list of files from S3 storage after unloading.');
   }
+
+  /**
+   * Returns an array of signed GCS URLs of the unloaded csv files.
+   */
+  protected async extractFilesFromGCS(
+    gcsConfig: GoogleStorageClientConfig,
+    bucketName: string,
+    tableName: string
+  ): Promise<string[]> {
+    const storage = new Storage({
+      credentials: gcsConfig.credentials,
+      projectId: gcsConfig.credentials.project_id
+    });
+    const bucket = storage.bucket(bucketName);
+    const [files] = await bucket.getFiles({ prefix: `${tableName}/` });
+    if (files.length) {
+      const csvFile = await Promise.all(files.map(async (file) => {
+        const [url] = await file.getSignedUrl({
+          action: 'read',
+          expires: new Date(new Date().getTime() + 60 * 60 * 1000)
+        });
+        return url;
+      }));
+      return csvFile;
+    } else {
+      return [];
+    }
+  }
+
+  protected async extractFilesFromAzure(
+    azureConfig: AzureStorageClientConfig,
+    bucketName: string,
+    tableName: string
+  ): Promise<string[]> {
+    const parts = bucketName.split('.blob.core.windows.net/');
+    const account = parts[0];
+    const container = parts[1].split('/')[0];
+    const credential = new StorageSharedKeyCredential(account, azureConfig.azureKey);
+    const url = `https://${account}.blob.core.windows.net`;
+    const blobServiceClient = azureConfig.sasToken ?
+      new BlobServiceClient(`${url}?${azureConfig.sasToken}`) :
+      new BlobServiceClient(url, credential);
+
+    const csvFiles: string[] = [];
+    const containerClient = blobServiceClient.getContainerClient(container);
+    const blobsList = containerClient.listBlobsFlat({ prefix: `${tableName}/` });
+    for await (const blob of blobsList) {
+      if (blob.name && (blob.name.endsWith('.csv.gz') || blob.name.endsWith('.csv'))) {
+        const sas = generateBlobSASQueryParameters(
+          {
+            containerName: container,
+            blobName: blob.name,
+            permissions: ContainerSASPermissions.parse('r'),
+            startsOn: new Date(new Date().valueOf()),
+            expiresOn:
+              new Date(new Date().valueOf() + 1000 * 60 * 60),
+            protocol: SASProtocol.Https,
+            version: '2020-08-04',
+          },
+          credential,
+        ).toString();
+        csvFiles.push(`${url}/${container}/${blob.name}?${sas}`);
+      }
+    }
+
+    return csvFiles;
+  }
 }

packages/cubejs-bigquery-driver/package.json

@@ -32,7 +32,7 @@
     "@cubejs-backend/dotenv": "^9.0.2",
     "@cubejs-backend/shared": "^0.36.0",
     "@google-cloud/bigquery": "^7.7.0",
-    "@google-cloud/storage": "^7.11.1",
+    "@google-cloud/storage": "^7.13.0",
     "ramda": "^0.27.2"
   },
   "devDependencies": {

packages/cubejs-bigquery-driver/src/BigQueryDriver.ts

@@ -325,6 +325,11 @@ export class BigQueryDriver extends BaseDriver implements DriverInterface {
     const bigQueryTable = this.bigquery.dataset(schema).table(tableName);
     const [job] = await bigQueryTable.createExtractJob(destination, { format: 'CSV', gzip: true });
     await this.waitForJobResult(job, { table }, false);
+    // There is an implementation for extracting and signing urls from S3
+    // @see BaseDriver->extractUnloadedFilesFromS3()
+    // Please use that if you need. Here is a different flow
+    // because bigquery requires storage/bucket object for other things,
+    // and there is no need to initiate another one (created in extractUnloadedFilesFromS3()).
     const [files] = await this.bucket.getFiles({ prefix: `${table}-` });
     const urls = await Promise.all(files.map(async file => {
       const [url] = await file.getSignedUrl({

packages/cubejs-databricks-jdbc-driver/package.json

@@ -28,9 +28,6 @@
     "bin"
   ],
   "dependencies": {
-    "@aws-sdk/client-s3": "^3.49.0",
-    "@aws-sdk/s3-request-presigner": "^3.49.0",
-    "@azure/storage-blob": "^12.9.0",
     "@cubejs-backend/base-driver": "^0.36.0",
     "@cubejs-backend/jdbc-driver": "^0.36.0",
     "@cubejs-backend/schema-compiler": "^0.36.4",