fix(snowflake-driver): Add support for IAM roles with IRSA for S3 export buckets

.github/workflows/drivers-tests.yml

@@ -225,6 +225,7 @@ jobs:
         snowflake
         snowflake-encrypted-pk
         snowflake-export-bucket-s3
+        snowflake-export-bucket-s3-iam-roles
         snowflake-export-bucket-s3-prefix
         snowflake-export-bucket-azure
         snowflake-export-bucket-azure-prefix
@@ -259,6 +260,7 @@ jobs:
           - snowflake
           - snowflake-encrypted-pk
           - snowflake-export-bucket-s3
+          - snowflake-export-bucket-s3-iam-roles
           - snowflake-export-bucket-s3-prefix
           - snowflake-export-bucket-azure
           - snowflake-export-bucket-azure-prefix

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

@@ -134,10 +134,11 @@ Storage][google-cloud-storage] for export bucket functionality.
 <InfoBox>
 
 Ensure the AWS credentials are correctly configured in IAM to allow reads and
-writes to the export bucket in S3 if you are not using storage integration.
-If you are using storage integration then you still need to configure access keys
-for Cube Store to be able to read from the export bucket.
-It's possible to authenticate with IAM roles instead of access keys for Cube Store.
+writes to the export bucket in S3. You can authenticate using:
+
+- **IAM user credentials** - Explicit AWS access keys and secrets
+- **Storage integration** - Snowflake-managed AWS integration (may still require credentials for Cube Store)
+- **IAM roles** - Use execution environment roles (IRSA, instance profiles) for both Snowflake and Cube Store
 
 </InfoBox>
 
