SNOW-1663726 make session config updates thread safe

src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py

@@ -113,17 +113,14 @@ def __init__(
         session: Session,
         query_generator: QueryGenerator,
         logical_plans: List[LogicalPlan],
+        complexity_bounds: Tuple[int, int],
     ) -> None:
         self.session = session
         self._query_generator = query_generator
         self.logical_plans = logical_plans
         self._parent_map = defaultdict(set)
-        self.complexity_score_lower_bound = (
-            session.large_query_breakdown_complexity_bounds[0]
-        )
-        self.complexity_score_upper_bound = (
-            session.large_query_breakdown_complexity_bounds[1]
-        )
+        self.complexity_score_lower_bound = complexity_bounds[0]
+        self.complexity_score_upper_bound = complexity_bounds[1]
 
     def apply(self) -> List[LogicalPlan]:
         if is_active_transaction(self.session):

src/snowflake/snowpark/_internal/compiler/plan_compiler.py

@@ -77,6 +77,7 @@ def should_start_query_compilation(self) -> bool:
 
     def compile(self) -> Dict[PlanQueryType, List[Query]]:
         if self.should_start_query_compilation():
+            session = self._plan.session
             # preparation for compilation
             # 1. make a copy of the original plan
             start_time = time.time()
@@ -93,7 +94,7 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
             # 3. apply each optimizations if needed
             # CTE optimization
             cte_start_time = time.time()
-            if self._plan.session.cte_optimization_enabled:
+            if session.cte_optimization_enabled:
                 repeated_subquery_eliminator = RepeatedSubqueryElimination(
                     logical_plans, query_generator
                 )
@@ -108,9 +109,12 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
                 plot_plan_if_enabled(plan, f"cte_optimized_plan_{i}")
 
             # Large query breakdown
-            if self._plan.session.large_query_breakdown_enabled:
+            if session.large_query_breakdown_enabled:
                 large_query_breakdown = LargeQueryBreakdown(
-                    self._plan.session, query_generator, logical_plans
+                    session,
+                    query_generator,
+                    logical_plans,
+                    session.large_query_breakdown_complexity_bounds,
                 )
                 logical_plans = large_query_breakdown.apply()
 
@@ -130,7 +134,6 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
             cte_time = cte_end_time - cte_start_time
             large_query_breakdown_time = large_query_breakdown_end_time - cte_end_time
             total_time = time.time() - start_time
-            session = self._plan.session
             summary_value = {
                 TelemetryField.CTE_OPTIMIZATION_ENABLED.value: session.cte_optimization_enabled,
                 TelemetryField.LARGE_QUERY_BREAKDOWN_ENABLED.value: session.large_query_breakdown_enabled,
@@ -151,8 +154,7 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
             return queries
         else:
             final_plan = self._plan
-            if self._plan.session.cte_optimization_enabled:
-                final_plan = final_plan.replace_repeated_subquery_with_cte()
+            final_plan = final_plan.replace_repeated_subquery_with_cte()
             return {
                 PlanQueryType.QUERIES: final_plan.queries,
                 PlanQueryType.POST_ACTIONS: final_plan.post_actions,

src/snowflake/snowpark/_internal/utils.py

@@ -17,6 +17,7 @@
 import random
 import re
 import string
+import threading
 import traceback
 import zipfile
 from enum import Enum
@@ -297,6 +298,15 @@ def normalize_path(path: str, is_local: bool) -> str:
     return f"'{path}'"
 
 
+def warn_session_config_update_in_multithreaded_mode(config) -> None:
+    if threading.active_count() > 1:
+        logger.warning(
+            "You might have more than one threads sharing the Session object trying to update "
+            f"{config}. Updating this while other tasks are running can potentially cause "
+            "unexpected behavior. Please update the session configuration before starting the threads."
+        )
+
+
 def normalize_remote_file_or_dir(name: str) -> str:
     return normalize_path(name, is_local=False)
 

src/snowflake/snowpark/session.py

@@ -114,6 +114,7 @@
     unwrap_single_quote,
     unwrap_stage_location_single_quote,
     validate_object_name,
+    warn_session_config_update_in_multithreaded_mode,
     warning,
     zip_file_or_directory_to_stream,
 )
@@ -339,38 +340,44 @@ def __init__(self, session: "Session", conf: Dict[str, Any]) -> None:
                 "use_constant_subquery_alias": True,
                 "flatten_select_after_filter_and_orderby": True,
             }  # For config that's temporary/to be removed soon
