test: migrate document indexing task tests to SQLAlchemy 2.0 select API (#35145)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py

@@ -9,6 +9,7 @@ The task is responsible for removing document segments from the search index whe
 from unittest.mock import MagicMock, patch
 
 from faker import Faker
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
@@ -471,9 +472,9 @@ class TestDisableSegmentsFromIndexTask:
                 db_session_with_containers.refresh(segments[1])
 
                 # Check that segments are re-enabled after error
-                updated_segments = (
-                    db_session_with_containers.query(DocumentSegment).where(DocumentSegment.id.in_(segment_ids)).all()
-                )
+                updated_segments = db_session_with_containers.scalars(
+                    select(DocumentSegment).where(DocumentSegment.id.in_(segment_ids))
+                ).all()
 
                 for segment in updated_segments:
                     assert segment.enabled is True

api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py

@@ -2,6 +2,7 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 from faker import Faker
+from sqlalchemy import func, select
 
 from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
@@ -123,13 +124,13 @@ class TestDocumentIndexingUpdateTask:
         db_session_with_containers.expire_all()
 
         # Assert document status updated before reindex
-        updated = db_session_with_containers.query(Document).where(Document.id == document.id).first()
+        updated = db_session_with_containers.scalar(select(Document).where(Document.id == document.id).limit(1))
         assert updated.indexing_status == IndexingStatus.PARSING
         assert updated.processing_started_at is not None
 
         # Segments should be deleted
-        remaining = (
-            db_session_with_containers.query(DocumentSegment).where(DocumentSegment.document_id == document.id).count()
+        remaining = db_session_with_containers.scalar(
+            select(func.count()).select_from(DocumentSegment).where(DocumentSegment.document_id == document.id)
         )
         assert remaining == 0
 
@@ -167,8 +168,8 @@ class TestDocumentIndexingUpdateTask:
         mock_external_dependencies["runner_instance"].run.assert_called_once()
 
         # Segments should remain (since clean failed before DB delete)
-        remaining = (
-            db_session_with_containers.query(DocumentSegment).where(DocumentSegment.document_id == document.id).count()
+        remaining = db_session_with_containers.scalar(
+            select(func.count()).select_from(DocumentSegment).where(DocumentSegment.document_id == document.id)
         )
         assert remaining > 0