Feat: add support for Trino's authorization session property (#4639)

Author: georgesittas

sqlmesh/core/engine_adapter/trino.py

@@ -1,4 +1,6 @@
 from __future__ import annotations
+
+import contextlib
 import re
 import typing as t
 from functools import lru_cache
@@ -28,7 +30,7 @@
 from sqlmesh.utils.date import TimeLike
 
 if t.TYPE_CHECKING:
-    from sqlmesh.core._typing import SchemaName, TableName
+    from sqlmesh.core._typing import SchemaName, SessionProperties, TableName
     from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF
 
 
@@ -87,6 +89,24 @@ def get_catalog_type(self, catalog: t.Optional[str]) -> str:
             )
         return seq_get(row, 0) or self.DEFAULT_CATALOG_TYPE
 
+    @contextlib.contextmanager
+    def session(self, properties: SessionProperties) -> t.Iterator[None]:
+        authorization = properties.get("authorization")
+        if not authorization:
+            yield
+            return
+
+        if not isinstance(authorization, exp.Expression):
+            authorization = exp.Literal.string(authorization)
+
+        authorization_sql = authorization.sql(dialect=self.dialect)
+
+        self.execute(f"SET SESSION AUTHORIZATION {authorization_sql}")
+        try:
+            yield
+        finally:
+            self.execute(f"RESET SESSION AUTHORIZATION")
+
     def _insert_overwrite_by_condition(
         self,
         table_name: TableName,

sqlmesh/core/model/meta.py

@@ -319,7 +319,9 @@ def session_properties_validator(cls, v: t.Any, info: ValidationInfo) -> t.Any:
             return parsed_session_properties
 
         for eq in parsed_session_properties:
-            if eq.name == "query_label":
+            prop_name = eq.left.name
+
+            if prop_name == "query_label":
                 query_label = eq.right
                 if not (
                     isinstance(query_label, exp.Array)
@@ -345,6 +347,12 @@ def session_properties_validator(cls, v: t.Any, info: ValidationInfo) -> t.Any:
                         raise ConfigError(
                             "Invalid entry in `session_properties.query_label`. Must be tuples of string literals with length 2."
                         )
+            elif prop_name == "authorization":
+                authorization = eq.right
+                if not (isinstance(authorization, exp.Literal) and authorization.is_string):
+                    raise ConfigError(
+                        "Invalid value for `session_properties.authorization`. Must be a string literal."
+                    )
 
         return parsed_session_properties
 

tests/core/engine_adapter/test_trino.py

@@ -591,3 +591,49 @@ def test_create_schema_sets_location(make_mocked_engine_adapter: t.Callable, moc
             'CREATE SCHEMA IF NOT EXISTS "landing"."transactions" WITH (LOCATION=\'s3://raw-data/landing/transactions\')',  # match '^landing\..*$'
         ]
     )
+
+
+def test_session_authorization(trino_mocked_engine_adapter: TrinoEngineAdapter):
+    adapter = trino_mocked_engine_adapter
+
+    # Test 1: No authorization property - should not execute any authorization commands
+    with adapter.session({}):
+        pass
+
+    assert to_sql_calls(adapter) == []
+
+    # Test 2: String authorization
+    with adapter.session({"authorization": "test_user"}):
+        adapter.execute("SELECT 1")
+
+    assert to_sql_calls(adapter) == [
+        "SET SESSION AUTHORIZATION 'test_user'",
+        "SELECT 1",
+        "RESET SESSION AUTHORIZATION",
+    ]
+
+    # Test 3: Expression authorization
+    adapter.cursor.execute.reset_mock()
+    with adapter.session({"authorization": exp.Literal.string("another_user")}):
+        adapter.execute("SELECT 2")
+
+    assert to_sql_calls(adapter) == [
+        "SET SESSION AUTHORIZATION 'another_user'",
+        "SELECT 2",
+        "RESET SESSION AUTHORIZATION",
+    ]
+
+    # Test 4: RESET is called even if exception occurs during session
+    adapter.cursor.execute.reset_mock()
+    try:
+        with adapter.session({"authorization": "test_user"}):
+            adapter.execute("SELECT 1")
+            raise RuntimeError("Test exception")
+    except RuntimeError:
+        pass
+
+    assert to_sql_calls(adapter) == [
+        "SET SESSION AUTHORIZATION 'test_user'",
+        "SELECT 1",
+        "RESET SESSION AUTHORIZATION",
+    ]

tests/core/test_model.py

@@ -4531,6 +4531,81 @@ def test_model_session_properties(sushi_context):
         )
 
 
+def test_session_properties_authorization_validation():
+    model = load_sql_based_model(
+        d.parse(
+            """
+        MODEL (
+            name test_schema.test_model,
+            session_properties (
+                authorization = 'test_user'
+            )
+        );
+        SELECT a FROM tbl;
+        """,
+            default_dialect="trino",
+        )
+    )
+    assert model.session_properties == {"authorization": "test_user"}
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid value for `session_properties.authorization`. Must be a string literal.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                session_properties (
+                    authorization = 123
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="trino",
+            )
+        )
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid value for `session_properties.authorization`. Must be a string literal.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                session_properties (
+                    authorization = some_column
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="trino",
+            )
+        )
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid value for `session_properties.authorization`. Must be a string literal.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                session_properties (
+                    authorization = CONCAT('user', '_suffix')
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="trino",
+            )
+        )
+
+
 def test_model_jinja_macro_rendering():
     expressions = d.parse(
         """