Rebase to upstream v3.0.2

octoenergy · Jan 31, 2024 · a92acb2 · a92acb2
1 parent faf13a4
commit a92acb2
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 165 deletions.
diff --git a/.github/workflows/code-quality-checks.yml b/.github/workflows/code-quality-checks.yml
diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
@@ -54,6 +54,10 @@
 DATABRICKS_REASON_HEADER = "x-databricks-reason-phrase"
 
 TIMESTAMP_AS_STRING_CONFIG = "spark.thriftserver.arrowBasedRowSet.timestampAsString"
+
+# HACK!
+THRIFT_SOCKET_TIMEOUT = os.getenv("THRIFT_SOCKET_TIMEOUT", None)
+
 DEFAULT_SOCKET_TIMEOUT = float(900)
 
 # see Connection.__init__ for parameter descriptions.
@@ -227,7 +231,12 @@ def __init__(
             **additional_transport_args,  # type: ignore
         )
 
-        timeout = 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))
 
@@ -637,7 +646,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: