SNOW-1663726 make session config updates thread safe

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py

@@ -312,15 +312,14 @@ def children_plan_nodes(self) -> List[Union["Selectable", "SnowflakePlan"]]:
         else:
             return []
 
-    def replace_repeated_subquery_with_cte(self) -> "SnowflakePlan":
+    def replace_repeated_subquery_with_cte(
+        self, cte_optimization_enabled: bool, query_compilation_stage_enabled: bool
+    ) -> "SnowflakePlan":
         # parameter protection
         # the common subquery elimination will be applied if cte_optimization is not enabled
         # and the new compilation stage is not enabled. When new compilation stage is enabled,
         # the common subquery elimination will be done through the new plan transformation.
-        if (
-            not self.session._cte_optimization_enabled
-            or self.session._query_compilation_stage_enabled
-        ):
+        if not cte_optimization_enabled or query_compilation_stage_enabled:
             return self
 
         # if source_plan or placeholder_query is none, it must be a leaf node,
@@ -355,8 +354,14 @@ def replace_repeated_subquery_with_cte(self) -> "SnowflakePlan":
         # create CTE query
         final_query = create_cte_query(self, duplicate_plan_set)
 
+        with self.session._lock:
+            # copy depends on the cte_optimization_enabled value. We should keep it
+            # consistent with the current context.
+            original_cte_optimization = self.session.cte_optimization_enabled
+            self.session.cte_optimization_enabled = cte_optimization_enabled
+            plan = copy.copy(self)
+            self.session.cte_optimization_enabled = original_cte_optimization
         # all other parts of query are unchanged, but just replace the original query
-        plan = copy.copy(self)
         plan.queries[-1].sql = final_query
         return plan
 
@@ -518,6 +523,56 @@ def add_aliases(self, to_add: Dict) -> None:
         self.expr_to_alias = {**self.expr_to_alias, **to_add}
 
 
