Athena - Get temporary credentials from the conn_id by octiva · Pull Request #758 · astronomer/astronomer-cosmos

octiva · 2023-12-11T01:54:27Z

Description

Passes the conn_id to the AwsGenericHook and uses get_credentials(), which handles the creation of a session, credentials, freezing of credentials & also masking.

See get_credentials() docs here

Related Issue(s)

#691

Breaking Change?

Checklist

I have made corresponding changes to the documentation (if required)
I have added tests that prove my fix is effective or that my feature works

netlify · 2023-12-11T01:54:32Z

👷 Deploy Preview for amazing-pothos-a3bca0 processing.

Name	Link
🔨 Latest commit	`6761548`
🔍 Latest deploy log	https://app.netlify.com/sites/amazing-pothos-a3bca0/deploys/657a59029dc4d1000891bfde

pixie79 · 2023-12-11T07:41:23Z

Do you mean dbt-Athena or dbt-athena-community. Dbt-Athena is the original module but the community version is much improved and optimised. https://github.com/dbt-athena/dbt-athena

octiva · 2023-12-11T08:39:50Z

@pixie79 the dbt-athena references the project optional dependency on line 58

dbt-athena = [
    "dbt-athena-community",
]

jbandoro

Thanks @octiva for the contribution!

I'm not familiar with Athena access keys, so if anyone from #691 wants to chime in here it would be helpful. My main comment below is if we should keep the current profile mapping of the Airflow connection, are there cases where users would still want to use the Airflow connection params, or is it broken and this PR fixes it?

octiva · 2023-12-11T23:50:36Z

@jbandoro Thanks for the feedback. If the user has an Airflow connection with no extra.role_arn provided, their credentials (secret + key id) used in the DBT profile will be unchanged, except there will now always be a aws_session_token.

octiva · 2023-12-12T01:16:14Z

Tests failing due to #761

benjamin-awd · 2023-12-13T08:52:31Z

LGTM, I tried this out today and it seems to work for my use-case -- unfortunately I don't have time to look into the memory issue and verify that this works at scale, but that's hopefully going to be resolved with the use of the AWS hook. Ran a couple of models concurrently and didn't run into any issues.

codecov · 2023-12-13T13:58:54Z

Codecov Report

Attention: 1 lines in your changes are missing coverage. Please review.

Comparison is base (9fcee8e) 93.22% compared to head (6761548) 93.22%.

Files	Patch %	Lines
cosmos/profiles/athena/access_key.py	94.44%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #758   +/-   ##
=======================================
  Coverage   93.22%   93.22%           
=======================================
  Files          55       55           
  Lines        2464     2481   +17     
=======================================
+ Hits         2297     2313   +16     
- Misses        167      168    +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

tatiana

Thanks for your contribution, @octiva! It is an excellent improvement to leverage the hook method instead of directly retrieving the values from the environment variable.

Before, we were obfuscating the secrets and setting them as environment variables. I suggest we keep this behaviour, even if the credentials are temporary. The description of your PR states that the Airflow get_credentials method masks those secrets. From what I understood from the docs:

By use this method also secret_key and token will mask in tasks logs.
https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/hooks/base_aws/index.html#airflow.providers.amazon.aws.hooks.base_aws.AwsGenericHook.get_credentials

In other words, it seems they are masked only in the logs - not in the profiles.yml file that Cosmos creates.

Assuming the value we wrote in the profiles.yml file was masked, the dbt commands using them would fail to run - since they would try to use invalid credentials.

octiva · 2023-12-13T23:26:12Z

@tatiana I am unable to write a test for the _get_temporary_credentialsmethod, as I am unable to import airflow.providers.amazon via the test dependencies as we get some dependency conflicts when using Python 3.8/3.9 + Airflow 2.3.

The amazon provider at 8.0.0 is compatible with 2.3 and python down to 3.7.

https://pypi.org/project/apache-airflow-providers-amazon/8.0.0/

I tried many ways to fix this on the weekend, but could not. I ended up resolving this by removing the dependency and mocking the class function.

This has the unfortunately side-affect of reducing our test coverage

