SNOW-1663726 make session config updates thread safe

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

@@ -0,0 +1,37 @@
+#
+# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
+#
+
+from typing import Any
+
+
+class ConfigContext:
+    """Class to help reading a snapshot of configuration attributes from a session object.
+
+    On instantiation, this object stores the configuration from the session object
+    and returns the stored configuration attributes when requested.
+    """
+
+    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",
+        }
+        self._create_snapshot()
+
+    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 _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

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

@@ -23,6 +23,7 @@
     Union,
 )
 
+from snowflake.snowpark._internal.analyzer.config_context import ConfigContext
 from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import (
     PlanNodeCategory,
     sum_node_complexities,
@@ -312,14 +313,16 @@ 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, config_context: ConfigContext
+    ) -> "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
+            not config_context.cte_optimization_enabled
+            or config_context._query_compilation_stage_enabled
         ):
             return self
 
@@ -355,8 +358,16 @@ 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 = (
+                config_context.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
 
@@ -531,6 +542,10 @@ 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: Optional[ConfigContext] = None
+
+    def set_config_context(self, config_context: ConfigContext) -> None:
+        self._config_context = config_context
 
     @SnowflakePlan.Decorator.wrap_exception
     def build(
@@ -564,9 +579,10 @@ def build(
             ), "No schema query is available in child SnowflakePlan"
             new_schema_query = schema_query or sql_generator(child.schema_query)
 
+        config_context = self._config_context or ConfigContext(self.session)
         placeholder_query = (
             sql_generator(select_child._id)
-            if self.session._cte_optimization_enabled and select_child._id is not None
+            if config_context.cte_optimization_enabled and select_child._id is not None
             else None
         )
 
@@ -603,9 +619,10 @@ def build_binary(
             right_schema_query = schema_value_statement(select_right.attributes)
             schema_query = sql_generator(left_schema_query, right_schema_query)
 
+        config_context = self._config_context or ConfigContext(self.session)
         placeholder_query = (
             sql_generator(select_left._id, select_right._id)
-            if self.session._cte_optimization_enabled
+            if config_context.cte_optimization_enabled
             and select_left._id is not None
             and select_right._id is not None
             else None
@@ -637,8 +654,8 @@ def build_binary(
 
         referenced_ctes: Set[str] = set()
         if (
-            self.session.cte_optimization_enabled
-            and self.session._query_compilation_stage_enabled
+            config_context.cte_optimization_enabled
+            and 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 +945,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._config_context or ConfigContext(self.session)
+        )
 
         def get_create_table_as_select_plan(child: SnowflakePlan, replace, error):
             return self.build(
@@ -1116,7 +1135,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._config_context or ConfigContext(self.session)
+        )
         return self.build(
             lambda x: create_or_replace_view_statement(name, x, is_temp, comment),
             child,
@@ -1159,7 +1180,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._config_context or ConfigContext(self.session)
+        )
         return self.build(
             lambda x: create_or_replace_dynamic_table_statement(
                 name=name,
@@ -1462,7 +1485,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._config_context or ConfigContext(self.session)
+        )
         return self.build(
             lambda x: copy_into_location(
                 query=x,
@@ -1489,7 +1514,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._config_context or ConfigContext(self.session)
+            )
             return self.build(
                 lambda x: update_statement(
                     table_name,
@@ -1520,7 +1547,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._config_context or ConfigContext(self.session)
+            )
             return self.build(
                 lambda x: delete_statement(
                     table_name,
@@ -1549,7 +1578,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._config_context or ConfigContext(self.session)
+        )
         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

@@ -6,6 +6,7 @@
 import time
 from typing import Dict, List
 
+from snowflake.snowpark._internal.analyzer.config_context import ConfigContext
 from snowflake.snowpark._internal.analyzer.query_plan_analysis_utils import (
     get_complexity_score,
 )
@@ -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,10 +67,10 @@ 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
             )
         )
 
@@ -84,12 +86,12 @@ def compile(self) -> Dict[PlanQueryType, List[Query]]:
             deep_copy_end_time = time.time()
 
             # 2. create a code generator with the original plan
-            query_generator = create_query_generator(self._plan)
+            query_generator = create_query_generator(self._plan, self._config_context)
 
             # 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 +104,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 +129,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 +148,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._config_context
+            )
             return {
                 PlanQueryType.QUERIES: final_plan.queries,
                 PlanQueryType.POST_ACTIONS: final_plan.post_actions,

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

@@ -8,6 +8,7 @@
 from snowflake.snowpark._internal.analyzer.expression import Attribute
 from snowflake.snowpark._internal.analyzer.select_statement import Selectable
 from snowflake.snowpark._internal.analyzer.snowflake_plan import (
+    ConfigContext,
     CreateViewCommand,
     PlanQueryType,
     Query,
@@ -49,6 +50,7 @@ class QueryGenerator(Analyzer):
     def __init__(
         self,
         session: Session,
+        config_context: ConfigContext,
         snowflake_create_table_plan_info: Optional[SnowflakeCreateTablePlanInfo] = None,
     ) -> None:
         super().__init__(session)
@@ -58,6 +60,7 @@ def __init__(
         self._snowflake_create_table_plan_info: Optional[
             SnowflakeCreateTablePlanInfo
         ] = snowflake_create_table_plan_info
+        self._config_context = config_context
         # Records the definition of all the with query blocks encountered during the code generation.
         # This information will be used to generate the final query of a SnowflakePlan with the
         # correct CTE definition.
@@ -105,18 +108,22 @@ def generate_queries(
             PlanQueryType.POST_ACTIONS: post_actions,
         }
 
+    def resolve(self, logical_plan: LogicalPlan) -> SnowflakePlan:
+        return super().resolve(logical_plan, self._config_context)
+
     def do_resolve_with_resolved_children(
         self,
         logical_plan: LogicalPlan,
         resolved_children: Dict[LogicalPlan, SnowflakePlan],
         df_aliased_col_name_to_real_col_name: DefaultDict[str, Dict[str, str]],
+        config_context: ConfigContext,
     ) -> SnowflakePlan:
         if isinstance(logical_plan, SnowflakePlan):
             if logical_plan.queries is None:
                 assert logical_plan.source_plan is not None
                 # when encounter a SnowflakePlan with no queries, try to re-resolve
                 # the source plan to construct the result
-                res = self.do_resolve(logical_plan.source_plan)
+                res = self.do_resolve(logical_plan.source_plan, config_context)
                 resolved_children[logical_plan] = res
                 resolved_plan = res
             else:
@@ -204,7 +211,10 @@ def do_resolve_with_resolved_children(
             copied_resolved_child.queries = final_queries[PlanQueryType.QUERIES]
             resolved_children[logical_plan.children[0]] = copied_resolved_child
             resolved_plan = super().do_resolve_with_resolved_children(
-                logical_plan, resolved_children, df_aliased_col_name_to_real_col_name
+                logical_plan,
+                resolved_children,
+                df_aliased_col_name_to_real_col_name,
+                config_context,
             )
 
         elif isinstance(logical_plan, Selectable):
@@ -228,7 +238,10 @@ def do_resolve_with_resolved_children(
 
         else:
             resolved_plan = super().do_resolve_with_resolved_children(
-                logical_plan, resolved_children, df_aliased_col_name_to_real_col_name
+                logical_plan,
+                resolved_children,
+                df_aliased_col_name_to_real_col_name,
+                config_context,
             )
 
         resolved_plan._is_valid_for_replacement = True