Repository: dbt-athena/dbt-athena
Pull request #88 (2022-12-08 16:54:28)

fix: prune static s3 locations before CTA
### Description

I took a first stab at https://github.com/dbt-athena/dbt-athena/issues/58 as a follow-up on the discussion on [slack](https://getdbt.slack.com/archives/C013MLFR7BQ/p1669050058685359):

This is resolving the endless loop of `HIVE_PATH_ALREADY_EXISTS` errors, that occur whenever dbt receives a 503 Slow Down Error from s3, during a CTA statement after it has already written partial files to s3. For typical table materializations with static s3 locations, i.e. when using `external_location` or `s3_data_naming in ['table_table', 'schema_table']`, retrying keeps failing as the table was already dropped from Glue, while the partially written files never get deleted. This is 
- adding a new `prune_s3_table_location()` method to the AthenaAdapter
- adding a new method that handles s3 deletions, that also checks whether an s3_path exists
- proposing to use `urllib.parse` to split an s3_path, rather than using regex

I refactored the different methods to use this new s3 deletion method to make it more consistent. As a consequence, I had to remove one trailing `'/'` in one of the assertions as I'm passing the untreated s3 path to the debug log. So far, all existing tests succeeded, but it might be worth adding more tests.

Are there other materializations that should run the pruning step?