+            self._lock = self._session._lock
             for key, val in conf.items():
                 if self.is_mutable(key):
                     self.set(key, val)
 
         def get(self, key: str, default=None) -> Any:
-            if hasattr(Session, key):
-                return getattr(self._session, key)
-            if hasattr(self._session._conn._conn, key):
-                return getattr(self._session._conn._conn, key)
-            return self._conf.get(key, default)
+            with self._lock:
+                if hasattr(Session, key):
+                    return getattr(self._session, key)
+                if hasattr(self._session._conn._conn, key):
+                    return getattr(self._session._conn._conn, key)
+                return self._conf.get(key, default)
 
         def is_mutable(self, key: str) -> bool:
-            if hasattr(Session, key) and isinstance(getattr(Session, key), property):
-                return getattr(Session, key).fset is not None
-            if hasattr(SnowflakeConnection, key) and isinstance(
-                getattr(SnowflakeConnection, key), property
-            ):
-                return getattr(SnowflakeConnection, key).fset is not None
-            return key in self._conf
+            with self._lock:
+                if hasattr(Session, key) and isinstance(
+                    getattr(Session, key), property
+                ):
+                    return getattr(Session, key).fset is not None
+                if hasattr(SnowflakeConnection, key) and isinstance(
+                    getattr(SnowflakeConnection, key), property
+                ):
+                    return getattr(SnowflakeConnection, key).fset is not None
+                return key in self._conf
 
         def set(self, key: str, value: Any) -> None:
