diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py index e3d029a17d..bb671822e3 100644 --- a/sqlmesh/core/engine_adapter/databricks.py +++ b/sqlmesh/core/engine_adapter/databricks.py @@ -411,3 +411,26 @@ def _build_column_defs( return super()._build_column_defs( target_columns_to_types, column_descriptions, is_view, materialized ) + + 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_name").eq(table.name), + exp.column("table_schema").eq(table.db), + exp.column("table_catalog").eq(column_catalog), + ) + ) + .order_by("ordinal_position ASC") + ) + + result = self.cursor.fetchall(query) + + return {row[0]: exp.DataType.build(row[1], dialect=self.dialect) for row in result} diff --git a/tests/core/engine_adapter/test_databricks.py b/tests/core/engine_adapter/test_databricks.py index de91fd3b70..9312616a54 100644 --- a/tests/core/engine_adapter/test_databricks.py +++ b/tests/core/engine_adapter/test_databricks.py @@ -526,3 +526,61 @@ def test_drop_data_object_materialized_view_calls_correct_drop(mocker: MockFixtu drop_view_mock.assert_called_once_with( mv_data_object.to_table(), ignore_if_not_exists=True, materialized=True ) + + +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 = [ + ("bigint_col", "bigint"), + ("binary_col", "binary"), + ("boolean_col", "boolean"), + ("date_col", "date"), + ("decimal_col", "decimal(38,4)"), + ("double_col", "double"), + ("float_col", "float"), + ("int_col", "int"), + ("small_int", "smallint"), + ("string_col", "string"), + ("timestamp_col", "timestamp"), + ("timestamp_ntz_col", "timestamp_ntz"), + ("tinyint_col", "tinyint"), + ("array_col", "array"), + ("simple_struct_col", "struct"), + ("long_struct_col", f"struct<{','.join(long_struct_cols)}>"), + ] + + resp = adapter.columns("test_db.test_table") + assert resp == { + "bigint_col": exp.DataType.build("bigint", dialect=adapter.dialect), + "binary_col": exp.DataType.build("binary", dialect=adapter.dialect), + "boolean_col": exp.DataType.build("boolean", dialect=adapter.dialect), + "date_col": exp.DataType.build("date", dialect=adapter.dialect), + "decimal_col": exp.DataType.build("decimal(38,4)", dialect=adapter.dialect), + "double_col": exp.DataType.build("double", dialect=adapter.dialect), + "float_col": exp.DataType.build("float", dialect=adapter.dialect), + "int_col": exp.DataType.build("int", dialect=adapter.dialect), + "small_int": exp.DataType.build("smallint", dialect=adapter.dialect), + "string_col": exp.DataType.build("string", dialect=adapter.dialect), + "timestamp_col": exp.DataType.build("timestamp", dialect=adapter.dialect), + "timestamp_ntz_col": exp.DataType.build("timestamp_ntz", dialect=adapter.dialect), + "tinyint_col": exp.DataType.build("tinyint", dialect=adapter.dialect), + "array_col": exp.DataType.build("array", dialect=adapter.dialect), + "simple_struct_col": exp.DataType.build("struct", dialect=adapter.dialect), + "long_struct_col": exp.DataType.build( + 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", + ) + )