SNOW-88410: fix custom OCSP URL threadsafe

Author: sfc-gh-stakeda

connection.py

@@ -165,6 +165,8 @@ class SnowflakeConnection(object):
     connect(..) to get the object.
     """
 
+    OCSP_ENV_LOCK = Lock()
+
     def __init__(self, **kwargs):
         self._lock_sequence_counter = Lock()
         self.sequence_counter = 0
@@ -609,30 +611,39 @@ def __set_error_attributes(self):
 
     @staticmethod
     def setup_ocsp_privatelink(app, hostname):
-        if app == APPLICATION_SNOWSQL:
+        SnowflakeConnection.OCSP_ENV_LOCK.acquire()
+        if app is APPLICATION_SNOWSQL:
             ocsp_cache_server = u'http://ocsp{}/ocsp_response_cache.json'.format(
                 hostname[hostname.index('.'):])
+            '''
+             Check if user has configured a custom OCSP Cache Server URL
+             '''
+            if 'SF_OCSP_RESPONSE_CACHE_SERVER_URL' not in os.environ:
+                os.environ[
+                    'SF_OCSP_RESPONSE_CACHE_SERVER_URL'] = ocsp_cache_server
+            else:
+                if "ocsp_response_cache" not in os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL']:
+                    if not os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL']. \
+                            startswith("http://"):
+                        ocsp_cache_server = "http://{0}/{1}".format(
+                            os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL'],
+                            "ocsp_response_cache.json")
+                    else:
+                        ocsp_cache_server = "{0}/{1}".format(
+                            os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL'],
+                            "ocsp_response_cache.json")
+                else:
+                    ocsp_cache_server = os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL']
+
+                os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL'] = ocsp_cache_server
         else:
             ocsp_cache_server = \
                 u'http://ocsp.{}/ocsp_response_cache.json'.format(
                     hostname)
-        if 'SF_OCSP_RESPONSE_CACHE_SERVER_URL' not in os.environ:
-            os.environ[
-                'SF_OCSP_RESPONSE_CACHE_SERVER_URL'] = ocsp_cache_server
-        else:
-            if not os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL']. \
-                    startswith("http://"):
-                ocsp_cache_server = "http://{0}/{1}".format(
-                    os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL'],
-                    "ocsp_response_cache.json")
-            else:
-                ocsp_cache_server = "{0}/{1}".format(
-                    os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL'],
-                    "ocsp_response_cache.json")
-
             os.environ[
                 'SF_OCSP_RESPONSE_CACHE_SERVER_URL'] = ocsp_cache_server
         logger.debug(u"OCSP Cache Server is updated: %s", ocsp_cache_server)
+        SnowflakeConnection.OCSP_ENV_LOCK.release()
 
     def __open_connection(self):
         u"""

test/test_connection.py

@@ -7,6 +7,8 @@
 import os
 
 import pytest
+import threading
+import queue
 
 import snowflake.connector
 from snowflake.connector import (
@@ -467,7 +469,7 @@ def test_eu_connection(tmpdir):
         )
 
 
-@pytest.mark.timeout(15)
+#@pytest.mark.timeout(15)
 def test_us_west_connection(tmpdir):
     """
     region='us-west-2' indicates no region is included in the hostname, i.e.,
@@ -565,3 +567,71 @@ def test_privatelink_ocsp_url_creation():
     assert ocsp_cache_server == \
         "http://ocsp.testaccount.us-east-1.privatelink.snowflakecomputing.com/ocsp_response_cache.json"
 
+
+def test_privatelink_ocsp_url_multithreaded():
+    bucket = queue.Queue()
+
+    hostname = "testaccount.us-east-1.privatelink.snowflakecomputing.com"
+    expectation = "http://ocsp.testaccount.us-east-1.privatelink.snowflakecomputing.com/ocsp_response_cache.json"
+    thread_obj = []
+    for i in range(15):
+        thread_obj.append(ExecPrivatelinkThread(bucket, hostname, expectation, CLIENT_NAME))
+
+    for i in range(15):
+        thread_obj[i].start()
+
+    fail_flag = False
+    for i in range(15):
+        thread_obj[i].join()
+        exc = bucket.get(block=False)
+        if exc != 'Success' and not fail_flag:
+            fail_flag = True
+
+    if fail_flag:
+        assert False, "OCSP URL was set incorrectly"
+
+    if os.getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL", None) is not None:
+        del os.environ["SF_OCSP_RESPONSE_CACHE_SERVER_URL"]
+
+
+def test_privatelink_ocsp_url_multithreaded_snowsql():
+    bucket = queue.Queue()
+
+    hostname = "testaccount.us-east-1.privatelink.snowflakecomputing.com"
+    expectation = "http://ocsp.us-east-1.privatelink.snowflakecomputing.com/ocsp_response_cache.json"
+    thread_obj = []
+    for i in range(15):
+        thread_obj.append(ExecPrivatelinkThread(bucket, hostname, expectation, APPLICATION_SNOWSQL))
+
+    for i in range(15):
+        thread_obj[i].start()
+
+    fail_flag = False
+    for i in range(15):
+        thread_obj[i].join()
+        exc = bucket.get(block=False)
+        if exc != 'Success' and not fail_flag:
+            fail_flag = True
+
+    if fail_flag:
+        assert False, "OCSP URL was set incorrectly"
+
+
+class ExecPrivatelinkThread(threading.Thread):
+
+    def __init__(self, bucket, hostname, expectation, client_name):
+        threading.Thread.__init__(self)
+        self.bucket = bucket
+        self.hostname = hostname
+        self.expectation = expectation
+        self.client_name = client_name
+
+    def run(self):
+        SnowflakeConnection.setup_ocsp_privatelink(self.client_name, self.hostname)
+        ocsp_cache_server = os.getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL", None)
+        if ocsp_cache_server is not None and ocsp_cache_server !=\
+                self.expectation:
+            print("Got {0} Expected {1}".format(ocsp_cache_server, self.expectation))
+            self.bucket.put("Fail")
+        else:
+            self.bucket.put("Success")