remove merge artifacts, account for result set

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -956,18 +956,6 @@ def execute_command(
             thrift_handle,
         )
 
-        logger.debug(
-            "ThriftBackend.execute_command(operation=%s, session_handle=%s)",
-            operation,
-            thrift_handle,
-        )
-
-        logger.debug(
-            "ThriftBackend.execute_command(operation=%s, session_handle=%s)",
-            operation,
-            session_handle,
-        )
-
         spark_arrow_types = ttypes.TSparkArrowTypes(
             timestampAsArrow=self._use_arrow_native_timestamps,
             decimalAsArrow=self._use_arrow_native_decimals,

src/databricks/sql/client.py

@@ -776,7 +776,7 @@ def execute(
 
         self._check_not_closed()
         self._close_and_clear_active_result_set()
-        execute_response = self.backend.execute_command(
+        self.active_result_set = self.backend.execute_command(
             operation=prepared_operation,
             session_id=self.connection.session.get_session_id(),
             max_rows=self.arraysize,
@@ -788,14 +788,6 @@ def execute(
             async_op=False,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
         )
-        self.active_result_set = ResultSet(
-            self.connection,
-            execute_response,
-            self.backend,
-            self.buffer_size_bytes,
-            self.arraysize,
-            self.connection.use_cloud_fetch,
-        )
 
         if self.active_result_set.is_staging_operation:
             self._handle_staging_operation(
@@ -930,7 +922,7 @@ def catalogs(self) -> "Cursor":
         self._check_not_closed()
         self._close_and_clear_active_result_set()
         self.active_result_set = self.backend.get_catalogs(
-            session_handle=self.connection.session._session_handle,
+            session_id=self.connection.session.get_session_id(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,
@@ -949,7 +941,7 @@ def schemas(
         self._check_not_closed()
         self._close_and_clear_active_result_set()
         self.active_result_set = self.backend.get_schemas(
-            session_handle=self.connection.session._session_handle,
+            session_id=self.connection.session.get_session_id(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,
@@ -975,7 +967,7 @@ def tables(
         self._close_and_clear_active_result_set()
 
         self.active_result_set = self.backend.get_tables(
-            session_handle=self.connection.session._session_handle,
+            session_id=self.connection.session.get_session_id(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,
@@ -1003,7 +995,7 @@ def columns(
         self._close_and_clear_active_result_set()
 
         self.active_result_set = self.backend.get_columns(
-            session_handle=self.connection.session._session_handle,
+            session_id=self.connection.session.get_session_id(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,

src/databricks/sql/result_set.py

@@ -109,7 +109,7 @@ def __init__(
         super().__init__(connection, thrift_client, arraysize, buffer_size_bytes)
 
         # Initialize ThriftResultSet-specific attributes
-        self.command_id = execute_response.command_handle
+        self.command_id = execute_response.command_id
         self.op_state = execute_response.status
         self.has_been_closed_server_side = execute_response.has_been_closed_server_side
         self.has_more_rows = execute_response.has_more_rows
@@ -127,7 +127,7 @@ def __init__(
 
     def _fill_results_buffer(self):
         results, has_more_rows = self.backend.fetch_results(
-            op_handle=self.command_id,
+            command_id=self.command_id,
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             expected_row_start_offset=self._next_row_index,

tests/unit/test_client.py

@@ -133,7 +133,6 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
 
         result_set = ThriftResultSet(
             connection=mock_connection,
-            backend=mock_backend,
             execute_response=Mock(),
             thrift_client=mock_backend,
         )
@@ -582,69 +581,71 @@ def test_cursor_context_manager_handles_exit_exception(self):
         """Test that cursor's context manager handles exceptions during __exit__."""
         mock_backend = Mock()
         mock_connection = Mock()
-        
+
         cursor = client.Cursor(mock_connection, mock_backend)
         original_close = cursor.close
         cursor.close = Mock(side_effect=Exception("Test error during close"))
-        
+
         try:
             with cursor:
                 raise ValueError("Test error inside context")
         except ValueError:
             pass
-        
+
         cursor.close.assert_called_once()
 
     def test_connection_close_handles_cursor_close_exception(self):
         """Test that _close handles exceptions from cursor.close() properly."""
         cursors_closed = []
-        
+
         def mock_close_with_exception():
             cursors_closed.append(1)
             raise Exception("Test error during close")
-        
+
         cursor1 = Mock()
         cursor1.close = mock_close_with_exception
-        
+
         def mock_close_normal():
             cursors_closed.append(2)
-        
+
         cursor2 = Mock()
         cursor2.close = mock_close_normal
-        
+
         mock_backend = Mock()
         mock_session_handle = Mock()
-        
+
         try:
             for cursor in [cursor1, cursor2]:
                 try:
                     cursor.close()
                 except Exception:
                     pass
-                    
+
             mock_backend.close_session(mock_session_handle)
         except Exception as e:
             self.fail(f"Connection close should handle exceptions: {e}")
-        
-        self.assertEqual(cursors_closed, [1, 2], "Both cursors should have close called")
+
+        self.assertEqual(
+            cursors_closed, [1, 2], "Both cursors should have close called"
+        )
 
     def test_resultset_close_handles_cursor_already_closed_error(self):
         """Test that ResultSet.close() handles CursorAlreadyClosedError properly."""
         result_set = client.ThriftResultSet.__new__(client.ThriftResultSet)
         result_set.backend = Mock()
-        result_set.backend.CLOSED_OP_STATE = 'CLOSED'
+        result_set.backend.CLOSED_OP_STATE = "CLOSED"
         result_set.connection = Mock()
         result_set.connection.open = True
-        result_set.op_state = 'RUNNING'
+        result_set.op_state = "RUNNING"
         result_set.has_been_closed_server_side = False
         result_set.command_id = Mock()
 
         class MockRequestError(Exception):
             def __init__(self):
                 self.args = ["Error message", CursorAlreadyClosedError()]
-        
+
         result_set.backend.close_command.side_effect = MockRequestError()
-        
+
         original_close = client.ResultSet.close
         try:
             try:
@@ -660,11 +661,13 @@ def __init__(self):
             finally:
                 result_set.has_been_closed_server_side = True
                 result_set.op_state = result_set.backend.CLOSED_OP_STATE
-                
-            result_set.backend.close_command.assert_called_once_with(result_set.command_id)
-            
+
+            result_set.backend.close_command.assert_called_once_with(
+                result_set.command_id
+            )
+
             assert result_set.has_been_closed_server_side is True
-            
+
             assert result_set.op_state == result_set.backend.CLOSED_OP_STATE
         finally:
             pass