refactor

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

Author: varun-edachali-dbx

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

@@ -24,6 +24,15 @@
 logger = logging.getLogger(__name__)
 
 
+# Column-to-column data mapping for metadata queries
+# Maps target column -> source column to get data from
+COLUMN_DATA_MAPPING = {
+    "DATA_TYPE": "TYPE_NAME",  # DATA_TYPE calculated from TYPE_NAME
+    "NULLABLE": "IS_NULLABLE",  # NULLABLE calculated from IS_NULLABLE
+    "BUFFER_LENGTH": "TYPE_NAME",  # BUFFER_LENGTH calculated from TYPE_NAME
+}
+
+
 class SeaResultSet(ResultSet):
     """ResultSet implementation for SEA backend."""
 
@@ -292,6 +301,40 @@ def prepare_metadata_columns(self, metadata_columns: List[ResultColumn]) -> None
         self._metadata_columns = metadata_columns
         self._prepare_column_mapping()
 
+    def _populate_columns_from_others(
+        self, result_column: ResultColumn, row_data: Any
+    ) -> Any:
+        """
+        Helper function to populate column data from other columns based on COLUMN_DATA_MAPPING.
+
+        Args:
+            result_column: The result column that needs data
+            row_data: Row data (list for JSON, PyArrow table for Arrow)
+
+        Returns:
+            The value to use for this column, or None if not found
+        """
+        target_column = result_column.column_name
+        if target_column not in COLUMN_DATA_MAPPING:
+            return None
+
+        source_column = COLUMN_DATA_MAPPING[target_column]
+
+        # Find the source column index
+        for idx, col in enumerate(self._metadata_columns):
+            if col.column_name == source_column:
+                source_idx = self._column_index_mapping.get(idx)
+                if source_idx is not None:
+                    # Handle Arrow table format
+                    if hasattr(row_data, "column"):  # PyArrow table
+                        return row_data.column(source_idx).to_pylist()
+                    # Handle JSON row format
+                    else:
+                        return row_data[source_idx]
+                break
+
+        return None
+
     def _prepare_column_mapping(self) -> None:
         """
         Prepare column index mapping for metadata queries.
@@ -359,44 +402,24 @@ def _transform_arrow_table(self, table: "pyarrow.Table") -> "pyarrow.Table":
 
         for new_idx, result_column in enumerate(self._metadata_columns):
             old_idx = self._column_index_mapping.get(new_idx)
-            
-            # Get the source data 
+
+            # Get the source data
             if old_idx is not None:
                 column = table.column(old_idx)
                 values = column.to_pylist()
             else:
                 values = None
-                
+
             # Special handling for columns that need data from other columns
-            if result_column.column_name == "DATA_TYPE" and result_column.result_set_column_name is None:
-                # Get TYPE_NAME column value for DATA_TYPE calculation
-                for idx, col in enumerate(self._metadata_columns):
-                    if col.column_name == "TYPE_NAME":
-                        type_idx = self._column_index_mapping.get(idx)
-                        if type_idx is not None:
-                            values = table.column(type_idx).to_pylist()
-                        break
-            elif result_column.column_name == "NULLABLE" and result_column.result_set_column_name is None:
-                # Get IS_NULLABLE column value for NULLABLE calculation
-                for idx, col in enumerate(self._metadata_columns):
-                    if col.column_name == "IS_NULLABLE":
-                        nullable_idx = self._column_index_mapping.get(idx)
-                        if nullable_idx is not None:
-                            values = table.column(nullable_idx).to_pylist()
-                        break
-            elif result_column.column_name == "BUFFER_LENGTH" and result_column.result_set_column_name is None:
-                # Get TYPE_NAME column value for BUFFER_LENGTH calculation
-                for idx, col in enumerate(self._metadata_columns):
-                    if col.column_name == "TYPE_NAME":
-                        type_idx = self._column_index_mapping.get(idx)
-                        if type_idx is not None:
-                            values = table.column(type_idx).to_pylist()
-                        break
-            
+            if result_column.result_set_column_name is None:
+                values = self._populate_columns_from_others(result_column, table)
+
             # Apply transformation and create column
             if values is not None:
                 if result_column.transform_value:
-                    transformed_values = [result_column.transform_value(v) for v in values]
+                    transformed_values = [
+                        result_column.transform_value(v) for v in values
+                    ]
                     column = pyarrow.array(transformed_values)
                 else:
                     column = pyarrow.array(values)
@@ -409,7 +432,7 @@ def _transform_arrow_table(self, table: "pyarrow.Table") -> "pyarrow.Table":
                 else:
                     null_array = pyarrow.nulls(table.num_rows)
                 new_columns.append(null_array)
-                
+
             column_names.append(result_column.column_name)
 
         return pyarrow.Table.from_arrays(new_columns, names=column_names)
@@ -428,37 +451,15 @@ def _transform_json_rows(self, rows: List[List[str]]) -> List[List[Any]]:
                     value = row[old_idx]
                 else:
                     value = None
-                    
+
                 # Special handling for columns that need data from other columns
-                if result_column.column_name == "DATA_TYPE" and result_column.result_set_column_name is None:
-                    # Get TYPE_NAME column value for DATA_TYPE calculation
-                    for idx, col in enumerate(self._metadata_columns):
-                        if col.column_name == "TYPE_NAME":
-                            type_idx = self._column_index_mapping.get(idx)
-                            if type_idx is not None and type_idx < len(row):
-                                value = row[type_idx]
-                            break
-                elif result_column.column_name == "NULLABLE" and result_column.result_set_column_name is None:
-                    # Get IS_NULLABLE column value for NULLABLE calculation
-                    for idx, col in enumerate(self._metadata_columns):
-                        if col.column_name == "IS_NULLABLE":
-                            nullable_idx = self._column_index_mapping.get(idx)
-                            if nullable_idx is not None and nullable_idx < len(row):
-                                value = row[nullable_idx]
-                            break
-                elif result_column.column_name == "BUFFER_LENGTH" and result_column.result_set_column_name is None:
-                    # Get TYPE_NAME column value for BUFFER_LENGTH calculation
-                    for idx, col in enumerate(self._metadata_columns):
-                        if col.column_name == "TYPE_NAME":
-                            type_idx = self._column_index_mapping.get(idx)
-                            if type_idx is not None and type_idx < len(row):
-                                value = row[type_idx]
-                            break
-                
+                if result_column.result_set_column_name is None:
+                    value = self._populate_columns_from_others(result_column, row)
+
                 # Apply transformation if defined
                 if result_column.transform_value:
                     value = result_column.transform_value(value)
-                    
+
                 new_row.append(value)
             transformed_rows.append(new_row)
         return transformed_rows

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

@@ -150,7 +150,7 @@ def convert_value(
         Returns:
             The converted value in the appropriate Python type
         """
