Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add aws_session_token for Athena mapping #663

Merged
merged 14 commits into from
Nov 10, 2023
Merged

Add aws_session_token for Athena mapping #663

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
ca65d10
Intial change for Snowflake encrypted private key
DanMawdsleyBA Nov 4, 2023
32c7f3d
Updating test description
DanMawdsleyBA Nov 4, 2023
ba8af7d
Work around for user/password mapping
DanMawdsleyBA Nov 4, 2023
c5668a7
🎨 [pre-commit.ci] Auto format from pre-commit.com hooks
pre-commit-ci[bot] Nov 4, 2023
feaae3e
Adding conditions on claiming connections
DanMawdsleyBA Nov 4, 2023
6c8ec36
🎨 [pre-commit.ci] Auto format from pre-commit.com hooks
pre-commit-ci[bot] Nov 4, 2023
8d95a49
Updating import on tests
DanMawdsleyBA Nov 4, 2023
643f03f
Or condition for user pass
DanMawdsleyBA Nov 4, 2023
0734166
🎨 [pre-commit.ci] Auto format from pre-commit.com hooks
pre-commit-ci[bot] Nov 4, 2023
8f46b55
Adding private key to env test
DanMawdsleyBA Nov 4, 2023
188fe56
Fixing typo
DanMawdsleyBA Nov 4, 2023
9cd46d2
Add `operator_args` `full_refresh` as a templated field (#623)
joppevos Nov 8, 2023
3c85872
Add aws_session_token for Athena mapping
benjamin-awd Nov 9, 2023
06aa561
Merge branch 'main' into add-athena-session-token
tatiana Nov 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions cosmos/profiles/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ def get_automatic_profile_mapping(


__all__ = [
"AthenaAccessKeyProfileMapping",
"BaseProfileMapping",
"GoogleCloudServiceAccountFileProfileMapping",
"GoogleCloudServiceAccountDictProfileMapping",
Expand Down
8 changes: 4 additions & 4 deletions cosmos/profiles/athena/access_key.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,11 @@ class AthenaAccessKeyProfileMapping(BaseProfileMapping):
"s3_staging_dir",
"schema",
]
secret_fields = [
"aws_secret_access_key",
]
secret_fields = ["aws_secret_access_key", "aws_session_token"]
airflow_param_mapping = {
"aws_access_key_id": "login",
"aws_secret_access_key": "password",
"aws_session_token": "extra.aws_session_token",
"aws_profile_name": "extra.aws_profile_name",
"database": "extra.database",
"debug_query_state": "extra.debug_query_state",
Expand All @@ -53,7 +52,8 @@ def profile(self) -> dict[str, Any | None]:
profile = {
**self.mapped_params,
**self.profile_args,
# aws_secret_access_key should always get set as env var
# aws_secret_access_key and aws_session_token should always get set as env var
"aws_secret_access_key": self.get_env_var_format("aws_secret_access_key"),
"aws_session_token": self.get_env_var_format("aws_session_token"),
}
return self.filter_null(profile)
8 changes: 7 additions & 1 deletion tests/profiles/athena/test_athena_access_key.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ def mock_athena_conn(): # type: ignore
password="my_aws_secret_key",
extra=json.dumps(
{
"aws_session_token": "token123",
"database": "my_database",
"region_name": "my_region",
"s3_staging_dir": "s3://my_bucket/dbt/",
Expand Down Expand Up @@ -107,6 +108,7 @@ def test_athena_profile_args(
"type": "athena",
"aws_access_key_id": mock_athena_conn.login,
"aws_secret_access_key": "{{ env_var('COSMOS_CONN_AWS_AWS_SECRET_ACCESS_KEY') }}",
"aws_session_token": "{{ env_var('COSMOS_CONN_AWS_AWS_SESSION_TOKEN') }}",
"database": mock_athena_conn.extra_dejson.get("database"),
"region_name": mock_athena_conn.extra_dejson.get("region_name"),
"s3_staging_dir": mock_athena_conn.extra_dejson.get("s3_staging_dir"),
Expand All @@ -122,17 +124,20 @@ def test_athena_profile_args_overrides(
"""
profile_mapping = get_automatic_profile_mapping(
mock_athena_conn.conn_id,
profile_args={"schema": "my_custom_schema", "database": "my_custom_db"},
profile_args={"schema": "my_custom_schema", "database": "my_custom_db", "aws_session_token": "override_token"},
)

assert profile_mapping.profile_args == {
"schema": "my_custom_schema",
"database": "my_custom_db",
"aws_session_token": "override_token",
}

assert profile_mapping.profile == {
"type": "athena",
"aws_access_key_id": mock_athena_conn.login,
"aws_secret_access_key": "{{ env_var('COSMOS_CONN_AWS_AWS_SECRET_ACCESS_KEY') }}",
"aws_session_token": "{{ env_var('COSMOS_CONN_AWS_AWS_SESSION_TOKEN') }}",
"database": "my_custom_db",
"region_name": mock_athena_conn.extra_dejson.get("region_name"),
"s3_staging_dir": mock_athena_conn.extra_dejson.get("s3_staging_dir"),
Expand All @@ -151,4 +156,5 @@ def test_athena_profile_env_vars(
)
assert profile_mapping.env_vars == {
"COSMOS_CONN_AWS_AWS_SECRET_ACCESS_KEY": mock_athena_conn.password,
"COSMOS_CONN_AWS_AWS_SESSION_TOKEN": mock_athena_conn.extra_dejson.get("aws_session_token"),
}
Loading