diff --git a/backend/api_v2/api_deployment_views.py b/backend/api_v2/api_deployment_views.py index d1b9ac543f..45b09e57dc 100644 --- a/backend/api_v2/api_deployment_views.py +++ b/backend/api_v2/api_deployment_views.py @@ -29,6 +29,7 @@ from utils.hubspot_notify import notify_hubspot_event from utils.pagination import CustomPagination from workflow_manager.workflow_v2.dto import ExecutionResponse +from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.models.execution import WorkflowExecution from api_v2.api_deployment_dto_registry import ApiDeploymentDTORegistry @@ -171,7 +172,10 @@ def post( # Determine response status based on execution result execution_status = response.get("execution_status", "") - has_error = response.get("error") or execution_status == "ERROR" + has_error = response.get("error") or execution_status in ( + ExecutionStatus.ERROR.value, + ExecutionStatus.STOPPED.value, + ) if has_error: # Check for tool not found in registry error - return 500 Internal Server Error @@ -253,6 +257,10 @@ def get( response_status = status.HTTP_422_UNPROCESSABLE_ENTITY if execution_status_value == CeleryTaskState.COMPLETED.value: response_status = status.HTTP_200_OK + if execution_status_value in ( + CeleryTaskState.COMPLETED.value, + ExecutionStatus.STOPPED.value, + ): DeploymentHelper.process_completed_execution( response=response, deployment_execution_dto=deployment_execution_dto, diff --git a/backend/api_v2/deployment_helper.py b/backend/api_v2/deployment_helper.py index a04fd79ff8..76585cc91a 100644 --- a/backend/api_v2/deployment_helper.py +++ b/backend/api_v2/deployment_helper.py @@ -296,10 +296,13 @@ def execute_workflow( ) except Exception as error: # Isolate the DB error-marking so cleanup still runs if it raises. + stored_status = ExecutionStatus.ERROR.value try: - WorkflowExecutionServiceHelper.update_execution_err( + marked = WorkflowExecutionServiceHelper.update_execution_err( str(execution_id), str(error) ) + # An aborted run stays STOPPED; answer with what is stored. + stored_status = marked.status if marked else stored_status except Exception: logger.exception(f"Failed to mark execution {execution_id} as ERROR") @@ -317,7 +320,7 @@ def execute_workflow( ExecutionResponse( workflow_id=workflow_id, execution_id=execution_id, - execution_status=ExecutionStatus.ERROR.value, + execution_status=stored_status, error=str(error), ) ).data @@ -360,6 +363,11 @@ def execute_workflow( f"Failed to acknowledge results for execution {execution_id}" ) + # An aborted run was already reported by the abort. + if ( + execution is not None + and execution.status == ExecutionStatus.COMPLETED.value + ): try: # Terminalising here bypasses the usual notifier, so subscribers # would otherwise hear nothing. diff --git a/backend/api_v2/openapi_schema.py b/backend/api_v2/openapi_schema.py index a82e92739d..5178c490b3 100644 --- a/backend/api_v2/openapi_schema.py +++ b/backend/api_v2/openapi_schema.py @@ -246,7 +246,9 @@ def _get_error_response_examples(self) -> list[OpenApiExample]: "later call for that execution answers 406. Keep the payload of the call " "that returns it — it cannot be fetched again.\n\n" "A still-running execution answers 422 carrying its current `status`, so a " - "polling loop should treat 422 as the normal reply and stop on 200." + "polling loop should treat 422 as the normal reply and stop on 200, or when " + "`status` is `ERROR` or `STOPPED`. Those also answer 422 but the run has ended, " + "and that read is the only one that carries its results." ) @@ -263,7 +265,9 @@ def _get_error_response_examples(self) -> list[OpenApiExample]: responses={ 200: ExecuteResponse, 422: OpenApiResponse( - ExecuteResponse, description="The execution finished with an error." + ExecuteResponse, + description="The execution finished with an error, or was aborted " + "(`STOPPED`).", ), 500: OpenApiResponse( ExecuteResponse, @@ -289,8 +293,8 @@ def _get_error_response_examples(self) -> list[OpenApiExample]: ), 422: OpenApiResponse( StatusResponse, - description="The execution is still running, or it finished with an " - "error; read `status` to tell them apart.", + description="The execution is still running, finished with an error, " + "or was aborted (`STOPPED`); read `status` to tell them apart.", ), 500: OpenApiResponse( StatusResponse, diff --git a/backend/api_v2/tests/test_deployment_aborted.py b/backend/api_v2/tests/test_deployment_aborted.py new file mode 100644 index 0000000000..9a7caee0c2 --- /dev/null +++ b/backend/api_v2/tests/test_deployment_aborted.py @@ -0,0 +1,111 @@ +"""API deployment calls report an aborted run as an error (UN-4318).""" + +from __future__ import annotations + +from unittest.mock import patch + +from account_v2.models import Organization +from django.core.files.uploadedfile import SimpleUploadedFile +from django.test import TestCase +from rest_framework.response import Response +from rest_framework.test import APIRequestFactory +from utils.user_context import UserContext +from workflow_manager.workflow_v2.dto import ExecutionResponse +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow + +from api_v2.api_deployment_views import DeploymentExecution +from api_v2.models import APIDeployment, APIKey + +ORG_ID = "org-aborted" + + +class SyncCallForAbortedRunTest(TestCase): + def setUp(self) -> None: + Organization.objects.create( + name=ORG_ID, display_name="Org", organization_id=ORG_ID + ) + UserContext.set_organization_identifier(ORG_ID) + self.workflow = Workflow.objects.create(workflow_name="wf", is_active=True) + self.api = APIDeployment.objects.create(api_name="api", workflow=self.workflow) + self.key = APIKey.objects.create(api=self.api) + + def _post_with_status(self, execution_status: str) -> Response: + request = APIRequestFactory().post( + f"/deployment/api/{ORG_ID}/api/", + {"files": SimpleUploadedFile("doc.txt", b"hello")}, + format="multipart", + HTTP_AUTHORIZATION=f"Bearer {self.key.api_key}", + ) + with ( + patch( + "api_v2.api_deployment_views.APIDeploymentRateLimiter.check_and_acquire", + return_value=(True, {}), + ), + patch( + "api_v2.api_deployment_views.DeploymentHelper.execute_workflow", + return_value={"execution_status": execution_status}, + ), + ): + return DeploymentExecution.as_view()( + request, org_name=ORG_ID, api_name="api" + ) + + def test_status_codes(self) -> None: + for execution_status, code in ( + ("STOPPED", 422), + ("ERROR", 422), + ("COMPLETED", 200), + ): + with self.subTest(execution_status=execution_status): + self.assertEqual( + self._post_with_status(execution_status).status_code, code + ) + + def _get_stopped_result(self) -> Response: + item = { + "file": "doc.txt", + "metadata": {"workflow_id": "w", "highlight_data": [1]}, + "result": { + "output": {"a": 1}, + "metadata": {"inner": 1}, + "metrics": {"llm": 1}, + }, + } + execution = WorkflowExecution.objects.create( + workflow=self.workflow, + pipeline_id=self.api.id, + execution_mode=WorkflowExecution.Mode.QUEUE, + execution_method=WorkflowExecution.Method.DIRECT, + execution_type=WorkflowExecution.Type.COMPLETE, + status="STOPPED", + ) + stopped = ExecutionResponse( + workflow_id="w", + execution_id="e", + execution_status="STOPPED", + result=[item], + ) + request = APIRequestFactory().get( + f"/deployment/api/{ORG_ID}/api/", + {"execution_id": str(execution.id)}, + HTTP_AUTHORIZATION=f"Bearer {self.key.api_key}", + ) + helper = "api_v2.api_deployment_views.DeploymentHelper" + with ( + patch(f"{helper}.get_execution_status", return_value=stopped), + patch(f"{helper}._enrich_result_with_workflow_metadata"), + ): + return DeploymentExecution.as_view()( + request, org_name=ORG_ID, api_name="api" + ) + + def test_get_filters_a_stopped_runs_results_like_post(self) -> None: + response = self._get_stopped_result() + + self.assertEqual(response.status_code, 422) + item = response.data["message"][0] + self.assertNotIn("highlight_data", item["metadata"]) + self.assertNotIn("metrics", item["result"]) + self.assertNotIn("metadata", item["result"]) + self.assertEqual(item["result"]["output"], {"a": 1}) diff --git a/backend/api_v2/tests/test_deployment_helper.py b/backend/api_v2/tests/test_deployment_helper.py index 2e8c3ec7d3..f3f37f8634 100644 --- a/backend/api_v2/tests/test_deployment_helper.py +++ b/backend/api_v2/tests/test_deployment_helper.py @@ -92,6 +92,24 @@ def test_staging_failure_cleanup_survives_db_marking_error(collaborators) -> Non collaborators["DestinationConnector"].delete_api_storage_dir.assert_called_once() +def test_staging_failure_after_an_abort_answers_stopped( + collaborators: dict[str, MagicMock], +) -> None: + """The answer matches the stored status, which the abort keeps STOPPED.""" + collaborators["WorkflowExecutionServiceHelper"].update_execution_err.return_value = ( + MagicMock(status="STOPPED") + ) + + response = dh.DeploymentHelper.execute_workflow( + organization_name="org", + api=_api(), + file_objs=[], + timeout=-1, + ) + + assert response["execution_status"] == "STOPPED" + + @pytest.fixture def staging_rejects_everything(): """Patch execute_workflow's collaborators; staging returns no dispatchable files.""" @@ -195,6 +213,27 @@ def test_all_files_rejected_acknowledges_and_notifies( ) +def test_an_aborted_run_is_not_reported_again( + staging_rejects_everything: dict[str, MagicMock], +) -> None: + """The abort already notified subscribers of a run it stopped.""" + mocks = staging_rejects_everything + stopped_row = mocks[ + "WorkflowExecutionServiceHelper" + ].update_execution_completed.return_value + stopped_row.status = "STOPPED" + + response = dh.DeploymentHelper.execute_workflow( + organization_name="org", + api=_api(), + file_objs=[MagicMock()], + timeout=-1, + ) + + mocks["PipelineUtils"].update_pipeline_status.assert_not_called() + assert response["execution_status"] == "STOPPED" + + def test_notification_survives_a_failed_acknowledgement( staging_rejects_everything, ) -> None: diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 48fd799996..d31f25c89e 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -407,6 +407,10 @@ def validate_agent_kv_path_prefix(raw, path_prefix, deployment_prefix): INSTANT_WF_POLLING_TIMEOUT = int( os.environ.get("INSTANT_WF_POLLING_TIMEOUT", "300") ) # 5 minutes +# A PENDING/EXECUTING run with no activity for this long is listed as stuck. +STUCK_EXECUTION_THRESHOLD_SECONDS = int( + os.environ.get("STUCK_EXECUTION_THRESHOLD_SECONDS", "9000") +) # ETL Pipeline minimum schedule interval (in seconds) # Default: 1800 seconds (30 minutes) diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index e3e16bf711..81338e1ae0 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -55,9 +55,9 @@ class MCPTool: consumes_result: True when *reading* is itself destructive — the result store is one-shot, so a successful read acknowledges the execution and the payload cannot be fetched again. ``getExecutionStatus`` on - a COMPLETED execution is the case: it neither writes configuration - nor spends money, so ``writes``/``billable`` are both false, but - calling it twice loses data. Without this flag the derivation below + a finished execution is the case: it neither writes + configuration nor spends money, so ``writes``/``billable`` are both + false, but calling it twice loses data. Without this flag the derivation below would advertise ``readOnlyHint: true`` and tell an agent it is safe to poll freely — the opposite of the truth. preflight: Optional ``preflight(context, **arguments)`` run *before* the @@ -266,14 +266,14 @@ def build_deployment_registry() -> MCPToolRegistry: "Fetch the status and, once available, the result of an " "extraction previously started by `extractDocument`. Pass the " "`execution_id` that call returned.\n\n" - "Returns an `execution_status` of PENDING, EXECUTING, COMPLETED " - "or ERROR. Poll while it is PENDING or EXECUTING, leaving a few " + "Returns an `execution_status` of PENDING, EXECUTING, COMPLETED, " + "ERROR or STOPPED. Poll while it is PENDING or EXECUTING, leaving a few " "seconds between calls." ), input_schema=get_execution_status_schema(), handler=get_execution_status, - # Reading a COMPLETED execution acknowledges it and the result - # store is one-shot, so the payload cannot be fetched again. + # Reading a finished execution acknowledges it and the + # result store is one-shot, so the payload cannot be fetched again. # Not `writes` (no configuration changes) and not `billable` # (no spend), but emphatically not read-only either. consumes_result=True, @@ -714,14 +714,14 @@ def build_platform_registry() -> MCPToolRegistry: "Poll for the result of an extraction started by " "extractDocument. Free to call.\n\n" "If extractDocument already returned execution_status " - "COMPLETED, the result is in that response and no polling is " - "needed. Otherwise poll this until COMPLETED or ERROR, " + "COMPLETED, ERROR or STOPPED, that response is final and no " + "polling is needed. Otherwise poll this until one of those, " "pausing a few seconds between calls." ), input_schema=platform_execution_status_schema(), handler=get_platform_execution_status, - # Reading a COMPLETED execution acknowledges it and the result - # store is one-shot, so the payload cannot be fetched again. + # Reading a finished execution acknowledges it and the + # result store is one-shot, so the payload cannot be fetched again. # Not `writes` (no configuration changes) and not `billable` # (no spend), but emphatically not read-only either. consumes_result=True, diff --git a/backend/mcp_server/tests/test_tool_execution.py b/backend/mcp_server/tests/test_tool_execution.py index 94c2872e82..917e409698 100644 --- a/backend/mcp_server/tests/test_tool_execution.py +++ b/backend/mcp_server/tests/test_tool_execution.py @@ -314,6 +314,20 @@ def test_completed_execution_is_enriched_and_returned(self) -> None: assert result["result"][0]["result"]["output"]["total"] == 42 assert result["execution_id"] == EXECUTION_ID + def test_stopped_execution_is_enriched_too(self) -> None: + response = ExecutionResponse( + workflow_id="wf", + execution_id=EXECUTION_ID, + execution_status="STOPPED", + result=[{"file": "invoice.pdf", "result": {"output": {"total": 42}}}], + ) + + result, process = self._run(response, include_metrics=True) + + process.assert_called_once() + assert process.call_args.kwargs["include_metrics"] is True + assert result["execution_status"] == "STOPPED" + def test_pending_execution_is_not_enriched(self) -> None: response = ExecutionResponse( workflow_id="wf", diff --git a/backend/mcp_server/tools/execution.py b/backend/mcp_server/tools/execution.py index 5ac468091d..d2bca3eee8 100644 --- a/backend/mcp_server/tools/execution.py +++ b/backend/mcp_server/tools/execution.py @@ -14,6 +14,7 @@ from tags.serializers import TagParamsSerializer from utils.enums import CeleryTaskState from workflow_manager.workflow_v2.dto import ExecutionResponse +from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.exceptions import ExecutionDoesNotExistError from workflow_manager.workflow_v2.models.execution import WorkflowExecution @@ -375,7 +376,11 @@ def get_execution_status( ), } - if response.execution_status == CeleryTaskState.COMPLETED.value: + # A stopped run returns its results so far, filtered the same way. + if response.execution_status in ( + CeleryTaskState.COMPLETED.value, + ExecutionStatus.STOPPED.value, + ): deployment_execution_dto = DeploymentExecutionDTO( api=context.api, api_key=context.api_key ) diff --git a/backend/mcp_server/tools/info.py b/backend/mcp_server/tools/info.py index e5bbbc2807..aa32278e62 100644 --- a/backend/mcp_server/tools/info.py +++ b/backend/mcp_server/tools/info.py @@ -47,10 +47,11 @@ def read_me_first(context: MCPContext) -> dict[str, Any]: "1. Call getApiInfo to confirm the deployment is active and see " "what it is meant to process.", "2. Call extractDocument with the URL(s) of the document(s).", - "3. If the response has execution_status COMPLETED, the result is " - "already in the response — you are done.", + "3. If the response has execution_status COMPLETED, ERROR or " + "STOPPED, it is final and any result is already in it — you are " + "done.", "4. Otherwise poll getExecutionStatus with the returned " - "execution_id until the status is COMPLETED or ERROR, pausing a " + "execution_id until the status is COMPLETED, ERROR or STOPPED, pausing a " "few seconds between polls.", ], "notes": [ diff --git a/backend/mcp_server/tools/platform.py b/backend/mcp_server/tools/platform.py index c5826f8bb1..5c61d36316 100644 --- a/backend/mcp_server/tools/platform.py +++ b/backend/mcp_server/tools/platform.py @@ -178,8 +178,8 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "to_run_an_extraction": ( "Call listApiDeployments for an api_name, then extractDocument " "with that name and your document URLs. If the response is not " - "already COMPLETED, poll getExecutionStatus with the returned " - "execution_id.\n\n" + "already COMPLETED, ERROR or STOPPED, poll getExecutionStatus with " + "the returned execution_id.\n\n" "The same extraction is also available on a deployment-scoped " "server at /deployment/api///mcp, which takes that " "one deployment's own API key and can reach nothing else. Both " diff --git a/backend/platform_audit/README.md b/backend/platform_audit/README.md index 3030b01336..fbfd252244 100644 --- a/backend/platform_audit/README.md +++ b/backend/platform_audit/README.md @@ -197,7 +197,7 @@ Expected emitters outside this PR: | Task | Action | Notes | |---|---|---| | PB-T05 (UN-4360) | `key.expiry_set` | target per key type; cloud review keys may register their own `review_key.*` actions | -| PB-T08 (UN-4318) | `execution.aborted` | target `execution`; actor is the owner or admin who aborted | +| PB-T08 (UN-4318) | `execution.aborted` | target `execution` | ## Action vocabulary diff --git a/backend/sample.env b/backend/sample.env index ce439c444d..0b96fdb503 100644 --- a/backend/sample.env +++ b/backend/sample.env @@ -295,6 +295,8 @@ EXECUTION_RESULT_TTL_SECONDS=86400 EXECUTION_CACHE_TTL_SECONDS=86400 # Instant workflow polling timeout in seconds (5 minutes) INSTANT_WF_POLLING_TIMEOUT=300 +# A PENDING/EXECUTING run with no activity for this long is listed as stuck (2.5 hours) +STUCK_EXECUTION_THRESHOLD_SECONDS=9000 # Maximum number of batches (i.e., parallel tasks) created for a single workflow execution (1 file at a time) MAX_PARALLEL_FILE_BATCHES=1 diff --git a/backend/utils/organization_utils.py b/backend/utils/organization_utils.py index b81252899d..f0cb742f4d 100644 --- a/backend/utils/organization_utils.py +++ b/backend/utils/organization_utils.py @@ -6,6 +6,7 @@ from typing import Any from account_v2.models import Organization +from django.db.models import QuerySet from django.shortcuts import get_object_or_404 logger = logging.getLogger(__name__) @@ -71,7 +72,9 @@ def get_organization_context(organization: Organization) -> dict[str, Any]: } -def filter_queryset_by_organization(queryset, request, organization_field="organization"): +def filter_queryset_by_organization( + queryset: QuerySet, request: Any, organization_field: str = "organization" +) -> QuerySet: """Filter a Django queryset by the request's organization context. Fails closed. For every caller, this function is the only tenant boundary diff --git a/backend/workflow_manager/execution/README.md b/backend/workflow_manager/execution/README.md new file mode 100644 index 0000000000..51b6ee50db --- /dev/null +++ b/backend/workflow_manager/execution/README.md @@ -0,0 +1,205 @@ +# Executions: stuck runs and abort + +Finding executions that stopped making progress, and stopping running ones +(PB-T08, UN-4318). This README is the published contract, backed by the +`@extend_schema` annotations in `views/stuck_abort.py` and `views/execution.py`. +Changes to anything below are contract changes. + +- [Endpoints](#endpoints) +- [What counts as stuck](#what-counts-as-stuck) +- [Listing stuck executions](#listing-stuck-executions) +- [Counting stuck executions](#counting-stuck-executions) +- [Aborting executions](#aborting-executions) +- [Deploying](#deploying) +- [Limits](#limits) + +## Endpoints + +Paths are relative to the organisation path: `/api/v1/unstract/{org_id}/`. + +| Method | Path | Returns | +|---|---|---| +| GET | `execution/?stuck=true` | The existing execution list, narrowed to stuck runs | +| GET | `execution/stuck-count/` | Stuck runs per entity, and the threshold | +| POST | `execution/abort/` | One outcome per requested id | + +Everything here is additive. The existing list, detail and log endpoints keep +their paths, parameters, response fields and permissions. Without `stuck`, or +with `stuck=false`, the list response is the same as before. Aborted runs are the +one new case the existing endpoints see; see +[How the work stops](#how-the-work-stops). + +## What counts as stuck + +An execution is stuck when its status is `PENDING` or `EXECUTING` and its last +activity is older than the threshold. Last activity is the later of the +execution row's `modified_at` and the newest `modified_at` of its file rows. + +The threshold is `STUCK_EXECUTION_THRESHOLD_SECONDS` (env, default **9000**, +2.5 hours). It is returned as `threshold_seconds` by `stuck-count/`. + +## Listing stuck executions + +`GET execution/?stuck=true` takes the list's existing filters, ordering and +pagination. Each row has the usual fields plus: + +| Field | Type | Meaning | +|---|---|---| +| `stuck_since` | ISO datetime | The time of last activity | + +`stuck` accepts `true`/`1` and `false`/`0`, case-insensitive. Any other value +is ignored, as if `stuck` were absent. Visibility is the same as the list's. + +## Counting stuck executions + +`GET execution/stuck-count/` returns 200: + +```json +{ + "threshold_seconds": 9000, + "total": 3, + "counts": {"API": 1, "ETL": 2, "TASK": 0, "WF": 0} +} +``` + +- `counts` uses the same keys as the list's `execution_entity` filter: `API` + deployments, `ETL` and `TASK` pipelines, and `WF` for workflow runs without a + pipeline. +- `total` counts every stuck run the caller can see. It can exceed the sum of + `counts` when a run's pipeline was deleted or is of another type. +- Visibility is the same as the list's. + +## Aborting executions + +`POST execution/abort/` with: + +```json +{"ids": ["3f0d6c1e-5a0b-4a39-9a7e-0d6f3c2b1a90", "..."]} +``` + +`ids` holds 1 to 100 execution ids; duplicates are answered once. It returns +200, with results in request order: + +```json +{ + "results": [ + {"id": "3f0d6c1e-5a0b-4a39-9a7e-0d6f3c2b1a90", "outcome": "aborted", "status": "STOPPED"}, + {"id": "57a4f9a2-2a4e-4a3c-8f0e-7c1d2b3a4f5e", "outcome": "already_finished", "status": "COMPLETED"}, + {"id": "9b1c2d3e-4f5a-4b6c-8d7e-0f1a2b3c4d5e", "outcome": "not_permitted", "status": null} + ], + "aborted": 1, + "already_finished": 1, + "not_permitted": 1, + "failed": 0 +} +``` + +| Outcome | When | `status` | +|---|---|---| +| `aborted` | The run was `PENDING` or `EXECUTING` and is now `STOPPED` | `STOPPED` | +| `already_finished` | The run was already `COMPLETED`, `ERROR` or `STOPPED`; nothing changed | its status | +| `not_permitted` | The caller may not abort it, or the id does not exist | `null` | +| `failed` | The abort hit an unexpected error for this id; retry it | `null` | + +Unknown ids and ids the caller cannot abort get the same answer, so the +endpoint does not reveal which ids exist. + +**Who may abort.** The owner of what started the run: an OWNER of its API +deployment or pipeline, or of its workflow when the run has no pipeline. Org +admins and service accounts (platform API keys with write access) may abort any +run in the organisation. Shared users who are not owners get `not_permitted`. + +**What abort does.** In one transaction per id: + +1. Sets the execution to `STOPPED` with the error `Execution aborted`, which + also frees an API deployment's rate-limit slot and sends the existing + "was stopped" inbox notice, unless the inbox item already shows a newer run. +2. Sets its `PENDING` and `EXECUTING` file rows to `STOPPED` with the same error, + and stamps their execution time. Finished file rows, file history, destination + writes and usage are kept. +3. Writes the run's file counts from its file rows: completed files are + successful, and every other file is failed. + +Once that commits, it: + +1. Sets the pipeline's last-run status to failure, unless the pipeline has a newer + run. +2. Sends the notifications of the run's pipeline (status `FAILURE`) or API + deployment (status `STOPPED`), paused pipelines included. They go out even if + step 1 fails. A failure in either step is logged with the run's id and does + not change the answer. +3. Records an `execution.aborted` event in the platform audit log, with target + `execution`, the workflow name, the summary `was ` and + `previous_status`, `pipeline_id` and `entity` in its metadata. This is best + effort: a lost event does not fail the abort and is not retried. + +### How the work stops + +- Queued batches of a stopped run are skipped when a worker picks them up, and + still count down on the barrier. So does a batch that fails after the abort, + unless its status write fails. +- A re-run ignores the active-file markers of a stopped run, so it keeps those + files even if no worker has picked up the queued batch yet. +- A running batch checks the status before each file and stops once it is + `STOPPED`: its remaining files are marked `STOPPED` and counted as failed. +- The file being processed when the abort lands finishes, and its own result + replaces `STOPPED` on its row. +- The last batch fires the aggregating callback. It keeps the status `STOPPED`, + saves the final file counts (the backend counts a stopped run's file rows, + whatever the worker sends), commits usage for completed files, publishes the + final log and cleans up the execution directory. It leaves the pipeline status + and notifications to the abort. +- A caller waiting on a synchronous API deployment `POST` gets 422 for an aborted + run, the same as polling with `GET`. Both apply the same result filters. +- Later status writes, from workers or the backend, cannot undo the abort or + replace its error. Code whose write is refused leaves the pipeline status and + notifications to the abort. + +### Errors + +400 uses the project's standard envelope (drf-standardized-errors). 401 and 403 +come from the auth middleware as `{"message": "..."}`. + +| Status | When | +|---|---| +| 400 | `ids` missing, empty, longer than 100, or holding a value that is not a UUID | +| 401 | Not authenticated | +| 403 | Authenticated but refused: a user outside the organisation, or a platform API key without write access | + +A failure on one id answers `failed` for that id, and the other ids still run. +Each id is aborted in its own transaction, so the others are kept (unless +`ATOMIC_REQUESTS` is on, which can roll the whole request back). Retrying is safe: +ids already aborted answer `already_finished`. + +## Deploying + +Deploy the backend before the workers. Workers read the status from the +internal `GET /internal/v1/workflow-execution//status/` endpoint. An older +backend answers 404, and the worker then processes the file anyway (fail +open), as it did before this change. Until the backend is updated, each file +also logs an `ERROR` and a `WARNING` for the failed check. + +Expose abort to users only once the workers are updated too. An older worker's +callback reports an aborted run with a completed file as completed. + +## Limits + +- The file being processed when the abort lands is not interrupted. A hung worker + is covered in [`pg_queue/README.md`](../../../workers/queue_backend/pg_queue/README.md). +- **Re-running right after an abort can process a file twice.** The file still + being processed no longer counts as active, so a manual or scheduled re-run can + pick it up while the old worker finishes it: double LLM spend, and two writes to + a database destination. Wait for that file to finish before re-running. +- **A dead worker's run never gets its callback**, and neither does a run whose + batch fails and then can't write its status, or a workflow or ETL run aborted + before a worker picks it up (the worker then skips it). The pipeline status, + notifications and file counts are still set by the abort, but usage for its + completed files is not committed and its execution directory, which can hold + uploaded documents, is not cleaned up. It stays in storage until something else + deletes it. The reaper's `ERROR` path has the same gap. +- The file being processed when the abort lands is billed once it finishes, but an + API caller can't get its result: the first read after the abort returns the + results so far and discards them, and later reads answer 406. Its processing cost + was already spent, the same as for a caller who never collects a result. +- A running batch stops between files only for `STOPPED`. A run that becomes + `ERROR` or `COMPLETED` mid-batch keeps processing its files, as before. diff --git a/backend/workflow_manager/execution/abort.py b/backend/workflow_manager/execution/abort.py new file mode 100644 index 0000000000..bdf0105a3e --- /dev/null +++ b/backend/workflow_manager/execution/abort.py @@ -0,0 +1,252 @@ +"""Abort running executions: stop the run, keep processed files, audit it.""" + +import logging +from dataclasses import dataclass +from typing import Any +from uuid import UUID + +from api_v2.models import APIDeployment +from api_v2.utils import APIDeploymentUtils +from django.db import transaction +from django.db.models import Count, DurationField, ExpressionWrapper, F, Q, QuerySet +from django.db.models.functions import Extract, Now +from django.utils import timezone +from permissions.roles import ResourceRole +from pipeline_v2.models import Pipeline +from pipeline_v2.notification import PipelineNotification +from platform_audit.actions import Action, TargetType +from platform_audit.recorder import record_platform_event +from tenant_account_v2.organization_member_service import OrganizationMemberService +from tenant_account_v2.sharing_helpers import resources_visible_via_memberships + +from workflow_manager.execution.entity import entity_annotation +from workflow_manager.execution.enum import AbortOutcome +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models import Workflow, WorkflowExecution + +logger = logging.getLogger(__name__) + +ABORT_ERROR_MESSAGE = "Execution aborted" +MAX_ABORT_IDS = 100 + + +@dataclass(frozen=True) +class AbortResult: + """Outcome for one id. ``status`` is ``None`` when not permitted or failed.""" + + id: UUID + outcome: AbortOutcome + status: str | None + + +@dataclass(frozen=True) +class AbortReport: + """Results in request order, with a count per outcome.""" + + results: list[AbortResult] + + def _count(self, outcome: AbortOutcome) -> int: + return sum(1 for result in self.results if result.outcome == outcome) + + @property + def aborted(self) -> int: + return self._count(AbortOutcome.ABORTED) + + @property + def already_finished(self) -> int: + return self._count(AbortOutcome.ALREADY_FINISHED) + + @property + def not_permitted(self) -> int: + return self._count(AbortOutcome.NOT_PERMITTED) + + @property + def failed(self) -> int: + return self._count(AbortOutcome.FAILED) + + +def abort_executions(request: Any, ids: list[UUID]) -> AbortReport: + """Abort each id the caller may abort; one transaction per id.""" + unique_ids = list(dict.fromkeys(ids)) + permitted = _abortable_entities(request.user, unique_ids) + results = [ + _abort_or_fail(request, execution_id, permitted[execution_id]) + if execution_id in permitted + else AbortResult(execution_id, AbortOutcome.NOT_PERMITTED, None) + for execution_id in unique_ids + ] + return AbortReport(results) + + +def _abortable_entities(user: Any, ids: list[UUID]) -> dict[UUID, str | None]: + """Map each id ``user`` may abort to its entity, in one query. + + Owners of the run's pipeline or API deployment (or of its workflow, for a + run without one), org admins and service accounts may abort. + """ + rows = ( + _abortable(user) + .filter(pk__in=ids) + .annotate(entity=entity_annotation()) + .values_list("pk", "entity") + ) + return dict(rows) + + +def _abortable(user: Any) -> QuerySet: + visible = WorkflowExecution.objects.for_user(user) + if getattr(user, "is_service_account", False): + return visible + if OrganizationMemberService.is_user_organization_admin(user): + return visible + return visible.filter(_owner_condition(user)) + + +def _owner_condition(user: Any) -> Q: + def owned(model: Any) -> QuerySet: + return resources_visible_via_memberships( + model, user, role=str(ResourceRole.OWNER) + ) + + return ( + Q(pipeline_id__isnull=True, workflow_id__in=owned(Workflow)) + | Q(pipeline_id__in=owned(APIDeployment)) + | Q(pipeline_id__in=owned(Pipeline)) + ) + + +def _abort_or_fail(request: Any, execution_id: UUID, entity: str | None) -> AbortResult: + """Abort one id, so its failure cannot lose the other ids' outcomes.""" + try: + return _abort_one(request, execution_id, entity) + except Exception: + logger.exception("Failed to abort execution %s", execution_id) + return AbortResult(execution_id, AbortOutcome.FAILED, None) + + +def _abort_one(request: Any, execution_id: UUID, entity: str | None) -> AbortResult: + """Stop one execution under its row lock, then record the audit event.""" + with transaction.atomic(): + execution = ( + WorkflowExecution.objects.select_for_update(no_key=True, of=("self",)) + .select_related("workflow") + .filter(pk=execution_id) + .first() + ) + if execution is None: + return AbortResult(execution_id, AbortOutcome.NOT_PERMITTED, None) + if execution.status in ExecutionStatus.terminal_values(): + return AbortResult( + execution_id, AbortOutcome.ALREADY_FINISHED, execution.status + ) + previous_status = execution.status + execution.update_execution( + status=ExecutionStatus.STOPPED, error=ABORT_ERROR_MESSAGE, row_locked=True + ) + _stop_open_files(execution) + write_file_counts(execution) + # After commit, so a rolled-back abort reports nothing. + transaction.on_commit(lambda: _report_abort(execution), robust=True) + _record_abort(request, execution, previous_status, entity) + return AbortResult(execution_id, AbortOutcome.ABORTED, execution.status) + + +def _stop_open_files(execution: WorkflowExecution) -> None: + """Mark file rows not yet terminal STOPPED; finished rows keep their status.""" + elapsed = ExpressionWrapper(Now() - F("created_at"), output_field=DurationField()) + stopped = ( + WorkflowFileExecution.objects.filter(workflow_execution=execution) + .exclude(status__in=ExecutionStatus.terminal_values()) + .update( + status=ExecutionStatus.STOPPED.value, + execution_error=ABORT_ERROR_MESSAGE, + execution_time=Extract(elapsed, "epoch"), + ) + ) + logger.info( + "Aborted execution %s; stopped %s open file execution(s)", execution.id, stopped + ) + + +def write_file_counts(execution: WorkflowExecution) -> None: + """Save the run's file counts from its rows; anything not completed is failed.""" + counts = WorkflowFileExecution.objects.filter(workflow_execution=execution).aggregate( + total=Count("pk"), + completed=Count("pk", filter=Q(status=ExecutionStatus.COMPLETED.value)), + ) + execution.total_files = max(execution.total_files, counts["total"]) + execution.successful_files = counts["completed"] + execution.failed_files = execution.total_files - counts["completed"] + execution.save(update_fields=["total_files", "successful_files", "failed_files"]) + + +def _report_abort(execution: WorkflowExecution) -> None: + """Report the ended run to its pipeline or API deployment. + + The abort does this because a dead worker's run gets no callback. + """ + try: + pipeline = Pipeline.objects.filter(pk=execution.pipeline_id).first() + if pipeline is not None: + _report_to_pipeline(pipeline, execution) + return + deployment = APIDeployment.objects.filter(pk=execution.pipeline_id).first() + if deployment is not None: + APIDeploymentUtils.send_notification( + api=deployment, workflow_execution=execution + ) + except Exception: + logger.exception( + "Failed to report aborted execution %s to %s", + execution.id, + execution.pipeline_id, + ) + + +def _report_to_pipeline(pipeline: Pipeline, execution: WorkflowExecution) -> None: + """Notify subscribers; the pipeline's own status follows only its newest run.""" + try: + _save_pipeline_status(pipeline, execution) + except Exception: + # Subscribers are still notified. + logger.exception( + "Failed to save pipeline %s status for aborted execution %s", + pipeline.id, + execution.id, + ) + PipelineNotification( + pipeline=pipeline, + execution_id=execution.id, + error_message=execution.error_message, + ).send() + + +def _save_pipeline_status(pipeline: Pipeline, execution: WorkflowExecution) -> None: + """Set the last-run status to failure, unless the pipeline has a newer run.""" + pipeline.last_run_status = Pipeline.PipelineStatus.FAILURE + newer_runs = WorkflowExecution.objects.filter( + pipeline_id=execution.pipeline_id, created_at__gt=execution.created_at + ) + if not newer_runs.exists(): + pipeline.last_run_time = timezone.now() + pipeline.save(update_fields=["last_run_status", "last_run_time"]) + + +def _record_abort( + request: Any, execution: WorkflowExecution, previous_status: str, entity: str | None +) -> None: + pipeline_id = str(execution.pipeline_id) if execution.pipeline_id else None + record_platform_event( + action=Action.EXECUTION_ABORTED, + target_type=TargetType.EXECUTION, + target_id=execution.id, + target_name=execution.workflow_name or "", + request=request, + summary=f"was {previous_status}", + metadata={ + "previous_status": previous_status, + "pipeline_id": pipeline_id, + "entity": entity, + }, + ) diff --git a/backend/workflow_manager/execution/entity.py b/backend/workflow_manager/execution/entity.py new file mode 100644 index 0000000000..0afdcc62c4 --- /dev/null +++ b/backend/workflow_manager/execution/entity.py @@ -0,0 +1,40 @@ +"""Which kind of source started a run: API deployment, ETL/TASK pipeline or workflow.""" + +from api_v2.models import APIDeployment +from django.db.models import Case, CharField, Exists, OuterRef, Q, Value, When +from pipeline_v2.models import Pipeline + +from workflow_manager.execution.enum import ExecutionEntity + +_PIPELINE_TYPES = { + ExecutionEntity.ETL: Pipeline.PipelineType.ETL, + ExecutionEntity.TASK: Pipeline.PipelineType.TASK, +} + + +def entity_condition(entity: ExecutionEntity) -> Exists | Q: + """Condition matching executions started by ``entity``. + + EXISTS subqueries keep the check in SQL instead of fetching all ids. + """ + if entity is ExecutionEntity.API: + return Exists(APIDeployment.objects.filter(id=OuterRef("pipeline_id"))) + if entity is ExecutionEntity.WORKFLOW: + return Q(pipeline_id__isnull=True) + return Exists( + Pipeline.objects.filter( + id=OuterRef("pipeline_id"), pipeline_type=_PIPELINE_TYPES[entity] + ) + ) + + +def entity_annotation() -> Case: + """Annotation naming each execution's entity; ``None`` when it matches none.""" + return Case( + *( + When(entity_condition(entity), then=Value(entity.value)) + for entity in ExecutionEntity + ), + default=Value(None), + output_field=CharField(), + ) diff --git a/backend/workflow_manager/execution/enum.py b/backend/workflow_manager/execution/enum.py index 27900508c0..48faa92b28 100644 --- a/backend/workflow_manager/execution/enum.py +++ b/backend/workflow_manager/execution/enum.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import Enum, StrEnum class ExecutionEntity(Enum): @@ -6,3 +6,12 @@ class ExecutionEntity(Enum): API = "API" TASK = "TASK" WORKFLOW = "WF" + + +class AbortOutcome(StrEnum): + """Per-id answer of the abort endpoint.""" + + ABORTED = "aborted" + ALREADY_FINISHED = "already_finished" + NOT_PERMITTED = "not_permitted" + FAILED = "failed" diff --git a/backend/workflow_manager/execution/filter.py b/backend/workflow_manager/execution/filter.py index d62a96432a..9dbae68f88 100644 --- a/backend/workflow_manager/execution/filter.py +++ b/backend/workflow_manager/execution/filter.py @@ -1,13 +1,23 @@ -from api_v2.models import APIDeployment -from django.db.models import Exists, OuterRef, QuerySet +from typing import Any + +from django.db.models import QuerySet from django_filters import rest_framework as filters -from pipeline_v2.models import Pipeline +from django_filters.widgets import BooleanWidget from utils.date import DateRangePresets, DateTimeProcessor +from workflow_manager.execution.entity import entity_condition from workflow_manager.execution.enum import ExecutionEntity +from workflow_manager.execution.stuck import stuck_executions from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.models import WorkflowExecution +STUCK_PARAM = "stuck" + + +def stuck_requested(query_params: Any) -> bool: + """True when ``stuck`` parses as true, by the same rule as the filter.""" + return BooleanWidget().value_from_datadict(query_params, None, STUCK_PARAM) is True + class ExecutionFilter(filters.FilterSet): id = filters.UUIDFilter() @@ -29,6 +39,8 @@ class ExecutionFilter(filters.FilterSet): choices=DateRangePresets.choices(), method="filter_by_date_range", ) + # Documented on ExecutionViewSet's schema. + stuck = filters.BooleanFilter(method="filter_stuck") class Meta: model = WorkflowExecution @@ -40,38 +52,13 @@ def filter_execution_entity( """Filter executions by entity type efficiently. The queryset is already filtered by user permissions via get_queryset(), - so we use EXISTS subqueries to check entity type without fetching all IDs. - This is more efficient than values_list() for large datasets. + so EXISTS subqueries check the entity type without fetching all IDs. """ - if value == ExecutionEntity.API.value: - # Filter for API deployments using EXISTS for efficiency - return queryset.filter( - Exists(APIDeployment.objects.filter(id=OuterRef("pipeline_id"))) - ) - elif value == ExecutionEntity.ETL.value: - # Filter for ETL pipelines using EXISTS - return queryset.filter( - Exists( - Pipeline.objects.filter( - id=OuterRef("pipeline_id"), - pipeline_type=Pipeline.PipelineType.ETL, - ) - ) - ) - elif value == ExecutionEntity.TASK.value: - # Filter for TASK pipelines using EXISTS - return queryset.filter( - Exists( - Pipeline.objects.filter( - id=OuterRef("pipeline_id"), - pipeline_type=Pipeline.PipelineType.TASK, - ) - ) - ) - elif value == ExecutionEntity.WORKFLOW.value: - # Filter for workflow-level executions (no pipeline) - return queryset.filter(pipeline_id__isnull=True) - return queryset + return queryset.filter(entity_condition(ExecutionEntity(value))) + + def filter_stuck(self, queryset: QuerySet, name: str, value: bool) -> QuerySet: + """Narrow to stuck executions when true; false leaves the queryset as is.""" + return stuck_executions(queryset) if value else queryset def filter_by_date_range(self, queryset: QuerySet, name: str, value: str) -> QuerySet: """Filters Usages based on the provided date range.""" diff --git a/backend/workflow_manager/execution/serializer/stuck_abort.py b/backend/workflow_manager/execution/serializer/stuck_abort.py new file mode 100644 index 0000000000..611c75729f --- /dev/null +++ b/backend/workflow_manager/execution/serializer/stuck_abort.py @@ -0,0 +1,51 @@ +from rest_framework import serializers + +from workflow_manager.execution.abort import MAX_ABORT_IDS +from workflow_manager.execution.enum import AbortOutcome +from workflow_manager.execution.serializer.execution import ExecutionSerializer + + +class StuckExecutionSerializer(ExecutionSerializer): + """An execution row plus ``stuck_since``, served for ``?stuck=true``.""" + + stuck_since = serializers.DateTimeField( + read_only=True, help_text="Time of last activity on the run or its files." + ) + + +class StuckExecutionCountsSerializer(serializers.Serializer): + API = serializers.IntegerField() + ETL = serializers.IntegerField() + TASK = serializers.IntegerField() + WF = serializers.IntegerField() + + +class StuckExecutionCountSerializer(serializers.Serializer): + threshold_seconds = serializers.IntegerField() + total = serializers.IntegerField() + counts = StuckExecutionCountsSerializer() + + +class ExecutionAbortRequestSerializer(serializers.Serializer): + ids = serializers.ListField( + child=serializers.UUIDField(), + min_length=1, + max_length=MAX_ABORT_IDS, + help_text=f"1 to {MAX_ABORT_IDS} execution ids; duplicates are removed.", + ) + + +class ExecutionAbortResultSerializer(serializers.Serializer): + id = serializers.UUIDField() + outcome = serializers.ChoiceField(choices=[o.value for o in AbortOutcome]) + status = serializers.CharField( + allow_null=True, help_text="The run's status; null when not_permitted or failed." + ) + + +class ExecutionAbortResponseSerializer(serializers.Serializer): + results = ExecutionAbortResultSerializer(many=True) + aborted = serializers.IntegerField() + already_finished = serializers.IntegerField() + not_permitted = serializers.IntegerField() + failed = serializers.IntegerField() diff --git a/backend/workflow_manager/execution/stuck.py b/backend/workflow_manager/execution/stuck.py new file mode 100644 index 0000000000..292d9f9a46 --- /dev/null +++ b/backend/workflow_manager/execution/stuck.py @@ -0,0 +1,65 @@ +"""Stuck executions: PENDING/EXECUTING runs with no activity past a threshold.""" + +from dataclasses import dataclass +from datetime import timedelta + +from django.conf import settings +from django.db.models import Count, OuterRef, QuerySet, Subquery +from django.db.models.functions import Greatest +from django.utils import timezone + +from workflow_manager.execution.entity import entity_condition +from workflow_manager.execution.enum import ExecutionEntity +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models import WorkflowExecution + +ACTIVE_STATUSES = (ExecutionStatus.PENDING.value, ExecutionStatus.EXECUTING.value) + + +@dataclass(frozen=True) +class StuckCount: + """Stuck executions visible to a caller, in total and per entity.""" + + threshold_seconds: int + total: int + counts: dict[str, int] + + +def stuck_threshold_seconds() -> int: + return int(settings.STUCK_EXECUTION_THRESHOLD_SECONDS) + + +def stuck_executions(queryset: QuerySet) -> QuerySet: + """Narrow ``queryset`` to stuck runs, annotated with ``stuck_since``. + + ``stuck_since`` is the last activity: the later of the run's own + ``modified_at`` and its newest file row's ``modified_at``. + """ + cutoff = timezone.now() - timedelta(seconds=stuck_threshold_seconds()) + last_file_at = ( + WorkflowFileExecution.objects.filter(workflow_execution=OuterRef("pk")) + .order_by("-modified_at") + .values("modified_at")[:1] + ) + # Postgres GREATEST skips NULL, so a run with no files uses its own time. + return ( + queryset.filter(status__in=ACTIVE_STATUSES, modified_at__lt=cutoff) + .annotate(stuck_since=Greatest("modified_at", Subquery(last_file_at))) + .filter(stuck_since__lt=cutoff) + ) + + +def count_stuck(queryset: QuerySet) -> StuckCount: + """Count the stuck runs in ``queryset``, in one query.""" + per_entity = { + entity.value: Count("pk", filter=entity_condition(entity)) + for entity in ExecutionEntity + } + counts = WorkflowExecution.objects.filter( + pk__in=stuck_executions(queryset).values("pk") + ).aggregate(total=Count("pk"), **per_entity) + total = counts.pop("total") + return StuckCount( + threshold_seconds=stuck_threshold_seconds(), total=total, counts=counts + ) diff --git a/backend/workflow_manager/execution/tests/stuck_abort_base.py b/backend/workflow_manager/execution/tests/stuck_abort_base.py new file mode 100644 index 0000000000..a3fe305381 --- /dev/null +++ b/backend/workflow_manager/execution/tests/stuck_abort_base.py @@ -0,0 +1,122 @@ +"""Fixtures for the stuck-execution and abort tests (UN-4318). + +``self.owner`` owns the workflow, deployments and pipelines built here; +``self.admin`` is the only org admin; ``self.outsider`` has no share. +""" + +import secrets +from datetime import timedelta +from typing import Any +from unittest.mock import patch + +from api_v2.models import APIDeployment +from django.test import override_settings +from django.utils import timezone +from permissions.roles import ResourceRole +from pipeline_v2.models import Pipeline +from rest_framework.test import APIRequestFactory, force_authenticate +from tenant_account_v2.tests import GroupSharingTestBase + +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution + +ADMIN_PREDICATE = ( + "tenant_account_v2.organization_member_service." + "OrganizationMemberService.is_user_organization_admin" +) +# Past the 9000 s threshold set below, and well inside it. +OLD = timedelta(hours=3) +RECENT = timedelta(minutes=1) + + +@override_settings(STUCK_EXECUTION_THRESHOLD_SECONDS=9000) +class StuckAbortTestBase(GroupSharingTestBase): + def setUp(self) -> None: + super().setUp() + patcher = patch( + ADMIN_PREDICATE, + side_effect=lambda user: getattr(user, "email", None) == self.admin.email, + ) + patcher.start() + self.addCleanup(patcher.stop) + self.factory = APIRequestFactory() + + def deployment(self) -> APIDeployment: + deployment: APIDeployment = APIDeployment.objects.create( + api_name=f"api-{secrets.token_hex(4)}", + workflow=self.workflow, + organization=self.org, + created_by=self.owner, + ) + deployment.memberships.create(user=self.owner, role=ResourceRole.OWNER) + return deployment + + def pipeline(self, pipeline_type: Any = Pipeline.PipelineType.ETL) -> Pipeline: + pipeline: Pipeline = Pipeline.objects.create( + pipeline_name=f"pipe-{secrets.token_hex(4)}", + workflow=self.workflow, + organization=self.org, + created_by=self.owner, + pipeline_type=pipeline_type, + ) + pipeline.memberships.create(user=self.owner, role=ResourceRole.OWNER) + return pipeline + + def execution( + self, + resource: Any = None, + status: ExecutionStatus = ExecutionStatus.EXECUTING, + idle: timedelta = OLD, + queue_message_id: int | None = 1, + ) -> WorkflowExecution: + """A run whose last own write was ``idle`` ago. + + Scheduled and API-triggered pipeline runs have no ``queue_message_id``. + """ + execution: WorkflowExecution = WorkflowExecution.objects.create( + workflow=self.workflow, + pipeline_id=resource.id if resource else None, + execution_mode=WorkflowExecution.Mode.QUEUE, + execution_method=WorkflowExecution.Method.DIRECT, + execution_type=WorkflowExecution.Type.COMPLETE, + status=status.value, + queue_message_id=queue_message_id, + ) + self.backdate(WorkflowExecution, execution.pk, idle) + execution.refresh_from_db() + return execution + + def file( + self, execution: WorkflowExecution, status: ExecutionStatus, idle: timedelta + ) -> WorkflowFileExecution: + row: WorkflowFileExecution = WorkflowFileExecution.objects.create( + workflow_execution=execution, file_name="f.pdf", status=status.value + ) + self.backdate(WorkflowFileExecution, row.pk, idle) + row.refresh_from_db() + return row + + @staticmethod + def backdate(model: Any, pk: Any, idle: timedelta) -> None: + # update() skips auto_now, so modified_at keeps the backdated value. + model.objects.filter(pk=pk).update(modified_at=timezone.now() - idle) + + def call( + self, + view: Any, + user: Any, + method: str = "get", + data: Any = None, + **url_kwargs: Any, + ) -> Any: + """Run ``view`` as ``user`` and every on-commit callback it registers.""" + if method == "get": + request = self.factory.get("/", data or {}) + else: + request = self.factory.post("/", data, format="json") + force_authenticate(request, user=user) + with self.captureOnCommitCallbacks(execute=True): + response = view(request, **url_kwargs) + response.render() + return response diff --git a/backend/workflow_manager/execution/tests/test_abort_executions.py b/backend/workflow_manager/execution/tests/test_abort_executions.py new file mode 100644 index 0000000000..38622294a0 --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_abort_executions.py @@ -0,0 +1,507 @@ +"""``POST execution/abort/`` (UN-4318): owner-or-admin, per-id outcomes, audit.""" + +import uuid +from collections.abc import Callable +from datetime import timedelta +from typing import Any +from unittest.mock import MagicMock, patch + +from account_v2.models import Organization +from django.db import connection +from django.test.utils import CaptureQueriesContext +from django.utils import timezone +from permissions.roles import ResourceRole +from pipeline_v2.models import Pipeline +from platform_audit.actions import Action, TargetType +from platform_audit.models import PlatformAuditEvent +from tenant_account_v2.tests import _add_viewers, _make_user + +from workflow_manager.execution.abort import ABORT_ERROR_MESSAGE, MAX_ABORT_IDS +from workflow_manager.execution.tests import test_update_status_locking as locking +from workflow_manager.execution.tests.stuck_abort_base import RECENT, StuckAbortTestBase +from workflow_manager.execution.views import ExecutionAbortView +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.file_execution.serializers import FileCentricExecutionSerializer +from workflow_manager.internal_views import WorkflowExecutionInternalViewSet +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.execution import WorkflowExecutionServiceHelper +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow +from workflow_manager.workflow_v2.workflow_helper import WorkflowHelper + +_abort = ExecutionAbortView.as_view() + + +class AbortExecutionTests(StuckAbortTestBase): + def _abort(self, user: Any, *ids: Any, raw: dict[str, Any] | None = None) -> Any: + body = raw if raw is not None else {"ids": [str(i) for i in ids]} + return self.call(_abort, user, method="post", data=body) + + def _outcomes(self, response: Any) -> dict[str, tuple[str, str | None]]: + self.assertEqual(response.status_code, 200, response.data) + return { + row["id"]: (row["outcome"], row["status"]) + for row in response.data["results"] + } + + def _status(self, execution: WorkflowExecution) -> str: + execution.refresh_from_db() + status: str = execution.status + return status + + def test_abort_stops_the_run_and_its_open_files_only(self) -> None: + execution = self.execution(self.deployment(), idle=RECENT) + done = self.file(execution, ExecutionStatus.COMPLETED, idle=RECENT) + running = self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + queued = self.file(execution, ExecutionStatus.PENDING, idle=RECENT) + failed = self.file(execution, ExecutionStatus.ERROR, idle=RECENT) + WorkflowFileExecution.objects.filter(pk=failed.pk).update(execution_error="boom") + + response = self._abort(self.owner, execution.id) + + self.assertEqual( + self._outcomes(response), {str(execution.id): ("aborted", "STOPPED")} + ) + self.assertEqual( + (response.data["aborted"], response.data["already_finished"]), (1, 0) + ) + execution.refresh_from_db() + self.assertEqual(execution.status, ExecutionStatus.STOPPED.value) + self.assertEqual(execution.error_message, ABORT_ERROR_MESSAGE) + for row, expected in ( + (done, ExecutionStatus.COMPLETED), + (running, ExecutionStatus.STOPPED), + (queued, ExecutionStatus.STOPPED), + (failed, ExecutionStatus.ERROR), + ): + row.refresh_from_db() + self.assertEqual(row.status, expected.value) + self.assertIsNone(done.execution_error) + self.assertEqual(failed.execution_error, "boom") + self.assertEqual(running.execution_error, ABORT_ERROR_MESSAGE) + self.assertEqual(queued.execution_error, ABORT_ERROR_MESSAGE) + + def test_stopped_file_row_shows_the_abort_and_a_fixed_time(self) -> None: + execution = self.execution(idle=RECENT) + row = self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + started = timezone.now() - timedelta(hours=1) + WorkflowFileExecution.objects.filter(pk=row.pk).update(created_at=started) + + self._abort(self.owner, execution.id) + + row.refresh_from_db() + self.assertAlmostEqual(row.execution_time, 3600, delta=60) + self.assertEqual( + FileCentricExecutionSerializer().get_status_msg(row), ABORT_ERROR_MESSAGE + ) + + def test_abort_writes_the_file_counts_counting_the_unfinished_as_failed(self) -> None: + for known_total, total in ((5, 5), (0, 3)): + with self.subTest(known_total=known_total): + execution = self.execution(idle=RECENT) + WorkflowExecution.objects.filter(pk=execution.pk).update( + total_files=known_total + ) + self.file(execution, ExecutionStatus.COMPLETED, idle=RECENT) + self.file(execution, ExecutionStatus.ERROR, idle=RECENT) + self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + + self._abort(self.owner, execution.id) + + execution.refresh_from_db() + self.assertEqual( + ( + execution.total_files, + execution.successful_files, + execution.failed_files, + ), + (total, 1, total - 1), + ) + + def test_abort_releases_the_slot_stamps_the_time_and_notifies_the_inbox(self) -> None: + execution = self.execution(self.deployment()) + started = timezone.now() - timedelta(hours=1) + WorkflowExecution.objects.filter(pk=execution.pk).update(created_at=started) + with ( + patch( + "api_v2.rate_limiter.APIDeploymentRateLimiter.release_slot" + ) as release, + patch("inbox.events.notify_execution_failed") as notify, + ): + self._abort(self.owner, execution.id) + release.assert_called_once_with(self.org.organization_id, str(execution.id)) + notify.assert_called_once_with(execution.id) + execution.refresh_from_db() + self.assertAlmostEqual(execution.execution_time, 3600, delta=60) + + def _abort_with_notification_patched(self, execution: WorkflowExecution) -> Any: + with patch("workflow_manager.execution.abort.PipelineNotification") as notify: + self._abort(self.owner, execution.id) + notify.return_value.send.assert_called_once() + return notify.call_args.kwargs + + def test_abort_marks_the_pipeline_failed_and_notifies(self) -> None: + for active in (True, False): + with self.subTest(active=active): + pipeline = self.pipeline() + Pipeline.objects.filter(pk=pipeline.pk).update(active=active) + execution = self.execution(pipeline) + + sent = self._abort_with_notification_patched(execution) + + pipeline.refresh_from_db() + self.assertEqual( + pipeline.last_run_status, Pipeline.PipelineStatus.FAILURE + ) + self.assertIsNotNone(pipeline.last_run_time) + self.assertEqual(sent["execution_id"], execution.id) + + def test_a_failed_pipeline_save_still_notifies(self) -> None: + execution = self.execution(self.pipeline()) + with ( + patch.object(Pipeline, "save", side_effect=RuntimeError("boom")), + self.assertLogs("workflow_manager.execution.abort", level="ERROR") as logs, + ): + sent = self._abort_with_notification_patched(execution) + + self.assertEqual(sent["execution_id"], execution.id) + self.assertIn(str(execution.id), logs.output[0]) + + def test_abort_of_an_old_run_notifies_but_keeps_the_pipeline_status(self) -> None: + pipeline = self.pipeline() + old_run = self.execution(pipeline) + self.execution(pipeline, status=ExecutionStatus.COMPLETED) + Pipeline.objects.filter(pk=pipeline.pk).update( + last_run_status=Pipeline.PipelineStatus.SUCCESS + ) + + sent = self._abort_with_notification_patched(old_run) + + self.assertEqual(sent["pipeline"].last_run_status, Pipeline.PipelineStatus.FAILURE) + pipeline.refresh_from_db() + self.assertEqual(pipeline.last_run_status, Pipeline.PipelineStatus.SUCCESS) + self.assertIsNone(pipeline.last_run_time) + + def test_abort_of_a_deployment_run_notifies_the_deployment(self) -> None: + execution = self.execution(self.deployment()) + with patch( + "workflow_manager.execution.abort.APIDeploymentUtils.send_notification" + ) as send: + self._abort(self.owner, execution.id) + send.assert_called_once() + self.assertEqual(send.call_args.kwargs["workflow_execution"].id, execution.id) + + def test_finished_run_announces_nothing(self) -> None: + execution = self.execution(self.pipeline(), status=ExecutionStatus.COMPLETED) + with patch("workflow_manager.execution.abort._report_abort") as report: + self._abort(self.owner, execution.id) + report.assert_not_called() + + def test_abort_takes_one_row_lock(self) -> None: + execution = self.execution() + # The inbox notice takes its own lock on commit. + with ( + patch("inbox.events.notify_execution_failed"), + CaptureQueriesContext(connection) as queries, + ): + self._abort(self.owner, execution.id) + lock_statements: Any = locking._lock_statements + _, no_key, full = lock_statements(queries.captured_queries) + self.assertEqual(len(no_key), 1) + self.assertEqual(full, []) + + def test_aborting_twice_answers_already_finished(self) -> None: + execution = self.execution() + self._abort(self.owner, execution.id) + outcome = self._outcomes(self._abort(self.owner, execution.id)) + self.assertEqual(outcome[str(execution.id)], ("already_finished", "STOPPED")) + self.assertEqual(PlatformAuditEvent.objects.count(), 1) + + def test_one_failing_id_answers_failed_in_a_200(self) -> None: + bad, good = self.execution(), self.execution() + + def stop(execution: WorkflowExecution) -> None: + if execution.pk == bad.pk: + raise RuntimeError("boom") + + with patch("workflow_manager.execution.abort._stop_open_files", side_effect=stop): + response = self._abort(self.owner, bad.id, good.id) + + self.assertEqual( + self._outcomes(response), + {str(bad.id): ("failed", None), str(good.id): ("aborted", "STOPPED")}, + ) + self.assertEqual((response.data["aborted"], response.data["failed"]), (1, 1)) + + def test_abort_is_audited(self) -> None: + deployment = self.deployment() + execution = self.execution(deployment) + + self._abort(self.owner, execution.id) + + event = PlatformAuditEvent.objects.get(action=Action.EXECUTION_ABORTED) + self.assertEqual(event.organization, self.org) + self.assertEqual(event.actor, self.owner) + self.assertEqual(event.target_type, TargetType.EXECUTION) + self.assertEqual(event.target_id, str(execution.id)) + self.assertEqual(event.target_name, self.workflow.workflow_name) + self.assertEqual(event.summary, "was EXECUTING") + self.assertEqual( + event.metadata, + { + "previous_status": "EXECUTING", + "pipeline_id": str(deployment.id), + "entity": "API", + }, + ) + + def test_finished_run_is_already_finished_and_not_audited(self) -> None: + execution = self.execution(status=ExecutionStatus.COMPLETED) + response = self._abort(self.owner, execution.id) + self.assertEqual( + self._outcomes(response), + {str(execution.id): ("already_finished", "COMPLETED")}, + ) + self.assertEqual(self._status(execution), ExecutionStatus.COMPLETED.value) + self.assertFalse(PlatformAuditEvent.objects.exists()) + + def test_mixed_batch_answers_each_id_in_one_200(self) -> None: + mine = self.execution(self.pipeline()) + finished = self.execution(status=ExecutionStatus.ERROR) + other = _make_user("other@example.com") + foreign_workflow_run = self._run_owned_by(other) + unknown = uuid.uuid4() + + response = self._abort( + self.owner, mine.id, finished.id, foreign_workflow_run.id, unknown + ) + + self.assertEqual( + self._outcomes(response), + { + str(mine.id): ("aborted", "STOPPED"), + str(finished.id): ("already_finished", "ERROR"), + str(foreign_workflow_run.id): ("not_permitted", None), + str(unknown): ("not_permitted", None), + }, + ) + self.assertEqual( + [r["id"] for r in response.data["results"]], + [str(mine.id), str(finished.id), str(foreign_workflow_run.id), str(unknown)], + ) + self.assertEqual( + ( + response.data["aborted"], + response.data["already_finished"], + response.data["not_permitted"], + ), + (1, 1, 2), + ) + self.assertEqual( + self._status(foreign_workflow_run), ExecutionStatus.EXECUTING.value + ) + + def _run_owned_by(self, user: Any) -> WorkflowExecution: + """A run of a deployment ``user`` owns, which ``self.owner`` only sees.""" + deployment = self.deployment() + deployment.memberships.filter(user=self.owner).delete() + deployment.memberships.create(user=user, role=ResourceRole.OWNER) + _add_viewers(deployment, self.owner) + return self.execution(deployment) + + def test_another_orgs_run_is_not_permitted_for_admin_or_service_account(self) -> None: + other_org = Organization.objects.create( + name="org-b", display_name="Org B", organization_id="org-b" + ) + other_workflow = Workflow.objects.create( + workflow_name="wf-b", organization=other_org, created_by=self.owner + ) + foreign = WorkflowExecution.objects.create( + workflow=other_workflow, + execution_mode=WorkflowExecution.Mode.QUEUE, + execution_method=WorkflowExecution.Method.DIRECT, + execution_type=WorkflowExecution.Type.COMPLETE, + status=ExecutionStatus.EXECUTING.value, + queue_message_id=1, + ) + service = _make_user("svc@example.com", is_service_account=True) + + for user in (self.admin, service): + with self.subTest(user=user.email): + outcome = self._outcomes(self._abort(user, foreign.id)) + self.assertEqual(outcome[str(foreign.id)], ("not_permitted", None)) + + self.assertEqual(self._status(foreign), ExecutionStatus.EXECUTING.value) + self.assertFalse(PlatformAuditEvent.objects.exists()) + + def test_shared_viewer_is_not_permitted(self) -> None: + deployment = self.deployment() + _add_viewers(deployment, self.outsider) + execution = self.execution(deployment) + # Control: the viewer can see the run. + self.assertTrue( + WorkflowExecution.objects.for_user(self.outsider) + .filter(pk=execution.pk) + .exists() + ) + + response = self._abort(self.outsider, execution.id) + + self.assertEqual( + self._outcomes(response), {str(execution.id): ("not_permitted", None)} + ) + self.assertEqual(self._status(execution), ExecutionStatus.EXECUTING.value) + + def test_workflow_owner_cannot_abort_a_pipeline_run_they_do_not_own(self) -> None: + pipeline = self.pipeline() + pipeline.memberships.filter(user=self.owner).delete() + Pipeline.objects.filter(pk=pipeline.pk).update(shared_to_org=True) + execution = self.execution(pipeline) + outcome = self._outcomes(self._abort(self.owner, execution.id)) + self.assertEqual(outcome[str(execution.id)][0], "not_permitted") + + def test_co_owner_of_the_workflow_aborts_a_workflow_run(self) -> None: + self.workflow.memberships.create(user=self.member, role=ResourceRole.OWNER) + execution = self.execution(status=ExecutionStatus.PENDING) + outcome = self._outcomes(self._abort(self.member, execution.id)) + self.assertEqual(outcome[str(execution.id)], ("aborted", "STOPPED")) + + def test_org_admin_can_abort_any_run(self) -> None: + execution = self.execution(self.deployment()) + outcome = self._outcomes(self._abort(self.admin, execution.id)) + self.assertEqual(outcome[str(execution.id)], ("aborted", "STOPPED")) + event = PlatformAuditEvent.objects.get(action=Action.EXECUTION_ABORTED) + self.assertEqual(event.actor, self.admin) + + def test_service_account_can_abort(self) -> None: + service = _make_user("svc@example.com", is_service_account=True) + execution = self.execution(self.pipeline()) + outcome = self._outcomes(self._abort(service, execution.id)) + self.assertEqual(outcome[str(execution.id)], ("aborted", "STOPPED")) + + def test_duplicate_ids_are_answered_once(self) -> None: + execution = self.execution() + response = self._abort(self.owner, execution.id, execution.id) + self.assertEqual(len(response.data["results"]), 1) + self.assertEqual(PlatformAuditEvent.objects.count(), 1) + + def test_invalid_bodies_are_400(self) -> None: + self.assertEqual(MAX_ABORT_IDS, 100) + too_many = [str(uuid.uuid4()) for _ in range(101)] + bodies: list[dict[str, Any]] = [ + {}, + {"ids": []}, + {"ids": ["not-a-uuid"]}, + {"ids": too_many}, + ] + for body in bodies: + with self.subTest(body=body): + response = self._abort(self.owner, raw=body) + self.assertEqual(response.status_code, 400) + + def test_requires_authentication(self) -> None: + request = self.factory.post("/", {"ids": [str(uuid.uuid4())]}, format="json") + response = _abort(request) + self.assertIn(response.status_code, (401, 403)) + + def test_later_writes_cannot_undo_the_abort(self) -> None: + # Scheduled and API-triggered pipeline runs have no queue_message_id. + for queue_message_id in (1, None): + for status in (ExecutionStatus.COMPLETED, ExecutionStatus.ERROR): + with self.subTest(queue_message_id=queue_message_id, status=status): + execution = self.execution( + self.pipeline(), queue_message_id=queue_message_id + ) + self._abort(self.owner, execution.id) + + out = self._write_status(execution, {"status": status.value}) + + self.assertEqual(out.data.get("reason"), "already_final") + self.assertEqual(out.data.get("new_status"), "STOPPED") + self.assertEqual( + self._status(execution), ExecutionStatus.STOPPED.value + ) + + def test_backend_writes_before_dispatch_cannot_undo_the_abort(self) -> None: + # Staging or dispatch fails, or stages nothing, after a PENDING run's abort. + helper = WorkflowExecutionServiceHelper + writes: dict[str, Callable[[str], Any]] = { + "error": lambda pk: helper.update_execution_err(pk, "staging failed"), + "completed": lambda pk: helper.update_execution_completed(pk, 2, 2), + } + for name, write in writes.items(): + with self.subTest(write=name): + execution = self.execution( + self.deployment(), + status=ExecutionStatus.PENDING, + queue_message_id=None, + ) + self._abort(self.owner, execution.id) + + write(str(execution.pk)) + + execution.refresh_from_db() + self.assertEqual(execution.status, ExecutionStatus.STOPPED.value) + self.assertEqual((execution.total_files, execution.failed_files), (0, 0)) + self.assertEqual(execution.error_message, ABORT_ERROR_MESSAGE) + + def test_a_failed_dispatch_after_the_abort_answers_stopped(self) -> None: + execution = self.execution( + self.deployment(), status=ExecutionStatus.PENDING, queue_message_id=None + ) + self._abort(self.owner, execution.id) + + with patch.object( + WorkflowHelper, "_dispatch_orchestrator_task", side_effect=RuntimeError("down") + ): + response = WorkflowHelper.execute_workflow_async( + workflow_id=str(self.workflow.id), + execution_id=str(execution.id), + hash_values_of_files={}, + ) + + self.assertEqual(response.execution_status, ExecutionStatus.STOPPED.value) + self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) + + @staticmethod + def _write_status(execution: WorkflowExecution, data: dict[str, Any]) -> Any: + """A worker's status write through the internal API.""" + view = WorkflowExecutionInternalViewSet() + request = MagicMock() + request.data = data + fresh = WorkflowExecution.objects.get(pk=execution.pk) + with patch.object(view, "get_object", return_value=fresh): + return view.update_status(request, id=str(execution.id)) + + def test_a_stopped_write_saves_the_counts_of_the_file_rows(self) -> None: + execution = self.execution(self.pipeline()) + for status in (ExecutionStatus.COMPLETED, ExecutionStatus.COMPLETED): + self.file(execution, status, idle=RECENT) + in_flight = self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + self._abort(self.owner, execution.id) + # The file in progress at the abort finishes afterwards. + WorkflowFileExecution.objects.filter(pk=in_flight.pk).update( + status=ExecutionStatus.COMPLETED.value + ) + + # The worker reports zeros, as a failed callback's retry does. + out = self._write_status( + execution, + { + "status": ExecutionStatus.STOPPED.value, + "total_files": 0, + "successful_files": 0, + "failed_files": 0, + }, + ) + + self.assertEqual(out.data.get("status"), "updated") + execution.refresh_from_db() + self.assertEqual( + ( + execution.total_files, + execution.successful_files, + execution.failed_files, + ), + (3, 3, 0), + ) diff --git a/backend/workflow_manager/execution/tests/test_abort_transactions.py b/backend/workflow_manager/execution/tests/test_abort_transactions.py new file mode 100644 index 0000000000..750141564b --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_abort_transactions.py @@ -0,0 +1,155 @@ +"""Abort on real transactions: the row lock and one commit per id (UN-4318).""" + +import threading +from typing import Any +from unittest.mock import patch + +from account_v2.models import Organization +from api_v2.models import APIDeployment +from django.db import transaction +from django.test import TransactionTestCase +from permissions.roles import ResourceRole +from platform_audit.models import PlatformAuditEvent +from rest_framework.test import APIRequestFactory, force_authenticate +from tenant_account_v2.tests import _make_user +from utils.user_context import UserContext + +from workflow_manager.execution.tests import test_update_status_locking as locking +from workflow_manager.execution.views import ExecutionAbortView +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow + +_ABORT = "workflow_manager.execution.abort" +_abort = ExecutionAbortView.as_view() + + +class AbortTransactionTests(TransactionTestCase): + def setUp(self) -> None: + self.org = Organization.objects.create( + name="org-a", display_name="Org A", organization_id="org-a" + ) + UserContext.set_organization_identifier(self.org.organization_id) + self.owner = _make_user("owner@example.com") + self.workflow = Workflow.objects.create( + workflow_name="wf-1", organization=self.org, created_by=self.owner + ) + self.workflow.memberships.create(user=self.owner, role=ResourceRole.OWNER) + + def _execution(self, pipeline_id: Any = None) -> WorkflowExecution: + execution: WorkflowExecution = WorkflowExecution.objects.create( + workflow=self.workflow, + pipeline_id=pipeline_id, + execution_mode=WorkflowExecution.Mode.QUEUE, + execution_method=WorkflowExecution.Method.DIRECT, + execution_type=WorkflowExecution.Type.COMPLETE, + status=ExecutionStatus.EXECUTING.value, + queue_message_id=1, + ) + return execution + + def _post(self, *ids: Any) -> dict[str, tuple[str, str | None]]: + body = {"ids": [str(i) for i in ids]} + request = APIRequestFactory().post("/", body, format="json") + force_authenticate(request, user=self.owner) + response = _abort(request) + response.render() + self.assertEqual(response.status_code, 200, response.data) + return {r["id"]: (r["outcome"], r["status"]) for r in response.data["results"]} + + def _status(self, execution: WorkflowExecution) -> str: + execution.refresh_from_db() + status: str = execution.status + return status + + def test_abort_outside_any_transaction_succeeds(self) -> None: + execution = self._execution() + self.assertEqual( + self._post(execution.id), {str(execution.id): ("aborted", "STOPPED")} + ) + self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) + + def test_a_failure_part_way_keeps_the_others_and_rolls_back_its_own(self) -> None: + first, bad, last = self._execution(), self._execution(), self._execution() + + def stop(execution: WorkflowExecution) -> None: + if execution.pk == bad.pk: + raise RuntimeError("boom") + + with ( + patch(f"{_ABORT}._stop_open_files", side_effect=stop), + patch(f"{_ABORT}._report_abort") as report, + ): + outcomes = self._post(first.id, bad.id, last.id) + + self.assertEqual( + outcomes, + { + str(first.id): ("aborted", "STOPPED"), + str(bad.id): ("failed", None), + str(last.id): ("aborted", "STOPPED"), + }, + ) + self.assertEqual(self._status(first), ExecutionStatus.STOPPED.value) + self.assertEqual(self._status(bad), ExecutionStatus.EXECUTING.value) + self.assertEqual(self._status(last), ExecutionStatus.STOPPED.value) + self.assertEqual(PlatformAuditEvent.objects.count(), 2) + self.assertEqual(report.call_count, 2) + + def test_a_failing_report_keeps_the_abort_and_its_audit_event(self) -> None: + deployment = APIDeployment.objects.create( + api_name="api-1", + workflow=self.workflow, + organization=self.org, + created_by=self.owner, + ) + deployment.memberships.create(user=self.owner, role=ResourceRole.OWNER) + execution = self._execution(deployment.id) + + with ( + patch( + f"{_ABORT}.APIDeploymentUtils.send_notification", + side_effect=RuntimeError("boom"), + ) as send, + self.assertLogs(_ABORT, level="ERROR") as logs, + ): + outcomes = self._post(execution.id) + + send.assert_called_once() + self.assertIn(str(execution.id), logs.output[0]) + self.assertEqual(outcomes, {str(execution.id): ("aborted", "STOPPED")}) + self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) + self.assertEqual(PlatformAuditEvent.objects.count(), 1) + + def test_a_run_finishing_during_the_abort_is_not_overwritten(self) -> None: + execution = self._execution() + locked, release = threading.Event(), threading.Event() + concurrent: Any = locking.ConcurrentStatusWriteTests + + def finish_the_run() -> None: + with transaction.atomic(): + WorkflowExecution.objects.select_for_update().get(pk=execution.pk) + locked.set() + release.wait(locking.WAIT_SECONDS) + WorkflowExecution.objects.filter(pk=execution.pk).update( + status=ExecutionStatus.COMPLETED.value + ) + + def abort() -> dict[str, tuple[str, str | None]]: + UserContext.set_organization_identifier(self.org.organization_id) + return self._post(execution.id) + + holder, _ = concurrent._run(finish_the_run) + self.assertTrue(locked.wait(locking.WAIT_SECONDS)) + thread, box = concurrent._run(abort) + concurrent._wait_until_a_backend_waits_on_a_lock(thread, box) + release.set() + holder.join(locking.WAIT_SECONDS) + thread.join(locking.WAIT_SECONDS) + + self.assertFalse(thread.is_alive()) + self.assertEqual( + box.get("result"), {str(execution.id): ("already_finished", "COMPLETED")} + ) + self.assertEqual(self._status(execution), ExecutionStatus.COMPLETED.value) + self.assertFalse(PlatformAuditEvent.objects.exists()) diff --git a/backend/workflow_manager/execution/tests/test_execution_contract.py b/backend/workflow_manager/execution/tests/test_execution_contract.py new file mode 100644 index 0000000000..04bae87d54 --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_execution_contract.py @@ -0,0 +1,221 @@ +"""The stuck/abort additions leave the existing execution endpoints unchanged, +and their published contract generates (UN-4318).""" + +import uuid +from typing import Any + +from account_v2.models import Organization +from django.test import SimpleTestCase +from django.urls import Resolver404, resolve, reverse +from drf_spectacular.generators import SchemaGenerator +from pipeline_v2.models import Pipeline +from rest_framework.test import APIRequestFactory + +from workflow_manager.execution import urls +from workflow_manager.execution.tests.stuck_abort_base import StuckAbortTestBase +from workflow_manager.execution.views import ExecutionViewSet +from workflow_manager.internal_views import WorkflowExecutionInternalViewSet +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution + +_URLCONF = "workflow_manager.execution.urls" +_list = ExecutionViewSet.as_view({"get": "list"}) +_detail = ExecutionViewSet.as_view({"get": "retrieve"}) +_status = WorkflowExecutionInternalViewSet.as_view({"get": "execution_status"}) +_retrieve = WorkflowExecutionInternalViewSet.as_view({"get": "retrieve"}) + +# The default row, pinned: a new key here is a contract change. +_EXECUTION_KEYS = { + "aggregated_total_pages_processed", + "attempts", + "created_at", + "dispatched_at", + "error_message", + "execution_method", + "execution_mode", + "execution_time", + "failed_files", + "id", + "modified_at", + "pipeline_id", + "pipeline_name", + "queue_message_id", + "result_acknowledged", + "status", + "successful_files", + "tags", + "total_files", + "workflow", + "workflow_name", +} + + +class ExistingResponsesUnchangedTests(StuckAbortTestBase): + def test_list_and_detail_keep_their_keys(self) -> None: + execution = self.execution() + listed = self.call(_list, self.owner) + self.assertEqual(set(listed.data["results"][0]), _EXECUTION_KEYS) + detail = self.call(_detail, self.owner, pk=str(execution.id)) + self.assertEqual(set(detail.data), _EXECUTION_KEYS) + + def test_stuck_false_is_the_same_response_as_no_parameter(self) -> None: + self.execution() + self.execution(status=ExecutionStatus.COMPLETED) + # A zero time renders as seconds since creation, which would tick between calls. + WorkflowExecution.objects.update(execution_time=1) + plain = self.call(_list, self.owner) + for value in ("false", "0", ""): + with self.subTest(stuck=value): + filtered = self.call(_list, self.owner, data={"stuck": value}) + self.assertEqual(filtered.content, plain.content) + + def test_entity_filter_still_selects_each_entity(self) -> None: + runs = { + "API": self.execution(self.deployment()), + "ETL": self.execution(self.pipeline(Pipeline.PipelineType.ETL)), + "TASK": self.execution(self.pipeline(Pipeline.PipelineType.TASK)), + "WF": self.execution(), + } + for entity, run in runs.items(): + with self.subTest(entity=entity): + response = self.call( + _list, self.owner, data={"execution_entity": entity} + ) + ids = [row["id"] for row in response.data["results"]] + self.assertEqual(ids, [str(run.id)]) + + def test_stuck_rows_add_only_stuck_since(self) -> None: + self.execution() + for value in ("true", "1", "TRUE"): + with self.subTest(stuck=value): + listed = self.call(_list, self.owner, data={"stuck": value}) + self.assertEqual( + set(listed.data["results"][0]), _EXECUTION_KEYS | {"stuck_since"} + ) + + def test_unrecognised_stuck_value_is_ignored(self) -> None: + self.execution(status=ExecutionStatus.COMPLETED) + listed = self.call(_list, self.owner, data={"stuck": "garbage"}) + self.assertEqual(set(listed.data["results"][0]), _EXECUTION_KEYS) + + +class InternalStatusActionTests(StuckAbortTestBase): + def _get(self, execution_id: Any, organization_id: str | None) -> Any: + request = APIRequestFactory().get("/") + request.organization_id = organization_id + return _status(request, id=str(execution_id)) + + def test_returns_the_status(self) -> None: + execution = self.execution(status=ExecutionStatus.STOPPED) + response = self._get(execution.id, self.org.organization_id) + self.assertEqual(response.status_code, 200) + self.assertEqual( + response.data, {"id": str(execution.id), "status": "STOPPED"} + ) + + def test_other_org_unknown_and_malformed_ids_are_404(self) -> None: + execution = self.execution() + Organization.objects.create( + name="org-b", display_name="Org B", organization_id="org-b" + ) + cases = ( + (execution.id, "org-b"), + (execution.id, "some-other-org"), + (execution.id, None), + (uuid.uuid4(), self.org.organization_id), + ("not-a-uuid", self.org.organization_id), + ) + for execution_id, organization_id in cases: + with self.subTest(execution_id=execution_id, org=organization_id): + response = self._get(execution_id, organization_id) + self.assertEqual(response.status_code, 404) + + +class InternalExecutionPayloadTests(StuckAbortTestBase): + def test_carries_the_saved_file_counts(self) -> None: + # A stopped run's final log uses these counts. + execution = self.execution(status=ExecutionStatus.STOPPED) + WorkflowExecution.objects.filter(pk=execution.pk).update( + total_files=3, successful_files=3, failed_files=0 + ) + request = APIRequestFactory().get("/", {"file_execution": "false"}) + request.organization_id = self.org.organization_id + + response = _retrieve(request, id=str(execution.id)) + + self.assertEqual(response.status_code, 200, response.data) + saved = response.data["execution"] + self.assertEqual((saved["successful_files"], saved["failed_files"]), (3, 0)) + + +class WorkerStatusRouteTests(StuckAbortTestBase): + def test_worker_status_url_reaches_the_status_action(self) -> None: + match = resolve(f"/internal/v1/workflow-execution/{uuid.uuid4()}/status/") + self.assertEqual(match.func.actions, {"get": "execution_status"}) + self.assertIs(match.func.cls, WorkflowExecutionInternalViewSet) + + +class RoutesUnchangedTests(SimpleTestCase): + def test_existing_names_reverse_to_the_same_paths(self) -> None: + pk = uuid.uuid4() + expected = { + ("execution-list", None): "/", + ("execution-detail", pk): f"/{pk}/", + ("execution-log", pk): f"/{pk}/logs/", + ("execution-log-export", pk): f"/{pk}/logs/export/", + ("execution-stuck-count", None): "/stuck-count/", + ("execution-abort", None): "/abort/", + } + for (name, key), path in expected.items(): + kwargs = {"pk": key} if key else None + with self.subTest(name=name): + self.assertEqual(reverse(name, urlconf=_URLCONF, kwargs=kwargs), path) + + def test_new_paths_take_no_format_suffix(self) -> None: + for path in ("/abort.json", "/stuck-count.json"): + with self.subTest(path=path), self.assertRaises(Resolver404): + resolve(path, urlconf=_URLCONF) + self.assertEqual( + resolve(f"/{uuid.uuid4()}.json", urlconf=_URLCONF).url_name, + "execution-detail", + ) + + def test_new_paths_do_not_shadow_detail(self) -> None: + self.assertEqual( + resolve("/stuck-count/", urlconf=_URLCONF).url_name, + "execution-stuck-count", + ) + self.assertEqual( + resolve(f"/{uuid.uuid4()}/", urlconf=_URLCONF).url_name, + "execution-detail", + ) + + +class SchemaContractTests(SimpleTestCase): + """The ``@extend_schema`` annotations and README.md are the contract.""" + + def test_schema_documents_the_new_operations(self) -> None: + schema = SchemaGenerator(patterns=urls.urlpatterns).get_schema( + request=None, public=True + ) + paths = schema["paths"] + count = paths["/stuck-count/"]["get"] + self.assertEqual(count["operationId"], "execution_stuck_count") + self.assertTrue({"200", "401", "403"} <= set(count["responses"])) + + abort = paths["/abort/"]["post"] + self.assertEqual(abort["operationId"], "execution_abort") + self.assertTrue({"200", "400", "401", "403"} <= set(abort["responses"])) + components = schema["components"]["schemas"] + self.assertIn("ids", components["ExecutionAbortRequest"]["properties"]) + self.assertEqual( + set(components["ExecutionAbortResponse"]["properties"]), + {"results", "aborted", "already_finished", "not_permitted", "failed"}, + ) + self.assertEqual( + set(components["StuckExecutionCount"]["properties"]), + {"threshold_seconds", "total", "counts"}, + ) + + list_params = {p["name"] for p in paths["/"]["get"]["parameters"]} + self.assertIn("stuck", list_params) diff --git a/backend/workflow_manager/execution/tests/test_stuck_executions.py b/backend/workflow_manager/execution/tests/test_stuck_executions.py new file mode 100644 index 0000000000..b338f66367 --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_stuck_executions.py @@ -0,0 +1,147 @@ +"""Stuck executions: the ``stuck=true`` list filter and ``stuck-count/`` (UN-4318).""" + +from datetime import timedelta +from typing import Any + +from django.test import override_settings +from django.utils.dateparse import parse_datetime +from pipeline_v2.models import Pipeline + +from workflow_manager.execution.stuck import count_stuck +from workflow_manager.execution.tests.stuck_abort_base import ( + OLD, + RECENT, + StuckAbortTestBase, +) +from workflow_manager.execution.views import ExecutionViewSet, StuckExecutionCountView +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution + +_list = ExecutionViewSet.as_view({"get": "list"}) +_count = StuckExecutionCountView.as_view() + + +class StuckExecutionTests(StuckAbortTestBase): + def _stuck_rows(self, user: Any) -> dict[str, Any]: + response = self.call(_list, user, data={"stuck": "true"}) + self.assertEqual(response.status_code, 200) + return {row["id"]: row for row in response.data["results"]} + + def _count(self, user: Any) -> dict[str, Any]: + response = self.call(_count, user) + self.assertEqual(response.status_code, 200) + data: dict[str, Any] = response.data + return data + + def test_stuck_run_is_listed_with_stuck_since_and_counted(self) -> None: + execution = self.execution() + latest = self.file( + execution, ExecutionStatus.EXECUTING, idle=OLD - timedelta(minutes=10) + ) + self.file(execution, ExecutionStatus.COMPLETED, idle=OLD) + + rows = self._stuck_rows(self.owner) + + self.assertEqual(list(rows), [str(execution.id)]) + stuck_since = parse_datetime(rows[str(execution.id)]["stuck_since"]) + self.assertEqual(stuck_since, latest.modified_at) + self.assertEqual( + self._count(self.owner), + { + "threshold_seconds": 9000, + "total": 1, + "counts": {"API": 0, "ETL": 0, "TASK": 0, "WF": 1}, + }, + ) + + def test_run_without_files_is_stuck_since_its_own_last_write(self) -> None: + execution = self.execution(status=ExecutionStatus.PENDING) + rows = self._stuck_rows(self.owner) + self.assertEqual( + parse_datetime(rows[str(execution.id)]["stuck_since"]), + execution.modified_at, + ) + + def test_run_newer_than_its_files_is_stuck_since_its_own_last_write(self) -> None: + execution = self.execution(idle=OLD - timedelta(minutes=10)) + self.file(execution, ExecutionStatus.EXECUTING, idle=OLD) + rows = self._stuck_rows(self.owner) + self.assertEqual( + parse_datetime(rows[str(execution.id)]["stuck_since"]), + execution.modified_at, + ) + + @override_settings(STUCK_EXECUTION_THRESHOLD_SECONDS=3600) + def test_cutoff_is_the_threshold(self) -> None: + just_over = self.execution(idle=timedelta(minutes=61)) + self.execution(idle=timedelta(minutes=59)) + self.assertEqual(list(self._stuck_rows(self.owner)), [str(just_over.id)]) + + def test_recent_file_activity_is_not_stuck(self) -> None: + execution = self.execution() + self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + self.assertEqual(self._stuck_rows(self.owner), {}) + self.assertEqual(self._count(self.owner)["total"], 0) + + def test_recently_updated_run_is_not_stuck(self) -> None: + self.execution(idle=RECENT) + self.assertEqual(self._stuck_rows(self.owner), {}) + + def test_finished_runs_are_not_stuck(self) -> None: + for status in ( + ExecutionStatus.COMPLETED, + ExecutionStatus.ERROR, + ExecutionStatus.STOPPED, + ): + self.execution(status=status) + self.assertEqual(self._stuck_rows(self.owner), {}) + self.assertEqual(self._count(self.owner)["total"], 0) + + def test_runs_the_caller_cannot_see_are_excluded(self) -> None: + self.execution(self.deployment()) + self.assertEqual(self._stuck_rows(self.outsider), {}) + self.assertEqual(self._count(self.outsider)["total"], 0) + # Control: the owner does see it. + self.assertEqual(self._count(self.owner)["total"], 1) + + def test_counts_split_by_entity(self) -> None: + self.execution(self.deployment()) + self.execution(self.pipeline(Pipeline.PipelineType.ETL)) + self.execution(self.pipeline(Pipeline.PipelineType.ETL)) + self.execution(self.pipeline(Pipeline.PipelineType.TASK)) + self.execution() + + report = self._count(self.owner) + + self.assertEqual(report["total"], 5) + self.assertEqual(report["counts"], {"API": 1, "ETL": 2, "TASK": 1, "WF": 1}) + + def test_count_is_one_query(self) -> None: + self.execution(self.deployment()) + self.execution() + with self.assertNumQueries(1): + report = count_stuck(WorkflowExecution.objects.all()) + self.assertEqual(report.total, 2) + self.assertEqual(report.counts["API"], 1) + + def test_stuck_combines_with_the_entity_filter(self) -> None: + api_run = self.execution(self.deployment()) + self.execution() + response = self.call( + _list, self.owner, data={"stuck": "true", "execution_entity": "API"} + ) + ids = [row["id"] for row in response.data["results"]] + self.assertEqual(ids, [str(api_run.id)]) + + @override_settings(STUCK_EXECUTION_THRESHOLD_SECONDS=30) + def test_threshold_comes_from_settings(self) -> None: + execution = self.execution(idle=RECENT) + self.assertIn(str(execution.id), self._stuck_rows(self.owner)) + report = self._count(self.owner) + self.assertEqual(report["threshold_seconds"], 30) + self.assertEqual(report["total"], 1) + + def test_count_requires_authentication(self) -> None: + request = self.factory.get("/") + response = _count(request) + self.assertIn(response.status_code, (401, 403)) diff --git a/backend/workflow_manager/execution/urls.py b/backend/workflow_manager/execution/urls.py index c5bf5cc307..435c86ecc5 100644 --- a/backend/workflow_manager/execution/urls.py +++ b/backend/workflow_manager/execution/urls.py @@ -1,7 +1,11 @@ from django.urls import path from rest_framework.urlpatterns import format_suffix_patterns -from workflow_manager.execution.views import ExecutionViewSet +from workflow_manager.execution.views import ( + ExecutionAbortView, + ExecutionViewSet, + StuckExecutionCountView, +) from workflow_manager.workflow_v2.execution_log_view import ( WorkflowExecutionLogViewSet as ExecutionLogViewSet, ) @@ -26,4 +30,12 @@ name="execution-log-export", ), ] -) +) + [ + # Outside the suffix patterns: these views take no format argument. + path( + "stuck-count/", + StuckExecutionCountView.as_view(), + name="execution-stuck-count", + ), + path("abort/", ExecutionAbortView.as_view(), name="execution-abort"), +] diff --git a/backend/workflow_manager/execution/views/__init__.py b/backend/workflow_manager/execution/views/__init__.py index 90a34404ac..e59f7e7297 100644 --- a/backend/workflow_manager/execution/views/__init__.py +++ b/backend/workflow_manager/execution/views/__init__.py @@ -1 +1,2 @@ from .execution import ExecutionViewSet # noqa: F401 +from .stuck_abort import ExecutionAbortView, StuckExecutionCountView # noqa: F401 diff --git a/backend/workflow_manager/execution/views/execution.py b/backend/workflow_manager/execution/views/execution.py index a80a1e354f..369ceeb590 100644 --- a/backend/workflow_manager/execution/views/execution.py +++ b/backend/workflow_manager/execution/views/execution.py @@ -1,13 +1,20 @@ import logging from django_filters.rest_framework import DjangoFilterBackend +from drf_spectacular.types import OpenApiTypes +from drf_spectacular.utils import OpenApiParameter, extend_schema, extend_schema_view from rest_framework import viewsets from rest_framework.permissions import BasePermission, IsAuthenticated from utils.filters.ordering_filter import DeterministicOrderingFilter from utils.pagination import CustomPagination -from workflow_manager.execution.filter import ExecutionFilter +from workflow_manager.execution.filter import ( + STUCK_PARAM, + ExecutionFilter, + stuck_requested, +) from workflow_manager.execution.serializer import ExecutionSerializer +from workflow_manager.execution.serializer.stuck_abort import StuckExecutionSerializer from workflow_manager.workflow_v2.models import WorkflowExecution logger = logging.getLogger(__name__) @@ -30,6 +37,21 @@ def has_object_permission(self, request, view, obj): return WorkflowExecution.objects.for_user(request.user).filter(id=obj.id).exists() +@extend_schema_view( + list=extend_schema( + parameters=[ + OpenApiParameter( + STUCK_PARAM, + OpenApiTypes.BOOL, + description=( + "true: only PENDING/EXECUTING runs with no activity for longer " + "than the stuck threshold; each row then also carries " + "stuck_since. See execution/README.md." + ), + ) + ] + ) +) class ExecutionViewSet(viewsets.ReadOnlyModelViewSet): permission_classes = [IsAuthenticated, UserWorkflowExecutionPermission] serializer_class = ExecutionSerializer @@ -42,3 +64,10 @@ class ExecutionViewSet(viewsets.ReadOnlyModelViewSet): def get_queryset(self): # Use the custom manager method to filter executions for the current user return WorkflowExecution.objects.for_user(self.request.user) + + def get_serializer_class(self) -> type[ExecutionSerializer]: + # Only ?stuck=true rows carry the stuck_since annotation. + request = getattr(self, "request", None) + if request is not None and stuck_requested(request.query_params): + return StuckExecutionSerializer + return ExecutionSerializer diff --git a/backend/workflow_manager/execution/views/stuck_abort.py b/backend/workflow_manager/execution/views/stuck_abort.py new file mode 100644 index 0000000000..76ac21cb50 --- /dev/null +++ b/backend/workflow_manager/execution/views/stuck_abort.py @@ -0,0 +1,72 @@ +"""Stuck-execution count and abort. See ../README.md for the contract.""" + +from drf_spectacular.utils import OpenApiExample, OpenApiResponse, extend_schema +from drf_standardized_errors.openapi_serializers import ValidationErrorResponseSerializer +from rest_framework.permissions import IsAuthenticated +from rest_framework.request import Request +from rest_framework.response import Response +from rest_framework.views import APIView + +from workflow_manager.execution.abort import abort_executions +from workflow_manager.execution.serializer.stuck_abort import ( + ExecutionAbortRequestSerializer, + ExecutionAbortResponseSerializer, + StuckExecutionCountSerializer, +) +from workflow_manager.execution.stuck import count_stuck +from workflow_manager.workflow_v2.models import WorkflowExecution + +_TAGS = ["execution"] +# The auth middleware answers these before DRF runs, as {"message": ...}. +_AUTH_ERRORS = { + 401: OpenApiResponse(description='Not authenticated: {"message": ...}.'), + 403: OpenApiResponse( + description='Authenticated but refused access: {"message": ...}.' + ), +} + + +class StuckExecutionCountView(APIView): + permission_classes = [IsAuthenticated] + + @extend_schema( + operation_id="execution_stuck_count", + tags=_TAGS, + description="Stuck executions visible to the caller, per entity.", + responses={200: StuckExecutionCountSerializer, **_AUTH_ERRORS}, + examples=[ + OpenApiExample( + "Counts", + value={ + "threshold_seconds": 9000, + "total": 3, + "counts": {"API": 1, "ETL": 2, "TASK": 0, "WF": 0}, + }, + response_only=True, + ) + ], + ) + def get(self, request: Request) -> Response: + report = count_stuck(WorkflowExecution.objects.for_user(request.user)) + return Response(StuckExecutionCountSerializer(report).data) + + +class ExecutionAbortView(APIView): + permission_classes = [IsAuthenticated] + + @extend_schema( + operation_id="execution_abort", + tags=_TAGS, + description="Stop running executions; every id gets its own outcome in a 200.", + request=ExecutionAbortRequestSerializer, + responses={ + 200: ExecutionAbortResponseSerializer, + 400: ValidationErrorResponseSerializer, + **_AUTH_ERRORS, + }, + ) + def post(self, request: Request) -> Response: + body = ExecutionAbortRequestSerializer(data=request.data) + body.is_valid(raise_exception=True) + report = abort_executions(request, body.validated_data["ids"]) + return Response(ExecutionAbortResponseSerializer(report).data) diff --git a/backend/workflow_manager/file_execution/serializers.py b/backend/workflow_manager/file_execution/serializers.py index a40f80b04b..59b62aa4f0 100644 --- a/backend/workflow_manager/file_execution/serializers.py +++ b/backend/workflow_manager/file_execution/serializers.py @@ -34,6 +34,8 @@ def get_status_msg(self, obj: FileExecution) -> dict[str, any] | None: return self.INIT_STATUS_MSG elif obj.status == ExecutionStatus.ERROR.value: return obj.execution_error + elif obj.status == ExecutionStatus.STOPPED.value and obj.execution_error: + return obj.execution_error # Use pre-annotated field from viewset queryset to avoid N+1 queries if hasattr(obj, "latest_log_data") and obj.latest_log_data: diff --git a/backend/workflow_manager/internal_serializers.py b/backend/workflow_manager/internal_serializers.py index 9127763873..fbe9e54b53 100644 --- a/backend/workflow_manager/internal_serializers.py +++ b/backend/workflow_manager/internal_serializers.py @@ -96,6 +96,8 @@ class Meta: "status", "result_acknowledged", "total_files", + "successful_files", + "failed_files", "error_message", "attempts", "execution_time", diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index 879a1b01e5..1d26d41c45 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -10,6 +10,7 @@ from django.shortcuts import get_object_or_404 from rest_framework import status, viewsets from rest_framework.decorators import action +from rest_framework.request import Request from rest_framework.response import Response from rest_framework.views import APIView from tool_instance_v2.models import ToolInstance @@ -19,6 +20,7 @@ # Import new dataclasses for WorkflowDefinitionAPIView from unstract.core.data_models import ( + ALREADY_FINAL_REASON, ConnectionType, ConnectorInstanceData, WorkflowDefinitionResponseData, @@ -515,6 +517,27 @@ def _get_organization_context(self, execution: WorkflowExecution) -> dict: "settings": {}, } + # Named execution_status because `status` is the imported DRF module. + @action(detail=True, methods=["get"], url_path="status") + def execution_status(self, request: Request, id: str | None = None) -> Response: + """Lightweight status read; workers poll it between files.""" + try: + execution_id = uuid.UUID(str(id)) + except ValueError: + execution_id = None + executions = filter_queryset_by_organization( + WorkflowExecution.objects.filter(id=execution_id), + request, + "workflow__organization", + ) + execution_status = executions.values_list("status", flat=True).first() + if execution_status is None: + return Response( + {"error": "WorkflowExecution not found"}, + status=status.HTTP_404_NOT_FOUND, + ) + return Response({"id": str(execution_id), "status": execution_status}) + @action(detail=True, methods=["post"]) def update_status(self, request, id=None): """Update workflow execution status.""" @@ -548,70 +571,54 @@ def update_status(self, request, id=None): # with failed_files=None, which is_failure_run() reads as a success and # silently bypasses notify_on_failures subscribers. with transaction.atomic(): - row_locked = False - # Concurrency guard for PG-transport executions, identified by the - # durable ``queue_message_id`` marker — NOT a re-evaluated Flipt flag. - # Keying on the row's own transport keeps the invariant correct even - # if the rollout flag is toggled off or Flipt is unreachable while a - # PG execution is still finalizing (per the "transport carried on the - # row, never re-read from Flipt" rule). - # - # Protect COMPLETED and STOPPED — the two statuses that assert a - # deliberate final outcome. COMPLETED is set exclusively by a - # successful callback (its redelivery is skipped, see - # _callback_already_ran); STOPPED records an explicit user/operator - # stop. Under the PG path's high concurrency ~15 call sites write - # status via unlocked read-modify-write, so a stale writer can commit - # AFTER one of these and clobber it (lost update → execution stranded - # in EXECUTING, a confusing success→failed flip, or a straggler - # callback erasing a user's stop). Lock the row and refuse to overwrite - # either with a different status. + # Lock the row and refuse to overwrite a deliberate final outcome: + # COMPLETED (set only by a successful callback) and STOPPED (an + # explicit stop). Under high concurrency ~15 call sites write status + # via unlocked read-modify-write, so a stale writer can commit AFTER + # one of these and clobber it (lost update → execution stranded in + # EXECUTING, a confusing success→failed flip, or a straggler callback + # erasing a user's stop). See _protected_statuses for which rows + # protect which status. # # ERROR is deliberately NOT protected: other paths (upstream error, # the reaper) can set it BEFORE the first real callback, which then # legitimately corrects the row to COMPLETED when the files actually # succeeded — blocking that would freeze a genuinely successful run at - # a premature ERROR. Celery executions (queue_message_id NULL) keep the - # legacy path. + # a premature ERROR. # # ONE lock per call: the row read here is handed to update_execution() - # with row_locked=True, so it does not lock the same row a second time. The - # Celery path takes no lock here and lets update_execution() take its - # own — still one. FOR NO KEY UPDATE (no_key=True) keeps writers - # mutually exclusive without blocking the FOR KEY SHARE that every - # child-row INSERT takes on this row for its foreign-key check. - if execution.queue_message_id is not None: - execution = WorkflowExecution.objects.select_for_update( - no_key=True - ).get(pk=execution.pk) - row_locked = True - protected = { - ExecutionStatus.COMPLETED.value, - ExecutionStatus.STOPPED.value, - } - if ( - execution.status in protected - and status_enum.value != execution.status - ): - logger.warning( - f"update_status: refusing to overwrite protected status " - f"'{execution.status}' with '{status_enum.value}' " - f"for execution {id}" - ) - return Response( - { - "status": "skipped", - "reason": "already_final", - "execution_id": str(execution.id), - "new_status": execution.status, - } - ) + # with row_locked=True, so it does not lock the same row a second + # time. FOR NO KEY UPDATE (no_key=True) keeps writers mutually + # exclusive without blocking the FOR KEY SHARE that every child-row + # INSERT takes on this row for its foreign-key check. + execution = WorkflowExecution.objects.select_for_update(no_key=True).get( + pk=execution.pk + ) + if ( + execution.status in self._protected_statuses(execution) + and status_enum.value != execution.status + ): + logger.warning( + "update_status: refusing to overwrite protected status " + "'%s' with '%s' for execution %s", + execution.status, + status_enum.value, + id, + ) + return Response( + { + "status": "skipped", + "reason": ALREADY_FINAL_REASON, + "execution_id": str(execution.id), + "new_status": execution.status, + } + ) execution.update_execution( status=status_enum, error=error_message, increment_attempt=increment_attempt, - row_locked=row_locked, + row_locked=True, ) # File aggregates are not handled by update_execution; persist @@ -623,7 +630,13 @@ def update_status(self, request, id=None): # either way: Django's foreign keys are DEFERRABLE INITIALLY DEFERRED, # so an open insert holds no lock on this row until it commits, and # FOR UPDATE never waited for it.) - self._update_file_aggregates(execution, validated_data) + if status_enum == ExecutionStatus.STOPPED: + from workflow_manager.execution.abort import write_file_counts + + # A stopped run's counts come from its file rows, not the worker. + write_file_counts(execution) + else: + self._update_file_aggregates(execution, validated_data) # Reaper recovery: cascade the terminal status to any non-terminal # file executions so the execution and its files agree (atomic with @@ -872,6 +885,20 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str: logger.exception("recover_stuck_pg_executions: failed to recover %s", exec_id) return "failed" + @staticmethod + def _protected_statuses(execution: WorkflowExecution) -> set[str]: + """Statuses a different status write must not overwrite on this row. + + Any run can be stopped, so STOPPED is protected on every row. COMPLETED + stays protected only on rows carrying the PG dispatch marker + ``queue_message_id``, keyed on the row and never on a re-read Flipt flag. + """ + from workflow_manager.workflow_v2.enums import ExecutionStatus + + if execution.queue_message_id is None: + return {ExecutionStatus.STOPPED.value} + return {ExecutionStatus.COMPLETED.value, ExecutionStatus.STOPPED.value} + @staticmethod def _cascade_terminal_files(execution, status_enum, error_message) -> None: """Mark this execution's non-terminal (PENDING/EXECUTING) file executions to diff --git a/backend/workflow_manager/workflow_v2/execution.py b/backend/workflow_manager/workflow_v2/execution.py index dd177e50d3..22731c5f46 100644 --- a/backend/workflow_manager/workflow_v2/execution.py +++ b/backend/workflow_manager/workflow_v2/execution.py @@ -175,8 +175,8 @@ def update_execution( increment_attempt: bool = False, ) -> None: # Delegate to the model method, which owns the single, robust terminal-one-way - # guard (atomic select_for_update, PG-scoped, field-scoped writes). Duplicating - # the guard here risked drift + its own TOCTOU/dropped-field bugs. + # guard (atomic select_for_update, PG and STOPPED rows, field-scoped writes). + # Duplicating the guard here risked drift + its own TOCTOU/dropped-field bugs. execution = WorkflowExecution.objects.get(pk=self.execution_id) execution.update_execution( status=status, error=error, increment_attempt=increment_attempt @@ -405,23 +405,23 @@ def update_execution_completed( execution = WorkflowExecution.objects.get(pk=execution_id) # One transaction, so on-commit hooks of the status write see the counts. with transaction.atomic(): - # Routed through the model method so the PG terminal-one-way guard - # applies; the legacy transport has no such guard. + # Routed through the model method so its guard applies. execution.update_execution(status=ExecutionStatus.COMPLETED) - execution.total_files = total_files - execution.successful_files = 0 - execution.failed_files = failed_files - # Field-scoped for the columns, but save() still republishes this - # object's status to the execution cache. Safe only because the row - # is undispatched here, so no worker can have moved it on. - execution.save( - update_fields=[ - "total_files", - "successful_files", - "failed_files", - "modified_at", - ] - ) + # A refused write (e.g. the run was aborted) keeps the stored counts. + if execution.status == ExecutionStatus.COMPLETED.value: + execution.total_files = total_files + execution.successful_files = 0 + execution.failed_files = failed_files + # Field-scoped, but save() still republishes this object's status + # to the cache; safe under the lock this block's write holds. + execution.save( + update_fields=[ + "total_files", + "successful_files", + "failed_files", + "modified_at", + ] + ) # The guard may have refused the status change without raising; re-read # so the returned row reflects what is actually stored. execution.refresh_from_db() diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index 536fda11f2..da1d2db49f 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -435,7 +435,7 @@ def update_execution( # self.queue_message_id — unless the caller passes row_locked=True, i.e. # this instance was itself read under the lock (see the docstring): a PG row snapshotted before dispatch recorded its # handle, mis-routed to an unlocked legacy full save(), would write NULL - # back into queue_message_id and permanently disable every guard. Under + # back into queue_message_id and permanently drop the COMPLETED guard. Under # this lock a concurrent _record_dispatch_handle serializes after us. # # no_key=True takes FOR NO KEY UPDATE rather than FOR UPDATE. Status @@ -460,7 +460,11 @@ def update_execution( ) if locked is None: return - if locked["queue_message_id"] is None: + # A stopped row takes the guarded write too, so nothing replaces the stop. + if ( + locked["queue_message_id"] is None + and locked["status"] != ExecutionStatus.STOPPED.value + ): should_release_rate_limit = self._apply_legacy_update( status, error, increment_attempt ) @@ -492,8 +496,9 @@ def _apply_legacy_update( error: str | None, increment_attempt: bool, ) -> bool: - """Celery-path (queue_message_id NULL) — unchanged full save(); runs inside the - caller's locked transaction. Returns whether to release the rate-limit slot. + """Celery-path (queue_message_id NULL, not STOPPED) — unchanged full save(); + runs inside the caller's locked transaction. Returns whether to release the + rate-limit slot. """ should_release = False if status is not None: @@ -520,12 +525,13 @@ def _apply_pg_guarded_update( error: str | None, increment_attempt: bool, ) -> bool: - """PG-path terminal-one-way guarded write, inside the caller's row lock. Refuses - a revert of a protected-terminal (COMPLETED/STOPPED) status — ERROR stays - correctable — while still applying error / increment_attempt. Field-scoped - (``update_fields``) so a stale object can never clobber counters or the marker, - and the post-save cache publishes the status actually persisted. Returns whether - to release the rate-limit slot. + """Terminal-one-way guarded write for PG and STOPPED rows, inside the caller's + row lock. Refuses a revert of a protected-terminal (COMPLETED/STOPPED) status + — ERROR stays correctable — while still applying error / increment_attempt + (a STOPPED row keeps its error text). + Field-scoped (``update_fields``) so a stale object can never clobber counters or + the marker, and the post-save cache publishes the status actually persisted. + Returns whether to release the rate-limit slot. """ committed = locked["status"] should_release = False @@ -535,7 +541,8 @@ def _apply_pg_guarded_update( ExecutionStatus(status), committed, error, increment_attempt ) update_fields.extend(status_fields) - if error: + # A stopped run keeps the abort's message. + if error and committed != ExecutionStatus.STOPPED.value: self.error_message = error[:EXECUTION_ERROR_LENGTH] update_fields.append("error_message") if increment_attempt: diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index e6ceb8d6c8..2b323fd458 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/backend/workflow_manager/workflow_v2/workflow_helper.py @@ -783,7 +783,9 @@ def execute_workflow_async( return ExecutionResponse( workflow_id, execution_id, ExecutionStatus.EXECUTING.value ) - WorkflowExecutionServiceHelper.update_execution_err(execution_id, str(error)) + marked = WorkflowExecutionServiceHelper.update_execution_err( + execution_id, str(error) + ) logger.exception( f"Error while enqueuing async job for WF '{workflow_id}', " f"execution '{execution_id}'" @@ -791,7 +793,8 @@ def execute_workflow_async( return ExecutionResponse( workflow_id, execution_id, - ExecutionStatus.ERROR.value, + # An aborted run stays STOPPED. + marked.status if marked else ExecutionStatus.ERROR.value, error=str(error), ) diff --git a/specs/docstudio-oss.json b/specs/docstudio-oss.json index 08b71f8e15..534ede8228 100644 --- a/specs/docstudio-oss.json +++ b/specs/docstudio-oss.json @@ -639,7 +639,7 @@ }, "/deployment/api/{org_name}/{api_name}/": { "get": { - "description": "Reads a previously started execution, taking the same deployment key that ran it. The read is one-shot: the first call that observes a completed execution acknowledges it and the stored result is discarded, so every later call for that execution answers 406. Keep the payload of the call that returns it \u2014 it cannot be fetched again.\n\nA still-running execution answers 422 carrying its current `status`, so a polling loop should treat 422 as the normal reply and stop on 200.", + "description": "Reads a previously started execution, taking the same deployment key that ran it. The read is one-shot: the first call that observes a completed execution acknowledges it and the stored result is discarded, so every later call for that execution answers 406. Keep the payload of the call that returns it \u2014 it cannot be fetched again.\n\nA still-running execution answers 422 carrying its current `status`, so a polling loop should treat 422 as the normal reply and stop on 200, or when `status` is `ERROR` or `STOPPED`. Those also answer 422 but the run has ended, and that read is the only one that carries its results.", "operationId": "status", "parameters": [ { @@ -833,7 +833,7 @@ } } }, - "description": "The execution is still running, or it finished with an error; read `status` to tell them apart." + "description": "The execution is still running, finished with an error, or was aborted (`STOPPED`); read `status` to tell them apart." }, "500": { "content": { @@ -1043,7 +1043,7 @@ } } }, - "description": "The execution finished with an error." + "description": "The execution finished with an error, or was aborted (`STOPPED`)." }, "429": { "content": { diff --git a/unstract/core/src/unstract/core/data_models.py b/unstract/core/src/unstract/core/data_models.py index 2f4f58a54d..70424621e2 100644 --- a/unstract/core/src/unstract/core/data_models.py +++ b/unstract/core/src/unstract/core/data_models.py @@ -582,6 +582,9 @@ def _protected_terminal_values(cls) -> frozenset[str]: ExecutionStatus.protected_terminal_values = classmethod(_protected_terminal_values) +# Reason the backend gives when it refuses a status write over a protected status. +ALREADY_FINAL_REASON = "already_final" + # Add the is_completed method as a class method def _is_completed(cls, status: str) -> bool: diff --git a/workers/callback/tasks.py b/workers/callback/tasks.py index cceb0aef04..79029faaf4 100644 --- a/workers/callback/tasks.py +++ b/workers/callback/tasks.py @@ -12,6 +12,7 @@ # Import shared worker infrastructure from shared.api import InternalAPIClient +from shared.clients.execution_client import refused_as_stopped # Import from shared worker modules from shared.enums import PipelineType @@ -56,7 +57,7 @@ class CallbackContext: """Container for callback processing context data.""" - def __init__(self): + def __init__(self) -> None: self.task_id: str = "" self.execution_id: str = "" self.pipeline_id: str | None = None @@ -255,6 +256,13 @@ def _get_performance_stats() -> dict: return stats +def _was_stopped(execution_context: dict[str, Any] | None) -> bool: + """True when the run was already STOPPED (aborted) when this callback fetched it.""" + execution = (execution_context or {}).get("execution") or {} + stopped: bool = execution.get("status") == ExecutionStatus.STOPPED.value + return stopped + + def _determine_execution_status_unified( file_batch_results: list[dict[str, Any]], api_client: InternalAPIClient, @@ -324,6 +332,16 @@ def _determine_execution_status_unified( total_files = aggregated_results.get("total_files", 0) successful_files = aggregated_results.get("successful_files", 0) + # An abort ended this run, whatever its file counts say. + if _was_stopped(execution_context): + logger.warning( + "Execution %s was stopped - %s of %s files completed", + execution_id, + successful_files, + total_files, + ) + return aggregated_results, ExecutionStatus.STOPPED.value, expected_files + # Step 5: Unified status determination with timeout failure detection # Detect timeout failures: expected files but processed none (likely SoftTimeLimitExceeded) @@ -395,16 +413,29 @@ def _update_execution_status_unified( successful_files = aggregated_results.get("successful_files", 0) failed_files = aggregated_results.get("failed_files", 0) + def post_status(status: str, error: str | None) -> Any: + return api_client.update_workflow_execution_status( + execution_id=execution_id, + status=status, + total_files=total_files, + successful_files=successful_files, + failed_files=failed_files, + organization_id=organization_id, + error_message=error, + ) + # Make the unified API call - api_client.update_workflow_execution_status( - execution_id=execution_id, - status=final_status, - total_files=total_files, - successful_files=successful_files, - failed_files=failed_files, - organization_id=organization_id, - error_message=error_message, - ) + response = post_status(final_status, error_message) + if refused_as_stopped(response): + logger.warning( + "Execution %s was stopped during this callback; " + "reporting STOPPED instead of %s", + execution_id, + final_status, + ) + final_status = ExecutionStatus.STOPPED.value + # The refused write saved no file counts. + post_status(final_status, None) # Return consistent result structure update_result = { @@ -478,6 +509,15 @@ def _handle_pipeline_updates_unified( "pipeline_id": context.pipeline_id, } + # The abort wrote it, and only if this is the pipeline's newest run. + if final_status == ExecutionStatus.STOPPED.value: + return { + "status": "skipped", + "reason": "handled_by_abort", + "message": "The abort handles a stopped run's pipeline status", + "pipeline_id": context.pipeline_id, + } + # Non-API workflows (ETL/TASK/APP) need pipeline status updates if not context.pipeline_id: logger.info("No pipeline_id provided - skipping pipeline status update") @@ -605,6 +645,17 @@ def _handle_notifications_unified( Returns: Notification result dictionary """ + # The abort handles these, so a run with a dead worker notifies too. + if status == ExecutionStatus.STOPPED.value: + logger.info( + "Execution %s was stopped - the abort handles notifying", execution_id + ) + return { + "status": "skipped", + "reason": "handled_by_abort", + "message": "The abort handles a stopped run's notifications", + } + try: if not pipeline_id: logger.warning("No pipeline_id provided - skipping notifications") @@ -1305,11 +1356,12 @@ def _track_subscription_usage_if_available( """Track subscription usage for successful file executions if plugin available. This is a non-blocking operation - errors are logged but do not fail the workflow. - Only tracks usage for completed workflows with successful file executions. + Only tracks usage for completed or stopped workflows with successful file + executions; files that completed before an abort are billed. Args: context: Callback execution context with organization and execution IDs - execution_status: Final execution status (should be COMPLETED for tracking) + execution_status: Final execution status (COMPLETED or STOPPED for tracking) successful_files: Successful file count from the aggregated batch results, used only to flag a fetch that came back without file-execution rows @@ -1317,8 +1369,11 @@ def _track_subscription_usage_if_available( dict with status, committed_count, message if tracking attempted None if plugin not available or workflow not completed """ - # Only track for completed workflows - if execution_status != ExecutionStatus.COMPLETED.value: + # Only track for workflows whose finished files stand + if execution_status not in ( + ExecutionStatus.COMPLETED.value, + ExecutionStatus.STOPPED.value, + ): logger.debug( f"Skipping subscription tracking for execution {context.execution_id} " f"with status {execution_status}" @@ -1509,6 +1564,7 @@ def _process_batch_callback_core( organization_id=context.organization_id, error_message=None, ) + execution_status = execution_update_result["final_status"] # Handle pipeline updates using unified function (non-API deployment) pipeline_result = _handle_pipeline_updates_unified( context, execution_status, is_api_deployment=False @@ -1775,6 +1831,7 @@ def process_batch_callback_api( aggregated_results=aggregated_results, organization_id=organization_id, ) + execution_status = execution_update_result["final_status"] # Create minimal context for unified pipeline handling context = CallbackContext() @@ -1898,6 +1955,26 @@ def process_batch_callback_api( logger.warning("api_client.close() failed during cleanup: %s", e) +def _final_log_counts( + aggregated_results: dict[str, Any], + execution_data: dict[str, Any] | None, + execution_status: str, +) -> tuple[int, int, int]: + """Total, successful and failed files for the final log line. + + A stopped run logs the counts the backend saved from its file rows. + """ + counts = aggregated_results + if execution_status == ExecutionStatus.STOPPED.value: + saved = (execution_data or {}).get("execution") or {} + counts = {**aggregated_results, **saved} + return ( + counts.get("total_files") or 0, + counts.get("successful_files") or 0, + counts.get("failed_files") or 0, + ) + + def _publish_final_workflow_ui_logs( context: "CallbackContext", aggregated_results: dict[str, Any], @@ -1914,11 +1991,6 @@ def _publish_final_workflow_ui_logs( execution_status: Final execution status """ try: - # Extract file statistics from aggregated results - total_files = aggregated_results.get("total_files", 0) - successful_files = aggregated_results.get("successful_files", 0) - failed_files = aggregated_results.get("failed_files", 0) - # Get execution data to extract cost information (with cost calculation) execution_response = context.api_client.get_workflow_execution( context.execution_id, include_cost=True, file_execution=False @@ -1931,6 +2003,9 @@ def _publish_final_workflow_ui_logs( # Cost data is at the top level of response when include_cost=True aggregated_usage_cost = execution_response.data.get("aggregated_usage_cost") + total_files, successful_files, failed_files = _final_log_counts( + aggregated_results, execution_response.data, execution_status + ) # Create workflow logger for UI feedback # Use general workflow logger since this is called from general workflow callback @@ -1954,6 +2029,7 @@ def _publish_final_workflow_ui_logs( total_files=total_files, successful_files=successful_files, failed_files=failed_files, + stopped=execution_status == ExecutionStatus.STOPPED.value, ) logger.info( @@ -1988,11 +2064,6 @@ def _publish_final_workflow_ui_logs_api( execution_status: Final execution status """ try: - # Extract file statistics from aggregated results - total_files = aggregated_results.get("total_files", 0) - successful_files = aggregated_results.get("successful_files", 0) - failed_files = aggregated_results.get("failed_files", 0) - # Get execution data to extract cost information (with cost calculation) execution_response = context.api_client.get_workflow_execution( context.execution_id, include_cost=True, file_execution=False @@ -2005,6 +2076,9 @@ def _publish_final_workflow_ui_logs_api( # Cost data is at the top level of response when include_cost=True aggregated_usage_cost = execution_response.data.get("aggregated_usage_cost") + total_files, successful_files, failed_files = _final_log_counts( + aggregated_results, execution_response.data, execution_status + ) # Create workflow logger for UI feedback # Use API workflow logger since this is called from API workflow callback @@ -2028,6 +2102,7 @@ def _publish_final_workflow_ui_logs_api( total_files=total_files, successful_files=successful_files, failed_files=failed_files, + stopped=execution_status == ExecutionStatus.STOPPED.value, ) logger.info( diff --git a/workers/file_processing/tasks.py b/workers/file_processing/tasks.py index 764953834b..c808805d42 100644 --- a/workers/file_processing/tasks.py +++ b/workers/file_processing/tasks.py @@ -19,6 +19,8 @@ # Import shared worker infrastructure from shared.api import InternalAPIClient +from shared.clients.base_client import InternalAPIClientError +from shared.clients.execution_client import refused_as_stopped # Import from shared worker modules from shared.constants import Account @@ -289,7 +291,9 @@ def _raise_if_execution_terminal( ) -def _terminal_skip_result(batch_data: FileBatchData) -> dict[str, Any]: +def _terminal_skip_result( + batch_data: FileBatchData, status: str | None = None +) -> dict[str, Any]: """Batch result for a batch skipped because its execution is already terminal. The files were never attempted; they are counted as failed (not left @@ -297,21 +301,86 @@ def _terminal_skip_result(batch_data: FileBatchData) -> dict[str, Any]: so reporting ``failed=total`` keeps a consumer that ever aggregates this result from rendering the files as perpetually in-progress. The ``SKIPPED_TERMINAL_EXECUTION_KEY`` marker tells ``run_batch_with_barrier`` to - bypass the barrier decrement (a terminal execution's barrier is by definition - already torn down). + bypass the barrier decrement. A STOPPED run gets no marker: the batch counts + down and the last one fires the callback, which cleans up. """ - n_files = len(batch_data.files) - result = BatchExecutionResult( + result = _all_failed_result( + len(batch_data.files), batch_data.file_data.organization_id + ) + if status != ExecutionStatus.STOPPED.value: + result[SKIPPED_TERMINAL_EXECUTION_KEY] = True + return result + + +def _all_failed_result(n_files: int, organization_id: str | None) -> dict[str, Any]: + """Batch result that counts every file as failed.""" + result: dict[str, Any] = BatchExecutionResult( total_files=n_files, successful_files=0, failed_files=n_files, execution_time=0.0, - organization_id=batch_data.file_data.organization_id, + organization_id=organization_id, ).to_dict() - result[SKIPPED_TERMINAL_EXECUTION_KEY] = True return result +def _run_stages_under_barrier( + file_batch_data: dict[str, Any], celery_task_id: str, batch_index: int +) -> dict[str, Any]: + """Run the batch stages; a stopped run's failed batch still counts down. + + Otherwise the barrier is torn down and the run's callback never fires. + """ + try: + return _run_batch_stages(file_batch_data, celery_task_id) + except Exception: + file_data = file_batch_data.get("file_data") or {} + execution_id = file_data.get("execution_id") or "" + organization_id = file_data.get("organization_id") or "" + api_client = create_api_client(organization_id) + if not _stopped_before_error_mark( + api_client, execution_id, organization_id, batch_index + ): + raise + logger.exception( + "[exec:%s] Batch failed after the run was stopped; counting it down", + execution_id, + ) + return _all_failed_result( + len(file_batch_data.get("files") or []), organization_id + ) + + +def _stopped_before_error_mark( + api_client: InternalAPIClient, + execution_id: str, + organization_id: str | None, + batch_index: int, +) -> bool: + """Mark the run ERROR; True when the backend refuses because it is STOPPED. + + The backend orders this write and an abort under the run's row lock, so an + abort landing just before the batch failed is not missed. The write matches + the barrier's own ERROR mark, so the two agree. + """ + try: + response = api_client.update_workflow_execution_status( + execution_id=execution_id, + status=ExecutionStatus.ERROR.value, + error_message=f"[pg-barrier-abort] batch {batch_index} failed.", + organization_id=organization_id, + cascade_terminal_files=True, + ) + except InternalAPIClientError as error: + logger.warning( + "[exec:%s] Marking the failed batch's run ERROR failed: %s", + execution_id, + error, + ) + return False + return refused_as_stopped(response) + + def _run_batch_stages( file_batch_data: dict[str, Any], celery_task_id: str ) -> dict[str, Any]: @@ -324,19 +393,20 @@ def _run_batch_stages( # Step 1: Validate and parse input data batch_data = _validate_and_parse_batch_data(file_batch_data) - # Step 2: Setup execution context. The terminal guard inside it may skip a - # stale/redelivered batch for an already-terminal execution — see + # Step 2: Setup execution context. The terminal guard inside it skips a + # batch whose execution is already terminal — see # :class:`_TerminalExecutionSkip`. try: context = _setup_execution_context(batch_data, celery_task_id) except _TerminalExecutionSkip as skip: logger.warning( - f"[exec:{skip.execution_id}] Skipping batch of " - f"{len(batch_data.files)} file(s): execution already terminal " - f"({skip.status}) — stale/redelivered after reaper-recovery; not " - f"reprocessing." + "[exec:%s] Skipping batch of %s file(s): execution already terminal " + "(%s); not processing.", + skip.execution_id, + len(batch_data.files), + skip.status, ) - return _terminal_skip_result(batch_data) + return _terminal_skip_result(batch_data, skip.status) # Step 3: Handle manual review logic # context = _handle_manual_review_logic(context) @@ -398,7 +468,9 @@ def _process_file_batch_core( # the barrier in-body / self-chain the callback. return run_batch_with_barrier( barrier_context, - lambda: _run_batch_stages(file_batch_data, celery_task_id), + lambda: _run_stages_under_barrier( + file_batch_data, celery_task_id, int(barrier_context["batch_index"]) + ), ) @@ -482,8 +554,7 @@ def _setup_execution_context( Raises: ValueError: If required context fields are missing - _TerminalExecutionSkip: If the execution is already terminal (a - stale/redelivered batch after reaper-recovery). + _TerminalExecutionSkip: If the execution is already terminal. """ # Extract context using dataclass file_data = batch_data.file_data @@ -523,9 +594,9 @@ def _setup_execution_context( execution_context = execution_response.data workflow_execution = execution_context.get("execution", {}) - # Validate-first terminal guard — skip a stale/redelivered batch - # for an already-terminal execution, before we touch status or process - # anything. See :class:`_TerminalExecutionSkip`. + # Validate-first terminal guard — skip a batch whose execution is already + # terminal, before we touch status or process anything. See + # :class:`_TerminalExecutionSkip`. _raise_if_execution_terminal(workflow_execution, execution_id) # Set LOG_EVENTS_ID in StateStore for WebSocket messaging (critical for UI logs) @@ -675,6 +746,97 @@ def _refactored_pre_create_file_executions( return context +def _file_item_name(file_item: Any) -> str | None: + """File name from a batch file item in any of its formats.""" + name: str | None = None + if isinstance(file_item, (list, tuple)) and len(file_item) == 2: + name = file_item[0] + elif isinstance(file_item, dict): + name = file_item.get("file_name") + return name + + +def _file_item_hash(file_item: Any) -> dict[str, Any]: + """Hash data from a batch file item in any of its formats; empty if none.""" + data = file_item + if isinstance(file_item, (list, tuple)) and len(file_item) == 2: + data = file_item[1] + return data if isinstance(data, dict) else {} + + +def _file_item_identifier(file_item: Any) -> str | None: + """The ``uuid:path`` id the duplicate skip lists use; None without both parts.""" + data = _file_item_hash(file_item) + uuid, path = data.get("provider_file_uuid"), data.get("file_path") + return f"{uuid}:{path}" if uuid and path else None + + +def _is_execution_stopped(api_client: InternalAPIClient, execution_id: str) -> bool: + """True when the execution was stopped (aborted) while this batch runs. + + Fails open on an API failure: the next file is processed. + """ + try: + status = api_client.get_workflow_execution_status(execution_id) + except InternalAPIClientError as error: + logger.warning( + "[exec:%s] status check before the next file failed (%s); " + "processing it anyway", + execution_id, + error, + ) + return False + if status not in {s.value for s in ExecutionStatus}: + logger.warning( + "[exec:%s] status check before the next file returned %r; " + "processing it anyway", + execution_id, + status, + ) + stopped: bool = status == ExecutionStatus.STOPPED.value + return stopped + + +def _mark_file_stopped(api_client: InternalAPIClient, file_execution_id: str) -> None: + """Mark one unprocessed file row STOPPED, logging a failed write.""" + # An APIResponse, whatever the facade's annotation says. + response: Any = api_client.update_file_execution_status( + file_execution_id=file_execution_id, + status=ExecutionStatus.STOPPED.value, + error_message="Not processed: execution was stopped", + ) + if not response.success: + logger.warning( + "Failed to mark file execution %s STOPPED: %s", + file_execution_id, + response.error, + ) + + +def _stop_remaining_files( + context: WorkflowContextData, remaining: list[Any], result: FileBatchResult +) -> None: + """Count the batch's unprocessed files as failed and mark their rows STOPPED. + + Files already completed in this execution stay uncounted, as in the main loop. + """ + api_client = context.organization_context.api_client + already_completed = (context.metadata or {}).get("skipped_already_completed", []) + for file_item in remaining: + if _file_item_identifier(file_item) in already_completed: + continue + result.increment_failure() + name = _file_item_name(file_item) + pre_created = context.pre_created_file_executions.get(name) if name else None + if pre_created: + _mark_file_stopped(api_client, pre_created.id) + logger.warning( + "[exec:%s] Execution was stopped; skipped the batch's %s remaining file(s)", + context.execution_id, + len(remaining), + ) + + def _process_individual_files(context: WorkflowContextData) -> WorkflowContextData: """Process each file individually through the workflow. @@ -703,6 +865,11 @@ def _process_individual_files(context: WorkflowContextData) -> WorkflowContextDa # Process each file - handle list, tuple, and dictionary formats for file_number, file_item in enumerate(files, 1): + # Stop between files once the run is aborted. + if _is_execution_stopped(api_client, context.execution_id): + _stop_remaining_files(context, files[file_number - 1 :], result) + break + # Handle Django list format (from asdict serialization), tuple format, and dictionary format if isinstance(file_item, list): # Django backend format after asdict(): [file_name, file_hash_dict] diff --git a/workers/general/tasks.py b/workers/general/tasks.py index d1ffe422f0..69637ee4fb 100644 --- a/workers/general/tasks.py +++ b/workers/general/tasks.py @@ -17,6 +17,7 @@ # Import shared worker infrastructure using new structure from shared.api import InternalAPIClient +from shared.clients.execution_client import refused_as_stopped # Import worker-specific data models from shared.data.models import ( @@ -304,12 +305,14 @@ def async_execute_bin_general( # to test) a redelivery could reach here, win a fresh claim, and re-arm + # re-dispatch a finished execution — duplicate (costly) destination # writes when use_file_history=False. Skip if already terminal, keeping - # the freshly-won claim so it re-establishes the tombstone. + # the freshly-won claim so it re-establishes the tombstone. A run aborted + # before pickup is skipped here too. if current_status in ExecutionStatus.terminal_values(): logger.warning( - f"Execution {execution_id} already terminal ({current_status}) " - f"at orchestration entry — skipping re-orchestration (a " - f"redelivery of a finished execution)." + "Execution %s already terminal (%s) at orchestration entry — " + "skipping (a redelivery, or a run aborted before pickup).", + execution_id, + current_status, ) return { "status": "skipped_terminal_execution", @@ -424,14 +427,15 @@ def async_execute_bin_general( try: with InternalAPIClient(config) as api_client: api_client.set_organization_context(schema_name) - api_client.update_workflow_execution_status( + reply = api_client.update_workflow_execution_status( execution_id=execution_id, status=ExecutionStatus.ERROR.value, error_message=str(e), ) - # CRITICAL FIX: Also update pipeline status to FAILED for consistency - if pipeline_id: + # CRITICAL FIX: Also update pipeline status to FAILED for consistency. + # An aborted run's pipeline status and notifications are the abort's. + if pipeline_id and not refused_as_stopped(reply): try: api_client.update_pipeline_status( pipeline_id=pipeline_id, @@ -776,30 +780,11 @@ def _execute_general_workflow( "💤 No files to process - all files have been filtered out (already processed, active, or no matches)", ) - # Send completion log to UI - workflow_logger.publish_execution_complete( - successful_files=0, failed_files=0, total_time=0.0 - ) - - # Complete immediately with no files - try: - api_client.update_workflow_execution_status( - execution_id=execution_id, - status=ExecutionStatus.COMPLETED.value, - total_files=0, + if _complete_without_files(api_client, execution_id, pipeline_id): + # Send completion log to UI + workflow_logger.publish_execution_complete( + successful_files=0, failed_files=0, total_time=0.0 ) - except Exception as status_error: - logger.warning(f"Failed to update execution status: {status_error}") - - # Update pipeline status if needed - if pipeline_id: - try: - api_client.update_pipeline_status( - pipeline_id=pipeline_id, - status=PipelineStatus.SUCCESS.value, - ) - except Exception as pipeline_error: - logger.warning(f"Failed to update pipeline status: {pipeline_error}") execution_time = time.time() - start_time response = WorkerTaskResponse.success_response( @@ -881,6 +866,34 @@ def _execute_general_workflow( raise +def _complete_without_files( + api_client: InternalAPIClient, execution_id: str, pipeline_id: str | None +) -> bool: + """Mark a run with no files COMPLETED and its pipeline SUCCESS. + + Returns False for a run aborted meanwhile: it stays STOPPED, and the abort + sets its pipeline status. + """ + try: + reply = api_client.update_workflow_execution_status( + execution_id=execution_id, + status=ExecutionStatus.COMPLETED.value, + total_files=0, + ) + if refused_as_stopped(reply): + return False + except Exception as status_error: + logger.warning("Failed to update execution status: %s", status_error) + if pipeline_id: + try: + api_client.update_pipeline_status( + pipeline_id=pipeline_id, status=PipelineStatus.SUCCESS.value + ) + except Exception as pipeline_error: + logger.warning("Failed to update pipeline status: %s", pipeline_error) + return True + + def _orchestrate_file_processing_general( api_client: InternalAPIClient, workflow_id: str, diff --git a/workers/queue_backend/pg_barrier.py b/workers/queue_backend/pg_barrier.py index 08e0a1d464..a85e6ac541 100644 --- a/workers/queue_backend/pg_barrier.py +++ b/workers/queue_backend/pg_barrier.py @@ -1305,11 +1305,10 @@ def _mark_execution_error_on_abort( ) -# A batch whose execution is already terminal returns its result with this -# marker set. run_batch_with_barrier bypasses the barrier decrement for it: the -# reaper has by definition already torn the barrier down, so decrementing would -# only log a spurious "decrement found no row" ERROR (false alert noise), and an -# already-terminal execution must not have its aggregating callback fired. +# A batch skipped because its execution is terminal returns its result with this +# marker set, and run_batch_with_barrier bypasses the barrier decrement for it, so +# that run's aggregating callback does not fire. Set for every terminal status +# except STOPPED: an aborted run's batches count down so the callback can finish it. SKIPPED_TERMINAL_EXECUTION_KEY = "skipped_terminal_execution" @@ -1380,12 +1379,8 @@ def run_batch_with_barrier( try: result = work_fn() if result.get(SKIPPED_TERMINAL_EXECUTION_KEY): - # The batch's execution is already terminal, so the reaper has by - # definition already torn the barrier down. Bypass the decrement — - # decrementing a gone barrier only logs a spurious "decrement found - # no row" ERROR (false alert noise) — and do NOT abort (nothing to - # tear down; and an already-terminal execution must not have its - # aggregating callback re-fired). + # Bypass the decrement and do NOT abort: this run's aggregating + # callback must not fire. logger.info( f"[exec:{execution_id}] batch {batch_index} skipped — execution " f"already terminal; barrier decrement bypassed." diff --git a/workers/shared/api/internal_client.py b/workers/shared/api/internal_client.py index 17b49ea833..bda7cc068f 100644 --- a/workers/shared/api/internal_client.py +++ b/workers/shared/api/internal_client.py @@ -498,6 +498,14 @@ def get_workflow_execution( response_dict = response return WorkflowExecutionResponse.from_api_response(response_dict) + def get_workflow_execution_status( + self, execution_id: str | uuid.UUID, organization_id: str | None = None + ) -> str | None: + """Current execution status, without the execution context. Raises on failure.""" + return self.execution_client.get_workflow_execution_status( + execution_id, organization_id + ) + @with_cache( CacheType.WORKFLOW, lambda self, workflow_id, organization_id=None: str(workflow_id), diff --git a/workers/shared/clients/execution_client.py b/workers/shared/clients/execution_client.py index e021e939a1..30bbdb8e80 100644 --- a/workers/shared/clients/execution_client.py +++ b/workers/shared/clients/execution_client.py @@ -13,10 +13,10 @@ import logging import uuid -from typing import Any +from typing import Any, cast from uuid import UUID -from unstract.core.data_models import ExecutionStatus +from unstract.core.data_models import ALREADY_FINAL_REASON, ExecutionStatus from ..data.models import ( BatchOperationRequest, @@ -88,6 +88,15 @@ def wrapper(*args, **kwargs): logger = logging.getLogger(__name__) +def refused_as_stopped(response: Any) -> bool: + """True when a status write was refused because an abort got there first.""" + data = getattr(response, "data", None) or {} + return bool( + data.get("reason") == ALREADY_FINAL_REASON + and data.get("new_status") == ExecutionStatus.STOPPED.value + ) + + class ExecutionAPIClient(BaseAPIClient): """Specialized API client for workflow execution operations. @@ -147,6 +156,19 @@ def get_workflow_execution( status_code=getattr(e, "status_code", None), ) + def get_workflow_execution_status( + self, execution_id: str | uuid.UUID, organization_id: str | None = None + ) -> str | None: + """Current status of an execution, from the lightweight status endpoint. + + Raises on a failed request; the caller decides how to fail. + """ + response = self.get( + self._build_url("workflow_execution", f"{execution_id}/status/"), + organization_id=organization_id, + ) + return response.get("status") if isinstance(response, dict) else None + def get_workflow_definition( self, workflow_id: str | uuid.UUID, organization_id: str | None = None ) -> ExecutionResponse: @@ -331,7 +353,10 @@ def update_workflow_execution_status( ) # Convert dict response to consistent APIResponse - return convert_dict_response(response, APIResponse) + api_response = cast(APIResponse, convert_dict_response(response, APIResponse)) + # The body is flat, and a refused write says why only there. + api_response.data = response + return api_response def recover_stuck_pg_executions( self, diff --git a/workers/shared/infrastructure/logging/workflow_logger.py b/workers/shared/infrastructure/logging/workflow_logger.py index 1f7f572232..e2107a5521 100644 --- a/workers/shared/infrastructure/logging/workflow_logger.py +++ b/workers/shared/infrastructure/logging/workflow_logger.py @@ -469,6 +469,7 @@ def publish_final_workflow_logs( total_files: int, successful_files: int, failed_files: int, + stopped: bool = False, ) -> None: """Publish final workflow execution logs. @@ -479,14 +480,20 @@ def publish_final_workflow_logs( total_files: Total number of files processed successful_files: Number of successfully processed files failed_files: Number of files that failed processing + stopped: The run was aborted, so it did not execute successfully """ # Publish workflow end status to UI status bar self.publish_update_log(LogState.END_WORKFLOW, "1", LogComponent.STATUS_BAR) # Publish workflow completion status - self.publish_update_log( - LogState.SUCCESS, "Executed successfully", LogComponent.WORKFLOW - ) + if stopped: + self.publish_update_log( + LogState.ERROR, "Execution stopped", LogComponent.WORKFLOW + ) + else: + self.publish_update_log( + LogState.SUCCESS, "Executed successfully", LogComponent.WORKFLOW + ) # Publish detailed execution summary summary_message = ( diff --git a/workers/shared/processing/filter_pipeline.py b/workers/shared/processing/filter_pipeline.py index 7a47bc7d42..79bd901218 100644 --- a/workers/shared/processing/filter_pipeline.py +++ b/workers/shared/processing/filter_pipeline.py @@ -11,6 +11,7 @@ from ..api.internal_client import InternalAPIClient from ..cache.cache_backends import RedisCacheBackend +from ..clients.base_client import InternalAPIClientError from ..infrastructure.logging import WorkerLogger from ..workflow.execution.active_file_manager import ActiveFileManager @@ -529,7 +530,7 @@ def get_name(self) -> str: class ActiveFileFilter(FileFilter): """Filter files that are currently being processed by other executions.""" - def __init__(self): + def __init__(self) -> None: self._cache_checked: set[str] = ( set() ) # Track composite identifiers we've already checked @@ -610,6 +611,31 @@ def apply( return filtered + @staticmethod + def _marker_run_is_live( + run_id: str, identifier: str, api_client: Any, live_runs: dict[str, bool] + ) -> bool: + """False for a stopped run, whose queued files no longer hold their markers. + + A failed status read keeps the marker in force. + """ + if run_id not in live_runs: + try: + status = api_client.get_workflow_execution_status(run_id) + except InternalAPIClientError as error: + logger.warning( + "[ActiveFileFilter] Status read for run %s failed: %s", run_id, error + ) + status = None + live_runs[run_id] = status != ExecutionStatus.STOPPED.value + if not live_runs[run_id]: + logger.info( + "[ActiveFileFilter] Ignoring the marker of stopped run %s for file %s", + run_id, + identifier, + ) + return live_runs[run_id] + def _check_active_files_batch( self, file_identifiers: dict[str, dict[str, Any]], @@ -640,6 +666,7 @@ def _check_active_files_batch( # 1. Check Redis cache for active files using precise cache keys try: cache = RedisCacheBackend() + live_runs: dict[str, bool] = {} for identifier in identifiers_to_check: data = file_identifiers[identifier] uuid = data["uuid"] @@ -653,7 +680,13 @@ def _check_active_files_batch( if cached_data and isinstance(cached_data, dict): cached_exec_id = cached_data.get("data", {}).get("execution_id") - if cached_exec_id and cached_exec_id != execution_id: + if ( + cached_exec_id + and cached_exec_id != execution_id + and self._marker_run_is_live( + cached_exec_id, identifier, api_client, live_runs + ) + ): active_identifiers.add(identifier) logger.debug( f"[ActiveFileFilter] File {identifier} active in cache (exec: {cached_exec_id})" diff --git a/workers/tests/test_abort_between_files.py b/workers/tests/test_abort_between_files.py new file mode 100644 index 0000000000..d96d7584ec --- /dev/null +++ b/workers/tests/test_abort_between_files.py @@ -0,0 +1,302 @@ +"""A running batch stops between files once its execution is STOPPED (UN-4318). + +Pins: the status is checked before each file; only STOPPED (abort) stops the +batch, with the remaining files counted as failed and their rows marked +STOPPED; any other status, or a failed API call, processes the file (fail open). +Also pins the client URL and the facade delegation. +""" + +from __future__ import annotations + +import types +from typing import Any +from unittest import mock + +import pytest +from celery.exceptions import SoftTimeLimitExceeded +from shared.api.internal_client import InternalAPIClient +from shared.clients.base_client import APIRequestError, AuthenticationError +from shared.data.models import APIResponse +from unstract.core.data_models import ExecutionStatus, FileBatchResult + +from file_processing.tasks import _file_item_name, _process_individual_files + +_TASKS = "file_processing.tasks" +FORMATS = ["list", "tuple", "dict"] + + +def _client() -> Any: + """A client that rejects a misnamed method or argument.""" + return mock.create_autospec(InternalAPIClient, instance=True) + + +def _item(index: int, file_format: str) -> Any: + name = f"f{index}.pdf" + data = { + "file_name": name, + "file_path": f"/in/{name}", + "provider_file_uuid": f"uuid-{index}", + } + return {"list": [name, data], "tuple": (name, data), "dict": data}[file_format] + + +def _context(api_client: Any, n_files: int = 3, file_format: str = "dict") -> Any: + files = {f"file_{i}": _item(i, file_format) for i in range(n_files)} + pre_created = { + f"f{i}.pdf": mock.Mock( + id=f"fe-{i}", + file_hash=mock.Mock(is_manualreview_required=False), + ) + for i in range(n_files) + } + return types.SimpleNamespace( + files=files, + execution_id="exec-1", + workflow_id="wf-1", + pre_created_file_executions=pre_created, + organization_context=types.SimpleNamespace(api_client=api_client), + get_setting=lambda key, default=None: default, + metadata={ + "file_data": mock.Mock(), + "workflow_execution": {}, + "result": FileBatchResult(), + "successful_files_for_manual_review": [], + "total_files": n_files, + "workflow_logger": None, + }, + ) + + +def _run( + api_client: Any, n_files: int = 3, file_format: str = "dict" +) -> tuple[FileBatchResult, Any]: + return _run_context(_context(api_client, n_files, file_format)) + + +def _run_context(context: Any) -> tuple[FileBatchResult, Any]: + """Run the loop with file processing stubbed to succeed.""" + + def _succeeded( + _result: Any, _name: Any, _start: Any, result: Any, *args: Any, **kwargs: Any + ) -> None: + result.increment_success() + + with ( + mock.patch(f"{_TASKS}._process_file") as process_file, + mock.patch( + f"{_TASKS}._handle_file_processing_result", side_effect=_succeeded + ), + mock.patch(f"{_TASKS}.log_file_processing_start"), + mock.patch(f"{_TASKS}.log_file_info"), + ): + _process_individual_files(context) + return context.metadata["result"], process_file + + +def _stopped_ids(api_client: Any) -> list[str]: + return [ + c.kwargs["file_execution_id"] + for c in api_client.update_file_execution_status.call_args_list + if c.kwargs["status"] == ExecutionStatus.STOPPED.value + ] + + +@pytest.mark.parametrize("file_format", FORMATS) +def test_stops_after_the_status_flips_to_stopped(file_format: str) -> None: + api_client = _client() + api_client.get_workflow_execution_status.side_effect = [ + ExecutionStatus.EXECUTING.value, + ExecutionStatus.STOPPED.value, + ] + + result, process_file = _run(api_client, file_format=file_format) + + assert process_file.call_count == 1 + assert api_client.get_workflow_execution_status.call_count == 2 + assert _stopped_ids(api_client) == ["fe-1", "fe-2"] + assert (result.successful_files, result.failed_files) == (1, 2) + + +def test_stopped_before_the_first_file_processes_nothing() -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = ( + ExecutionStatus.STOPPED.value + ) + + result, process_file = _run(api_client) + + process_file.assert_not_called() + assert _stopped_ids(api_client) == ["fe-0", "fe-1", "fe-2"] + assert (result.successful_files, result.failed_files) == (0, 3) + + +def test_a_failed_stopped_write_is_logged_and_the_rest_still_written() -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = ( + ExecutionStatus.STOPPED.value + ) + api_client.update_file_execution_status.return_value = APIResponse( + success=False, error="500" + ) + + with mock.patch(f"{_TASKS}.logger") as logger: + result, process_file = _run(api_client) + + process_file.assert_not_called() + assert _stopped_ids(api_client) == ["fe-0", "fe-1", "fe-2"] + assert (result.successful_files, result.failed_files) == (0, 3) + warned = [ + c.args[1] + for c in logger.warning.call_args_list + if c.args[0].startswith("Failed to mark") + ] + assert warned == ["fe-0", "fe-1", "fe-2"] + + +def test_a_file_without_a_row_is_counted_but_not_written() -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = ( + ExecutionStatus.STOPPED.value + ) + context = _context(api_client, file_format="list") + del context.pre_created_file_executions["f1.pdf"] + + result, _ = _run_context(context) + + assert _stopped_ids(api_client) == ["fe-0", "fe-2"] + assert result.failed_files == 3 + + +@pytest.mark.parametrize("file_format", FORMATS) +def test_a_file_already_completed_in_this_run_is_not_counted(file_format: str) -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = ( + ExecutionStatus.STOPPED.value + ) + context = _context(api_client, file_format=file_format) + context.metadata["skipped_already_completed"] = ["uuid-1:/in/f1.pdf"] + + result, _ = _run_context(context) + + assert _stopped_ids(api_client) == ["fe-0", "fe-2"] + assert result.failed_files == 2 + + +def test_error_mid_batch_keeps_processing_the_files() -> None: + """Only an abort stops a running batch; ERROR keeps today's behaviour.""" + api_client = _client() + api_client.get_workflow_execution_status.side_effect = [ + ExecutionStatus.EXECUTING.value, + ExecutionStatus.ERROR.value, + ExecutionStatus.ERROR.value, + ] + + result, process_file = _run(api_client) + + assert process_file.call_count == 3 + assert _stopped_ids(api_client) == [] + assert (result.successful_files, result.failed_files) == (3, 0) + + +@pytest.mark.parametrize( + "status", + [ + ExecutionStatus.PENDING.value, + ExecutionStatus.EXECUTING.value, + ExecutionStatus.ERROR.value, + ExecutionStatus.COMPLETED.value, + None, + ], +) +def test_any_other_status_processes_every_file(status: str | None) -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = status + + result, process_file = _run(api_client) + + assert process_file.call_count == 3 + assert _stopped_ids(api_client) == [] + assert (result.successful_files, result.failed_files) == (3, 0) + + +@pytest.mark.parametrize(("status", "warnings"), [(None, 3), ("RENAMED", 3), ("ERROR", 0)]) +def test_a_missing_or_unknown_status_is_logged(status: Any, warnings: int) -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = status + + with mock.patch(f"{_TASKS}.logger") as logger: + _, process_file = _run(api_client) + + assert process_file.call_count == 3 + logged = [c.args[0] for c in logger.warning.call_args_list] + assert sum("status check before the next file returned" in m for m in logged) == warnings + + +@pytest.mark.parametrize( + "error", [APIRequestError("404"), AuthenticationError("401")], ids=["404", "401"] +) +def test_a_failed_api_call_processes_the_file_anyway(error: Exception) -> None: + api_client = _client() + api_client.get_workflow_execution_status.side_effect = error + + result, process_file = _run(api_client) + + assert process_file.call_count == 3 + assert _stopped_ids(api_client) == [] + assert result.successful_files == 3 + + +@pytest.mark.parametrize("error", [AttributeError("bug"), SoftTimeLimitExceeded()]) +def test_other_errors_in_the_check_are_not_hidden(error: Exception) -> None: + api_client = _client() + api_client.get_workflow_execution_status.side_effect = error + + with pytest.raises(type(error)): + _run(api_client) + + +@pytest.mark.parametrize( + ("item", "name"), + [ + (["a.pdf", {}], "a.pdf"), + (("b.pdf", {}), "b.pdf"), + ({"file_name": "c.pdf"}, "c.pdf"), + (["too", "many", "parts"], None), + (object(), None), + ], +) +def test_file_item_name_reads_every_format(item: Any, name: str | None) -> None: + assert _file_item_name(item) == name + + +def test_client_reads_the_status_endpoint() -> None: + from shared.clients.execution_client import ExecutionAPIClient + + client = ExecutionAPIClient.__new__(ExecutionAPIClient) + body = {"id": "exec-1", "status": "STOPPED"} + with mock.patch.object(client, "get", return_value=body) as get: + assert client.get_workflow_execution_status("exec-1", "org-1") == "STOPPED" + get.assert_called_once_with( + "v1/workflow-execution/exec-1/status/", organization_id="org-1" + ) + + +def test_client_returns_none_for_a_non_dict_body() -> None: + from shared.clients.execution_client import ExecutionAPIClient + + client = ExecutionAPIClient.__new__(ExecutionAPIClient) + with mock.patch.object(client, "get", return_value=None): + assert client.get_workflow_execution_status("exec-1") is None + + +def test_facade_delegates_to_the_execution_client() -> None: + from shared.api.internal_client import InternalAPIClient + + client = InternalAPIClient.__new__(InternalAPIClient) + client.execution_client = mock.MagicMock() + client.execution_client.get_workflow_execution_status.return_value = "ERROR" + + assert client.get_workflow_execution_status("exec-1", "org-1") == "ERROR" + client.execution_client.get_workflow_execution_status.assert_called_once_with( + "exec-1", "org-1" + ) diff --git a/workers/tests/test_abort_finishes_the_run.py b/workers/tests/test_abort_finishes_the_run.py new file mode 100644 index 0000000000..4fc18747dd --- /dev/null +++ b/workers/tests/test_abort_finishes_the_run.py @@ -0,0 +1,606 @@ +"""An aborted run still finishes: skipped and failed batches count down, and the +callback keeps STOPPED, bills completed files and leaves the pipeline status and +notifications to the abort. A re-run ignores the stopped run's markers (UN-4318). +""" + +from __future__ import annotations + +from contextlib import ExitStack +from types import SimpleNamespace +from typing import Any +from unittest import mock + +import callback.tasks as cb +import pytest +from shared.clients.base_client import APIRequestError +from shared.clients.execution_client import ExecutionAPIClient, refused_as_stopped +from shared.infrastructure.logging.workflow_logger import WorkerWorkflowLogger +from shared.processing.filter_pipeline import ActiveFileFilter +from unstract.core.data_models import ExecutionStatus +from unstract.workflow_execution.enums import LogState + +from file_processing.tasks import ( + _process_file_batch_core, + _run_batch_stages, + _run_stages_under_barrier, + _terminal_skip_result, + _TerminalExecutionSkip, +) +from queue_backend import pg_barrier +from queue_backend.pg_barrier import SKIPPED_TERMINAL_EXECUTION_KEY + +STOPPED = ExecutionStatus.STOPPED.value +ERROR = ExecutionStatus.ERROR.value +COMPLETED = ExecutionStatus.COMPLETED.value +_TASKS = "file_processing.tasks" +_REFUSAL = {"status": "skipped", "reason": "already_final", "new_status": STOPPED} +_COMPLETED_REFUSAL = {**_REFUSAL, "new_status": COMPLETED} +_UPDATED = {"status": "updated"} + + +def _batch_data() -> Any: + def data(index: int) -> dict[str, str]: + return {"file_name": f"f{index}.pdf", "provider_file_uuid": f"uuid-{index}"} + + return SimpleNamespace( + files=[["f0.pdf", data(0)], ("f1.pdf", data(1)), data(2), {"file_name": "x"}], + file_data=SimpleNamespace(organization_id="org-1", workflow_id="wf-1"), + ) + + +def _skip_batch(status: str) -> dict[str, Any]: + with ( + mock.patch(f"{_TASKS}._validate_and_parse_batch_data", return_value=_batch_data()), + mock.patch( + f"{_TASKS}._setup_execution_context", + side_effect=_TerminalExecutionSkip("exec-1", status), + ), + ): + return _run_batch_stages({"any": "payload"}, "task-1") + + +def test_a_stopped_runs_skipped_batch_counts_down() -> None: + result = _skip_batch(STOPPED) + + assert SKIPPED_TERMINAL_EXECUTION_KEY not in result + assert result["failed_files"] == 4 + + +@pytest.mark.parametrize("status", [ERROR, COMPLETED]) +def test_other_terminal_statuses_keep_the_marker(status: str) -> None: + assert _skip_batch(status)[SKIPPED_TERMINAL_EXECUTION_KEY] is True + + +@pytest.mark.parametrize(("status", "decrements"), [(STOPPED, 1), (ERROR, 0)]) +def test_only_a_stopped_runs_batch_decrements(status: str, decrements: int) -> None: + result = _terminal_skip_result(_batch_data(), status) + context: Any = {"execution_id": "e", "batch_index": 0, "callback_descriptor": {}} + with ( + mock.patch.object(pg_barrier, "claim_batch", return_value=True), + mock.patch.object(pg_barrier, "_barrier_pg_decrement") as decrement, + mock.patch.object(pg_barrier, "_abort_barrier_in_body") as abort, + ): + pg_barrier.run_batch_with_barrier(context, lambda: result) + + assert decrement.call_count == decrements + abort.assert_not_called() + + +_PAYLOAD = { + "files": [["f0.pdf", {}], ["f1.pdf", {}]], + "file_data": {"execution_id": "e", "organization_id": "org-1"}, +} + + +def _failing_batch(write: Any) -> tuple[Any, Any, Any]: + """A batch whose stages raise; ``write`` answers its run's ERROR mark.""" + client = mock.Mock() + client.update_workflow_execution_status.side_effect = [write] + return ( + mock.patch(f"{_TASKS}._run_batch_stages", side_effect=RuntimeError("boom")), + mock.patch(f"{_TASKS}.create_api_client", return_value=client), + client, + ) + + +def test_a_batch_failing_after_the_abort_counts_down() -> None: + """Also covers an abort landing after the batch's last status check.""" + stages, factory, client = _failing_batch(SimpleNamespace(data=_REFUSAL)) + context: Any = {"execution_id": "e", "batch_index": 3, "callback_descriptor": {}} + with ( + stages, + factory, + mock.patch.object(pg_barrier, "claim_batch", return_value=True), + mock.patch.object(pg_barrier, "_barrier_pg_decrement") as decrement, + mock.patch.object(pg_barrier, "_abort_barrier_in_body") as abort, + ): + _process_file_batch_core(mock.Mock(), _PAYLOAD, context) + + mark = client.update_workflow_execution_status.call_args.kwargs + assert (mark["execution_id"], mark["status"]) == ("e", ERROR) + # The same write as the barrier's own mark (see test_pg_barrier). + assert mark["error_message"] == "[pg-barrier-abort] batch 3 failed." + assert mark["cascade_terminal_files"] is True + assert decrement.call_args.args[0]["failed_files"] == 2 + assert SKIPPED_TERMINAL_EXECUTION_KEY not in decrement.call_args.args[0] + abort.assert_not_called() + + +@pytest.mark.parametrize( + "write", + [ + SimpleNamespace(data=_UPDATED), + SimpleNamespace(data=_COMPLETED_REFUSAL), + APIRequestError("500"), + ], + ids=["up", "completed", "down"], +) +def test_a_batch_failing_in_a_live_run_still_raises(write: Any) -> None: + stages, factory, _ = _failing_batch(write) + with stages, factory, pytest.raises(RuntimeError, match="boom"): + _run_stages_under_barrier(_PAYLOAD, "task-1", 0) + + +def _status_for(run_status: str, successful: int, failed: int) -> str: + counts = { + "total_files": successful + failed, + "successful_files": successful, + "failed_files": failed, + "total_files_processed": successful + failed, + } + with ( + mock.patch.object(cb, "aggregate_file_batch_results", return_value=counts), + mock.patch.object( + cb.WallClockTimeCalculator, "calculate_execution_time", return_value=1.0 + ), + ): + _, status, _ = cb._determine_execution_status_unified( + file_batch_results=[], + api_client=mock.Mock(), + execution_id="e1", + organization_id="org-1", + execution_context={"execution": {"status": run_status}}, + ) + return status + + +@pytest.mark.parametrize(("successful", "failed"), [(1, 9), (0, 10), (10, 0)]) +def test_a_stopped_run_stays_stopped_whatever_the_counts( + successful: int, failed: int +) -> None: + assert _status_for(STOPPED, successful, failed) == STOPPED + + +@pytest.mark.parametrize( + ("successful", "failed", "expected"), [(1, 9, COMPLETED), (0, 10, ERROR)] +) +def test_a_running_run_is_still_judged_by_its_counts( + successful: int, failed: int, expected: str +) -> None: + executing = ExecutionStatus.EXECUTING.value + assert _status_for(executing, successful, failed) == expected + + +def test_a_missing_context_is_judged_by_its_counts() -> None: + assert cb._was_stopped(None) is False + assert cb._was_stopped({}) is False + + +def test_an_abort_after_the_callbacks_fetch_is_reported_as_stopped() -> None: + api_client = mock.Mock() + api_client.update_workflow_execution_status.side_effect = [ + SimpleNamespace(data=_REFUSAL), + SimpleNamespace(data={"status": "updated"}), + ] + + result = cb._update_execution_status_unified( + api_client, "e1", COMPLETED, {"total_files": 2, "successful_files": 2}, "org-1" + ) + + calls = api_client.update_workflow_execution_status.call_args_list + assert [c.kwargs["status"] for c in calls] == [COMPLETED, STOPPED] + assert calls[1].kwargs["successful_files"] == 2 + assert calls[1].kwargs["error_message"] is None + assert result["final_status"] == STOPPED + + +def test_an_accepted_write_keeps_its_status() -> None: + api_client = mock.Mock() + api_client.update_workflow_execution_status.return_value = SimpleNamespace( + data={"status": "updated"} + ) + + result = cb._update_execution_status_unified(api_client, "e1", COMPLETED, {}, "org-1") + + assert api_client.update_workflow_execution_status.call_count == 1 + assert result["final_status"] == COMPLETED + + +def test_the_client_keeps_the_body_of_a_refused_write() -> None: + client = ExecutionAPIClient.__new__(ExecutionAPIClient) + with mock.patch.object(client, "post", return_value=dict(_REFUSAL)): + response = client.update_workflow_execution_status("e1", COMPLETED) + + assert response.success is True + assert refused_as_stopped(response) is True + + +_DOWNSTREAM = ( + "_handle_pipeline_updates_unified", + "_track_subscription_usage_if_available", + "_publish_final_workflow_ui_logs", + "_publish_final_workflow_ui_logs_api", + "_cleanup_execution_resources", + "_handle_notifications_unified", +) + + +def _refusing_client() -> Any: + """Refuses the callback's COMPLETED write, then accepts STOPPED.""" + api_client = mock.Mock() + api_client.update_workflow_execution_status.side_effect = [ + SimpleNamespace(data=_REFUSAL), + SimpleNamespace(data=_UPDATED), + ] + api_client.get_workflow_execution.return_value = SimpleNamespace( + success=True, + data={"execution": {"status": "EXECUTING", "workflow_id": "wf"}, "workflow": {}}, + ) + return api_client + + +def _run_callback(run: Any) -> dict[str, Any]: + """Run a callback whose COMPLETED write is refused; return the stubbed steps.""" + determined = ({"total_files": 1, "successful_files": 1}, COMPLETED, 0) + with ExitStack() as stack: + stack.enter_context( + mock.patch.object( + cb, "_determine_execution_status_unified", return_value=determined + ) + ) + stack.enter_context(mock.patch.object(cb, "_initialize_performance_managers")) + steps = { + name: stack.enter_context(mock.patch.object(cb, name)) for name in _DOWNSTREAM + } + run() + return steps + + +def _statuses(steps: dict[str, Any]) -> set[Any]: + """Every status the callback handed to its later steps.""" + seen = set() + for step in steps.values(): + for call in step.call_args_list: + seen.update(call.args[1:2]) + seen.update( + value + for key, value in call.kwargs.items() + if key in ("final_status", "execution_status", "status") + ) + return seen + + +def test_the_general_callback_carries_on_as_stopped() -> None: + context = cb.CallbackContext() + context.execution_id = "e1" + context.organization_id = "org-1" + context.api_client = _refusing_client() + context.execution_status = "EXECUTING" + + def run() -> None: + with mock.patch.object(cb, "_extract_callback_parameters", return_value=context): + cb._process_batch_callback_core(mock.Mock(), []) + + steps = _run_callback(run) + + assert steps["_handle_pipeline_updates_unified"].called + assert steps["_handle_notifications_unified"].called + assert _statuses(steps) == {STOPPED} + + +def test_the_api_callback_carries_on_as_stopped() -> None: + api_client = _refusing_client() + + def run() -> None: + with ( + mock.patch.object(cb, "create_api_client", return_value=api_client), + mock.patch.object( + cb, "_fetch_pipeline_data_simplified", return_value=("api", "API") + ), + ): + # Celery binds the task itself as ``self``. + cb.process_batch_callback_api( + [], execution_id="e1", organization_id="org-1", pipeline_id="p1" + ) + + steps = _run_callback(run) + + assert steps["_handle_pipeline_updates_unified"].called + assert steps["_handle_notifications_unified"].called + assert _statuses(steps) == {STOPPED} + + +@pytest.mark.parametrize(("status", "calls"), [(STOPPED, 0), (COMPLETED, 1)]) +def test_only_the_abort_writes_a_stopped_runs_pipeline_status( + status: str, calls: int +) -> None: + context = cb.CallbackContext() + context.pipeline_id = "3f0d6c1e-5a0b-4a39-9a7e-0d6f3c2b1a90" + context.execution_id = "e1" + context.organization_id = "org-1" + context.api_client = mock.Mock() + with mock.patch.object(cb, "_update_pipeline_directly") as update: + cb._handle_pipeline_updates_unified(context, status, is_api_deployment=False) + + assert update.call_count == calls + + +@pytest.mark.parametrize(("status", "calls"), [(STOPPED, 0), (COMPLETED, 1)]) +def test_a_stopped_run_is_not_notified_twice(status: str, calls: int) -> None: + with mock.patch.object(cb, "handle_status_notifications") as notify: + result = cb._handle_notifications_unified( + api_client=mock.Mock(), + status=status, + organization_id="org-1", + execution_id="e1", + pipeline_id="p1", + ) + + assert notify.call_count == calls + if status == STOPPED: + assert result["reason"] == "handled_by_abort" + else: + assert result["status"] == "completed" + + +@pytest.mark.parametrize(("status", "commits"), [(STOPPED, 1), (ERROR, 0)]) +def test_a_stopped_run_bills_only_its_finished_files(status: str, commits: int) -> None: + context = cb.CallbackContext() + context.execution_id = "e1" + context.organization_id = "org-1" + context.file_executions = [ + {"id": "done", "status": COMPLETED}, + {"id": "unfinished", "status": STOPPED}, + ] + plugin = mock.MagicMock() + with mock.patch("client_plugin_registry.get_client_plugin", return_value=plugin): + cb._track_subscription_usage_if_available(context, status) + + commit = plugin.commit_batch_subscription_usage + assert commit.call_count == commits + if commits: + assert commit.call_args.kwargs["file_execution_ids"] == ["done"] + + +@pytest.mark.parametrize( + ("publish", "creator"), + [ + (cb._publish_final_workflow_ui_logs, "create_for_general_workflow"), + (cb._publish_final_workflow_ui_logs_api, "create_for_api_workflow"), + ], + ids=["general", "api"], +) +@pytest.mark.parametrize( + ("status", "stopped", "counts"), + [(STOPPED, True, (3, 3, 0)), (COMPLETED, False, (3, 2, 1))], +) +def test_both_callbacks_pass_the_stop_to_the_final_log( + publish: Any, creator: str, status: str, stopped: bool, counts: tuple[int, int, int] +) -> None: + context = cb.CallbackContext() + context.execution_id = "e1" + context.organization_id = "org-1" + context.pipeline_id = "p1" + context.api_client = mock.Mock() + saved = {"total_files": 3, "successful_files": 3, "failed_files": 0} + context.api_client.get_workflow_execution.return_value = SimpleNamespace( + success=True, data={"aggregated_usage_cost": 0, "execution": saved} + ) + worker = {"total_files": 3, "successful_files": 2, "failed_files": 1} + workflow_logger = mock.Mock() + with mock.patch.object( + cb.WorkerWorkflowLogger, creator, return_value=workflow_logger + ): + publish(context=context, aggregated_results=worker, execution_status=status) + + final = workflow_logger.publish_final_workflow_logs.call_args.kwargs + assert final["stopped"] is stopped + logged = (final["total_files"], final["successful_files"], final["failed_files"]) + assert logged == counts + + +@pytest.mark.parametrize( + ("reply", "pipeline_writes"), + [(_REFUSAL, 0), (_UPDATED, 1), (_COMPLETED_REFUSAL, 1)], + ids=["stopped", "live", "completed"], +) +def test_a_run_with_no_files_leaves_a_stopped_pipeline_to_the_abort( + reply: dict[str, Any], pipeline_writes: int +) -> None: + from general import tasks as general_tasks + + api_client = mock.Mock() + api_client.update_workflow_execution_status.return_value = SimpleNamespace( + data=reply + ) + + general_tasks._complete_without_files(api_client, "exec-1", "pipe-1") + + assert api_client.update_pipeline_status.call_count == pipeline_writes + + +def _run_with_no_files(reply: dict[str, Any]) -> tuple[Any, Any]: + from general import tasks as general_tasks + + api_client = mock.Mock(organization_id="org-1") + api_client.update_workflow_execution_status.return_value = SimpleNamespace( + data=reply + ) + source = mock.Mock() + source.list_files_from_source.return_value = ({}, 0) + workflow_logger = mock.Mock() + with ExitStack() as stack: + for name in ("create_organization_context", "WorkflowContextData"): + stack.enter_context(mock.patch.object(general_tasks, name)) + stack.enter_context( + mock.patch.object(general_tasks, "WorkerSourceConnector", return_value=source) + ) + stack.enter_context( + mock.patch.object( + general_tasks.WorkerWorkflowLogger, + "create_for_general_workflow", + return_value=workflow_logger, + ) + ) + general_tasks._execute_general_workflow( + api_client=api_client, + execution_context={"execution": {"id": "exec-1", "workflow_id": "wf-1"}}, + file_batch_results=[], + pipeline_id="pipe-1", + execution_mode=None, + use_file_history=True, + scheduled=False, + schema_name="org-1", + ) + return api_client, workflow_logger + + +@pytest.mark.parametrize( + ("reply", "reported"), [(_REFUSAL, 0), (_UPDATED, 1)], ids=["stopped", "live"] +) +def test_a_stopped_run_with_no_files_reports_nothing( + reply: dict[str, Any], reported: int +) -> None: + api_client, workflow_logger = _run_with_no_files(reply) + + assert api_client.update_pipeline_status.call_count == reported + assert workflow_logger.publish_execution_complete.call_count == reported + + +def _fail_orchestration(reply: dict[str, Any]) -> tuple[Any, Any]: + from general import tasks as general_tasks + + task = general_tasks.async_execute_bin_general + run = task + while hasattr(run, "__wrapped__"): + run = run.__wrapped__ + client = mock.Mock() + client.get_workflow_execution.return_value = SimpleNamespace( + success=False, error="boom" + ) + error_client = mock.MagicMock() + error_client.__enter__.return_value = error_client + error_client.update_workflow_execution_status.return_value = SimpleNamespace( + data=reply + ) + notify = mock.Mock() + last_attempt = SimpleNamespace( + request=SimpleNamespace(id="t-1", retries=task.max_retries), + max_retries=task.max_retries, + ) + with ExitStack() as stack: + stack.enter_context( + mock.patch.dict( + run.__globals__, + { + "_should_skip_duplicate_orchestration": lambda *_: False, + "release_orchestration_claim": lambda _eid: None, + "InternalAPIClient": lambda _config: error_client, + "notify_execution_failure": notify, + }, + ) + ) + stack.enter_context( + mock.patch.object( + general_tasks.WorkerExecutionContext, + "setup_execution_context", + return_value=(mock.Mock(), client), + ) + ) + with pytest.raises(Exception, match="boom"): + run(last_attempt, "org-1", "wf-1", "exec-1", {}, pipeline_id="pipe-1") + return error_client, notify + + +@pytest.mark.parametrize( + ("reply", "reported"), [(_REFUSAL, 0), (_UPDATED, 1)], ids=["stopped", "live"] +) +def test_a_failed_orchestration_leaves_a_stopped_run_to_the_abort( + reply: dict[str, Any], reported: int +) -> None: + error_client, notify = _fail_orchestration(reply) + + assert error_client.update_pipeline_status.call_count == reported + assert notify.call_count == reported + + +def _active_files(status: Any) -> tuple[set[str], Any]: + api_client = mock.Mock() + api_client.get_workflow_execution_status.side_effect = [status] + api_client.check_files_active_processing.return_value = SimpleNamespace( + success=True, data={} + ) + cache = mock.Mock() + cache.get.return_value = {"data": {"execution_id": "old-run"}} + files = { + "u1:/a": {"uuid": "u1", "path": "/a"}, + "u2:/b": {"uuid": "u2", "path": "/b"}, + } + with mock.patch( + "shared.processing.filter_pipeline.RedisCacheBackend", return_value=cache + ): + found = ActiveFileFilter()._check_active_files_batch( + files, "wf-1", "new-run", api_client + ) + return found, api_client + + +def test_a_rerun_ignores_a_stopped_runs_markers() -> None: + with mock.patch("shared.processing.filter_pipeline.logger") as logger: + found, api_client = _active_files(STOPPED) + + assert found == set() + api_client.get_workflow_execution_status.assert_called_once_with("old-run") + ignored = [c.args[2] for c in logger.info.call_args_list if "stopped run" in c.args[0]] + assert ignored == ["u1:/a", "u2:/b"] + + +@pytest.mark.parametrize( + "status", [ExecutionStatus.EXECUTING.value, APIRequestError("404")], ids=["live", "down"] +) +def test_other_runs_markers_stay_in_force(status: Any) -> None: + found, _ = _active_files(status) + + assert found == {"u1:/a", "u2:/b"} + + +@pytest.mark.parametrize( + ("stopped", "state", "message"), + [ + (True, LogState.ERROR, "Execution stopped"), + (False, LogState.SUCCESS, "Executed successfully"), + ], +) +def test_the_final_log_tells_a_stopped_run_apart( + stopped: bool, state: LogState, message: str +) -> None: + workflow_logger = WorkerWorkflowLogger.__new__(WorkerWorkflowLogger) + with ( + mock.patch.object(workflow_logger, "publish_update_log") as update, + mock.patch.object(workflow_logger, "publish_log"), + ): + workflow_logger.publish_final_workflow_logs(3, 1, 2, stopped=stopped) + + assert mock.call(state, message, mock.ANY) in update.call_args_list + + +@pytest.mark.parametrize( + ("status", "expected"), [(STOPPED, (3, 3, 0)), (COMPLETED, (3, 2, 1))] +) +def test_a_stopped_runs_final_log_uses_the_saved_counts( + status: str, expected: tuple[int, int, int] +) -> None: + worker = {"total_files": 3, "successful_files": 2, "failed_files": 1} + saved = {"execution": {"total_files": 3, "successful_files": 3, "failed_files": 0}} + + assert cb._final_log_counts(worker, saved, status) == expected