Skip to content

Commit

Permalink
fix(ingest/athena): detect decimal type correctly
Browse files Browse the repository at this point in the history
  • Loading branch information
bossenti committed Nov 19, 2023
1 parent b7df9e0 commit eab48b7
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 2 deletions.
4 changes: 4 additions & 0 deletions metadata-ingestion/src/datahub/ingestion/source/sql/athena.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,10 @@ def _get_column_type(
# can also be returned, so we need to extend the handling here as well
elif type_name in ["bigint", "long"]:
detected_col_type = types.BIGINT
elif type_name in ["decimal"]:
detected_col_type = types.DECIMAL
precision, scale = type_meta_information.split(",")
args = [int(precision), int(scale)]
else:
return super()._get_column_type(type_name)
return detected_col_type(*args)
Expand Down
10 changes: 8 additions & 2 deletions metadata-ingestion/tests/unit/test_athena_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,6 @@ def test_get_column_type_map():


def test_column_type_struct():

result = CustomAthenaRestDialect()._get_column_type(type_="struct<test:string>")

assert isinstance(result, STRUCT)
Expand All @@ -175,8 +174,15 @@ def test_column_type_struct():
assert isinstance(result._STRUCT_fields[0][1], types.String)


def test_column_type_complex_combination():
def test_column_type_decimal():
result = CustomAthenaRestDialect()._get_column_type(type_="decimal(10,2)")

assert isinstance(result, types.DECIMAL)
assert 10 == result.precision
assert 2 == result.scale


def test_column_type_complex_combination():
result = CustomAthenaRestDialect()._get_column_type(
type_="struct<id:string,name:string,choices:array<struct<id:string,label:string>>>"
)
Expand Down

0 comments on commit eab48b7

Please sign in to comment.