jbandoro · 2023-12-14T00:14:22Z

Thanks @octiva addressing all the feedback! For your question below on test coverage:

I tried many ways to fix this on the weekend, but could not. I ended up resolving this by removing the dependency and mocking the class function.

This has the unfortunately side-affect of reducing our test coverage

You can get test coverage for the _get_temporary_credential method by instead patching the provider module that can't be imported, like in the examples here and here.

It might be easier to do this patching in a fixture and reuse it in all of tests that require it.

octiva · 2023-12-14T00:56:58Z

reposting from correct account @jbandoro Thanks for the hint. I had attempted this, but was doing it at the wrong level, and now ive got it working quite nicely. Let me know what you think 🚀

jbandoro

Thanks @octiva! Mocking the hook looks much better. Approving with some minor comments below to address before merging.

@woogakoki

**Features** * Add new parsing method ``LoadMode.DBT_LS_FILE`` by @woogakoki in #733 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/parsing-methods.html#dbt-ls-file)). * Add support to select using (some) graph operators when using ``LoadMode.CUSTOM`` and ``LoadMode.DBT_MANIFEST`` by @tatiana in #728 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/selecting-excluding.html#using-select-and-exclude)) * Add support for dbt ``selector`` arg for DAG parsing by @jbandoro in #755, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/render-config.html#render-config)). * Add ``ProfileMapping`` for Vertica by @perttus in #540, #688 and #741, as ([documentation](https://astronomer.github.io/astronomer-cosmos/profiles/VerticaUserPassword.html)). * Add ``ProfileMapping`` for Snowflake encrypted private key path by @ivanstillfront in #608, as ([documentation]( https://astronomer.github.io/astronomer-cosmos/profiles/SnowflakeEncryptedPrivateKeyFilePem.html)). * Add support for Snowflake encrypted private key environment variable by @DanMawdsleyBA in #649 * Add ``DbtDocsGCSOperator`` for uploading dbt docs to GCS by @jbandoro in #616, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/generating-docs.html#upload-to-gcs)). * Add cosmos/propagate_logs Airflow config support for disabling log propagation by @agreenburg in #648, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/logging.html)). * Add operator_args ``full_refresh`` as a templated field by @joppevos in #623 * Expose environment variables and dbt variables in ``ProjectConfig`` by @jbandoro in #735 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/project-config.html#project-config-example)). * Support disabling event tracking when using Cosmos profile mapping by @jbandoro in #768, ([documentation](https://astronomer.github.io/astronomer-cosmos/profiles/index.html#disabling-dbt-event-tracking)). **Enhancements** * Make Pydantic an optional dependency by @pixie79 in #736 * Create a symbolic link to ``dbt_packages`` when ``dbt_deps`` is False when using ``LoadMode.DBT_LS`` by @DanMawdsleyBA in #730 * Add ``aws_session_token`` for Athena mapping by @benjamin-awd in #663 * Retrieve temporary credentials from ``conn_id`` for Athena by @octiva in #758 * Extend ``DbtDocsLocalOperator`` with static flag by @joppevos in #759 **Bug fixes** * Remove Pydantic upper version restriction so Cosmos can be used with Airflow 2.8 by @jlaneve in #772 **Others** * Replace flake8 for Ruff by @joppevos in #743 * Reduce code complexity to 8 by @joppevos in #738 * Speed up integration tests by @jbandoro in #732 * Fix README quickstart link in by @RNHTTR in #776 * Add package location to work with hatchling 1.19.0 by @jbandoro in #761 * Fix type check error in ``DbtKubernetesBaseOperator.build_env_args`` by @jbandoro in #766 * Improve ``DBT_MANIFEST`` documentation by @dwreeves in #757 * Update conflict matrix between Airflow and dbt versions by @tatiana in #731 and #779 * pre-commit updates in #775, #770, #762

@woogakoki

**Features** * Add new parsing method ``LoadMode.DBT_LS_FILE`` by @woogakoki in astronomer#733 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/parsing-methods.html#dbt-ls-file)). * Add support to select using (some) graph operators when using ``LoadMode.CUSTOM`` and ``LoadMode.DBT_MANIFEST`` by @tatiana in astronomer#728 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/selecting-excluding.html#using-select-and-exclude)) * Add support for dbt ``selector`` arg for DAG parsing by @jbandoro in astronomer#755, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/render-config.html#render-config)). * Add ``ProfileMapping`` for Vertica by @perttus in astronomer#540, astronomer#688 and astronomer#741, as ([documentation](https://astronomer.github.io/astronomer-cosmos/profiles/VerticaUserPassword.html)). * Add ``ProfileMapping`` for Snowflake encrypted private key path by @ivanstillfront in astronomer#608, as ([documentation]( https://astronomer.github.io/astronomer-cosmos/profiles/SnowflakeEncryptedPrivateKeyFilePem.html)). * Add support for Snowflake encrypted private key environment variable by @DanMawdsleyBA in astronomer#649 * Add ``DbtDocsGCSOperator`` for uploading dbt docs to GCS by @jbandoro in astronomer#616, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/generating-docs.html#upload-to-gcs)). * Add cosmos/propagate_logs Airflow config support for disabling log propagation by @agreenburg in astronomer#648, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/logging.html)). * Add operator_args ``full_refresh`` as a templated field by @joppevos in astronomer#623 * Expose environment variables and dbt variables in ``ProjectConfig`` by @jbandoro in astronomer#735 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/project-config.html#project-config-example)). * Support disabling event tracking when using Cosmos profile mapping by @jbandoro in astronomer#768, ([documentation](https://astronomer.github.io/astronomer-cosmos/profiles/index.html#disabling-dbt-event-tracking)). **Enhancements** * Make Pydantic an optional dependency by @pixie79 in astronomer#736 * Create a symbolic link to ``dbt_packages`` when ``dbt_deps`` is False when using ``LoadMode.DBT_LS`` by @DanMawdsleyBA in astronomer#730 * Add ``aws_session_token`` for Athena mapping by @benjamin-awd in astronomer#663 * Retrieve temporary credentials from ``conn_id`` for Athena by @octiva in astronomer#758 * Extend ``DbtDocsLocalOperator`` with static flag by @joppevos in astronomer#759 **Bug fixes** * Remove Pydantic upper version restriction so Cosmos can be used with Airflow 2.8 by @jlaneve in astronomer#772 **Others** * Replace flake8 for Ruff by @joppevos in astronomer#743 * Reduce code complexity to 8 by @joppevos in astronomer#738 * Speed up integration tests by @jbandoro in astronomer#732 * Fix README quickstart link in by @RNHTTR in astronomer#776 * Add package location to work with hatchling 1.19.0 by @jbandoro in astronomer#761 * Fix type check error in ``DbtKubernetesBaseOperator.build_env_args`` by @jbandoro in astronomer#766 * Improve ``DBT_MANIFEST`` documentation by @dwreeves in astronomer#757 * Update conflict matrix between Airflow and dbt versions by @tatiana in astronomer#731 and astronomer#779 * pre-commit updates in astronomer#775, astronomer#770, astronomer#762

## Description  Passes the `conn_id` to the `AwsGenericHook` and uses `get_credentials()`, which handles the creation of a session, credentials, freezing of credentials & also masking. [See get_credentials() docs here](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/hooks/base_aws/index.html#airflow.providers.amazon.aws.hooks.base_aws.AwsGenericHook.get_credentials) ## Related Issue(s) Closes: astronomer#691 Co-authored-by: Spencer horton <spencer.horton@workcoverqld.com.au>

@woogakoki

