From 7c22750230884db4bd3e2df448bf3736f4abd7c5 Mon Sep 17 00:00:00 2001
From: Afroz Alam <afroz.alam@snowflake.com>
Date: Mon, 14 Oct 2024 14:31:57 -0700
Subject: [PATCH] SNOW-1418523 Make Session thread safe (#2312)

---
 .github/pull_request_template.md              |   3 +-
 .github/workflows/threadsafe-check.yml        |  36 ++
 CHANGELOG.md                                  |   5 +-
 .../compiler/large_query_breakdown.py         |   9 +-
 .../_internal/compiler/plan_compiler.py       |  14 +-
 .../snowpark/_internal/server_connection.py   |  63 +-
 src/snowflake/snowpark/_internal/telemetry.py |  18 +
 .../_internal/temp_table_auto_cleaner.py      |  22 +-
 src/snowflake/snowpark/_internal/udf_utils.py |  12 +-
 src/snowflake/snowpark/_internal/utils.py     |  10 +
 src/snowflake/snowpark/dataframe.py           |   5 +-
 src/snowflake/snowpark/session.py             | 356 ++++++-----
 src/snowflake/snowpark/stored_procedure.py    |  11 +-
 src/snowflake/snowpark/udaf.py                |   9 +-
 src/snowflake/snowpark/udf.py                 |   9 +-
 src/snowflake/snowpark/udtf.py                |   9 +-
 tests/integ/conftest.py                       |  23 +-
 .../integ/scala/test_file_operation_suite.py  |  17 +-
 tests/integ/test_multithreading.py            | 582 ++++++++++++++++++
 tests/integ/test_telemetry.py                 |  26 +
 tests/integ/test_temp_table_cleanup.py        |  13 +-
 .../compiler/test_large_query_breakdown.py    |   7 +-
 tests/unit/test_server_connection.py          |   1 +
 tests/unit/test_udf_utils.py                  |   2 +
 24 files changed, 1033 insertions(+), 229 deletions(-)
 create mode 100644 .github/workflows/threadsafe-check.yml
 create mode 100644 tests/integ/test_multithreading.py

diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
index f8741e619c3..0f191f6b06c 100644
--- a/.github/pull_request_template.md
+++ b/.github/pull_request_template.md
@@ -6,7 +6,7 @@ Please answer these questions before creating your pull request. Thanks!
 
    <!---
    In this section, please add a Snowflake Jira issue number.
-   
+
    Note that if a corresponding GitHub issue exists, you should still include
    the Snowflake Jira issue number. For example, for GitHub issue
    https://github.com/snowflakedb/snowpark-python/issues/1400, you should
@@ -24,6 +24,7 @@ Please answer these questions before creating your pull request. Thanks!
    - [ ] I am adding new credentials
    - [ ] I am adding a new dependency
    - [ ] If this is a new feature/behavior, I'm adding the Local Testing parity changes.
+   - [ ] I acknowledge that I have ensured my changes to be thread-safe. Follow the link for more information: [Thread-safe Developer Guidelines](https://docs.google.com/document/d/162d_i4zZ2AfcGRXojj0jByt8EUq-DrSHPPnTa4QvwbA/edit#bookmark=id.e82u4nekq80k)
 
 3. Please describe how your code solves the related issue.
 
diff --git a/.github/workflows/threadsafe-check.yml b/.github/workflows/threadsafe-check.yml
new file mode 100644
index 00000000000..0e90abcb7ef
--- /dev/null
+++ b/.github/workflows/threadsafe-check.yml
@@ -0,0 +1,36 @@
+name: Threadsafe Check
+
+on:
+  pull_request:
+    types: [opened, synchronize, labeled, unlabeled, edited]
+
+jobs:
+    check_threadsafety:
+        runs-on: ubuntu-latest
+        steps:
+            - name: Checkout
+              uses: actions/checkout@v4
+              with:
+                  fetch-depth: 0
+
+            - name: Check for modified files
+              id: changed_files
+              run: |
+                echo "changed_files=$(git diff --name-only --diff-filter=ACMRT ${{ github.event.pull_request.base.sha }} ${{ github.sha }} | xargs)" >> $GITHUB_OUTPUT
+
+            - name: Verify threadsafety acknowledgement
+              run: |
+                CHANGED_FILES="${{ steps.changed_files.outputs.changed_files }}"
+                # Check if changed files are in snowpark/_internal, snowpark/mock, or snowpark/*.py. We exclude snowpark/modin in this check.
+                if echo "$CHANGED_FILES" | grep -qE '(src/snowflake/snowpark/_internal|src/snowflake/snowpark/mock|src/snowflake/snowpark/[^/]+\.py)'; then
+                    echo "Checking PR description for thread-safety acknowledgment..."
+                    if [[ "${{ github.event.pull_request.body }}" != *"[x] I acknowledge that I have ensured my changes to be thread-safe"* ]]; then
+                        echo "Thread-safety acknowledgment not found in PR description."
+                        echo "Please acknowledge the threadsafety implications of your changes by adding '[x] I acknowledge that I have ensured my changes to be thread-safe' to the PR description."
+                        exit 1
+                    else
+                        echo "Thread-safety acknowledgment found in PR description."
+                    fi
+                else
+                    echo "No critical files modified; skipping threadsafety check."
+                fi
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e9a034657ef..270d14ddc8e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,11 +4,14 @@
 
 ### Snowpark Python API Updates
 
+- Updated `Session` class to be thread-safe. This allows concurrent dataframe transformations, dataframe actions, UDF and store procedure registration, and concurrent file uploads.
+
 #### New Features
 
 ### Snowpark pandas API Updates
 
 #### New Features
+
 - Added numpy compatibility support for `np.float_power`, `np.mod`, `np.remainder`, `np.greater`, `np.greater_equal`, `np.less`, `np.less_equal`, `np.not_equal`, and `np.equal`.
 
 ## 1.23.0 (2024-10-09)
@@ -23,7 +26,7 @@
 - Added support for file writes. This feature is currently in private preview.
 - Added `thread_id` to `QueryRecord` to track the thread id submitting the query history.
 - Added support for `Session.stored_procedure_profiler`.
-- Added support for 'Service' domain to `session.lineage.trace` API. 
+- Added support for 'Service' domain to `session.lineage.trace` API.
 
 #### Improvements
 
diff --git a/src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py b/src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py
index 038b2e57714..98a989b2ac5 100644
--- a/src/snowflake/snowpark/_internal/compiler/large_query_breakdown.py
+++ b/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):
diff --git a/src/snowflake/snowpark/_internal/compiler/plan_compiler.py b/src/snowflake/snowpark/_internal/compiler/plan_compiler.py
index 66aaafc4454..db085420524 100644
--- a/src/snowflake/snowpark/_internal/compiler/plan_compiler.py
+++ b/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
                 )
@@ -112,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 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()
 
@@ -133,7 +137,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,
@@ -156,8 +159,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,
diff --git a/src/snowflake/snowpark/_internal/server_connection.py b/src/snowflake/snowpark/_internal/server_connection.py
index 864c88fd3bf..1b57697ad6d 100644
--- a/src/snowflake/snowpark/_internal/server_connection.py
+++ b/src/snowflake/snowpark/_internal/server_connection.py
@@ -155,6 +155,8 @@ def __init__(
         options: Dict[str, Union[int, str]],
         conn: Optional[SnowflakeConnection] = None,
     ) -> None:
+        self._lock = threading.RLock()
+        self._thread_store = 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)
@@ -171,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
@@ -184,6 +185,15 @@ def __init__(
             "_skip_upload_on_content_match" in signature.parameters
         )
 
