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

Jump to bottom

feat: skip_workgroup_check setting to reduce AWS throttling #713

Merged
merged 7 commits into from
Sep 12, 2024
Merged

feat: skip_workgroup_check setting to reduce AWS throttling #713

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
9ddb520
Introduce work_group_enforced
amacal Sep 5, 2024
d34558f
Add work_group_enforced into AthenaCredentials
amacal Sep 5, 2024
ddcfb01
Return predefined enforced value
amacal Sep 5, 2024
bef5b9f
Fix mypy boolean warning
amacal Sep 5, 2024
be76cdf
Add suggested impl with tests.
amacal Sep 6, 2024
7972cf2
Add documentation line
amacal Sep 6, 2024
1295374
Merge dicts using 3.8 syntax
amacal Sep 6, 2024
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
2 changes: 2 additions & 0 deletions dbt/adapters/athena/connections.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ class AthenaCredentials(Credentials):
region_name: str
endpoint_url: Optional[str] = None
work_group: Optional[str] = None
work_group_enforced: Optional[bool] = None
nicor88 marked this conversation as resolved.
Show resolved Hide resolved
nicor88 marked this conversation as resolved.
Show resolved Hide resolved
aws_profile_name: Optional[str] = None
aws_access_key_id: Optional[str] = None
aws_secret_access_key: Optional[str] = None
Expand Down Expand Up @@ -91,6 +92,7 @@ def _connection_keys(self) -> Tuple[str, ...]:
return (
"s3_staging_dir",
"work_group",
"work_group_enforced",
"region_name",
"database",
"schema",
Expand Down
5 changes: 5 additions & 0 deletions dbt/adapters/athena/impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ class AthenaConfig(AdapterConfig):

Args:
work_group: Identifier of Athena workgroup.
work_group_enforced: Assummed value of the WorkGroup location enforcement.
s3_staging_dir: S3 location to store Athena query results and metadata.
external_location: If set, the full S3 path in which the table will be saved.
partitioned_by: An array list of columns by which the table will be partitioned.
Expand All @@ -102,6 +103,7 @@ class AthenaConfig(AdapterConfig):
"""

work_group: Optional[str] = None
work_group_enforced: Optional[bool] = None
s3_staging_dir: Optional[str] = None
external_location: Optional[str] = None
partitioned_by: Optional[str] = None
Expand Down Expand Up @@ -240,6 +242,9 @@ def is_work_group_output_location_enforced(self) -> bool:
conn = self.connections.get_thread_connection()
creds = conn.credentials

if creds.work_group_enforced is not None:
nicor88 marked this conversation as resolved.
Show resolved Hide resolved
return creds.work_group_enforced
nicor88 marked this conversation as resolved.
Show resolved Hide resolved

if creds.work_group:
work_group = self._get_work_group(creds.work_group)
output_location = (
Expand Down
Loading