Feat: add support for Trino's authorization session property

Author: georgesittas

sqlmesh/core/engine_adapter/trino.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
+
+import contextlib
+import pandas as pd
 import re
 import typing as t
 from functools import lru_cache
-import pandas as pd
 from pandas.api.types import is_datetime64_any_dtype  # type: ignore
 from sqlglot import exp
 from sqlglot.helper import seq_get
@@ -25,11 +27,12 @@
     SourceQuery,
     set_catalog,
 )
+from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.core.schema_diff import SchemaDiffer
 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
 
 
@@ -88,6 +91,26 @@ 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 isinstance(authorization, str):
+            authorization = exp.Literal.string(authorization)
+        if not (isinstance(authorization, exp.Expression) and authorization.is_string):
+            raise SQLMeshError(f"Invalid authorization: '{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,

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]
@@ -591,3 +592,55 @@ 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: Invalid authorization (non-string expression)
+    adapter.cursor.execute.reset_mock()
+    with pytest.raises(SQLMeshError, match="Invalid authorization"):
+        with adapter.session({"authorization": exp.Literal.number(123)}):
+            pass
+
+    # Test 5: 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",
+    ]