From 57996049fb46f488473baa00844988bb8d4f1b15 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marcin=20Szyma=C5=84ski?= <ms32035@gmail.com>
Date: Fri, 9 Sep 2022 22:07:00 +0100
Subject: [PATCH] fix(ingest): sql - add missing trino types (#5870)

Co-authored-by: Shirshanka Das <shirshanka@apache.org>
---
 .../src/datahub/ingestion/source/sql/sql_types.py            | 5 ++++-
 metadata-ingestion/tests/integration/dbt/test_dbt.py         | 2 ++
 2 files changed, 6 insertions(+), 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
index da5c5aa8d03f47..ac9e636c57ca8a 100644
--- 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
index 0a642fc6127415..20abea5b91509a 100644
--- 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):