Snow 916662 design retry timeout config in python connector (#1759)

Author: sfc-gh-yqiu

DESCRIPTION.md

@@ -12,6 +12,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
   - Added support for `use_logical_type` in `write_pandas`.
   - Removed dependencies on pycryptodomex and oscrypto. All connections now go through OpenSSL via the cryptography library, which was already a dependency.
+  - Added the `backoff_policy` argument to `snowflake.connector.connect` allowing for configurable backoff policy between retries of failed requests. See available implementations in the `backoff_policies` module.
+  - Added the `socket_timeout` argument to `snowflake.connector.connect` specifying socket read and connect timeout.
+  - Fixed `login_timeout` and `network_timeout` behaviour. Retries of login and network requests are now properly halted after these timeouts expire.
 
 - v3.3.1(October 16,2023)
 

src/snowflake/connector/auth/_auth.py

@@ -127,8 +127,9 @@ def base_auth_data(
         internal_application_name,
         internal_application_version,
         ocsp_mode,
-        login_timeout,
-        network_timeout=None,
+        login_timeout: int | None = None,
+        network_timeout: int | None = None,
+        socket_timeout: int | None = None,
     ):
         return {
             "data": {
@@ -148,6 +149,7 @@ def base_auth_data(
                     "TRACING": logger.getEffectiveLevel(),
                     "LOGIN_TIMEOUT": login_timeout,
                     "NETWORK_TIMEOUT": network_timeout,
+                    "SOCKET_TIMEOUT": socket_timeout,
                 },
             },
         }
@@ -166,10 +168,14 @@ def authenticate(
         mfa_callback: Callable[[], None] | None = None,
         password_callback: Callable[[], str] | None = None,
         session_parameters: dict[Any, Any] | None = None,
-        timeout: int = 120,
+        # max time waiting for MFA response, currently unused
+        timeout: int | None = None,
     ) -> dict[str, str | int | bool]:
         logger.debug("authenticate")
 
+        if timeout is None:
+            timeout = auth_instance.timeout
+
         if session_parameters is None:
             session_parameters = {}
 
@@ -194,6 +200,7 @@ def authenticate(
             self._rest._connection._ocsp_mode(),
             self._rest._connection._login_timeout,
             self._rest._connection._network_timeout,
+            self._rest._connection._socket_timeout,
         )
 
         body = copy.deepcopy(body_template)
@@ -239,20 +246,12 @@ def authenticate(
             {k: v for (k, v) in body["data"].items() if k != "PASSWORD"},
         )
 
-        # accommodate any authenticator specific timeout requirements here.
-        # login_timeout comes from user configuration.
-        # Between login timeout and auth specific
-        # timeout use whichever value is smaller
-        auth_timeout = min(self._rest._connection.login_timeout, auth_instance.timeout)
-        logger.debug(f"Timeout set to {auth_timeout}")
-
         try:
             ret = self._rest._post_request(
                 url,
                 headers,
                 json.dumps(body),
-                timeout=auth_timeout,
-                socket_timeout=auth_timeout,
+                socket_timeout=auth_instance._socket_timeout,
             )
         except ForbiddenError as err:
             # HTTP 403
@@ -293,7 +292,10 @@ def authenticate(
             def post_request_wrapper(self, url, headers, body) -> None:
                 # get the MFA response
                 self.ret = self._rest._post_request(
-                    url, headers, body, timeout=self._rest._connection.login_timeout
+                    url,
+                    headers,
+                    body,
+                    socket_timeout=auth_instance._socket_timeout,
                 )
 
             # send new request to wait until MFA is approved
@@ -307,6 +309,7 @@ def post_request_wrapper(self, url, headers, body) -> None:
                 while not self.ret or self.ret.get("message") == "Timeout":
                     next(c)
             else:
+                # _post_request should already terminate on timeout, so this is just a safeguard
                 t.join(timeout=timeout)
 
             ret = self.ret
@@ -322,8 +325,7 @@ def post_request_wrapper(self, url, headers, body) -> None:
                     url,
                     headers,
                     json.dumps(body),
-                    timeout=self._rest._connection.login_timeout,
-                    socket_timeout=self._rest._connection.login_timeout,
+                    socket_timeout=auth_instance._socket_timeout,
                 )
             elif not ret or not ret["data"] or not ret["data"].get("token"):
                 # not token is returned.
@@ -363,8 +365,7 @@ def post_request_wrapper(self, url, headers, body) -> None:
                     url,
                     headers,
                     json.dumps(body),
-                    timeout=self._rest._connection.login_timeout,
-                    socket_timeout=self._rest._connection.login_timeout,
+                    socket_timeout=auth_instance._socket_timeout,
                 )
 
         logger.debug("completed authentication")

