Improve error logging + messages

1. Change logging base to only log information about duration, number of retries etc if relevant;
2. Now log `x-databricks-error-or-redirect-message` and `x-databricks-reason-phrase` as well as thrift server error message (we only expect one or the other to be set it seems)
3. Forward port as it seemed the simplest so we could just re-use logging code

* New unit tests
* Re-running integration tests:

Author: NiallEgan

cmdexec/clients/python/src/databricks/sql/client.py

@@ -79,17 +79,8 @@ def __init__(self,
                                                    kwargs.get("_user_agent_entry"))
 
         base_headers = [("User-Agent", useragent_header)] + authorization_header
-        self.thrift_backend = ThriftBackend(
-            self.host,
-            self.port,
-            http_path, (metadata or []) + base_headers,
-            _tls_verify_hostname=kwargs.get("_tls_verify_hostname"),
-            _tls_trusted_ca_file=kwargs.get("_tls_trusted_ca_file"),
-            _tls_client_cert_file=kwargs.get("_tls_client_cert_file"),
-            _tls_client_cert_key_file=kwargs.get("_tls_client_cert_key_file"),
-            _tls_client_cert_key_password=kwargs.get("_tls_client_cert_key_password"),
-            _connection_uri=kwargs.get("_connection_uri"),
-            _retry_stop_after_attempts_count=kwargs.get("_retry_stop_after_attempts_count", 24))
+        self.thrift_backend = ThriftBackend(self.host, self.port, http_path,
+                                            (metadata or []) + base_headers, **kwargs)
 
         self._session_handle = self.thrift_backend.open_session()
         self.open = True

cmdexec/clients/python/src/databricks/sql/thrift_backend.py

@@ -1,5 +1,6 @@
 from decimal import Decimal
 import logging
+import math
 import time
 import threading
 from uuid import uuid4
@@ -14,16 +15,29 @@
 
 from databricks.sql.thrift_api.TCLIService import TCLIService, ttypes
 from databricks.sql import *
-from databricks.sql.utils import ArrowQueue, ExecuteResponse
+from databricks.sql.utils import ArrowQueue, ExecuteResponse, _bound, RequestErrorInfo, NoRetryReason
 
 logger = logging.getLogger(__name__)
 
+THRIFT_ERROR_MESSAGE_HEADER = "x-thriftserver-error-message"
+DATABRICKS_ERROR_OR_REDIRECT_HEADER = "x-databricks-error-or-redirect-message"
+DATABRICKS_REASON_HEADER = "x-databricks-reason-phrase"
+
+# see Connection.__init__ for parameter descriptions.
+# - Min/Max avoids unsustainable configs (sane values are far more constrained)
+# - 900s attempts-duration lines up w ODBC/JDBC drivers (for cluster startup > 10 mins)
+_retry_policy = {                          # (type, default, min, max)
+    "_retry_delay_min":                     (float, 1, 0.1, 60),
+    "_retry_delay_max":                     (float, 60, 5, 3600),
+    "_retry_stop_after_attempts_count":     (int, 30, 1, 60),
+    "_retry_stop_after_attempts_duration":  (float, 900, 1, 86400),
+}
+
 
 class ThriftBackend:
     CLOSED_OP_STATE = ttypes.TOperationState.CLOSED_STATE
     ERROR_OP_STATE = ttypes.TOperationState.ERROR_STATE
     BIT_MASKS = [1, 2, 4, 8, 16, 32, 64, 128]
-    ERROR_MSG_HEADER = "X-Thriftserver-Error-Message"
 
     def __init__(self, server_hostname: str, port, http_path: str, http_headers, **kwargs):
         # Internal arguments in **kwargs:
@@ -43,8 +57,18 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
         #   See https://docs.python.org/3/library/ssl.html#ssl.SSLContext.load_cert_chain
         # _connection_uri
         #   Overrides server_hostname and http_path.
-        # _retry_stop_after_attempts_count
-        #  The maximum number of times we should retry retryable requests (defaults to 24)
+        # RETRY/ATTEMPT POLICY
+        # _retry_delay_min                      (default: 1)
+        # _retry_delay_max                      (default: 60)
+        #   {min,max} pre-retry delay bounds
+        # _retry_stop_after_attempts_count      (default: 30)
+        #   total max attempts during retry sequence
+        # _retry_stop_after_attempts_duration   (default: 900)
+        #   total max wait duration during retry sequence
+        #   (Note this will stop _before_ intentionally exceeding; thus if the
+        #   next calculated pre-retry delay would go past
+        #   _retry_stop_after_attempts_duration, stop now.)
+        #
 
         port = port or 443
         if kwargs.get("_connection_uri"):
