diff --git a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py index b54192f..48b7d9e 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py +++ b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py @@ -24,6 +24,7 @@ from dve.core_engine.backends.utilities import DEFAULT_ISO_FORMATS, datetime_format_to_regex from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME from dve.core_engine.type_hints import URI, EntityName +from dve.metadata_parser.utilities import resilient_get from dve.parser.file_handling.service import LocalFilesystemImplementation, _get_implementation @@ -451,23 +452,27 @@ def get_duckdb_cast_statement_from_annotation( raise ValueError(f"dict must be `typing.TypedDict` subclass, got {type_annotation!r}") for type_ in type_annotation.mro(): - _date_format: str = getattr( # type: ignore - type_, "DATE_FORMAT", DEFAULT_ISO_FORMATS.get(type_, DEFAULT_ISO_FORMATS.get(datetime)) - ) - dt_cast_statement = rf"CASE WHEN REGEXP_FULL_MATCH(TRIM({quoted_name}), '{datetime_format_to_regex(_date_format)}') THEN TRY_STRPTIME(TRIM({quoted_name}), '{_date_format}') ELSE NULL END" # pylint: disable=C0301 - - # datetime is subclass of date, so needs to be handled first - if issubclass(type_, datetime): - stmt = rf"TRY_CAST({dt_cast_statement} as TIMESTAMP)" - return stmt - if issubclass(type_, date): - stmt = rf"TRY_CAST({dt_cast_statement} as DATE)" - return stmt - if issubclass(type_, time): - stmt = rf"TRY_CAST({dt_cast_statement} as TIME)" - return stmt - duck_type = get_duckdb_type_from_annotation(type_) - if duck_type: - stmt = f"TRIM({quoted_name})" - return _cast_as_ddb_type(stmt, type_) if parent_element else stmt + if issubclass(type_, (date, time)): + _date_format: str = resilient_get( + type_, "DATE_FORMAT", "TIME_FORMAT" + ) or DEFAULT_ISO_FORMATS.get( + type_, DEFAULT_ISO_FORMATS.get(datetime) + ) # type: ignore + dt_cast_statement = rf"CASE WHEN REGEXP_FULL_MATCH(TRIM({quoted_name}), '{datetime_format_to_regex(_date_format)}') THEN TRY_STRPTIME(TRIM({quoted_name}), '{_date_format}') ELSE NULL END" # pylint: disable=C0301 + + # datetime is subclass of date, so needs to be handled first + if issubclass(type_, datetime): + stmt = rf"TRY_CAST({dt_cast_statement} as TIMESTAMP)" + return stmt + if issubclass(type_, date): + stmt = rf"TRY_CAST({dt_cast_statement} as DATE)" + return stmt + if issubclass(type_, time): + stmt = rf"TRY_CAST({dt_cast_statement} as TIME)" + return stmt + else: + duck_type = get_duckdb_type_from_annotation(type_) + if duck_type: + stmt = f"TRIM({quoted_name})" + return _cast_as_ddb_type(stmt, type_) if parent_element else stmt raise ValueError(f"No equivalent DuckDB type for {type_annotation!r}") diff --git a/src/dve/core_engine/backends/implementations/spark/spark_helpers.py b/src/dve/core_engine/backends/implementations/spark/spark_helpers.py index ae714a6..375135c 100644 --- a/src/dve/core_engine/backends/implementations/spark/spark_helpers.py +++ b/src/dve/core_engine/backends/implementations/spark/spark_helpers.py @@ -593,29 +593,31 @@ def get_spark_cast_statement_from_annotation( raise ValueError(f"dict must be `typing.TypedDict` subclass, got {type_annotation!r}") for type_ in type_annotation.mro(): - _date_format: str = getattr( # type: ignore - type_, - "DATE_FORMAT", - DEFAULT_ISO_FORMATS.get(type_, DEFAULT_ISO_FORMATS.get(dt.datetime)), - ) - - # pylint: disable=C0301 - dt_cast_statement = f"CASE WHEN REGEXP(TRIM({quoted_name}), '{datetime_format_to_regex(_date_format)}') THEN TRY_TO_TIMESTAMP(TRIM({quoted_name}), \"{python_to_java_datetime_format(_date_format)}\") ELSE NULL END" # pylint: disable=C0301 - # datetime is subclass of date, so needs to be handled first - if issubclass(type_, dt.datetime): - return ( - _cast_as_spark_type(dt_cast_statement, type_) - if parent_element - else dt_cast_statement - ) if issubclass(type_, dt.date): - return ( - _cast_as_spark_type(dt_cast_statement, type_) - if parent_element - else dt_cast_statement + _date_format: str = getattr( # type: ignore + type_, + "DATE_FORMAT", + DEFAULT_ISO_FORMATS.get(type_, DEFAULT_ISO_FORMATS.get(dt.datetime)), ) - spark_type = get_type_from_annotation(type_) - if spark_type: - stmt = f"TRIM({quoted_name})" - return _cast_as_spark_type(stmt, type_) if parent_element else stmt + + # pylint: disable=C0301 + dt_cast_statement = f"CASE WHEN REGEXP(TRIM({quoted_name}), '{datetime_format_to_regex(_date_format)}') THEN TRY_TO_TIMESTAMP(TRIM({quoted_name}), \"{python_to_java_datetime_format(_date_format)}\") ELSE NULL END" # pylint: disable=C0301 + # datetime is subclass of date, so needs to be handled first + if issubclass(type_, dt.datetime): + return ( + _cast_as_spark_type(dt_cast_statement, type_) + if parent_element + else dt_cast_statement + ) + if issubclass(type_, dt.date): + return ( + _cast_as_spark_type(dt_cast_statement, type_) + if parent_element + else dt_cast_statement + ) + else: + spark_type = get_type_from_annotation(type_) + if spark_type: + stmt = f"TRIM({quoted_name})" + return _cast_as_spark_type(stmt, type_) if parent_element else stmt raise ValueError(f"No equivalent Spark type for {type_annotation!r}") diff --git a/src/dve/metadata_parser/utilities.py b/src/dve/metadata_parser/utilities.py index 0efa078..b20cee1 100644 --- a/src/dve/metadata_parser/utilities.py +++ b/src/dve/metadata_parser/utilities.py @@ -52,3 +52,23 @@ def chain_get( return result raise exc.TypeNotFoundError(f"Callable or type ({item!r}) not found") + + +def resilient_get(item: object, *attribute_names: str) -> Any: + """Given a number of attribute names, try to get attribute value + sequentially. Returns the first value found, and if no attributes found + returns None. + + Args: + item (object): The object to obtain attributes from (where possible) + attribute_names (str): The attribute names to search for + + Returns: + Any: The first found attribute, otherwise None + """ + for attr in attribute_names: + try: + return getattr(item, attr) + except AttributeError: + continue + return None diff --git a/tests/test_parser/test_utils.py b/tests/test_parser/test_utils.py new file mode 100644 index 0000000..bf0c700 --- /dev/null +++ b/tests/test_parser/test_utils.py @@ -0,0 +1,24 @@ +import pytest +from dve.metadata_parser.utilities import resilient_get + +class MyParent: + cls_attr = "hello" + def __init__(self, my_attr:str, another_attr:str): + self.my_attr = my_attr + self.another_attr = another_attr + +class MyObject(MyParent): + sub_attr = "bye" + def __init__(self, extra_attr:int): + self.extra_attr = extra_attr + super().__init__("from", "child") + + +@pytest.mark.parametrize("obj,attrs,expected", [(MyParent, ("cls_attr",), "hello"), + (MyObject, ("cls_attr", "sub_attr"), "hello"), + (MyObject, ("sub_attr", "cls_attr"), "bye"), + (MyParent, ("my_attr",), None), + (MyParent("this", "test"), ("extra_attr", "my_attr"), "this"), + (MyObject("this"), ("daft_attr", "another_daft_attr", "yet_another", "another_attr"), "child")]) +def test_resilient_get(obj, attrs, expected): + assert resilient_get(obj, *attrs) == expected \ No newline at end of file