From fb4ecf6fc8ca6cfce343559006798ce8d4fddec9 Mon Sep 17 00:00:00 2001
From: Afroz Alam <afroz.alam@snowflake.com>
Date: Wed, 11 Sep 2024 14:04:36 -0700
Subject: [PATCH] add locks

---
 src/snowflake/snowpark/session.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py
index 8da0794f139..607cd047f2b 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
@@ -499,6 +500,8 @@ 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()
         self._query_tag = None
         self._import_paths: Dict[str, Tuple[Optional[str], Optional[str]]] = {}
         self._packages: Dict[str, str] = {}