SNOW-1848793: Enable large query breakdown by min version (#2730)

snowflakedb · Dec 9, 2024 · e57cb78 · e57cb78
1 parent 0a64ef8
commit e57cb78
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py
@@ -247,8 +247,8 @@
 _PYTHON_SNOWPARK_REDUCE_DESCRIBE_QUERY_ENABLED = (
     "PYTHON_SNOWPARK_REDUCE_DESCRIBE_QUERY_ENABLED"
 )
-_PYTHON_SNOWPARK_USE_LARGE_QUERY_BREAKDOWN_OPTIMIZATION = (
-    "PYTHON_SNOWPARK_USE_LARGE_QUERY_BREAKDOWN_OPTIMIZATION"
+_PYTHON_SNOWPARK_USE_LARGE_QUERY_BREAKDOWN_OPTIMIZATION_VERSION = (
+    "PYTHON_SNOWPARK_USE_LARGE_QUERY_BREAKDOWN_OPTIMIZATION_VERSION"
 )
 _PYTHON_SNOWPARK_LARGE_QUERY_BREAKDOWN_COMPLEXITY_UPPER_BOUND = (
     "PYTHON_SNOWPARK_LARGE_QUERY_BREAKDOWN_COMPLEXITY_UPPER_BOUND"
@@ -619,10 +619,8 @@ def __init__(
                 _PYTHON_SNOWPARK_ENABLE_QUERY_COMPILATION_STAGE, False
             )
         )
-        self._large_query_breakdown_enabled: bool = (
-            self._conn._get_client_side_session_parameter(
-                _PYTHON_SNOWPARK_USE_LARGE_QUERY_BREAKDOWN_OPTIMIZATION, False
-            )
+        self._large_query_breakdown_enabled: bool = self.is_feature_enabled_for_version(
+            _PYTHON_SNOWPARK_USE_LARGE_QUERY_BREAKDOWN_OPTIMIZATION_VERSION
         )
         self._ast_enabled: bool = self._conn._get_client_side_session_parameter(
             _PYTHON_SNOWPARK_USE_AST, _PYTHON_SNOWPARK_USE_AST_DEFAULT_VALUE

diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
@@ -630,7 +630,12 @@ def test_session_builder_app_name_existing_invalid_json_query_tag():
     ],
 )
 @pytest.mark.parametrize(
-    "parameter_name", ["_auto_clean_up_temp_table_enabled", "_cte_optimization_enabled"]
+    "parameter_name",
+    [
+        "_auto_clean_up_temp_table_enabled",
+        "_cte_optimization_enabled",
+        "_large_query_breakdown_enabled",
+    ],
 )
 def test_parameter_version(version_value, expected_parameter_value, parameter_name):
     fake_server_connection = mock.create_autospec(ServerConnection)