Skip to content

Commit

Permalink
chore: get boto3 client in function and cover some edge cases (#143)
Browse files Browse the repository at this point in the history
  • Loading branch information
nicor88 authored Feb 6, 2023
1 parent 26e5f03 commit cddc1c8
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 9 deletions.
19 changes: 11 additions & 8 deletions dbt/adapters/athena/impl.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import posixpath as path
from itertools import chain
from threading import Lock
from typing import Any, Dict, Iterator, List, Optional, Set, Tuple
from typing import Dict, Iterator, List, Optional, Set, Tuple
from urllib.parse import urlparse
from uuid import uuid4

Expand Down Expand Up @@ -144,7 +144,9 @@ def clean_up_partitions(self, database_name: str, table_name: str, where_conditi
def clean_up_table(self, database_name: str, table_name: str):
table_location = self.get_table_location(database_name, table_name)

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

@available
Expand All @@ -161,7 +163,7 @@ def delete_from_s3(self, s3_path: str):
conn = self.connections.get_thread_connection()
client = conn.handle
bucket_name, prefix = self._parse_s3_path(s3_path)
if self._s3_path_exists(client, bucket_name, prefix):
if self._s3_path_exists(bucket_name, prefix):
s3_resource = client.session.resource("s3", region_name=client.region_name, config=get_boto3_config())
s3_bucket = s3_resource.Bucket(bucket_name)
logger.debug(f"Deleting table data: path='{s3_path}', bucket='{bucket_name}', prefix='{prefix}'")
Expand Down Expand Up @@ -195,12 +197,13 @@ def _parse_s3_path(s3_path: str) -> Tuple[str, str]:
prefix = o.path.lstrip("/").rstrip("/") + "/"
return bucket_name, prefix

@staticmethod
def _s3_path_exists(client: Any, s3_bucket: str, s3_prefix: str) -> bool:
def _s3_path_exists(self, s3_bucket: str, s3_prefix: str) -> bool:
"""Checks whether a given s3 path exists."""
response = client.session.client(
"s3", region_name=client.region_name, config=get_boto3_config()
).list_objects_v2(Bucket=s3_bucket, Prefix=s3_prefix)
conn = self.connections.get_thread_connection()
client = conn.handle
with boto3_client_lock:
s3_client = client.session.client("s3", region_name=client.region_name, config=get_boto3_config())
response = s3_client.list_objects_v2(Bucket=s3_bucket, Prefix=s3_prefix)
return True if "Contents" in response else False

def _join_catalog_table_owners(self, table: agate.Table, manifest: Manifest) -> agate.Table:
Expand Down
10 changes: 10 additions & 0 deletions tests/unit/test_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -322,6 +322,16 @@ def test_clean_up_table_table_does_not_exist(self, dbt_debug_caplog, aws_credent
assert result is None
assert "Table 'table' does not exists - Ignoring" in dbt_debug_caplog.getvalue()

@mock_glue
@mock_athena
def test_clean_up_table_view(self, dbt_debug_caplog, aws_credentials):
self.mock_aws_service.create_data_catalog()
self.mock_aws_service.create_database()
self.adapter.acquire_connection("dummy")
self.mock_aws_service.create_view("test_view")
result = self.adapter.clean_up_table(DATABASE_NAME, "test_view")
assert result is None

@mock_glue
@mock_s3
@mock_athena
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ def create_view(self, view_name: str):
"Type": "date",
},
],
"Location": f"s3://{BUCKET}/tables/{view_name}",
"Location": "",
},
"TableType": "VIRTUAL_VIEW",
},
Expand Down

0 comments on commit cddc1c8

Please sign in to comment.