-        
+
         # Handle None values directly
         if value is None:
             return None

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

@@ -87,7 +87,7 @@ def _filter_sea_result_set(
         )
 
         # Preserve metadata columns setup from original result set
-        if hasattr(result_set, '_metadata_columns') and result_set._metadata_columns:
+        if hasattr(result_set, "_metadata_columns") and result_set._metadata_columns:
             filtered_result_set._metadata_columns = result_set._metadata_columns
             filtered_result_set._column_index_mapping = result_set._column_index_mapping
             # Update the description to match the original prepared description

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

@@ -11,7 +11,7 @@
     always_null,
     always_null_int,
     always_null_smallint,
-    identity
+    identity,
 )
 
 
@@ -20,66 +20,132 @@ class MetadataColumnMappings:
 
     # Common columns used across multiple metadata queries
     # FIX 1: Catalog columns - swap the mappings
-    CATALOG_COLUMN = ResultColumn("TABLE_CAT", "catalogName", "string", transform_value=identity)
-    CATALOG_COLUMN_FOR_GET_CATALOGS = ResultColumn("TABLE_CAT", "catalog", "string", transform_value=identity)
+    CATALOG_COLUMN = ResultColumn(
+        "TABLE_CAT", "catalogName", "string", transform_value=identity
+    )
+    CATALOG_COLUMN_FOR_GET_CATALOGS = ResultColumn(
+        "TABLE_CAT", "catalog", "string", transform_value=identity
+    )
     # Remove CATALOG_COLUMN_FOR_TABLES - will use CATALOG_COLUMN instead
