SNOW-1418523: make udf and sproc registration thread safe

src/snowflake/snowpark/_internal/server_connection.py

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

src/snowflake/snowpark/_internal/udf_utils.py

@@ -980,8 +980,11 @@ 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._lock:
+                session_packages = session._packages.copy()
+            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):
@@ -1073,6 +1076,7 @@ def resolve_imports_and_packages(
             )
         )
 
+    all_urls = []
     if session is not None:
         import_only_stage = (
             unwrap_stage_location_single_quote(stage_location)
@@ -1086,7 +1090,6 @@ def resolve_imports_and_packages(
             else session.get_session_stage(statement_params=statement_params)
         )
 
-    if session:
         if imports:
             udf_level_imports = {}
             for udf_import in imports:
@@ -1114,22 +1117,15 @@ def resolve_imports_and_packages(
                 upload_and_import_stage,
                 statement_params=statement_params,
             )
-        else:
-            all_urls = []
-    else:
-        all_urls = []
 
     dest_prefix = get_udf_upload_prefix(udf_name)
 
     # Upload closure to stage if it is beyond inline closure size limit
     handler = inline_code = upload_file_stage_location = None
-    custom_python_runtime_version_allowed = False
+    # As cloudpickle is being used, we cannot allow a custom runtime
+    custom_python_runtime_version_allowed = not isinstance(func, Callable)
     if session is not None:
         if isinstance(func, Callable):
-            custom_python_runtime_version_allowed = (
-                False  # As cloudpickle is being used, we cannot allow a custom runtime
-            )
-
             # generate a random name for udf py file
             # and we compress it first then upload it
             udf_file_name_base = f"udf_py_{random_number()}"
@@ -1174,7 +1170,6 @@ def resolve_imports_and_packages(
                 upload_file_stage_location = None
                 handler = _DEFAULT_HANDLER_NAME
         else:
-            custom_python_runtime_version_allowed = True
             udf_file_name = os.path.basename(func[0])
             # for a compressed file, it might have multiple extensions
             # and we should remove all extensions
@@ -1199,11 +1194,6 @@ def resolve_imports_and_packages(
                     skip_upload_on_content_match=skip_upload_on_content_match,
                 )
                 all_urls.append(upload_file_stage_location)
-    else:
-        if isinstance(func, Callable):
-            custom_python_runtime_version_allowed = False
-        else:
-            custom_python_runtime_version_allowed = True
 
     # build imports and packages string
     all_imports = ",".join(
@@ -1246,11 +1236,13 @@ 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 = (
+        f"{sys.version_info[0]}.{sys.version_info[1]}"
+        if not runtime_version
+        else runtime_version
+    )
 
     if replace and if_not_exists:
         raise ValueError("options replace and if_not_exists are incompatible")

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