SNOW-1830203: Add a new parameter for CTE optimization (#2689)

snowflakedb · Nov 28, 2024 · ce1d22b · ce1d22b
1 parent 22cdea0
commit ce1d22b
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 22 deletions.
diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py
@@ -235,7 +235,9 @@
 _PYTHON_SNOWPARK_ENABLE_QUERY_COMPILATION_STAGE = (
     "PYTHON_SNOWPARK_COMPILATION_STAGE_ENABLED"
 )
-_PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION_STRING = "PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION"
+_PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION_VERSION = (
+    "PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION_VERSION"
+)
 _PYTHON_SNOWPARK_ELIMINATE_NUMERIC_SQL_VALUE_CAST_ENABLED = (
     "PYTHON_SNOWPARK_ELIMINATE_NUMERIC_SQL_VALUE_CAST_ENABLED"
 )
@@ -589,10 +591,8 @@ def __init__(
                 _PYTHON_SNOWPARK_USE_SQL_SIMPLIFIER_STRING, True
             )
         )
-        self._cte_optimization_enabled: bool = (
-            self._conn._get_client_side_session_parameter(
-                _PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION_STRING, False
-            )
+        self._cte_optimization_enabled: bool = self.is_feature_enabled_for_version(
+            _PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION_VERSION
         )
         self._use_logical_type_for_create_df: bool = (
             self._conn._get_client_side_session_parameter(
@@ -604,16 +604,10 @@ def __init__(
                 _PYTHON_SNOWPARK_ELIMINATE_NUMERIC_SQL_VALUE_CAST_ENABLED, False
             )
         )
-        auto_clean_up_temp_table_enabled_version = (
-            self._conn._get_client_side_session_parameter(
-                _PYTHON_SNOWPARK_AUTO_CLEAN_UP_TEMP_TABLE_ENABLED_VERSION, ""
-            )
-        )
         self._auto_clean_up_temp_table_enabled: bool = (
-            isinstance(auto_clean_up_temp_table_enabled_version, str)
-            and auto_clean_up_temp_table_enabled_version != ""
-            and pkg_resources.parse_version(self.version)
-            >= pkg_resources.parse_version(auto_clean_up_temp_table_enabled_version)
+            self.is_feature_enabled_for_version(
+                _PYTHON_SNOWPARK_AUTO_CLEAN_UP_TEMP_TABLE_ENABLED_VERSION
+            )
         )
         self._reduce_describe_query_enabled: bool = (
             self._conn._get_client_side_session_parameter(
@@ -698,6 +692,19 @@ def __str__(self):
             f"schema={self.get_current_schema()}, warehouse={self.get_current_warehouse()}>"
         )
 
+    def is_feature_enabled_for_version(self, parameter_name: str) -> bool:
+        """
+        This method checks if a feature is enabled for the current session based on
+        the server side parameter.
+        """
+        version = self._conn._get_client_side_session_parameter(parameter_name, "")
+        return (
+            isinstance(version, str)
+            and version != ""
+            and pkg_resources.parse_version(self.version)
+            >= pkg_resources.parse_version(version)
+        )
+
     def _generate_new_action_id(self) -> int:
         with self._lock:
             self._last_action_id += 1

diff --git a/tests/integ/test_session.py b/tests/integ/test_session.py
@@ -24,7 +24,6 @@
 )
 from snowflake.snowpark.session import (
     _PYTHON_SNOWPARK_ELIMINATE_NUMERIC_SQL_VALUE_CAST_ENABLED,
-    _PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION_STRING,
     _PYTHON_SNOWPARK_USE_SQL_SIMPLIFIER_STRING,
     _active_sessions,
     _get_active_session,
@@ -707,12 +706,13 @@ def test_cte_optimization_enabled_on_session(session, db_parameters):
         new_session.cte_optimization_enabled = default_value
         assert new_session.cte_optimization_enabled is default_value
 
-    parameters = db_parameters.copy()
-    parameters["session_parameters"] = {
-        _PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION_STRING: not default_value
-    }
-    with Session.builder.configs(parameters).create() as new_session2:
-        assert new_session2.cte_optimization_enabled is not default_value
+    # TODO SNOW-1830248: add back the test when the parameter is available on the server side
+    # parameters = db_parameters.copy()
+    # parameters["session_parameters"] = {
+    #     _PYTHON_SNOWPARK_USE_CTE_OPTIMIZATION_VERSION: get_version() if default_value else ""
+    # }
+    # with Session.builder.configs(parameters).create() as new_session2:
+    #     assert new_session2.cte_optimization_enabled is not default_value
 
 
 @pytest.mark.skipif(IS_IN_STORED_PROC, reason="Can't create a session in SP")

diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
@@ -629,7 +629,9 @@ def test_session_builder_app_name_existing_invalid_json_query_tag():
         (".".join([str(d + 5) for d in snowpark_version if d is not None]), False),
     ],
 )
-@pytest.mark.parametrize("parameter_name", ["_auto_clean_up_temp_table_enabled"])
+@pytest.mark.parametrize(
+    "parameter_name", ["_auto_clean_up_temp_table_enabled", "_cte_optimization_enabled"]
+)
 def test_parameter_version(version_value, expected_parameter_value, parameter_name):
     fake_server_connection = mock.create_autospec(ServerConnection)
     fake_server_connection._thread_safe_session_enabled = True