Moved S3Url util class to util.py

tests/functional/adapter/test_basic.py

@@ -1,8 +1,6 @@
 import pytest
 
-import boto3
 import os
-from urllib.parse import urlparse
 from dbt.tests.adapter.basic.test_base import BaseSimpleMaterializations
 from dbt.tests.adapter.basic.test_singular_tests import BaseSingularTests
 from dbt.tests.adapter.basic.test_singular_tests_ephemeral import BaseSingularTestsEphemeral
@@ -29,7 +27,7 @@
     check_relations_equal,
 )
 
-from tests.util import get_s3_location, get_region
+from tests.util import get_s3_location, get_region, cleanup_s3_location
 
 
 s3bucket = get_s3_location()
@@ -61,11 +59,6 @@
 base_materialized_var_sql = config_materialized_var + config_incremental_strategy + model_base
 
 
-def cleanup_s3_location():
-    client = boto3.client("s3", region_name=region)
-    S3Url(s3bucket + schema_name).delete_all_keys_v2(client)
-
-
 class TestSimpleMaterializationsGlue(BaseSimpleMaterializations):
     # all tests within this test has the same schema
     @pytest.fixture(scope="class")
@@ -92,7 +85,7 @@ def models(self):
 
     @pytest.fixture(scope='class', autouse=True)
     def cleanup(self):
-        cleanup_s3_location()
+        cleanup_s3_location(s3bucket + schema_name, region)
         yield
 
     pass
@@ -131,7 +124,7 @@ def models(self):
 
     @pytest.fixture(scope='class', autouse=True)
     def cleanup(self):
-        cleanup_s3_location()
+        cleanup_s3_location(s3bucket + schema_name, region)
         yield
 
     # test_ephemeral with refresh table
@@ -184,7 +177,7 @@ def unique_schema(request, prefix) -> str:
 class TestIncrementalGlue(BaseIncremental):
     @pytest.fixture(scope='class', autouse=True)
     def cleanup(self):
-        cleanup_s3_location()
+        cleanup_s3_location(s3bucket + schema_name, region)
         yield
 
     @pytest.fixture(scope="class")
@@ -250,7 +243,7 @@ def unique_schema(request, prefix) -> str:
 
     @pytest.fixture(scope='class', autouse=True)
     def cleanup(self):
-        cleanup_s3_location()
+        cleanup_s3_location(s3bucket + schema_name, region)
         yield
 
     def test_generic_tests(self, project):
@@ -291,33 +284,3 @@ def test_generic_tests(self, project):
 
 #class TestSnapshotTimestampGlue(BaseSnapshotTimestamp):
 #    pass
-
-class S3Url(object):
-    def __init__(self, url):
-        self._parsed = urlparse(url, allow_fragments=False)
-
-    @property
-    def bucket(self):
-        return self._parsed.netloc
-
-    @property
-    def key(self):
-        if self._parsed.query:
-            return self._parsed.path.lstrip("/") + "?" + self._parsed.query
-        else:
-            return self._parsed.path.lstrip("/")
-
-    @property
-    def url(self):
-        return self._parsed.geturl()
-
-    def delete_all_keys_v2(self, client):
-        bucket = self.bucket
-        prefix = self.key
-
-        for response in client.get_paginator('list_objects_v2').paginate(Bucket=bucket, Prefix=prefix):
-            if 'Contents' not in response:
-                continue
-            for content in response['Contents']:
-                print("Deleting: s3://" + bucket + "/" + content['Key'])
-                client.delete_object(Bucket=bucket, Key=content['Key'])

tests/util.py

@@ -1,4 +1,6 @@
 import os
+import boto3
+from urllib.parse import urlparse
 from dbt.config.project import PartialProject
 
 
@@ -110,3 +112,38 @@ def get_s3_location():
 def get_role_arn():
     return os.environ.get("DBT_GLUE_ROLE_ARN", f"arn:aws:iam::{get_account_id()}:role/GlueInteractiveSessionRole")
 
+
+def cleanup_s3_location(path, region):
+    client = boto3.client("s3", region_name=region)
+    S3Url(path).delete_all_keys_v2(client)
+
+
+class S3Url(object):
+    def __init__(self, url):
+        self._parsed = urlparse(url, allow_fragments=False)
+
+    @property
+    def bucket(self):
+        return self._parsed.netloc
+
+    @property
+    def key(self):
+        if self._parsed.query:
+            return self._parsed.path.lstrip("/") + "?" + self._parsed.query
+        else:
+            return self._parsed.path.lstrip("/")
+
+    @property
+    def url(self):
+        return self._parsed.geturl()
+
+    def delete_all_keys_v2(self, client):
+        bucket = self.bucket
+        prefix = self.key
+
+        for response in client.get_paginator('list_objects_v2').paginate(Bucket=bucket, Prefix=prefix):
+            if 'Contents' not in response:
+                continue
+            for content in response['Contents']:
+                print("Deleting: s3://" + bucket + "/" + content['Key'])
+                client.delete_object(Bucket=bucket, Key=content['Key'])