fix(dataset): clean deleted dataset in committed batches to avoid orphaned data

[shaba]

Summary

Fixes orphaned data left behind when a knowledge base is deleted (especially while its documents are still indexing). See #38518.

clean_dataset_task deleted all documents, segments and child chunks of a deleted dataset in a single transaction using large IN (...) statements (after loading every segment/document into memory). On a large dataset this times out or gets blocked by the still-running indexing transactions; the exception then hits session.rollback(), so nothing is deleted and the documents, document_segments, child_chunks rows and the embedding_vector_index_<uuid>_nod pgvector table are orphaned forever — the datasets row itself is already gone, so the cleanup never gets retried.

Real-world impact we hit: deleting 3 knowledge bases mid-indexing left ~2.4M orphaned document_segments, ~830k child_chunks and ~34 GB of pgvector tables that had to be removed by hand.

What changed

api/tasks/clean_dataset_task.py:

• Delete each large child table (document_segments, documents, child_chunks) in bounded, individually committed batches (_DELETE_BATCH_SIZE = 1000), so partial progress is durable and a timeout/lock no longer rolls back the entire cleanup.
• Split the steps (vector cleanup / segments / child chunks / documents / attachments / dataset-scoped rows) so one failing step can't orphan everything; the vector-store cleanup stays isolated in its own try/except as before.
• Explicitly delete child_chunks by dataset_id as a fallback for when the vector-store clean(..., delete_child_chunks=True) did not remove them.
• Preserve existing behaviour: inline segment image files and uploaded document source files are still deleted from storage before their UploadFile rows; attachments and the pipeline/workflow are still cleaned.

How to test

• Create a dataset with many documents, delete it, and confirm no orphaned rows remain in document_segments / child_chunks / documents and no embedding_vector_index_<id>_nod table survives.
• The existing test_dataset_service_delete_dataset service tests mock clean_dataset_task.delay, so they are unaffected.

Refs #38518

[shaba]

