SNOW-1663726 make temp table cleaner 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

@@ -49,6 +49,13 @@ class PlanCompiler:
 
     def __init__(self, plan: SnowflakePlan) -> None:
         self._plan = plan
+        session = plan.session
+        self.cte_optimization_enabled = session.cte_optimization_enabled
+        self.large_query_breakdown_enabled = session.large_query_breakdown_enabled
+        self.large_query_breakdown_complexity_bounds = (
+            session.large_query_breakdown_complexity_bounds
+        )
+        self.query_compilation_stage_enabled = session._query_compilation_stage_enabled
 
     def should_start_query_compilation(self) -> bool:
         """
@@ -68,15 +75,13 @@ def should_start_query_compilation(self) -> bool:
         return (
             not isinstance(current_session._conn, MockServerConnection)
             and (self._plan.source_plan is not None)
-            and current_session._query_compilation_stage_enabled
-            and (
-                current_session.cte_optimization_enabled
-                or current_session.large_query_breakdown_enabled
-            )
+            and self.query_compilation_stage_enabled
+            and (self.cte_optimization_enabled or self.large_query_breakdown_enabled)
         )
 
     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 +98,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 self.cte_optimization_enabled:
                 repeated_subquery_eliminator = RepeatedSubqueryElimination(
                     logical_plans, query_generator
                 )
@@ -108,9 +113,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 self.large_query_breakdown_enabled:
                 large_query_breakdown = LargeQueryBreakdown(
-                    self._plan.session, query_generator, logical_plans
+                    session,
+                    query_generator,
+                    logical_plans,
+                    self.large_query_breakdown_complexity_bounds,
                 )
                 logical_plans = large_query_breakdown.apply()
 
@@ -130,11 +138,10 @@ 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,
-                CompilationStageTelemetryField.COMPLEXITY_SCORE_BOUNDS.value: session.large_query_breakdown_complexity_bounds,
+                TelemetryField.CTE_OPTIMIZATION_ENABLED.value: self.cte_optimization_enabled,
+                TelemetryField.LARGE_QUERY_BREAKDOWN_ENABLED.value: self.large_query_breakdown_enabled,
+                CompilationStageTelemetryField.COMPLEXITY_SCORE_BOUNDS.value: self.large_query_breakdown_complexity_bounds,
                 CompilationStageTelemetryField.TIME_TAKEN_FOR_COMPILATION.value: total_time,
                 CompilationStageTelemetryField.TIME_TAKEN_FOR_DEEP_COPY_PLAN.value: deep_copy_time,
                 CompilationStageTelemetryField.TIME_TAKEN_FOR_CTE_OPTIMIZATION.value: cte_time,
@@ -151,8 +158,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/temp_table_auto_cleaner.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
 #
 import logging
+import threading
 import weakref
 from collections import defaultdict
 from typing import TYPE_CHECKING, Dict
@@ -31,9 +32,12 @@ def __init__(self, session: "Session") -> None:
         # to its reference count for later temp table management
         # this dict will still be maintained even if the cleaner is stopped (`stop()` is called)
         self.ref_count_map: Dict[str, int] = defaultdict(int)
+        # Lock to protect the ref_count_map
+        self.lock = threading.RLock()
 
     def add(self, table: SnowflakeTable) -> None:
-        self.ref_count_map[table.name] += 1
+        with self.lock:
+            self.ref_count_map[table.name] += 1
         # the finalizer will be triggered when it gets garbage collected
         # and this table will be dropped finally
         _ = weakref.finalize(table, self._delete_ref_count, table.name)
@@ -43,13 +47,15 @@ def _delete_ref_count(self, name: str) -> None:  # pragma: no cover
         Decrements the reference count of a temporary table,
         and if the count reaches zero, puts this table in the queue for cleanup.
         """
-        self.ref_count_map[name] -= 1
-        if self.ref_count_map[name] == 0:
+        with self.lock:
+            self.ref_count_map[name] -= 1
+            current_ref_count = self.ref_count_map[name]
+        if current_ref_count == 0:
             if self.session.auto_clean_up_temp_table_enabled:
                 self.drop_table(name)
-        elif self.ref_count_map[name] < 0:
+        elif current_ref_count < 0:
             logging.debug(
-                f"Unexpected reference count {self.ref_count_map[name]} for table {name}"
+                f"Unexpected reference count {current_ref_count} for table {name}"
             )
 
     def drop_table(self, name: str) -> None:  # pragma: no cover
@@ -89,9 +95,11 @@ def stop(self) -> None:
 
     @property
     def num_temp_tables_created(self) -> int:
-        return len(self.ref_count_map)
+        with self.lock:
+            return len(self.ref_count_map)
 
     @property
     def num_temp_tables_cleaned(self) -> int:
         # TODO SNOW-1662536: we may need a separate counter for the number of tables cleaned when parameter is enabled
-        return sum(v == 0 for v in self.ref_count_map.values())
+        with self.lock:
+            return sum(v == 0 for v in self.ref_count_map.values())

src/snowflake/snowpark/session.py

@@ -339,38 +339,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:
         """
@@ -538,11 +544,6 @@ def __init__(
         self._udtf_registration = UDTFRegistration(self)
         self._udaf_registration = UDAFRegistration(self)
 
-        self._plan_builder = (
-            SnowflakePlanBuilder(self)
-            if isinstance(self._conn, ServerConnection)
-            else MockSnowflakePlanBuilder(self)
-        )
         self._last_action_id = 0
         self._last_canceled_id = 0
         self._use_scoped_temp_objects: bool = (
@@ -637,6 +638,16 @@ def _analyzer(self) -> Analyzer:
             )
         return self._thread_store.analyzer
 
+    @property
+    def _plan_builder(self):
+        if not hasattr(self._thread_store, "plan_builder"):
+            self._thread_store.plan_builder = (
+                SnowflakePlanBuilder(self)
+                if isinstance(self._conn, ServerConnection)
+                else MockSnowflakePlanBuilder(self)
+            )
+        return self._thread_store.plan_builder
+
     def close(self) -> None:
         """Close this session."""
         if is_in_stored_procedure():
@@ -770,36 +781,45 @@ 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}"
+        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
+        if threading.active_count() > 1:
+            # TODO (SNOW-1541096): Remove the limitation once old cte implementation is removed.
+            _logger.warning(
+                "Setting cte_optimization_enabled is not currently thread-safe. Ignoring the update"
             )
-        self._cte_optimization_enabled = value
+            return
+        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"""
 
         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!"
@@ -810,10 +830,11 @@ def eliminate_numeric_sql_value_cast_enabled(self, value: bool) -> None:
     def auto_clean_up_temp_table_enabled(self, value: bool) -> None:
         """Set the value for 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
-            )
-            self._auto_clean_up_temp_table_enabled = value
+            with self._lock:
+                self._conn._telemetry_client.send_auto_clean_up_temp_table_telemetry(
+                    self._session_id, value
+                )
+                self._auto_clean_up_temp_table_enabled = value
         else:
             raise ValueError(
                 "value for auto_clean_up_temp_table_enabled must be True or False!"
@@ -829,10 +850,11 @@ def large_query_breakdown_enabled(self, value: bool) -> None:
         """
 
         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!"
@@ -850,16 +872,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 +1484,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():