@@ -163,14 +164,21 @@ CUBEJS_DB_EXPORT_BUCKET_AWS_SECRET=<AWS_SECRET>
 CUBEJS_DB_EXPORT_BUCKET_AWS_REGION=<AWS_REGION>
 ```
 
-Using Storage Integration to write to export bocket and IAM role to read from Cube Store:
+Using Storage Integration to write to export bucket and IAM role to read from Cube Store:
 ```dotenv
 CUBEJS_DB_EXPORT_BUCKET_TYPE=s3
 CUBEJS_DB_EXPORT_BUCKET=my.bucket.on.s3
 CUBEJS_DB_EXPORT_INTEGRATION=aws_int
 CUBEJS_DB_EXPORT_BUCKET_AWS_REGION=<AWS_REGION>
 ```
 
+Using IAM roles for both Snowflake and Cube Store (no credentials required):
+```dotenv
+CUBEJS_DB_EXPORT_BUCKET_TYPE=s3
+CUBEJS_DB_EXPORT_BUCKET=my.bucket.on.s3
+CUBEJS_DB_EXPORT_BUCKET_AWS_REGION=<AWS_REGION>
+```
+
 
 #### Google Cloud Storage
 

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

@@ -106,8 +106,8 @@ const SnowflakeToGenericType: Record<string, GenericDataBaseType> = {
 interface SnowflakeDriverExportAWS {
   bucketType: 's3',
   bucketName: string,
-  keyId: string,
-  secretKey: string,
+  keyId?: string,
+  secretKey?: string,
   region: string,
   integrationName?: string,
 }
@@ -328,14 +328,17 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
     if (bucketType === 's3') {
       // integrationName is optional for s3
       const integrationName = getEnv('dbExportIntegration', { dataSource });
+      // keyId and secretKey are optional for s3 if IAM role is used
+      const keyId = getEnv('dbExportBucketAwsKey', { dataSource });
+      const secretKey = getEnv('dbExportBucketAwsSecret', { dataSource });
 
       return {
         bucketType,
         bucketName: getEnv('dbExportBucket', { dataSource }),
-        keyId: getEnv('dbExportBucketAwsKey', { dataSource }),
-        secretKey: getEnv('dbExportBucketAwsSecret', { dataSource }),
         region: getEnv('dbExportBucketAwsRegion', { dataSource }),
         ...(integrationName !== undefined && { integrationName }),
+        ...(keyId !== undefined && { keyId }),
+        ...(secretKey !== undefined && { secretKey }),
       };
     }
 
@@ -387,6 +390,23 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
     );
   }
 
+  private getRequiredExportBucketKeys(
+    exportBucket: SnowflakeDriverExportBucket,
+    emptyKeys: string[]
+  ): string[] {
+    if (exportBucket.bucketType === 's3') {
+      const s3Config = exportBucket as SnowflakeDriverExportAWS;
+      const hasCredentials = s3Config.keyId && s3Config.secretKey;
+      const hasIntegration = s3Config.integrationName;
+
+      if (!hasCredentials && !hasIntegration) {
+        return emptyKeys.filter(key => key !== 'keyId' && key !== 'secretKey');
+      }
+    }
+
+    return emptyKeys;
+  }
+
   protected getExportBucket(
     dataSource: string,
   ): SnowflakeDriverExportBucket | undefined {
@@ -402,9 +422,11 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
 
       const emptyKeys = Object.keys(exportBucket)
         .filter((key: string) => exportBucket[<keyof SnowflakeDriverExportBucket>key] === undefined);
-      if (emptyKeys.length) {
+      const keysToValidate = this.getRequiredExportBucketKeys(exportBucket, emptyKeys);
+
+      if (keysToValidate.length) {
         throw new Error(
-          `Unsupported configuration exportBucket, some configuration keys are empty: ${emptyKeys.join(',')}`
+          `Unsupported configuration exportBucket, some configuration keys are empty: ${keysToValidate.join(',')}`
         );
       }
 
@@ -731,7 +753,7 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
       // Storage integration export flow takes precedence over direct auth if it is defined
       if (conf.integrationName) {
         optionsToExport.STORAGE_INTEGRATION = conf.integrationName;
-      } else {
+      } else if (conf.keyId && conf.secretKey) {
         optionsToExport.CREDENTIALS = `(AWS_KEY_ID = '${conf.keyId}' AWS_SECRET_KEY = '${conf.secretKey}')`;
       }
     } else if (bucketType === 'gcs') {
@@ -771,14 +793,18 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
       const { bucketName, path } = this.parseBucketUrl(this.config.exportBucket!.bucketName);
       const exportPrefix = path ? `${path}/${tableName}` : tableName;
 
+      const s3Config: any = { region };
+      if (keyId && secretKey) {
+        // If access key and secret are provided, use them as credentials
+        // Otherwise, let the SDK use the default credential chain (IRSA, instance profile, etc.)
+        s3Config.credentials = {
+          accessKeyId: keyId,
+          secretAccessKey: secretKey,
+        };
+      }
+
       return this.extractUnloadedFilesFromS3(
-        {
-          credentials: {
-            accessKeyId: keyId,
-            secretAccessKey: secretKey,
-          },
-          region,
-        },
+        s3Config,
         bucketName,
         exportPrefix,
       );

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

@@ -22,6 +22,15 @@
         }
       }
     },
+    "export-bucket-s3-iam-roles": {
+      "cube": {
+        "environment": {
+          "CUBEJS_DB_EXPORT_BUCKET_TYPE": "s3",
+          "CUBEJS_DB_EXPORT_BUCKET": "snowflake-drivers-tests-preaggs",
+          "CUBEJS_DB_EXPORT_BUCKET_AWS_REGION": "us-west-1"
+        }
+      }
+    },
     "export-bucket-azure": {
       "cube": {
         "environment": {

packages/cubejs-testing-drivers/package.json

@@ -49,6 +49,7 @@
     "snowflake-full": "yarn test-driver -i dist/test/snowflake-full.test.js",
     "snowflake-encrypted-pk-full": "yarn test-driver -i dist/test/snowflake-encrypted-pk-full.test.js",
     "snowflake-export-bucket-s3-full": "yarn test-driver -i dist/test/snowflake-export-bucket-s3-full.test.js",
+    "snowflake-export-bucket-s3-iam-roles-full": "yarn test-driver -i dist/test/snowflake-export-bucket-s3-iam-roles-full.test.js",
     "snowflake-export-bucket-s3-prefix-full": "yarn test-driver -i dist/test/snowflake-export-bucket-s3-prefix-full.test.js",
     "snowflake-export-bucket-azure-full": "yarn test-driver -i dist/test/snowflake-export-bucket-azure-full.test.js",
     "snowflake-export-bucket-azure-prefix-full": "yarn test-driver -i dist/test/snowflake-export-bucket-azure-prefix-full.test.js",
@@ -59,7 +60,7 @@
     "redshift-core": "yarn test-driver -i dist/test/redshift-core.test.js",
     "redshift-full": "yarn test-driver -i dist/test/redshift-full.test.js",
     "redshift-export-bucket-s3-full": "yarn test-driver -i dist/test/redshift-export-bucket-s3-full.test.js",
-    "update-all-snapshots-local": "yarn run athena-export-bucket-s3-full --mode=local -u; yarn run bigquery-export-bucket-gcs-full --mode=local -u; yarn run clickhouse-full --mode=local -u; yarn run clickhouse-export-bucket-s3-full --mode=local -u; yarn run clickhouse-export-bucket-s3-prefix-full --mode=local -u; yarn run databricks-jdbc-export-bucket-azure-full --mode=local -u; yarn run databricks-jdbc-export-bucket-azure-prefix-full --mode=local -u; yarn run databricks-jdbc-export-bucket-gcs-full --mode=local -u; yarn run databricks-jdbc-export-bucket-gcs-prefix-full --mode=local -u; yarn run databricks-jdbc-export-bucket-s3-full --mode=local -u; yarn run databricks-jdbc-export-bucket-s3-prefix-full --mode=local -u; yarn run databricks-jdbc-full --mode=local -u; yarn run mssql-full --mode=local -u; yarn run mysql-full --mode=local -u; yarn run postgres-full --mode=local -u; yarn run redshift-export-bucket-s3-full --mode=local -u; yarn run redshift-full --mode=local -u; yarn run snowflake-encrypted-pk-full --mode=local -u; yarn run snowflake-export-bucket-azure-full --mode=local -u; yarn run snowflake-export-bucket-azure-prefix-full --mode=local -u; yarn run snowflake-export-bucket-azure-via-storage-integration-full --mode=local -u; yarn run snowflake-export-bucket-gcs-full --mode=local -u; yarn run snowflake-export-bucket-gcs-prefix-full --mode=local -u; yarn run snowflake-export-bucket-s3-full --mode=local -u; yarn run snowflake-export-bucket-s3-prefix-full --mode=local -u; yarn run snowflake-export-bucket-azure-prefix-full --mode=local -u; yarn run snowflake-export-bucket-azure-full --mode=local -u; yarn run snowflake-full --mode=local -u",
+    "update-all-snapshots-local": "yarn run athena-export-bucket-s3-full --mode=local -u; yarn run bigquery-export-bucket-gcs-full --mode=local -u; yarn run clickhouse-full --mode=local -u; yarn run clickhouse-export-bucket-s3-full --mode=local -u; yarn run clickhouse-export-bucket-s3-prefix-full --mode=local -u; yarn run databricks-jdbc-export-bucket-azure-full --mode=local -u; yarn run databricks-jdbc-export-bucket-azure-prefix-full --mode=local -u; yarn run databricks-jdbc-export-bucket-gcs-full --mode=local -u; yarn run databricks-jdbc-export-bucket-gcs-prefix-full --mode=local -u; yarn run databricks-jdbc-export-bucket-s3-full --mode=local -u; yarn run databricks-jdbc-export-bucket-s3-prefix-full --mode=local -u; yarn run databricks-jdbc-full --mode=local -u; yarn run mssql-full --mode=local -u; yarn run mysql-full --mode=local -u; yarn run postgres-full --mode=local -u; yarn run redshift-export-bucket-s3-full --mode=local -u; yarn run redshift-full --mode=local -u; yarn run snowflake-encrypted-pk-full --mode=local -u; yarn run snowflake-export-bucket-azure-full --mode=local -u; yarn run snowflake-export-bucket-azure-prefix-full --mode=local -u; yarn run snowflake-export-bucket-azure-via-storage-integration-full --mode=local -u; yarn run snowflake-export-bucket-gcs-full --mode=local -u; yarn run snowflake-export-bucket-gcs-prefix-full --mode=local -u; yarn run snowflake-export-bucket-s3-full --mode=local -u; yarn run snowflake-export-bucket-s3-iam-roles-full --mode=local -u; yarn run snowflake-export-bucket-s3-prefix-full --mode=local -u; yarn run snowflake-export-bucket-azure-prefix-full --mode=local -u; yarn run snowflake-export-bucket-azure-full --mode=local -u; yarn run snowflake-full --mode=local -u",
     "tst": "clear && yarn tsc && yarn bigquery-core"
   },
   "files": [

packages/cubejs-testing-drivers/test/snowflake-export-bucket-s3-iam-roles-full.test.ts

@@ -0,0 +1,9 @@
+import { testQueries } from '../src/tests/testQueries';
+
+testQueries('snowflake', {
+  // NOTICE: It's enough to turn on this flag only once for any one
+  // cloud storage integration. Please do not turn it on for every integration test!
+  includeIncrementalSchemaSuite: false,
+  includeHLLSuite: false,
+  extendedEnv: 'export-bucket-s3-iam-roles'
+});