Cursor.description collapses TIMESTAMP_NTZ to 'timestamp' on the SELECT path (#786)

Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com>

Author: peco-engineer-bot[bot]

src/databricks/sql/backend/thrift_backend.py

@@ -777,8 +777,10 @@ def _col_to_description(col, field=None, host_url=None):
                     sql_type = field.metadata.get(b"Spark:DataType:SqlName")
                     if sql_type == b"VARIANT":
                         cleaned_type = "variant"
+                    elif sql_type == b"TIMESTAMP_NTZ":
+                        cleaned_type = "timestamp_ntz"
             except Exception as e:
-                logger.debug(f"Could not extract variant type from field: {e}")
+                logger.debug(f"Could not extract type from field metadata: {e}")
 
         return col.columnName, cleaned_type, None, None, precision, scale, None
 

src/databricks/sql/utils.py

@@ -771,7 +771,7 @@ def convert_to_assigned_datatypes_in_column_table(column_table, description):
             converted_column_table.append(
                 tuple(v if v is None else datetime.date.fromisoformat(v) for v in col)
             )
-        elif description[i][1] == "timestamp":
+        elif description[i][1] in ("timestamp", "timestamp_ntz"):
             converted_column_table.append(
                 tuple((v if v is None else parser.parse(v)) for v in col)
             )

tests/e2e/test_driver.py

@@ -1043,6 +1043,22 @@ def test_timezone_with_timestamp(self):
                 assert arrow_result_table.field(0).type == ts_type
                 assert arrow_result_value == expected.timestamp() * 1000000
 
+    def test_timestamp_ntz_description_type_code(self):
+        # See issue #786: cursor.description must distinguish TIMESTAMP_NTZ
+        # from TIMESTAMP. Both arrive over Thrift as TTypeId.TIMESTAMP_TYPE,
+        # so the type_code must be recovered from the Arrow field metadata.
+        with self.cursor() as cursor:
+            cursor.execute(
+                "SELECT "
+                "  CAST('2024-10-07 12:00:00' AS TIMESTAMP)     AS tz_aware, "
+                "  CAST('2024-10-07 12:00:00' AS TIMESTAMP_NTZ) AS tz_naive"
+            )
+            description = cursor.description
+            assert description[0][0] == "tz_aware"
+            assert description[0][1] == "timestamp"
+            assert description[1][0] == "tz_naive"
+            assert description[1][1] == "timestamp_ntz"
+
     @skipUnless(pysql_supports_arrow(), "arrow test needs arrow support")
     def test_can_flip_compression(self):
         with self.cursor() as cursor:

tests/unit/test_thrift_backend.py

@@ -2454,21 +2454,44 @@ def test_execute_command_sets_complex_type_fields_correctly(
     @unittest.skipIf(pyarrow is None, "Requires pyarrow")
     def test_col_to_description(self):
         test_cases = [
-            ("variant_col", {b"Spark:DataType:SqlName": b"VARIANT"}, "variant"),
-            ("normal_col", {}, "string"),
+            (
+                "variant_col",
+                ttypes.TTypeId.STRING_TYPE,
+                {b"Spark:DataType:SqlName": b"VARIANT"},
+                "variant",
+            ),
+            (
+                "timestamp_ntz_col",
+                ttypes.TTypeId.TIMESTAMP_TYPE,
+                {b"Spark:DataType:SqlName": b"TIMESTAMP_NTZ"},
+                "timestamp_ntz",
+            ),
+            (
+                "timestamp_col",
+                ttypes.TTypeId.TIMESTAMP_TYPE,
+                {b"Spark:DataType:SqlName": b"TIMESTAMP"},
+                "timestamp",
+            ),
+            ("normal_col", ttypes.TTypeId.STRING_TYPE, {}, "string"),
             (
                 "weird_field",
+                ttypes.TTypeId.STRING_TYPE,
                 {b"Spark:DataType:SqlName": b"Some unexpected value"},
                 "string",
             ),
-            ("missing_field", None, "string"),  # None field case
+            (
+                "missing_field",
+                ttypes.TTypeId.STRING_TYPE,
+                None,
+                "string",
+            ),  # None field case
         ]
 
-        for column_name, field_metadata, expected_type in test_cases:
+        for column_name, primitive_type, field_metadata, expected_type in test_cases:
             with self.subTest(column_name=column_name, expected_type=expected_type):
                 col = ttypes.TColumnDesc(
                     columnName=column_name,
-                    typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
+                    typeDesc=self._make_type_desc(primitive_type),
                 )
 
                 field = (

tests/unit/test_util.py

@@ -30,7 +30,7 @@ def get_column_table_and_description(self):
             ("decimal_column", "decimal", None, None, 10, 2, None),
             ("date_column", "date", None, None, None, None, None),
             ("timestamp_column", "timestamp", None, None, None, None, None),
-            ("timestamp_ntz_column", "timestamp", None, None, None, None, None),
+            ("timestamp_ntz_column", "timestamp_ntz", None, None, None, None, None),
             ("timestamp_column_2", "timestamp", None, None, None, None, None),
             ("timestamp_column_3", "timestamp", None, None, None, None, None),
             ("timestamp_column_4", "timestamp", None, None, None, None, None),