src/snowflake/connector/auth/by_plugin.py

@@ -18,61 +18,31 @@
 from abc import ABC, abstractmethod
 from enum import Enum, unique
 from os import getenv
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Iterator
 
 from ..errorcode import ER_FAILED_TO_CONNECT_TO_DB
 from ..errors import DatabaseError, Error, OperationalError
 from ..sqlstate import SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED
-from ..time_util import DecorrelateJitterBackoff
+from ..time_util import TimeoutBackoffCtx
 
 if TYPE_CHECKING:
     from .. import SnowflakeConnection
 
 logger = logging.getLogger(__name__)
 
+"""
+Default value for max retry is 1 because
+Python requests module already tries twice
+by default. Unlike JWT where we need to refresh
+token every 10 seconds, general authenticators
+wait for 60 seconds before connection timeout
+per attempt totaling a 240 sec wait time for a non
+JWT based authenticator which is more than enough.
+This can be changed ofcourse using MAX_CNXN_RETRY_ATTEMPTS
+env variable.
+"""
 DEFAULT_MAX_CON_RETRY_ATTEMPTS = 1
-
-
-class AuthRetryCtx:
-    def __init__(self) -> None:
-        self._current_retry_count = 0
-        self._max_retry_attempts = int(
-            getenv("MAX_CON_RETRY_ATTEMPTS", DEFAULT_MAX_CON_RETRY_ATTEMPTS)
-        )
-        self._backoff = DecorrelateJitterBackoff(1, 16)
-        self._current_sleep_time = 1
-
-    def get_current_retry_count(self) -> int:
-        return self._current_retry_count
-
-    def increment_retry(self) -> None:
-        self._current_retry_count += 1
-
-    def should_retry(self) -> bool:
-        """Decides whether to retry connection.
-
-        Default value for max retry is 1 because
-        Python requests module already tries twice
-        by default. Unlike JWT where we need to refresh
-        token every 10 seconds, general authenticators
-        wait for 60 seconds before connection timeout
-        per attempt totaling a 240 sec wait time for a non
-        JWT based authenticator which is more than enough.
-        This can be changed ofcourse using MAX_CNXN_RETRY_ATTEMPTS
-        env variable.
-        """
-        return self._current_retry_count < self._max_retry_attempts
-
-    def next_sleep_duration(self) -> int:
-        self._current_sleep_time = self._backoff.next_sleep(
-            self._current_retry_count, self._current_sleep_time
-        )
-        logger.debug(f"Sleeping for {self._current_sleep_time} seconds")
-        return self._current_sleep_time
-
-    def reset(self) -> None:
-        self._current_retry_count = 0
-        self._current_sleep_time = 1
+DEFAULT_AUTH_CLASS_TIMEOUT = 120
 
 
 @unique
@@ -89,18 +59,38 @@ class AuthType(Enum):
 class AuthByPlugin(ABC):
     """External Authenticator interface."""
 
-    def __init__(self) -> None:
-        self._retry_ctx = AuthRetryCtx()
+    def __init__(
+        self,
+        timeout: int | None = None,
+        backoff_generator: Iterator | None = None,
+        **kwargs,
+    ) -> None:
         self.consent_cache_id_token = False
