fix:ensure that format for times can be accessed and supplied in duckdb casting (#135)

Author: stevenhsd

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}")

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}")

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

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