+class ConfigContext:
+    """Class to manage reading of configuration values from session in the context of
+    plan building, analysis and resolution.
+
+    Behavior:
+        - Inside an active context, the configuration will be read based on snapshot taken
+        at the context creation stage and reset when the context is exited.
+        - When no active context is present, the configuration will be read from the
+        session object directly.
+    """
+
+    def __init__(self, session) -> None:
+        self.session = session
+        self.configs = {
+            "_query_compilation_stage_enabled",
+            "cte_optimization_enabled",
+            "eliminate_numeric_sql_value_cast_enabled",
+            "large_query_breakdown_complexity_bounds",
+            "large_query_breakdown_enabled",
+        }
+
+    def __getattr__(self, name: str) -> Any:
+        if name in self.configs:
+            return getattr(self.session, name)
+        raise AttributeError(f"ConfigContext has no attribute {name}")
+
+    def __enter__(self) -> "ConfigContext":
+        self._create_snapshot()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self._reset()
+
+    def _create_snapshot(self) -> "ConfigContext":
+        """Reads the configuration attributes from the session object and stores them
+        in the context object.
+        """
+        for name in self.configs:
+            setattr(self, name, getattr(self.session, name))
+        return self
+
+    def _reset(self) -> None:
+        """Removes the configuration attributes from the context object."""
+        for name in self.configs:
+            try:
+                delattr(self, name)
+            except AttributeError:
+                pass
+
+
 class SnowflakePlanBuilder:
     def __init__(
         self,
@@ -531,6 +586,7 @@ def __init__(
         # on the optimized plan. During the final query generation, no schema query is needed,
         # this helps reduces un-necessary overhead for the describing call.
         self._skip_schema_query = skip_schema_query
+        self.config_context: ConfigContext = ConfigContext(session)
 
     @SnowflakePlan.Decorator.wrap_exception
     def build(
@@ -566,7 +622,8 @@ def build(
 
         placeholder_query = (
             sql_generator(select_child._id)
-            if self.session._cte_optimization_enabled and select_child._id is not None
+            if self.config_context.cte_optimization_enabled
+            and select_child._id is not None
             else None
         )
 
@@ -605,7 +662,7 @@ def build_binary(
 
         placeholder_query = (
             sql_generator(select_left._id, select_right._id)
-            if self.session._cte_optimization_enabled
+            if self.config_context.cte_optimization_enabled
             and select_left._id is not None
             and select_right._id is not None
             else None
@@ -637,8 +694,8 @@ def build_binary(
 
         referenced_ctes: Set[str] = set()
         if (
-            self.session.cte_optimization_enabled
-            and self.session._query_compilation_stage_enabled
+            self.config_context.cte_optimization_enabled
+            and self.config_context._query_compilation_stage_enabled
         ):
             # When the cte optimization and the new compilation stage is enabled,
             # the referred cte tables are propagated from left and right can have
@@ -928,7 +985,10 @@ def save_as_table(
                 column_definition_with_hidden_columns,
             )
 
-        child = child.replace_repeated_subquery_with_cte()
+        child = child.replace_repeated_subquery_with_cte(
+            self.config_context.cte_optimization_enabled,
+            self.config_context._query_compilation_stage_enabled,
+        )
 
         def get_create_table_as_select_plan(child: SnowflakePlan, replace, error):
             return self.build(
@@ -1116,7 +1176,10 @@ def create_or_replace_view(
         if not is_sql_select_statement(child.queries[0].sql.lower().strip()):
             raise SnowparkClientExceptionMessages.PLAN_CREATE_VIEWS_FROM_SELECT_ONLY()
 
-        child = child.replace_repeated_subquery_with_cte()
+        child = child.replace_repeated_subquery_with_cte(
+            self.config_context.cte_optimization_enabled,
+            self.config_context._query_compilation_stage_enabled,
+        )
         return self.build(
             lambda x: create_or_replace_view_statement(name, x, is_temp, comment),
             child,
@@ -1159,7 +1222,10 @@ def create_or_replace_dynamic_table(
             # should never reach here
             raise ValueError(f"Unknown create mode: {create_mode}")  # pragma: no cover
 
-        child = child.replace_repeated_subquery_with_cte()
+        child = child.replace_repeated_subquery_with_cte(
+            self.config_context.cte_optimization_enabled,
+            self.config_context._query_compilation_stage_enabled,
+        )
         return self.build(
             lambda x: create_or_replace_dynamic_table_statement(
                 name=name,
@@ -1462,7 +1528,10 @@ def copy_into_location(
         header: bool = False,
         **copy_options: Optional[Any],
     ) -> SnowflakePlan:
-        query = query.replace_repeated_subquery_with_cte()
+        query = query.replace_repeated_subquery_with_cte(
+            self.config_context.cte_optimization_enabled,
+            self.config_context._query_compilation_stage_enabled,
+        )
         return self.build(
             lambda x: copy_into_location(
                 query=x,
@@ -1489,7 +1558,10 @@ def update(
         source_plan: Optional[LogicalPlan],
     ) -> SnowflakePlan:
         if source_data:
-            source_data = source_data.replace_repeated_subquery_with_cte()
+            source_data = source_data.replace_repeated_subquery_with_cte(
+                self.config_context.cte_optimization_enabled,
+                self.config_context._query_compilation_stage_enabled,
+            )
             return self.build(
                 lambda x: update_statement(
                     table_name,
@@ -1520,7 +1592,10 @@ def delete(
         source_plan: Optional[LogicalPlan],
     ) -> SnowflakePlan:
         if source_data:
-            source_data = source_data.replace_repeated_subquery_with_cte()
+            source_data = source_data.replace_repeated_subquery_with_cte(
+                self.config_context.cte_optimization_enabled,
+                self.config_context._query_compilation_stage_enabled,
+            )
             return self.build(
                 lambda x: delete_statement(
                     table_name,
@@ -1549,7 +1624,10 @@ def merge(
         clauses: List[str],
         source_plan: Optional[LogicalPlan],
     ) -> SnowflakePlan:
-        source_data = source_data.replace_repeated_subquery_with_cte()
+        source_data = source_data.replace_repeated_subquery_with_cte(
+            self.config_context.cte_optimization_enabled,
+            self.config_context._query_compilation_stage_enabled,
+        )
         return self.build(
             lambda x: merge_statement(table_name, x, join_expr, clauses),
             source_data,

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

@@ -10,6 +10,7 @@
     get_complexity_score,
 )
 from snowflake.snowpark._internal.analyzer.snowflake_plan import (
+    ConfigContext,
     PlanQueryType,
     Query,
     SnowflakePlan,
@@ -46,6 +47,7 @@ class PlanCompiler:
 
     def __init__(self, plan: SnowflakePlan) -> None:
         self._plan = plan
+        self.config_context = ConfigContext(self._plan.session)
 
     def should_start_query_compilation(self) -> bool:
         """
@@ -65,14 +67,18 @@ 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 self.config_context._query_compilation_stage_enabled
             and (
-                current_session.cte_optimization_enabled
-                or current_session.large_query_breakdown_enabled
+                self.config_context.cte_optimization_enabled
+                or self.config_context.large_query_breakdown_enabled
             )
         )
 
     def compile(self) -> Dict[PlanQueryType, List[Query]]:
+        with self.config_context:
+            return self._compile()
+
+    def _compile(self) -> Dict[PlanQueryType, List[Query]]:
         if self.should_start_query_compilation():
             # preparation for compilation
             # 1. make a copy of the original plan
@@ -89,7 +95,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.config_context.cte_optimization_enabled:
                 repeated_subquery_eliminator = RepeatedSubqueryElimination(
                     logical_plans, query_generator
                 )
@@ -102,9 +108,12 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
             ]
 
             # Large query breakdown
-            if self._plan.session.large_query_breakdown_enabled:
+            if self.config_context.large_query_breakdown_enabled:
                 large_query_breakdown = LargeQueryBreakdown(
-                    self._plan.session, query_generator, logical_plans
+                    self._plan.session,
+                    query_generator,
+                    logical_plans,
+                    self.config_context.large_query_breakdown_complexity_bounds,
                 )
                 logical_plans = large_query_breakdown.apply()
 
@@ -124,9 +133,9 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
             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.config_context.cte_optimization_enabled,
+                TelemetryField.LARGE_QUERY_BREAKDOWN_ENABLED.value: self.config_context.large_query_breakdown_enabled,
+                CompilationStageTelemetryField.COMPLEXITY_SCORE_BOUNDS.value: self.config_context.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,
@@ -143,8 +152,10 @@ 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(
+                self.config_context.cte_optimization_enabled,
+                self.config_context._query_compilation_stage_enabled,
+            )
             return {
                 PlanQueryType.QUERIES: final_plan.queries,
                 PlanQueryType.POST_ACTIONS: final_plan.post_actions,

src/snowflake/snowpark/_internal/server_connection.py

@@ -8,6 +8,7 @@
 import inspect
 import os
 import sys
+import threading
 import time
 from logging import getLogger
 from typing import (
@@ -154,6 +155,8 @@ def __init__(
         options: Dict[str, Union[int, str]],
         conn: Optional[SnowflakeConnection] = None,
     ) -> None:
+        self._lock = threading.RLock()
+        self._thread_stored = threading.local()
         self._lower_case_parameters = {k.lower(): v for k, v in options.items()}
         self._add_application_parameters()
         self._conn = conn if conn else connect(**self._lower_case_parameters)
@@ -170,7 +173,6 @@ def __init__(
 
         if "password" in self._lower_case_parameters:
             self._lower_case_parameters["password"] = None
-        self._cursor = self._conn.cursor()
         self._telemetry_client = TelemetryClient(self._conn)
         self._query_listener: Set[QueryHistory] = set()
         # The session in this case refers to a Snowflake session, not a
@@ -183,6 +185,12 @@ def __init__(
             "_skip_upload_on_content_match" in signature.parameters
         )
 
+    @property
+    def _cursor(self) -> SnowflakeCursor:
+        if not hasattr(self._thread_stored, "cursor"):
+            self._thread_stored.cursor = self._conn.cursor()
+        return self._thread_stored.cursor
+
     def _add_application_parameters(self) -> None:
         if PARAM_APPLICATION not in self._lower_case_parameters:
             # Mirrored from snowflake-connector-python/src/snowflake/connector/connection.py#L295
@@ -210,10 +218,12 @@ def _add_application_parameters(self) -> None:
             ] = get_version()
 
     def add_query_listener(self, listener: QueryHistory) -> None:
-        self._query_listener.add(listener)
+        with self._lock:
+            self._query_listener.add(listener)
 
     def remove_query_listener(self, listener: QueryHistory) -> None:
-        self._query_listener.remove(listener)
+        with self._lock:
+            self._query_listener.remove(listener)
 
     def close(self) -> None:
         if self._conn:
@@ -360,8 +370,9 @@ def upload_stream(
                 raise ex
 
     def notify_query_listeners(self, query_record: QueryRecord) -> None:
-        for listener in self._query_listener:
-            listener._add_query(query_record)
+        with self._lock:
+            for listener in self._query_listener:
+                listener._add_query(query_record)
 
     def execute_and_notify_query_listener(
         self, query: str, **kwargs: Any