Athena - Get temporary credentials from the conn_id

[octiva]

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]

👷 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]

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]

@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]

@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]

Tests failing due to #761

[benjamin-awd]

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]

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]

@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]

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]

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.