init col normalisation

varun-edachali-dbx · varun-edachali-dbx · commit ad52ed4fc20c · 2025-07-25T11:55:50.000Z
Signed-off-by: varun-edachali-dbx &lt;varun.edachali@databricks.com&gt;
diff --git a/src/databricks/sql/backend/sea/backend.py b/src/databricks/sql/backend/sea/backend.py
@@ -694,6 +694,14 @@ def get_catalogs(
             enforce_embedded_schema_correctness=False,
         )
         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
+
+        result.description = normalize_metadata_description(
+            result.description, CATALOG_COLUMNS
+        )
+
         return result
 
     def get_schemas(
@@ -727,6 +735,14 @@ def get_schemas(
             enforce_embedded_schema_correctness=False,
         )
         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
+
+        result.description = normalize_metadata_description(
+            result.description, SCHEMA_COLUMNS
+        )
+
         return result
 
     def get_tables(
@@ -769,6 +785,13 @@ 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
+
+        result.description = normalize_metadata_description(
+            result.description, TABLE_COLUMNS
+        )
+
         # Apply client-side filtering by table_types
         from databricks.sql.backend.sea.utils.filters import ResultSetFilter
 
@@ -815,4 +838,10 @@ def get_columns(
             enforce_embedded_schema_correctness=False,
         )
         assert result is not None, "execute_command returned None in synchronous mode"
+
+        # Normalize column names to match JDBC/thrift backend
+        from .metadata_constants import normalize_columns_metadata_description
+
+        result.description = normalize_columns_metadata_description(result.description)
+
         return result
diff --git a/src/databricks/sql/backend/sea/metadata_constants.py b/src/databricks/sql/backend/sea/metadata_constants.py
@@ -0,0 +1,144 @@
+"""
+Column normalization constants for SEA backend metadata queries.
+
+This module defines column mappings to normalize SEA backend metadata results 
+to match JDBC DatabaseMetaData standards and thrift backend behavior.
+"""
+
+from typing import List, Tuple, Dict, Any
+
+
+# Columns for catalogs() - matching JDBC CATALOG_COLUMNS exactly
+CATALOG_COLUMNS: List[Tuple[str, str]] = [
+    ("TABLE_CAT", "catalog"),  # CATALOG_COLUMN_FOR_GET_CATALOGS
+]
+
+# Columns for schemas() - matching JDBC SCHEMA_COLUMNS exactly
+SCHEMA_COLUMNS: List[Tuple[str, str]] = [
+    ("TABLE_SCHEM", "databaseName"),  # SCHEMA_COLUMN_FOR_GET_SCHEMA
+    ("TABLE_CATALOG", "catalogName"),  # CATALOG_FULL_COLUMN
+]
+
+# Columns for tables() - matching JDBC TABLE_COLUMNS exactly
+TABLE_COLUMNS: List[Tuple[str, str]] = [
+    ("TABLE_CAT", "catalogName"),  # CATALOG_COLUMN
+    ("TABLE_SCHEM", "namespace"),  # SCHEMA_COLUMN
+    ("TABLE_NAME", "tableName"),  # TABLE_NAME_COLUMN
+    ("TABLE_TYPE", "tableType"),  # TABLE_TYPE_COLUMN
+    ("REMARKS", "remarks"),  # REMARKS_COLUMN
+    ("TYPE_CAT", "TYPE_CATALOG_COLUMN"),  # TYPE_CATALOG_COLUMN (likely None in data)
+    ("TYPE_SCHEM", "TYPE_SCHEMA_COLUMN"),  # TYPE_SCHEMA_COLUMN (likely None in data)
+    ("TYPE_NAME", "TYPE_NAME"),  # TYPE_NAME_COLUMN (likely None in data)
+    (
+        "SELF_REFERENCING_COL_NAME",
+        "SELF_REFERENCING_COLUMN_NAME",
+    ),  # (likely None in data)
+    (
+        "REF_GENERATION",
+        "REF_GENERATION_COLUMN",
+    ),  # REF_GENERATION_COLUMN (likely None in data)
+]
+
+# Columns for columns() - matching JDBC COLUMN_COLUMNS exactly
+COLUMN_COLUMNS: List[Tuple[str, str]] = [
+    ("TABLE_CAT", "catalogName"),  # CATALOG_COLUMN
+    ("TABLE_SCHEM", "namespace"),  # SCHEMA_COLUMN
+    ("TABLE_NAME", "tableName"),  # TABLE_NAME_COLUMN
+    ("COLUMN_NAME", "col_name"),  # COL_NAME_COLUMN
+    ("DATA_TYPE", "dataType"),  # DATA_TYPE_COLUMN
+    ("TYPE_NAME", "columnType"),  # COLUMN_TYPE_COLUMN
+    ("COLUMN_SIZE", "columnSize"),  # COLUMN_SIZE_COLUMN
+    ("BUFFER_LENGTH", "bufferLength"),  # BUFFER_LENGTH_COLUMN
+    ("DECIMAL_DIGITS", "decimalDigits"),  # DECIMAL_DIGITS_COLUMN
+    ("NUM_PREC_RADIX", "radix"),  # NUM_PREC_RADIX_COLUMN
+    ("NULLABLE", "Nullable"),  # NULLABLE_COLUMN
+    ("REMARKS", "remarks"),  # REMARKS_COLUMN
+    ("COLUMN_DEF", "columnType"),  # COLUMN_DEF_COLUMN (same source as TYPE_NAME)
+    ("SQL_DATA_TYPE", "SQLDataType"),  # SQL_DATA_TYPE_COLUMN
+    ("SQL_DATETIME_SUB", "SQLDateTimeSub"),  # SQL_DATETIME_SUB_COLUMN
+    ("CHAR_OCTET_LENGTH", "CharOctetLength"),  # CHAR_OCTET_LENGTH_COLUMN
+    ("ORDINAL_POSITION", "ordinalPosition"),  # ORDINAL_POSITION_COLUMN
+    ("IS_NULLABLE", "isNullable"),  # IS_NULLABLE_COLUMN
+    ("SCOPE_CATALOG", "ScopeCatalog"),  # SCOPE_CATALOG_COLUMN
+    ("SCOPE_SCHEMA", "ScopeSchema"),  # SCOPE_SCHEMA_COLUMN
+    ("SCOPE_TABLE", "ScopeTable"),  # SCOPE_TABLE_COLUMN
+    ("SOURCE_DATA_TYPE", "SourceDataType"),  # SOURCE_DATA_TYPE_COLUMN
+    ("IS_AUTOINCREMENT", "isAutoIncrement"),  # IS_AUTO_INCREMENT_COLUMN
+    ("IS_GENERATEDCOLUMN", "isGenerated"),  # IS_GENERATED_COLUMN
+]
+
+# 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
+
+
+# Helper functions to work with column definitions
+def get_column_names(columns: List[Tuple[str, str]]) -> List[str]:
+    """Extract JDBC column names from column definitions."""
+    return [jdbc_name for jdbc_name, _ in columns]
+
+
+def get_column_mapping(columns: List[Tuple[str, str]]) -> Dict[str, str]:
+    """Create mapping dict from SEA names to JDBC names."""
+    return {
+        sea_name: jdbc_name for jdbc_name, sea_name in columns if sea_name is not None
+    }
+
+
+def normalize_metadata_description(
+    original_description: List[Tuple], column_definitions: List[Tuple[str, str]]
+) -> List[Tuple]:
+    """
+    Transform result set description to use JDBC-standard column names.
+
+    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
+
+    Returns:
+        Normalized description with JDBC column names
+    """
+    if not original_description:
+        return original_description
+
+    # Build mapping from SEA column names to their indices
+    sea_col_to_idx = {}
+    for idx, col_desc in enumerate(original_description):
+        sea_col_to_idx[col_desc[0]] = idx
+
+    # Build new description based on column definitions
+    normalized_description = []
+
+    for jdbc_name, sea_name in column_definitions:
+        if sea_name and sea_name in sea_col_to_idx:
+            # 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:]
+            normalized_description.append(new_desc)
+        else:
+            # Column doesn't exist, add with default metadata
+            # Use VARCHAR type and nullable=None as defaults
+            normalized_description.append(
+                (jdbc_name, "string", None, None, None, None, None)
+            )
+
+    return normalized_description
+
+
+def normalize_columns_metadata_description(
+    original_description: List[Tuple],
+) -> List[Tuple]:
+    """
+    Normalization for columns() metadata.
+
+    Args:
+        original_description: Original description from SEA backend
+
+    Returns:
+        Normalized description matching JDBC COLUMN_COLUMNS
+    """
+    # COLUMN_DEF and TYPE_NAME both map to "columnType" so no special handling needed
+    return normalize_metadata_description(original_description, COLUMN_COLUMNS)
diff --git a/tests/unit/test_metadata_constants.py b/tests/unit/test_metadata_constants.py
@@ -0,0 +1,180 @@
+"""
+Unit tests for metadata column normalization constants and functions.
+"""
+
+import unittest
+from databricks.sql.backend.sea.metadata_constants import (
+    CATALOG_COLUMNS,
+    SCHEMA_COLUMNS,
+    TABLE_COLUMNS,
+    COLUMN_COLUMNS,
+    get_column_names,
+    get_column_mapping,
+    normalize_metadata_description,
+    normalize_columns_metadata_description,
+)
+
+
+class TestMetadataConstants(unittest.TestCase):
+    """Test metadata column constants and helper functions."""
+
+    def test_catalog_columns_structure(self):
+        """Test CATALOG_COLUMNS has correct structure."""
+        self.assertEqual(len(CATALOG_COLUMNS), 1)
+        self.assertEqual(CATALOG_COLUMNS[0], ("TABLE_CAT", "catalog"))
+
+    def test_schema_columns_structure(self):
+        """Test SCHEMA_COLUMNS has correct structure."""
+        self.assertEqual(len(SCHEMA_COLUMNS), 2)
+        self.assertEqual(SCHEMA_COLUMNS[0], ("TABLE_SCHEM", "databaseName"))
+        self.assertEqual(SCHEMA_COLUMNS[1], ("TABLE_CATALOG", "catalogName"))
+
+    def test_table_columns_structure(self):
+        """Test TABLE_COLUMNS has correct structure and count."""
+        self.assertEqual(len(TABLE_COLUMNS), 10)
+        # Check key columns
+        self.assertEqual(TABLE_COLUMNS[0], ("TABLE_CAT", "catalogName"))
+        self.assertEqual(TABLE_COLUMNS[1], ("TABLE_SCHEM", "namespace"))
+        self.assertEqual(TABLE_COLUMNS[2], ("TABLE_NAME", "tableName"))
+        self.assertEqual(TABLE_COLUMNS[3], ("TABLE_TYPE", "tableType"))
+        self.assertEqual(TABLE_COLUMNS[4], ("REMARKS", "remarks"))
+
+    def test_column_columns_structure(self):
+        """Test COLUMN_COLUMNS has correct structure and count."""
+        self.assertEqual(len(COLUMN_COLUMNS), 24)
+        # Check key columns
+        self.assertEqual(COLUMN_COLUMNS[0], ("TABLE_CAT", "catalogName"))
+        self.assertEqual(COLUMN_COLUMNS[1], ("TABLE_SCHEM", "namespace"))
+        self.assertEqual(COLUMN_COLUMNS[2], ("TABLE_NAME", "tableName"))
+        self.assertEqual(COLUMN_COLUMNS[3], ("COLUMN_NAME", "col_name"))
+        self.assertEqual(COLUMN_COLUMNS[4], ("DATA_TYPE", "dataType"))
+        self.assertEqual(COLUMN_COLUMNS[5], ("TYPE_NAME", "columnType"))
+        # Check that COLUMN_DEF also maps to columnType (same source)
+        self.assertEqual(COLUMN_COLUMNS[12], ("COLUMN_DEF", "columnType"))
+
+    def test_get_column_names(self):
+        """Test get_column_names helper function."""
+        test_columns = [("JDBC_NAME1", "sea_name1"), ("JDBC_NAME2", "sea_name2")]
+        result = get_column_names(test_columns)
+        self.assertEqual(result, ["JDBC_NAME1", "JDBC_NAME2"])
+
+    def test_get_column_mapping(self):
+        """Test get_column_mapping helper function."""
+        test_columns = [
+            ("JDBC_NAME1", "sea_name1"),
+            ("JDBC_NAME2", "sea_name2"),
+            ("JDBC_NAME3", None),  # Should be excluded
+        ]
+        result = get_column_mapping(test_columns)
+        expected = {"sea_name1": "JDBC_NAME1", "sea_name2": "JDBC_NAME2"}
+        self.assertEqual(result, expected)
+
+    def test_normalize_metadata_description_basic(self):
+        """Test basic metadata description normalization."""
+        # Mock original description
+        original_desc = [
+            ("catalog", "string", None, None, None, None, True),
+        ]
+
+        result = normalize_metadata_description(original_desc, CATALOG_COLUMNS)
+
+        expected = [
+            ("TABLE_CAT", "string", None, None, None, None, True),
+        ]
+        self.assertEqual(result, expected)
+
+    def test_normalize_metadata_description_with_missing_columns(self):
+        """Test normalization when some columns are missing from source."""
+        # Original description has only one column
+        original_desc = [
+            ("databaseName", "string", None, None, None, None, True),
+        ]
+
+        result = normalize_metadata_description(original_desc, SCHEMA_COLUMNS)
+
+        expected = [
+            ("TABLE_SCHEM", "string", None, None, None, None, True),
+            (
+                "TABLE_CATALOG",
+                "string",
+                None,
+                None,
+                None,
+                None,
+                None,
+            ),  # Missing column gets defaults
+        ]
+        self.assertEqual(result, expected)
+
+    def test_normalize_metadata_description_empty_input(self):
+        """Test normalization with empty input."""
+        result = normalize_metadata_description([], CATALOG_COLUMNS)
+        self.assertEqual(result, [])
+
+    def test_normalize_columns_metadata_description(self):
+        """Test columns-specific normalization function."""
+        # Mock original description with key columns
+        original_desc = [
+            ("catalogName", "string", None, None, None, None, True),
+            ("namespace", "string", None, None, None, None, True),
+            ("tableName", "string", None, None, None, None, True),
+            ("col_name", "string", None, None, None, None, True),
+            ("dataType", "int", None, None, None, None, True),
+            ("columnType", "string", None, None, None, None, True),
+        ]
+
+        result = normalize_columns_metadata_description(original_desc)
+
+        # Should have 24 columns total
+        self.assertEqual(len(result), 24)
+
+        # Check that key columns are mapped correctly
+        self.assertEqual(result[0][0], "TABLE_CAT")  # catalogName -> TABLE_CAT
+        self.assertEqual(result[1][0], "TABLE_SCHEM")  # namespace -> TABLE_SCHEM
+        self.assertEqual(result[5][0], "TYPE_NAME")  # columnType -> TYPE_NAME
+        self.assertEqual(
+            result[12][0], "COLUMN_DEF"
+        )  # columnType -> COLUMN_DEF (same source)
+
+        # Both TYPE_NAME and COLUMN_DEF should have same metadata (except name)
+        self.assertEqual(result[5][1:], result[12][1:])
+
+    def test_normalize_metadata_description_preserves_metadata(self):
+        """Test that normalization preserves non-name metadata."""
+        original_desc = [
+            ("catalog", "varchar", 100, 50, 10, 2, False),
+        ]
+
+        result = normalize_metadata_description(original_desc, CATALOG_COLUMNS)
+
+        expected = [
+            ("TABLE_CAT", "varchar", 100, 50, 10, 2, False),
+        ]
+        self.assertEqual(result, expected)
+
+    def test_columns_with_duplicate_source_mapping(self):
+        """Test that TYPE_NAME and COLUMN_DEF both map to columnType correctly."""
+        original_desc = [
+            ("columnType", "string", None, None, None, None, True),
+        ]
+
+        # Create a subset of column definitions that includes both TYPE_NAME and COLUMN_DEF
+        test_columns = [
+            ("TYPE_NAME", "columnType"),
+            ("COLUMN_DEF", "columnType"),
+        ]
+
+        result = normalize_metadata_description(original_desc, test_columns)
+
+        expected = [
+            ("TYPE_NAME", "string", None, None, None, None, True),
+            ("COLUMN_DEF", "string", None, None, None, None, True),
+        ]
+        self.assertEqual(result, expected)
+
+        # Both should have identical metadata except for the name
+        self.assertEqual(result[0][1:], result[1][1:])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit/test_sea_backend.py b/tests/unit/test_sea_backend.py