+    @property
+    def _cursor(self) -> SnowflakeCursor:
+        if not hasattr(self._thread_store, "cursor"):
+            self._thread_store.cursor = self._conn.cursor()
+            self._telemetry_client.send_cursor_created_telemetry(
+                self.get_session_id(), threading.get_ident()
+            )
+        return self._thread_store.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
@@ -211,10 +221,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:
@@ -253,17 +265,21 @@ def _run_new_describe(
     ) -> Union[List[ResultMetadata], List["ResultMetadataV2"]]:
         result_metadata = run_new_describe(cursor, query)
 
-        for listener in filter(
-            lambda listener: hasattr(listener, "include_describe")
-            and listener.include_describe,
-            self._query_listener,
-        ):
-            query_record = QueryRecord(cursor.sfqid, query, True)
-            if getattr(listener, "include_thread_id", False):
+        with self._lock:
+            for listener in filter(
+                lambda listener: hasattr(listener, "include_describe")
+                and listener.include_describe,
+                self._query_listener,
+            ):
+                thread_id = (
+                    threading.get_ident()
+                    if getattr(listener, "include_thread_id", False)
+                    else None
+                )
                 query_record = QueryRecord(
-                    cursor.sfqid, query, True, threading.get_ident()
+                    cursor.sfqid, query, True, thread_id=thread_id
                 )
-            listener._add_query(query_record)
+                listener._add_query(query_record)
 
         return result_metadata
 
@@ -380,17 +396,18 @@ def upload_stream(
                 raise ex
 
     def notify_query_listeners(self, query_record: QueryRecord) -> None:
-        for listener in self._query_listener:
-            if getattr(listener, "include_thread_id", False):
-                new_record = QueryRecord(
-                    query_record.query_id,
-                    query_record.sql_text,
-                    query_record.is_describe,
-                    thread_id=threading.get_ident(),
-                )
-                listener._add_query(new_record)
-            else:
-                listener._add_query(query_record)
+        with self._lock:
+            for listener in self._query_listener:
+                if getattr(listener, "include_thread_id", False):
+                    new_record = QueryRecord(
+                        query_record.query_id,
+                        query_record.sql_text,
+                        query_record.is_describe,
+                        thread_id=threading.get_ident(),
+                    )
+                    listener._add_query(new_record)
+                else:
+                    listener._add_query(query_record)
 
     def execute_and_notify_query_listener(
         self, query: str, **kwargs: Any
diff --git a/src/snowflake/snowpark/_internal/telemetry.py b/src/snowflake/snowpark/_internal/telemetry.py
index 19ab8f06525..62813038133 100644
--- a/src/snowflake/snowpark/_internal/telemetry.py
+++ b/src/snowflake/snowpark/_internal/telemetry.py
@@ -4,6 +4,7 @@
 #
 
 import functools
+import threading
 from enum import Enum, unique
 from typing import Any, Dict, List, Optional
 
@@ -38,6 +39,7 @@ class TelemetryField(Enum):
     TYPE_PERFORMANCE_DATA = "snowpark_performance_data"
     TYPE_FUNCTION_USAGE = "snowpark_function_usage"
     TYPE_SESSION_CREATED = "snowpark_session_created"
+    TYPE_CURSOR_CREATED = "snowpark_cursor_created"
     TYPE_SQL_SIMPLIFIER_ENABLED = "snowpark_sql_simplifier_enabled"
     TYPE_CTE_OPTIMIZATION_ENABLED = "snowpark_cte_optimization_enabled"
     # telemetry for optimization that eliminates the extra cast expression generated for expressions
@@ -90,6 +92,8 @@ class TelemetryField(Enum):
     TEMP_TABLE_CLEANUP_ABNORMAL_EXCEPTION_MESSAGE = (
         "temp_table_cleanup_abnormal_exception_message"
     )
+    # multi-threading
+    THREAD_IDENTIFIER = "thread_ident"
 
 
 # These DataFrame APIs call other DataFrame APIs
@@ -196,6 +200,7 @@ def wrap(*args, **kwargs):
                 key.value: value
                 for key, value in plan.cumulative_node_complexity.items()
             }
+            api_calls[0][TelemetryField.THREAD_IDENTIFIER.value] = threading.get_ident()
         except Exception:
             pass
         args[0]._session._conn._telemetry_client.send_function_usage_telemetry(
@@ -343,6 +348,7 @@ def send_upload_file_perf_telemetry(
                 TelemetryField.KEY_CATEGORY.value: TelemetryField.PERF_CAT_UPLOAD_FILE.value,
                 TelemetryField.KEY_FUNC_NAME.value: func_name,
                 TelemetryField.KEY_DURATION.value: duration,
+                TelemetryField.THREAD_IDENTIFIER.value: threading.get_ident(),
             },
         }
         self.send(message)
@@ -537,3 +543,15 @@ def send_large_query_breakdown_update_complexity_bounds(
             },
         }
         self.send(message)
+
+    def send_cursor_created_telemetry(self, session_id: int, thread_id: int):
+        message = {
+            **self._create_basic_telemetry_data(
+                TelemetryField.TYPE_CURSOR_CREATED.value
+            ),
+            TelemetryField.KEY_DATA.value: {
+                TelemetryField.SESSION_ID.value: session_id,
+                TelemetryField.THREAD_IDENTIFIER.value: thread_id,
+            },
+        }
+        self.send(message)
diff --git a/src/snowflake/snowpark/_internal/temp_table_auto_cleaner.py b/src/snowflake/snowpark/_internal/temp_table_auto_cleaner.py
index 6f655cd66e7..69453d48596 100644
--- a/src/snowflake/snowpark/_internal/temp_table_auto_cleaner.py
+++ b/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,8 +47,10 @@ 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
                 # if the session is already closed before garbage collection,
@@ -52,9 +58,9 @@ def _delete_ref_count(self, name: str) -> None:  # pragma: no cover
                 and not self.session._conn.is_closed()
             ):
                 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
@@ -95,9 +101,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())
diff --git a/src/snowflake/snowpark/_internal/udf_utils.py b/src/snowflake/snowpark/_internal/udf_utils.py
index cd8a6f34536..e5a437925d2 100644
--- a/src/snowflake/snowpark/_internal/udf_utils.py
+++ b/src/snowflake/snowpark/_internal/udf_utils.py
@@ -980,8 +980,10 @@ def add_snowpark_package_to_sproc_packages(
     if packages is None:
         if session is None:
             packages = [this_package]
-        elif package_name not in session._packages:
-            packages = list(session._packages.values()) + [this_package]
+        else:
+            with session._package_lock:
+                if package_name not in session._packages:
+                    packages = list(session._packages.values()) + [this_package]
     else:
         package_names = [p if isinstance(p, str) else p.__name__ for p in packages]
         if not any(p.startswith(package_name) for p in package_names):
@@ -1234,11 +1236,9 @@ def create_python_udf_or_sp(
     statement_params: Optional[Dict[str, str]] = None,
     comment: Optional[str] = None,
     native_app_params: Optional[Dict[str, Any]] = None,
+    runtime_version: Optional[str] = None,
 ) -> None:
-    if session is not None and session._runtime_version_from_requirement:
-        runtime_version = session._runtime_version_from_requirement
-    else:
-        runtime_version = f"{sys.version_info[0]}.{sys.version_info[1]}"
+    runtime_version = runtime_version or f"{sys.version_info[0]}.{sys.version_info[1]}"
 
     if replace and if_not_exists:
         raise ValueError("options replace and if_not_exists are incompatible")
diff --git a/src/snowflake/snowpark/_internal/utils.py b/src/snowflake/snowpark/_internal/utils.py
index aff2c6d7ad1..8783faa39d6 100644
--- a/src/snowflake/snowpark/_internal/utils.py
+++ b/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)
 
diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py
index 25ce987fd78..20cf07334f5 100644
--- a/src/snowflake/snowpark/dataframe.py
+++ b/src/snowflake/snowpark/dataframe.py
@@ -2027,12 +2027,13 @@ def _union_by_name_internal(
         ]
 
         names = right_project_list + not_found_attrs
-        if self._session.sql_simplifier_enabled and other._select_statement:
+        sql_simplifier_enabled = self._session.sql_simplifier_enabled
+        if sql_simplifier_enabled and other._select_statement:
             right_child = self._with_plan(other._select_statement.select(names))
         else:
             right_child = self._with_plan(Project(names, other._plan))
 
-        if self._session.sql_simplifier_enabled:
+        if sql_simplifier_enabled:
             df = self._with_plan(
                 self._select_statement.set_operator(
                     right_child._select_statement
diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py
index 779b849e927..b1920eb7e8d 100644
--- a/src/snowflake/snowpark/session.py
+++ b/src/snowflake/snowpark/session.py
@@ -13,6 +13,7 @@
 import re
 import sys
 import tempfile
+import threading
 import warnings
 from array import array
 from functools import reduce
@@ -114,6 +115,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,
 )
@@ -340,38 +342,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:
         """
@@ -509,6 +517,13 @@ def __init__(
         if len(_active_sessions) >= 1 and is_in_stored_procedure():
             raise SnowparkClientExceptionMessages.DONT_CREATE_SESSION_IN_SP()
         self._conn = conn
+        self._thread_store = threading.local()
+        self._lock = threading.RLock()
+
+        # this lock is used to protect _packages. We use introduce a new lock because add_packages
+        # launches a query to snowflake to get all version of packages available in snowflake. This
+        # query can be slow and prevent other threads from moving on waiting for _lock.
+        self._package_lock = threading.RLock()
         self._query_tag = None
         self._import_paths: Dict[str, Tuple[Optional[str], Optional[str]]] = {}
         self._packages: Dict[str, str] = {}
@@ -547,9 +562,6 @@ def __init__(
         )
         self._file = FileOperation(self)
         self._lineage = Lineage(self)
-        self._analyzer = (
-            Analyzer(self) if isinstance(conn, ServerConnection) else MockAnalyzer(self)
-        )
         self._sql_simplifier_enabled: bool = (
             self._conn._get_client_side_session_parameter(
                 _PYTHON_SNOWPARK_USE_SQL_SIMPLIFIER_STRING, True
@@ -622,8 +634,19 @@ def __str__(self):
         )
 
     def _generate_new_action_id(self) -> int:
-        self._last_action_id += 1
-        return self._last_action_id
+        with self._lock:
+            self._last_action_id += 1
+            return self._last_action_id
+
+    @property
+    def _analyzer(self) -> Analyzer:
+        if not hasattr(self._thread_store, "analyzer"):
+            self._thread_store.analyzer = (
+                Analyzer(self)
+                if isinstance(self._conn, ServerConnection)
+                else MockAnalyzer(self)
+            )
+        return self._thread_store.analyzer
 
     def close(self) -> None:
         """Close this session."""
@@ -758,36 +781,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!"
@@ -797,11 +830,16 @@ 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
-            )
-            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!"
@@ -815,12 +853,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!"
@@ -829,6 +871,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(
@@ -838,16 +883,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:
         """
@@ -855,7 +904,8 @@ def cancel_all(self) -> None:
         This does not affect any action methods called in the future.
         """
         _logger.info("Canceling all running queries")
-        self._last_canceled_id = self._last_action_id
+        with self._lock:
+            self._last_canceled_id = self._last_action_id
         if not isinstance(self._conn, MockServerConnection):
             self._conn.run_query(
                 f"select system$cancel_all_queries({self._session_id})"
@@ -866,7 +916,8 @@ def get_imports(self) -> List[str]:
         Returns a list of imports added for user defined functions (UDFs).
         This list includes any Python or zip files that were added automatically by the library.
         """
-        return list(self._import_paths.keys())
+        with self._lock:
+            return list(self._import_paths.keys())
 
     def add_import(
         self,
@@ -947,7 +998,8 @@ def add_import(
         path, checksum, leading_path = self._resolve_import_path(
             path, import_path, chunk_size, whole_file_hash
         )
-        self._import_paths[path] = (checksum, leading_path)
+        with self._lock:
+            self._import_paths[path] = (checksum, leading_path)
 
     def remove_import(self, path: str) -> None:
         """
@@ -974,10 +1026,11 @@ def remove_import(self, path: str) -> None:
             if not trimmed_path.startswith(STAGE_PREFIX)
             else trimmed_path
         )
-        if abs_path not in self._import_paths:
-            raise KeyError(f"{abs_path} is not found in the existing imports")
-        else:
-            self._import_paths.pop(abs_path)
+        with self._lock:
+            if abs_path not in self._import_paths:
+                raise KeyError(f"{abs_path} is not found in the existing imports")
+            else:
+                self._import_paths.pop(abs_path)
 
     def clear_imports(self) -> None:
         """
@@ -986,10 +1039,11 @@ def clear_imports(self) -> None:
         if isinstance(self._conn, MockServerConnection):
             self.udf._clear_session_imports()
             self.sproc._clear_session_imports()
-        self._import_paths.clear()
+        with self._lock:
+            self._import_paths.clear()
 
+    @staticmethod
     def _resolve_import_path(
-        self,
         path: str,
         import_path: Optional[str] = None,
         chunk_size: int = 8192,
@@ -1077,7 +1131,8 @@ def _resolve_imports(
             upload_and_import_stage
         )
 
-        import_paths = udf_level_import_paths or self._import_paths
+        with self._lock:
+            import_paths = udf_level_import_paths or self._import_paths.copy()
         for path, (prefix, leading_path) in import_paths.items():
             # stage file
             if path.startswith(STAGE_PREFIX):
@@ -1159,7 +1214,8 @@ def get_packages(self) -> Dict[str, str]:
         The key of this ``dict`` is the package name and the value of this ``dict``
         is the corresponding requirement specifier.
         """
-        return self._packages.copy()
+        with self._package_lock:
+            return self._packages.copy()
 
     def add_packages(
         self, *packages: Union[str, ModuleType, Iterable[Union[str, ModuleType]]]
@@ -1250,16 +1306,18 @@ def remove_package(self, package: str) -> None:
             0
         """
         package_name = pkg_resources.Requirement.parse(package).key
-        if package_name in self._packages:
-            self._packages.pop(package_name)
-        else:
-            raise ValueError(f"{package_name} is not in the package list")
+        with self._package_lock:
+            if package_name in self._packages:
+                self._packages.pop(package_name)
+            else:
+                raise ValueError(f"{package_name} is not in the package list")
 
     def clear_packages(self) -> None:
         """
         Clears all third-party packages of a user-defined function (UDF).
         """
-        self._packages.clear()
+        with self._package_lock:
+            self._packages.clear()
 
     def add_requirements(self, file_path: str) -> None:
         """
@@ -1437,7 +1495,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():
@@ -1607,23 +1666,26 @@ def _resolve_packages(
         if isinstance(self._conn, MockServerConnection):
             # in local testing we don't resolve the packages, we just return what is added
             errors = []
-            result_dict = self._packages.copy()
-            for pkg_name, _, pkg_req in package_dict.values():
-                if pkg_name in result_dict and str(pkg_req) != result_dict[pkg_name]:
-                    errors.append(
-                        ValueError(
-                            f"Cannot add package '{str(pkg_req)}' because {result_dict[pkg_name]} "
-                            "is already added."
+            with self._package_lock:
+                result_dict = self._packages
+                for pkg_name, _, pkg_req in package_dict.values():
+                    if (
+                        pkg_name in result_dict
+                        and str(pkg_req) != result_dict[pkg_name]
+                    ):
+                        errors.append(
+                            ValueError(
+                                f"Cannot add package '{str(pkg_req)}' because {result_dict[pkg_name]} "
+                                "is already added."
+                            )
                         )
-                    )
-                else:
-                    result_dict[pkg_name] = str(pkg_req)
-            if len(errors) == 1:
-                raise errors[0]
-            elif len(errors) > 0:
-                raise RuntimeError(errors)
+                    else:
+                        result_dict[pkg_name] = str(pkg_req)
+                if len(errors) == 1:
+                    raise errors[0]
+                elif len(errors) > 0:
+                    raise RuntimeError(errors)
 
-            self._packages.update(result_dict)
             return list(result_dict.values())
 
         package_table = "information_schema.packages"
@@ -1638,46 +1700,47 @@ def _resolve_packages(
         #  'python-dateutil': 'python-dateutil==2.8.2'}
         # Add to packages dictionary. Make a copy of existing packages
         # dictionary to avoid modifying it during intermediate steps.
-        result_dict = (
-            existing_packages_dict.copy() if existing_packages_dict is not None else {}
-        )
-
-        # Retrieve list of dependencies that need to be added
-        dependency_packages = self._get_dependency_packages(
-            package_dict,
-            validate_package,
-            package_table,
-            result_dict,
-            statement_params=statement_params,
-        )
+        with self._package_lock:
+            result_dict = (
+                existing_packages_dict if existing_packages_dict is not None else {}
+            )
 
-        # Add dependency packages
-        for package in dependency_packages:
-            name = package.name
-            version = package.specs[0][1] if package.specs else None
+            # Retrieve list of dependencies that need to be added
+            dependency_packages = self._get_dependency_packages(
+                package_dict,
+                validate_package,
+                package_table,
+                result_dict,
+                statement_params=statement_params,
+            )
 
-            if name in result_dict:
-                if version is not None:
-                    added_package_has_version = "==" in result_dict[name]
-                    if added_package_has_version and result_dict[name] != str(package):
-                        raise ValueError(
-                            f"Cannot add dependency package '{name}=={version}' "
-                            f"because {result_dict[name]} is already added."
-                        )
+            # Add dependency packages
+            for package in dependency_packages:
+                name = package.name
+                version = package.specs[0][1] if package.specs else None
+
+                if name in result_dict:
+                    if version is not None:
+                        added_package_has_version = "==" in result_dict[name]
+                        if added_package_has_version and result_dict[name] != str(
+                            package
+                        ):
+                            raise ValueError(
+                                f"Cannot add dependency package '{name}=={version}' "
+                                f"because {result_dict[name]} is already added."
+                            )
+                        result_dict[name] = str(package)
+                else:
                     result_dict[name] = str(package)
-            else:
-                result_dict[name] = str(package)
 
-        # Always include cloudpickle
-        extra_modules = [cloudpickle]
-        if include_pandas:
-            extra_modules.append("pandas")
+            # Always include cloudpickle
+            extra_modules = [cloudpickle]
+            if include_pandas:
+                extra_modules.append("pandas")
 
-        if existing_packages_dict is not None:
-            existing_packages_dict.update(result_dict)
-        return list(result_dict.values()) + self._get_req_identifiers_list(
-            extra_modules, result_dict
-        )
+            return list(result_dict.values()) + self._get_req_identifiers_list(
+                extra_modules, result_dict
+            )
 
     def _upload_unsupported_packages(
         self,
@@ -1957,11 +2020,12 @@ def query_tag(self) -> Optional[str]:
 
     @query_tag.setter
     def query_tag(self, tag: str) -> None:
-        if tag:
-            self._conn.run_query(f"alter session set query_tag = {str_to_sql(tag)}")
-        else:
-            self._conn.run_query("alter session unset query_tag")
-        self._query_tag = tag
+        with self._lock:
+            if tag:
+                self._conn.run_query(f"alter session set query_tag = {str_to_sql(tag)}")
+            else:
+                self._conn.run_query("alter session unset query_tag")
+            self._query_tag = tag
 
     def _get_remote_query_tag(self) -> None:
         """
@@ -2354,18 +2418,19 @@ def get_session_stage(
             Therefore, if you switch database or schema during the session, the stage will not be re-created
             in the new database or schema, and still references the stage in the old database or schema.
         """
-        if not self._session_stage:
-            full_qualified_stage_name = self.get_fully_qualified_name_if_possible(
-                random_name_for_temp_object(TempObjectType.STAGE)
-            )
-            self._run_query(
-                f"create {get_temp_type_for_object(self._use_scoped_temp_objects, True)} \
-                stage if not exists {full_qualified_stage_name}",
-                is_ddl_on_temp_object=True,
-                statement_params=statement_params,
-            )
-            # set the value after running the query to ensure atomicity
-            self._session_stage = full_qualified_stage_name
+        with self._lock:
+            if not self._session_stage:
+                full_qualified_stage_name = self.get_fully_qualified_name_if_possible(
+                    random_name_for_temp_object(TempObjectType.STAGE)
+                )
+                self._run_query(
+                    f"create {get_temp_type_for_object(self._use_scoped_temp_objects, True)} \
+                    stage if not exists {full_qualified_stage_name}",
+                    is_ddl_on_temp_object=True,
+                    statement_params=statement_params,
+                )
+                # set the value after running the query to ensure atomicity
+                self._session_stage = full_qualified_stage_name
         return f"{STAGE_PREFIX}{self._session_stage}"
 
     def _write_modin_pandas_helper(
@@ -3067,8 +3132,9 @@ def get_fully_qualified_name_if_possible(self, name: str) -> str:
         """
         Returns the fully qualified object name if current database/schema exists, otherwise returns the object name
         """
-        database = self.get_current_database()
-        schema = self.get_current_schema()
+        with self._lock:
+            database = self.get_current_database()
+            schema = self.get_current_schema()
         if database and schema:
             return f"{database}.{schema}.{name}"
 
diff --git a/src/snowflake/snowpark/stored_procedure.py b/src/snowflake/snowpark/stored_procedure.py
index 92e63f6a76a..c2fcb024196 100644
--- a/src/snowflake/snowpark/stored_procedure.py
+++ b/src/snowflake/snowpark/stored_procedure.py
@@ -820,11 +820,15 @@ def _do_register_sp(
             force_inline_code=force_inline_code,
         )
 
-        if (not custom_python_runtime_version_allowed) and (self._session is not None):
-            check_python_runtime_version(
+        runtime_version_from_requirement = None
+        if self._session is not None:
+            runtime_version_from_requirement = (
                 self._session._runtime_version_from_requirement
             )
 
+        if not custom_python_runtime_version_allowed:
+            check_python_runtime_version(runtime_version_from_requirement)
+
         anonymous_sp_sql = None
         if anonymous:
             anonymous_sp_sql = generate_anonymous_python_sp_sql(
@@ -838,7 +842,7 @@ def _do_register_sp(
                 raw_imports=imports,
                 inline_python_code=code,
                 strict=strict,
-                runtime_version=self._session._runtime_version_from_requirement,
+                runtime_version=runtime_version_from_requirement,
                 external_access_integrations=external_access_integrations,
                 secrets=secrets,
                 native_app_params=native_app_params,
@@ -870,6 +874,7 @@ def _do_register_sp(
                     statement_params=statement_params,
                     comment=comment,
                     native_app_params=native_app_params,
+                    runtime_version=runtime_version_from_requirement,
                 )
             # an exception might happen during registering a stored procedure
             # (e.g., a dependency might not be found on the stage),
diff --git a/src/snowflake/snowpark/udaf.py b/src/snowflake/snowpark/udaf.py
index fdc3d555281..889b5b62915 100644
--- a/src/snowflake/snowpark/udaf.py
+++ b/src/snowflake/snowpark/udaf.py
@@ -680,11 +680,15 @@ def _do_register_udaf(
             is_permanent=is_permanent,
         )
 
-        if (not custom_python_runtime_version_allowed) and (self._session is not None):
-            check_python_runtime_version(
+        runtime_version_from_requirement = None
+        if self._session is not None:
+            runtime_version_from_requirement = (
                 self._session._runtime_version_from_requirement
             )
 
+        if not custom_python_runtime_version_allowed:
+            check_python_runtime_version(runtime_version_from_requirement)
+
         raised = False
         try:
             create_python_udf_or_sp(
@@ -710,6 +714,7 @@ def _do_register_udaf(
                 statement_params=statement_params,
                 comment=comment,
                 native_app_params=native_app_params,
+                runtime_version=runtime_version_from_requirement,
             )
         # an exception might happen during registering a udaf
         # (e.g., a dependency might not be found on the stage),
diff --git a/src/snowflake/snowpark/udf.py b/src/snowflake/snowpark/udf.py
index 74278e48d30..b71a40263fd 100644
--- a/src/snowflake/snowpark/udf.py
+++ b/src/snowflake/snowpark/udf.py
@@ -873,11 +873,15 @@ def _do_register_udf(
             is_permanent=is_permanent,
         )
 
-        if (not custom_python_runtime_version_allowed) and (self._session is not None):
-            check_python_runtime_version(
+        runtime_version_from_requirement = None
+        if self._session is not None:
+            runtime_version_from_requirement = (
                 self._session._runtime_version_from_requirement
             )
 
+        if not custom_python_runtime_version_allowed:
+            check_python_runtime_version(runtime_version_from_requirement)
+
         raised = False
         try:
             create_python_udf_or_sp(
@@ -905,6 +909,7 @@ def _do_register_udf(
                 statement_params=statement_params,
                 comment=comment,
                 native_app_params=native_app_params,
+                runtime_version=runtime_version_from_requirement,
             )
         # an exception might happen during registering a udf
         # (e.g., a dependency might not be found on the stage),
diff --git a/src/snowflake/snowpark/udtf.py b/src/snowflake/snowpark/udtf.py
index 03b71aa5f75..856007cdb64 100644
--- a/src/snowflake/snowpark/udtf.py
+++ b/src/snowflake/snowpark/udtf.py
@@ -934,11 +934,15 @@ def _do_register_udtf(
             is_permanent=is_permanent,
         )
 
-        if (not custom_python_runtime_version_allowed) and (self._session is not None):
-            check_python_runtime_version(
+        runtime_version_from_requirement = None
+        if self._session is not None:
+            runtime_version_from_requirement = (
                 self._session._runtime_version_from_requirement
             )
 
+        if not custom_python_runtime_version_allowed:
+            check_python_runtime_version(runtime_version_from_requirement)
+
         raised = False
         try:
             create_python_udf_or_sp(
@@ -966,6 +970,7 @@ def _do_register_udtf(
                 statement_params=statement_params,
                 comment=comment,
                 native_app_params=native_app_params,
+                runtime_version=runtime_version_from_requirement,
             )
         # an exception might happen during registering a udtf
         # (e.g., a dependency might not be found on the stage),
diff --git a/tests/integ/conftest.py b/tests/integ/conftest.py
index 8c0d855d0da..4760095c5b8 100644
--- a/tests/integ/conftest.py
+++ b/tests/integ/conftest.py
@@ -13,7 +13,13 @@
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.mock._connection import MockServerConnection
 from tests.parameters import CONNECTION_PARAMETERS
-from tests.utils import TEST_SCHEMA, Utils, running_on_jenkins, running_on_public_ci
+from tests.utils import (
+    TEST_SCHEMA,
+    TestFiles,
+    Utils,
+    running_on_jenkins,
+    running_on_public_ci,
+)
 
 
 def print_help() -> None:
@@ -268,3 +274,18 @@ def temp_schema(connection, session, local_testing_mode) -> None:
             )
             yield temp_schema_name
             cursor.execute(f"DROP SCHEMA IF EXISTS {temp_schema_name}")
+
+
+@pytest.fixture(scope="module")
+def temp_stage(session, resources_path, local_testing_mode):
+    tmp_stage_name = Utils.random_stage_name()
+    test_files = TestFiles(resources_path)
+
+    if not local_testing_mode:
+        Utils.create_stage(session, tmp_stage_name, is_temporary=True)
+    Utils.upload_to_stage(
+        session, tmp_stage_name, test_files.test_file_parquet, compress=False
+    )
+    yield tmp_stage_name
+    if not local_testing_mode:
+        Utils.drop_stage(session, tmp_stage_name)
diff --git a/tests/integ/scala/test_file_operation_suite.py b/tests/integ/scala/test_file_operation_suite.py
index 2dc424dde09..82a1722a729 100644
--- a/tests/integ/scala/test_file_operation_suite.py
+++ b/tests/integ/scala/test_file_operation_suite.py
@@ -14,7 +14,7 @@
     SnowparkSQLException,
     SnowparkUploadFileException,
 )
-from tests.utils import IS_IN_STORED_PROC, IS_WINDOWS, TestFiles, Utils
+from tests.utils import IS_IN_STORED_PROC, IS_WINDOWS, Utils
 
 
 def random_alphanumeric_name():
@@ -74,21 +74,6 @@ def path4(temp_source_directory):
     yield filename
 
 
-@pytest.fixture(scope="module")
-def temp_stage(session, resources_path, local_testing_mode):
-    tmp_stage_name = Utils.random_stage_name()
-    test_files = TestFiles(resources_path)
-
-    if not local_testing_mode:
-        Utils.create_stage(session, tmp_stage_name, is_temporary=True)
-    Utils.upload_to_stage(
-        session, tmp_stage_name, test_files.test_file_parquet, compress=False
-    )
-    yield tmp_stage_name
-    if not local_testing_mode:
-        Utils.drop_stage(session, tmp_stage_name)
-
-
 def test_put_with_one_file(
     session, temp_stage, path1, path2, path3, local_testing_mode
 ):
diff --git a/tests/integ/test_multithreading.py b/tests/integ/test_multithreading.py
new file mode 100644
index 00000000000..a6d13584bcf
--- /dev/null
+++ b/tests/integ/test_multithreading.py
@@ -0,0 +1,582 @@
+#
+# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
+#
+
+import gc
+import hashlib
+import logging
+import os
+import tempfile
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List, Tuple  # noqa: F401
+from unittest.mock import patch
+
+import pytest
+
+from snowflake.snowpark.session import Session
+from snowflake.snowpark.types import IntegerType
+from tests.integ.test_temp_table_cleanup import wait_for_drop_table_sql_done
+
+try:
+    import dateutil
+
+    # six is the dependency of dateutil
+    import six
+
+    is_dateutil_available = True
+except ImportError:
+    is_dateutil_available = False
+
+from snowflake.snowpark.functions import lit
+from snowflake.snowpark.row import Row
+from tests.utils import IS_IN_STORED_PROC, IS_LINUX, IS_WINDOWS, TestFiles, Utils
+
+
+def test_concurrent_select_queries(session):
+    def run_select(session_, thread_id):
+        df = session_.sql(f"SELECT {thread_id} as A")
+        assert df.collect()[0][0] == thread_id
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        for i in range(10):
+            executor.submit(run_select, session, i)
+
+
+def test_concurrent_dataframe_operations(session):
+    try:
+        table_name = Utils.random_table_name()
+        data = [(i, 11 * i) for i in range(10)]
+        df = session.create_dataframe(data, ["A", "B"])
+        df.write.save_as_table(table_name, table_type="temporary")
+
+        def run_dataframe_operation(session_, thread_id):
+            df = session_.table(table_name)
+            df = df.filter(df.a == lit(thread_id))
+            df = df.with_column("C", df.b + 100 * df.a)
+            df = df.rename(df.a, "D").limit(1)
+            return df
+
+        dfs = []
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            df_futures = [
+                executor.submit(run_dataframe_operation, session, i) for i in range(10)
+            ]
+
+            for future in as_completed(df_futures):
+                dfs.append(future.result())
+
+        main_df = dfs[0]
+        for df in dfs[1:]:
+            main_df = main_df.union(df)
+
+        Utils.check_answer(
+            main_df, [Row(D=i, B=11 * i, C=11 * i + 100 * i) for i in range(10)]
+        )
+
+    finally:
+        Utils.drop_table(session, table_name)
+
+
+@pytest.mark.xfail(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="SQL query and query listeners are not supported",
+    run=False,
+)
+def test_query_listener(session):
+    def run_select(session_, thread_id):
+        session_.sql(f"SELECT {thread_id} as A").collect()
+
+    with session.query_history() as history:
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            for i in range(10):
+                executor.submit(run_select, session, i)
+
+    queries_sent = [query.sql_text for query in history.queries]
+    assert len(queries_sent) == 10
+    for i in range(10):
+        assert f"SELECT {i} as A" in queries_sent
+
+
+@pytest.mark.xfail(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="Query tag is a SQL feature",
+    run=False,
+)
+@pytest.mark.skipif(
+    IS_IN_STORED_PROC, reason="show parameters is not supported in stored procedure"
+)
+def test_query_tagging(session):
+    def set_query_tag(session_, thread_id):
+        session_.query_tag = f"tag_{thread_id}"
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        for i in range(10):
+            executor.submit(set_query_tag, session, i)
+
+    actual_query_tag = session.sql("SHOW PARAMETERS LIKE 'QUERY_TAG'").collect()[0][1]
+    assert actual_query_tag == session.query_tag
+
+
+@pytest.mark.xfail(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="SQL query is not supported",
+    run=False,
+)
+def test_session_stage_created_once(session):
+    with patch.object(
+        session._conn, "run_query", wraps=session._conn.run_query
+    ) as patched_run_query:
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            for _ in range(10):
+                executor.submit(session.get_session_stage)
+
+        assert patched_run_query.call_count == 1
+
+
+def test_action_ids_are_unique(session):
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        action_ids = set()
+        futures = [executor.submit(session._generate_new_action_id) for _ in range(10)]
+
+        for future in as_completed(futures):
+            action_ids.add(future.result())
+
+    assert len(action_ids) == 10
+
+
+@pytest.mark.parametrize("use_stream", [True, False])
+def test_file_io(session, resources_path, temp_stage, use_stream):
+    stage_prefix = f"prefix_{Utils.random_alphanumeric_str(10)}"
+    stage_with_prefix = f"@{temp_stage}/{stage_prefix}/"
+    test_files = TestFiles(resources_path)
+
+    resources_files = [
+        test_files.test_file_csv,
+        test_files.test_file2_csv,
+        test_files.test_file_json,
+        test_files.test_file_csv_header,
+        test_files.test_file_csv_colon,
+        test_files.test_file_csv_quotes,
+        test_files.test_file_csv_special_format,
+        test_files.test_file_json_special_format,
+        test_files.test_file_csv_quotes_special,
+        test_files.test_concat_file1_csv,
+        test_files.test_concat_file2_csv,
+    ]
+
+    def get_file_hash(fd):
+        return hashlib.md5(fd.read()).hexdigest()
+
+    def put_and_get_file(upload_file_path, download_dir):
+        if use_stream:
+            with open(upload_file_path, "rb") as fd:
+                results = session.file.put_stream(
+                    fd, stage_with_prefix, auto_compress=False, overwrite=False
+                )
+        else:
+            results = session.file.put(
+                upload_file_path,
+                stage_with_prefix,
+                auto_compress=False,
+                overwrite=False,
+            )
+        # assert file is uploaded successfully
+        assert len(results) == 1
+        assert results[0].status == "UPLOADED"
+
+        stage_file_name = f"{stage_with_prefix}{os.path.basename(upload_file_path)}"
+        if use_stream:
+            fd = session.file.get_stream(stage_file_name, download_dir)
+            with open(upload_file_path, "rb") as upload_fd:
+                assert get_file_hash(upload_fd) == get_file_hash(fd)
+
+        else:
+            results = session.file.get(stage_file_name, download_dir)
+            # assert file is downloaded successfully
+            assert len(results) == 1
+            assert results[0].status == "DOWNLOADED"
+            download_file_path = results[0].file
+            # assert two files are identical
+            with open(upload_file_path, "rb") as upload_fd, open(
+                download_file_path, "rb"
+            ) as download_fd:
+                assert get_file_hash(upload_fd) == get_file_hash(download_fd)
+
+    with tempfile.TemporaryDirectory() as download_dir:
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            for file_path in resources_files:
+                executor.submit(put_and_get_file, file_path, download_dir)
+
+        if not use_stream:
+            # assert all files are downloaded
+            assert set(os.listdir(download_dir)) == {
+                os.path.basename(file_path) for file_path in resources_files
+            }
+
+
+def test_concurrent_add_packages(session):
+    # this is a list of packages available in snowflake anaconda. If this
+    # test fails due to packages not being available, please update the list
+    package_list = {
+        "graphviz",
+        "numpy",
+        "pandas",
+        "scipy",
+        "scikit-learn",
+        "matplotlib",
+    }
+
+    try:
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            futures = [
+                executor.submit(session.add_packages, package)
+                for package in package_list
+            ]
+
+            for future in as_completed(futures):
+                future.result()
+
+            assert session.get_packages() == {
+                package: package for package in package_list
+            }
+    finally:
+        session.clear_packages()
+
+
+def test_concurrent_remove_package(session):
+    def remove_package(session_, package_name):
+        try:
+            session_.remove_package(package_name)
+            return True
+        except ValueError:
+            return False
+        except Exception as e:
+            raise e
+
+    try:
+        session.add_packages("numpy")
+        with ThreadPoolExecutor(max_workers=10) as executor:
+
+            futures = [
+                executor.submit(remove_package, session, "numpy") for _ in range(10)
+            ]
+            success_count, failure_count = 0, 0
+            for future in as_completed(futures):
+                if future.result():
+                    success_count += 1
+                else:
+                    failure_count += 1
+
+            # assert that only one thread was able to remove the package
+            assert success_count == 1
+            assert failure_count == 9
+    finally:
+        session.clear_packages()
+
+
+@pytest.mark.skipif(not is_dateutil_available, reason="dateutil is not available")
+def test_concurrent_add_import(session, resources_path):
+    test_files = TestFiles(resources_path)
+    import_files = [
+        test_files.test_udf_py_file,
+        os.path.relpath(test_files.test_udf_py_file),
+        test_files.test_udf_directory,
+        os.path.relpath(test_files.test_udf_directory),
+        six.__file__,
+        os.path.relpath(six.__file__),
+        os.path.dirname(dateutil.__file__),
+    ]
+    try:
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            for file in import_files:
+                executor.submit(
+                    session.add_import,
+                    file,
+                )
+
+        assert set(session.get_imports()) == {
+            os.path.abspath(file) for file in import_files
+        }
+    finally:
+        session.clear_imports()
+
+
+def test_concurrent_remove_import(session, resources_path):
+    test_files = TestFiles(resources_path)
+
+    def remove_import(session_, import_file):
+        try:
+            session_.remove_import(import_file)
+            return True
+        except KeyError:
+            return False
+        except Exception as e:
+            raise e
+
+    try:
+        session.add_import(test_files.test_udf_py_file)
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            futures = [
+                executor.submit(remove_import, session, test_files.test_udf_py_file)
+                for _ in range(10)
+            ]
+
+            success_count, failure_count = 0, 0
+            for future in as_completed(futures):
+                if future.result():
+                    success_count += 1
+                else:
+                    failure_count += 1
+
+            # assert that only one thread was able to remove the import
+            assert success_count == 1
+            assert failure_count == 9
+    finally:
+        session.clear_imports()
+
+
+def test_concurrent_sp_register(session, tmpdir):
+    try:
+        session.add_packages("snowflake-snowpark-python")
+
+        def register_and_test_sp(session_, thread_id):
+            prefix = Utils.random_alphanumeric_str(10)
+            sp_file_path = os.path.join(tmpdir, f"{prefix}_add_{thread_id}.py")
+            sproc_body = f"""
+from snowflake.snowpark import Session
+from snowflake.snowpark.functions import (
+    col,
+    lit
+)
+def add_{thread_id}(session_: Session, x: int) -> int:
+    return (
+        session_.create_dataframe([[x, ]], schema=["x"])
+        .select(col("x") + lit({thread_id}))
+        .collect()[0][0]
+    )
+"""
+            with open(sp_file_path, "w") as f:
+                f.write(sproc_body)
+                f.flush()
+
+            add_sp_from_file = session_.sproc.register_from_file(
+                sp_file_path, f"add_{thread_id}"
+            )
+            add_sp = session_.sproc.register(
+                lambda sess_, x: sess_.sql(f"select {x} + {thread_id}").collect()[0][0],
+                return_type=IntegerType(),
+                input_types=[IntegerType()],
+            )
+
+            assert add_sp_from_file(1) == thread_id + 1
+            assert add_sp(1) == thread_id + 1
+
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            for i in range(10):
+                executor.submit(register_and_test_sp, session, i)
+    finally:
+        session.clear_packages()
+
+
+def test_concurrent_udf_register(session, tmpdir):
+    df = session.range(-5, 5).to_df("a")
+
+    def register_and_test_udf(session_, thread_id):
+        prefix = Utils.random_alphanumeric_str(10)
+        file_path = os.path.join(tmpdir, f"{prefix}_add_{thread_id}.py")
+        with open(file_path, "w") as f:
+            func = f"""
+def add_{thread_id}(x: int) -> int:
+    return x + {thread_id}
+"""
+            f.write(func)
+            f.flush()
+        add_i_udf_from_file = session_.udf.register_from_file(
+            file_path, f"add_{thread_id}"
+        )
+        add_i_udf = session_.udf.register(
+            lambda x: x + thread_id,
+            return_type=IntegerType(),
+            input_types=[IntegerType()],
+        )
+
+        Utils.check_answer(
+            df.select(add_i_udf(df.a), add_i_udf_from_file(df.a)),
+            [(thread_id + i, thread_id + i) for i in range(-5, 5)],
+        )
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        for i in range(10):
+            executor.submit(register_and_test_udf, session, i)
+
+
+@pytest.mark.xfail(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="UDTFs is not supported in local testing mode",
+    run=False,
+)
+def test_concurrent_udtf_register(session, tmpdir):
+    def register_and_test_udtf(session_, thread_id):
+        udtf_body = f"""
+from typing import List, Tuple
+
+class UDTFEcho:
+    def process(
+        self,
+        num: int,
+    ) -> List[Tuple[int]]:
+        return [(num + {thread_id},)]
+"""
+        prefix = Utils.random_alphanumeric_str(10)
+        file_path = os.path.join(tmpdir, f"{prefix}_udtf_echo_{thread_id}.py")
+        with open(file_path, "w") as f:
+            f.write(udtf_body)
+            f.flush()
+
+        d = {}
+        exec(udtf_body, {**globals(), **locals()}, d)
+        echo_udtf_from_file = session_.udtf.register_from_file(
+            file_path, "UDTFEcho", output_schema=["num"]
+        )
+        echo_udtf = session_.udtf.register(d["UDTFEcho"], output_schema=["num"])
+
+        df_local = session.table_function(echo_udtf(lit(1)))
+        df_from_file = session.table_function(echo_udtf_from_file(lit(1)))
+        assert df_local.collect() == [(thread_id + 1,)]
+        assert df_from_file.collect() == [(thread_id + 1,)]
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        for i in range(10):
+            executor.submit(register_and_test_udtf, session, i)
+
+
+@pytest.mark.xfail(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="UDAFs is not supported in local testing mode",
+    run=False,
+)
+def test_concurrent_udaf_register(session: Session, tmpdir):
+    df = session.create_dataframe([[1, 3], [1, 4], [2, 5], [2, 6]]).to_df("a", "b")
+
+    def register_and_test_udaf(session_, thread_id):
+        udaf_body = f"""
+class OffsetSumUDAFHandler:
+    def __init__(self) -> None:
+        self._sum = 0
+
+    @property
+    def aggregate_state(self):
+        return self._sum
+
+    def accumulate(self, input_value):
+        self._sum += input_value
+
+    def merge(self, other_sum):
+        self._sum += other_sum
+
+    def finish(self):
+        return self._sum + {thread_id}
+    """
+        prefix = Utils.random_alphanumeric_str(10)
+        file_path = os.path.join(tmpdir, f"{prefix}_udaf_{thread_id}.py")
+        with open(file_path, "w") as f:
+            f.write(udaf_body)
+            f.flush()
+        d = {}
+        exec(udaf_body, {**globals(), **locals()}, d)
+
+        offset_sum_udaf_from_file = session_.udaf.register_from_file(
+            file_path,
+            "OffsetSumUDAFHandler",
+            return_type=IntegerType(),
+            input_types=[IntegerType()],
+        )
+        offset_sum_udaf = session_.udaf.register(
+            d["OffsetSumUDAFHandler"],
+            return_type=IntegerType(),
+            input_types=[IntegerType()],
+        )
+
+        Utils.check_answer(
+            df.agg(offset_sum_udaf_from_file(df.a)), [Row(6 + thread_id)]
+        )
+        Utils.check_answer(df.agg(offset_sum_udaf(df.a)), [Row(6 + thread_id)])
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        for i in range(10):
+            executor.submit(register_and_test_udaf, session, i)
+
+
+@pytest.mark.xfail(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="session.sql is not supported in local testing mode",
+    run=False,
+)
+def test_auto_temp_table_cleaner(session, caplog):
+    session._temp_table_auto_cleaner.ref_count_map.clear()
+    original_auto_clean_up_temp_table_enabled = session.auto_clean_up_temp_table_enabled
+    session.auto_clean_up_temp_table_enabled = True
+
+    def create_temp_table(session_, thread_id):
+        df = session.sql(f"select {thread_id} as A").cache_result()
+        table_name = df.table_name
+        del df
+        return table_name
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        futures = []
+        table_names = []
+        for i in range(10):
+            futures.append(executor.submit(create_temp_table, session, i))
+
+        for future in as_completed(futures):
+            table_names.append(future.result())
+
+    gc.collect()
+    wait_for_drop_table_sql_done(session, caplog, expect_drop=True)
+
+    try:
+        for table_name in table_names:
+            assert session._temp_table_auto_cleaner.ref_count_map[table_name] == 0
+        assert session._temp_table_auto_cleaner.num_temp_tables_created == 10
+        assert session._temp_table_auto_cleaner.num_temp_tables_cleaned == 10
+    finally:
+        session.auto_clean_up_temp_table_enabled = (
+            original_auto_clean_up_temp_table_enabled
+        )
+
+
+@pytest.mark.skipif(
+    IS_LINUX or IS_WINDOWS,
+    reason="Linux and Windows test show multiple active threads when no threadpool is enabled",
+)
+@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
+    )
diff --git a/tests/integ/test_telemetry.py b/tests/integ/test_telemetry.py
index f243f37d3cb..4708ad8a80e 100644
--- a/tests/integ/test_telemetry.py
+++ b/tests/integ/test_telemetry.py
@@ -5,6 +5,7 @@
 
 import decimal
 import sys
+import threading
 import uuid
 from functools import partial
 from typing import Any, Dict, Tuple
@@ -595,12 +596,14 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
     query_plan_height = 2 if sql_simplifier_enabled else 3
     filter = 1 if sql_simplifier_enabled else 2
     low_impact = 3 if sql_simplifier_enabled else 2
+    thread_ident = threading.get_ident()
 
     assert df._plan.api_calls == [
         {
             "name": "Session.range",
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "plan_uuid": df._plan.uuid,
+            "thread_ident": thread_ident,
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_num_selects_with_complexity_merged": 0,
@@ -625,6 +628,7 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "name": "Session.range",
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "plan_uuid": df._plan.uuid,
+            "thread_ident": thread_ident,
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_num_selects_with_complexity_merged": 0,
@@ -649,6 +653,7 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "name": "Session.range",
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "plan_uuid": df._plan.uuid,
+            "thread_ident": thread_ident,
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_num_selects_with_complexity_merged": 0,
@@ -673,6 +678,7 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "name": "Session.range",
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "plan_uuid": df._plan.uuid,
+            "thread_ident": thread_ident,
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_num_selects_with_complexity_merged": 0,
@@ -697,6 +703,7 @@ def test_execute_queries_api_calls(session, sql_simplifier_enabled):
             "name": "Session.range",
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "plan_uuid": df._plan.uuid,
+            "thread_ident": thread_ident,
             "query_plan_height": query_plan_height,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_num_selects_with_complexity_merged": 0,
@@ -844,11 +851,13 @@ def test_dataframe_stat_functions_api_calls(session):
     # therefore we can't predict it. We check that the uuid for crosstab is same as
     # that for df.
     uuid = df._plan.api_calls[0]["plan_uuid"]
+    thread_ident = threading.get_ident()
     assert crosstab._plan.api_calls == [
         {
             "name": "Session.create_dataframe[values]",
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "plan_uuid": uuid,
+            "thread_ident": thread_ident,
             "query_plan_height": 4,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_num_selects_with_complexity_merged": 0,
@@ -870,6 +879,7 @@ def test_dataframe_stat_functions_api_calls(session):
             "name": "Session.create_dataframe[values]",
             "sql_simplifier_enabled": session.sql_simplifier_enabled,
             "plan_uuid": uuid,
+            "thread_ident": thread_ident,
             "query_plan_height": 4,
             "query_plan_num_duplicate_nodes": 0,
             "query_plan_num_selects_with_complexity_merged": 0,
@@ -1281,3 +1291,19 @@ def send_telemetry():
     data, type_, _ = telemetry_tracker.extract_telemetry_log_data(-1, send_telemetry)
     assert data == expected_data
     assert type_ == "snowpark_temp_table_cleanup_abnormal_exception"
+
+
+def test_cursor_created_telemetry(session):
+    client = session._conn._telemetry_client
+    telemetry_tracker = TelemetryDataTracker(session)
+
+    def send_telemetry():
+        client.send_cursor_created_telemetry(session_id=session.session_id, thread_id=1)
+
+    expected_data = {
+        "session_id": session.session_id,
+        "thread_ident": 1,
+    }
+    data, type_, _ = telemetry_tracker.extract_telemetry_log_data(-1, send_telemetry)
+    assert data == expected_data
+    assert type_ == "snowpark_cursor_created"
diff --git a/tests/integ/test_temp_table_cleanup.py b/tests/integ/test_temp_table_cleanup.py
index 739f84fa878..514d0a080fc 100644
--- a/tests/integ/test_temp_table_cleanup.py
+++ b/tests/integ/test_temp_table_cleanup.py
@@ -37,12 +37,14 @@ def setup(session):
 def wait_for_drop_table_sql_done(session: Session, caplog, expect_drop: bool) -> None:
     # Loop through captured logs and search for the pattern
     pattern = r"Dropping .* with query id ([0-9a-f\-]+)"
+    matches = []
     for record in caplog.records:
         match = re.search(pattern, record.message)
         if match:
             query_id = match.group(1)
-            break
-    else:
+            matches.append(query_id)
+
+    if len(matches) == 0:
         if expect_drop:
             pytest.fail("No drop table sql found in logs")
         else:
@@ -50,9 +52,10 @@ def wait_for_drop_table_sql_done(session: Session, caplog, expect_drop: bool) ->
             return
 
     caplog.clear()
-    async_job = session.create_async_job(query_id)
-    # Wait for the async job to finish
-    _ = async_job.result()
+    for query_id in matches:
+        async_job = session.create_async_job(query_id)
+        # Wait for the async job to finish
+        _ = async_job.result()
 
 
 def test_basic(session, caplog):
diff --git a/tests/unit/compiler/test_large_query_breakdown.py b/tests/unit/compiler/test_large_query_breakdown.py
index d040ca25f49..5c9e140694f 100644
--- a/tests/unit/compiler/test_large_query_breakdown.py
+++ b/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 (
diff --git a/tests/unit/test_server_connection.py b/tests/unit/test_server_connection.py
index 72ccb6f6c42..cf10dc9d29a 100644
--- a/tests/unit/test_server_connection.py
+++ b/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",
diff --git a/tests/unit/test_udf_utils.py b/tests/unit/test_udf_utils.py
index 12a49443539..09e389d0c24 100644
--- a/tests/unit/test_udf_utils.py
+++ b/tests/unit/test_udf_utils.py
@@ -5,6 +5,7 @@
 import logging
 import os
 import pickle
+import threading
 from unittest import mock
 
 import pytest
@@ -249,6 +250,7 @@ def test_add_snowpark_package_to_sproc_packages_to_session():
         "random_package_one": "random_package_one",
         "random_package_two": "random_package_two",
     }
+    fake_session._package_lock = threading.RLock()
     result = add_snowpark_package_to_sproc_packages(session=fake_session, packages=None)
 
     major, minor, patch = VERSION