fix(ingest): sql - add missing trino types (#5870)

Co-authored-by: Shirshanka Das <[email protected]>
datahub-project · Sep 9, 2022 · 5799604 · 5799604
1 parent 5074035
commit 5799604
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_types.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_types.py
@@ -7,6 +7,7 @@
     BytesType,
     DateType,
     EnumType,
+    MapType,
     NullType,
     NumberType,
     RecordType,
@@ -321,7 +322,7 @@ def resolve_trino_modified_type(type_string: str) -> Any:
 
 # https://trino.io/docs/current/language/types.html
 # https://github.com/trinodb/trino-python-client/blob/master/trino/sqlalchemy/datatype.py#L75
-TRINO_SQL_TYPES_MAP = {
+TRINO_SQL_TYPES_MAP: Dict[str, Any] = {
     "boolean": BooleanType,
     "tinyint": NumberType,
     "smallint": NumberType,
@@ -339,4 +340,6 @@ def resolve_trino_modified_type(type_string: str) -> Any:
     "time": TimeType,
     "timestamp": TimeType,
     "row": RecordType,
+    "map": MapType,
+    "array": ArrayType,
 }
diff --git a/metadata-ingestion/tests/integration/dbt/test_dbt.py b/metadata-ingestion/tests/integration/dbt/test_dbt.py
@@ -668,6 +668,8 @@ def test_dbt_stateful_tests(pytestconfig, tmp_path, mock_time, mock_datahub_grap
         ("timestamp", "timestamp"),
         ("timestamp(3)", "timestamp"),
         ("row(x bigint, y double)", "row"),
+        ("array(row(x bigint, y double))", "array"),
+        ("map(varchar, varchar)", "map"),
     ],
 )
 def test_resolve_trino_modified_type(data_type, expected_data_type):