## Checklist
- [x] You followed [contributing section](https://github.com/dbt-athena/dbt-athena#contributing)
- [ ] You added unit testing when necessary
- [ ] You added functional testing when necessary


nicor88: @antruo this is a relevant PR. Could you please rebase main in your branch? and be sure that the CI pass.
mattiamatrix: Please get the latest changes from `main` and update the PR title to align with conventional commits style
nicor88: @Jrmyy could you check this one? to me looks fine.

Patch:
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
index 0b780ba..157a90b 100755
--- a/dbt/adapters/athena/impl.py
+++ b/dbt/adapters/athena/impl.py
@@ -1,8 +1,8 @@
-import re
 from itertools import chain
 from os import path
 from threading import Lock
-from typing import Dict, Iterator, List, Optional, Set
+from typing import Any, Dict, Iterator, List, Optional, Set, Tuple
+from urllib.parse import urlparse
 from uuid import uuid4
 
 import agate
@@ -100,7 +100,6 @@ class AthenaAdapter(SQLAdapter):
 
         with boto3_client_lock:
             glue_client = client.session.client("glue", region_name=client.region_name, config=get_boto3_config())
-        s3_resource = client.session.resource("s3", region_name=client.region_name, config=get_boto3_config())
         paginator = glue_client.get_paginator("get_partitions")
         partition_params = {
             "DatabaseName": database_name,
@@ -110,32 +109,8 @@ class AthenaAdapter(SQLAdapter):
         }
         partition_pg = paginator.paginate(**partition_params)
         partitions = partition_pg.build_full_result().get("Partitions")
-        s3_rg = re.compile("s3://([^/]*)/(.*)")
         for partition in partitions:
-            logger.debug(
-                f"Deleting objects for partition '{partition['Values']}' "
-                f"at '{partition['StorageDescriptor']['Location']}'"
-            )
-            m = s3_rg.match(partition["StorageDescriptor"]["Location"])
-            if m is not None:
-                bucket_name = m.group(1)
-                prefix = m.group(2)
-                s3_bucket = s3_resource.Bucket(bucket_name)
-                response = s3_bucket.objects.filter(Prefix=prefix).delete()
-                is_all_successful = True
-                for res in response:
-                    if "Errors" in res:
-                        for err in res["Errors"]:
-                            is_all_successful = False
-                            logger.error(
-                                "Failed to clean up partitions: Key='{}', Code='{}', Message='{}', s3_bucket='{}'",
-                                err["Key"],
-                                err["Code"],
-                                err["Message"],
-                                bucket_name,
-                            )
-                if is_all_successful is False:
-                    raise RuntimeException("Failed to clean up table partitions.")
+            self._delete_from_s3(client, partition["StorageDescriptor"]["Location"])
 
     @available
     def clean_up_table(self, database_name: str, table_name: str):
@@ -152,20 +127,74 @@ class AthenaAdapter(SQLAdapter):
                 return
 
         if table is not None:
-            p = re.compile("s3://([^/]*)/(.*)")
-            m = p.match(table["Table"]["StorageDescriptor"]["Location"])
-            if m is not None:
-                bucket_name = m.group(1)
-                prefix = m.group(2).rstrip("/") + "/"
-                logger.debug(f"Deleting table data from 's3://{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)
-                s3_bucket.objects.filter(Prefix=prefix).delete()
+            s3_location = table["Table"]["StorageDescriptor"]["Location"]
+            self._delete_from_s3(client, s3_location)
+
+    @available
+    def prune_s3_table_location(self, s3_table_location: str):
+        """
+        Prunes an s3 table location.
+        This is ncessary resolve the HIVE_PARTITION_ALREADY_EXISTS error
+        that occurs during retrying after receiving a 503 Slow Down error
+        during a CTA command, if partial files have already been written to s3.
+        """
+        conn = self.connections.get_thread_connection()
+        client = conn.handle
+        self._delete_from_s3(client, s3_table_location)
 
     @available
     def quote_seed_column(self, column: str, quote_config: Optional[bool]) -> str:
         return super().quote_seed_column(column, False)
 
+    def _delete_from_s3(self, client: Any, s3_path: str):
+        """
+        Deletes files from s3.
+        Additionally parses the response from the s3 delete request and raises
+        a RunTimeException in case it included errors.
+        """
+        bucket_name, prefix = self._parse_s3_path(s3_path)
+        if self._s3_path_exists(client, 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}'")
+            response = s3_bucket.objects.filter(Prefix=prefix).delete()
+            is_all_successful = True
+            for res in response:
+                if "Errors" in res:
+                    for err in res["Errors"]:
+                        is_all_successful = False
+                        logger.error(
+                            "Failed to delete files: Key='{}', Code='{}', Message='{}', s3_bucket='{}'",
+                            err["Key"],
+                            err["Code"],
+                            err["Message"],
+                            bucket_name,
+                        )
+            if is_all_successful is False:
+                raise RuntimeException("Failed to delete files from S3.")
+        else:
+            logger.debug("S3 path does not exist")
+
+    @staticmethod
+    def _parse_s3_path(s3_path: str) -> Tuple[str, str]:
+        """
+        Parses and splits an s3 path into bucket name and prefix.
+        This assumes that s3_path is a prefix instead of a URI. It adds a
+        trailing slash to the prefix, if there is none.
+        """
+        o = urlparse(s3_path, allow_fragments=False)
+        bucket_name = o.netloc
+        prefix = o.path.lstrip("/").rstrip("/") + "/"
+        return bucket_name, prefix
+
+    @staticmethod
+    def _s3_path_exists(client: Any, 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)
+        return True if "Contents" in response else False
+
     def _join_catalog_table_owners(self, table: agate.Table, manifest: Manifest) -> agate.Table:
         owners = []
         # Get the owner for each model from the manifest
diff --git a/dbt/include/athena/macros/materializations/models/table/create_table_as.sql b/dbt/include/athena/macros/materializations/models/table/create_table_as.sql
index fbd88cc..39c4af5 100644
--- a/dbt/include/athena/macros/materializations/models/table/create_table_as.sql
+++ b/dbt/include/athena/macros/materializations/models/table/create_table_as.sql
@@ -13,6 +13,7 @@
 
   {%- set location_property = 'external_location' -%}
   {%- set partition_property = 'partitioned_by' -%}
+  {%- set location = adapter.s3_table_location(s3_data_dir, s3_data_naming, relation.schema, relation.identifier, external_location, temporary) -%}
 
   {%- if table_type == 'iceberg' -%}
     {%- set location_property = 'location' -%}
@@ -35,12 +36,16 @@
     {%- endif -%}
   {%- endif %}
 
+{%- if table_type != 'iceberg' -%}
+    {% do adapter.prune_s3_table_location(location) %}
+{%- endif -%}
+
   create table
     {{ relation }}
   with (
     table_type='{{ table_type }}',
     is_external={%- if table_type == 'iceberg' -%}false{%- else -%}true{%- endif %},
-    {{ location_property }}='{{ adapter.s3_table_location(s3_data_dir, s3_data_naming, relation.schema, relation.identifier, external_location, temporary) }}',
+    {{ location_property }}='{{ location }}',
   {%- if partitioned_by is not none %}
     {{ partition_property }}=ARRAY{{ partitioned_by | tojson | replace('\"', '\'') }},
   {%- endif %}


Test patch:
diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
index b873c89..2ce13a8 100644
--- a/tests/unit/test_adapter.py
+++ b/tests/unit/test_adapter.py
@@ -279,8 +279,10 @@ class TestAthenaAdapter:
         self.adapter.acquire_connection("dummy")
         self.adapter.clean_up_partitions(DATABASE_NAME, table_name, "dt < '2022-01-03'")
         assert (
-            "Deleting objects for partition '['2022-01-01']' at "
-            "'s3://test-dbt-athena-test-delete-partitions/tables/table/dt=2022-01-01'" in caplog.text
+            "Deleting table data: path="
+            "'s3://test-dbt-athena-test-delete-partitions/tables/table/dt=2022-01-01', "
+            "bucket='test-dbt-athena-test-delete-partitions', "
+            "prefix='tables/table/dt=2022-01-01/'" in caplog.text
         )
         assert (
             "Calling s3:delete_objects with {'Bucket': 'test-dbt-athena-test-delete-partitions', "
@@ -288,8 +290,10 @@ class TestAthenaAdapter:
             "{'Key': 'tables/table/dt=2022-01-01/data2.parquet'}]}}" in caplog.text
         )
         assert (
-            "Deleting objects for partition '['2022-01-02']' at "
-            "'s3://test-dbt-athena-test-delete-partitions/tables/table/dt=2022-01-02'" in caplog.text
+            "Deleting table data: path="
+            "'s3://test-dbt-athena-test-delete-partitions/tables/table/dt=2022-01-02', "
+            "bucket='test-dbt-athena-test-delete-partitions', "
+            "prefix='tables/table/dt=2022-01-02/'" in caplog.text
         )
         assert (
             "Calling s3:delete_objects with {'Bucket': 'test-dbt-athena-test-delete-partitions', "
@@ -320,7 +324,11 @@ class TestAthenaAdapter:
         self.mock_aws_service.add_data_in_table("table")
         self.adapter.acquire_connection("dummy")
         self.adapter.clean_up_table(DATABASE_NAME, "table")
-        assert "Deleting table data from 's3://test-dbt-athena-test-delete-partitions/tables/table/'" in caplog.text
+        assert (
+            "Deleting table data: path='s3://test-dbt-athena-test-delete-partitions/tables/table', "
+            "bucket='test-dbt-athena-test-delete-partitions', "
+            "prefix='tables/table/'" in caplog.text
+        )
         s3 = boto3.client("s3", region_name=AWS_REGION)
         objs = s3.list_objects_v2(Bucket=BUCKET)
         assert objs["KeyCount"] == 0
@@ -463,6 +471,16 @@ class TestAthenaAdapter:
         self.adapter.list_relations_without_caching(schema_relation)
         parent_list_relations_without_caching.assert_called_once_with(schema_relation)
 
+    @pytest.mark.parametrize(
+        "s3_path,expected",
+        [
+            ("s3://my-bucket/test-dbt/tables/schema/table", ("my-bucket", "test-dbt/tables/schema/table/")),
+            ("s3://my-bucket/test-dbt/tables/schema/table/", ("my-bucket", "test-dbt/tables/schema/table/")),
+        ],
+    )
+    def test_parse_s3_path(self, s3_path, expected):
+        assert self.adapter._parse_s3_path(s3_path) == expected
+
 
 class TestAthenaFilterCatalog:
     def test__catalog_filter_table(self):