From 46c6ce0cc98c26eb594364bb05d1a855a310bc76 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" <3815206+peco-engineer-bot[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:43:44 +0000 Subject: [PATCH] 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> --- src/databricks/sql/backend/thrift_backend.py | 4 ++- src/databricks/sql/utils.py | 2 +- tests/e2e/test_driver.py | 16 ++++++++++ tests/unit/test_thrift_backend.py | 33 +++++++++++++++++--- tests/unit/test_util.py | 2 +- 5 files changed, 49 insertions(+), 8 deletions(-) diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index e047aedf6..f631125e8 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/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 diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 7e2abc07d..ef3dd4f8b 100644 --- a/src/databricks/sql/utils.py +++ b/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) ) diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 5fe3db037..8523d1eb7 100644 --- a/tests/e2e/test_driver.py +++ b/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: diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 4746b18ff..7abb3befb 100644 --- a/tests/unit/test_thrift_backend.py +++ b/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 = ( diff --git a/tests/unit/test_util.py b/tests/unit/test_util.py index 94a96b54a..4230fcaef 100644 --- a/tests/unit/test_util.py +++ b/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),