formatting (black)

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -7,6 +7,7 @@
 
 # Forward reference for type hints
 from typing import TYPE_CHECKING
+
 if TYPE_CHECKING:
     from databricks.sql.result_set import ResultSet
 

src/databricks/sql/backend/thrift_backend.py

@@ -845,14 +845,14 @@ def get_execution_result(self, op_handle, cursor):
             description=description,
             arrow_schema_bytes=schema_bytes,
         )
-        
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
-            use_cloud_fetch=cursor.connection.use_cloud_fetch
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -951,14 +951,14 @@ def execute_command(
             return None
         else:
             execute_response = self._handle_execute_response(resp, cursor)
-            
+
             return ThriftResultSet(
                 connection=cursor.connection,
                 execute_response=execute_response,
                 thrift_client=self,
                 buffer_size_bytes=max_bytes,
                 arraysize=max_rows,
-                use_cloud_fetch=use_cloud_fetch
+                use_cloud_fetch=use_cloud_fetch,
             )
 
     def get_catalogs(self, session_handle, max_rows, max_bytes, cursor):
@@ -971,16 +971,16 @@ def get_catalogs(self, session_handle, max_rows, max_bytes, cursor):
             ),
         )
         resp = self.make_request(self._client.GetCatalogs, req)
-        
+
         execute_response = self._handle_execute_response(resp, cursor)
-        
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
-            use_cloud_fetch=cursor.connection.use_cloud_fetch
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
         )
 
     def get_schemas(
@@ -1003,16 +1003,16 @@ def get_schemas(
             schemaName=schema_name,
         )
         resp = self.make_request(self._client.GetSchemas, req)
-        
+
         execute_response = self._handle_execute_response(resp, cursor)
-        
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
-            use_cloud_fetch=cursor.connection.use_cloud_fetch
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
         )
 
     def get_tables(
@@ -1039,16 +1039,16 @@ def get_tables(
             tableTypes=table_types,
         )
         resp = self.make_request(self._client.GetTables, req)
-        
+
         execute_response = self._handle_execute_response(resp, cursor)
-        
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
-            use_cloud_fetch=cursor.connection.use_cloud_fetch
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
         )
 
     def get_columns(
@@ -1075,16 +1075,16 @@ def get_columns(
             columnName=column_name,
         )
         resp = self.make_request(self._client.GetColumns, req)
-        
+
         execute_response = self._handle_execute_response(resp, cursor)
-        
+
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
-            use_cloud_fetch=cursor.connection.use_cloud_fetch
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
         )
 
     def _handle_execute_response(self, resp, cursor):

src/databricks/sql/client.py

@@ -754,7 +754,7 @@ def execute(
 
         self._check_not_closed()
         self._close_and_clear_active_result_set()
-        
+
         self.active_result_set = self.backend.execute_command(
             operation=prepared_operation,
             session_handle=self.connection.session._session_handle,
@@ -1124,5 +1124,3 @@ def setinputsizes(self, sizes):
     def setoutputsize(self, size, column=None):
         """Does nothing by default"""
         pass
-
-

src/databricks/sql/result_set.py

@@ -21,10 +21,10 @@
 class ResultSet(ABC):
     """
     Abstract base class for result sets returned by different backend implementations.
-    
+
     This class defines the interface that all concrete result set implementations must follow.
     """
-    
+
     def __init__(self, connection, backend, arraysize: int, buffer_size_bytes: int):
         """Initialize the base ResultSet with common properties."""
         self.connection = connection
@@ -33,50 +33,50 @@ def __init__(self, connection, backend, arraysize: int, buffer_size_bytes: int):
         self.buffer_size_bytes = buffer_size_bytes
         self._next_row_index = 0
         self.description = None
-    
+
     def __iter__(self):
         while True:
             row = self.fetchone()
             if row:
                 yield row
             else:
                 break
-    
+
     @property
     def rownumber(self):
         return self._next_row_index
-    
+
     # Define abstract methods that concrete implementations must implement
     @abstractmethod
     def _fill_results_buffer(self):
         """Fill the results buffer from the backend."""
         pass
-    
+
     @abstractmethod
     def fetchone(self) -> Optional[Row]:
         """Fetch the next row of a query result set."""
         pass
-    
+
     @abstractmethod
     def fetchmany(self, size: int) -> List[Row]:
         """Fetch the next set of rows of a query result."""
         pass
-    
+
     @abstractmethod
     def fetchall(self) -> List[Row]:
         """Fetch all remaining rows of a query result."""
         pass
-    
+
     @abstractmethod
     def fetchmany_arrow(self, size: int) -> Any:
         """Fetch the next set of rows as an Arrow table."""
         pass
-    
+
     @abstractmethod
     def fetchall_arrow(self) -> Any:
         """Fetch all remaining rows as an Arrow table."""
         pass
-    
+
     @abstractmethod
     def close(self) -> None:
         """Close the result set and release any resources."""
@@ -85,7 +85,7 @@ def close(self) -> None:
 
 class ThriftResultSet(ResultSet):
     """ResultSet implementation for the Thrift backend."""
-    
+
     def __init__(
         self,
         connection,
@@ -97,7 +97,7 @@ def __init__(
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
-        
+
         Args:
             connection: The parent connection
             execute_response: Response from the execute command
@@ -107,7 +107,7 @@ def __init__(
             use_cloud_fetch: Whether to use cloud fetch for retrieving results
         """
         super().__init__(connection, thrift_client, arraysize, buffer_size_bytes)
-        
+
         # Initialize ThriftResultSet-specific attributes
         self.command_id = execute_response.command_handle
         self.op_state = execute_response.status
@@ -124,7 +124,7 @@ def __init__(
             self.results = execute_response.arrow_queue
         else:
             self._fill_results_buffer()
-    
+
     def _fill_results_buffer(self):
         """Fill the results buffer using the ThriftDatabricksClient directly."""
         # Use the thrift_client (backend) directly to fetch results
@@ -140,7 +140,7 @@ def _fill_results_buffer(self):
         )
         self.results = results
         self.has_more_rows = has_more_rows
-    
+
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]
         ResultRow = Row(*column_names)
@@ -352,4 +352,4 @@ def close(self) -> None:
                 logger.info("Operation was canceled by a prior request")
         finally:
             self.has_been_closed_server_side = True
-            self.op_state = ttypes.TOperationState.CLOSED_STATE
+            self.op_state = ttypes.TOperationState.CLOSED_STATE