impl: support profile_name when accessing Glue, S3

Tomme · Feb 16, 2022 · a9acaf0 · a9acaf0
1 parent 5991fc4
commit a9acaf0
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
@@ -1,7 +1,7 @@
 from uuid import uuid4
 import agate
 import re
-import boto3
+import boto3.session
 from botocore.exceptions import ClientError
 from typing import Optional
 
@@ -50,10 +50,11 @@ def clean_up_partitions(
     ):
         # Look up Glue partitions & clean up
         conn = self.connections.get_thread_connection()
-        client = conn.handle
+        creds = conn.credentials
+        session = boto3.session.Session(region_name=creds.region_name, profile_name=creds.aws_profile_name)
 
-        glue_client = boto3.client('glue', region_name=client.region_name)
-        s3_resource = boto3.resource('s3', region_name=client.region_name)
+        glue_client = session.client('glue')
+        s3_resource = session.resource('s3')
         partitions = glue_client.get_partitions(
             # CatalogId='123456789012', # Need to make this configurable if it is different from default AWS Account ID
             DatabaseName=database_name,
@@ -76,8 +77,10 @@ def clean_up_table(
     ):
         # Look up Glue partitions & clean up
         conn = self.connections.get_thread_connection()
-        client = conn.handle
-        glue_client = boto3.client('glue', region_name=client.region_name)
+        creds = conn.credentials
+        session = boto3.session.Session(region_name=creds.region_name, profile_name=creds.aws_profile_name)
+
+        glue_client = session.client('glue')
         try:
             table = glue_client.get_table(
                 DatabaseName=database_name,
@@ -95,7 +98,7 @@ def clean_up_table(
             if m is not None:
                 bucket_name = m.group(1)
                 prefix = m.group(2)
-                s3_resource = boto3.resource('s3', region_name=client.region_name)
+                s3_resource = session.resource('s3')
                 s3_bucket = s3_resource.Bucket(bucket_name)
                 s3_bucket.objects.filter(Prefix=prefix).delete()