SNOW-1867961: Fix from_json not working the TimestampType that contains tzinfo (#2825)

Author: sfc-gh-jrose

CHANGELOG.md

@@ -20,6 +20,10 @@
 
 - Updated README.md to include instructions on how to verify package signatures using `cosign`.
 
+#### Bug Fixes
+
+- Fixed a bug in StructField.from_json that prevented TimestampTypes with tzinfo from being parsed correctly.
+
 ### Snowpark pandas API Updates
 
 #### New Features

src/snowflake/snowpark/types.py

@@ -943,10 +943,20 @@ def _fill_ast(self, ast: proto.SpDataType) -> None:
     IntegerType,
     LongType,
     DateType,
-    TimestampType,
     NullType,
 ]
+
+_timestamp_types: List[DataType] = [
+    TimestampType(),
+    TimestampType(timezone=TimestampTimeZone.NTZ),
+    TimestampType(timezone=TimestampTimeZone.LTZ),
+    TimestampType(timezone=TimestampTimeZone.TZ),
+]
+
 _all_atomic_types: Dict[str, Type[DataType]] = {t.typeName(): t for t in _atomic_types}
+_all_timestamp_types: Dict[str, DataType] = {
+    t.json_value(): t for t in _timestamp_types
+}
 
 _complex_types: List[Type[Union[ArrayType, MapType, StructType]]] = [
     ArrayType,
@@ -964,8 +974,10 @@ def _fill_ast(self, ast: proto.SpDataType) -> None:
 
 def _parse_datatype_json_value(json_value: Union[dict, str]) -> DataType:
     if not isinstance(json_value, dict):
-        if json_value in _all_atomic_types.keys():
+        if json_value in _all_atomic_types:
             return _all_atomic_types[json_value]()
+        if json_value in _all_timestamp_types:
+            return TimestampType(timezone=_all_timestamp_types[json_value].tz)
         elif json_value == "decimal":
             return DecimalType()
         elif _FIXED_DECIMAL_PATTERN.match(json_value):

tests/unit/test_types.py

@@ -1431,6 +1431,23 @@ def test_from_json_wrong_data_type():
         StructField.from_json(wrong_json)
 
 
+def test_timestamp_json_round_trip():
+    timestamp_types = [
+        "timestamp",
+        "timestamp_tz",
+        "timestamp_ntz",
+        "timestamp_ltz",
+    ]
+
+    for ts in timestamp_types:
+        assert (
+            StructField.from_json(
+                {"name": "TS", "type": ts, "nullable": True}
+            ).json_value()["type"]
+            == ts
+        )
+
+
 def test_maptype_alias():
     expected_key = StringType()
     expected_value = IntegerType()