Minor refactor and adding mocked current catalog to unit test

Signed-off-by: Bjarke Enkelund <47357343+MisterWheatley@users.noreply.github.com>

Author: MisterWheatley

sqlmesh/core/engine_adapter/databricks.py

@@ -416,14 +416,16 @@ def columns(
         self, table_name: TableName, include_pseudo_columns: bool = False
     ) -> t.Dict[str, exp.DataType]:
         table = exp.to_table(table_name)
+
+        column_catalog = table.catalog or self.get_current_catalog()
         query = (
             exp.select("columns.column_name", "columns.full_data_type")
             .from_("system.information_schema.columns")
             .where(
                 exp.and_(
-                    exp.column("table_catalog").eq(table.catalog),
-                    exp.column("table_schema").eq(table.db),
                     exp.column("table_name").eq(table.name),
+                    exp.column("table_schema").eq(table.db),
+                    exp.column("table_catalog").eq(column_catalog),
                 )
             )
             .order_by("ordinal_position ASC")

tests/core/engine_adapter/test_databricks.py

@@ -531,6 +531,10 @@ def test_drop_data_object_materialized_view_calls_correct_drop(mocker: MockFixtu
 def test_columns(mocker: MockFixture, make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
 
+    # Override/mock get_current_catalog to return default
+    current_catalog_mock = mocker.patch.object(
+        adapter, "get_current_catalog", return_value="test_catalog"
+    )
     # create long struct columns datatype
     long_struct_cols = [f"a_{i}:int" for i in range(50)]
     adapter.cursor.fetchall.return_value = [
@@ -573,3 +577,10 @@ def test_columns(mocker: MockFixture, make_mocked_engine_adapter: t.Callable):
             f"struct<{','.join(long_struct_cols)}>", dialect=adapter.dialect
         ),
     }
+
+    adapter.cursor.fetchall.assert_called_once_with(
+        parse_one(
+            """SELECT columns.column_name, columns.full_data_type FROM system.information_schema.columns WHERE table_name = 'test_table' AND table_schema = 'test_db' AND table_catalog = 'test_catalog' ORDER BY ordinal_position ASC""",
+            dialect="databricks",
+        )
+    )