diff --git a/workers/callback/tasks.py b/workers/callback/tasks.py index 8259fff591..9dbeb36e00 100644 --- a/workers/callback/tasks.py +++ b/workers/callback/tasks.py @@ -71,6 +71,9 @@ def __init__(self): # here) — lets the PG duplicate guard skip a redelivered callback without a # second round-trip. self.execution_status: str | None = None + # The full execution context from that same fetch, reused by the status + # determination step instead of fetching the execution again. + self.execution_context: dict[str, Any] | None = None def _initialize_performance_managers(): @@ -257,6 +260,7 @@ def _determine_execution_status_unified( api_client: InternalAPIClient, execution_id: str, organization_id: str, + execution_context: dict[str, Any], ) -> tuple[dict[str, Any], str, int]: """Unified status determination logic with timeout detection for all callback types. @@ -269,6 +273,8 @@ def _determine_execution_status_unified( api_client: Internal API client for workflow execution queries execution_id: Workflow execution ID organization_id: Organization context + execution_context: Execution context the callback already fetched, + reused here instead of fetching the execution again Returns: Tuple of (aggregated_results, final_status, expected_files) @@ -282,6 +288,7 @@ def _determine_execution_status_unified( execution_id, organization_id, aggregated_results.get("file_results", []), + execution_context=execution_context, ) # Update aggregated results with wall-clock time @@ -297,23 +304,10 @@ def _determine_execution_status_unified( logger.warning(f"Execution time is 0! File results for execution {execution_id}") # Step 3: Get expected file count from workflow execution details for timeout detection - expected_files = 0 - try: - execution_response = api_client.get_workflow_execution( - execution_id, file_execution=False - ) - if execution_response.success: - execution_data = execution_response.data - expected_files = execution_data.get("total_files", 0) - logger.info( - f"Expected files from execution details: {expected_files} for execution {execution_id}" - ) - else: - logger.warning( - f"Could not fetch execution details for {execution_id}: {execution_response.error}" - ) - except Exception as e: - logger.warning(f"Could not fetch execution details for {execution_id}: {e}") + expected_files = (execution_context or {}).get("total_files", 0) + logger.info( + f"Expected files from execution details: {expected_files} for execution {execution_id}" + ) # Step 4: Extract file processing metrics for status determination total_files_processed = aggregated_results.get("total_files_processed", 0) @@ -728,8 +722,9 @@ def _extract_callback_parameters( f"Fetching complete context from workflow execution {context.execution_id}" ) + # file_execution=True: subscription billing reads context.file_executions. execution_response = api_client.get_workflow_execution( - context.execution_id, file_execution=False + context.execution_id, file_execution=True ) if not execution_response.success: @@ -762,9 +757,10 @@ def _extract_callback_parameters( temp_config = WorkerConfig() temp_api_client = InternalAPIClient(temp_config) try: + # file_execution=True: subscription billing reads context.file_executions. execution_response = temp_api_client.get_workflow_execution( context.execution_id, - file_execution=False, + file_execution=True, ) if not execution_response.success: @@ -810,6 +806,7 @@ def _extract_callback_parameters( ) or workflow_definition.get("workflow_id") # Status from the same fetch, for the PG duplicate guard (no extra call). context.execution_status = execution_info.get("status") + context.execution_context = execution_data # Use existing API detection from source_config (no additional API calls needed) is_api_deployment = source_config.get("is_api", False) @@ -1479,6 +1476,7 @@ def _process_batch_callback_core( api_client=context.api_client, execution_id=context.execution_id, organization_id=context.organization_id, + execution_context=context.execution_context, ) ) # Update workflow execution status using unified function @@ -1668,8 +1666,9 @@ def process_batch_callback_api( logger.info(f"Created organization-scoped API client: {organization_id}") try: + # file_execution=True: subscription billing reads context.file_executions. execution_response = api_client.get_workflow_execution( - execution_id, file_execution=False + execution_id, file_execution=True ) if not execution_response.success: raise RuntimeError( @@ -1742,6 +1741,7 @@ def process_batch_callback_api( api_client=api_client, execution_id=execution_id, organization_id=organization_id, + execution_context=execution_context, ) ) diff --git a/workers/shared/api/internal_client.py b/workers/shared/api/internal_client.py index e832c609d0..c3fcf6ac0c 100644 --- a/workers/shared/api/internal_client.py +++ b/workers/shared/api/internal_client.py @@ -472,7 +472,7 @@ def get_workflow_execution( execution_id: str | uuid.UUID, organization_id: str | None = None, include_cost: bool = False, - file_execution: bool = True, + file_execution: bool = False, ) -> WorkflowExecutionResponse: """Get workflow execution with context. @@ -480,6 +480,7 @@ def get_workflow_execution( execution_id: Workflow execution ID organization_id: Optional organization ID override include_cost: Whether to include aggregated usage cost (expensive operation) + file_execution: Whether to include the execution's file-execution rows Returns: WorkflowExecutionResponse containing workflow execution data diff --git a/workers/shared/clients/execution_client.py b/workers/shared/clients/execution_client.py index ef98f8daec..aa4c906c25 100644 --- a/workers/shared/clients/execution_client.py +++ b/workers/shared/clients/execution_client.py @@ -104,7 +104,7 @@ def get_workflow_execution( execution_id: str | uuid.UUID, organization_id: str | None = None, include_cost: bool = False, - file_execution: bool = True, + file_execution: bool = False, ) -> ExecutionResponse: """Get workflow execution with context. @@ -112,20 +112,24 @@ def get_workflow_execution( execution_id: Workflow execution ID organization_id: Optional organization ID override include_cost: Whether to include aggregated usage cost (expensive operation) + file_execution: Whether to include every file-execution row of the + execution (expensive for large executions; only the callback's + billing step needs them) Returns: ExecutionResponse containing workflow execution data """ try: - # Build URL with optional cost parameter url = self._build_url("workflow_execution", f"{str(execution_id)}/") + # file_execution is always sent: the backend defaults it to true, so + # omitting it would load every file-execution row. + params = {"file_execution": "true" if file_execution else "false"} if include_cost: - url += "?include_cost=true" - if file_execution: - url += "?file_execution=true" + params["include_cost"] = "true" response = self.get( url, + params=params, organization_id=organization_id, ) return ExecutionResponse.success_response( diff --git a/workers/shared/processing/files/time_utils.py b/workers/shared/processing/files/time_utils.py index f045455272..2a568ddd5d 100644 --- a/workers/shared/processing/files/time_utils.py +++ b/workers/shared/processing/files/time_utils.py @@ -20,6 +20,7 @@ def calculate_execution_time( execution_id: str, organization_id: str, fallback_results: list[dict[str, Any]] = None, + execution_context: dict[str, Any] | None = None, ) -> float: """Calculate wall-clock execution time with multiple fallback strategies. @@ -28,6 +29,7 @@ def calculate_execution_time( execution_id: Workflow execution ID organization_id: Organization context fallback_results: List of file results for summing as fallback + execution_context: Already-fetched execution context to reuse, if any Returns: Execution time in seconds @@ -35,7 +37,7 @@ def calculate_execution_time( try: # Primary: Get workflow execution start time from backend return WallClockTimeCalculator._get_wall_clock_time( - api_client, execution_id, organization_id + api_client, execution_id, organization_id, execution_context ) except Exception as e: logger.error(f"Error calculating wall-clock time: {e}") @@ -44,30 +46,39 @@ def calculate_execution_time( @staticmethod def _get_wall_clock_time( - api_client: InternalAPIClient, execution_id: str, organization_id: str + api_client: InternalAPIClient, + execution_id: str, + organization_id: str, + execution_context: dict[str, Any] | None = None, ) -> float: - """Get wall-clock time from execution created_at timestamp.""" - execution_response = api_client.get_workflow_execution( - execution_id, organization_id - ) + """Get wall-clock time from execution created_at timestamp. + + ``execution_context`` is an already-fetched execution context; when given, + no API call is made (created_at never changes). + """ + if not execution_context: + execution_response = api_client.get_workflow_execution( + execution_id, organization_id + ) - if not (execution_response.success and execution_response.data): - raise ValueError("Failed to get execution data from API") + if not (execution_response.success and execution_response.data): + raise ValueError("Failed to get execution data from API") + execution_context = execution_response.data # Extract execution data from the nested structure - execution_data = execution_response.data.get("execution", {}) - if not execution_data: + execution = execution_context.get("execution", {}) + if not execution: logger.error( - f"No 'execution' key in API response. Available keys: {list(execution_response.data.keys())}" + f"No 'execution' key in API response. Available keys: {list(execution_context.keys())}" ) raise ValueError("No execution data found in API response") # Get created_at from the execution data - created_at_str = execution_data.get("created_at") + created_at_str = execution.get("created_at") if not created_at_str: logger.error( - f"Missing timestamp field in API response. Available fields: {list(execution_response.data.keys())}" + f"Missing timestamp field in API response. Available fields: {list(execution_context.keys())}" ) # Don't raise error, let it fall back to file timing calculation raise ValueError("No created_at timestamp found in execution data") diff --git a/workers/shared/workflow/execution/service.py b/workers/shared/workflow/execution/service.py index ebf8bce17c..c28e10b3e0 100644 --- a/workers/shared/workflow/execution/service.py +++ b/workers/shared/workflow/execution/service.py @@ -70,6 +70,9 @@ def __init__(self, api_client: InternalAPIClient = None): self.api_client = api_client self.logger = logger self._last_execution_error = None + # (execution_id, context) of the file being processed, reused by the + # source/destination config lookups. + self._execution_context: tuple[str, dict] | None = None def execute_workflow_for_file( self, @@ -556,19 +559,38 @@ def _cleanup_tool_execution_tracker( def _get_workflow_execution_context( self, execution_id: str, workflow_id: str, organization_id: str ) -> tuple[dict, list]: - """Get workflow execution context and tool instances.""" - execution_response = self.api_client.get_workflow_execution(execution_id) - if not execution_response.success: - raise Exception( - f"Failed to get workflow execution: {execution_response.error}" - ) + """Get workflow execution context and tool instances. + + The context is fetched fresh for each file (connector settings can change + during a long batch) and kept for this file's config lookups. + """ + execution_context = self._fetch_execution_context(execution_id) + self._execution_context = (str(execution_id), execution_context) tool_instances_response = self.api_client.get_tool_instances_by_workflow( workflow_id=workflow_id, organization_id=organization_id, ) - return execution_response.data, tool_instances_response.tool_instances + return execution_context, tool_instances_response.tool_instances + + def _fetch_execution_context(self, execution_id: str) -> dict: + """Fetch the execution context, without its file-execution rows.""" + execution_response = self.api_client.get_workflow_execution(execution_id) + if not execution_response.success: + raise Exception( + f"Failed to get workflow execution: {execution_response.error}" + ) + return execution_response.data + + def _get_cached_execution_context(self, execution_id: str) -> dict: + """Execution context already loaded for this file, else fetch it.""" + cached = self._execution_context + if cached is not None and cached[0] == str(execution_id): + return cached[1] + execution_context = self._fetch_execution_context(execution_id) + self._execution_context = (str(execution_id), execution_context) + return execution_context def _get_workflow( self, workflow_id: str, organization_id: str @@ -1706,12 +1728,7 @@ def _get_destination_config( """Get destination configuration for the workflow via API.""" try: # Get workflow execution context which includes destination config - execution_response = self.api_client.get_workflow_execution(execution_id) - if not execution_response.success: - raise Exception( - f"Failed to get execution context: {execution_response.error}" - ) - execution_context = execution_response.data + execution_context = self._get_cached_execution_context(execution_id) destination_config = execution_context.get("destination_config", {}) if not destination_config: @@ -1737,12 +1754,7 @@ def _get_source_config( """Get source configuration for the workflow via API.""" try: # Get workflow execution context which includes source config - execution_response = self.api_client.get_workflow_execution(execution_id) - if not execution_response.success: - raise Exception( - f"Failed to get execution context: {execution_response.error}" - ) - execution_context = execution_response.data + execution_context = self._get_cached_execution_context(execution_id) source_config = execution_context.get("source_config", {}) if not source_config: diff --git a/workers/tests/test_workflow_execution_fetch_reuse.py b/workers/tests/test_workflow_execution_fetch_reuse.py new file mode 100644 index 0000000000..b3f9dc0bc2 --- /dev/null +++ b/workers/tests/test_workflow_execution_fetch_reuse.py @@ -0,0 +1,274 @@ +"""GET workflow-execution: fetched once per step, file-execution rows only on demand. + +The internal endpoint defaults ``file_execution`` to true and then loads every +file-execution row of the execution. Workers called it about 8 times per document, +and the client never sent ``file_execution=false``, so every call paid for those +rows (quadratic in the file count for an ETL run). These tests pin: + +* the client always sends ``file_execution`` explicitly and defaults it to false; +* only the callback's context fetch asks for the rows (billing reads them); +* the callback's status step reuses that fetch instead of two more GETs; +* a file fetches once, and its source config lookup reuses that fetch. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from unittest.mock import MagicMock, patch + +import callback.tasks as cb +import pytest +from shared.clients.execution_client import ExecutionAPIClient +from shared.processing.files.time_utils import WallClockTimeCalculator +from shared.workflow.execution.service import WorkerWorkflowExecutionService +from unstract.core.data_models import ExecutionStatus + + +def _execution_context(**execution): + execution.setdefault("status", ExecutionStatus.EXECUTING.value) + execution.setdefault("workflow_id", "wf") + return { + "execution": execution, + "workflow_definition": {"workflow_id": "wf"}, + "source_config": {"type": "FILESYSTEM", "connector_id": "c1"}, + "destination_config": {"connection_type": "DATABASE"}, + "organization_context": {"organization_id": "org-1"}, + "file_executions": [], + } + + +class TestClientQueryParams: + """``ExecutionAPIClient.get_workflow_execution`` builds its query via params.""" + + def _client(self): + client = ExecutionAPIClient.__new__(ExecutionAPIClient) + client.get = MagicMock(return_value={"execution": {}}) + return client + + def _params(self, client): + assert client.get.call_count == 1 + url = client.get.call_args.args[0] + assert "?" not in url, "query must go through params, not the URL" + return client.get.call_args.kwargs["params"] + + def test_default_sends_file_execution_false(self): + client = self._client() + client.get_workflow_execution("e1") + assert self._params(client) == {"file_execution": "false"} + + def test_file_execution_true_is_sent(self): + client = self._client() + client.get_workflow_execution("e1", file_execution=True) + assert self._params(client) == {"file_execution": "true"} + + def test_include_cost_combines_with_file_execution(self): + # Before: "?include_cost=true?file_execution=true" — the backend read + # include_cost as "true?file_execution=true" and skipped the cost. + client = self._client() + client.get_workflow_execution("e1", include_cost=True, file_execution=True) + assert self._params(client) == { + "file_execution": "true", + "include_cost": "true", + } + + def test_organization_id_passed_through(self): + client = self._client() + client.get_workflow_execution("e1", organization_id="org-9") + assert client.get.call_args.kwargs["organization_id"] == "org-9" + + +class TestCallbackFetchesOnce: + def _api(self, data): + api = MagicMock() + api.get_workflow_execution.return_value = MagicMock(success=True, data=data) + return api + + def test_extract_requests_file_executions_and_keeps_context(self): + data = _execution_context(pipeline_id="p") + data["file_executions"] = [{"id": "f1", "status": "COMPLETED"}] + api = self._api(data) + task = MagicMock() + task.request.id = "t1" + with ( + patch.object(cb, "create_api_client", return_value=api), + patch.object(cb, "_fetch_pipeline_name_from_api", return_value="pipe"), + ): + context = cb._extract_callback_parameters( + task, [], {"execution_id": "e1", "organization_id": "org-1"} + ) + + api.get_workflow_execution.assert_called_once_with("e1", file_execution=True) + assert context.execution_context is data + # Billing still sees the rows it commits usage for. + assert context.file_executions == data["file_executions"] + + def test_status_determination_reuses_execution_context(self): + created = (datetime.now(UTC) - timedelta(seconds=30)).isoformat() + data = _execution_context(created_at=created) + api = MagicMock() + api.get_workflow_execution.side_effect = AssertionError("no extra GET") + + _, status, expected_files = cb._determine_execution_status_unified( + file_batch_results=[], + api_client=api, + execution_id="e1", + organization_id="org-1", + execution_context=data, + ) + + api.get_workflow_execution.assert_not_called() + # Unchanged semantics: total_files is read from the top level, as before. + assert expected_files == 0 + assert status == ExecutionStatus.COMPLETED.value + + def test_api_callback_passes_its_context_to_status_step(self): + data = _execution_context() + data["file_executions"] = [{"id": "f1", "status": "COMPLETED"}] + api = self._api(data) + task = MagicMock() + task.request.id = "t1" + seen = {} + + def determine(**kwargs): + seen.update(kwargs) + raise RuntimeError("stop after status step") + + with ( + patch.object(cb, "create_api_client", return_value=api), + patch.object( + cb, "_fetch_pipeline_data_simplified", return_value=("api", "API") + ), + patch.object( + cb, "_determine_execution_status_unified", side_effect=determine + ), + pytest.raises(RuntimeError, match="stop after status step"), + ): + cb.process_batch_callback_api( + task, [], execution_id="e1", organization_id="org-1", pipeline_id="p" + ) + + api.get_workflow_execution.assert_called_once_with("e1", file_execution=True) + assert seen["execution_context"] is data + + +class TestWallClockReuse: + def test_uses_given_context_without_fetching(self): + created = (datetime.now(UTC) - timedelta(seconds=30)).isoformat() + api = MagicMock() + api.get_workflow_execution.side_effect = AssertionError("no extra GET") + + elapsed = WallClockTimeCalculator.calculate_execution_time( + api, "e1", "org-1", [], execution_context=_execution_context(created_at=created) + ) + + assert 29 <= elapsed <= 60 + api.get_workflow_execution.assert_not_called() + + def test_falls_back_to_file_times_when_context_lacks_created_at(self): + api = MagicMock() + elapsed = WallClockTimeCalculator.calculate_execution_time( + api, + "e1", + "org-1", + [{"processing_time": 2.0}, {"processing_time": 3.0}], + execution_context=_execution_context(), + ) + assert elapsed == pytest.approx(5.0) + api.get_workflow_execution.assert_not_called() + + +class TestPerFileReuse: + """Each file fetches its execution context once, without file-execution rows. + + It is fetched per file, not taken from the batch, so a long batch still sees + connector settings (auth tokens) as they are when each file starts. + """ + + def _service(self): + api = MagicMock() + api.get_workflow_execution.return_value = MagicMock( + success=True, data=_execution_context(pipeline_id="p") + ) + api.get_tool_instances_by_workflow.return_value = MagicMock( + tool_instances=[{"id": "t"}] + ) + return WorkerWorkflowExecutionService(api_client=api), api + + def test_config_lookups_reuse_the_files_fetch(self): + service, api = self._service() + + context, tools = service._get_workflow_execution_context("e1", "wf", "org-1") + source = service._get_source_config("wf", "e1") + destination = service._get_destination_config("wf", "e1") + + api.get_workflow_execution.assert_called_once_with("e1") + assert tools == [{"id": "t"}] + assert source == context["source_config"] + assert destination == context["destination_config"] + + def test_each_file_fetches_fresh(self): + # One service per file (processor.py); a second file must not see the + # first file's context. + first, api = self._service() + first._get_workflow_execution_context("e1", "wf", "org-1") + second = WorkerWorkflowExecutionService(api_client=api) + second._get_workflow_execution_context("e1", "wf", "org-1") + assert api.get_workflow_execution.call_count == 2 + + def test_source_config_alone_fetches(self): + service, api = self._service() + assert service._get_source_config("wf", "e1") == { + "type": "FILESYSTEM", + "connector_id": "c1", + } + api.get_workflow_execution.assert_called_once_with("e1") + + def test_failed_fetch_keeps_none_result(self): + service, api = self._service() + api.get_workflow_execution.return_value = MagicMock(success=False, error="x") + assert service._get_source_config("wf", "e1") is None + + def test_failed_fetch_is_retried_by_a_later_lookup(self): + service, api = self._service() + api.get_workflow_execution.return_value = MagicMock(success=False, error="x") + assert service._get_source_config("wf", "e1") is None + api.get_workflow_execution.return_value = MagicMock( + success=True, data=_execution_context() + ) + assert service._get_source_config("wf", "e1") is not None + + def test_cached_context_is_not_served_for_another_execution(self): + service, api = self._service() + service._get_workflow_execution_context("e1", "wf", "org-1") + service._get_source_config("wf", "e2") + assert [c.args[0] for c in api.get_workflow_execution.call_args_list] == [ + "e1", + "e2", + ] + + +class TestMissingContextFallbacks: + def test_status_step_tolerates_missing_context(self): + api = MagicMock() + api.get_workflow_execution.return_value = MagicMock(success=False, error="x") + _, status, expected_files = cb._determine_execution_status_unified( + file_batch_results=[], + api_client=api, + execution_id="e1", + organization_id="org-1", + execution_context=None, + ) + assert expected_files == 0 + assert status == ExecutionStatus.COMPLETED.value + + def test_wall_clock_refetches_on_empty_context(self): + created = (datetime.now(UTC) - timedelta(seconds=30)).isoformat() + api = MagicMock() + api.get_workflow_execution.return_value = MagicMock( + success=True, data=_execution_context(created_at=created) + ) + elapsed = WallClockTimeCalculator.calculate_execution_time( + api, "e1", "org-1", [], execution_context={} + ) + assert 29 <= elapsed <= 60 + api.get_workflow_execution.assert_called_once()