add static validation + documentation + more tests

Signed-off-by: Christian Troelsen <christian.troelsen@tryg.dk>

Author: WhammyLeaf

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`:

sqlmesh/core/engine_adapter/databricks.py

@@ -36,7 +36,7 @@ def _query_tags(
     if not query_tags:
         return None
 
-    if not isinstance(query_tags, exp.Map):
+    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")

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
 

tests/core/engine_adapter/test_databricks.py

@@ -135,7 +135,13 @@ def test_session_query_tags(mocker: MockFixture, make_mocked_engine_adapter: t.C
     )
     adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
 
-    with adapter.session({"query_tags": _query_tags_map("team", "data-eng", "app", "sqlmesh")}):
+    with adapter.session(
+        {
+            "query_tags": d.parse_one(
+                "MAP('team', 'data-eng', 'app', 'sqlmesh')", dialect="databricks"
+            )
+        }
+    ):
         adapter.execute("SELECT 1")
 
     adapter.cursor.execute.assert_called_with(

tests/core/test_model.py

@@ -5190,6 +5190,90 @@ def test_session_properties_authorization_validation():
         )
 
 
+def test_session_properties_query_tags_validation():
+    model = load_sql_based_model(
+        d.parse(
+            """
+        MODEL (
+            name test_schema.test_model,
+            dialect databricks,
+            session_properties (
+                query_tags = MAP('team', 'data-eng', 'app', 'sqlmesh', 'feature', NULL)
+            )
+        );
+        SELECT a FROM tbl;
+        """,
+            default_dialect="databricks",
+        )
+    )
+    assert model.session_properties == {
+        "query_tags": parse_one(
+            "MAP('team', 'data-eng', 'app', 'sqlmesh', 'feature', NULL)",
+            dialect="databricks",
+        )
+    }
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid value for `session_properties.query_tags`. Must be a map.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                dialect databricks,
+                session_properties (
+                    query_tags = 'invalid value'
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="databricks",
+            )
+        )
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid key in `session_properties.query_tags`. Keys must be string literals.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                dialect databricks,
+                session_properties (
+                    query_tags = MAP(1, 'data-eng')
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="databricks",
+            )
+        )
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid value in `session_properties.query_tags`. Values must be string literals or NULL.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                dialect databricks,
+                session_properties (
+                    query_tags = MAP('team', 1)
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="databricks",
+            )
+        )
+
+
 def test_model_jinja_macro_rendering():
     expressions = d.parse(
         """
@@ -11719,6 +11803,35 @@ def test_authorization_macro(evaluator):
     }
 
 
+def test_query_tags_macro() -> None:
+    @macro()
+    def test_query_tags_macro(evaluator):
+        return "MAP('team', 'data-eng')"
+
+    expressions = d.parse(
+        """
+        MODEL (
+           name db.table,
+           dialect databricks,
+           session_properties (
+            query_tags = @test_query_tags_macro()
+           )
+        );
+
+        SELECT 1 AS c;
+        """
+    )
+
+    model = load_sql_based_model(expressions)
+    assert model.session_properties == {
+        "query_tags": d.parse_one("@test_query_tags_macro()"),
+    }
+
+    assert model.render_session_properties() == {
+        "query_tags": d.parse_one("MAP('team', 'data-eng')", dialect="databricks"),
+    }
+
+
 def test_boolean_property_validation() -> None:
     expressions = d.parse(
         """