Add timeout hack to mitigate timeouts

Author: capitancambio

src/databricks/sql/auth/thrift_http_client.py

@@ -3,6 +3,7 @@
 
 
 import thrift
+import thrift.transport.THttpClient
 
 import urllib.parse, six, base64
 

src/databricks/sql/thrift_backend.py

@@ -1,34 +1,27 @@
-from decimal import Decimal
 import errno
 import logging
 import math
-import time
+import os
 import threading
-import lz4.frame
+import time
+from decimal import Decimal
 from ssl import CERT_NONE, CERT_REQUIRED, create_default_context
 from typing import List, Union
 
+import databricks.sql.auth.thrift_http_client
+import lz4.frame
 import pyarrow
-import thrift.transport.THttpClient
 import thrift.protocol.TBinaryProtocol
+import thrift.transport.THttpClient
 import thrift.transport.TSocket
 import thrift.transport.TTransport
-
-import databricks.sql.auth.thrift_http_client
+from databricks.sql import *
 from databricks.sql.auth.authenticators import AuthProvider
 from databricks.sql.thrift_api.TCLIService import TCLIService, ttypes
-from databricks.sql import *
-from databricks.sql.thrift_api.TCLIService.TCLIService import (
-    Client as TCLIServiceClient,
-)
-
-from databricks.sql.utils import (
-    ArrowQueue,
-    ExecuteResponse,
-    _bound,
-    RequestErrorInfo,
-    NoRetryReason,
-)
+from databricks.sql.thrift_api.TCLIService.TCLIService import \
+    Client as TCLIServiceClient
+from databricks.sql.utils import (ArrowQueue, ExecuteResponse, NoRetryReason,
+                                  RequestErrorInfo, _bound)
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +31,9 @@
 
 TIMESTAMP_AS_STRING_CONFIG = "spark.thriftserver.arrowBasedRowSet.timestampAsString"
 
+# HACK!
+THRIFT_SOCKET_TIMEOUT = os.getenv("THRIFT_SOCKET_TIMEOUT", None)
+
 # 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)
@@ -114,13 +110,9 @@ def __init__(
 
         self.staging_allowed_local_path = staging_allowed_local_path
         self._initialize_retry_args(kwargs)
-        self._use_arrow_native_complex_types = kwargs.get(
-            "_use_arrow_native_complex_types", True
-        )
+        self._use_arrow_native_complex_types = kwargs.get("_use_arrow_native_complex_types", True)
         self._use_arrow_native_decimals = kwargs.get("_use_arrow_native_decimals", True)
-        self._use_arrow_native_timestamps = kwargs.get(
-            "_use_arrow_native_timestamps", True
-        )
+        self._use_arrow_native_timestamps = kwargs.get("_use_arrow_native_timestamps", True)
 
         # Configure tls context
         ssl_context = create_default_context(cafile=kwargs.get("_tls_trusted_ca_file"))
@@ -152,7 +144,10 @@ def __init__(
             ssl_context=ssl_context,
         )
 
-        timeout = kwargs.get("_socket_timeout")
+        # HACK!
+        timeout = kwargs.get("_socket_timeout") or THRIFT_SOCKET_TIMEOUT
+        logger.info(f"Setting timeout HACK! to {timeout}")
+
         # setTimeout defaults to None (i.e. no timeout), and is expected in ms
         self._transport.setTimeout(timeout and (float(timeout) * 1000.0))
 
@@ -175,15 +170,11 @@ def _initialize_retry_args(self, kwargs):
             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)
-            )
+            logger.debug("retry parameter: {} given_or_default {}".format(key, given_or_default))
             if bound != given_or_default:
                 logger.warning(
                     "Override out of policy retry parameter: "
-                    + "{} given {}, restricted to {}".format(
-                        key, given_or_default, bound
-                    )
+                    + "{} given {}, restricted to {}".format(key, given_or_default, bound)
                 )
 
         # Fail on retry delay min > max; consider later adding fail on min > duration?
@@ -211,9 +202,7 @@ def _extract_error_message_from_headers(headers):
         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
+            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]
                 )
@@ -406,10 +395,7 @@ def _check_initial_namespace(self, catalog, schema, response):
         if not (catalog or schema):
             return
 