@@ -55,7 +79,7 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
         else:
             raise ValueError("No valid connection settings.")
 
-        self._retry_stop_after_attempts_count = kwargs.get("_retry_stop_after_attempts_count", 24)
+        self._initialize_retry_args(kwargs)
 
         # Configure tls context
         ssl_context = create_default_context(cafile=kwargs.get("_tls_trusted_ca_file"))
@@ -95,55 +119,156 @@ def __init__(self, server_hostname: str, port, http_path: str, http_headers, **k
 
         self._request_lock = threading.RLock()
 
+    def _initialize_retry_args(self, kwargs):
+        # Configure retries & timing: use user-settings or defaults, and bound
+        # by policy. Log.warn when given param gets restricted.
+        for (key, (type_, default, min, max)) in _retry_policy.items():
+            given_or_default = type_(kwargs.get(key, default))
+            bound = _bound(min, max, given_or_default)
+            setattr(self, key, bound)
+            logger.debug('retry parameter: {} given_or_default {}'.format(key, given_or_default))
+            if bound != given_or_default:
+                logger.warn('Override out of policy retry parameter: ' +
+                            '{} given {}, restricted to {}'.format(key, given_or_default, bound))
+
+        # Fail on retry delay min > max; consider later adding fail on min > duration?
+        if self._retry_stop_after_attempts_count > 1 \
+                and self._retry_delay_min > self._retry_delay_max:
+            raise ValueError(
+                "Invalid configuration enables retries with retry delay min(={}) > max(={})".format(
+                    self._retry_delay_min, self._retry_delay_max))
+
     @staticmethod
     def _check_response_for_error(response):
         if response.status and response.status.statusCode in \
                 [ttypes.TStatusCode.ERROR_STATUS, ttypes.TStatusCode.INVALID_HANDLE_STATUS]:
             raise DatabaseError(response.status.errorMessage)
 
+    @staticmethod
+    def _extract_error_message_from_headers(headers):
+        err_msg = ""
+        if THRIFT_ERROR_MESSAGE_HEADER in headers:
+            err_msg = headers[THRIFT_ERROR_MESSAGE_HEADER]
+        if DATABRICKS_ERROR_OR_REDIRECT_HEADER in headers:
+            if err_msg:  # We don't expect both to be set, but log both here just in case
+                err_msg = "Thriftserver error: {}, Databricks error: {}".format(
+                    err_msg, headers[DATABRICKS_ERROR_OR_REDIRECT_HEADER])
+            else:
+                err_msg = headers[DATABRICKS_ERROR_OR_REDIRECT_HEADER]
+            if DATABRICKS_REASON_HEADER in headers:
+                err_msg += ": " + headers[DATABRICKS_REASON_HEADER]
+        return err_msg
+
+    def _handle_request_error(self, error_info, attempt, elapsed):
+        max_attempts = self._retry_stop_after_attempts_count
+        max_duration_s = self._retry_stop_after_attempts_duration
+
+        if error_info.retry_delay is not None and elapsed + error_info.retry_delay > max_duration_s:
+            no_retry_reason = NoRetryReason.OUT_OF_TIME
+        elif error_info.retry_delay is not None and attempt >= max_attempts:
+            no_retry_reason = NoRetryReason.OUT_OF_ATTEMPTS
+        elif error_info.retry_delay is None:
+            no_retry_reason = NoRetryReason.NOT_RETRYABLE
+        else:
+            no_retry_reason = None
+
+        full_error_info_str = error_info.full_info_logging_str(
+            no_retry_reason, attempt, max_attempts, elapsed, max_duration_s)
+
+        if no_retry_reason is not None:
+            user_friendly_error_message = error_info.user_friendly_error_message(
+                no_retry_reason, attempt, elapsed)
+            logger.info("{}: {}".format(user_friendly_error_message, full_error_info_str))
+
+            raise OperationalError(user_friendly_error_message, error_info.error)
+
+        logger.info("Retrying request after error in {} seconds: {}".format(
+            error_info.retry_delay, full_error_info_str))
+        time.sleep(error_info.retry_delay)
+
     # FUTURE: Consider moving to https://github.com/litl/backoff or
-    # https://github.com/jd/tenacity for retry logic. Otherwise, copy from
-    # v1 client.
-    def make_request(self, method, request, attempt_number=1):
-        try:
+    # https://github.com/jd/tenacity for retry logic.
+    def make_request(self, method, request):
+        """Execute given request, attempting retries when receiving HTTP 429/503.
+
+        For delay between attempts, honor the given Retry-After header, but with bounds.
+        Use lower bound of expontial-backoff based on _retry_delay_min,
+        and upper bound of _retry_delay_max.
+        Will stop retry attempts if total elapsed time + next retry delay would exceed
+        _retry_stop_after_attempts_duration.
+        """
+        # basic strategy: build range iterator rep'ing number of available
+        # retries. bounds can be computed from there. iterate over it with
+        # retries until success or final failure achieved.
+
+        t0 = time.time()
+
+        def get_elapsed():
+            return time.time() - t0
+
+        def extract_retry_delay(attempt):
+            # encapsulate retry checks, returns None || delay-in-secs
+            # Retry IFF 429/503 code + Retry-After header set
+            http_code = getattr(self._transport, "code", None)
+            retry_after = getattr(self._transport, "headers", {}).get("Retry-After")
+            if http_code in [429, 503] and retry_after:
+                # bound delay (seconds) by [min_delay*1.5^(attempt-1), max_delay]
+                delay = int(retry_after)
+                delay = max(delay, self._retry_delay_min * math.pow(1.5, attempt - 1))
+                delay = min(delay, self._retry_delay_max)
+                return delay
+            return None
+
+        def attempt_request(attempt):
+            # splits out lockable attempt, from delay & retry loop
+            # returns tuple: (method_return, delay_fn(), error, error_message)
+            # - non-None method_return -> success, return and be done
+            # - non-None retry_delay -> sleep delay before retry
+            # - error, error_message always set when available
+            try:
+                logger.debug("Sending request: {}".format(request))
+                response = method(request)
+                logger.debug("Received response: {}".format(response))
+                return response
+            except Exception as error:
+                retry_delay = extract_retry_delay(attempt)
+                error_message = ThriftBackend._extract_error_message_from_headers(
+                    getattr(self._transport, "headers", {}))
+                return RequestErrorInfo(
+                    error=error,
+                    error_message=error_message,
+                    retry_delay=retry_delay,
+                    http_code=getattr(self._transport, "code", None),
+                    method=method.__name__,
+                    request=request)
+
+        # The real work:
+        # - for each available attempt:
+        #       lock-and-attempt
+        #       return on success
+        #       if available: bounded delay and retry
+        #       if not: raise error
+        max_attempts = self._retry_stop_after_attempts_count
+
+        # use index-1 counting for logging/human consistency
+        for attempt in range(1, max_attempts + 1):
             # We have a lock here because .cancel can be called from a separate thread.
             # We do not want threads to be simultaneously sharing the Thrift Transport
             # because we use its state to determine retries
-            self._request_lock.acquire()
-            logger.warning("Sending request: {}".format(request))
-            response = method(request)
-            logger.warning("Received response: {}".format(response))
-            ThriftBackend._check_response_for_error(response)
-            return response
-        except Exception as error:
-            # _transport.code isn't necessarily set :(
-            code_and_headers_is_set = hasattr(self._transport, 'code') \
-                                      and hasattr(self._transport, 'headers')
-            # We only retry if a Retry-After header is set
-            if code_and_headers_is_set and self._transport.code in [503, 429] and \
-                    "Retry-After" in self._transport.headers and \
-                    attempt_number <= self._retry_stop_after_attempts_count - 1:
-                retry_time_seconds = int(self._transport.headers["Retry-After"])
-                if self.ERROR_MSG_HEADER in self._transport.headers:
-                    error_message = self._transport.headers[self.ERROR_MSG_HEADER]
-                else:
-                    error_message = str(error)
-                logger.warning("Received retryable error during {}. Request: {} Error: {}".format(
-                    method, request, error_message))
-                logger.warning("Retrying in {} seconds. This is attempt number {}".format(
-                    retry_time_seconds, attempt_number))
-                time.sleep(retry_time_seconds)
-                return self.make_request(method, request, attempt_number + 1)
-            else:
-                logger.error("Received error when issuing: {}".format(request))
-                if hasattr(self._transport, "headers") and \
-                   self.ERROR_MSG_HEADER in self._transport.headers:
-                    error_message = self._transport.headers[self.ERROR_MSG_HEADER]
-                    raise OperationalError("Error during Thrift request: {}".format(error_message))
-                else:
-                    raise OperationalError("Error during Thrift request", error)
-        finally:
-            self._request_lock.release()
+            with self._request_lock:
+                response_or_error_info = attempt_request(attempt)
+            elapsed = get_elapsed()
+
+            # conditions: success, non-retry-able, no-attempts-left, no-time-left, delay+retry
+            if not isinstance(response_or_error_info, RequestErrorInfo):
+                # log nothing here, presume that main request logging covers
+                response = response_or_error_info
+                ThriftBackend._check_response_for_error(response)
+                return response
+
+            error_info = response_or_error_info
+            # The error handler will either sleep or throw an exception
+            self._handle_request_error(error_info, attempt, elapsed)
 
     def _check_protocol_version(self, t_open_session_resp):
         protocol_version = t_open_session_resp.serverProtocolVersion

cmdexec/clients/python/src/databricks/sql/utils.py

@@ -1,4 +1,5 @@
-from collections import namedtuple
+from collections import namedtuple, OrderedDict
+from enum import Enum
 
 import pyarrow
 
@@ -32,3 +33,78 @@ def remaining_rows(self) -> pyarrow.Table:
 ExecuteResponse = namedtuple(
     'ExecuteResponse', 'status has_been_closed_server_side has_more_rows description '
     'command_handle arrow_queue arrow_schema')
+
+
+def _bound(min_x, max_x, x):
+    """Bound x by [min_x, max_x]
+
+    min_x or max_x being None means unbounded in that respective side.
+    """
+    if min_x is None and max_x is None:
+        return x
+    if min_x is None:
+        return min(max_x, x)
+    if max_x is None:
+        return max(min_x, x)
+    return min(max_x, max(min_x, x))
+
+
+class NoRetryReason(Enum):
+    OUT_OF_TIME = "out of time"
+    OUT_OF_ATTEMPTS = "out of attempts"
+    NOT_RETRYABLE = "non-retryable error"
+
+
+class RequestErrorInfo(
+        namedtuple("RequestErrorInfo_",
+                   "error error_message retry_delay http_code method request")):
+    @property
+    def request_session_id(self):
+        if hasattr(self.request, "sessionHandle"):
+            return self.request.sessionHandle.sessionId.guid
+        else:
+            return None
+
+    @property
+    def request_query_id(self):
+        if hasattr(self.request, "operationHandle"):
+            return self.request.operationHandle.operationId.guid
+        else:
+            return None
+
+    def full_info_logging_str(self, no_retry_reason, attempt, max_attempts, elapsed, max_duration):
+        log_base_data_dict = OrderedDict([
+            ("Method", self.method),
+            ("Session-id", self.request_session_id),
+            ("Query-id", self.request_query_id),
+            ("HTTP-code", self.http_code),
+            ("Error-message", self.error_message),
+            ("Original-exception", self.error),
+        ])
+
+        if no_retry_reason is not None:
+            log_base_data_dict["No-retry-reason"] = no_retry_reason.value
+        else:
+            log_base_data_dict["Bounded-retry-delay"] = self.retry_delay
+            log_base_data_dict["Attempt"] = "{}/{}".format(attempt, max_attempts)
+            log_base_data_dict["Elapsed-seconds"] = "{}/{}".format(elapsed, max_duration)
+
+        log_base = "; ".join(["{}: {}".format(k, v) for k, v in log_base_data_dict.items()])
+
+        return log_base
+
+    def user_friendly_error_message(self, no_retry_reason, attempt, elapsed):
+        # This should be kept at the level that is appropriate to return to a Redash user
+        user_friendly_error_message = "Error during request to server"
+        if self.error_message:
+            user_friendly_error_message = "{}: {}".format(user_friendly_error_message,
+                                                          self.error_message)
+
+        if no_retry_reason is NoRetryReason.OUT_OF_ATTEMPTS:
+            user_friendly_error_message = "{}: After {} retry attempts, retries are exhausted".format(
+                user_friendly_error_message, attempt)
+        elif no_retry_reason is NoRetryReason.OUT_OF_TIME:
+            user_friendly_error_message = "{}: After {} seconds, maximum retry duration will be exceeded".format(
+                user_friendly_error_message, elapsed)
+
+        return user_friendly_error_message