**Features** * Add new parsing method ``LoadMode.DBT_LS_FILE`` by @woogakoki in astronomer#733 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/parsing-methods.html#dbt-ls-file)). * Add support to select using (some) graph operators when using ``LoadMode.CUSTOM`` and ``LoadMode.DBT_MANIFEST`` by @tatiana in astronomer#728 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/selecting-excluding.html#using-select-and-exclude)) * Add support for dbt ``selector`` arg for DAG parsing by @jbandoro in astronomer#755, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/render-config.html#render-config)). * Add ``ProfileMapping`` for Vertica by @perttus in astronomer#540, astronomer#688 and astronomer#741, as ([documentation](https://astronomer.github.io/astronomer-cosmos/profiles/VerticaUserPassword.html)). * Add ``ProfileMapping`` for Snowflake encrypted private key path by @ivanstillfront in astronomer#608, as ([documentation]( https://astronomer.github.io/astronomer-cosmos/profiles/SnowflakeEncryptedPrivateKeyFilePem.html)). * Add support for Snowflake encrypted private key environment variable by @DanMawdsleyBA in astronomer#649 * Add ``DbtDocsGCSOperator`` for uploading dbt docs to GCS by @jbandoro in astronomer#616, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/generating-docs.html#upload-to-gcs)). * Add cosmos/propagate_logs Airflow config support for disabling log propagation by @agreenburg in astronomer#648, ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/logging.html)). * Add operator_args ``full_refresh`` as a templated field by @joppevos in astronomer#623 * Expose environment variables and dbt variables in ``ProjectConfig`` by @jbandoro in astronomer#735 ([documentation](https://astronomer.github.io/astronomer-cosmos/configuration/project-config.html#project-config-example)). * Support disabling event tracking when using Cosmos profile mapping by @jbandoro in astronomer#768, ([documentation](https://astronomer.github.io/astronomer-cosmos/profiles/index.html#disabling-dbt-event-tracking)). **Enhancements** * Make Pydantic an optional dependency by @pixie79 in astronomer#736 * Create a symbolic link to ``dbt_packages`` when ``dbt_deps`` is False when using ``LoadMode.DBT_LS`` by @DanMawdsleyBA in astronomer#730 * Add ``aws_session_token`` for Athena mapping by @benjamin-awd in astronomer#663 * Retrieve temporary credentials from ``conn_id`` for Athena by @octiva in astronomer#758 * Extend ``DbtDocsLocalOperator`` with static flag by @joppevos in astronomer#759 **Bug fixes** * Remove Pydantic upper version restriction so Cosmos can be used with Airflow 2.8 by @jlaneve in astronomer#772 **Others** * Replace flake8 for Ruff by @joppevos in astronomer#743 * Reduce code complexity to 8 by @joppevos in astronomer#738 * Speed up integration tests by @jbandoro in astronomer#732 * Fix README quickstart link in by @RNHTTR in astronomer#776 * Add package location to work with hatchling 1.19.0 by @jbandoro in astronomer#761 * Fix type check error in ``DbtKubernetesBaseOperator.build_env_args`` by @jbandoro in astronomer#766 * Improve ``DBT_MANIFEST`` documentation by @dwreeves in astronomer#757 * Update conflict matrix between Airflow and dbt versions by @tatiana in astronomer#731 and astronomer#779 * pre-commit updates in astronomer#775, astronomer#770, astronomer#762

octiva requested a review from a team as a code owner December 11, 2023 01:54

octiva requested a review from a team December 11, 2023 01:54

dosubot Bot added the size:L This PR changes 100-499 lines, ignoring generated files. label Dec 11, 2023

octiva had a problem deploying to external December 11, 2023 01:54 — with GitHub Actions Error

octiva force-pushed the athena-assume-role branch from 92f2df4 to 5a4de6a Compare December 11, 2023 02:46

octiva had a problem deploying to external December 11, 2023 02:46 — with GitHub Actions Error

octiva mentioned this pull request Dec 11, 2023

Support automatic/short-lived token generation for AWS Athena #691

Closed

jbandoro reviewed Dec 11, 2023

View reviewed changes

Comment thread cosmos/profiles/athena/access_key.py Outdated

Comment thread dev/dags/example_athena_profile.py Outdated

octiva force-pushed the athena-assume-role branch from 5a4de6a to 182b536 Compare December 12, 2023 01:08

octiva had a problem deploying to external December 12, 2023 01:08 — with GitHub Actions Error

octiva force-pushed the athena-assume-role branch from 182b536 to 92ba184 Compare December 12, 2023 01:48

