helpers

Author: hectorcast-db

databricks/sdk/service/_internal.py

@@ -1,6 +1,9 @@
 import datetime
 import urllib.parse
-from typing import Callable, Dict, Generic, Optional, Type, TypeVar
+from typing import Callable, Dict, Generic, List, Optional, Type, TypeVar
+from google.protobuf.timestamp_pb2 import Timestamp
+from google.protobuf.duration_pb2 import Duration
+from databricks.sdk.common.types.fieldmask import FieldMask
 
 
 def _from_dict(d: Dict[str, any], field: str, cls: Type) -> any:
@@ -46,6 +49,59 @@ def _escape_multi_segment_path_parameter(param: str) -> str:
     return urllib.parse.quote(param)
 
 
+def _timestamp(d: Dict[str, any], field: str) -> Optional[Timestamp]:
+    if field not in d or not d[field]:
+        return None
+    ts = Timestamp()
+    ts.FromJsonString(d[field])
+    return ts
+
+def _repeated_timestamp(d: Dict[str, any], field: str) -> Optional[List[Timestamp]]:
+    if field not in d or not d[field]:
+        return None
+    result = []
+    for v in d[field]:
+        ts = Timestamp()
+        ts.FromJsonString(v)
+        result.append(ts)
+    return result
+
+def _duration(d: Dict[str, any], field: str) -> Optional[Duration]:
+    if field not in d or not d[field]:
+        return None
+    dur = Duration()
+    dur.FromJsonString(d[field])
+    return dur
+
+def _repeated_duration(d: Dict[str, any], field: str) -> Optional[List[Duration]]:
+    if field not in d or not d[field]:
+        return None
+    result = []
+    for v in d[field]:
+        dur = Duration()
+        dur.FromJsonString(v)
+        result.append(dur)
+    return result
+
+def _fieldmask(d: Dict[str, any], field: str) -> Optional[FieldMask]:
+    if field not in d or not d[field]:
+        return None
+    fm = FieldMask()
+    fm.FromJsonString(d[field])
+    return fm
+
+def _repeated_fieldmask(d: Dict[str, any], field: str) -> Optional[List[FieldMask]]:
+    if field not in d or not d[field]:
+        return None
+    result = []
+    for v in d[field]:
+        fm = FieldMask()
+        fm.FromJsonString(v)
+        result.append(fm)
+    return result
+
+
+
 ReturnType = TypeVar("ReturnType")
 
 

tests/test_internal.py

@@ -1,9 +1,13 @@
 from dataclasses import dataclass
 from enum import Enum
 
+from google.protobuf.duration_pb2 import Duration
+from google.protobuf.timestamp_pb2 import Timestamp
+from databricks.sdk.common.types.fieldmask import FieldMask
+
 from databricks.sdk.service._internal import (
-    _enum, _escape_multi_segment_path_parameter, _from_dict, _repeated_dict,
-    _repeated_enum)
+    _duration, _enum, _escape_multi_segment_path_parameter, _fieldmask, _from_dict, _repeated_dict,
+    _repeated_duration, _repeated_enum, _repeated_fieldmask, _repeated_timestamp, _timestamp)
 
 
 class A(Enum):
@@ -52,3 +56,141 @@ def test_escape_multi_segment_path_parameter():
     assert _escape_multi_segment_path_parameter("a/b") == "a/b"
     assert _escape_multi_segment_path_parameter("a?b") == "a%3Fb"
     assert _escape_multi_segment_path_parameter("a#b") == "a%23b"
