Skip to content

Commit cddc1c8

Browse files
authored
chore: get boto3 client in function and cover some edge cases (#143)
1 parent 26e5f03 commit cddc1c8

File tree

3 files changed

+22
-9
lines changed

3 files changed

+22
-9
lines changed

dbt/adapters/athena/impl.py

+11-8
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import posixpath as path
22
from itertools import chain
33
from threading import Lock
4-
from typing import Any, Dict, Iterator, List, Optional, Set, Tuple
4+
from typing import Dict, Iterator, List, Optional, Set, Tuple
55
from urllib.parse import urlparse
66
from uuid import uuid4
77

@@ -144,7 +144,9 @@ def clean_up_partitions(self, database_name: str, table_name: str, where_conditi
144144
def clean_up_table(self, database_name: str, table_name: str):
145145
table_location = self.get_table_location(database_name, table_name)
146146

147-
if table_location is not None:
147+
# this check avoid issues for when the table location is an empty string
148+
# or when the table do not exist and table location is None
149+
if table_location:
148150
self.delete_from_s3(table_location)
149151

150152
@available
@@ -161,7 +163,7 @@ def delete_from_s3(self, s3_path: str):
161163
conn = self.connections.get_thread_connection()
162164
client = conn.handle
163165
bucket_name, prefix = self._parse_s3_path(s3_path)
164-
if self._s3_path_exists(client, bucket_name, prefix):
166+
if self._s3_path_exists(bucket_name, prefix):
165167
s3_resource = client.session.resource("s3", region_name=client.region_name, config=get_boto3_config())
166168
s3_bucket = s3_resource.Bucket(bucket_name)
167169
logger.debug(f"Deleting table data: path='{s3_path}', bucket='{bucket_name}', prefix='{prefix}'")
@@ -195,12 +197,13 @@ def _parse_s3_path(s3_path: str) -> Tuple[str, str]:
195197
prefix = o.path.lstrip("/").rstrip("/") + "/"
196198
return bucket_name, prefix
197199

198-
@staticmethod
199-
def _s3_path_exists(client: Any, s3_bucket: str, s3_prefix: str) -> bool:
200+
def _s3_path_exists(self, s3_bucket: str, s3_prefix: str) -> bool:
200201
"""Checks whether a given s3 path exists."""
201-
response = client.session.client(
202-
"s3", region_name=client.region_name, config=get_boto3_config()
203-
).list_objects_v2(Bucket=s3_bucket, Prefix=s3_prefix)
202+
conn = self.connections.get_thread_connection()
203+
client = conn.handle
204+
with boto3_client_lock:
205+
s3_client = client.session.client("s3", region_name=client.region_name, config=get_boto3_config())
206+
response = s3_client.list_objects_v2(Bucket=s3_bucket, Prefix=s3_prefix)
204207
return True if "Contents" in response else False
205208

206209
def _join_catalog_table_owners(self, table: agate.Table, manifest: Manifest) -> agate.Table:

tests/unit/test_adapter.py

+10
Original file line numberDiff line numberDiff line change
@@ -322,6 +322,16 @@ def test_clean_up_table_table_does_not_exist(self, dbt_debug_caplog, aws_credent
322322
assert result is None
323323
assert "Table 'table' does not exists - Ignoring" in dbt_debug_caplog.getvalue()
324324

325+
@mock_glue
326+
@mock_athena
327+
def test_clean_up_table_view(self, dbt_debug_caplog, aws_credentials):
328+
self.mock_aws_service.create_data_catalog()
329+
self.mock_aws_service.create_database()
330+
self.adapter.acquire_connection("dummy")
331+
self.mock_aws_service.create_view("test_view")
332+
result = self.adapter.clean_up_table(DATABASE_NAME, "test_view")
333+
assert result is None
334+
325335
@mock_glue
326336
@mock_s3
327337
@mock_athena

tests/unit/utils.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ def create_view(self, view_name: str):
161161
"Type": "date",
162162
},
163163
],
164-
"Location": f"s3://{BUCKET}/tables/{view_name}",
164+
"Location": "",
165165
},
166166
"TableType": "VIRTUAL_VIEW",
167167
},

0 commit comments

Comments
 (0)