fix: mapper return type (#16964)

(cherry picked from commit 2ad53a3)
open-metadata · Jul 10, 2024 · 5accefe · 5accefe
1 parent 3550ad8
commit 5accefe
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 7 deletions.
diff --git a/ingestion/src/metadata/profiler/orm/converter/bigquery/converter.py b/ingestion/src/metadata/profiler/orm/converter/bigquery/converter.py
@@ -15,6 +15,10 @@
 """
 
 
+from typing import Dict, Set
+
+from sqlalchemy.sql.sqltypes import TypeEngine
+
 from metadata.generated.schema.entity.data.database import databaseService
 from metadata.generated.schema.entity.data.table import Column, DataType
 from metadata.profiler.orm.converter.common import CommonMapTypes
@@ -31,12 +35,12 @@ def return_custom_type(self, col: Column, table_service_type):
         return super().return_custom_type(col, table_service_type)
 
     @staticmethod
-    def map_sqa_to_om_types() -> dict:
+    def map_sqa_to_om_types() -> Dict[TypeEngine, Set[DataType]]:
         """returns an ORM type"""
         # pylint: disable=import-outside-toplevel
         from sqlalchemy_bigquery import STRUCT
 
         return {
             **CommonMapTypes.map_sqa_to_om_types(),
-            STRUCT: DataType.STRUCT,
+            STRUCT: {DataType.STRUCT},
         }
diff --git a/ingestion/src/metadata/profiler/orm/converter/common.py b/ingestion/src/metadata/profiler/orm/converter/common.py
@@ -13,7 +13,10 @@
 Common Class For Profiler Converter.
 """
 
+from typing import Dict, Set
+
 import sqlalchemy
+from sqlalchemy.sql.sqltypes import TypeEngine
 
 from metadata.generated.schema.entity.data.table import Column, DataType
 from metadata.ingestion.source import sqa_types
@@ -81,7 +84,7 @@ def return_custom_type(self, col: Column, _):
         return self._TYPE_MAP.get(col.dataType)
 
     @staticmethod
-    def map_sqa_to_om_types() -> dict:
+    def map_sqa_to_om_types() -> Dict[TypeEngine, Set[DataType]]:
         """returns an ORM type"""
         return {
             sqlalchemy.NUMERIC: {DataType.NUMBER, DataType.NUMERIC},

diff --git a/ingestion/src/metadata/profiler/orm/converter/redshift/converter.py b/ingestion/src/metadata/profiler/orm/converter/redshift/converter.py
@@ -14,6 +14,10 @@
 to an SQLAlchemy ORM class.
 """
 
+from typing import Dict, Set
+
+from sqlalchemy.sql.sqltypes import TypeEngine
+
 from metadata.generated.schema.entity.data.table import Column, DataType
 from metadata.generated.schema.entity.services import databaseService
 from metadata.profiler.orm.converter.common import CommonMapTypes
@@ -37,12 +41,12 @@ def return_custom_type(self, col: Column, table_service_type):
         return super().return_custom_type(col, table_service_type)
 
     @staticmethod
-    def map_sqa_to_om_types() -> dict:
+    def map_sqa_to_om_types() -> Dict[TypeEngine, Set[DataType]]:
         """returns an ORM type"""
         # pylint: disable=import-outside-toplevel
         from sqlalchemy_redshift.dialect import GEOMETRY
 
         return {
             **CommonMapTypes.map_sqa_to_om_types(),
-            GEOMETRY: DataType.GEOMETRY,
+            GEOMETRY: {DataType.GEOMETRY},
         }
diff --git a/ingestion/src/metadata/profiler/orm/converter/snowflake/converter.py b/ingestion/src/metadata/profiler/orm/converter/snowflake/converter.py
@@ -15,6 +15,10 @@
 """
 
 
+from typing import Dict, Set
+
+from sqlalchemy.sql.sqltypes import TypeEngine
+
 from metadata.generated.schema.entity.data.database import databaseService
 from metadata.generated.schema.entity.data.table import Column, DataType
 from metadata.profiler.orm.converter.common import CommonMapTypes
@@ -39,12 +43,12 @@ def return_custom_type(self, col: Column, table_service_type):
         return super().return_custom_type(col, table_service_type)
 
     @staticmethod
-    def map_sqa_to_om_types() -> dict:
+    def map_sqa_to_om_types() -> Dict[TypeEngine, Set[DataType]]:
         """returns an ORM type"""
         # pylint: disable=import-outside-toplevel
         from snowflake.sqlalchemy import VARIANT
 
         return {
             **CommonMapTypes.map_sqa_to_om_types(),
-            VARIANT: DataType.JSON,
+            VARIANT: {DataType.JSON},
         }