refactor: migrate session.query to select API in clean dataset task (#34815)

api/tests/unit_tests/tasks/test_clean_dataset_task.py

@@ -60,12 +60,6 @@ def mock_db_session():
         cm.__exit__.return_value = None
         mock_sf.create_session.return_value = cm
 
-        # Setup query chain
-        mock_query = MagicMock()
-        mock_session.query.return_value = mock_query
-        mock_query.where.return_value = mock_query
-        mock_query.delete.return_value = 0
-
         # Setup scalars for select queries
         mock_session.scalars.return_value.all.return_value = []
 
@@ -220,11 +214,6 @@ class TestPipelineAndWorkflowDeletion:
         - Pipeline record is deleted
         - Related workflow record is deleted
         """
-        # Arrange
-        mock_query = mock_db_session.session.query.return_value
-        mock_query.where.return_value = mock_query
-        mock_query.delete.return_value = 1
-
         # Act
         clean_dataset_task(
             dataset_id=dataset_id,
@@ -236,9 +225,9 @@ class TestPipelineAndWorkflowDeletion:
             pipeline_id=pipeline_id,
         )
 
-        # Assert - verify delete was called for pipeline-related queries
-        # The actual count depends on total queries, but pipeline deletion should add 2 more
-        assert mock_query.delete.call_count >= 7  # 5 base + 2 pipeline/workflow
+        # Assert - verify execute was called for delete operations
+        # 1 attachment JOIN query + 5 base deletes + 2 pipeline/workflow deletes = 8
+        assert mock_db_session.session.execute.call_count >= 8
 
     def test_clean_dataset_task_without_pipeline_id(
         self,
@@ -256,11 +245,6 @@ class TestPipelineAndWorkflowDeletion:
         Expected behavior:
         - Pipeline and workflow deletion queries are not executed
         """
-        # Arrange
-        mock_query = mock_db_session.session.query.return_value
-        mock_query.where.return_value = mock_query
-        mock_query.delete.return_value = 1
-
         # Act
         clean_dataset_task(
             dataset_id=dataset_id,
@@ -272,8 +256,9 @@ class TestPipelineAndWorkflowDeletion:
             pipeline_id=None,
         )
 
-        # Assert - verify delete was called only for base queries (5 times)
-        assert mock_query.delete.call_count == 5
+        # Assert - verify execute was called for delete operations
+        # 1 attachment JOIN query + 5 base deletes = 6
+        assert mock_db_session.session.execute.call_count == 6
 
 
 # ============================================================================