SNOW-1506571: make connector domain agnostic (#1994)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -11,7 +11,8 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 - v3.12.0(TBD)
   - Set default connection timeout of 10 seconds and socket read timeout of 10 minutes for HTTP calls in file transfer.
   - Optimized `to_pandas()` performance by fully parallel downloading logic.
-  - Fixed a bug that specifying client_session_keep_alive_heartbeat_frequency in snowflake-sqlalchemy could crash the connector
+  - Fixed a bug that specifying client_session_keep_alive_heartbeat_frequency in snowflake-sqlalchemy could crash the connector.
+  - Added support for connectivity to multiple domains.
 
 - v3.11.0(June 17,2024)
   - Added support for `token_file_path` connection parameter to read an OAuth token from a file when connecting to Snowflake.

src/snowflake/connector/connection.py

@@ -52,6 +52,7 @@
 from .config_manager import CONFIG_MANAGER, _get_default_connection_params
 from .connection_diagnostic import ConnectionDiagnostic
 from .constants import (
+    _DOMAIN_NAME_MAP,
     ENV_VAR_PARTNER,
     PARAMETER_AUTOCOMMIT,
     PARAMETER_CLIENT_PREFETCH_THREADS,
@@ -107,6 +108,7 @@
 from .telemetry import TelemetryClient, TelemetryData, TelemetryField
 from .telemetry_oob import TelemetryService
 from .time_util import HeartBeatTimer, get_time_millis
+from .url_util import extract_top_level_domain_from_hostname
 from .util_text import construct_hostname, parse_account, split_statements
 
 DEFAULT_CLIENT_PREFETCH_THREADS = 4
@@ -935,7 +937,7 @@ def __open_connection(self):
                 os.environ["SF_OCSP_RESPONSE_CACHE_SERVER_URL"],
             )
 
-        if self.host.endswith(".privatelink.snowflakecomputing.com"):
+        if ".privatelink.snowflakecomputing." in self.host:
             SnowflakeConnection.setup_ocsp_privatelink(self.application, self.host)
         else:
             if "SF_OCSP_RESPONSE_CACHE_SERVER_URL" in os.environ:
@@ -1182,6 +1184,10 @@ def __config(self, **kwargs):
             if "protocol" not in kwargs:
                 self._protocol = "https"
 
+        logger.info(
+            f"Connecting to {_DOMAIN_NAME_MAP.get(extract_top_level_domain_from_hostname(self._host), 'GLOBAL')} Snowflake domain"
+        )
+
         # If using a custom auth class, we should set the authenticator
         # type to be the same as the custom auth class
         if self._auth_class:

src/snowflake/connector/connection_diagnostic.py

@@ -23,6 +23,7 @@
 
 from .compat import IS_WINDOWS, urlparse
 from .cursor import SnowflakeCursor
+from .url_util import extract_top_level_domain_from_hostname
 from .vendored import urllib3
 
 logger = getLogger(__name__)
@@ -89,15 +90,21 @@ def __init__(
         self.__append_message(
             host_type, f"Host based on specified account: {self.host}"
         )
-        if ".com.snowflakecomputing.com" in self.host:
-            self.host = host.split(".com.snow", 1)[0] + ".com"
+
+        top_level_domain = extract_top_level_domain_from_hostname(host)
+        if (
+            f".{top_level_domain}.snowflakecomputing.{top_level_domain}" in self.host
+        ):  # repeated domain name pattern
+            self.host = (
+                host.split(f".{top_level_domain}.snow", 1)[0] + f".{top_level_domain}"
+            )
             logger.warning(
-                f"Account should not have snowflakecomputing.com in it. You provided {host}.  "
+                f"Account should not have snowflakecomputing.{top_level_domain} in it. You provided {host}.  "
                 f"Continuing with fixed host."
             )
             self.__append_message(
                 host_type,
-                f"We removed extra .snowflakecomputing.com and will continue with host: "
+                f"We removed extra .snowflakecomputing.{top_level_domain} and will continue with host: "
                 f"{self.host}",
             )
         else:
@@ -183,7 +190,7 @@ def __init__(
             self.ocsp_urls.append(f"ocsp.{self.host}")
             self.allowlist_sql = "select system$allowlist_privatelink();"
         else:
-            self.ocsp_urls.append("ocsp.snowflakecomputing.com")
+            self.ocsp_urls.append(f"ocsp.snowflakecomputing.{top_level_domain}")
 
         self.allowlist_retrieval_success: bool = False
         self.cursor: SnowflakeCursor | None = None

src/snowflake/connector/constants.py

@@ -36,6 +36,11 @@
 DBAPI_TYPE_NUMBER = 2
 DBAPI_TYPE_TIMESTAMP = 3
 
+_DEFAULT_HOSTNAME_TLD = "com"
+_CHINA_HOSTNAME_TLD = "cn"
+_TOP_LEVEL_DOMAIN_REGEX = r"\.[a-zA-Z]{1,63}$"
+_SNOWFLAKE_HOST_SUFFIX_REGEX = r"snowflakecomputing(\.[a-zA-Z]{1,63}){1,2}$"
+
 
 class FieldType(NamedTuple):
     name: str
@@ -420,3 +425,6 @@ class IterUnit(Enum):
 # TODO: all env variables definitions should be here
 ENV_VAR_PARTNER = "SF_PARTNER"
 ENV_VAR_TEST_MODE = "SNOWFLAKE_TEST_MODE"
+
+
+_DOMAIN_NAME_MAP = {_DEFAULT_HOSTNAME_TLD: "GLOBAL", _CHINA_HOSTNAME_TLD: "CHINA"}

src/snowflake/connector/network.py

@@ -11,6 +11,7 @@
 import itertools
 import json
 import logging
+import re
 import time
 import traceback
 import uuid
@@ -46,6 +47,7 @@
     urlparse,
 )
 from .constants import (
+    _SNOWFLAKE_HOST_SUFFIX_REGEX,
     HTTP_HEADER_ACCEPT,
     HTTP_HEADER_CONTENT_TYPE,
     HTTP_HEADER_SERVICE_NAME,
@@ -156,6 +158,7 @@
 REQUEST_GUID = "request_guid"
 SNOWFLAKE_HOST_SUFFIX = ".snowflakecomputing.com"
 
+
 SNOWFLAKE_CONNECTOR_VERSION = SNOWFLAKE_CONNECTOR_VERSION
 PYTHON_VERSION = PYTHON_VERSION
 OPERATING_SYSTEM = OPERATING_SYSTEM
@@ -856,7 +859,7 @@ def add_retry_params(self, full_url: str) -> str:
     def add_request_guid(full_url: str) -> str:
         """Adds request_guid parameter for HTTP request tracing."""
         parsed_url = urlparse(full_url)
-        if not parsed_url.hostname.endswith(SNOWFLAKE_HOST_SUFFIX):
+        if not re.search(_SNOWFLAKE_HOST_SUFFIX_REGEX, parsed_url.hostname):
             return full_url
         request_guid = str(uuid.uuid4())
         suffix = urlencode({REQUEST_GUID: request_guid})

src/snowflake/connector/ocsp_snowflake.py

@@ -62,7 +62,7 @@
 from .backoff_policies import exponential_backoff
 from .cache import SFDictCache, SFDictFileCache
 from .telemetry import TelemetryField, generate_telemetry_data_dict
-from .url_util import url_encode_str
+from .url_util import extract_top_level_domain_from_hostname, url_encode_str
 
 
 class OCSPResponseValidationResult(NamedTuple):
@@ -268,15 +268,20 @@ def generate_telemetry_data(
 class OCSPServer:
     MAX_RETRY = int(os.getenv("OCSP_MAX_RETRY", "3"))
 
-    def __init__(self) -> None:
-        self.DEFAULT_CACHE_SERVER_URL = "http://ocsp.snowflakecomputing.com"
+    def __init__(self, **kwargs) -> None:
+        top_level_domain = kwargs.pop(
+            "top_level_domain", constants._DEFAULT_HOSTNAME_TLD
+        )
+        self.DEFAULT_CACHE_SERVER_URL = (
+            f"http://ocsp.snowflakecomputing.{top_level_domain}"
+        )
         """
         The following will change to something like
         http://ocspssd.snowflakecomputing.com/ocsp/
         once the endpoint is up in the backend
         """
         self.NEW_DEFAULT_CACHE_SERVER_BASE_URL = (
-            "https://ocspssd.snowflakecomputing.com/ocsp/"
+            f"https://ocspssd.snowflakecomputing.{top_level_domain}/ocsp/"
         )
         if not OCSPServer.is_enabled_new_ocsp_endpoint():
             self.CACHE_SERVER_URL = os.getenv(
@@ -307,12 +312,13 @@ def reset_ocsp_endpoint(self, hname) -> None:
         on the hostname the customer is trying to connect to. The deployment or in case of client failover, the
         replication ID is copied from the hostname.
         """
-        if hname.endswith("privatelink.snowflakecomputing.com"):
+        top_level_domain = extract_top_level_domain_from_hostname(hname)
+        if "privatelink.snowflakecomputing." in hname:
             temp_ocsp_endpoint = "".join(["https://ocspssd.", hname, "/ocsp/"])
-        elif hname.endswith("global.snowflakecomputing.com"):
+        elif "global.snowflakecomputing." in hname:
             rep_id_begin = hname[hname.find("-") :]
             temp_ocsp_endpoint = "".join(["https://ocspssd", rep_id_begin, "/ocsp/"])
-        elif not hname.endswith("snowflakecomputing.com"):
+        elif not hname.endswith(f"snowflakecomputing.{top_level_domain}"):
             temp_ocsp_endpoint = self.NEW_DEFAULT_CACHE_SERVER_BASE_URL
         else:
             hname_wo_acc = hname[hname.find(".") :]
@@ -832,8 +838,8 @@ class SnowflakeOCSP:
 
     OCSP_WHITELIST = re.compile(
         r"^"
-        r"(.*\.snowflakecomputing\.com$"
-        r"|(?:|.*\.)s3.*\.amazonaws\.com$"  # start with s3 or .s3 in the middle
+        r"(.*\.snowflakecomputing(\.[a-zA-Z]{1,63}){1,2}$"
+        r"|(?:|.*\.)s3.*\.amazonaws(\.[a-zA-Z]{1,63}){1,2}$"  # start with s3 or .s3 in the middle
         r"|.*\.okta\.com$"
         r"|(?:|.*\.)storage\.googleapis\.com$"
         r"|.*\.blob\.core\.windows\.net$"
@@ -881,14 +887,19 @@ def __init__(
         use_ocsp_cache_server=None,
         use_post_method: bool = True,
         use_fail_open: bool = True,
+        **kwargs,
     ) -> None:
         self.test_mode = os.getenv("SF_OCSP_TEST_MODE", None)
 
         if self.test_mode == "true":
             logger.debug("WARNING - DRIVER CONFIGURED IN TEST MODE")
 
         self._use_post_method = use_post_method
-        self.OCSP_CACHE_SERVER = OCSPServer()
+        self.OCSP_CACHE_SERVER = OCSPServer(
+            top_level_domain=extract_top_level_domain_from_hostname(
+                kwargs.pop("hostname", None)
+            )
+        )
 
         self.debug_ocsp_failure_url = None
 

src/snowflake/connector/ssl_wrap_socket.py

@@ -87,6 +87,7 @@ def ssl_wrap_socket_with_ocsp(*args: Any, **kwargs: Any) -> WrappedSocket:
         v = SFOCSP(
             ocsp_response_cache_uri=FEATURE_OCSP_RESPONSE_CACHE_FILE_NAME,
             use_fail_open=FEATURE_OCSP_MODE == OCSPMode.FAIL_OPEN,
+            hostname=server_hostname,
         ).validate(server_hostname, ret.connection)
         if not v:
             raise OperationalError(

src/snowflake/connector/url_util.py

@@ -8,6 +8,8 @@
 import urllib.parse
 from logging import getLogger
 
+from .constants import _TOP_LEVEL_DOMAIN_REGEX
+
 logger = getLogger(__name__)
 
 
@@ -41,3 +43,11 @@ def url_encode_str(target: str | None) -> str:
         logger.debug("The string to be URL encoded is None")
         return ""
     return urllib.parse.quote_plus(target, safe="")
+
+
+def extract_top_level_domain_from_hostname(hostname: str | None = None) -> str:
+    if not hostname:
+        return "com"
+    # RFC1034 for TLD spec, and https://data.iana.org/TLD/tlds-alpha-by-domain.txt for full TLD list
+    match = re.search(_TOP_LEVEL_DOMAIN_REGEX, hostname)
+    return (match.group(0)[1:] if match else "com").lower()

src/snowflake/connector/util_text.py

@@ -240,7 +240,12 @@ def construct_hostname(region: str | None, account: str) -> str:
     if region:
         if account.find(".") > 0:
             account = account[0 : account.find(".")]
-        host = f"{account}.{region}.snowflakecomputing.com"
+        top_level_domain = (
+            "com"
+            if not any(substring in region for substring in ["cn-", "CN-"])
+            else "cn"
+        )
+        host = f"{account}.{region}.snowflakecomputing.{top_level_domain}"
     else:
         host = f"{account}.snowflakecomputing.com"
     return host

test/integ/test_connection.py

@@ -11,6 +11,7 @@
 import pathlib
 import queue
 import stat
+import tempfile
 import threading
 import warnings
 import weakref
@@ -39,6 +40,7 @@
 from snowflake.connector.telemetry import TelemetryField
 
 from ..randomize import random_string
+from .conftest import RUNNING_ON_GH
 
 try:  # pragma: no cover
     from ..parameters import CONNECTION_PARAMETERS_ADMIN
@@ -1370,3 +1372,35 @@ def test_token_file_path(tmp_path, db_parameters):
     assert conn._token == fake_token
     conn = snowflake.connector.connect(**db_parameters, token_file_path=token_file_path)
     assert conn._token == fake_token
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.skipif(not RUNNING_ON_GH, reason="no ocsp in the environment")
+def test_mock_non_existing_server(conn_cnx, caplog):
+    from snowflake.connector.cache import SFDictCache
+
+    # disabling local cache and pointing ocsp cache server to a non-existing url
+    # connection should still work as it will directly validate the certs against CA servers
+    with tempfile.NamedTemporaryFile() as tmp, caplog.at_level(logging.DEBUG):
+        with mock.patch(
+            "snowflake.connector.url_util.extract_top_level_domain_from_hostname",
+            return_value="nonexistingtopleveldomain",
+        ):
+            with mock.patch(
+                "snowflake.connector.ocsp_snowflake.OCSP_RESPONSE_VALIDATION_CACHE",
+                SFDictCache(),
+            ):
+                with mock.patch(
+                    "snowflake.connector.ocsp_snowflake.OCSPCache.OCSP_RESPONSE_CACHE_FILE_NAME",
+                    tmp.name,
+                ):
+                    with conn_cnx():
+                        pass
+        assert all(
+            s in caplog.text
+            for s in [
+                "Failed to read OCSP response cache file",
+                "It will validate with OCSP server.",
+                "writing OCSP response cache file to",
+            ]
+        )