remove hardcoding in SqlType

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/result_set.py

@@ -313,17 +313,17 @@ def _prepare_column_mapping(self) -> None:
         for new_idx, result_column in enumerate(self._metadata_columns or []):
             # Find the corresponding SEA column
             if (
-                result_column.result_set_column_name
-                and result_column.result_set_column_name in sea_column_indices
+                result_column.sea_col_name
+                and result_column.sea_col_name in sea_column_indices
             ):
-                old_idx = sea_column_indices[result_column.result_set_column_name]
+                old_idx = sea_column_indices[result_column.sea_col_name]
                 self._column_index_mapping[new_idx] = old_idx
                 # Use the original column metadata but with JDBC name
                 old_col = self.description[old_idx]
                 new_description.append(
                     (
-                        result_column.column_name,  # JDBC name
-                        result_column.column_type,  # Expected type
+                        result_column.thrift_col_name,  # JDBC name
+                        result_column.thrift_col_type,  # Expected type
                         old_col[2],  # display_size
                         old_col[3],  # internal_size
                         old_col[4],  # precision
@@ -335,8 +335,8 @@ def _prepare_column_mapping(self) -> None:
                 # Column doesn't exist in SEA - add with None values
                 new_description.append(
                     (
-                        result_column.column_name,
-                        result_column.column_type,
+                        result_column.thrift_col_name,
+                        result_column.thrift_col_type,
                         None,
                         None,
                         None,
@@ -380,7 +380,7 @@ def _normalise_arrow_metadata_cols(self, table: "pyarrow.Table") -> "pyarrow.Tab
                 null_array = pyarrow.nulls(table.num_rows)
                 new_columns.append(null_array)
 
-            column_names.append(result_column.column_name)
+            column_names.append(result_column.thrift_col_name)
 
         return pyarrow.Table.from_arrays(new_columns, names=column_names)
 

src/databricks/sql/backend/sea/utils/conversion.py

@@ -11,6 +11,8 @@
 from dateutil import parser
 from typing import Callable, Dict, Optional
 
+from databricks.sql.thrift_api.TCLIService import ttypes
+
 logger = logging.getLogger(__name__)
 
 
@@ -56,43 +58,49 @@ class SqlType:
     after normalize_sea_type_to_thrift processing (lowercase, without _TYPE suffix).
     """
 
+    @staticmethod
+    def _get_type_name(thrift_type_id: int) -> str:
+        type_name = ttypes.TTypeId._VALUES_TO_NAMES[thrift_type_id]
+        type_name = type_name.lower()
+        if type_name.endswith("_type"):
+            type_name = type_name[:-5]
+        return type_name
+
     # Numeric types
-    TINYINT = "tinyint"  # Maps to TTypeId.TINYINT_TYPE
-    SMALLINT = "smallint"  # Maps to TTypeId.SMALLINT_TYPE
-    INT = "int"  # Maps to TTypeId.INT_TYPE
-    BIGINT = "bigint"  # Maps to TTypeId.BIGINT_TYPE
-    FLOAT = "float"  # Maps to TTypeId.FLOAT_TYPE
-    DOUBLE = "double"  # Maps to TTypeId.DOUBLE_TYPE
-    DECIMAL = "decimal"  # Maps to TTypeId.DECIMAL_TYPE
+    TINYINT = _get_type_name(ttypes.TTypeId.TINYINT_TYPE)
+    SMALLINT = _get_type_name(ttypes.TTypeId.SMALLINT_TYPE)
+    INT = _get_type_name(ttypes.TTypeId.INT_TYPE)
+    BIGINT = _get_type_name(ttypes.TTypeId.BIGINT_TYPE)
+    FLOAT = _get_type_name(ttypes.TTypeId.FLOAT_TYPE)
+    DOUBLE = _get_type_name(ttypes.TTypeId.DOUBLE_TYPE)
+    DECIMAL = _get_type_name(ttypes.TTypeId.DECIMAL_TYPE)
 
     # Boolean type
-    BOOLEAN = "boolean"  # Maps to TTypeId.BOOLEAN_TYPE
+    BOOLEAN = _get_type_name(ttypes.TTypeId.BOOLEAN_TYPE)
 
     # Date/Time types
-    DATE = "date"  # Maps to TTypeId.DATE_TYPE
-    TIMESTAMP = "timestamp"  # Maps to TTypeId.TIMESTAMP_TYPE
-    INTERVAL_YEAR_MONTH = (
-        "interval_year_month"  # Maps to TTypeId.INTERVAL_YEAR_MONTH_TYPE
-    )
-    INTERVAL_DAY_TIME = "interval_day_time"  # Maps to TTypeId.INTERVAL_DAY_TIME_TYPE
+    DATE = _get_type_name(ttypes.TTypeId.DATE_TYPE)
+    TIMESTAMP = _get_type_name(ttypes.TTypeId.TIMESTAMP_TYPE)
+    INTERVAL_YEAR_MONTH = _get_type_name(ttypes.TTypeId.INTERVAL_YEAR_MONTH_TYPE)
+    INTERVAL_DAY_TIME = _get_type_name(ttypes.TTypeId.INTERVAL_DAY_TIME_TYPE)
 
     # String types
-    CHAR = "char"  # Maps to TTypeId.CHAR_TYPE
-    VARCHAR = "varchar"  # Maps to TTypeId.VARCHAR_TYPE
-    STRING = "string"  # Maps to TTypeId.STRING_TYPE
+    CHAR = _get_type_name(ttypes.TTypeId.CHAR_TYPE)
+    VARCHAR = _get_type_name(ttypes.TTypeId.VARCHAR_TYPE)
+    STRING = _get_type_name(ttypes.TTypeId.STRING_TYPE)
 
     # Binary type
-    BINARY = "binary"  # Maps to TTypeId.BINARY_TYPE
+    BINARY = _get_type_name(ttypes.TTypeId.BINARY_TYPE)
 
     # Complex types
-    ARRAY = "array"  # Maps to TTypeId.ARRAY_TYPE
-    MAP = "map"  # Maps to TTypeId.MAP_TYPE
-    STRUCT = "struct"  # Maps to TTypeId.STRUCT_TYPE
+    ARRAY = _get_type_name(ttypes.TTypeId.ARRAY_TYPE)
+    MAP = _get_type_name(ttypes.TTypeId.MAP_TYPE)
+    STRUCT = _get_type_name(ttypes.TTypeId.STRUCT_TYPE)
 
     # Other types
-    NULL = "null"  # Maps to TTypeId.NULL_TYPE
-    UNION = "union"  # Maps to TTypeId.UNION_TYPE
-    USER_DEFINED = "user_defined"  # Maps to TTypeId.USER_DEFINED_TYPE
+    NULL = _get_type_name(ttypes.TTypeId.NULL_TYPE)
+    UNION = _get_type_name(ttypes.TTypeId.UNION_TYPE)
+    USER_DEFINED = _get_type_name(ttypes.TTypeId.USER_DEFINED_TYPE)
 
 
 class SqlTypeConverter:

src/databricks/sql/backend/sea/utils/result_column.py

@@ -8,12 +8,11 @@ class ResultColumn:
     Represents a mapping between JDBC specification column names and actual result set column names.
 
     Attributes:
-        column_name: JDBC specification column name (e.g., "TABLE_CAT")
-        result_set_column_name: Server result column name from SEA (e.g., "catalog")
-        column_type: SQL type code from databricks.sql.types
-        transform_value: Optional function to transform values for this column
+        thrift_col_name: Column name as returned by Thrift (e.g., "TABLE_CAT")
+        sea_col_name: Server result column name from SEA (e.g., "catalog")
+        thrift_col_type: SQL type name
     """
 
-    column_name: str
-    result_set_column_name: Optional[str]  # None if SEA doesn't return this column
-    column_type: str
+    thrift_col_name: str
+    sea_col_name: Optional[str]  # None if SEA doesn't return this column
+    thrift_col_type: str