Repository: dbt-athena/dbt-athena
Pull request #661 (2024-05-28 12:38:20)

fix: Fixed table_type for GOVERNED tables
# Description

with following table

```
{'Name': 'test', 'DatabaseName': 'raw', 'Owner': 'hadoop', 'CreateTime': datetime.datetime(2022, 3, 8, 15, 49, 45, tzinfo=tzlocal()), 'UpdateTime': datetime.datetime(2022, 3, 8, 15, 49, 45, tzinfo=tzlocal()), 'LastAccessTime': datetime.datetime(1970, 1, 1, 1, 0, tzinfo=tzlocal()), 'Retention': 0, 'StorageDescriptor': {'Columns': [{'Name': 'brand', 'Type': 'string'}, {'Name': 'dt', 'Type': 'timestamp'}], 'Location': 's3://stage-dwh-transformed/brand_governed', 'InputFormat': 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat', 'OutputFormat': 'org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat', 'Compressed': False, 'NumberOfBuckets': -1, 'SerdeInfo': {'SerializationLibrary': 'org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe', 'Parameters': {'serialization.format': '1'}}, 'BucketColumns': [], 'SortColumns': [], 'Parameters': {}, 'SkewedInfo': {'SkewedColumnNames': [], 'SkewedColumnValues': [], 'SkewedColumnValueLocationMaps': {}}, 'StoredAsSubDirectories': False}, 'PartitionKeys': [], 'TableType': 'GOVERNED', 'Parameters': {'transient_lastDdlTime': '1646750985', 'classification': 'parquet', 'table_type': 'LAKEFORMATION_GOVERNED'}, 'CreatedBy': 'xxx', 'IsRegisteredWithLakeFormation': False, 'CatalogId': 'xxx', 'IsMultiDialectView': False}
```

I've got an error on `dbt docs generate`:

```
Encountered an error while generating catalog: 'GOVERNED'
```

I've added GOVERNED table type and improved logging as well

## 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


svdimchenko: > @svdimchenko have a look at the CI, there are some issues that needs to be addressed.

I see there is problem with https://github.com/dbt-athena/dbt-athena/blob/49896c21361ea74bb28aa1ee184c19f01d746776/tests/unit/test_adapter.py#L417-L417

do we have a use case when we need to create table without type ?
nicor88: As far as I remember there are specific scenarios (like external catalogs) where you can have tables without type, so yes, we should cover that scenarios.
nicor88: @svdimchenko final comment, https://github.com/dbt-athena/dbt-athena/pull/661#discussion_r1617360774 the rest looks good, thanks for the refactor.

Patch:
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
index c549ccb..95ee0fc 100755
--- a/dbt/adapters/athena/impl.py
+++ b/dbt/adapters/athena/impl.py
@@ -47,7 +47,6 @@ from dbt.adapters.athena.lakeformation import (
 )
 from dbt.adapters.athena.python_submissions import AthenaPythonJobHelper
 from dbt.adapters.athena.relation import (
-    RELATION_TYPE_MAP,
     AthenaRelation,
     AthenaSchemaSearchMap,
     TableType,
@@ -542,12 +541,13 @@ class AthenaAdapter(SQLAdapter):
         response = s3_client.list_objects_v2(Bucket=s3_bucket, Prefix=s3_prefix)
         return True if "Contents" in response else False
 
-    def _get_one_table_for_catalog(self, table: TableTypeDef, database: str) -> List[Dict[str, Any]]:
+    @staticmethod
+    def _get_one_table_for_catalog(table: TableTypeDef, database: str) -> List[Dict[str, Any]]:
         table_catalog = {
             "table_database": database,
             "table_schema": table["DatabaseName"],
             "table_name": table["Name"],
-            "table_type": RELATION_TYPE_MAP[table.get("TableType", "EXTERNAL_TABLE")].value,
+            "table_type": get_table_type(table).value,
             "table_comment": table.get("Parameters", {}).get("comment", table.get("Description", "")),
         }
         return [
@@ -563,14 +563,13 @@ class AthenaAdapter(SQLAdapter):
             for idx, col in enumerate(table["StorageDescriptor"]["Columns"] + table.get("PartitionKeys", []))
         ]
 
-    def _get_one_table_for_non_glue_catalog(
-        self, table: TableTypeDef, schema: str, database: str
-    ) -> List[Dict[str, Any]]:
+    @staticmethod
+    def _get_one_table_for_non_glue_catalog(table: TableTypeDef, schema: str, database: str) -> List[Dict[str, Any]]:
         table_catalog = {
             "table_database": database,
             "table_schema": schema,
             "table_name": table["Name"],
-            "table_type": RELATION_TYPE_MAP[table.get("TableType", "EXTERNAL_TABLE")].value,
+            "table_type": get_table_type(table).value,
             "table_comment": table.get("Parameters", {}).get("comment", ""),
         }
         return [
@@ -583,6 +582,7 @@ class AthenaAdapter(SQLAdapter):
                     "column_comment": col.get("Comment", ""),
                 },
             }
+            # TODO: review this code part as TableTypeDef class does not contain "Columns" attribute
             for idx, col in enumerate(table["Columns"] + table.get("PartitionKeys", []))
         ]
 
