mirror of
https://github.com/langgenius/dify.git
synced 2026-10-16 15:35:32 +08:00
fix(workflow): recover legacy paused task IDs before upgrades
Add a bounded, repeatable upgrade backfill through the existing task-control service. Validate saved invocation ownership and index only the same unchanged pause, then fail explicitly if active legacy runs remain unresolved. Share the snapshot validation with event replay and document the required worker drain.
This commit is contained in:
@@ -56,6 +56,18 @@
|
||||
|
||||
Leave `SECRET_KEY` empty to let Dify generate a persistent key in the storage directory. To manage it explicitly, generate a value with `openssl rand -base64 42` and set `SECRET_KEY` in `api/.env`.
|
||||
|
||||
## Upgrading workflow task cancellation
|
||||
|
||||
When upgrading from a version that did not persist `workflow_runs.task_id`, stop new workflow admissions and drain the old workers before switching versions. Old running tasks have no recoverable task-to-run index; do not terminate them by guessing a run ID or using an unscoped Redis fallback.
|
||||
|
||||
After applying database migrations, keep API traffic and workflow consumers stopped and recover paused task IDs from their existing storage snapshots:
|
||||
|
||||
```bash
|
||||
uv run --project "$PWD/api" --directory api flask backfill-workflow-task-ids
|
||||
```
|
||||
|
||||
Run this with the deployment's database and storage configuration. The command processes bounded batches and can be rerun. It exits unsuccessfully if legacy running records or missing/invalid pause snapshots remain; inspect the reported run and repair its state before retrying. It preserves run ownership and skips pauses changed by another worker. Once it succeeds, start the new API and workers. Fresh installations need no backfill.
|
||||
|
||||
## Testing
|
||||
|
||||
Run from the repository root:
|
||||
|
||||
@@ -44,6 +44,7 @@ from .retention import (
|
||||
)
|
||||
from .storage import clear_orphaned_file_records, file_usage, migrate_oss, remove_orphaned_files_on_storage
|
||||
from .system import (
|
||||
backfill_workflow_task_ids,
|
||||
convert_to_agent_apps,
|
||||
fix_app_site_missing,
|
||||
reset_encrypt_key_pair,
|
||||
@@ -63,6 +64,7 @@ __all__ = [
|
||||
"archive_workflow_runs_plan",
|
||||
"backfill_plugin_auto_upgrade",
|
||||
"backfill_workflow_run_archive_bundles",
|
||||
"backfill_workflow_task_ids",
|
||||
"clean_expired_messages",
|
||||
"clean_workflow_runs",
|
||||
"cleanup_orphaned_draft_variables",
|
||||
|
||||
@@ -8,6 +8,7 @@ from sqlalchemy.orm import sessionmaker
|
||||
from configs import dify_config
|
||||
from enums import DeploymentEdition
|
||||
from events.app_event import app_was_created
|
||||
from extensions.ext_application_services import application_services
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from libs.db_migration_lock import DbMigrationAutoRenewLock
|
||||
@@ -133,6 +134,17 @@ def convert_to_agent_apps():
|
||||
click.echo(click.style(f"Conversion complete. Converted {len(proceeded_app_ids)} agent apps.", fg="green"))
|
||||
|
||||
|
||||
@click.command(
|
||||
"backfill-workflow-task-ids", help="Recover task IDs from legacy workflow pauses after draining old workers."
|
||||
)
|
||||
def backfill_workflow_task_ids() -> None:
|
||||
try:
|
||||
indexed = application_services().app_tasks.backfill_legacy_workflow_task_ids()
|
||||
except ValueError as error:
|
||||
raise click.ClickException(str(error)) from error
|
||||
click.echo(f"Indexed {indexed} legacy workflow tasks.")
|
||||
|
||||
|
||||
@click.command("upgrade-db", help="Upgrade the database")
|
||||
def upgrade_db():
|
||||
click.echo("Preparing database migration...")
|
||||
|
||||
@@ -8,6 +8,7 @@ def init_app(app: DifyApp):
|
||||
archive_workflow_runs_plan,
|
||||
backfill_plugin_auto_upgrade,
|
||||
backfill_workflow_run_archive_bundles,
|
||||
backfill_workflow_task_ids,
|
||||
clean_expired_messages,
|
||||
clean_workflow_runs,
|
||||
cleanup_orphaned_draft_variables,
|
||||
@@ -81,6 +82,7 @@ def init_app(app: DifyApp):
|
||||
archive_workflow_runs_plan,
|
||||
archive_workflow_runs,
|
||||
backfill_workflow_run_archive_bundles,
|
||||
backfill_workflow_task_ids,
|
||||
delete_archived_workflow_runs,
|
||||
restore_workflow_runs,
|
||||
clean_workflow_runs,
|
||||
|
||||
@@ -195,8 +195,8 @@ class LogstoreAPIWorkflowRunRepository(APIWorkflowRunRepository):
|
||||
@override
|
||||
def backfill_workflow_pause_task_id(
|
||||
self, *, tenant_id: str, app_id: str, workflow_run_id: str, pause_id: str, task_id: str
|
||||
) -> None:
|
||||
self._control_repository.backfill_workflow_pause_task_id(
|
||||
) -> bool:
|
||||
return self._control_repository.backfill_workflow_pause_task_id(
|
||||
tenant_id=tenant_id, app_id=app_id, workflow_run_id=workflow_run_id, pause_id=pause_id, task_id=task_id
|
||||
)
|
||||
|
||||
|
||||
@@ -551,8 +551,8 @@ class APIWorkflowRunRepository(Protocol):
|
||||
|
||||
def backfill_workflow_pause_task_id(
|
||||
self, *, tenant_id: str, app_id: str, workflow_run_id: str, pause_id: str, task_id: str
|
||||
) -> None:
|
||||
"""Index a legacy pause after its already-authorized snapshot is read."""
|
||||
) -> bool:
|
||||
"""Index an unchanged legacy pause; return whether its task ID was written."""
|
||||
...
|
||||
|
||||
def get_workflow_pause(self, workflow_run_id: str) -> WorkflowPauseEntity | None:
|
||||
|
||||
@@ -31,7 +31,7 @@ import sqlalchemy as sa
|
||||
from pydantic import ValidationError
|
||||
from sqlalchemy import and_, delete, func, null, or_, select, tuple_, update
|
||||
from sqlalchemy.engine import CursorResult
|
||||
from sqlalchemy.orm import Session, selectinload, sessionmaker
|
||||
from sqlalchemy.orm import Session, load_only, selectinload, sessionmaker
|
||||
|
||||
from core.workflow.human_input_forms import load_form_tokens_by_form_id
|
||||
from core.workflow.nodes.human_input.entities import FormDefinition
|
||||
@@ -1193,16 +1193,44 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
|
||||
.values(task_id=task_id)
|
||||
)
|
||||
|
||||
def list_legacy_workflow_runs(self, *, after_id: str | None, limit: int) -> Sequence[WorkflowRun]:
|
||||
"""Page unresolved control records across tenants for the upgrade backfill."""
|
||||
stmt = (
|
||||
select(WorkflowRun)
|
||||
.options(
|
||||
load_only(
|
||||
WorkflowRun.id,
|
||||
WorkflowRun.tenant_id,
|
||||
WorkflowRun.app_id,
|
||||
WorkflowRun.task_id,
|
||||
WorkflowRun.status,
|
||||
WorkflowRun.created_by,
|
||||
)
|
||||
)
|
||||
.where(
|
||||
WorkflowRun.task_id.is_(None),
|
||||
WorkflowRun.type.in_([WorkflowType.WORKFLOW, WorkflowType.CHAT]),
|
||||
WorkflowRun.status.in_([WorkflowExecutionStatus.RUNNING, WorkflowExecutionStatus.PAUSED]),
|
||||
)
|
||||
.order_by(WorkflowRun.id)
|
||||
.limit(limit)
|
||||
)
|
||||
if after_id is not None:
|
||||
stmt = stmt.where(WorkflowRun.id > after_id)
|
||||
with self._session_maker() as session:
|
||||
return session.scalars(stmt).all()
|
||||
|
||||
@override
|
||||
def backfill_workflow_pause_task_id(
|
||||
self, *, tenant_id: str, app_id: str, workflow_run_id: str, pause_id: str, task_id: str
|
||||
) -> None:
|
||||
) -> bool:
|
||||
current_pause = select(WorkflowPause.workflow_run_id).where(
|
||||
WorkflowPause.id == pause_id,
|
||||
WorkflowPause.workflow_run_id == workflow_run_id,
|
||||
WorkflowPause.resumed_at.is_(None),
|
||||
)
|
||||
with self._session_maker() as session, session.begin():
|
||||
session.execute(
|
||||
result = session.execute(
|
||||
update(WorkflowRun)
|
||||
.where(
|
||||
WorkflowRun.id.in_(current_pause),
|
||||
@@ -1213,6 +1241,7 @@ class DifyAPISQLAlchemyWorkflowRunRepository(APIWorkflowRunRepository):
|
||||
)
|
||||
.values(task_id=task_id)
|
||||
)
|
||||
return cast(CursorResult, result).rowcount == 1
|
||||
|
||||
@override
|
||||
def get_workflow_pause(
|
||||
|
||||
@@ -15,8 +15,9 @@ from sqlalchemy.orm import sessionmaker
|
||||
from core.app.apps.base_app_queue_manager import AppQueueManager
|
||||
from core.app.apps.execution_coordinator import send_abort_command, set_app_task_stop_flag
|
||||
from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom
|
||||
from core.app.entities.task_entities import WorkflowFinishStreamResponse
|
||||
from core.app.layers.pause_state_persist_layer import WorkflowResumptionContext
|
||||
from core.repositories.factory import DifyCoreRepositoryFactory
|
||||
from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository
|
||||
from extensions.ext_database import db
|
||||
@@ -28,6 +29,8 @@ from models import Account, EndUser, WorkflowRun
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import AppMode
|
||||
from models.workflow import WorkflowNodeExecutionModel
|
||||
from repositories.api_workflow_run_repository import APIWorkflowRunRepository
|
||||
from repositories.entities.workflow_pause import WorkflowPauseEntity
|
||||
from repositories.sqlalchemy_api_workflow_run_repository import (
|
||||
DifyAPISQLAlchemyWorkflowRunRepository,
|
||||
WorkflowTaskOwnerMismatchError,
|
||||
@@ -82,6 +85,40 @@ class AppTaskControlService:
|
||||
set_app_task_stop_flag(task_id, redis=self._redis_client)
|
||||
return True
|
||||
|
||||
def backfill_legacy_workflow_task_ids(self) -> int:
|
||||
"""Index saved pauses before admitting traffic after a legacy-worker drain."""
|
||||
repository = DifyAPISQLAlchemyWorkflowRunRepository(sessionmaker(bind=db.engine))
|
||||
after_id: str | None = None
|
||||
indexed = 0
|
||||
while runs := repository.list_legacy_workflow_runs(after_id=after_id, limit=100):
|
||||
for run in runs:
|
||||
if run.status != WorkflowExecutionStatus.PAUSED:
|
||||
continue
|
||||
try:
|
||||
pause = repository.get_workflow_pause(run.id)
|
||||
if pause is None:
|
||||
continue
|
||||
context = WorkflowResumptionContext.loads(pause.get_state().decode())
|
||||
indexed += AppTaskService.index_workflow_pause_task(
|
||||
workflow_run=run,
|
||||
pause_entity=pause,
|
||||
resumption_context=context,
|
||||
workflow_run_repo=repository,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to index legacy workflow task, workflow_run_id=%s", run.id)
|
||||
after_id = runs[-1].id
|
||||
|
||||
unresolved = repository.list_legacy_workflow_runs(after_id=None, limit=1)
|
||||
if unresolved:
|
||||
run = unresolved[0]
|
||||
raise ValueError(
|
||||
f"Legacy workflow task remains unindexed: {run.id} ({run.status}). "
|
||||
"Drain legacy running executions and repair missing or invalid pause snapshots "
|
||||
"before admitting traffic."
|
||||
)
|
||||
return indexed
|
||||
|
||||
def stop_workflow_task_no_user_check(self, *, task_id: str) -> None:
|
||||
"""Stop a workflow after app admission, without consulting the user ownership cache.
|
||||
|
||||
@@ -114,6 +151,42 @@ class AppTaskControlService:
|
||||
class AppTaskService:
|
||||
"""Compatibility entry point for callers outside ApplicationServices."""
|
||||
|
||||
@staticmethod
|
||||
def index_workflow_pause_task(
|
||||
*,
|
||||
workflow_run: WorkflowRun,
|
||||
pause_entity: WorkflowPauseEntity,
|
||||
resumption_context: WorkflowResumptionContext,
|
||||
workflow_run_repo: APIWorkflowRunRepository,
|
||||
) -> bool:
|
||||
"""Recover a missing task ID only from the same run's saved invocation."""
|
||||
if workflow_run.task_id is not None:
|
||||
return False
|
||||
generate_entity = resumption_context.get_generate_entity()
|
||||
try:
|
||||
task_id = str(UUID(generate_entity.task_id))
|
||||
except ValueError:
|
||||
return False
|
||||
run_id = (
|
||||
generate_entity.workflow_run_id
|
||||
if isinstance(generate_entity, AdvancedChatAppGenerateEntity)
|
||||
else generate_entity.workflow_execution_id
|
||||
)
|
||||
if (
|
||||
run_id != workflow_run.id
|
||||
or generate_entity.app_config.tenant_id != workflow_run.tenant_id
|
||||
or generate_entity.app_config.app_id != workflow_run.app_id
|
||||
or generate_entity.user_id != workflow_run.created_by
|
||||
):
|
||||
return False
|
||||
return workflow_run_repo.backfill_workflow_pause_task_id(
|
||||
tenant_id=workflow_run.tenant_id,
|
||||
app_id=workflow_run.app_id,
|
||||
workflow_run_id=workflow_run.id,
|
||||
pause_id=pause_entity.id,
|
||||
task_id=task_id,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def stop_workflow_task(
|
||||
*,
|
||||
|
||||
@@ -9,7 +9,6 @@ import time
|
||||
from collections.abc import Generator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from sqlalchemy import desc, select
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
@@ -53,6 +52,7 @@ from models.workflow import WorkflowNodeExecutionTriggeredFrom, WorkflowRun
|
||||
from repositories.api_workflow_node_execution_repository import WorkflowNodeExecutionSnapshot
|
||||
from repositories.entities.workflow_pause import WorkflowPauseEntity
|
||||
from repositories.factory import DifyAPIRepositoryFactory
|
||||
from services.app_task_service import AppTaskService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -99,32 +99,18 @@ def build_workflow_event_stream(
|
||||
pause_entity = None
|
||||
|
||||
resumption_context = _load_resumption_context(pause_entity)
|
||||
if pause_entity is not None and resumption_context is not None:
|
||||
saved_generate_entity = resumption_context.get_generate_entity()
|
||||
try:
|
||||
UUID(saved_generate_entity.task_id)
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
saved_run_id = (
|
||||
saved_generate_entity.workflow_run_id
|
||||
if isinstance(saved_generate_entity, AdvancedChatAppGenerateEntity)
|
||||
else saved_generate_entity.workflow_execution_id
|
||||
)
|
||||
if (
|
||||
saved_run_id == workflow_run.id
|
||||
and saved_generate_entity.app_config.tenant_id == tenant_id
|
||||
and saved_generate_entity.app_config.app_id == app_id
|
||||
):
|
||||
# The authorized snapshot is already loaded. Index its trusted
|
||||
# task ID without scanning or re-reading historical state files.
|
||||
workflow_run_repo.backfill_workflow_pause_task_id(
|
||||
tenant_id=tenant_id,
|
||||
app_id=app_id,
|
||||
workflow_run_id=workflow_run.id,
|
||||
pause_id=pause_entity.id,
|
||||
task_id=saved_generate_entity.task_id,
|
||||
)
|
||||
if (
|
||||
pause_entity is not None
|
||||
and resumption_context is not None
|
||||
and workflow_run.tenant_id == tenant_id
|
||||
and workflow_run.app_id == app_id
|
||||
):
|
||||
AppTaskService.index_workflow_pause_task(
|
||||
workflow_run=workflow_run,
|
||||
pause_entity=pause_entity,
|
||||
resumption_context=resumption_context,
|
||||
workflow_run_repo=workflow_run_repo,
|
||||
)
|
||||
message_context: MessageContext | None = None
|
||||
if app_mode == AppMode.ADVANCED_CHAT:
|
||||
if workflow_run.status == WorkflowExecutionStatus.PAUSED:
|
||||
|
||||
@@ -0,0 +1,304 @@
|
||||
"""Upgrade backfill restores trusted task IDs without reviving old executions."""
|
||||
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import event, select
|
||||
from sqlalchemy.engine import Connection, Engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
import services.app_task_service as task_module
|
||||
from core.app.app_config.entities import WorkflowUIBasedAppConfig
|
||||
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity
|
||||
from core.app.layers.pause_state_persist_layer import (
|
||||
WorkflowResumptionContext,
|
||||
_AdvancedChatAppGenerateEntityWrapper,
|
||||
_WorkflowGenerateEntityWrapper,
|
||||
)
|
||||
from graphon.enums import WorkflowExecutionStatus, WorkflowType
|
||||
from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
|
||||
from models.model import AppMode
|
||||
from models.workflow import WorkflowPause, WorkflowRun
|
||||
from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository
|
||||
from services.app_task_service import AppTaskControlService, AppTaskService
|
||||
from tests.unit_tests.services.test_app_task_service import _StopRedis
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pause_states(monkeypatch: pytest.MonkeyPatch) -> dict[str, bytes]:
|
||||
states: dict[str, bytes] = {}
|
||||
|
||||
def load(key: str) -> bytes:
|
||||
try:
|
||||
return states[key]
|
||||
except KeyError as error:
|
||||
raise FileNotFoundError(key) from error
|
||||
|
||||
monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", load)
|
||||
return states
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def backfill_service(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> AppTaskControlService:
|
||||
monkeypatch.setattr(type(task_module.db), "engine", property(lambda _: sqlite_engine))
|
||||
return AppTaskControlService(redis_client=_StopRedis())
|
||||
|
||||
|
||||
def _seed_pause(
|
||||
session: Session, states: dict[str, bytes], app_mode: AppMode = AppMode.WORKFLOW
|
||||
) -> tuple[WorkflowRun, WorkflowPause, WorkflowResumptionContext]:
|
||||
run = WorkflowRun(
|
||||
id=str(uuid4()),
|
||||
tenant_id="tenant-1",
|
||||
app_id="app-1",
|
||||
workflow_id="workflow-1",
|
||||
type=WorkflowType.WORKFLOW if app_mode == AppMode.WORKFLOW else WorkflowType.CHAT,
|
||||
triggered_from=WorkflowRunTriggeredFrom.APP_RUN,
|
||||
version="1",
|
||||
status=WorkflowExecutionStatus.PAUSED,
|
||||
created_by_role=CreatorUserRole.END_USER,
|
||||
created_by="user-1",
|
||||
created_at=datetime(2026, 10, 10),
|
||||
elapsed_time=0,
|
||||
total_tokens=0,
|
||||
total_steps=1,
|
||||
exceptions_count=0,
|
||||
)
|
||||
pause = WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key=f"pause/{run.id}")
|
||||
config = WorkflowUIBasedAppConfig(
|
||||
tenant_id=run.tenant_id, app_id=run.app_id, app_mode=app_mode, workflow_id=run.workflow_id
|
||||
)
|
||||
if app_mode == AppMode.WORKFLOW:
|
||||
wrapper = _WorkflowGenerateEntityWrapper(
|
||||
entity=WorkflowAppGenerateEntity(
|
||||
task_id=str(uuid4()),
|
||||
app_config=config,
|
||||
inputs={},
|
||||
files=[],
|
||||
user_id=run.created_by,
|
||||
stream=True,
|
||||
invoke_from=InvokeFrom.WEB_APP,
|
||||
workflow_execution_id=run.id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
wrapper = _AdvancedChatAppGenerateEntityWrapper(
|
||||
entity=AdvancedChatAppGenerateEntity(
|
||||
task_id=str(uuid4()),
|
||||
app_config=config,
|
||||
inputs={},
|
||||
files=[],
|
||||
user_id=run.created_by,
|
||||
stream=True,
|
||||
invoke_from=InvokeFrom.WEB_APP,
|
||||
workflow_run_id=run.id,
|
||||
query="Approve this run",
|
||||
conversation_id="conversation-1",
|
||||
)
|
||||
)
|
||||
context = WorkflowResumptionContext(generate_entity=wrapper, serialized_graph_runtime_state="{}")
|
||||
states[pause.state_object_key] = context.dumps().encode()
|
||||
session.add_all([run, pause])
|
||||
session.commit()
|
||||
return run, pause, context
|
||||
|
||||
|
||||
@pytest.mark.parametrize("app_mode", [AppMode.WORKFLOW, AppMode.ADVANCED_CHAT])
|
||||
def test_backfill_restores_the_public_task_and_is_idempotent(
|
||||
sqlite_session: Session,
|
||||
pause_states: dict[str, bytes],
|
||||
backfill_service: AppTaskControlService,
|
||||
app_mode: AppMode,
|
||||
) -> None:
|
||||
run, pause, context = _seed_pause(sqlite_session, pause_states, app_mode)
|
||||
task_id = context.get_generate_entity().task_id
|
||||
assert task_id != run.id
|
||||
|
||||
assert backfill_service.backfill_legacy_workflow_task_ids() == 1
|
||||
sqlite_session.refresh(run)
|
||||
assert run.task_id == task_id
|
||||
assert run.status == WorkflowExecutionStatus.PAUSED
|
||||
assert run.stop_requested_at is None
|
||||
assert sqlite_session.get(WorkflowPause, pause.id) is not None
|
||||
sqlite_session.commit()
|
||||
pause_states.clear()
|
||||
assert backfill_service.backfill_legacy_workflow_task_ids() == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["task_id", "run_id", "tenant_id", "app_id", "user_id"])
|
||||
def test_backfill_rejects_snapshot_identity_mismatches(
|
||||
sqlite_session: Session,
|
||||
pause_states: dict[str, bytes],
|
||||
backfill_service: AppTaskControlService,
|
||||
field: str,
|
||||
) -> None:
|
||||
run, pause, _ = _seed_pause(sqlite_session, pause_states)
|
||||
state = json.loads(pause_states[pause.state_object_key])
|
||||
entity = state["generate_entity"]["entity"]
|
||||
if field in {"tenant_id", "app_id"}:
|
||||
entity["app_config"][field] = "other-owner"
|
||||
else:
|
||||
entity["workflow_execution_id" if field == "run_id" else field] = "invalid-task-or-owner"
|
||||
pause_states[pause.state_object_key] = json.dumps(state).encode()
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
backfill_service.backfill_legacy_workflow_task_ids()
|
||||
|
||||
sqlite_session.refresh(run)
|
||||
assert run.task_id is None
|
||||
assert run.status == WorkflowExecutionStatus.PAUSED
|
||||
assert sqlite_session.get(WorkflowPause, pause.id) is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("missing", ["pause", "state", "corrupt_state"])
|
||||
def test_backfill_reports_unrecoverable_pause_without_changing_it(
|
||||
sqlite_session: Session,
|
||||
pause_states: dict[str, bytes],
|
||||
backfill_service: AppTaskControlService,
|
||||
missing: str,
|
||||
) -> None:
|
||||
run, pause, _ = _seed_pause(sqlite_session, pause_states)
|
||||
if missing == "pause":
|
||||
sqlite_session.delete(pause)
|
||||
sqlite_session.commit()
|
||||
elif missing == "state":
|
||||
pause_states.clear()
|
||||
else:
|
||||
pause_states[pause.state_object_key] = b"invalid JSON"
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
backfill_service.backfill_legacy_workflow_task_ids()
|
||||
|
||||
sqlite_session.refresh(run)
|
||||
assert run.task_id is None
|
||||
assert run.status == WorkflowExecutionStatus.PAUSED
|
||||
assert run.stop_requested_at is None
|
||||
|
||||
|
||||
def test_backfill_requires_old_running_work_to_drain_but_indexes_valid_pauses(
|
||||
sqlite_session: Session,
|
||||
pause_states: dict[str, bytes],
|
||||
backfill_service: AppTaskControlService,
|
||||
) -> None:
|
||||
running, running_pause, _ = _seed_pause(sqlite_session, pause_states)
|
||||
running.status = WorkflowExecutionStatus.RUNNING
|
||||
sqlite_session.delete(running_pause)
|
||||
sqlite_session.commit()
|
||||
paused, _, context = _seed_pause(sqlite_session, pause_states)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
backfill_service.backfill_legacy_workflow_task_ids()
|
||||
|
||||
sqlite_session.refresh(running)
|
||||
sqlite_session.refresh(paused)
|
||||
assert running.task_id is None
|
||||
assert running.status == WorkflowExecutionStatus.RUNNING
|
||||
assert running.stop_requested_at is None
|
||||
assert paused.task_id == context.get_generate_entity().task_id
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["replaced_pause", "resumed", "already_bound"])
|
||||
def test_indexing_does_not_overwrite_a_concurrent_transition(
|
||||
sqlite_session: Session,
|
||||
sqlite_session_factory: sessionmaker[Session],
|
||||
pause_states: dict[str, bytes],
|
||||
change: str,
|
||||
) -> None:
|
||||
run, pause, context = _seed_pause(sqlite_session, pause_states)
|
||||
repository = DifyAPISQLAlchemyWorkflowRunRepository(sqlite_session_factory)
|
||||
saved_pause = repository.get_workflow_pause(run.id)
|
||||
assert saved_pause is not None
|
||||
task_id = str(uuid4()) if change == "already_bound" else None
|
||||
if change == "replaced_pause":
|
||||
sqlite_session.delete(pause)
|
||||
sqlite_session.flush()
|
||||
sqlite_session.add(
|
||||
WorkflowPause(workflow_id=run.workflow_id, workflow_run_id=run.id, state_object_key="replacement")
|
||||
)
|
||||
elif change == "resumed":
|
||||
run.status = WorkflowExecutionStatus.RUNNING
|
||||
pause.resumed_at = datetime(2026, 10, 10, 1)
|
||||
else:
|
||||
run.task_id = task_id
|
||||
sqlite_session.commit()
|
||||
|
||||
assert not AppTaskService.index_workflow_pause_task(
|
||||
workflow_run=run,
|
||||
pause_entity=saved_pause,
|
||||
resumption_context=context,
|
||||
workflow_run_repo=repository,
|
||||
)
|
||||
|
||||
sqlite_session.refresh(run)
|
||||
assert run.task_id == task_id
|
||||
|
||||
|
||||
def test_backfill_skips_new_writer_rows_terminal_runs_and_rag_pipelines(
|
||||
sqlite_session: Session,
|
||||
pause_states: dict[str, bytes],
|
||||
backfill_service: AppTaskControlService,
|
||||
) -> None:
|
||||
bound, _, _ = _seed_pause(sqlite_session, pause_states)
|
||||
task_id = str(uuid4())
|
||||
bound.task_id = task_id
|
||||
finished, _, _ = _seed_pause(sqlite_session, pause_states)
|
||||
finished.status = WorkflowExecutionStatus.SUCCEEDED
|
||||
pipeline, _, _ = _seed_pause(sqlite_session, pause_states)
|
||||
pipeline.type = WorkflowType.RAG_PIPELINE
|
||||
sqlite_session.commit()
|
||||
pause_states.clear()
|
||||
|
||||
assert backfill_service.backfill_legacy_workflow_task_ids() == 0
|
||||
sqlite_session.refresh(bound)
|
||||
sqlite_session.refresh(finished)
|
||||
sqlite_session.refresh(pipeline)
|
||||
assert bound.task_id == task_id
|
||||
assert finished.task_id is None
|
||||
assert pipeline.task_id is None
|
||||
assert pipeline.status == WorkflowExecutionStatus.PAUSED
|
||||
|
||||
|
||||
def test_backfill_reads_every_page_even_as_previous_pages_leave_the_query(
|
||||
sqlite_session: Session,
|
||||
pause_states: dict[str, bytes],
|
||||
backfill_service: AppTaskControlService,
|
||||
) -> None:
|
||||
expected = {}
|
||||
for _ in range(101):
|
||||
run, _, context = _seed_pause(sqlite_session, pause_states)
|
||||
expected[run.id] = context.get_generate_entity().task_id
|
||||
|
||||
assert backfill_service.backfill_legacy_workflow_task_ids() == 101
|
||||
assert dict(sqlite_session.execute(select(WorkflowRun.id, WorkflowRun.task_id)).all()) == expected
|
||||
|
||||
|
||||
def test_snapshot_storage_reads_do_not_hold_a_database_transaction(
|
||||
sqlite_session: Session,
|
||||
sqlite_engine: Engine,
|
||||
pause_states: dict[str, bytes],
|
||||
backfill_service: AppTaskControlService,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_seed_pause(sqlite_session, pause_states)
|
||||
active_transactions: set[Connection] = set()
|
||||
|
||||
def load(key: str) -> bytes:
|
||||
assert not active_transactions
|
||||
return pause_states[key]
|
||||
|
||||
monkeypatch.setattr("repositories.sqlalchemy_api_workflow_run_repository.storage.load", load)
|
||||
listeners: tuple[tuple[str, Callable[[Connection], None]], ...] = (
|
||||
("begin", active_transactions.add),
|
||||
("commit", active_transactions.discard),
|
||||
("rollback", active_transactions.discard),
|
||||
)
|
||||
for name, listener in listeners:
|
||||
event.listen(sqlite_engine, name, listener)
|
||||
try:
|
||||
assert backfill_service.backfill_legacy_workflow_task_ids() == 1
|
||||
finally:
|
||||
for name, listener in listeners:
|
||||
event.remove(sqlite_engine, name, listener)
|
||||
Reference in New Issue
Block a user