Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add detailed_table_type property to list_relations_without_caching function #513

2 changes: 2 additions & 0 deletions dbt/adapters/athena/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -712,6 +712,7 @@ def list_relations_without_caching(self, schema_relation: AthenaRelation) -> Lis
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:
Expand All @@ -724,6 +725,7 @@ def list_relations_without_caching(self, schema_relation: AthenaRelation) -> Lis
identifier=table["Name"],
quote_policy=quote_policy,
type=_type,
detailed_table_type=_detailed_table_type,
)
)
except ClientError as e:
Expand Down
1 change: 1 addition & 0 deletions dbt/adapters/athena/relation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down
44 changes: 44 additions & 0 deletions tests/functional/adapter/test_detailed_table_type.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
import re

import pytest

from dbt.tests.util import run_dbt, run_dbt_and_capture

get_detailed_table_type_sql = """
nicor88 marked this conversation as resolved.
Show resolved Hide resolved
{% 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"
17 changes: 13 additions & 4 deletions tests/unit/test_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -857,18 +857,25 @@ def test__get_data_catalog(self, mock_aws_service):
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
Expand All @@ -880,6 +887,7 @@ def test_list_relations_without_caching_with_awsdatacatalog(self, mock_aws_servi
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,
Expand All @@ -897,6 +905,7 @@ def test_list_relations_without_caching_with_other_glue_data_catalog(self, mock_
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,
Expand Down
8 changes: 7 additions & 1 deletion tests/unit/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,9 @@ def create_view(self, view_name: str):
"Location": "",
},
"TableType": "VIRTUAL_VIEW",
"Parameters": {
"TableOwner": "John Doe",
},
},
)

Expand Down Expand Up @@ -329,7 +332,7 @@ def create_iceberg_table(self, table_name: str):
"TableType": "EXTERNAL_TABLE",
"Parameters": {
"metadata_location": f"s3://{BUCKET}/tables/metadata/{table_name}/123.json",
"table_type": "iceberg",
"table_type": "ICEBERG",
},
},
)
Expand All @@ -349,6 +352,9 @@ def create_table_without_table_type(self, table_name: str):
],
"Location": f"s3://{BUCKET}/tables/{table_name}",
},
"Parameters": {
"TableOwner": "John Doe",
},
},
)

Expand Down