fix(thrift): report geospatial cursor metadata

Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com>

Author: cathleeny

src/databricks/sql/backend/thrift_backend.py

@@ -3,6 +3,7 @@
 import errno
 import logging
 import math
+import re
 import time
 import threading
 from typing import Dict, List, Optional, Union, Any, TYPE_CHECKING
@@ -788,16 +789,20 @@ def _col_to_description(col, field=None, host_url=None):
         else:
             precision, scale = None, None
 
-        # Extract variant type from field if available
+        # Extract logical types that Thrift surfaces as strings from Arrow metadata.
         if field is not None:
             try:
-                # Check for variant type in metadata
                 if field.metadata and b"Spark:DataType:SqlName" in field.metadata:
                     sql_type = field.metadata.get(b"Spark:DataType:SqlName")
-                    if sql_type == b"VARIANT":
-                        cleaned_type = "variant"
+                    logical_type = re.fullmatch(
+                        rb"(variant|geometry|geography)(?:\([^()]*\))?",
+                        sql_type.strip(),
+                        re.IGNORECASE,
+                    )
+                    if logical_type:
+                        cleaned_type = logical_type.group(1).decode("ascii").lower()
             except Exception as e:
-                logger.debug(f"Could not extract variant type from field: {e}")
+                logger.debug(f"Could not extract logical type from field: {e}")
 
         return col.columnName, cleaned_type, None, None, precision, scale, None
 

tests/unit/test_thrift_backend.py

@@ -2534,12 +2534,37 @@ def test_execute_command_sets_complex_type_fields_correctly(
     def test_col_to_description(self):
         test_cases = [
             ("variant_col", {b"Spark:DataType:SqlName": b"VARIANT"}, "variant"),
+            (
+                "geometry_col",
+                {b"Spark:DataType:SqlName": b"GEOMETRY(0)"},
+                "geometry",
+            ),
+            (
+                "geography_col",
+                {b"Spark:DataType:SqlName": b"GEOGRAPHY(4326)"},
+                "geography",
+            ),
+            (
+                "geometry_any_col",
+                {b"Spark:DataType:SqlName": b"geometry(ANY)"},
+                "geometry",
+            ),
+            (
+                "unqualified_geography_col",
+                {b"Spark:DataType:SqlName": b"GEOGRAPHY"},
+                "geography",
+            ),
             ("normal_col", {}, "string"),
             (
                 "weird_field",
                 {b"Spark:DataType:SqlName": b"Some unexpected value"},
                 "string",
             ),
+            (
+                "malformed_geometry_col",
+                {b"Spark:DataType:SqlName": b"GEOMETRY("},
+                "string",
+            ),
             ("missing_field", None, "string"),  # None field case
         ]
 
@@ -2575,12 +2600,27 @@ def test_hive_schema_to_description(self):
                 [
                     ("regular_col", ttypes.TTypeId.STRING_TYPE),
                     ("variant_col", ttypes.TTypeId.STRING_TYPE),
+                    ("geometry_col", ttypes.TTypeId.STRING_TYPE),
+                    ("geography_col", ttypes.TTypeId.STRING_TYPE),
                 ],
                 [
                     ("regular_col", {}),
                     ("variant_col", {b"Spark:DataType:SqlName": b"VARIANT"}),
+                    (
+                        "geometry_col",
+                        {b"Spark:DataType:SqlName": b"GEOMETRY(0)"},
+                    ),
+                    (
+                        "geography_col",
+                        {b"Spark:DataType:SqlName": b"GEOGRAPHY(4326)"},
+                    ),
+                ],
+                [
+                    ("regular_col", "string"),
+                    ("variant_col", "variant"),
+                    ("geometry_col", "geometry"),
+                    ("geography_col", "geography"),
                 ],
-                [("regular_col", "string"), ("variant_col", "variant")],
             ),
             (
                 [("regular_col", ttypes.TTypeId.STRING_TYPE)],