Repository: dbt-athena/dbt-athena
Pull request #513 (2023-11-17 16:01:41)

feat: add `detailed_table_type` property to `list_relations_without_caching` function
# Description

In `dbt/adapters/athena/impl.py`, I've added new parameter to additionally handle the `table_type` attribute for Iceberg tables. The update add new optional parameter `detailed_table_type` which gets from Table["Parameters"]["table_type"] dictionary if it not exists - set `detailed_table_type` to `None`.

I've extended the AthenaRelation class in `dbt/adapters/athena/relation.py` to include a new optional attribute `detailed_table_type`. This attribute holds the specific type of table as described in the table's parameters in the Glue Catalog, providing more detailed context about the Athena relation.

Resolves: #511 

## Checklist

- [x] You followed [contributing section](https://github.com/dbt-athena/dbt-athena#contributing)
- [x] You kept your Pull Request small and focused on a single feature or bug fix.
- [x] You added unit testing when necessary
- [x] You added functional testing when necessary


nicor88: @sanromeo nice, could you add some unit tests and functional tests if possible? For the functional tests you can add 2 models, one that is iceberg, there other calls the macro and return the detailed table type 
sanromeo: > @sanromeo nice, could you add some unit tests and functional tests if possible? For the functional tests you can add 2 models, one that is iceberg, there other calls the macro and return the detailed table type

@nicor88 no problem 😉 I will add tests


nicor88: @sanromeo thanks, also have a look at the ci failures, seems that unit tests are broken
sanromeo: @nicor88 I've added functional test for `detailed_table_type` parameter.
Could you review PR again?

Patch:
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
index a74326a..ec1bf99 100755
--- a/dbt/adapters/athena/impl.py
+++ b/dbt/adapters/athena/impl.py
@@ -712,6 +712,7 @@ class AthenaAdapter(SQLAdapter):
                         LOGGER.debug(f"Table '{table['Name']}' has no TableType attribute - Ignoring")
                         continue
                     _type = table["TableType"]
+                    _detailed_table_type = table["Parameters"].get("table_type", "")
                     if _type == "VIRTUAL_VIEW":
                         _type = self.Relation.View
                     else:
@@ -724,6 +725,7 @@ class AthenaAdapter(SQLAdapter):
                             identifier=table["Name"],
                             quote_policy=quote_policy,
                             type=_type,
+                            detailed_table_type=_detailed_table_type,
                         )
                     )
         except ClientError as e:
diff --git a/dbt/adapters/athena/relation.py b/dbt/adapters/athena/relation.py
index 680f9e0..ca1bb28 100644
--- a/dbt/adapters/athena/relation.py
+++ b/dbt/adapters/athena/relation.py
@@ -31,6 +31,7 @@ class AthenaRelation(BaseRelation):
     quote_character: str = '"'  # Presto quote character
     include_policy: Policy = field(default_factory=lambda: AthenaIncludePolicy())
     s3_path_table_part: Optional[str] = None