-        if (
-            response.serverProtocolVersion
-            < ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4
-        ):
+        if response.serverProtocolVersion < ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4:
             raise InvalidServerResponseError(
                 "Setting initial namespace not supported by the DBR version, "
                 "Please use a Databricks SQL endpoint or a cluster with DBR >= 9.0."
@@ -424,10 +410,7 @@ def _check_initial_namespace(self, catalog, schema, response):
 
     def _check_session_configuration(self, session_configuration):
         # This client expects timetampsAsString to be false, so we do not allow users to modify that
-        if (
-            session_configuration.get(TIMESTAMP_AS_STRING_CONFIG, "false").lower()
-            != "false"
-        ):
+        if session_configuration.get(TIMESTAMP_AS_STRING_CONFIG, "false").lower() != "false":
             raise Error(
                 "Invalid session configuration: {} cannot be changed "
                 "while using the Databricks SQL connector, it must be false not {}".format(
@@ -439,18 +422,14 @@ def _check_session_configuration(self, session_configuration):
     def open_session(self, session_configuration, catalog, schema):
         try:
             self._transport.open()
-            session_configuration = {
-                k: str(v) for (k, v) in (session_configuration or {}).items()
-            }
+            session_configuration = {k: str(v) for (k, v) in (session_configuration or {}).items()}
             self._check_session_configuration(session_configuration)
             # We want to receive proper Timestamp arrow types.
             # We set it also in confOverlay in TExecuteStatementReq on a per query basic,
             # but it doesn't hurt to also set for the whole session.
             session_configuration[TIMESTAMP_AS_STRING_CONFIG] = "false"
             if catalog or schema:
-                initial_namespace = ttypes.TNamespace(
-                    catalogName=catalog, schemaName=schema
-                )
+                initial_namespace = ttypes.TNamespace(catalogName=catalog, schemaName=schema)
             else:
                 initial_namespace = None
 
@@ -476,9 +455,7 @@ def close_session(self, session_handle) -> None:
         finally:
             self._transport.close()
 
-    def _check_command_not_in_error_or_closed_state(
-        self, op_handle, get_operations_resp
-    ):
+    def _check_command_not_in_error_or_closed_state(self, op_handle, get_operations_resp):
         if get_operations_resp.operationState == ttypes.TOperationState.ERROR_STATE:
             if get_operations_resp.displayMessage:
                 raise ServerOperationError(
@@ -513,17 +490,11 @@ def _poll_for_status(self, op_handle):
 
     def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, description):
         if t_row_set.columns is not None:
-            (
-                arrow_table,
-                num_rows,
-            ) = ThriftBackend._convert_column_based_set_to_arrow_table(
+            (arrow_table, num_rows,) = ThriftBackend._convert_column_based_set_to_arrow_table(
                 t_row_set.columns, description
             )
         elif t_row_set.arrowBatches is not None:
-            (
-                arrow_table,
-                num_rows,
-            ) = ThriftBackend._convert_arrow_based_set_to_arrow_table(
+            (arrow_table, num_rows,) = ThriftBackend._convert_arrow_based_set_to_arrow_table(
                 t_row_set.arrowBatches, lz4_compressed, schema_bytes
             )
         else:
@@ -534,9 +505,7 @@ def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, descripti
     def _convert_decimals_in_arrow_table(table, description):
         for (i, col) in enumerate(table.itercolumns()):
             if description[i][1] == "decimal":
-                decimal_col = col.to_pandas().apply(
-                    lambda v: v if v is None else Decimal(v)
-                )
+                decimal_col = col.to_pandas().apply(lambda v: v if v is None else Decimal(v))
                 precision, scale = description[i][4], description[i][5]
                 assert scale is not None
                 assert precision is not None
@@ -549,9 +518,7 @@ def _convert_decimals_in_arrow_table(table, description):
         return table
 
     @staticmethod
-    def _convert_arrow_based_set_to_arrow_table(
-        arrow_batches, lz4_compressed, schema_bytes
-    ):
+    def _convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema_bytes):
         ba = bytearray()
         ba += schema_bytes
         n_rows = 0
@@ -597,9 +564,7 @@ def _convert_column_to_arrow_array(t_col):
         for field in field_name_to_arrow_type.keys():
             wrapper = getattr(t_col, field)
             if wrapper:
-                return ThriftBackend._create_arrow_array(
-                    wrapper, field_name_to_arrow_type[field]
-                )
+                return ThriftBackend._create_arrow_array(wrapper, field_name_to_arrow_type[field])
 
         raise OperationalError("Empty TColumn instance {}".format(t_col))
 
@@ -654,9 +619,7 @@ def map_type(t_type_entry):
             else:
                 # Current thriftserver implementation should always return a primitiveEntry,
                 # even for complex types
-                raise OperationalError(
-                    "Thrift protocol error: t_type_entry not a primitiveEntry"
-                )
+                raise OperationalError("Thrift protocol error: t_type_entry not a primitiveEntry")
 
         def convert_col(t_column_desc):
             return pyarrow.field(
@@ -674,9 +637,7 @@ def _col_to_description(col):
             # Drop _TYPE suffix
             cleaned_type = (name[:-5] if name.endswith("_TYPE") else name).lower()
         else:
-            raise OperationalError(
-                "Thrift protocol error: t_type_entry not a primitiveEntry"
-            )
+            raise OperationalError("Thrift protocol error: t_type_entry not a primitiveEntry")
 
         if type_entry.primitiveEntry.type == ttypes.TTypeId.DECIMAL_TYPE:
             qualifiers = type_entry.primitiveEntry.typeQualifiers.qualifiers
@@ -697,9 +658,7 @@ def _col_to_description(col):
 
     @staticmethod
     def _hive_schema_to_description(t_table_schema):
-        return [
-            ThriftBackend._col_to_description(col) for col in t_table_schema.columns
-        ]
+        return [ThriftBackend._col_to_description(col) for col in t_table_schema.columns]
 
     def _results_message_to_execute_response(self, resp, operation_state):
         if resp.directResults and resp.directResults.resultSetMetadata:
@@ -726,9 +685,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             or (not direct_results.resultSet)
             or direct_results.resultSet.hasMoreRows
         )
-        description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema
-        )
+        description = self._hive_schema_to_description(t_result_set_metadata_resp.schema)
         schema_bytes = (
             t_result_set_metadata_resp.arrowSchema
             or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
@@ -768,8 +725,7 @@ def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
                 op_handle, initial_operation_status_resp
             )
         operation_state = (
-            initial_operation_status_resp
-            and initial_operation_status_resp.operationState
+            initial_operation_status_resp and initial_operation_status_resp.operationState
         )
         while not operation_state or operation_state in [
             ttypes.TOperationState.RUNNING_STATE,
@@ -784,21 +740,13 @@ def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
     def _check_direct_results_for_error(t_spark_direct_results):
         if t_spark_direct_results:
             if t_spark_direct_results.operationStatus:
-                ThriftBackend._check_response_for_error(
-                    t_spark_direct_results.operationStatus
-                )
+                ThriftBackend._check_response_for_error(t_spark_direct_results.operationStatus)
             if t_spark_direct_results.resultSetMetadata:
-                ThriftBackend._check_response_for_error(
-                    t_spark_direct_results.resultSetMetadata
-                )
+                ThriftBackend._check_response_for_error(t_spark_direct_results.resultSetMetadata)
             if t_spark_direct_results.resultSet:
-                ThriftBackend._check_response_for_error(
-                    t_spark_direct_results.resultSet
-                )
+                ThriftBackend._check_response_for_error(t_spark_direct_results.resultSet)
             if t_spark_direct_results.closeOperation:
-                ThriftBackend._check_response_for_error(
-                    t_spark_direct_results.closeOperation
-                )
+                ThriftBackend._check_response_for_error(t_spark_direct_results.closeOperation)
 
     def execute_command(
         self, operation, session_handle, max_rows, max_bytes, lz4_compression, cursor
@@ -817,9 +765,7 @@ def execute_command(
             sessionHandle=session_handle,
             statement=operation,
             runAsync=True,
-            getDirectResults=ttypes.TSparkGetDirectResults(
-                maxRows=max_rows, maxBytes=max_bytes
-            ),
+            getDirectResults=ttypes.TSparkGetDirectResults(maxRows=max_rows, maxBytes=max_bytes),
             canReadArrowResult=True,
             canDecompressLZ4Result=lz4_compression,
             canDownloadResult=False,
@@ -837,9 +783,7 @@ def get_catalogs(self, session_handle, max_rows, max_bytes, cursor):
 
         req = ttypes.TGetCatalogsReq(
             sessionHandle=session_handle,
-            getDirectResults=ttypes.TSparkGetDirectResults(
-                maxRows=max_rows, maxBytes=max_bytes
-            ),
+            getDirectResults=ttypes.TSparkGetDirectResults(maxRows=max_rows, maxBytes=max_bytes),
         )
         resp = self.make_request(self._client.GetCatalogs, req)
         return self._handle_execute_response(resp, cursor)
@@ -857,9 +801,7 @@ def get_schemas(
 
         req = ttypes.TGetSchemasReq(
             sessionHandle=session_handle,
-            getDirectResults=ttypes.TSparkGetDirectResults(
-                maxRows=max_rows, maxBytes=max_bytes
-            ),
+            getDirectResults=ttypes.TSparkGetDirectResults(maxRows=max_rows, maxBytes=max_bytes),
             catalogName=catalog_name,
             schemaName=schema_name,
         )
@@ -881,9 +823,7 @@ def get_tables(
 
         req = ttypes.TGetTablesReq(
             sessionHandle=session_handle,
-            getDirectResults=ttypes.TSparkGetDirectResults(
-                maxRows=max_rows, maxBytes=max_bytes
-            ),
+            getDirectResults=ttypes.TSparkGetDirectResults(maxRows=max_rows, maxBytes=max_bytes),
             catalogName=catalog_name,
             schemaName=schema_name,
             tableName=table_name,
@@ -907,9 +847,7 @@ def get_columns(
 
         req = ttypes.TGetColumnsReq(
             sessionHandle=session_handle,
-            getDirectResults=ttypes.TSparkGetDirectResults(
-                maxRows=max_rows, maxBytes=max_bytes
-            ),
+            getDirectResults=ttypes.TSparkGetDirectResults(maxRows=max_rows, maxBytes=max_bytes),
             catalogName=catalog_name,
             schemaName=schema_name,
             tableName=table_name,