SNOW-82268: Privatelink OCSP URL Update

Author: sfc-gh-stakeda

connection.py

@@ -68,6 +68,8 @@
     HeartBeatTimer, get_time_millis)
 from .util_text import split_statements, construct_hostname, parse_account
 
+from snowflake.connector.network import APPLICATION_SNOWSQL
+
 
 def DefaultConverterClass():
     if PY_ISSUE_23517 or IS_WINDOWS:
@@ -599,6 +601,33 @@ def __set_error_attributes(self):
                   callable(getattr(errors, method))]:
             setattr(self, m, getattr(errors, m))
 
+    @staticmethod
+    def setup_ocsp_privatelink(app, hostname):
+        if app == APPLICATION_SNOWSQL:
+            ocsp_cache_server = u'http://ocsp{}/ocsp_response_cache.json'.format(
+                hostname[hostname.index('.'):])
+        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)
+
     def __open_connection(self):
         u"""
         Opens a new network connection
@@ -627,26 +656,7 @@ def __open_connection(self):
                 os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL'])
 
         if self.host.endswith(u".privatelink.snowflakecomputing.com"):
-            ocsp_cache_server = \
-                u'http://ocsp.{}/ocsp_response_cache.json'.format(
-                    self.host)
-            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.setup_ocsp_privatelink(self.application, self.host)
         else:
             if 'SF_OCSP_RESPONSE_CACHE_SERVER_URL' in os.environ:
                 del os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL']

test/test_connection.py

@@ -19,6 +19,9 @@
 except:
     CONNECTION_PARAMETERS_ADMIN = {}
 
+from snowflake.connector.network import APPLICATION_SNOWSQL, CLIENT_NAME
+from snowflake.connector.connection import SnowflakeConnection
+
 
 def test_basic(conn_testaccount):
     """
@@ -543,4 +546,21 @@ def test_disable_request_pooling(db_parameters):
     try:
         assert cnx.disable_request_pooling
     finally:
-        cnx.close()
+        cnx.close()
+
+
+def test_privatelink_ocsp_url_creation():
+    hostname = "testaccount.us-east-1.privatelink.snowflakecomputing.com"
+    SnowflakeConnection.setup_ocsp_privatelink(APPLICATION_SNOWSQL, hostname)
+
+    ocsp_cache_server = os.getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL", None)
+    assert ocsp_cache_server == \
+        "http://ocsp.us-east-1.privatelink.snowflakecomputing.com/ocsp_response_cache.json"
+
+    del os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL']
+
+    SnowflakeConnection.setup_ocsp_privatelink(CLIENT_NAME, hostname)
+    ocsp_cache_server = os.getenv("SF_OCSP_RESPONSE_CACHE_SERVER_URL", None)
+    assert ocsp_cache_server == \
+        "http://ocsp.testaccount.us-east-1.privatelink.snowflakecomputing.com/ocsp_response_cache.json"
+