+    detailed_table_type: Optional[str] = None  # table_type option from the table Parameters in Glue Catalog
 
     def render_hive(self) -> str:
         """


Test patch:
diff --git a/tests/functional/adapter/test_detailed_table_type.py b/tests/functional/adapter/test_detailed_table_type.py
new file mode 100644
index 0000000..ba6c885
--- /dev/null
+++ b/tests/functional/adapter/test_detailed_table_type.py
@@ -0,0 +1,44 @@
+import re
+
+import pytest
+
+from dbt.tests.util import run_dbt, run_dbt_and_capture
+
+get_detailed_table_type_sql = """
+{% macro get_detailed_table_type(schema) %}
+    {% if execute %}
+    {% set relation = api.Relation.create(database="awsdatacatalog", schema=schema) %}
+    {% set schema_tables = adapter.list_relations_without_caching(schema_relation = relation) %}
+    {% for rel in schema_tables %}
+        {% do log('Detailed Table Type: ' ~ rel.detailed_table_type, info=True) %}
+    {% endfor %}
+    {% endif %}
+{% endmacro %}
+"""
+
+# Model SQL for an Iceberg table
+iceberg_model_sql = """
+  select 1 as id, 'iceberg' as name
+  {{ config(materialized='table', table_type='iceberg') }}
+"""
+
+
+@pytest.mark.usefixtures("project")
+class TestDetailedTableType:
+    @pytest.fixture(scope="class")
+    def macros(self):
+        return {"get_detailed_table_type.sql": get_detailed_table_type_sql}
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {"iceberg_model.sql": iceberg_model_sql}
+
+    def test_detailed_table_type(self, project):
+        # Run the models
+        run_results = run_dbt(["run"])
+        assert len(run_results) == 1  # Ensure model ran successfully
+
+        args_str = f'{{"schema": "{project.test_schema}"}}'
+        run_macro, stdout = run_dbt_and_capture(["run-operation", "get_detailed_table_type", "--args", args_str])
+        iceberg_table_type = re.search(r"Detailed Table Type: (\w+)", stdout).group(1)
+        assert iceberg_table_type == "ICEBERG"
diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
index 22a4d24..d1cf69e 100644
--- a/tests/unit/test_adapter.py
+++ b/tests/unit/test_adapter.py
@@ -857,18 +857,25 @@ class TestAthenaAdapter:
     def _test_list_relations_without_caching(self, schema_relation):
         self.adapter.acquire_connection("dummy")
         relations = self.adapter.list_relations_without_caching(schema_relation)
-        assert len(relations) == 3
+        assert len(relations) == 4
         assert all(isinstance(rel, AthenaRelation) for rel in relations)
         relations.sort(key=lambda rel: rel.name)
-        other = relations[0]
-        table = relations[1]
-        view = relations[2]
+        iceberg_table = relations[0]
+        other = relations[1]
+        table = relations[2]
+        view = relations[3]
+        assert iceberg_table.name == "iceberg"
+        assert iceberg_table.type == "table"
+        assert iceberg_table.detailed_table_type == "ICEBERG"
         assert other.name == "other"
         assert other.type == "table"
+        assert other.detailed_table_type == ""
         assert table.name == "table"
         assert table.type == "table"
+        assert table.detailed_table_type == ""
         assert view.name == "view"
         assert view.type == "view"
+        assert view.detailed_table_type == ""
 
     @mock_athena
     @mock_glue
@@ -880,6 +887,7 @@ class TestAthenaAdapter:
         mock_aws_service.create_table("other")
         mock_aws_service.create_view("view")
         mock_aws_service.create_table_without_table_type("without_table_type")
+        mock_aws_service.create_iceberg_table("iceberg")
         schema_relation = self.adapter.Relation.create(
             database=DATA_CATALOG_NAME,
             schema=DATABASE_NAME,
@@ -897,6 +905,7 @@ class TestAthenaAdapter:
         mock_aws_service.create_table("other")
         mock_aws_service.create_view("view")
         mock_aws_service.create_table_without_table_type("without_table_type")
+        mock_aws_service.create_iceberg_table("iceberg")
         schema_relation = self.adapter.Relation.create(
             database=data_catalog_name,
             schema=DATABASE_NAME,
diff --git a/tests/unit/utils.py b/tests/unit/utils.py
index 097f6ce..791181c 100644
--- a/tests/unit/utils.py
+++ b/tests/unit/utils.py
@@ -188,6 +188,9 @@ class MockAWSService:
                     "Location": "",
                 },
                 "TableType": "VIRTUAL_VIEW",
+                "Parameters": {
+                    "TableOwner": "John Doe",
+                },
             },
         )
 
@@ -329,7 +332,7 @@ class MockAWSService:
                 "TableType": "EXTERNAL_TABLE",
                 "Parameters": {
                     "metadata_location": f"s3://{BUCKET}/tables/metadata/{table_name}/123.json",
-                    "table_type": "iceberg",
+                    "table_type": "ICEBERG",
                 },
             },
         )
@@ -349,6 +352,9 @@ class MockAWSService:
                     ],
                     "Location": f"s3://{BUCKET}/tables/{table_name}",
                 },
+                "Parameters": {
+                    "TableOwner": "John Doe",
+                },
             },
         )