-            if self.is_mutable(key):
-                if hasattr(Session, key):
-                    setattr(self._session, key, value)
-                if hasattr(SnowflakeConnection, key):
-                    setattr(self._session._conn._conn, key, value)
-                if key in self._conf:
-                    self._conf[key] = value
-            else:
-                raise AttributeError(
-                    f'Configuration "{key}" does not exist or is not mutable in runtime'
-                )
+            with self._lock:
+                if self.is_mutable(key):
+                    if hasattr(Session, key):
+                        setattr(self._session, key, value)
+                    if hasattr(SnowflakeConnection, key):
+                        setattr(self._session._conn._conn, key, value)
+                    if key in self._conf:
+                        self._conf[key] = value
+                else:
+                    raise AttributeError(
+                        f'Configuration "{key}" does not exist or is not mutable in runtime'
+                    )
 
     class SessionBuilder:
         """
@@ -770,36 +777,46 @@ def custom_package_usage_config(self) -> Dict:
 
     @sql_simplifier_enabled.setter
     def sql_simplifier_enabled(self, value: bool) -> None:
-        self._conn._telemetry_client.send_sql_simplifier_telemetry(
-            self._session_id, value
-        )
-        try:
-            self._conn._cursor.execute(
-                f"alter session set {_PYTHON_SNOWPARK_USE_SQL_SIMPLIFIER_STRING} = {value}"
+        warn_session_config_update_in_multithreaded_mode("sql_simplifier_enabled")
+
+        with self._lock:
+            self._conn._telemetry_client.send_sql_simplifier_telemetry(
+                self._session_id, value
             )
-        except Exception:
-            pass
-        self._sql_simplifier_enabled = value
+            try:
+                self._conn._cursor.execute(
+                    f"alter session set {_PYTHON_SNOWPARK_USE_SQL_SIMPLIFIER_STRING} = {value}"
+                )
+            except Exception:
+                pass
+            self._sql_simplifier_enabled = value
 
     @cte_optimization_enabled.setter
     @experimental_parameter(version="1.15.0")
     def cte_optimization_enabled(self, value: bool) -> None:
-        if value:
-            self._conn._telemetry_client.send_cte_optimization_telemetry(
-                self._session_id
-            )
-        self._cte_optimization_enabled = value
+        warn_session_config_update_in_multithreaded_mode("cte_optimization_enabled")
+
+        with self._lock:
+            if value:
+                self._conn._telemetry_client.send_cte_optimization_telemetry(
+                    self._session_id
+                )
+            self._cte_optimization_enabled = value
 
     @eliminate_numeric_sql_value_cast_enabled.setter
     @experimental_parameter(version="1.20.0")
     def eliminate_numeric_sql_value_cast_enabled(self, value: bool) -> None:
         """Set the value for eliminate_numeric_sql_value_cast_enabled"""
+        warn_session_config_update_in_multithreaded_mode(
+            "eliminate_numeric_sql_value_cast_enabled"
+        )
 
         if value in [True, False]:
-            self._conn._telemetry_client.send_eliminate_numeric_sql_value_cast_telemetry(
-                self._session_id, value
-            )
-            self._eliminate_numeric_sql_value_cast_enabled = value
+            with self._lock:
+                self._conn._telemetry_client.send_eliminate_numeric_sql_value_cast_telemetry(
+                    self._session_id, value
+                )
+                self._eliminate_numeric_sql_value_cast_enabled = value
         else:
             raise ValueError(
                 "value for eliminate_numeric_sql_value_cast_enabled must be True or False!"
@@ -809,6 +826,10 @@ def eliminate_numeric_sql_value_cast_enabled(self, value: bool) -> None:
     @experimental_parameter(version="1.21.0")
     def auto_clean_up_temp_table_enabled(self, value: bool) -> None:
         """Set the value for auto_clean_up_temp_table_enabled"""
+        warn_session_config_update_in_multithreaded_mode(
+            "auto_clean_up_temp_table_enabled"
+        )
+
         if value in [True, False]:
             self._conn._telemetry_client.send_auto_clean_up_temp_table_telemetry(
                 self._session_id, value
@@ -827,12 +848,16 @@ def large_query_breakdown_enabled(self, value: bool) -> None:
         materialize the partitions, and then combine them to execute the query to improve
         overall performance.
         """
+        warn_session_config_update_in_multithreaded_mode(
+            "large_query_breakdown_enabled"
+        )
 
         if value in [True, False]:
-            self._conn._telemetry_client.send_large_query_breakdown_telemetry(
-                self._session_id, value
-            )
-            self._large_query_breakdown_enabled = value
+            with self._lock:
+                self._conn._telemetry_client.send_large_query_breakdown_telemetry(
+                    self._session_id, value
+                )
+                self._large_query_breakdown_enabled = value
         else:
             raise ValueError(
                 "value for large_query_breakdown_enabled must be True or False!"
@@ -841,6 +866,9 @@ def large_query_breakdown_enabled(self, value: bool) -> None:
     @large_query_breakdown_complexity_bounds.setter
     def large_query_breakdown_complexity_bounds(self, value: Tuple[int, int]) -> None:
         """Set the lower and upper bounds for the complexity score used in large query breakdown optimization."""
+        warn_session_config_update_in_multithreaded_mode(
+            "large_query_breakdown_complexity_bounds"
+        )
 
         if len(value) != 2:
             raise ValueError(
@@ -850,16 +878,20 @@ def large_query_breakdown_complexity_bounds(self, value: Tuple[int, int]) -> Non
             raise ValueError(
                 f"Expecting a tuple of lower and upper bound with the lower bound less than the upper bound. Got (lower, upper) = ({value[0], value[1]})"
             )
-        self._conn._telemetry_client.send_large_query_breakdown_update_complexity_bounds(
-            self._session_id, value[0], value[1]
-        )
+        with self._lock:
+            self._conn._telemetry_client.send_large_query_breakdown_update_complexity_bounds(
+                self._session_id, value[0], value[1]
+            )
 
-        self._large_query_breakdown_complexity_bounds = value
+            self._large_query_breakdown_complexity_bounds = value
 
     @custom_package_usage_config.setter
     @experimental_parameter(version="1.6.0")
     def custom_package_usage_config(self, config: Dict) -> None:
-        self._custom_package_usage_config = {k.lower(): v for k, v in config.items()}
+        with self._lock:
+            self._custom_package_usage_config = {
+                k.lower(): v for k, v in config.items()
+            }
 
     def cancel_all(self) -> None:
         """
@@ -1458,7 +1490,8 @@ def _get_dependency_packages(
             statement_params=statement_params,
         )
 
-        custom_package_usage_config = self._custom_package_usage_config.copy()
+        with self._lock:
+            custom_package_usage_config = self._custom_package_usage_config.copy()
 
         unsupported_packages: List[str] = []
         for package, package_info in package_dict.items():

tests/integ/test_multithreading.py

@@ -3,6 +3,7 @@
 #
 
 import hashlib
+import logging
 import os
 import tempfile
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -26,7 +27,7 @@
 
 from snowflake.snowpark.functions import lit
 from snowflake.snowpark.row import Row
-from tests.utils import IS_IN_STORED_PROC, TestFiles, Utils
+from tests.utils import IS_IN_STORED_PROC, IS_LINUX, IS_WINDOWS, TestFiles, Utils
 
 
 def test_concurrent_select_queries(session):
@@ -502,3 +503,38 @@ def finish(self):
     with ThreadPoolExecutor(max_workers=10) as executor:
         for i in range(10):
             executor.submit(register_and_test_udaf, session, i)
+
+
+@pytest.mark.skipif(
+    IS_LINUX or IS_WINDOWS, reason="Linux and Windows behave badly for this test"
+)
+@pytest.mark.parametrize(
+    "config,value",
+    [
+        ("cte_optimization_enabled", True),
+        ("sql_simplifier_enabled", True),
+        ("eliminate_numeric_sql_value_cast_enabled", True),
+        ("auto_clean_up_temp_table_enabled", True),
+        ("large_query_breakdown_enabled", True),
+        ("large_query_breakdown_complexity_bounds", (20, 30)),
+    ],
+)
+def test_concurrent_update_on_sensitive_configs(session, config, value, caplog):
+    def change_config_value(session_):
+        session_.conf.set(config, value)
+
+    caplog.clear()
+    change_config_value(session)
+    assert (
+        f"You might have more than one threads sharing the Session object trying to update {config}"
+        not in caplog.text
+    )
+
+    with caplog.at_level(logging.WARNING):
+        with ThreadPoolExecutor(max_workers=5) as executor:
+            for _ in range(5):
+                executor.submit(change_config_value, session)
+    assert (
+        f"You might have more than one threads sharing the Session object trying to update {config}"
+        in caplog.text
+    )

tests/unit/compiler/test_large_query_breakdown.py

@@ -105,7 +105,12 @@
     ],
 )
 def test_pipeline_breaker_node(mock_session, mock_analyzer, node_generator, expected):
-    large_query_breakdown = LargeQueryBreakdown(mock_session, mock_analyzer, [])
+    large_query_breakdown = LargeQueryBreakdown(
+        mock_session,
+        mock_analyzer,
+        [],
+        mock_session.large_query_breakdown_complexity_bounds,
+    )
     node = node_generator(mock_analyzer)
 
     assert (

tests/unit/test_server_connection.py

@@ -119,6 +119,7 @@ def test_get_result_set_exception(mock_server_connection):
     fake_session._last_canceled_id = 100
     fake_session._conn = mock_server_connection
     fake_session._cte_optimization_enabled = False
+    fake_session._query_compilation_stage_enabled = False
     fake_plan = SnowflakePlan(
         queries=[Query("fake query 1"), Query("fake query 2")],
         schema_query="fake schema query",