sqlmesh 0.188.2.dev19__tar.gz → 0.188.2.dev21__tar.gz

{sqlmesh-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.188.2.dev19
+Version: 0.188.2.dev21
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.188.2.dev19'
-__version_tuple__ = version_tuple = (0, 188, 2, 'dev19')
+__version__ = version = '0.188.2.dev21'
+__version_tuple__ = version_tuple = (0, 188, 2, 'dev21')

{sqlmesh-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/sqlmesh/core/engine_adapter/snowflake.py

@@ -93,8 +93,10 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
             return
 
         self.execute(f"USE WAREHOUSE {warehouse_sql}")
-        yield
-        self.execute(f"USE WAREHOUSE {current_warehouse_sql}")
+        try:
+            yield
+        finally:
+            self.execute(f"USE WAREHOUSE {current_warehouse_sql}")
 
     @property
     def _current_warehouse(self) -> exp.Identifier:

{sqlmesh-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/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.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
 
 
@@ -87,6 +89,24 @@ class TrinoEngineAdapter(
             )
         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-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/sqlmesh/core/model/meta.py

@@ -319,7 +319,9 @@ class ModelMeta(_Node):
             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 @@ class ModelMeta(_Node):
                         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
 

{sqlmesh-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/sqlmesh.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.188.2.dev19
+Version: 0.188.2.dev21
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/tests/core/engine_adapter/test_snowflake.py

@@ -98,6 +98,7 @@ def test_session(
     adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
     adapter.cursor.fetchone.return_value = (current_warehouse,)
 
+    # Test normal execution
     with adapter.session({"warehouse": configured_warehouse}):
         pass
 
@@ -114,6 +115,27 @@ def test_session(
 
     assert to_sql_calls(adapter) == expected_calls
 
+    # Test exception handling - warehouse should still be reset
+    if should_change:
+        adapter.cursor.execute.reset_mock()
+        adapter.cursor.fetchone.return_value = (current_warehouse,)
+
+        try:
+            with adapter.session({"warehouse": configured_warehouse}):
+                adapter.execute("SELECT 1")
+                raise RuntimeError("Test exception")
+        except RuntimeError:
+            pass
+
+        expected_exception_calls = [
+            "SELECT CURRENT_WAREHOUSE()",
+            f"USE WAREHOUSE {configured_warehouse_exp}",
+            "SELECT 1",
+            f"USE WAREHOUSE {current_warehouse_exp}",
+        ]
+
+        assert to_sql_calls(adapter) == expected_exception_calls
+
 
 def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
     adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)

{sqlmesh-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/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",
+    ]

{sqlmesh-0.188.2.dev19 → sqlmesh-0.188.2.dev21}/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(
         """