+
+
+def test_timestamp():
+    # Test valid timestamp
+    result = _timestamp({"field": "2023-01-01T12:00:00Z"}, "field")
+    assert isinstance(result, Timestamp)
+    assert result.ToJsonString() == "2023-01-01T12:00:00Z"
+
+
+def test_timestamp_missing_field():
+    # Test missing field
+    assert _timestamp({}, "field") is None
+
+
+def test_timestamp_none_value():
+    # Test None value
+    assert _timestamp({"field": None}, "field") is None
+
+
+def test_timestamp_empty_value():
+    # Test empty value
+    assert _timestamp({"field": ""}, "field") is None
+
+
+def test_repeated_timestamp():
+    # Test valid repeated timestamps
+    result = _repeated_timestamp({"field": ["2023-01-01T12:00:00Z", "2023-01-02T12:00:00Z"]}, "field")
+    assert len(result) == 2
+    assert all(isinstance(ts, Timestamp) for ts in result)
+    assert result[0].ToJsonString() == "2023-01-01T12:00:00Z"
+    assert result[1].ToJsonString() == "2023-01-02T12:00:00Z"
+
+
+def test_repeated_timestamp_missing_field():
+    # Test missing field
+    assert _repeated_timestamp({}, "field") is None
+
+
+def test_repeated_timestamp_none_value():
+    # Test None value
+    assert _repeated_timestamp({"field": None}, "field") is None
+
+
+def test_repeated_timestamp_empty_list():
+    # Test empty list
+    assert _repeated_timestamp({"field": []}, "field") is None
+
+
+def test_duration():
+    # Test valid duration
+    result = _duration({"field": "3600s"}, "field")
+    assert isinstance(result, Duration)
+    assert result.ToJsonString() == "3600s"
+
+
+def test_duration_missing_field():
+    # Test missing field
+    assert _duration({}, "field") is None
+
+
+def test_duration_none_value():
+    # Test None value
+    assert _duration({"field": None}, "field") is None
+
+
+def test_duration_empty_value():
+    # Test empty value
+    assert _duration({"field": ""}, "field") is None
+
+
+def test_repeated_duration():
+    # Test valid repeated durations
+    result = _repeated_duration({"field": ["3600s", "7200s"]}, "field")
+    assert len(result) == 2
+    assert all(isinstance(dur, Duration) for dur in result)
+    assert result[0].ToJsonString() == "3600s"
+    assert result[1].ToJsonString() == "7200s"
+
+
+def test_repeated_duration_missing_field():
+    # Test missing field
+    assert _repeated_duration({}, "field") is None
+
+
+def test_repeated_duration_none_value():
+    # Test None value
+    assert _repeated_duration({"field": None}, "field") is None
+
+
+def test_repeated_duration_empty_list():
+    # Test empty list
+    assert _repeated_duration({"field": []}, "field") is None
+
+
+def test_fieldmask():
+    # Test valid fieldmask
+    result = _fieldmask({"field": "path1,path2"}, "field")
+    assert isinstance(result, FieldMask)
+    assert result.ToJsonString() == "path1,path2"
+
+
+def test_fieldmask_missing_field():
+    # Test missing field
+    assert _fieldmask({}, "field") is None
+
+
+def test_fieldmask_none_value():
+    # Test None value
+    assert _fieldmask({"field": None}, "field") is None
+
+
+def test_fieldmask_empty_value():
+    # Test empty value
+    assert _fieldmask({"field": ""}, "field") is None
+
+
+def test_repeated_fieldmask():
+    # Test valid repeated fieldmasks
+    result = _repeated_fieldmask({"field": ["path1,path2", "path3,path4"]}, "field")
+    assert len(result) == 2
+    assert all(isinstance(fm, FieldMask) for fm in result)
+    assert result[0].ToJsonString() == "path1,path2"
+    assert result[1].ToJsonString() == "path3,path4"
+
+
+def test_repeated_fieldmask_missing_field():
+    # Test missing field
+    assert _repeated_fieldmask({}, "field") is None
+
+
+def test_repeated_fieldmask_none_value():
+    # Test None value
+    assert _repeated_fieldmask({"field": None}, "field") is None
+
+
+def test_repeated_fieldmask_empty_list():
+    # Test empty list
+    assert _repeated_fieldmask({"field": []}, "field") is None