From 0a595d4459e1c23e69e67b24db1032e1d6bb69bf Mon Sep 17 00:00:00 2001 From: Keshav Mohta Date: Mon, 6 Jan 2025 15:45:35 +0530 Subject: [PATCH] refactor: optimize helper code --- .../source/database/cassandra/helpers.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/ingestion/src/metadata/ingestion/source/database/cassandra/helpers.py b/ingestion/src/metadata/ingestion/source/database/cassandra/helpers.py index a142506233a1..832d721876af 100644 --- 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)