From d6ea416178020a012f14069e2fab0267cb57a60c Mon Sep 17 00:00:00 2001 From: abhishek-pattern Date: Mon, 4 May 2026 15:41:51 +0530 Subject: [PATCH 1/6] fix: set Snowflake warehouse via USE WAREHOUSE query and bump version to 0.4.2 Passing warehouse in Snowflake connection parameters silently fails to set it. Switched to executing USE WAREHOUSE as a raw query after connection is established. Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 2 +- src/ds_platform_utils/metaflow/snowflake_connection.py | 4 +++- uv.lock | 2 +- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 83432b4..121b9cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ds-platform-utils" -version = "0.4.1" +version = "0.4.2" description = "Utility library for Pattern Data Science." readme = "README.md" authors = [ diff --git a/src/ds_platform_utils/metaflow/snowflake_connection.py b/src/ds_platform_utils/metaflow/snowflake_connection.py index 2b41107..047d6f1 100644 --- a/src/ds_platform_utils/metaflow/snowflake_connection.py +++ b/src/ds_platform_utils/metaflow/snowflake_connection.py @@ -80,11 +80,13 @@ def _create_snowflake_connection( conn: SnowflakeConnection = Snowflake( integration=SNOWFLAKE_INTEGRATION, client_session_keep_alive=True, - warehouse=warehouse, timezone="UTC" if use_utc else None, session_parameters={"QUERY_TAG": query_tag}, ).cn # type: ignore[attr-defined] + ## Doing this in the connection parameters result in silently failing to set the warehouse, so we have to execute a raw query to set it. + conn.execute_string("USE WAREHOUSE {}".format(warehouse)) + return conn diff --git a/uv.lock b/uv.lock index c15589d..a099e72 100644 --- a/uv.lock +++ b/uv.lock @@ -479,7 +479,7 @@ wheels = [ [[package]] name = "ds-platform-utils" -version = "0.4.1" +version = "0.4.2" source = { editable = "." } dependencies = [ { name = "jinja2" }, From 3d2871c8bab316f48f30f4bbd429ed2192223314 Mon Sep 17 00:00:00 2001 From: abhishek-pattern Date: Mon, 4 May 2026 15:45:25 +0530 Subject: [PATCH 2/6] Fix lint --- src/ds_platform_utils/metaflow/snowflake_connection.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/ds_platform_utils/metaflow/snowflake_connection.py b/src/ds_platform_utils/metaflow/snowflake_connection.py index 047d6f1..1b20e62 100644 --- a/src/ds_platform_utils/metaflow/snowflake_connection.py +++ b/src/ds_platform_utils/metaflow/snowflake_connection.py @@ -84,8 +84,9 @@ def _create_snowflake_connection( session_parameters={"QUERY_TAG": query_tag}, ).cn # type: ignore[attr-defined] - ## Doing this in the connection parameters result in silently failing to set the warehouse, so we have to execute a raw query to set it. - conn.execute_string("USE WAREHOUSE {}".format(warehouse)) + # Doing this in the connection parameters result in silently failing to set the warehouse, + # so we have to execute a raw query to set it. + conn.execute_string("USE WAREHOUSE {}".format(warehouse)) return conn From ffc7a6db473a7610b8c3dad050542b68f0009fdd Mon Sep 17 00:00:00 2001 From: abhishek-pattern Date: Mon, 4 May 2026 15:45:42 +0530 Subject: [PATCH 3/6] fix lint --- src/ds_platform_utils/metaflow/snowflake_connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ds_platform_utils/metaflow/snowflake_connection.py b/src/ds_platform_utils/metaflow/snowflake_connection.py index 1b20e62..8a17a8b 100644 --- a/src/ds_platform_utils/metaflow/snowflake_connection.py +++ b/src/ds_platform_utils/metaflow/snowflake_connection.py @@ -84,7 +84,7 @@ def _create_snowflake_connection( session_parameters={"QUERY_TAG": query_tag}, ).cn # type: ignore[attr-defined] - # Doing this in the connection parameters result in silently failing to set the warehouse, + # Doing this in the connection parameters result in silently failing to set the warehouse, # so we have to execute a raw query to set it. conn.execute_string("USE WAREHOUSE {}".format(warehouse)) From ba03beab83bd1e3921917708dc43e946bd11e57d Mon Sep 17 00:00:00 2001 From: abhishek-pattern Date: Mon, 11 May 2026 15:51:31 +0530 Subject: [PATCH 4/6] fix: add error handling when setting Snowflake warehouse in connection --- src/ds_platform_utils/metaflow/snowflake_connection.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/ds_platform_utils/metaflow/snowflake_connection.py b/src/ds_platform_utils/metaflow/snowflake_connection.py index 8a17a8b..e9828a5 100644 --- a/src/ds_platform_utils/metaflow/snowflake_connection.py +++ b/src/ds_platform_utils/metaflow/snowflake_connection.py @@ -86,7 +86,10 @@ def _create_snowflake_connection( # Doing this in the connection parameters result in silently failing to set the warehouse, # so we have to execute a raw query to set it. - conn.execute_string("USE WAREHOUSE {}".format(warehouse)) + try: + conn.execute_string("USE WAREHOUSE {}".format(warehouse)) + except Exception as e: + raise RuntimeError(f"Failed to set Snowflake warehouse to {warehouse}: {e}") from e return conn From b20bb9d25f87394f27b6924e1c05244031895990 Mon Sep 17 00:00:00 2001 From: abhishek-pattern Date: Mon, 11 May 2026 16:07:21 +0530 Subject: [PATCH 5/6] fix: update Snowflake connection import and usage in tests --- tests/unit_tests/snowflake/test__execute_sql.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/snowflake/test__execute_sql.py b/tests/unit_tests/snowflake/test__execute_sql.py index b328988..111e287 100644 --- a/tests/unit_tests/snowflake/test__execute_sql.py +++ b/tests/unit_tests/snowflake/test__execute_sql.py @@ -7,7 +7,7 @@ from snowflake.connector import SnowflakeConnection from ds_platform_utils._snowflake.run_query import _execute_sql -from ds_platform_utils.metaflow.snowflake_connection import _create_snowflake_connection +from ds_platform_utils.metaflow.snowflake_connection import get_snowflake_connection @pytest.fixture(scope="module") @@ -30,7 +30,7 @@ def patched_current() -> Generator[MagicMock, None, None]: @pytest.fixture(scope="module") def snowflake_conn(patched_current) -> Generator[SnowflakeConnection, None, None]: """Get a Snowflake connection for testing.""" - yield _create_snowflake_connection(warehouse=None, use_utc=True) + yield get_snowflake_connection(warehouse=None, use_utc=True) def test_execute_sql_empty_string(snowflake_conn): From 96ce66d60ee48bf36fc7fa7094590bbe50efa608 Mon Sep 17 00:00:00 2001 From: abhishek-pattern Date: Mon, 11 May 2026 16:26:35 +0530 Subject: [PATCH 6/6] fix: simplify Snowflake connection fixture by removing unnecessary patching --- .../unit_tests/snowflake/test__execute_sql.py | 23 ++++--------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/tests/unit_tests/snowflake/test__execute_sql.py b/tests/unit_tests/snowflake/test__execute_sql.py index 111e287..1c8ce7b 100644 --- a/tests/unit_tests/snowflake/test__execute_sql.py +++ b/tests/unit_tests/snowflake/test__execute_sql.py @@ -1,7 +1,6 @@ """Functional test for _execute_sql.""" from typing import Generator -from unittest.mock import MagicMock import pytest from snowflake.connector import SnowflakeConnection @@ -11,25 +10,11 @@ @pytest.fixture(scope="module") -def patched_current() -> Generator[MagicMock, None, None]: - """Patch Metaflow `current` object for modules used in this test file.""" - mock_current = MagicMock("metaflow.current") - mock_current.tags = ["ds.domain:testing", "ds.project:unit-tests"] - mock_current.flow_name = "DummyFlow" - mock_current.project_name = "dummy-project" - mock_current.step_name = "dummy-step" - mock_current.run_id = "123" - mock_current.username = "tester" - mock_current.is_production = False - mock_current.namespace = "user:tester" - mock_current.is_running_flow = True - mock_current.card = [] - yield mock_current - - -@pytest.fixture(scope="module") -def snowflake_conn(patched_current) -> Generator[SnowflakeConnection, None, None]: +def snowflake_conn() -> Generator[SnowflakeConnection, None, None]: """Get a Snowflake connection for testing.""" + from metaflow import current + + current.is_production = False # Ensure we're in non-prod for testing yield get_snowflake_connection(warehouse=None, use_utc=True)