Merge branch 'sea-migration' into sea-http-client

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/backend.py

@@ -361,7 +361,7 @@ def _results_message_to_execute_response(
 
         # Check for compression
         lz4_compressed = (
-            response.manifest.result_compression == ResultCompression.LZ4_FRAME
+            response.manifest.result_compression == ResultCompression.LZ4_FRAME.value
         )
 
         execute_response = ExecuteResponse(

src/databricks/sql/backend/sea/queue.py

@@ -72,6 +72,7 @@ def build_queue(
             return JsonQueue(result_data.data)
         elif manifest.format == ResultFormat.ARROW_STREAM.value:
             if result_data.attachment is not None:
+                # direct results from Hybrid disposition
                 arrow_file = (
                     ResultSetDownloadHandler._decompress_data(result_data.attachment)
                     if lz4_compressed

src/databricks/sql/backend/sea/utils/filters.py

@@ -53,7 +53,7 @@ def _filter_sea_result_set(
         # Reuse the command_id from the original result set
         command_id = result_set.command_id
 
-        # Create an ExecuteResponse with the filtered data
+        # Create an ExecuteResponse for the filtered data
         execute_response = ExecuteResponse(
             command_id=command_id,
             status=result_set.status,

src/databricks/sql/result_set.py

@@ -55,7 +55,7 @@ def __init__(
 
         Parameters:
             :param connection: The parent connection that was used to execute this command
-            :param backend: The backend specialised backend client to be invoked in the fetch phase
+            :param backend: The specialised backend client to be invoked in the fetch phase
             :param arraysize: The max number of rows to fetch at a time (PEP-249)
             :param buffer_size_bytes: The size (in bytes) of the internal buffer + max fetch
             :param command_id: The command ID
@@ -334,7 +334,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
             n_remaining_rows -= partial_results.num_rows
             self._next_row_index += partial_results.num_rows
 
-        return pyarrow.concat_tables(partial_result_chunks)
+        return pyarrow.concat_tables(partial_result_chunks, use_threads=True)
 
     def fetchmany_columnar(self, size: int):
         """
@@ -385,7 +385,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
                 for name, col in zip(results.column_names, results.column_table)
             }
             return pyarrow.Table.from_pydict(data)
-        return pyarrow.concat_tables(partial_result_chunks)
+        return pyarrow.concat_tables(partial_result_chunks, use_threads=True)
 
     def fetchall_columnar(self):
         """Fetch all (remaining) rows of a query result, returning them as a Columnar table."""

src/databricks/sql/session.py

@@ -126,9 +126,10 @@ def open(self):
             catalog=self.catalog,
             schema=self.schema,
         )
+
         self.protocol_version = self.get_protocol_version(self._session_id)
         self.is_open = True
-        logger.info("Successfully opened session " + str(self.guid_hex))
+        logger.info("Successfully opened session %s", str(self.guid_hex))
 
     @staticmethod
     def get_protocol_version(session_id: SessionId):

src/databricks/sql/types.py

@@ -187,7 +187,6 @@ def __contains__(self, item: Any) -> bool:
     # let object acts like class
     def __call__(self, *args: Any) -> "Row":
         """create new Row object"""
-
         if len(args) > len(self):
             raise ValueError(
                 "Can not create Row with fields %s, expected %d values "
@@ -230,15 +229,13 @@ def __reduce__(
         self,
     ) -> Union[str, Tuple[Any, ...]]:
         """Returns a tuple so Python knows how to pickle Row."""
-
         if hasattr(self, "__fields__"):
             return (_create_row, (self.__fields__, tuple(self)))
         else:
             return tuple.__reduce__(self)
 
     def __repr__(self) -> str:
         """Printable representation of Row used in Python REPL."""
-
         if hasattr(self, "__fields__"):
             return "Row(%s)" % ", ".join(
                 "%s=%r" % (k, v) for k, v in zip(self.__fields__, tuple(self))

src/databricks/sql/utils.py

@@ -740,6 +740,7 @@ def convert_decimals_in_arrow_table(table, description) -> "pyarrow.Table":
 
 
 def convert_to_assigned_datatypes_in_column_table(column_table, description):
+
     converted_column_table = []
     for i, col in enumerate(column_table):
         if description[i][1] == "decimal":

tests/unit/test_client.py

@@ -105,17 +105,13 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
 
                 # Mock the execute response with controlled state
                 mock_execute_response = Mock(spec=ExecuteResponse)
-
-                mock_execute_response.command_id = Mock(spec=CommandId)
-                mock_execute_response.status = (
-                    CommandState.SUCCEEDED if not closed else CommandState.CLOSED
-                )
+                mock_execute_response.status = initial_state
                 mock_execute_response.has_been_closed_server_side = closed
                 mock_execute_response.is_staging_operation = False
                 mock_execute_response.command_id = Mock(spec=CommandId)
                 mock_execute_response.description = []
 
-                # Mock the backend that will be used by the real ThriftResultSet
+                # Mock the backend that will be used
                 mock_backend = Mock(spec=ThriftDatabricksClient)
                 mock_backend.staging_allowed_local_path = None
                 mock_backend.fetch_results.return_value = (Mock(), False, 0)
@@ -137,15 +133,13 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                 # Mock execute_command to return our real result set
                 cursor.backend.execute_command = Mock(return_value=real_result_set)
 
-                # Execute a command - this should set cursor.active_result_set to our real result set
+                # Execute a command
                 cursor.execute("SELECT 1")
 
-                # Close the connection - this should trigger the real close chain:
-                # connection.close() -> cursor.close() -> result_set.close()
+                # Close the connection
                 connection.close()
 
-                # Verify the REAL close logic worked through the chain:
-                # 1. has_been_closed_server_side should always be True after close()
+                # Verify the close logic worked:
                 assert real_result_set.has_been_closed_server_side is True
 
                 # 2. op_state should always be CLOSED after close()

tests/unit/test_cloud_fetch_queue.py

@@ -162,8 +162,6 @@ def test_next_n_rows_0_rows(self, mock_create_next_table):
         result = queue.next_n_rows(0)
         assert result.num_rows == 0
         assert queue.table_row_index == 0
-        # Instead of comparing tables directly, just check the row count
-        # This avoids issues with empty table schema differences
 
     @patch("databricks.sql.utils.ThriftCloudFetchQueue._create_next_table")
     def test_next_n_rows_partial_table(self, mock_create_next_table):