diff --git a/dbt/adapters/athena/relation.py b/dbt/adapters/athena/relation.py
index ca1bb28..3584fa3 100644
--- a/dbt/adapters/athena/relation.py
+++ b/dbt/adapters/athena/relation.py
@@ -81,6 +81,7 @@ class AthenaSchemaSearchMap(Dict[InformationSchema, Dict[str, Set[Optional[str]]
 RELATION_TYPE_MAP = {
     "EXTERNAL_TABLE": TableType.TABLE,
     "EXTERNAL": TableType.TABLE,  # type returned by federated query tables
+    "GOVERNED": TableType.TABLE,
     "MANAGED_TABLE": TableType.TABLE,
     "VIRTUAL_VIEW": TableType.VIEW,
     "table": TableType.TABLE,
@@ -91,16 +92,20 @@ RELATION_TYPE_MAP = {
 
 
 def get_table_type(table: TableTypeDef) -> TableType:
-    _type = RELATION_TYPE_MAP.get(table.get("TableType"))
-    _specific_type = table.get("Parameters", {}).get("table_type", "")
+    table_full_name = ".".join(filter(None, [table.get("CatalogId"), table.get("DatabaseName"), table["Name"]]))
 
-    if _specific_type.lower() == "iceberg":
-        _type = TableType.ICEBERG
+    input_table_type = table.get("TableType")
+    if input_table_type and input_table_type not in RELATION_TYPE_MAP:
+        raise ValueError(f"Table type {table['TableType']} is not supported for table {table_full_name}")
 
-    if _type is None:
-        raise ValueError("Table type cannot be None")
+    if table.get("Parameters", {}).get("table_type", "").lower() == "iceberg":
+        _type = TableType.ICEBERG
+    elif not input_table_type:
+        raise ValueError(f"Table type cannot be None for table {table_full_name}")
+    else:
+        _type = RELATION_TYPE_MAP[input_table_type]
 
-    LOGGER.debug(f"table_name : {table.get('Name')}")
+    LOGGER.debug(f"table_name : {table_full_name}")
     LOGGER.debug(f"table type : {_type}")
 
     return _type


Test patch:
diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
index fa4aa23..3ae756b 100644
--- a/tests/unit/test_adapter.py
+++ b/tests/unit/test_adapter.py
@@ -414,16 +414,11 @@ class TestAthenaAdapter:
         mock_aws_service.create_database("baz")
         mock_aws_service.create_table(table_name="bar", database_name="foo")
         mock_aws_service.create_table(table_name="bar", database_name="quux")
-        mock_aws_service.create_table_without_type(table_name="qux", database_name="baz")
         mock_information_schema = mock.MagicMock()
         mock_information_schema.database = "awsdatacatalog"
 
         self.adapter.acquire_connection("dummy")
-        actual = self.adapter._get_one_catalog(
-            mock_information_schema,
-            {"foo", "quux", "baz"},
-            self.used_schemas,
-        )
+        actual = self.adapter._get_one_catalog(mock_information_schema, {"foo", "quux"}, self.used_schemas)
 
         expected_column_names = (
             "table_database",
@@ -443,14 +438,16 @@ class TestAthenaAdapter:
             ("awsdatacatalog", "quux", "bar", "table", None, "id", 0, "string", None),
             ("awsdatacatalog", "quux", "bar", "table", None, "country", 1, "string", None),
             ("awsdatacatalog", "quux", "bar", "table", None, "dt", 2, "date", None),
-            ("awsdatacatalog", "baz", "qux", "table", None, "id", 0, "string", None),
-            ("awsdatacatalog", "baz", "qux", "table", None, "country", 1, "string", None),
         ]
         assert actual.column_names == expected_column_names
         assert len(actual.rows) == len(expected_rows)
         for row in actual.rows.values():
             assert row.values() in expected_rows
 
+        mock_aws_service.create_table_without_type(table_name="qux", database_name="baz")
+        with pytest.raises(ValueError):
+            self.adapter._get_one_catalog(mock_information_schema, {"baz"}, self.used_schemas)
+
     @mock_aws
     def test__get_one_catalog_by_relations(self, mock_aws_service):
         mock_aws_service.create_data_catalog()
diff --git a/tests/unit/test_relation.py b/tests/unit/test_relation.py
index fe28421..8da091a 100644
--- a/tests/unit/test_relation.py
+++ b/tests/unit/test_relation.py
@@ -8,21 +8,24 @@ TABLE_NAME = "test_table"
 
 
 class TestRelation:
-    def test__get_relation_type_table(self):
-        assert get_table_type({"Name": "name", "TableType": "table"}) == TableType.TABLE
+    @pytest.mark.parametrize(
+        ("table", "expected"),
+        [
+            ({"Name": "n", "TableType": "table"}, TableType.TABLE),
+            ({"Name": "n", "TableType": "VIRTUAL_VIEW"}, TableType.VIEW),
+            ({"Name": "n", "TableType": "EXTERNAL_TABLE", "Parameters": {"table_type": "ICEBERG"}}, TableType.ICEBERG),
+        ],
+    )
+    def test__get_relation_type(self, table, expected):
+        assert get_table_type(table) == expected
 
     def test__get_relation_type_with_no_type(self):
         with pytest.raises(ValueError):
             get_table_type({"Name": "name"})
 
-    def test__get_relation_type_view(self):
-        assert get_table_type({"Name": "name", "TableType": "VIRTUAL_VIEW"}) == TableType.VIEW
-
-    def test__get_relation_type_iceberg(self):
-        assert (
-            get_table_type({"Name": "name", "TableType": "EXTERNAL_TABLE", "Parameters": {"table_type": "ICEBERG"}})
-            == TableType.ICEBERG
-        )
+    def test__get_relation_type_with_unknown_type(self):
+        with pytest.raises(ValueError):
+            get_table_type({"Name": "name", "TableType": "test"})
 
 
 class TestAthenaRelation: