Merge branch 'main' into feature/DRM/revert-to-shared-connection-for-databricks-access-token-connections

davem-bis · web-flow · commit 150f41f0f4f0 · 2026-06-24T15:10:26.000+01:00
diff --git a/docs/concepts/models/overview.md b/docs/concepts/models/overview.md
@@ -184,7 +184,7 @@ This table lists each engine's support for `TABLE` and `VIEW` object comments:
 | DuckDB <=0.9  | N                | N               |
 | DuckDB >=0.10 | Y                | Y               |
 | MySQL         | Y                | Y               |
-| MSSQL         | N                | N               |
+| MSSQL         | Y                | Y               |
 | Postgres      | Y                | Y               |
 | GCP Postgres  | Y                | Y               |
 | Redshift      | Y                | N               |
diff --git a/docs/integrations/engines/databricks.md b/docs/integrations/engines/databricks.md
@@ -271,6 +271,28 @@ The only relevant SQLMesh configuration parameter is the optional `catalog` para
 | `disable_databricks_connect`         | When running locally, disable the use of Databricks Connect for all model operations (so use SQL Connector for all models)                                                                                                                          |  bool  |    N     |
 | `disable_spark_session`              | Do not use SparkSession if it is available (like when running in a notebook).                                                                                                                                                                       |  bool  |    N     |
 
+### Query tags
+
+Databricks SQL Connector supports per-query tags through the `query_tags` model session property. Specify tags as a `MAP(...)` of string keys to string or `NULL` values:
+
+```sql
+MODEL (
+    name sqlmesh_example.tagged_model,
+    dialect databricks,
+    session_properties (
+        query_tags = MAP(
+            'team', 'data-eng',
+            'app', 'sqlmesh',
+            'feature', NULL
+        )
+    )
+);
+
+SELECT 1 AS id;
+```
+
+Query tags are only applied when SQLMesh executes SQL through the Databricks SQL Connector. They are not applied when SQLMesh routes execution through Databricks Connect, a Databricks notebook SparkSession, or the Spark engine adapter.
+
 ## Model table properties to support altering tables
 
 If you are making a change to the structure of a table that is [forward only](../../guides/incremental_time.md#forward-only-models), then you may need to add the following to your model's `physical_properties`:
diff --git a/pyproject.toml b/pyproject.toml
@@ -51,7 +51,7 @@ bigquery = [
 # pinned an older SQLGlot which is incompatible with SQLMesh
 bigframes = ["bigframes>=1.32.0"]
 clickhouse = ["clickhouse-connect"]
-databricks = ["databricks-sql-connector[pyarrow]"]
+databricks = ["databricks-sql-connector[pyarrow]>=4.2.6"]
 dev = [
     "agate",
     "beautifulsoup4",
diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py
@@ -30,6 +30,43 @@
 logger = logging.getLogger(__name__)
 
 
+def _query_tags(
+    query_tags: t.Optional[t.Union[exp.Expr, str, int, float, bool]],
+) -> t.Optional[t.Dict[str, t.Optional[str]]]:
+    if not query_tags:
+        return None
+
+    if not isinstance(query_tags, (exp.Map, exp.VarMap)):
+        raise SQLMeshError("Invalid value for `session_properties.query_tags`. Must be a map.")
+
+    keys = query_tags.args.get("keys")
+    values = query_tags.args.get("values")
+    if not isinstance(keys, exp.Array) or not isinstance(values, exp.Array):
+        raise SQLMeshError(
+            "Invalid value for `session_properties.query_tags`. Must be a map with array "
+            "keys and array values."
+        )
+
+    tags: t.Dict[str, t.Optional[str]] = {}
+    for key, value in zip(keys.expressions, values.expressions):
+        if not isinstance(key, exp.Literal) or not key.is_string:
+            raise SQLMeshError(
+                "Invalid key in `session_properties.query_tags`. Keys must be string literals."
+            )
+
+        if isinstance(value, exp.Null):
+            tags[key.this] = None
+        elif isinstance(value, exp.Literal) and value.is_string:
+            tags[key.this] = value.this
+        else:
+            raise SQLMeshError(
+                "Invalid value in `session_properties.query_tags`. Values must be string "
+                "literals or NULL."
+            )
+
+    return tags
+
+
 class DatabricksEngineAdapter(SparkEngineAdapter, GrantsFromInfoSchemaMixin):
     DIALECT = "databricks"
     INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.REPLACE_WHERE
@@ -98,6 +135,12 @@ def _use_spark_session(self) -> bool:
     def is_spark_session_connection(self) -> bool:
         return isinstance(self.connection, SparkSessionConnection)
 
+    @property
+    def _is_databricks_sql_connector_connection(self) -> bool:
+        return not self.is_spark_session_connection and not self._connection_pool.get_attribute(
+            "use_spark_engine_adapter"
+        )
+
     def _set_spark_engine_adapter_if_needed(self) -> None:
         self._spark_engine_adapter = None
 
@@ -181,10 +224,23 @@ def _begin_session(self, properties: SessionProperties) -> t.Any:
         """Begin a new session."""
         # Align the different possible connectors to a single catalog
         self.set_current_catalog(self.default_catalog)  # type: ignore
+        self._connection_pool.set_attribute("query_tags", _query_tags(properties.get("query_tags")))
 
     def _end_session(self) -> None:
+        self._connection_pool.set_attribute("query_tags", None)
         self._connection_pool.set_attribute("use_spark_engine_adapter", False)
 
+    def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None:
+        query_tags = self._connection_pool.get_attribute("query_tags")
+        if (
+            query_tags
+            and "query_tags" not in kwargs
+            and self._is_databricks_sql_connector_connection
+        ):
+            kwargs["query_tags"] = query_tags
+
+        return super()._execute(sql, track_rows_processed, **kwargs)
+
     def _df_to_source_queries(
         self,
         df: DF,
diff --git a/sqlmesh/core/engine_adapter/mssql.py b/sqlmesh/core/engine_adapter/mssql.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from textwrap import dedent
 import typing as t
 import logging
 
@@ -53,8 +54,8 @@ class MSSQLEngineAdapter(
     SUPPORTS_TUPLE_IN = False
     SUPPORTS_MATERIALIZED_VIEWS = False
     CURRENT_CATALOG_EXPRESSION = exp.func("db_name")
-    COMMENT_CREATION_TABLE = CommentCreationTable.UNSUPPORTED
-    COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
+    COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
+    COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
     SUPPORTS_REPLACE_TABLE = False
     MAX_IDENTIFIER_LENGTH = 128
     SUPPORTS_QUERY_EXECUTION_TRACKING = True
@@ -457,3 +458,83 @@ def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> N
             )
 
         return super().delete_from(table_name, where)
+
+    def _build_create_comment_table_exp(
+        self, table: exp.Table, table_comment: str, table_kind: str = "TABLE"
+    ) -> exp.Comment | str:
+        template = dedent("""
+            DECLARE @comment sql_variant = {comment};
+            DECLARE @property_name VARCHAR(128) = 'MS_Description';
+            DECLARE @schema_name VARCHAR(128) = {schema_name};
+            DECLARE @object_name VARCHAR(128) = {object_name};
+            DECLARE @object_kind VARCHAR(128) = '{object_kind}';
+            DECLARE @existing sql_variant;
+
+            SELECT TOP 1 @existing = CAST(VALUE AS NVARCHAR) FROM fn_listextendedproperty(@property_name, 'schema', @schema_name, @object_kind, @object_name, DEFAULT, DEFAULT);
+
+            IF @comment IS NULL
+                BEGIN
+                    IF @existing IS NOT NULL
+                        EXEC sp_dropextendedproperty @property_name, 'schema', @schema_name, @object_kind, @object_name;
+                END
+            ELSE
+                BEGIN
+                    IF @existing IS NULL
+                        EXEC sp_addextendedproperty @property_name,@comment, 'schema', @schema_name, @object_kind, @object_name;
+                    ELSE IF @existing != @comment
+                        EXEC sp_updateextendedproperty @property_name, @comment, 'schema', @schema_name, @object_kind, @object_name;
+                END
+        """)
+        tsql_text = template.format(
+            comment=exp.Literal.string(table_comment or "NULL").sql(
+                dialect=self.dialect, identify=False
+            ),
+            schema_name=exp.Literal.string(table.db or "dbo").sql(
+                dialect=self.dialect, identify=False
+            ),
+            object_name=exp.Literal.string(table.name).sql(dialect=self.dialect, identify=False),
+            object_kind=table_kind,
+        )
+        return tsql_text
+
+    def _build_create_comment_column_exp(
+        self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE"
+    ) -> exp.Comment | str:
+        template = dedent("""
+            DECLARE @comment sql_variant = {comment};
+            DECLARE @property_name VARCHAR(128) = 'MS_Description';
+            DECLARE @schema_name VARCHAR(128) = {schema_name};
+            DECLARE @object_name VARCHAR(128) = {object_name};
+            DECLARE @object_kind VARCHAR(128) = '{object_kind}';
+            DECLARE @column_name VARCHAR(128) = {column_name};
+            DECLARE @existing sql_variant;
+
+            SELECT TOP 1 @existing = CAST(VALUE AS NVARCHAR) FROM fn_listextendedproperty(@property_name, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name);
+
+            IF @comment IS NULL
+                BEGIN
+                    IF @existing IS NOT NULL
+                        EXEC sp_dropextendedproperty @property_name, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
+                END
+            ELSE
+                BEGIN
+                    IF @existing IS NULL
+                        EXEC sp_addextendedproperty @property_name,@comment, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
+                    ELSE IF @existing != @comment
+                        EXEC sp_updateextendedproperty @property_name, @comment, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
+                END
+        """)
+
+        tsql_text = template.format(
+            comment=exp.Literal.string(column_comment or "NULL").sql(
+                dialect=self.dialect, identify=False
+            ),
+            schema_name=exp.Literal.string(table.db or "dbo").sql(
+                dialect=self.dialect, identify=False
+            ),
+            object_name=exp.Literal.string(table.name).sql(dialect=self.dialect, identify=False),
+            object_kind=table_kind,
+            column_name=exp.Literal.string(column_name).sql(dialect=self.dialect, identify=False),
+        )
+
+        return tsql_text
diff --git a/sqlmesh/core/model/meta.py b/sqlmesh/core/model/meta.py
@@ -396,6 +396,38 @@ def session_properties_validator(cls, v: t.Any, info: ValidationInfo) -> t.Any:
                     raise ConfigError(
                         "Invalid value for `session_properties.authorization`. Must be a string literal."
                     )
+            elif prop_name == "query_tags":
+                query_tags = eq.right
+                if isinstance(query_tags, (d.MacroFunc, d.MacroVar)):
+                    continue
+
+                if not isinstance(query_tags, (exp.Map, exp.VarMap)):
+                    raise ConfigError(
+                        "Invalid value for `session_properties.query_tags`. Must be a map."
+                    )
+
+                keys = query_tags.args.get("keys")
+                values = query_tags.args.get("values")
+                if not isinstance(keys, exp.Array) or not isinstance(values, exp.Array):
+                    raise ConfigError(
+                        "Invalid value for `session_properties.query_tags`. Must be a map with array "
+                        "keys and array values."
+                    )
+
+                for key, value in zip(keys.expressions, values.expressions):
+                    if not isinstance(key, exp.Literal) or not key.is_string:
+                        raise ConfigError(
+                            "Invalid key in `session_properties.query_tags`. Keys must be string literals."
+                        )
+
+                    if not (
+                        isinstance(value, exp.Null)
+                        or (isinstance(value, exp.Literal) and value.is_string)
+                    ):
+                        raise ConfigError(
+                            "Invalid value in `session_properties.query_tags`. Values must be string "
+                            "literals or NULL."
+                        )
 
         return parsed_session_properties
 
diff --git a/tests/core/engine_adapter/integration/__init__.py b/tests/core/engine_adapter/integration/__init__.py
@@ -526,6 +526,14 @@ def get_table_comment(
                     AND c.relkind = '{"v" if table_kind == "VIEW" else "r"}'
                 ;
             """
+        elif self.dialect == "tsql":
+            kind = "table" if table_kind == "BASE TABLE" else "view"
+            query = f"""
+                SELECT 
+                	ep.name,
+                    CAST(ep.value AS NVARCHAR(MAX)) comment 
+                FROM fn_listextendedproperty('MS_Description', 'schema', '{schema_name}', '{kind}', '{table_name}', DEFAULT, DEFAULT) ep
+            """
 
         result = self.engine_adapter.fetchall(query)
 
@@ -636,6 +644,16 @@ def get_column_comments(
                     AND c.relkind = '{"v" if table_kind == "VIEW" else "r"}'
                 ;
             """
+        elif self.dialect == "tsql":
+            kind = "table" if table_kind == "BASE TABLE" else "view"
+            query = f"""
+                SELECT
+                    col.COLUMN_NAME column_name,
+                    CAST(ep.value AS NVARCHAR(MAX)) comment 
+                FROM INFORMATION_SCHEMA.COLUMNS col
+                CROSS APPLY fn_listextendedproperty('MS_Description', 'schema', col.TABLE_SCHEMA, '{kind}', col.TABLE_NAME, 'column', col.COLUMN_NAME) ep
+                WHERE col.TABLE_SCHEMA = '{schema_name}' AND col.TABLE_NAME = '{table_name}'
+            """
 
         result = self.engine_adapter.fetchall(query)
 
diff --git a/tests/core/engine_adapter/test_databricks.py b/tests/core/engine_adapter/test_databricks.py
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
