run e2e tests for metadata queries, run transform regardless of None value

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

Author: varun-edachali-dbx

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

@@ -20,7 +20,12 @@
     MetadataCommands,
 )
 from databricks.sql.backend.sea.utils.metadata_mappings import MetadataColumnMappings
+from databricks.sql.backend.sea.utils.metadata_transforms import (
+    create_table_catalog_transform,
+)
 from databricks.sql.backend.sea.utils.normalize import normalize_sea_type_to_thrift
+from databricks.sql.backend.sea.utils.result_column import ResultColumn
+from databricks.sql.backend.sea.utils.conversion import SqlType
 from databricks.sql.thrift_api.TCLIService import ttypes
 
 if TYPE_CHECKING:
@@ -740,7 +745,23 @@ def get_schemas(
         assert isinstance(
             result, SeaResultSet
         ), "Expected SeaResultSet from SEA backend"
-        result.prepare_metadata_columns(MetadataColumnMappings.SCHEMA_COLUMNS)
+
+        # Create dynamic schema columns with catalog name bound to TABLE_CATALOG
+        schema_columns = []
+        for col in MetadataColumnMappings.SCHEMA_COLUMNS:
+            if col.thrift_col_name == "TABLE_CATALOG":
+                # Create a new column with the catalog transform bound
+                dynamic_col = ResultColumn(
+                    col.thrift_col_name,
+                    col.sea_col_name,
+                    col.thrift_col_type,
+                    create_table_catalog_transform(catalog_name),
+                )
+                schema_columns.append(dynamic_col)
+            else:
+                schema_columns.append(col)
+
+        result.prepare_metadata_columns(schema_columns)
         return result
 
     def get_tables(

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

@@ -361,14 +361,13 @@ def _normalise_arrow_metadata_cols(self, table: "pyarrow.Table") -> "pyarrow.Tab
                 column = pyarrow.nulls(table.num_rows)
             else:
                 column = table.column(old_idx)
-                # Apply transform if available
-                if result_column.transform_value:
-                    # Convert to list, apply transform, and convert back
-                    values = column.to_pylist()
-                    transformed_values = [
-                        result_column.transform_value(v) for v in values
-                    ]
-                    column = pyarrow.array(transformed_values)
+
+            # Apply transform if available
+            if result_column.transform_value:
+                # Convert to list, apply transform, and convert back
+                values = column.to_pylist()
+                transformed_values = [result_column.transform_value(v) for v in values]
+                column = pyarrow.array(transformed_values)
 
             new_columns.append(column)
             column_names.append(result_column.thrift_col_name)
@@ -392,7 +391,7 @@ def _normalise_json_metadata_cols(self, rows: List[List[str]]) -> List[List[Any]
 
                 value = None if old_idx is None else row[old_idx]
                 # Apply transform if available
-                if value is not None and result_column.transform_value:
+                if result_column.transform_value:
                     value = result_column.transform_value(value)
                 new_row.append(value)
             transformed_rows.append(new_row)

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

@@ -26,7 +26,9 @@ class MetadataColumnMappings:
     SCHEMA_COLUMN = ResultColumn("TABLE_SCHEM", "namespace", SqlType.STRING)
     TABLE_NAME_COLUMN = ResultColumn("TABLE_NAME", "tableName", SqlType.STRING)
     TABLE_TYPE_COLUMN = ResultColumn("TABLE_TYPE", "tableType", SqlType.STRING)
-    REMARKS_COLUMN = ResultColumn("REMARKS", "remarks", SqlType.STRING, transform_remarks)
+    REMARKS_COLUMN = ResultColumn(
+        "REMARKS", "remarks", SqlType.STRING, transform_remarks
+    )
     TYPE_CATALOG_COLUMN = ResultColumn("TYPE_CAT", None, SqlType.STRING)
     TYPE_SCHEM_COLUMN = ResultColumn("TYPE_SCHEM", None, SqlType.STRING)
     TYPE_NAME_COLUMN = ResultColumn("TYPE_NAME", None, SqlType.STRING)
@@ -54,7 +56,9 @@ class MetadataColumnMappings:
         transform_ordinal_position,
     )
 
-    NULLABLE_COLUMN = ResultColumn("NULLABLE", "isNullable", SqlType.INT, transform_nullable)
+    NULLABLE_COLUMN = ResultColumn(
+        "NULLABLE", "isNullable", SqlType.INT, transform_nullable
+    )
     COLUMN_DEF_COLUMN = ResultColumn("COLUMN_DEF", None, SqlType.STRING)
     SQL_DATA_TYPE_COLUMN = ResultColumn("SQL_DATA_TYPE", None, SqlType.INT)
     SQL_DATETIME_SUB_COLUMN = ResultColumn("SQL_DATETIME_SUB", None, SqlType.INT)

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

@@ -16,10 +16,12 @@ def transform_is_nullable(value):
         return "NO"
     return value
 
+
 def transform_remarks(value):
-    if value is None: 
-        return "" 
-    return value 
+    if value is None:
+        return ""
+    return value
+
 
 def transform_nullable(value):
     """Transform NULLABLE column: boolean/string to integer."""
@@ -69,3 +71,13 @@ def transform_ordinal_position(value):
     if isinstance(value, int):
         return value - 1
     return value
+
+
+def create_table_catalog_transform(catalog_name):
+    """Factory function to create TABLE_CATALOG transform with bound catalog name."""
+
+    def transform_table_catalog(value):
+        """Transform TABLE_CATALOG: return the catalog name for all rows."""
+        return catalog_name
+
+    return transform_table_catalog

tests/e2e/test_driver.py

@@ -386,8 +386,17 @@ def test_create_table_will_return_empty_result_set(self, extra_params):
             finally:
                 cursor.execute("DROP TABLE IF EXISTS {}".format(table_name))
 
-    def test_get_tables(self):
-        with self.cursor({}) as cursor:
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_get_tables(self, backend_params):
+        with self.cursor(backend_params) as cursor:
             table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
             table_names = [table_name + "_1", table_name + "_2"]
 
@@ -432,8 +441,17 @@ def test_get_tables(self):
                 for table in table_names:
                     cursor.execute("DROP TABLE IF EXISTS {}".format(table))
 
-    def test_get_columns(self):
-        with self.cursor({}) as cursor:
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_get_columns(self, backend_params):
+        with self.cursor(backend_params) as cursor:
             table_name = "table_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
             table_names = [table_name + "_1", table_name + "_2"]
 
@@ -544,8 +562,17 @@ def test_escape_single_quotes(self):
             rows = cursor.fetchall()
             assert rows[0]["col_1"] == "you're"
 
-    def test_get_schemas(self):
-        with self.cursor({}) as cursor:
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_get_schemas(self, backend_params):
+        with self.cursor(backend_params) as cursor:
             database_name = "db_{uuid}".format(uuid=str(uuid4()).replace("-", "_"))
             try:
                 cursor.execute("CREATE DATABASE IF NOT EXISTS {}".format(database_name))
@@ -562,8 +589,17 @@ def test_get_schemas(self):
             finally:
                 cursor.execute("DROP DATABASE IF EXISTS {}".format(database_name))
 
-    def test_get_catalogs(self):
-        with self.cursor({}) as cursor:
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_get_catalogs(self, backend_params):
+        with self.cursor(backend_params) as cursor:
             cursor.catalogs()
             cursor.fetchall()
             catalogs_desc = cursor.description