octiva had a problem deploying to external December 12, 2023 01:48 — with GitHub Actions Error

tatiana had a problem deploying to external December 12, 2023 15:38 — with GitHub Actions Error

benjamin-awd reviewed Dec 13, 2023

View reviewed changes

Comment thread cosmos/profiles/athena/access_key.py

octiva force-pushed the athena-assume-role branch from e609523 to 0d8871a Compare December 13, 2023 09:37

octiva had a problem deploying to external December 13, 2023 09:37 — with GitHub Actions Error

tatiana temporarily deployed to external December 13, 2023 13:36 — with GitHub Actions Inactive

tatiana reviewed Dec 13, 2023

View reviewed changes

Comment thread tests/profiles/athena/test_athena_access_key.py Outdated

Comment thread tests/profiles/athena/test_athena_access_key.py

Comment thread tests/profiles/athena/test_athena_access_key.py Outdated

octiva force-pushed the athena-assume-role branch from 78d01e7 to 74256ef Compare December 13, 2023 20:20

octiva had a problem deploying to external December 13, 2023 20:20 — with GitHub Actions Error

pre-commit-ci Bot had a problem deploying to external December 13, 2023 20:21 Error

octiva force-pushed the athena-assume-role branch from 832165c to 5f3a278 Compare December 13, 2023 21:06

octiva had a problem deploying to external December 13, 2023 21:06 — with GitHub Actions Error

pre-commit-ci Bot temporarily deployed to external December 13, 2023 21:07 Inactive

octiva force-pushed the athena-assume-role branch from e0e0b76 to db8720f Compare December 13, 2023 23:17

octiva had a problem deploying to external December 13, 2023 23:18 — with GitHub Actions Error

pre-commit-ci Bot had a problem deploying to external December 13, 2023 23:18 Error

octiva force-pushed the athena-assume-role branch from bb61140 to bd1b126 Compare December 14, 2023 00:47

octiva temporarily deployed to external December 14, 2023 00:47 — with GitHub Actions Inactive

jbandoro approved these changes Dec 14, 2023

View reviewed changes

Comment thread cosmos/profiles/athena/access_key.py Outdated

Comment thread tests/profiles/athena/test_athena_access_key.py

dosubot Bot added the lgtm This PR has been approved by a maintainer label Dec 14, 2023

spencer-horton-wcq added 2 commits December 14, 2023 11:23

feat: use AwsGenericHook to get creds

e786e0b

feat: add tests

6761548

octiva force-pushed the athena-assume-role branch from bd1b126 to 6761548 Compare December 14, 2023 01:23

octiva temporarily deployed to external December 14, 2023 01:23 — with GitHub Actions Inactive

jbandoro added this to the 1.3.0 milestone Dec 14, 2023

tatiana merged commit d062543 into astronomer:main Dec 14, 2023

octiva deleted the athena-assume-role branch December 15, 2023 01:31

tatiana mentioned this pull request Jan 4, 2024

Release 1.3.0 #781

Merged

pankajkoti mentioned this pull request Aug 4, 2025

AthenaAccessKeyProfileMapping does not work as expected locally #996

Closed

Conversation

octiva commented Dec 11, 2023

Description

Related Issue(s)

Breaking Change?

Checklist

Uh oh!

netlify Bot commented Dec 11, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

👷 Deploy Preview for amazing-pothos-a3bca0 processing.

Uh oh!

pixie79 commented Dec 11, 2023

Uh oh!

octiva commented Dec 11, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jbandoro left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

octiva commented Dec 11, 2023

Uh oh!

octiva commented Dec 12, 2023

Uh oh!

Uh oh!

benjamin-awd commented Dec 13, 2023

Uh oh!

codecov Bot commented Dec 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

tatiana left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

octiva commented Dec 13, 2023

Uh oh!

jbandoro commented Dec 14, 2023

Uh oh!

octiva commented Dec 14, 2023

Uh oh!

jbandoro left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

netlify Bot commented Dec 11, 2023 •

edited

Loading

octiva commented Dec 11, 2023 •

edited

Loading

codecov Bot commented Dec 13, 2023 •

edited

Loading