Add timeout hack to mitigate timeouts

Author: capitancambio

src/databricks/sql/auth/thrift_http_client.py

@@ -5,7 +5,6 @@
 
 import six
 import thrift
-import thrift.transport.THttpClient
 
 logger = logging.getLogger(__name__)
 
@@ -212,4 +211,4 @@ def set_retry_command_type(self, value: CommandType):
         else:
             logger.warning(
                 "DatabricksRetryPolicy is currently bypassed. The CommandType cannot be set."
-            )
+            )

src/databricks/sql/thrift_backend.py

@@ -8,8 +8,6 @@
 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.protocol.TBinaryProtocol
 import thrift.transport.THttpClient
@@ -23,7 +21,6 @@
 from databricks.sql.auth.authenticators import AuthProvider
 from databricks.sql.thrift_api.TCLIService import TCLIService, ttypes
 from databricks.sql import *
-from databricks.sql.exc import MaxRetryDurationError
 from databricks.sql.thrift_api.TCLIService.TCLIService import (
     Client as TCLIServiceClient,
 )
@@ -157,9 +154,13 @@ 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
+        )
 
         # Cloud fetch
         self.max_download_threads = kwargs.get("max_download_threads", 10)
@@ -229,7 +230,12 @@ def __init__(
             **additional_transport_args,  # type: ignore
         )
 
-        timeout = THRIFT_SOCKET_TIMEOUT or kwargs.get("_socket_timeout", DEFAULT_SOCKET_TIMEOUT)
+        timeout = THRIFT_SOCKET_TIMEOUT or kwargs.get(
+            "_socket_timeout", DEFAULT_SOCKET_TIMEOUT
+        )
+        # HACK!
+        logger.info(f"Setting timeout HACK! to {timeout}")
+
         # setTimeout defaults to 15 minutes and is expected in ms
         self._transport.setTimeout(timeout and (float(timeout) * 1000.0))
 
@@ -253,11 +259,15 @@ 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?
@@ -285,7 +295,9 @@ 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]
                 )
@@ -516,7 +528,10 @@ 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."
@@ -531,7 +546,10 @@ 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(
@@ -543,14 +561,18 @@ 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
 
@@ -576,7 +598,9 @@ 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(
@@ -621,7 +645,10 @@ def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, descripti
                 num_rows,
             ) = convert_column_based_set_to_arrow_table(t_row_set.columns, description)
         elif t_row_set.arrowBatches is not None:
-            (arrow_table, num_rows,) = convert_arrow_based_set_to_arrow_table(
+            (
+                arrow_table,
+                num_rows,
+            ) = convert_arrow_based_set_to_arrow_table(
                 t_row_set.arrowBatches, lz4_compressed, schema_bytes
             )
         else:
@@ -663,7 +690,9 @@ 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(
@@ -681,7 +710,9 @@ 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
@@ -702,7 +733,9 @@ 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:
@@ -730,7 +763,9 @@ 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)
@@ -771,7 +806,8 @@ 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,
@@ -786,13 +822,21 @@ 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,
@@ -819,7 +863,9 @@ 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=use_cloud_fetch,
@@ -838,7 +884,9 @@ 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)
@@ -856,7 +904,9 @@ 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,
         )
@@ -878,7 +928,9 @@ 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,
@@ -902,7 +954,9 @@ 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,