formatting

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/backend/thrift_backend.py

@@ -1022,9 +1022,11 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, is_direct_results, statement_id = self._handle_execute_response(
-                resp, cursor
-            )
+            (
+                execute_response,
+                is_direct_results,
+                statement_id,
+            ) = self._handle_execute_response(resp, cursor)
 
             t_row_set = None
             if resp.directResults and resp.directResults.resultSet:
@@ -1248,7 +1250,11 @@ def _handle_execute_response(self, resp, cursor):
             resp, final_operation_state
         )
 
-        return execute_response, is_direct_results, cursor.active_command_id.to_hex_guid()
+        return (
+            execute_response,
+            is_direct_results,
+            cursor.active_command_id.to_hex_guid(),
+        )
 
     def _handle_execute_response_async(self, resp, cursor):
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
@@ -1306,7 +1312,7 @@ def fetch_results(
             description=description,
             ssl_options=self._ssl_options,
             session_id_hex=self._session_id_hex,
-            statement_id=statement_id
+            statement_id=statement_id,
         )
 
         return queue, resp.hasMoreRows

src/databricks/sql/cloudfetch/download_manager.py

@@ -95,15 +95,17 @@ def _schedule_downloads(self):
         ):
             chunk_id, link = self._pending_links.pop(0)
             logger.debug(
-                "- chunk: {}, start: {}, row count: {}".format(chunk_id, link.startRowOffset, link.rowCount)
+                "- chunk: {}, start: {}, row count: {}".format(
+                    chunk_id, link.startRowOffset, link.rowCount
+                )
             )
             handler = ResultSetDownloadHandler(
                 settings=self._downloadable_result_settings,
                 link=link,
                 ssl_options=self._ssl_options,
                 chunk_id=chunk_id,
                 session_id_hex=self.session_id_hex,
-                statement_id=self.statement_id
+                statement_id=self.statement_id,
             )
             task = self._thread_pool.submit(handler.run)
             self._download_tasks.append(task)

src/databricks/sql/telemetry/latency_logger.py

@@ -65,8 +65,9 @@ def get_is_compressed(self) -> bool:
     def get_execution_result(self) -> ExecutionResultFormat:
         if self.active_result_set is None:
             return ExecutionResultFormat.FORMAT_UNSPECIFIED
-        
+
         from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
+
         if isinstance(self.active_result_set.results, ColumnQueue):
             return ExecutionResultFormat.COLUMNAR_INLINE
         elif isinstance(self.active_result_set.results, CloudFetchQueue):
@@ -76,13 +77,10 @@ def get_execution_result(self) -> ExecutionResultFormat:
         return ExecutionResultFormat.FORMAT_UNSPECIFIED
 
     def get_retry_count(self) -> int:
-        if (
-            hasattr(self.backend, "retry_policy")
-            and self.backend.retry_policy
-        ):
+        if hasattr(self.backend, "retry_policy") and self.backend.retry_policy:
             return len(self.backend.retry_policy.history)
         return 0
-    
+
     def get_chunk_id(self):
         return None
 
@@ -108,6 +106,7 @@ def get_is_compressed(self) -> bool:
 
     def get_execution_result(self) -> ExecutionResultFormat:
         from databricks.sql.utils import ColumnQueue, CloudFetchQueue, ArrowQueue
+
         if isinstance(self.results, ColumnQueue):
             return ExecutionResultFormat.COLUMNAR_INLINE
         elif isinstance(self.results, CloudFetchQueue):
@@ -123,7 +122,7 @@ def get_retry_count(self) -> int:
         ):
             return len(self.thrift_backend.retry_policy.history)
         return 0
-    
+
     def get_chunk_id(self):
         return None
 
@@ -132,6 +131,7 @@ class ResultSetDownloadHandlerExtractor(TelemetryExtractor):
     """
     Telemetry extractor specialized for ResultSetDownloadHandler objects.
     """
+
     def get_session_id_hex(self) -> Optional[str]:
         return self._obj.session_id_hex
 
@@ -150,7 +150,7 @@ def get_retry_count(self) -> Optional[int]:
 
     def get_chunk_id(self) -> Optional[int]:
         return self._obj.chunk_id
-    
+
 
 def get_extractor(obj):
     """
@@ -174,7 +174,7 @@ def get_extractor(obj):
         return CursorExtractor(obj)
     elif obj.__class__.__name__ == "ResultSet":
         return ResultSetExtractor(obj)
-    elif obj.__class__.__name__=="ResultSetDownloadHandler":
+    elif obj.__class__.__name__ == "ResultSetDownloadHandler":
         return ResultSetDownloadHandlerExtractor(obj)
     else:
         logger.debug("No extractor found for %s", obj.__class__.__name__)
@@ -233,7 +233,17 @@ def _safe_call(func_to_call):
                 duration_ms = int((end_time - start_time) * 1000)
 
                 extractor = get_extractor(self)
-                print("function name", func.__name__, "latency", duration_ms, "session_id_hex", extractor.get_session_id_hex(), "statement_id", extractor.get_statement_id(), flush=True)
+                print(
+                    "function name",
+                    func.__name__,
+                    "latency",
+                    duration_ms,
+                    "session_id_hex",
+                    extractor.get_session_id_hex(),
+                    "statement_id",
+                    extractor.get_statement_id(),
+                    flush=True,
+                )
 
                 if extractor is not None:
                     session_id_hex = _safe_call(extractor.get_session_id_hex)

src/databricks/sql/telemetry/models/event.py

@@ -131,6 +131,7 @@ class SqlExecutionEvent(JsonSerializableMixin):
     retry_count: Optional[int]
     chunk_id: Optional[int]
 
+
 @dataclass
 class TelemetryEvent(JsonSerializableMixin):
     """

src/databricks/sql/utils.py

@@ -239,7 +239,7 @@ def __init__(
         self._ssl_options = ssl_options
         self.session_id_hex = session_id_hex
         self.statement_id = statement_id
-        
+
         logger.debug(
             "Initialize CloudFetch loader, row set start offset: {}, file list:".format(
                 start_row_offset