Fix: Support macros for 'authorization' and 'query_label' keys in session properties (#4660)

Author: izeigerman

sqlmesh/core/engine_adapter/bigquery.py

@@ -199,6 +199,10 @@ def _begin_session(self, properties: SessionProperties) -> None:
                 (label_tuple.expressions[0].name, label_tuple.expressions[1].name)
                 for label_tuple in label_tuples
             )
+        elif query_label_property is not None:
+            raise SQLMeshError(
+                "Invalid value for `session_properties.query_label`. Must be an array or tuple."
+            )
 
         if parsed_query_label:
             query_label_str = ",".join([":".join(label) for label in parsed_query_label])

sqlmesh/core/engine_adapter/trino.py

@@ -27,6 +27,7 @@
     set_catalog,
 )
 from sqlmesh.core.schema_diff import SchemaDiffer
+from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.date import TimeLike
 
 if t.TYPE_CHECKING:
@@ -99,6 +100,11 @@ def session(self, properties: SessionProperties) -> t.Iterator[None]:
         if not isinstance(authorization, exp.Expression):
             authorization = exp.Literal.string(authorization)
 
+        if not authorization.is_string:
+            raise SQLMeshError(
+                "Invalid value for `session_properties.authorization`. Must be a string literal."
+            )
+
         authorization_sql = authorization.sql(dialect=self.dialect)
 
         self.execute(f"SET SESSION AUTHORIZATION {authorization_sql}")

sqlmesh/core/model/meta.py

@@ -323,10 +323,8 @@ def session_properties_validator(cls, v: t.Any, info: ValidationInfo) -> t.Any:
 
             if prop_name == "query_label":
                 query_label = eq.right
-                if not (
-                    isinstance(query_label, exp.Array)
-                    or isinstance(query_label, exp.Tuple)
-                    or isinstance(query_label, exp.Paren)
+                if not isinstance(
+                    query_label, (exp.Array, exp.Tuple, exp.Paren, d.MacroFunc, d.MacroVar)
                 ):
                     raise ConfigError(
                         "Invalid value for `session_properties.query_label`. Must be an array or tuple."
@@ -349,7 +347,9 @@ def session_properties_validator(cls, v: t.Any, info: ValidationInfo) -> t.Any:
                         )
             elif prop_name == "authorization":
                 authorization = eq.right
-                if not (isinstance(authorization, exp.Literal) and authorization.is_string):
+                if not (
+                    isinstance(authorization, exp.Literal) and authorization.is_string
+                ) and not isinstance(authorization, (d.MacroFunc, d.MacroVar)):
                     raise ConfigError(
                         "Invalid value for `session_properties.authorization`. Must be a string literal."
                     )

tests/core/engine_adapter/test_bigquery.py

@@ -14,6 +14,7 @@
 from sqlmesh.core.engine_adapter.bigquery import select_partitions_expr
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.utils import AttributeDict
+from sqlmesh.utils.errors import SQLMeshError
 
 pytestmark = [pytest.mark.bigquery, pytest.mark.engine]
 
@@ -564,6 +565,14 @@ def test_begin_end_session(mocker: MockerFixture):
     begin_new_session_call = connection_mock._client.query.call_args_list[5]
     assert begin_new_session_call[0][0] == 'SET @@query_label = "key1:value1";SELECT 1;'
 
+    # test invalid query_label value
+    with pytest.raises(
+        SQLMeshError,
+        match="Invalid value for `session_properties.query_label`. Must be an array or tuple.",
+    ):
+        with adapter.session({"query_label": parse_one("'key1:value1'")}):
+            adapter.execute("SELECT 6;")
+
 
 def _to_sql_calls(execute_mock: t.Any, identify: bool = True) -> t.List[str]:
     output = []

tests/core/engine_adapter/test_trino.py

@@ -11,6 +11,7 @@
 from sqlmesh.core.model import load_sql_based_model
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.core.dialect import schema_
+from sqlmesh.utils.errors import SQLMeshError
 from tests.core.engine_adapter import to_sql_calls
 
 pytestmark = [pytest.mark.engine, pytest.mark.trino]
@@ -632,6 +633,14 @@ def test_session_authorization(trino_mocked_engine_adapter: TrinoEngineAdapter):
     except RuntimeError:
         pass
 
+    # Test 5: Invalid authorization value
+    with pytest.raises(
+        SQLMeshError,
+        match="Invalid value for `session_properties.authorization`. Must be a string literal.",
+    ):
+        with adapter.session({"authorization": exp.Literal.number(1)}):
+            adapter.execute("SELECT 1")
+
     assert to_sql_calls(adapter) == [
         "SET SESSION AUTHORIZATION 'test_user'",
         "SELECT 1",

tests/core/test_model.py

@@ -10244,3 +10244,38 @@ def test_invalid_sql_model_query() -> None:
             match=r"^A query is required and must be a SELECT statement, a UNION statement, or a JINJA_QUERY block.*",
         ):
             load_sql_based_model(expressions)
+
+
+def test_query_label_and_authorization_macro():
+    @macro()
+    def test_query_label_macro(evaluator):
+        return "[('key', 'value')]"
+
+    @macro()
+    def test_authorization_macro(evaluator):
+        return exp.Literal.string("test_authorization")
+
+    expressions = d.parse(
+        """
+        MODEL (
+           name db.table,
+           session_properties (
+            query_label = @test_query_label_macro(),
+            authorization = @test_authorization_macro()
+           )
+        );
+
+        SELECT 1 AS c;
+        """
+    )
+
+    model = load_sql_based_model(expressions)
+    assert model.session_properties == {
+        "query_label": d.parse_one("@test_query_label_macro()"),
+        "authorization": d.parse_one("@test_authorization_macro()"),
+    }
+
+    assert model.render_session_properties() == {
+        "query_label": d.parse_one("[('key', 'value')]"),
+        "authorization": d.parse_one("'test_authorization'"),
+    }