normalise type codes

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

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

@@ -696,10 +696,10 @@ def get_catalogs(
         assert result is not None, "execute_command returned None in synchronous mode"
 
         # Normalize column names to match JDBC/thrift backend
-        from .metadata_constants import CATALOG_COLUMNS, normalize_metadata_description
+        from .metadata_constants import CATALOG_COLUMNS, CATALOG_TYPE_CODES, normalize_metadata_description
 
         result.description = normalize_metadata_description(
-            result.description, CATALOG_COLUMNS
+            result.description, CATALOG_COLUMNS, CATALOG_TYPE_CODES
         )
 
         return result
@@ -737,10 +737,10 @@ def get_schemas(
         assert result is not None, "execute_command returned None in synchronous mode"
 
         # Normalize column names to match JDBC/thrift backend
-        from .metadata_constants import SCHEMA_COLUMNS, normalize_metadata_description
+        from .metadata_constants import SCHEMA_COLUMNS, SCHEMA_TYPE_CODES, normalize_metadata_description
 
         result.description = normalize_metadata_description(
-            result.description, SCHEMA_COLUMNS
+            result.description, SCHEMA_COLUMNS, SCHEMA_TYPE_CODES
         )
 
         return result
@@ -786,10 +786,10 @@ def get_tables(
         assert result is not None, "execute_command returned None in synchronous mode"
 
         # Normalize column names to match JDBC/thrift backend
-        from .metadata_constants import TABLE_COLUMNS, normalize_metadata_description
+        from .metadata_constants import TABLE_COLUMNS, TABLE_TYPE_CODES, normalize_metadata_description
 
         result.description = normalize_metadata_description(
-            result.description, TABLE_COLUMNS
+            result.description, TABLE_COLUMNS, TABLE_TYPE_CODES
         )
 
         # Apply client-side filtering by table_types

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

@@ -5,7 +5,7 @@
 to match JDBC DatabaseMetaData standards and thrift backend behavior.
 """
 
-from typing import List, Tuple, Dict, Any
+from typing import List, Tuple, Dict, Any, Optional
 
 
 # Columns for catalogs() - matching JDBC CATALOG_COLUMNS exactly
@@ -70,6 +70,51 @@
 # Note: COLUMN_DEF and TYPE_NAME both map to "columnType" - no special handling needed
 # since they both reference the same source column in the data
 
+# Type code overrides for metadata columns to match thrift backend behavior
+# Format: column_name -> expected_type_code
+CATALOG_TYPE_CODES: Dict[str, str] = {
+    # All columns in catalogs() are strings
+}
+
+SCHEMA_TYPE_CODES: Dict[str, str] = {
+    # All columns in schemas() are strings
+}
+
+TABLE_TYPE_CODES: Dict[str, str] = {
+    # All columns in tables() are strings
+}
+
+COLUMN_TYPE_CODES: Dict[str, str] = {
+    # Integer types
+    "DATA_TYPE": "int",
+    "COLUMN_SIZE": "int",
+    "DECIMAL_DIGITS": "int",
+    "NUM_PREC_RADIX": "int",
+    "NULLABLE": "int",
+    "SQL_DATA_TYPE": "int",
+    "SQL_DATETIME_SUB": "int",
+    "CHAR_OCTET_LENGTH": "int",
+    "ORDINAL_POSITION": "int",
+    # Small integer types
+    "SOURCE_DATA_TYPE": "smallint",
+    # Tiny integer types
+    "BUFFER_LENGTH": "tinyint",
+    # String types (explicitly listed for clarity, though they don't need override)
+    "TABLE_CAT": "string",
+    "TABLE_SCHEM": "string",
+    "TABLE_NAME": "string",
+    "COLUMN_NAME": "string",
+    "TYPE_NAME": "string",
+    "REMARKS": "string",
+    "COLUMN_DEF": "string",
+    "IS_NULLABLE": "string",
+    "SCOPE_CATALOG": "string",
+    "SCOPE_SCHEMA": "string",
+    "SCOPE_TABLE": "string",
+    "IS_AUTOINCREMENT": "string",
+    "IS_GENERATEDCOLUMN": "string",
+}
+
 
 # Helper functions to work with column definitions
 def get_column_names(columns: List[Tuple[str, str]]) -> List[str]:
@@ -85,19 +130,22 @@ def get_column_mapping(columns: List[Tuple[str, str]]) -> Dict[str, str]:
 
 
 def normalize_metadata_description(
-    original_description: List[Tuple], column_definitions: List[Tuple[str, str]]
+    original_description: List[Tuple], 
+    column_definitions: List[Tuple[str, str]],
+    type_code_overrides: Optional[Dict[str, str]] = None
 ) -> List[Tuple]:
     """
-    Transform result set description to use JDBC-standard column names.
+    Transform result set description to use JDBC-standard column names and type codes.
 
     Args:
         original_description: Original PEP-249 description from SEA backend
                             Format: [(name, type_code, display_size, internal_size,
                                      precision, scale, null_ok), ...]
         column_definitions: List of (jdbc_name, sea_source_name) tuples defining mappings
+        type_code_overrides: Optional dict of column_name -> type_code overrides
 
     Returns:
-        Normalized description with JDBC column names
+        Normalized description with JDBC column names and corrected type codes
     """
     if not original_description:
         return original_description
@@ -115,14 +163,29 @@ def normalize_metadata_description(
             # Column exists in original description
             orig_idx = sea_col_to_idx[sea_name]
             orig_desc = original_description[orig_idx]
-            # Replace the column name, keep other metadata
-            new_desc = (jdbc_name,) + orig_desc[1:]
+            
+            # Check if we need to override the type code
+            if type_code_overrides and jdbc_name in type_code_overrides:
+                # Override the type code (second element)
+                new_desc = (
+                    jdbc_name,
+                    type_code_overrides[jdbc_name],
+                ) + orig_desc[2:]
+            else:
+                # Replace the column name, keep other metadata
+                new_desc = (jdbc_name,) + orig_desc[1:]
+            
             normalized_description.append(new_desc)
         else:
             # Column doesn't exist, add with default metadata
-            # Use VARCHAR type and nullable=None as defaults
+            # Check if there's a type override for this column
+            if type_code_overrides and jdbc_name in type_code_overrides:
+                type_code = type_code_overrides[jdbc_name]
+            else:
+                type_code = "string"  # Default type
+            
             normalized_description.append(
-                (jdbc_name, "string", None, None, None, None, None)
+                (jdbc_name, type_code, None, None, None, None, None)
             )
 
     return normalized_description
@@ -132,13 +195,15 @@ def normalize_columns_metadata_description(
     original_description: List[Tuple],
 ) -> List[Tuple]:
     """
-    Normalization for columns() metadata.
+    Normalization for columns() metadata with type code overrides.
 
     Args:
         original_description: Original description from SEA backend
 
     Returns:
-        Normalized description matching JDBC COLUMN_COLUMNS
+        Normalized description matching JDBC COLUMN_COLUMNS with correct type codes
     """
     # COLUMN_DEF and TYPE_NAME both map to "columnType" so no special handling needed
-    return normalize_metadata_description(original_description, COLUMN_COLUMNS)
+    return normalize_metadata_description(
+        original_description, COLUMN_COLUMNS, COLUMN_TYPE_CODES
+    )