feat(athena-driver): Add option for providing default database to use (#9735)

* feat(athena-driver): add session-token and database env vars

* update Athena driver docs

* fix comment

* chore: use CUBEJS_DB_NAME for Athena default database

* chore: revert doc changes

* chore: revert add env to aws access token

* update docs

* fix settings

---------

Co-authored-by: paulorsbrito <[email protected]>

Author: KSDaemon

docs/pages/product/configuration/data-sources/aws-athena.mdx

@@ -20,6 +20,7 @@ CUBEJS_AWS_SECRET=****************************************
 CUBEJS_AWS_REGION=us-east-1
 CUBEJS_AWS_S3_OUTPUT_LOCATION=s3://my-athena-output-bucket
 CUBEJS_AWS_ATHENA_WORKGROUP=primary
+CUBEJS_DB_NAME=my_non_default_athena_database
 CUBEJS_AWS_ATHENA_CATALOG=AwsDataCatalog
 ```
 
@@ -59,8 +60,9 @@ if [dedicated infrastructure][ref-dedicated-infra] is used. Check out the
 | `CUBEJS_AWS_S3_OUTPUT_LOCATION` | The S3 path to store query results made by the Cube deployment                                                         | A valid S3 path                                  |    ❌    |
 | `CUBEJS_AWS_ATHENA_WORKGROUP`   | The name of the workgroup in which the query is being started                                                          | [A valid Athena Workgroup][aws-athena-workgroup] |    ❌    |
 | `CUBEJS_AWS_ATHENA_CATALOG`     | The name of the catalog to use by default                                                                              | [A valid Athena Catalog name][awsdatacatalog]    |    ❌    |
+| `CUBEJS_DB_NAME`                | The name of the database to use by default                                                                             | A valid Athena Database name                     |    ❌    |
 | `CUBEJS_DB_SCHEMA`              | The name of the schema to use as `information_schema` filter. Reduces count of tables loaded during schema generation. | A valid schema name                              |    ❌    |
-| `CUBEJS_CONCURRENCY` | The number of [concurrent queries][ref-data-source-concurrency] to the data source | A valid number |    ❌    |
+| `CUBEJS_CONCURRENCY`            | The number of [concurrent queries][ref-data-source-concurrency] to the data source                                     | A valid number                                   |    ❌    |
 
 [ref-data-source-concurrency]: /product/configuration/concurrency#data-source-concurrency
 
@@ -148,4 +150,4 @@ connections are made over HTTPS.
 [ref-caching-using-preaggs-build-strats]:
   /product/caching/using-pre-aggregations#pre-aggregation-build-strategies
 [ref-schema-ref-types-formats-countdistinctapprox]: /product/data-modeling/reference/types-and-formats#count_distinct_approx
-[self-preaggs-batching]: #batching
+[self-preaggs-batching]: #batching

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

@@ -45,6 +45,7 @@ interface AthenaDriverOptions extends AthenaClientConfig {
   workGroup?: string
   catalog?: string
   schema?: string
+  database?: string
   S3OutputLocation?: string
   exportBucket?: string
   pollTimeout?: number
@@ -147,6 +148,9 @@ export class AthenaDriver extends BaseDriver implements DriverInterface {
       catalog:
         config.catalog ||
         getEnv('athenaAwsCatalog', { dataSource }),
+      database:
+        config.database ||
+        getEnv('dbName', { dataSource }),
       exportBucket:
         config.exportBucket ||
         getEnv('dbExportBucket', { dataSource }),
@@ -477,7 +481,12 @@ export class AthenaDriver extends BaseDriver implements DriverInterface {
       ResultConfiguration: {
         OutputLocation: this.config.S3OutputLocation
       },
-      ...(this.config.catalog != null ? { QueryExecutionContext: { Catalog: this.config.catalog } } : {})
+      ...(this.config.catalog || this.config.database ? {
+        QueryExecutionContext: {
+          Catalog: this.config.catalog,
+          Database: this.config.database
+        }
+      } : {})
     };
     const { QueryExecutionId } = await this.athena.startQueryExecution(request);
     return { QueryExecutionId: checkNonNullable('StartQueryExecution', QueryExecutionId) };