Scope / relationship to other work (see #38518):

This PR fixes the cleanup-robustness half — clean_dataset_task no longer deletes everything in one transaction that rolls back (and orphans documents, segments, child chunks and the pgvector index table) on a timeout or a lock held by indexing that is still running at delete time. Each concern is deleted in bounded, per-step-committed batches, each step is isolated so one failure doesn't skip the rest, and child_chunks is cleaned explicitly as a fallback.

It is complementary to #37773 (document-vs-dataset segment deletion race) — a different code path, no overlap/conflict. The remaining part called out in #38518 — adding a guard so a dataset can't be deleted while its documents are still indexing (cancel/await in-flight indexing in DatasetService.delete_dataset) — is intentionally left as a separate follow-up.

Existing unit tests (api/tests/unit_tests/tasks/test_clean_dataset_task.py) pass unchanged (7 passed); the service-level delete_dataset tests mock clean_dataset_task.delay, so they are unaffected.

[github-actions]

Pyrefly Diff

base → PR

--- /tmp/pyrefly_base.txt	2026-07-07 16:19:36.924089772 +0000
+++ /tmp/pyrefly_pr.txt	2026-07-07 16:19:27.596999387 +0000
@@ -2111,9 +2111,9 @@
 ERROR `SimpleNamespace` is not assignable to attribute `_current_tenant` with type `Tenant | None` [bad-assignment]
    --> tests/unit_tests/controllers/console/test_workspace_members.py:186:43
 ERROR Argument `Literal['app']` is not assignable to parameter `resource_type` with type `RBACResourceScope` in function `controllers.common.wraps._extract_resource_id` [bad-argument-type]
-   --> tests/unit_tests/controllers/console/test_wraps.py:305:41
+   --> tests/unit_tests/controllers/console/test_wraps.py:298:41
 ERROR Argument `Literal['app']` is not assignable to parameter `resource_type` with type `RBACResourceScope` in function `controllers.common.wraps._extract_resource_id` [bad-argument-type]
-   --> tests/unit_tests/controllers/console/test_wraps.py:313:41
+   --> tests/unit_tests/controllers/console/test_wraps.py:306:41
 ERROR `SimpleNamespace` is not assignable to attribute `db` with type `SQLAlchemy` [bad-assignment]
   --> tests/unit_tests/controllers/files/test_image_preview.py:18:17
 ERROR `SimpleNamespace` is not assignable to attribute `request` with type `Request` [bad-assignment]
@@ -7434,6 +7434,18 @@
     --> tests/unit_tests/services/test_annotation_service.py:1194:59
 ERROR TypedDict `AnnotationSettingDisabledDict` does not have key `embedding_model` [bad-typed-dict-key]
     --> tests/unit_tests/services/test_annotation_service.py:1231:27
+ERROR Argument `SimpleNamespace` is not assignable to parameter `session` with type `Session` in function `services.app_dsl_service.AppDslService.__init__` [bad-argument-type]
+ --> tests/unit_tests/services/test_app_dsl_service.py:9:37
+ERROR Argument `SimpleNamespace` is not assignable to parameter `account` with type `Account` in function `services.app_dsl_service.AppDslService.import_app` [bad-argument-type]
+  --> tests/unit_tests/services/test_app_dsl_service.py:12:17
+ERROR Argument `SimpleNamespace` is not assignable to parameter `session` with type `Session` in function `services.app_dsl_service.AppDslService.__init__` [bad-argument-type]
+  --> tests/unit_tests/services/test_app_dsl_service.py:27:37
+ERROR Argument `SimpleNamespace` is not assignable to parameter `account` with type `Account` in function `services.app_dsl_service.AppDslService.import_app` [bad-argument-type]
+  --> tests/unit_tests/services/test_app_dsl_service.py:30:17
+ERROR Argument `SimpleNamespace` is not assignable to parameter `session` with type `Session` in function `services.app_dsl_service.AppDslService.__init__` [bad-argument-type]
+  --> tests/unit_tests/services/test_app_dsl_service.py:44:37
+ERROR Argument `SimpleNamespace` is not assignable to parameter `account` with type `Account` in function `services.app_dsl_service.AppDslService.import_app` [bad-argument-type]
+  --> tests/unit_tests/services/test_app_dsl_service.py:47:17
 ERROR Class member `_TrackingRateLimit.exit` overrides a member in a parent class but is missing an `@override` decorator [missing-override-decorator]
    --> tests/unit_tests/services/test_app_generate_service.py:560:17
 ERROR Class member `_TrackingRateLimit.exit` overrides a member in a parent class but is missing an `@override` decorator [missing-override-decorator]
@@ -7689,13 +7701,13 @@
 ERROR Argument `None` is not assignable to parameter `plan` with type `CloudPlan` in function `RagPipelineTaskProxyTestDataFactory.create_mock_features` [bad-argument-type]
    --> tests/unit_tests/services/test_rag_pipeline_task_proxy.py:442:109
 ERROR Argument `dict[str, Any]` is not assignable to parameter `node_config` with type `NodeConfigDict` in function `services.trigger.schedule_service.ScheduleService.to_schedule_config` [bad-argument-type]
-   --> tests/unit_tests/services/test_schedule_service.py:536:61
+   --> tests/unit_tests/services/test_schedule_service.py:521:61
 ERROR Argument `dict[str, str | dict[str, str]]` is not assignable to parameter `node_config` with type `NodeConfigDict` in function `services.trigger.schedule_service.ScheduleService.to_schedule_config` [bad-argument-type]
-   --> tests/unit_tests/services/test_schedule_service.py:550:56
+   --> tests/unit_tests/services/test_schedule_service.py:535:56
 ERROR Argument `dict[str, str | dict[str, str | dict[str, str]]]` is not assignable to parameter `node_config` with type `NodeConfigDict` in function `services.trigger.schedule_service.ScheduleService.to_schedule_config` [bad-argument-type]
-   --> tests/unit_tests/services/test_schedule_service.py:567:61
+   --> tests/unit_tests/services/test_schedule_service.py:552:61
 ERROR Argument `dict[str, str | dict[str, str]]` is not assignable to parameter `node_config` with type `NodeConfigDict` in function `services.trigger.schedule_service.ScheduleService.to_schedule_config` [bad-argument-type]
-   --> tests/unit_tests/services/test_schedule_service.py:579:56
+   --> tests/unit_tests/services/test_schedule_service.py:564:56
 ERROR Argument `SimpleNamespace` is not assignable to parameter `session` with type `Session` in function `services.snippet_dsl_service.SnippetDslService.__init__` [bad-argument-type]
   --> tests/unit_tests/services/test_snippet_dsl_service.py:33:41
 ERROR Argument `SimpleNamespace` is not assignable to parameter `account` with type `Account` in function `services.snippet_dsl_service.SnippetDslService.import_snippet` [bad-argument-type]

[github-actions]

Pyrefly Type Coverage

Metric	Base	PR	Delta
Type coverage	52.41%	52.40%	-0.01%
Strict coverage	51.92%	51.91%	-0.01%
Typed symbols	32,201	32,187	-14
Untyped symbols	29,518	29,514	-4
Modules	2976	2976	0