Address comments

Author: maiarareinaldo

dbt/adapters/athena/impl.py

@@ -2,6 +2,7 @@
 import os
 import posixpath as path
 import tempfile
+from enum import Enum
 from itertools import chain
 from textwrap import dedent
 from threading import Lock
@@ -53,6 +54,12 @@
 boto3_client_lock = Lock()
 
 
+class AthenaCatalogType(Enum):
+    GLUE = "GLUE"
+    LAMBDA = "LAMBDA"
+    HIVE = "HIVE"
+
+
 class AthenaAdapter(SQLAdapter):
     BATCH_CREATE_PARTITION_API_LIMIT = 100
     BATCH_DELETE_PARTITION_API_LIMIT = 25
@@ -416,9 +423,8 @@ def _get_one_catalog(
         manifest: Manifest,
     ) -> agate.Table:
         data_catalog = self._get_data_catalog(information_schema.path.database)
-        catalog_id = get_catalog_id(data_catalog)
 
-        if data_catalog["Type"] == "GLUE":
+        if data_catalog["Type"] == AthenaCatalogType.GLUE.value:
             conn = self.connections.get_thread_connection()
             client = conn.handle
             with boto3_client_lock:
@@ -433,6 +439,7 @@ def _get_one_catalog(
                 }
                 # If the catalog is `awsdatacatalog` we don't need to pass CatalogId as boto3
                 # infers it from the account Id.
+                catalog_id = get_catalog_id(data_catalog)
                 if catalog_id:
                     kwargs["CatalogId"] = catalog_id
 
@@ -441,13 +448,15 @@ def _get_one_catalog(
                         if relations and table["Name"] in relations:
                             catalog.extend(self._get_one_table_for_catalog(table, information_schema.path.database))
             table = agate.Table.from_object(catalog)
-        else:
+        elif data_catalog["Type"] == AthenaCatalogType.LAMBDA.value:
             kwargs = {"information_schema": information_schema, "schemas": schemas}
             table = self.execute_macro(
                 GET_CATALOG_MACRO_NAME,
                 kwargs=kwargs,
                 manifest=manifest,
             )
+        else:
+            raise NotImplementedError(f"Type of catalog {data_catalog['Type']} not supported: {data_catalog['Name']}")
 
         filtered_table = self._catalog_filter_table(table, manifest)
         return self._join_catalog_table_owners(filtered_table, manifest)

tests/unit/test_adapter.py

@@ -706,6 +706,24 @@ def test__get_one_catalog_federated_query_catalog(self, mock_execute, mock_aws_s
         for row in actual.rows.values():
             assert row.values() in expected_rows
 
+    @mock_athena
+    def test__get_one_catalog_unsupported_type(self, mock_aws_service):
+        catalog_name = "example_hive_catalog"
+        catalog_type = "HIVE"
+        mock_aws_service.create_data_catalog(catalog_name=catalog_name, catalog_type=catalog_type)
+        mock_information_schema = mock.MagicMock()
+        mock_information_schema.path.database = catalog_name
+
+        self.adapter.acquire_connection("dummy")
+
+        with pytest.raises(NotImplementedError) as exc:
+            self.adapter._get_one_catalog(
+                mock_information_schema,
+                mock.MagicMock(),
+                self.mock_manifest,
+            )
+            assert exc.message == f"Type of catalog {catalog_type} not supported: {catalog_name}"
+
     def test__get_catalog_schemas(self):
         res = self.adapter._get_catalog_schemas(self.mock_manifest)
         assert len(res.keys()) == 3