-        self._timeout: int = 120
+
+        self._retry_ctx = TimeoutBackoffCtx(
+            timeout=timeout if timeout is not None else DEFAULT_AUTH_CLASS_TIMEOUT,
+            max_retry_attempts=kwargs.get(
+                "max_retry_attempts",
+                int(getenv("MAX_CON_RETRY_ATTEMPTS", DEFAULT_MAX_CON_RETRY_ATTEMPTS)),
+            ),
+            backoff_generator=backoff_generator,
+        )
+
+        # some authenticators may want to override socket level timeout
+        # for example, AuthByKeyPair will set this to ensure JWT tokens are refreshed in time
+        # if not None, this will override socket_timeout specified in connection
+        self._socket_timeout = None
 
     @property
     def timeout(self) -> int:
-        return self._timeout
+        """The timeout of _retry_ctx is guaranteed not to be None during AuthByPlugin initialization"""
+        return self._retry_ctx.timeout
 
     @timeout.setter
-    def timeout(self, value: Any) -> None:
-        self._timeout = int(value)
+    def timeout(self) -> None:
+        logger.warning(
+            "Attempting to mutate timeout of AuthByPlugin. Create a new instance with desired parameters instead."
+        )
 
     @property
     @abstractmethod
@@ -207,20 +197,23 @@ def handle_timeout(
         time ranges between 1 and 16 seconds.
         """
 
-        del authenticator, service_name, account, user, password
+        # Some authenticators may not want to delete the parameters to this function
+        # Currently, the only authenticator where this is the case is AuthByKeyPair
+        if kwargs.pop("delete_params", True):
+            del authenticator, service_name, account, user, password
+
         logger.debug("Default timeout handler invoked for authenticator")
-        if not self._retry_ctx.should_retry():
+        if not self._retry_ctx.should_retry:
             error = OperationalError(
-                msg=f"Could not connect to Snowflake backend after {self._retry_ctx.get_current_retry_count()} attempt(s)."
+                msg=f"Could not connect to Snowflake backend after {self._retry_ctx.current_retry_count + 1} attempt(s)."
                 "Aborting",
                 errno=ER_FAILED_TO_CONNECT_TO_DB,
             )
-            self._retry_ctx.reset()
             raise error
         else:
             logger.debug(
-                f"Hit connection timeout, attempt number {self._retry_ctx.get_current_retry_count()}."
+                f"Hit connection timeout, attempt number {self._retry_ctx.current_retry_count + 1}."
                 " Will retry in a bit..."
             )
-            self._retry_ctx.increment_retry()
-            time.sleep(self._retry_ctx.next_sleep_duration())
+            time.sleep(float(self._retry_ctx.current_sleep_time))
+            self._retry_ctx.increment()

src/snowflake/connector/auth/default.py

@@ -21,9 +21,9 @@ def type_(self) -> AuthType:
     def assertion_content(self) -> str:
         return "*********"
 
-    def __init__(self, password: str) -> None:
+    def __init__(self, password: str, **kwargs) -> None:
         """Initializes an instance with a password."""
-        super().__init__()
+        super().__init__(**kwargs)
         self._password: str | None = password
 
     def reset_secrets(self) -> None:

src/snowflake/connector/auth/idtoken.py

@@ -36,9 +36,10 @@ def __init__(
         protocol: str | None,
         host: str | None,
         port: str | None,
+        **kwargs,
     ) -> None:
         """Initialized an instance with an IdToken."""
-        super().__init__()
+        super().__init__(**kwargs)
         self._id_token: str | None = id_token
         self._application = application
         self._protocol = protocol
@@ -62,6 +63,8 @@ def reauthenticate(
             protocol=self._protocol,
             host=self._host,
             port=self._port,
+            timeout=conn.login_timeout,
+            backoff_generator=conn._backoff_generator,
         )
         conn._authenticate(conn.auth_class)
         conn._auth_class.reset_secrets()