SNOW-1663726 make session config updates thread safe

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

@@ -166,6 +166,13 @@ def __init__(self, session: "snowflake.snowpark.session.Session") -> None:
         self.generated_alias_maps = {}
         self.subquery_plans = []
         self.alias_maps_to_use: Optional[Dict[uuid.UUID, str]] = None
+        self._eliminate_numeric_sql_value_cast_enabled: Optional[bool] = None
+
+    @property
+    def eliminate_numeric_sql_value_cast_enabled(self) -> bool:
+        if self._eliminate_numeric_sql_value_cast_enabled is None:
+            return self.session.eliminate_numeric_sql_value_cast_enabled
+        return self._eliminate_numeric_sql_value_cast_enabled
 
     def analyze(
         self,
@@ -264,7 +271,7 @@ def analyze(
         if isinstance(expr, MultipleExpression):
             block_expressions = []
             for expression in expr.expressions:
-                if self.session.eliminate_numeric_sql_value_cast_enabled:
+                if self.eliminate_numeric_sql_value_cast_enabled:
                     resolved_expr = self.to_sql_try_avoid_cast(
                         expression,
                         df_aliased_col_name_to_real_col_name,
@@ -283,7 +290,7 @@ def analyze(
         if isinstance(expr, InExpression):
             in_values = []
             for expression in expr.values:
-                if self.session.eliminate_numeric_sql_value_cast_enabled:
+                if self.eliminate_numeric_sql_value_cast_enabled:
                     in_value = self.to_sql_try_avoid_cast(
                         expression,
                         df_aliased_col_name_to_real_col_name,
@@ -678,7 +685,7 @@ def binary_operator_extractor(
         df_aliased_col_name_to_real_col_name,
         parse_local_name=False,
     ) -> str:
-        if self.session.eliminate_numeric_sql_value_cast_enabled:
+        if self.eliminate_numeric_sql_value_cast_enabled:
             left_sql_expr = self.to_sql_try_avoid_cast(
                 expr.left, df_aliased_col_name_to_real_col_name, parse_local_name
             )
@@ -760,13 +767,29 @@ def resolve(self, logical_plan: LogicalPlan) -> SnowflakePlan:
         self.subquery_plans = []
         self.generated_alias_maps = {}
 
+        # To ensure that the context remain unchanged during resolving the plan, we
+        # read these values at the beginning and reset them at the end.
+        self.plan_builder._cte_optimization_enabled = (
+            self.session.cte_optimization_enabled
+        )
+        self.plan_builder._query_compilation_stage_enabled = (
+            self.session._query_compilation_stage_enabled
+        )
+        self._eliminate_numeric_sql_value_cast_enabled = (
+            self.session.eliminate_numeric_sql_value_cast_enabled
+        )
+
         result = self.do_resolve(logical_plan)
 
         result.add_aliases(self.generated_alias_maps)
 
         if self.subquery_plans:
             result = result.with_subqueries(self.subquery_plans)
 
+        self.plan_builder._cte_optimization_enabled = None
+        self.plan_builder._query_compilation_stage_enabled = None
+        self._eliminate_numeric_sql_value_cast_enabled = None
+
         return result
 
     def do_resolve(self, logical_plan: LogicalPlan) -> SnowflakePlan:

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,
@@ -452,7 +451,7 @@ def cumulative_node_complexity(self, value: Dict[PlanNodeCategory, int]):
         self._cumulative_node_complexity = value
 
     def __copy__(self) -> "SnowflakePlan":
-        if self.session._cte_optimization_enabled:
+        if self.cte_optimization_enabled:
             return SnowflakePlan(
                 copy.deepcopy(self.queries) if self.queries else [],
                 self.schema_query,
@@ -531,6 +530,21 @@ 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
+        # TODO: describe
+        self._cte_optimization_enabled: Optional[bool] = None
+        self._query_compilation_stage_enabled: Optional[bool] = None
+
+    @property
+    def cte_optimization_enabled(self) -> bool:
+        if self._cte_optimization_enabled is None:
+            return self.session.cte_optimization_enabled
+        return self._cte_optimization_enabled
+
+    @property
+    def query_compilation_stage_enabled(self) -> bool:
+        if self._query_compilation_stage_enabled is None:
+            return self.session._query_compilation_stage_enabled
+        return self._query_compilation_stage_enabled
 
     @SnowflakePlan.Decorator.wrap_exception
     def build(
@@ -566,7 +580,7 @@ def build(
 
         placeholder_query = (
             sql_generator(select_child._id)
-            if self.session._cte_optimization_enabled and select_child._id is not None
+            if self.cte_optimization_enabled and select_child._id is not None
             else None
         )
 
@@ -605,7 +619,7 @@ def build_binary(
 
         placeholder_query = (
             sql_generator(select_left._id, select_right._id)
-            if self.session._cte_optimization_enabled
+            if self.cte_optimization_enabled
             and select_left._id is not None
             and select_right._id is not None
             else None
@@ -636,10 +650,7 @@ def build_binary(
                 post_actions.append(copy.copy(post_action))
 
         referenced_ctes: Set[str] = set()
-        if (
-            self.session.cte_optimization_enabled
-            and self.session._query_compilation_stage_enabled
-        ):
+        if self.cte_optimization_enabled and self.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
             # duplicated queries if there is a common CTE block referenced by
@@ -928,7 +939,9 @@ def save_as_table(
                 column_definition_with_hidden_columns,
             )
 
-        child = child.replace_repeated_subquery_with_cte()
+        child = child.replace_repeated_subquery_with_cte(
+            self.cte_optimization_enabled, self.query_compilation_stage_enabled
+        )
 
         def get_create_table_as_select_plan(child: SnowflakePlan, replace, error):
             return self.build(
@@ -1116,7 +1129,9 @@ 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.cte_optimization_enabled, self.query_compilation_stage_enabled
+        )
         return self.build(
             lambda x: create_or_replace_view_statement(name, x, is_temp, comment),
             child,
@@ -1159,7 +1174,9 @@ 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.cte_optimization_enabled, self.query_compilation_stage_enabled
+        )
         return self.build(
             lambda x: create_or_replace_dynamic_table_statement(
                 name=name,
@@ -1462,7 +1479,9 @@ 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.cte_optimization_enabled, self.query_compilation_stage_enabled
+        )
         return self.build(
             lambda x: copy_into_location(
                 query=x,
@@ -1489,7 +1508,9 @@ 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.cte_optimization_enabled, self.query_compilation_stage_enabled
+            )
             return self.build(
                 lambda x: update_statement(
                     table_name,
@@ -1520,7 +1541,9 @@ 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.cte_optimization_enabled, self.query_compilation_stage_enabled
+            )
             return self.build(
                 lambda x: delete_statement(
                     table_name,
@@ -1549,7 +1572,9 @@ 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.cte_optimization_enabled, self.query_compilation_stage_enabled
+        )
         return self.build(
             lambda x: merge_statement(table_name, x, join_expr, clauses),
             source_data,

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

@@ -48,6 +48,14 @@ class PlanCompiler:
 
     def __init__(self, plan: SnowflakePlan) -> None:
         self._plan = plan
+        current_session = self._plan.session
+        self.cte_optimization_enabled = current_session.cte_optimization_enabled
+        self.large_query_breakdown_enabled = (
+            current_session.large_query_breakdown_enabled
+        )
+        self.query_compilation_stage_enabled = (
+            current_session._query_compilation_stage_enabled
+        )
 
     def should_start_query_compilation(self) -> bool:
         """
@@ -67,11 +75,8 @@ 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]]:
@@ -91,7 +96,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
                 )
@@ -104,7 +109,7 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
             ]
 
             # 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
                 )
@@ -126,8 +131,8 @@ 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,
+                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: (
                     COMPLEXITY_SCORE_LOWER_BOUND,
                     COMPLEXITY_SCORE_UPPER_BOUND,
@@ -148,8 +153,9 @@ 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.cte_optimization_enabled, self.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