remove excess changes

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

tests/unit/test_client.py

@@ -48,7 +48,7 @@ def new(cls):
             is_staging_operation=False,
             command_id=None,
             has_been_closed_server_side=True,
-            has_more_rows=True,
+            is_direct_results=True,
             lz4_compressed=True,
             arrow_schema_bytes=b"schema",
         )
@@ -104,6 +104,7 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                 # Mock the backend that will be used by the real ThriftResultSet
                 mock_backend = Mock(spec=ThriftDatabricksClient)
                 mock_backend.staging_allowed_local_path = None
+                mock_backend.fetch_results.return_value = (Mock(), False)
 
                 # Configure the decorator's mock to return our specific mock_backend
                 mock_thrift_client_class.return_value = mock_backend
@@ -184,6 +185,7 @@ def test_arraysize_buffer_size_passthrough(
     def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
         mock_connection = Mock()
         mock_backend = Mock()
+        mock_backend.fetch_results.return_value = (Mock(), False)
 
         result_set = ThriftResultSet(
             connection=mock_connection,
@@ -210,6 +212,7 @@ def test_closing_result_set_hard_closes_commands(self):
         mock_session.open = True
         type(mock_connection).session = PropertyMock(return_value=mock_session)
 
+        mock_thrift_backend.fetch_results.return_value = (Mock(), False)
         result_set = ThriftResultSet(
             mock_connection, mock_results_response, mock_thrift_backend
         )
@@ -254,7 +257,10 @@ def test_closed_cursor_doesnt_allow_operations(self):
             self.assertIn("closed", e.msg)
 
     def test_negative_fetch_throws_exception(self):
-        result_set = ThriftResultSet(Mock(), Mock(), Mock())
+        mock_backend = Mock()
+        mock_backend.fetch_results.return_value = (Mock(), False)
+
+        result_set = ThriftResultSet(Mock(), Mock(), mock_backend)
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)
@@ -472,7 +478,6 @@ def make_fake_row_slice(n_rows):
         mock_aq = Mock()
         mock_aq.next_n_rows.side_effect = make_fake_row_slice
         mock_thrift_backend.execute_command.return_value.arrow_queue = mock_aq
-        mock_thrift_backend.fetch_results.return_value = (mock_aq, True)
 
         cursor = client.Cursor(Mock(), mock_thrift_backend)
         cursor.execute("foo")

tests/unit/test_fetches.py

@@ -40,25 +40,30 @@ def make_dummy_result_set_from_initial_results(initial_results):
         # If the initial results have been set, then we should never try and fetch more
         schema, arrow_table = FetchTests.make_arrow_table(initial_results)
         arrow_queue = ArrowQueue(arrow_table, len(initial_results), 0)
+
+        # Create a mock backend that will return the queue when _fill_results_buffer is called
+        mock_thrift_backend = Mock(spec=ThriftDatabricksClient)
+        mock_thrift_backend.fetch_results.return_value = (arrow_queue, False)
+
+        num_cols = len(initial_results[0]) if initial_results else 0
+        description = [
+            (f"col{col_id}", "integer", None, None, None, None, None)
+            for col_id in range(num_cols)
+        ]
+
         rs = ThriftResultSet(
             connection=Mock(),
             execute_response=ExecuteResponse(
                 command_id=None,
                 status=None,
                 has_been_closed_server_side=True,
-                has_more_rows=False,
-                description=Mock(),
-                lz4_compressed=Mock(),
-                results_queue=arrow_queue,
+                description=description,
+                lz4_compressed=True,
                 is_staging_operation=False,
             ),
-            thrift_client=None,
+            thrift_client=mock_thrift_backend,
+            t_row_set=None,
         )
-        num_cols = len(initial_results[0]) if initial_results else 0
-        rs.description = [
-            (f"col{col_id}", "integer", None, None, None, None, None)
-            for col_id in range(num_cols)
-        ]
         return rs
 
     @staticmethod
@@ -85,19 +90,19 @@ def fetch_results(
         mock_thrift_backend.fetch_results = fetch_results
         num_cols = len(batch_list[0][0]) if batch_list and batch_list[0] else 0
 
+        description = [
+            (f"col{col_id}", "integer", None, None, None, None, None)
+            for col_id in range(num_cols)
+        ]
+
         rs = ThriftResultSet(
             connection=Mock(),
             execute_response=ExecuteResponse(
                 command_id=None,
                 status=None,
                 has_been_closed_server_side=False,
-                has_more_rows=True,
-                description=[
-                    (f"col{col_id}", "integer", None, None, None, None, None)
-                    for col_id in range(num_cols)
-                ],
-                lz4_compressed=Mock(),
-                results_queue=None,
+                description=description,
+                lz4_compressed=True,
                 is_staging_operation=False,
             ),
             thrift_client=mock_thrift_backend,

tests/unit/test_fetches_bench.py

@@ -36,7 +36,7 @@ def make_dummy_result_set_from_initial_results(arrow_table):
             execute_response=ExecuteResponse(
                 status=None,
                 has_been_closed_server_side=True,
-                has_more_rows=False,
+                is_direct_results=False,
                 description=Mock(),
                 command_id=None,
                 arrow_queue=arrow_queue,