-    
-    SCHEMA_COLUMN = ResultColumn("TABLE_SCHEM", "namespace", "string", transform_value=identity)
-    SCHEMA_COLUMN_FOR_GET_SCHEMA = ResultColumn("TABLE_SCHEM", "databaseName", "string", transform_value=identity)
-    TABLE_NAME_COLUMN = ResultColumn("TABLE_NAME", "tableName", "string", transform_value=identity)
-    TABLE_TYPE_COLUMN = ResultColumn("TABLE_TYPE", "tableType", "string", transform_value=transform_table_type)
-    REMARKS_COLUMN = ResultColumn("REMARKS", "remarks", "string", transform_value=transform_remarks_default)
+
+    SCHEMA_COLUMN = ResultColumn(
+        "TABLE_SCHEM", "namespace", "string", transform_value=identity
+    )
+    SCHEMA_COLUMN_FOR_GET_SCHEMA = ResultColumn(
+        "TABLE_SCHEM", "databaseName", "string", transform_value=identity
+    )
+    TABLE_NAME_COLUMN = ResultColumn(
+        "TABLE_NAME", "tableName", "string", transform_value=identity
+    )
+    TABLE_TYPE_COLUMN = ResultColumn(
+        "TABLE_TYPE", "tableType", "string", transform_value=transform_table_type
+    )
+    REMARKS_COLUMN = ResultColumn(
+        "REMARKS", "remarks", "string", transform_value=transform_remarks_default
+    )
 
     # Columns specific to getColumns()
-    COLUMN_NAME_COLUMN = ResultColumn("COLUMN_NAME", "col_name", "string", transform_value=identity)
+    COLUMN_NAME_COLUMN = ResultColumn(
+        "COLUMN_NAME", "col_name", "string", transform_value=identity
+    )
     DATA_TYPE_COLUMN = ResultColumn(
         "DATA_TYPE", None, "int", transform_value=calculate_data_type
     )  # Calculated from columnType
-    TYPE_NAME_COLUMN = ResultColumn("TYPE_NAME", "columnType", "string", transform_value=identity)
-    
+    TYPE_NAME_COLUMN = ResultColumn(
+        "TYPE_NAME", "columnType", "string", transform_value=identity
+    )
+
     # FIX 5: SEA actually provides these columns
-    COLUMN_SIZE_COLUMN = ResultColumn("COLUMN_SIZE", "columnSize", "int", transform_value=identity)
-    DECIMAL_DIGITS_COLUMN = ResultColumn("DECIMAL_DIGITS", "decimalDigits", "int", transform_value=transform_numeric_default_zero)
-    NUM_PREC_RADIX_COLUMN = ResultColumn("NUM_PREC_RADIX", "radix", "int", transform_value=transform_numeric_default_zero)
-    ORDINAL_POSITION_COLUMN = ResultColumn("ORDINAL_POSITION", "ordinalPosition", "int", transform_value=transform_ordinal_position_offset)
-    
-    NULLABLE_COLUMN = ResultColumn("NULLABLE", None, "int", transform_value=transform_nullable_to_int)  # Calculated from isNullable
+    COLUMN_SIZE_COLUMN = ResultColumn(
+        "COLUMN_SIZE", "columnSize", "int", transform_value=identity
+    )
+    DECIMAL_DIGITS_COLUMN = ResultColumn(
+        "DECIMAL_DIGITS",
+        "decimalDigits",
+        "int",
+        transform_value=transform_numeric_default_zero,
+    )
+    NUM_PREC_RADIX_COLUMN = ResultColumn(
+        "NUM_PREC_RADIX", "radix", "int", transform_value=transform_numeric_default_zero
+    )
+    ORDINAL_POSITION_COLUMN = ResultColumn(
+        "ORDINAL_POSITION",
+        "ordinalPosition",
+        "int",
+        transform_value=transform_ordinal_position_offset,
+    )
+
+    NULLABLE_COLUMN = ResultColumn(
+        "NULLABLE", None, "int", transform_value=transform_nullable_to_int
+    )  # Calculated from isNullable
     COLUMN_DEF_COLUMN = ResultColumn(
         "COLUMN_DEF", "columnType", "string", transform_value=identity
     )  # Note: duplicate mapping
-    SQL_DATA_TYPE_COLUMN = ResultColumn("SQL_DATA_TYPE", None, "int", transform_value=always_null_int)
-    SQL_DATETIME_SUB_COLUMN = ResultColumn("SQL_DATETIME_SUB", None, "int", transform_value=always_null_int)
-    CHAR_OCTET_LENGTH_COLUMN = ResultColumn("CHAR_OCTET_LENGTH", None, "int", transform_value=always_null_int)
-    IS_NULLABLE_COLUMN = ResultColumn("IS_NULLABLE", "isNullable", "string", transform_value=transform_is_nullable)
+    SQL_DATA_TYPE_COLUMN = ResultColumn(
+        "SQL_DATA_TYPE", None, "int", transform_value=always_null_int
+    )
+    SQL_DATETIME_SUB_COLUMN = ResultColumn(
+        "SQL_DATETIME_SUB", None, "int", transform_value=always_null_int
+    )
+    CHAR_OCTET_LENGTH_COLUMN = ResultColumn(
+        "CHAR_OCTET_LENGTH", None, "int", transform_value=always_null_int
+    )
+    IS_NULLABLE_COLUMN = ResultColumn(
+        "IS_NULLABLE", "isNullable", "string", transform_value=transform_is_nullable
+    )
 
     # Columns for getTables() that don't exist in SEA
-    TYPE_CAT_COLUMN = ResultColumn("TYPE_CAT", None, "string", transform_value=always_null)
-    TYPE_SCHEM_COLUMN = ResultColumn("TYPE_SCHEM", None, "string", transform_value=always_null)
-    TYPE_NAME_COLUMN = ResultColumn("TYPE_NAME", None, "string", transform_value=always_null)
+    TYPE_CAT_COLUMN = ResultColumn(
+        "TYPE_CAT", None, "string", transform_value=always_null
+    )
+    TYPE_SCHEM_COLUMN = ResultColumn(
+        "TYPE_SCHEM", None, "string", transform_value=always_null
+    )
+    TYPE_NAME_COLUMN = ResultColumn(
+        "TYPE_NAME", None, "string", transform_value=always_null
+    )
     SELF_REFERENCING_COL_NAME_COLUMN = ResultColumn(
         "SELF_REFERENCING_COL_NAME", None, "string", transform_value=always_null
     )
-    REF_GENERATION_COLUMN = ResultColumn("REF_GENERATION", None, "string", transform_value=always_null)
-    
+    REF_GENERATION_COLUMN = ResultColumn(
+        "REF_GENERATION", None, "string", transform_value=always_null
+    )
+
     # FIX 8: Scope columns (always null per JDBC)
-    SCOPE_CATALOG_COLUMN = ResultColumn("SCOPE_CATALOG", None, "string", transform_value=always_null)
-    SCOPE_SCHEMA_COLUMN = ResultColumn("SCOPE_SCHEMA", None, "string", transform_value=always_null)
-    SCOPE_TABLE_COLUMN = ResultColumn("SCOPE_TABLE", None, "string", transform_value=always_null)
-    SOURCE_DATA_TYPE_COLUMN = ResultColumn("SOURCE_DATA_TYPE", None, "smallint", transform_value=always_null_smallint)
-    
+    SCOPE_CATALOG_COLUMN = ResultColumn(
+        "SCOPE_CATALOG", None, "string", transform_value=always_null
+    )
+    SCOPE_SCHEMA_COLUMN = ResultColumn(
+        "SCOPE_SCHEMA", None, "string", transform_value=always_null
+    )
+    SCOPE_TABLE_COLUMN = ResultColumn(
+        "SCOPE_TABLE", None, "string", transform_value=always_null
+    )
+    SOURCE_DATA_TYPE_COLUMN = ResultColumn(
+        "SOURCE_DATA_TYPE", None, "smallint", transform_value=always_null_smallint
+    )
+
     # FIX 9 & 10: Auto increment and generated columns
-    IS_AUTO_INCREMENT_COLUMN = ResultColumn("IS_AUTOINCREMENT", "isAutoIncrement", "string", transform_value=identity)  # No underscore!
-    IS_GENERATED_COLUMN = ResultColumn("IS_GENERATEDCOLUMN", "isGenerated", "string", transform_value=identity)  # SEA provides this
-    
+    IS_AUTO_INCREMENT_COLUMN = ResultColumn(
+        "IS_AUTOINCREMENT", "isAutoIncrement", "string", transform_value=identity
+    )  # No underscore!
+    IS_GENERATED_COLUMN = ResultColumn(
+        "IS_GENERATEDCOLUMN", "isGenerated", "string", transform_value=identity
+    )  # SEA provides this
+
     # FIX 11: Buffer length column
-    BUFFER_LENGTH_COLUMN = ResultColumn("BUFFER_LENGTH", None, "int", transform_value=always_null_int)  # Always null per JDBC
+    BUFFER_LENGTH_COLUMN = ResultColumn(
+        "BUFFER_LENGTH", None, "int", transform_value=always_null_int
+    )  # Always null per JDBC
 
     # Column lists for each metadata operation
     CATALOG_COLUMNS = [CATALOG_COLUMN_FOR_GET_CATALOGS]  # Use specific catalog column
 
     SCHEMA_COLUMNS = [
         SCHEMA_COLUMN_FOR_GET_SCHEMA,
-        ResultColumn("TABLE_CATALOG", None, "string", transform_value=always_null),  # Will need special population logic
+        ResultColumn(
+            "TABLE_CATALOG", None, "string", transform_value=always_null
+        ),  # Will need special population logic
     ]
 
     TABLE_COLUMNS = [