refactor: optimize helper code

open-metadata · Jan 6, 2025 · 0a595d4 · 0a595d4
1 parent d23182a
commit 0a595d4
Showing 1 changed file with 10 additions and 11 deletions.
diff --git a/ingestion/src/metadata/ingestion/source/database/cassandra/helpers.py b/ingestion/src/metadata/ingestion/source/database/cassandra/helpers.py
@@ -57,6 +57,7 @@ def parse(cls, field) -> Column:
         data_type = None
         array_data_type = None
         raw_data_type = ""
+
         for letter in field.type:
             if letter == "<":
                 if raw_data_type in ("", "frozen"):
@@ -71,31 +72,29 @@ def parse(cls, field) -> Column:
                     array_data_type = cls.datatype_mapping.get(
                         raw_data_type.lower(), DataType.UNKNOWN
                     )
+
                 raw_data_type = ""
-                if data_type != DataType.ARRAY:
+                if data_type != DataType.ARRAY or array_data_type:
                     break
 
-            elif letter != ">":
-                raw_data_type += letter
-
             elif letter == ">":
                 if not array_data_type and data_type:
                     array_data_type = cls.datatype_mapping.get(
                         raw_data_type.lower(), DataType.UNKNOWN
                     )
                     break
-        else:
-            if not data_type:
-                data_type = cls.datatype_mapping.get(
-                    field.type.lower(), DataType.UNKNOWN
-                )
+
+            else:
+                raw_data_type += letter
+
+        if not data_type:
+            data_type = cls.datatype_mapping.get(field.type.lower(), DataType.UNKNOWN)
 
         column_def = {
             "name": field.column_name,
             "dataTypeDisplay": field.type,
             "dataType": data_type,
+            "arrayDataType": array_data_type,
         }
-        if array_data_type:
-            column_def["arrayDataType"] = array_data_type
 
         return Column(**column_def)