fix: avoid deleting S3 location because listing relations continues o…

…n error (dbt-labs#629)
octiva · Apr 25, 2024 · 9f5361d · 9f5361d
1 parent 107437c
commit 9f5361d
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 32 deletions.
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
@@ -694,48 +694,47 @@ def list_relations_without_caching(self, schema_relation: AthenaRelation) -> Lis
                 region_name=client.region_name,
                 config=get_boto3_config(num_retries=creds.effective_num_retries),
             )
-        paginator = glue_client.get_paginator("get_tables")
 
         kwargs = {
             "DatabaseName": schema_relation.schema,
         }
-        # If the catalog is `awsdatacatalog` we don't need to pass CatalogId as boto3 infers it from the account Id.
         if catalog_id := get_catalog_id(data_catalog):
             kwargs["CatalogId"] = catalog_id
-        page_iterator = paginator.paginate(**kwargs)
-
-        relations = []
-        quote_policy = {"database": True, "schema": True, "identifier": True}
-
+        paginator = glue_client.get_paginator("get_tables")
         try:
-            for page in page_iterator:
-                tables = page["TableList"]
-                for table in tables:
-                    if "TableType" not in table:
-                        LOGGER.debug(f"Table '{table['Name']}' has no TableType attribute - Ignoring")
-                        continue
-                    _type = table["TableType"]
-                    _detailed_table_type = table.get("Parameters", {}).get("table_type", "")
-                    if _type == "VIRTUAL_VIEW":
-                        _type = self.Relation.View
-                    else:
-                        _type = self.Relation.Table
-
-                    relations.append(
-                        self.Relation.create(
-                            schema=schema_relation.schema,
-                            database=schema_relation.database,
-                            identifier=table["Name"],
-                            quote_policy=quote_policy,
-                            type=_type,
-                            detailed_table_type=_detailed_table_type,
-                        )
-                    )
+            tables = paginator.paginate(**kwargs).build_full_result().get("TableList")
         except ClientError as e:
             # don't error out when schema doesn't exist
             # this allows dbt to create and manage schemas/databases
-            LOGGER.debug(f"Schema '{schema_relation.schema}' does not exist - Ignoring: {e}")
+            if e.response["Error"]["Code"] == "EntityNotFoundException":
+                LOGGER.debug(f"Schema '{schema_relation.schema}' does not exist - Ignoring: {e}")
+                return []
+            else:
+                raise e
 
+        relations: list[BaseRelation] = []
+        quote_policy = {"database": True, "schema": True, "identifier": True}
+        for table in tables:
+            if "TableType" not in table:
+                LOGGER.info(f"Table '{table['Name']}' has no TableType attribute - Ignoring")
+                continue
+            _type = table["TableType"]
+            _detailed_table_type = table.get("Parameters", {}).get("table_type", "")
+            if _type == "VIRTUAL_VIEW":
+                _type = self.Relation.View
+            else:
+                _type = self.Relation.Table
+
+            relations.append(
+                self.Relation.create(
+                    schema=schema_relation.schema,
+                    database=schema_relation.database,
+                    identifier=table["Name"],
+                    quote_policy=quote_policy,
+                    type=_type,
+                    detailed_table_type=_detailed_table_type,
+                )
+            )
         return relations
 
     def _get_one_catalog_by_relations(

diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
@@ -665,6 +665,17 @@ def test_list_relations_without_caching_with_other_glue_data_catalog(self, mock_
         )
         self._test_list_relations_without_caching(schema_relation)
 
+    @mock_aws
+    def test_list_relations_without_caching_on_unknown_schema(self, mock_aws_service):
+        schema_relation = self.adapter.Relation.create(
+            database=DATA_CATALOG_NAME,
+            schema="unknown_schema",
+            quote_policy=self.adapter.config.quoting,
+        )
+        self.adapter.acquire_connection("dummy")
+        relations = self.adapter.list_relations_without_caching(schema_relation)
+        assert relations == []
+
     @mock_aws
     @patch("dbt.adapters.athena.impl.SQLAdapter.list_relations_without_caching", return_value=[])
     def test_list_relations_without_caching_with_non_glue_data_catalog(

diff --git a/tests/unit/test_query_headers.py b/tests/unit/test_query_headers.py
@@ -1,7 +1,7 @@
 from unittest import mock
 
 from dbt.adapters.athena.query_headers import AthenaMacroQueryStringSetter
-from dbt.context.manifest import generate_query_header_context
+from dbt.context.query_header import generate_query_header_context
 
 from .constants import AWS_REGION, DATA_CATALOG_NAME, DATABASE_NAME
 from .utils import config_from_parts_or_dicts