Merge branch 'main' into add-databricks-query-tags-session-properties

Author: WhammyLeaf

docs/integrations/engines/duckdb.md

@@ -79,6 +79,7 @@ SQLMesh will place models with the explicit catalog "ephemeral", such as `epheme
               type: ducklake
               path: 'catalog.ducklake'
               data_path: data/ducklake
+              override_data_path: true
               encrypted: True
               data_inlining_row_limit: 10
               metadata_schema: main
@@ -105,6 +106,7 @@ SQLMesh will place models with the explicit catalog "ephemeral", such as `epheme
                             type="ducklake",
                             path="catalog.ducklake",
                             data_path="data/ducklake",
+                            override_data_path=False,
                             encrypted=True,
                             data_inlining_row_limit=10,
                             metadata_schema="main",
@@ -120,6 +122,7 @@ SQLMesh will place models with the explicit catalog "ephemeral", such as `epheme
 
 - `path`: Path to the DuckLake catalog file
 - `data_path`: Path where DuckLake data files are stored
+- `override_data_path`: Whether data_override_path option is set
 - `encrypted`: Whether to enable encryption for the catalog (default: `False`)
 - `data_inlining_row_limit`: Maximum number of rows to inline in the catalog (default: `0`)
 - `metadata_schema`: The schema in the catalog server in which to store the DuckLake metadata tables (default: `main`)
@@ -364,6 +367,7 @@ The `filesystems` accepts a list of file systems to register in the DuckDB conne
               type: ducklake
               path: myducklakecatalog.duckdb
               data_path: abfs://MyFabricWorkspace/MyFabricLakehouse.Lakehouse/Files/DuckLake.Files
+              override_data_path: False
         extensions:
           - ducklake
         filesystems:

sqlmesh/core/config/connection.py

@@ -238,6 +238,7 @@ class DuckDBAttachOptions(BaseConfig):
 
     # DuckLake specific options
     data_path: t.Optional[str] = None
+    override_data_path: t.Optional[bool] = False
     encrypted: bool = False
     data_inlining_row_limit: t.Optional[int] = None
     metadata_schema: t.Optional[str] = None
@@ -258,6 +259,8 @@ def to_sql(self, alias: str) -> str:
                 path = f"ducklake:{path}"
             if self.data_path is not None:
                 options.append(f"DATA_PATH '{self.data_path}'")
+                if self.override_data_path:
+                    options.append("OVERRIDE_DATA_PATH true")
             if self.encrypted:
                 options.append("ENCRYPTED")
             if self.data_inlining_row_limit is not None:
@@ -2097,6 +2100,7 @@ class ClickhouseConnectionConfig(ConnectionConfig):
     https_proxy: t.Optional[str] = None
     server_host_name: t.Optional[str] = None
     tls_mode: t.Optional[str] = None
+    secure: bool = False
 
     concurrent_tasks: int = 1
     register_comments: bool = True
@@ -2133,6 +2137,7 @@ def _connection_kwargs_keys(self) -> t.Set[str]:
             "https_proxy",
             "server_host_name",
             "tls_mode",
+            "secure",
         }
         return kwargs
 

sqlmesh/core/dialect.py

@@ -774,7 +774,8 @@ def format_model_expressions(
     if rewrite_casts:
 
         def cast_to_colon(node: exp.Expr) -> exp.Expr:
-            if isinstance(node, exp.Cast) and not any(
+            # Directly check type instead of isinstance to avoid rewriting subclasses of CAST, e.g. JSONCast
+            if type(node) is exp.Cast and not any(
                 # Only convert CAST into :: if it doesn't have additional args set, otherwise this
                 # conversion could alter the semantics (eg. changing SAFE_CAST in BigQuery to CAST)
                 arg

sqlmesh/core/engine_adapter/databricks.py

@@ -467,3 +467,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}

tests/core/engine_adapter/test_databricks.py

@@ -646,3 +646,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<int>"),
+        ("simple_struct_col", "struct<a:int,b:string>"),
+        ("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<int>", dialect=adapter.dialect),
+        "simple_struct_col": exp.DataType.build("struct<a:int,b:string>", 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",
+        )
+    )

tests/core/test_connection_config.py

@@ -810,6 +810,7 @@ def test_duckdb_attach_ducklake_catalog(make_config):
                 type="ducklake",
                 path="catalog.ducklake",
                 data_path="/tmp/ducklake_data",
+                override_data_path=False,
                 encrypted=True,
                 data_inlining_row_limit=10,
             ),

tests/core/test_dialect.py

@@ -207,6 +207,29 @@ def test_format_model_expressions():
   SAFE_CAST('bla' AS INT64) AS FOO"""
     )
 
+    x = format_model_expressions(
+        parse(
+            """
+            MODEL(name a.b, kind FULL, dialect clickhouse);
+            SELECT data.:String AS foo, CAST(1 AS INT) AS bar
+            """
+        ),
+        dialect="clickhouse",
+    )
+    # JSONCast (e.g. `.:` syntax in ClickHouse) must not be written to `::`
+    assert (
+        x
+        == """MODEL (
+  name a.b,
+  kind FULL,
+  dialect clickhouse
+);
+
+SELECT
+  data.:String AS foo,
+  1::Int32 AS bar"""
+    )
+
     x = format_model_expressions(
         parse(
             """