Repository: TobikoData/sqlmesh
Pull request #2912 (2024-07-17 00:02:07)

Feat: pass session parameters for Snowflake
This PR is to implement the features in this issue:
https://github.com/TobikoData/sqlmesh/issues/2829

Looking at the code it didn't seem like much was missing.
Tried it with this in my config:
            session_parameters:
              timezone: UTC
              query_tag: state_test_tag

Not sure how to test it, but I checked in my query history that the queries are getting tagged.
<img width="193" alt="image" src="https://github.com/user-attachments/assets/14e09e5e-faec-4de8-af0e-c529c3b4e077">

From what I understand of the code this should get passed straight through to the connector.
https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect
This would add support for all session parameters supported by Snowflake.
https://docs.snowflake.com/en/sql-reference/parameters

I see here it passes everything defined in the connection kwarg keys.
https://github.com/TobikoData/sqlmesh/blob/7c1ed838295a898da759e59379cdc361b687a22c/sqlmesh/core/config/connection.py#L90
Then that gets passed to Snowflake's connector function and it takes in session_parameters.

I put in a basic test with the others would that be enough?

izeigerman: @alexischicoinedatadeveloper , thanks a lot for the contribution! Make sure to run `make style` and this should be good to go.
CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=2912) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our [Contributor License Agreement](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=2912) before we can accept your contribution.<br/>**1** out of **2** committers have signed the CLA.<br/><br/>:white_check_mark: alexischicoinedatadeveloper<br/>:x: achicoine-coveo<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/TobikoData/sqlmesh?pullRequest=2912) it.</sub>

Patch:
diff --git a/docs/integrations/engines/snowflake.md b/docs/integrations/engines/snowflake.md
index 4aed0dd3..a0a44655 100644
--- a/docs/integrations/engines/snowflake.md
+++ b/docs/integrations/engines/snowflake.md
@@ -24,6 +24,7 @@ pip install "sqlmesh[snowflake]"
 | `private_key`            | The optional private key to use for authentication. Key can be Base64-encoded DER format (representing the key bytes), a plain-text PEM format, or bytes (Python config only). | string |    N     |
 | `private_key_path`       | The optional path to the private key to use for authentication. This would be used instead of `private_key`.                                                                   | string |    N     |
 | `private_key_passphrase` | The optional passphrase to use to decrypt `private_key` (if in PEM format) or `private_key_path`. Keys can be created without encryption so only provide this if needed.       | string |    N     |
+| `session_parameters`     | The optional session parameters to set for the connection.                                                                                                                     | dict   |    N     |
 
 
 #### Lowercase object names
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 343bfbb8..38d20f51 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -333,6 +333,7 @@ class SnowflakeConnectionConfig(ConnectionConfig):
         private_key_passphrase: The optional passphrase to use to decrypt `private_key` or `private_key_path`. Keys can be created without encryption so only provide this if needed.
         register_comments: Whether or not to register model comments with the SQL engine.
         pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
+        session_parameters: The optional session parameters to set for the connection.
     """
 
     account: str
@@ -353,6 +354,8 @@ class SnowflakeConnectionConfig(ConnectionConfig):
     register_comments: bool = True
     pre_ping: bool = False
 
+    session_parameters: t.Optional[dict] = None
+
     type_: Literal["snowflake"] = Field(alias="type", default="snowflake")
 
     _concurrent_tasks_validator = concurrent_tasks_validator
@@ -471,6 +474,7 @@ class SnowflakeConnectionConfig(ConnectionConfig):
             "authenticator",
             "token",
             "private_key",
+            "session_parameters",
         }
 
     @property


Test patch:
diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py
index 10171fa9..6fb158ce 100644
--- a/tests/core/test_connection_config.py
+++ b/tests/core/test_connection_config.py
@@ -183,6 +183,17 @@ def test_snowflake(make_config, snowflake_key_passphrase_bytes, snowflake_oauth_
             authenticator="oauth",
         )
 
+    # Can pass in session parameters
+    config = make_config(
+        type="snowflake",
+        account="test",
+        user="test",
+        password="test",
+        authenticator="externalbrowser",
+        session_parameters={"query_tag": "test", "timezone": "UTC"},
+    )
+    assert isinstance(config, SnowflakeConnectionConfig)
+
 
 @pytest.mark.parametrize(
     "key_path_fixture, key_pem_fixture, key_b64_fixture, key_bytes_fixture, key_passphrase",