Kimirock/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(api): extract workflow run archive application service (#41253)

Browse Source
This commit is contained in:
非法操作
2026-08-27 02:41:43 +00:00
committed by GitHub
parent 5a6191b96b
commit a9319c86ee
17 changed files with 1402 additions and 780 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
api/.importlinter
View File

@@ -459,3 +459,20 @@ forbidden_modules =
services.billing_service
sqlalchemy
werkzeug
[importlinter:contract:workflow-run-archive-application-boundary]
name = Workflow run archive application service and task domain are framework and persistence neutral
type = forbidden
source_modules =
services.retention.workflow_run.archive_download_task
services.retention.workflow_run.archive_log_service
forbidden_modules =
configs
controllers
extensions
flask
models
repositories
sqlalchemy
tasks
werkzeug

115
api/controllers/console/workflow_run_archive.py
View File

@@ -4,38 +4,33 @@ from http import HTTPStatus
from flask import redirect
from flask_restx import Resource
from pydantic import BaseModel, Field
from werkzeug.exceptions import Conflict, Forbidden, NotFound
from werkzeug.exceptions import Conflict, NotFound
from controllers.common.fields import RedirectResponse
from controllers.common.schema import register_response_schema_models, register_schema_models
from controllers.console import console_ns
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import (
account_initialization_required,
cloud_edition_billing_paid_plan_required,
model_validate,
only_edition_cloud,
setup_required,
)
from extensions.ext_database import db
from enums import DeploymentEdition
from extensions.ext_application_services import application_services
from fields.base import ResponseModel
from libs.archive_storage import get_export_storage
from libs.helper import dump_response
from libs.login import current_account_with_tenant, login_required
from models import TenantAccountRole
from services.retention.workflow_run.archive_download_preparation import ARCHIVE_DOWNLOAD_MIME_TYPE
from services.retention.workflow_run.archive_download_task_cache import (
from machinery.context import RequestContext
from models.account import TenantAccountRole
from services.retention.workflow_run.archive_download_task import (
WorkflowRunArchiveDownloadStatus,
)
from services.retention.workflow_run.archive_log_service import (
WorkflowRunArchiveDownloadNotReadyError,
WorkflowRunArchiveDownloadTaskNotFoundError,
WorkflowRunArchiveNotFoundError,
create_workflow_run_archive_download_task,
get_ready_workflow_run_archive_download_task,
get_workflow_run_archive_download_task,
list_workflow_run_archives,
)
_WORKFLOW_RUN_ARCHIVE_ALLOWED_ROLES = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN})
class WorkflowRunArchiveDownloadPayload(BaseModel):
"""Request body for preparing one monthly workflow-run archive download."""
@@ -95,36 +90,19 @@ register_response_schema_models(
)
def _current_owner_or_admin_ids() -> tuple[str, str]:
"""Return current Cloud workspace IDs for an owner or admin, independently of enterprise RBAC."""
current_user, current_tenant_id = current_account_with_tenant()
if not current_tenant_id:
raise NotFound("Current workspace not found")
if not TenantAccountRole.is_privileged_role(current_user.current_role):
raise Forbidden()
return current_tenant_id, current_user.id
def _presigned_url_expires_in(expires_at: datetime.datetime) -> int:
"""Keep the storage URL no longer-lived than the Redis task and cap it for browser downloads."""
expires_at_utc = expires_at if expires_at.tzinfo else expires_at.replace(tzinfo=datetime.UTC)
remaining_seconds = int((expires_at_utc - datetime.datetime.now(datetime.UTC)).total_seconds())
return max(1, min(3600, remaining_seconds))
@console_ns.route("/workflow-run-archives")
class WorkflowRunArchivesApi(Resource):
@console_ns.doc("list_workflow_run_archives")
@console_ns.doc(description="List monthly workflow-run archive metadata for the current workspace")
@console_ns.response(200, "Success", console_ns.models[WorkflowRunArchiveListResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@console_account_admission(
editions=frozenset({DeploymentEdition.CLOUD}),
allowed_roles=_WORKFLOW_RUN_ARCHIVE_ALLOWED_ROLES,
)
@cloud_edition_billing_paid_plan_required
def get(self):
tenant_id, _ = _current_owner_or_admin_ids()
return dump_response(WorkflowRunArchiveListResponse, list_workflow_run_archives(db.session(), tenant_id))
def get(self, request_context: RequestContext):
archives = application_services().workflow_run_archives.list_archives(request_context)
return dump_response(WorkflowRunArchiveListResponse, archives)
@console_ns.route("/workflow-run-archives/downloads")
@@ -137,19 +115,16 @@ class WorkflowRunArchiveDownloadsApi(Resource):
"Download task accepted",
console_ns.models[WorkflowRunArchiveDownloadTaskResponse.__name__],
)
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@console_account_admission(
editions=frozenset({DeploymentEdition.CLOUD}),
allowed_roles=_WORKFLOW_RUN_ARCHIVE_ALLOWED_ROLES,
)
@cloud_edition_billing_paid_plan_required
@model_validate(WorkflowRunArchiveDownloadPayload)
def post(self, req_data: WorkflowRunArchiveDownloadPayload):
tenant_id, account_id = _current_owner_or_admin_ids()
def post(self, req_data: WorkflowRunArchiveDownloadPayload, request_context: RequestContext):
try:
task = create_workflow_run_archive_download_task(
db.session(),
tenant_id=tenant_id,
requested_by=account_id,
task = application_services().workflow_run_archives.create_download(
request_context,
year=req_data.year,
month=req_data.month,
)
@@ -163,15 +138,17 @@ class WorkflowRunArchiveDownloadApi(Resource):
@console_ns.doc("get_workflow_run_archive_download")
@console_ns.doc(description="Get a temporary workflow-run archive download task")
@console_ns.response(200, "Success", console_ns.models[WorkflowRunArchiveDownloadTaskResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@console_account_admission(
editions=frozenset({DeploymentEdition.CLOUD}),
allowed_roles=_WORKFLOW_RUN_ARCHIVE_ALLOWED_ROLES,
)
@cloud_edition_billing_paid_plan_required
def get(self, download_id: str):
tenant_id, _ = _current_owner_or_admin_ids()
def get(self, request_context: RequestContext, download_id: str):
try:
task = get_workflow_run_archive_download_task(tenant_id=tenant_id, download_id=download_id)
task = application_services().workflow_run_archives.get_download(
request_context,
download_id=download_id,
)
except WorkflowRunArchiveDownloadTaskNotFoundError as exc:
raise NotFound(str(exc)) from exc
return dump_response(WorkflowRunArchiveDownloadTaskResponse, task)
@@ -187,29 +164,19 @@ class WorkflowRunArchiveDownloadFileApi(Resource):
console_ns.models[RedirectResponse.__name__],
)
@console_ns.response(409, "Download task is not ready")
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@console_account_admission(
editions=frozenset({DeploymentEdition.CLOUD}),
allowed_roles=_WORKFLOW_RUN_ARCHIVE_ALLOWED_ROLES,
)
@cloud_edition_billing_paid_plan_required
def get(self, download_id: str):
tenant_id, _ = _current_owner_or_admin_ids()
def get(self, request_context: RequestContext, download_id: str):
try:
task = get_ready_workflow_run_archive_download_task(tenant_id=tenant_id, download_id=download_id)
presigned_url = application_services().workflow_run_archives.get_download_url(
request_context,
download_id=download_id,
)
except WorkflowRunArchiveDownloadTaskNotFoundError as exc:
raise NotFound(str(exc)) from exc
except WorkflowRunArchiveDownloadNotReadyError as exc:
raise Conflict(str(exc)) from exc
storage_key = task.storage_key
if storage_key is None:
raise Conflict(f"Workflow run archive download is not ready: {download_id}")
storage = get_export_storage()
presigned_url = storage.generate_presigned_url(
storage_key,
expires_in=_presigned_url_expires_in(task.expires_at),
filename=task.file_name,
content_type=ARCHIVE_DOWNLOAD_MIME_TYPE,
)
return redirect(presigned_url, code=HTTPStatus.FOUND)

14
api/extensions/ext_application_services.py
View File

@@ -35,6 +35,7 @@ from repositories.tag_repository import TagRepository
from repositories.trial_app_query_repository import TrialAppQueryRepository
from repositories.trial_app_usage_repository import TrialAppUsageRepository
from repositories.webapp_access_query_repository import WebAppAccessQueryRepository
from repositories.workflow_run_archive_repository import WorkflowRunArchiveBundleQueryRepository
from repositories.workspace_member_query_repository import WorkspaceMemberQueryRepository
from repositories.workspace_query_repository import WorkspaceQueryRepository
from services.account_activation_adapters import (
@@ -101,6 +102,12 @@ from services.recommended_app_catalog_gateway import (
RemoteRecommendedAppCatalogGateway,
)
from services.recommended_app_query_service import RecommendedAppQueryService
from services.retention.workflow_run.archive_download_adapters import (
dispatch_workflow_run_archive_download_task,
sign_workflow_run_archive_download_url,
)
from services.retention.workflow_run.archive_download_task_cache import WorkflowRunArchiveDownloadTaskCache
from services.retention.workflow_run.archive_log_service import WorkflowRunArchiveService
from services.schema_definition_service import SchemaDefinitionService
from services.setup_adapters import RedisSetupLock, RegisterServiceAccountProvisioner
from services.setup_service import SetupService
@@ -172,6 +179,7 @@ class ApplicationServices:
partner_tenant_bindings: PartnerTenantBindingService
recommended_app_queries: RecommendedAppQueryService
trial_app_usage: TrialAppUsageRecorder
workflow_run_archives: WorkflowRunArchiveService
workspace_queries: WorkspaceQueryService
workspace_member_queries: WorkspaceMemberQueryService
tags: TagApplicationService
@@ -397,6 +405,12 @@ def build_application_services(
trial_enabled=trial_app_enabled,
),
trial_app_usage=TrialAppUsageRepository(session_factory=database_client),
workflow_run_archives=WorkflowRunArchiveService(
bundles=WorkflowRunArchiveBundleQueryRepository(session_factory=database_client),
tasks=WorkflowRunArchiveDownloadTaskCache(redis=redis),
dispatcher=dispatch_workflow_run_archive_download_task,
sign_download_url=sign_workflow_run_archive_download_url,
),
workspace_queries=WorkspaceQueryService(
workspaces=workspace_query_repository,
plans=DeploymentWorkspacePlanGateway(),

98
api/repositories/workflow_run_archive_repository.py Normal file
View File

@@ -0,0 +1,98 @@
"""Database repository for workflow-run archive bundle queries."""
from typing import override
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from models.workflow import WorkflowRunArchiveBundle
from services.retention.workflow_run.archive_log_service import (
WorkflowRunArchiveBundleQuery,
WorkflowRunArchiveBundleRecord,
)
class WorkflowRunArchiveBundleQueryRepository(WorkflowRunArchiveBundleQuery):
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
@override
def list_for_tenant(self, tenant_id: str) -> tuple[WorkflowRunArchiveBundleRecord, ...]:
stmt = (
select(
WorkflowRunArchiveBundle.year,
WorkflowRunArchiveBundle.month,
WorkflowRunArchiveBundle.shard,
WorkflowRunArchiveBundle.bundle_id,
WorkflowRunArchiveBundle.workflow_run_count,
WorkflowRunArchiveBundle.row_count,
WorkflowRunArchiveBundle.archive_bytes,
WorkflowRunArchiveBundle.archived_at,
)
.where(WorkflowRunArchiveBundle.tenant_id == tenant_id)
.order_by(
WorkflowRunArchiveBundle.year.desc(),
WorkflowRunArchiveBundle.month.desc(),
WorkflowRunArchiveBundle.shard,
WorkflowRunArchiveBundle.bundle_id,
)
)
with self._session_factory() as session:
rows = session.execute(stmt).all()
return tuple(
WorkflowRunArchiveBundleRecord(
year=row.year,
month=row.month,
shard=row.shard,
bundle_id=row.bundle_id,
workflow_run_count=row.workflow_run_count,
row_count=row.row_count,
archive_bytes=row.archive_bytes,
archived_at=row.archived_at,
)
for row in rows
)
@override
def list_for_tenant_month(
self,
tenant_id: str,
*,
year: int,
month: int,
) -> tuple[WorkflowRunArchiveBundleRecord, ...]:
stmt = (
select(
WorkflowRunArchiveBundle.year,
WorkflowRunArchiveBundle.month,
WorkflowRunArchiveBundle.shard,
WorkflowRunArchiveBundle.bundle_id,
WorkflowRunArchiveBundle.workflow_run_count,
WorkflowRunArchiveBundle.row_count,
WorkflowRunArchiveBundle.archive_bytes,
WorkflowRunArchiveBundle.archived_at,
)
.where(
WorkflowRunArchiveBundle.tenant_id == tenant_id,
WorkflowRunArchiveBundle.year == year,
WorkflowRunArchiveBundle.month == month,
)
.order_by(WorkflowRunArchiveBundle.shard, WorkflowRunArchiveBundle.bundle_id)
)
with self._session_factory() as session:
rows = session.execute(stmt).all()
return tuple(
WorkflowRunArchiveBundleRecord(
year=row.year,
month=row.month,
shard=row.shard,
bundle_id=row.bundle_id,
workflow_run_count=row.workflow_run_count,
row_count=row.row_count,
archive_bytes=row.archive_bytes,
archived_at=row.archived_at,
)
for row in rows
)

37
api/services/retention/workflow_run/archive_download_adapters.py Normal file
View File

@@ -0,0 +1,37 @@
"""Infrastructure adapters used by the workflow-run archive application service."""
from libs.archive_storage import get_export_storage
from services.retention.workflow_run.archive_download_task import WorkflowRunArchiveDownloadTask
_ARCHIVE_DOWNLOAD_MIME_TYPE = "application/zip"
def dispatch_workflow_run_archive_download_task(
task: WorkflowRunArchiveDownloadTask,
) -> None:
"""Enqueue ZIP preparation after the application service claims the Celery task id."""
from tasks.workflow_run_archive_download_tasks import prepare_workflow_run_archive_download_task
celery_task_id = task.celery_task_id
if celery_task_id is None:
raise ValueError("celery_task_id is required before dispatch")
prepare_workflow_run_archive_download_task.apply_async(
args=(task.tenant_id, task.download_id),
task_id=celery_task_id,
)
def sign_workflow_run_archive_download_url(
storage_key: str,
*,
expires_in: int,
filename: str,
) -> str:
"""Create a browser download URL without initializing export storage during app startup."""
return get_export_storage().generate_presigned_url(
storage_key,
expires_in=expires_in,
filename=filename,
content_type=_ARCHIVE_DOWNLOAD_MIME_TYPE,
)

79
api/services/retention/workflow_run/archive_download_preparation.py
View File

@@ -12,31 +12,31 @@ import io
import logging
import re
import zipfile
from collections.abc import Sequence
from collections.abc import Callable, Sequence
import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.csv as pa_csv
import pyarrow.parquet as pq
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from core.helper.csv_sanitizer import CSVSanitizer
from extensions.ext_database import db
from libs.archive_storage import ArchiveStorage, get_archive_storage, get_export_storage
from models.workflow import WorkflowRunArchiveBundle
from libs.archive_storage import ArchiveStorage
from services.retention.workflow_run.archive_bundle_index import (
ARCHIVE_BUNDLE_ROOT_PREFIX,
ArchiveBundleManifest,
ArchiveBundleTableManifestEntry,
decode_archive_bundle_manifest,
)
from services.retention.workflow_run.archive_download_task_cache import (
from services.retention.workflow_run.archive_download_task import (
WorkflowRunArchiveDownloadStatus,
WorkflowRunArchiveDownloadTask,
WorkflowRunArchiveDownloadTaskCache,
build_archive_download_id,
)
from services.retention.workflow_run.archive_log_service import (
WorkflowRunArchiveBundleQuery,
WorkflowRunArchiveBundleRecord,
WorkflowRunArchiveDownloadTaskStore,
)
from services.retention.workflow_run.constants import (
ARCHIVE_BUNDLE_FORMAT,
ARCHIVE_BUNDLE_MANIFEST_NAME,
@@ -46,7 +46,6 @@ from services.retention.workflow_run.constants import (
logger = logging.getLogger(__name__)
ARCHIVE_DOWNLOAD_ROOT_PREFIX = "workflow-runs/downloads/v1/"
ARCHIVE_DOWNLOAD_MIME_TYPE = "application/zip"
_CSV_FORMULA_PREFIX_PATTERN = f"^[{re.escape(''.join(CSVSanitizer.FORMULA_CHARS))}]"
@@ -60,24 +59,23 @@ class WorkflowRunArchiveDownloadPreparer:
without touching long-lived archives.
"""
archive_storage: ArchiveStorage | None
download_storage: ArchiveStorage | None
cache: WorkflowRunArchiveDownloadTaskCache
session_factory: sessionmaker[Session]
_archive_storage_provider: Callable[[], ArchiveStorage]
_download_storage_provider: Callable[[], ArchiveStorage]
cache: WorkflowRunArchiveDownloadTaskStore
_bundles: WorkflowRunArchiveBundleQuery
def __init__(
self,
*,
storage: ArchiveStorage | None = None,
archive_storage: ArchiveStorage | None = None,
download_storage: ArchiveStorage | None = None,
cache: WorkflowRunArchiveDownloadTaskCache | None = None,
session_factory: sessionmaker[Session] | None = None,
archive_storage_provider: Callable[[], ArchiveStorage],
download_storage_provider: Callable[[], ArchiveStorage],
bundles: WorkflowRunArchiveBundleQuery,
cache: WorkflowRunArchiveDownloadTaskStore,
) -> None:
self.archive_storage = archive_storage or storage
self.download_storage = download_storage or storage
self.cache = cache or WorkflowRunArchiveDownloadTaskCache()
self.session_factory = session_factory or sessionmaker(bind=db.engine, expire_on_commit=False)
self._archive_storage_provider = archive_storage_provider
self._download_storage_provider = download_storage_provider
self._bundles = bundles
self.cache = cache
def prepare(self, *, tenant_id: str, download_id: str) -> WorkflowRunArchiveDownloadTask | None:
"""Prepare a ZIP for an existing Redis task and persist terminal task state."""
@@ -91,8 +89,8 @@ class WorkflowRunArchiveDownloadPreparer:
return task
processing_task = self._mark_processing(task)
try:
archive_storage = self.archive_storage or get_archive_storage()
download_storage = self.download_storage or get_export_storage()
archive_storage = self._archive_storage_provider()
download_storage = self._download_storage_provider()
bundles = self._get_task_bundles(processing_task)
payload = self._build_zip_payload(archive_storage, processing_task, bundles)
storage_key = build_archive_download_storage_key(processing_task)
@@ -102,15 +100,19 @@ class WorkflowRunArchiveDownloadPreparer:
logger.exception("Failed to prepare workflow run archive download %s", download_id)
return self._mark_failed(processing_task, error=str(exc))
def _get_task_bundles(self, task: WorkflowRunArchiveDownloadTask) -> list[WorkflowRunArchiveBundle]:
with self.session_factory() as session:
return _list_task_bundles(session, task)
def _get_task_bundles(self, task: WorkflowRunArchiveDownloadTask) -> list[WorkflowRunArchiveBundleRecord]:
indexed_bundles = self._bundles.list_for_tenant_month(
task.tenant_id,
year=task.year,
month=task.month,
)
return _select_task_bundles(indexed_bundles, task)
def _build_zip_payload(
self,
storage: ArchiveStorage,
task: WorkflowRunArchiveDownloadTask,
bundles: Sequence[WorkflowRunArchiveBundle],
bundles: Sequence[WorkflowRunArchiveBundleRecord],
) -> bytes:
zip_root = f"workflow-run-logs-{task.year:04d}-{task.month:02d}"
csv_buffers: dict[str, io.BytesIO] = {}
@@ -218,17 +220,10 @@ def build_archive_download_storage_key(task: WorkflowRunArchiveDownloadTask) ->
)
def _list_task_bundles(session: Session, task: WorkflowRunArchiveDownloadTask) -> list[WorkflowRunArchiveBundle]:
stmt = (
select(WorkflowRunArchiveBundle)
.where(
WorkflowRunArchiveBundle.tenant_id == task.tenant_id,
WorkflowRunArchiveBundle.year == task.year,
WorkflowRunArchiveBundle.month == task.month,
)
.order_by(WorkflowRunArchiveBundle.shard, WorkflowRunArchiveBundle.bundle_id)
)
indexed_bundles = list(session.scalars(stmt))
def _select_task_bundles(
indexed_bundles: Sequence[WorkflowRunArchiveBundleRecord],
task: WorkflowRunArchiveDownloadTask,
) -> list[WorkflowRunArchiveBundleRecord]:
if task.bundle_refs:
requested_refs = [(ref.shard, ref.bundle_id) for ref in task.bundle_refs]
else:
@@ -260,7 +255,7 @@ def _list_task_bundles(session: Session, task: WorkflowRunArchiveDownloadTask) -
def _build_archive_bundle_object_prefix(
task: WorkflowRunArchiveDownloadTask,
bundle: WorkflowRunArchiveBundle,
bundle: WorkflowRunArchiveBundleRecord,
) -> str:
return (
f"{ARCHIVE_BUNDLE_ROOT_PREFIX}tenant_prefix={task.tenant_id[0].lower()}/tenant_id={task.tenant_id}/"
@@ -271,7 +266,7 @@ def _build_archive_bundle_object_prefix(
def _load_and_validate_manifest(
storage: ArchiveStorage,
task: WorkflowRunArchiveDownloadTask,
bundle: WorkflowRunArchiveBundle,
bundle: WorkflowRunArchiveBundleRecord,
object_prefix: str,
) -> tuple[bytes, ArchiveBundleManifest]:
manifest_key = f"{object_prefix}/{ARCHIVE_BUNDLE_MANIFEST_NAME}"
@@ -284,7 +279,7 @@ def _load_and_validate_manifest(
def _validate_manifest(
*,
task: WorkflowRunArchiveDownloadTask,
bundle: WorkflowRunArchiveBundle,
bundle: WorkflowRunArchiveBundleRecord,
manifest: ArchiveBundleManifest,
object_prefix: str,
) -> None:

121
api/services/retention/workflow_run/archive_download_task.py Normal file
View File

@@ -0,0 +1,121 @@
"""Domain objects and builders for workflow-run archive downloads."""
import datetime
import hashlib
import json
from collections.abc import Sequence
from enum import StrEnum
from pydantic import BaseModel, ConfigDict, Field
ARCHIVE_DOWNLOAD_FORMAT_VERSION = "v1"
DEFAULT_ARCHIVE_DOWNLOAD_TASK_TTL_SECONDS = 24 * 60 * 60
class WorkflowRunArchiveDownloadStatus(StrEnum):
"""Lifecycle state for an asynchronous archive download request."""
PENDING = "pending"
PROCESSING = "processing"
READY = "ready"
FAILED = "failed"
class WorkflowRunArchiveBundleRef(BaseModel):
"""Immutable object-store identity for one bundle included in a download task."""
model_config = ConfigDict(extra="forbid")
shard: str
bundle_id: str
class WorkflowRunArchiveDownloadTask(BaseModel):
"""Temporary state for a monthly archive download request."""
model_config = ConfigDict(extra="forbid")
download_id: str
tenant_id: str
requested_by: str
year: int = Field(ge=1)
month: int = Field(ge=1, le=12)
bundle_ids: list[str]
bundle_refs: list[WorkflowRunArchiveBundleRef] = Field(default_factory=list)
bundle_count: int = Field(ge=0)
archive_bytes: int = Field(ge=0)
status: WorkflowRunArchiveDownloadStatus
file_name: str | None = None
storage_key: str | None = None
file_size_bytes: int | None = Field(default=None, ge=0)
celery_task_id: str | None = None
error: str | None = None
created_at: datetime.datetime
updated_at: datetime.datetime
expires_at: datetime.datetime
started_at: datetime.datetime | None = None
finished_at: datetime.datetime | None = None
def build_pending_archive_download_task(
*,
tenant_id: str,
requested_by: str,
year: int,
month: int,
bundle_ids: Sequence[str],
bundle_refs: Sequence[tuple[str, str]] = (),
archive_bytes: int,
download_id: str,
ttl_seconds: int = DEFAULT_ARCHIVE_DOWNLOAD_TASK_TTL_SECONDS,
now: datetime.datetime | None = None,
) -> WorkflowRunArchiveDownloadTask:
"""Create the temporary state stored when the console starts an archive download."""
created_at = now or datetime.datetime.now(datetime.UTC)
if created_at.tzinfo is None:
created_at = created_at.replace(tzinfo=datetime.UTC)
normalized_bundle_ids = list(bundle_ids)
normalized_bundle_refs = [
WorkflowRunArchiveBundleRef(shard=shard, bundle_id=bundle_id) for shard, bundle_id in bundle_refs
]
return WorkflowRunArchiveDownloadTask(
download_id=download_id,
tenant_id=tenant_id,
requested_by=requested_by,
year=year,
month=month,
bundle_ids=normalized_bundle_ids,
bundle_refs=normalized_bundle_refs,
bundle_count=len(normalized_bundle_ids),
archive_bytes=archive_bytes,
status=WorkflowRunArchiveDownloadStatus.PENDING,
created_at=created_at,
updated_at=created_at,
expires_at=created_at + datetime.timedelta(seconds=ttl_seconds),
)
def build_archive_download_id(
*,
tenant_id: str,
year: int,
month: int,
bundle_refs: Sequence[tuple[str, str]],
download_format_version: str = ARCHIVE_DOWNLOAD_FORMAT_VERSION,
) -> str:
"""Build a stable id for the exact archive download content."""
if not bundle_refs:
raise ValueError("bundle_refs must not be empty")
normalized_refs = sorted(f"{shard}:{bundle_id}" for shard, bundle_id in bundle_refs)
payload = json.dumps(
{
"tenant_id": tenant_id,
"year": year,
"month": month,
"bundle_refs": normalized_refs,
"download_format_version": download_format_version,
},
sort_keys=True,
separators=(",", ":"),
)
return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:32]

122
api/services/retention/workflow_run/archive_download_task_cache.py
View File

@@ -1,76 +1,24 @@
"""Redis-backed temporary state for workflow-run archive downloads."""
import datetime
import hashlib
import json
import logging
from collections.abc import Sequence
from enum import StrEnum
from typing import Any
from pydantic import BaseModel, ConfigDict, Field
from extensions.ext_redis import RedisClientWrapper, redis_client
from extensions.ext_redis import RedisClientWrapper
from services.retention.workflow_run.archive_download_task import WorkflowRunArchiveDownloadTask
logger = logging.getLogger(__name__)
ARCHIVE_DOWNLOAD_FORMAT_VERSION = "v1"
DEFAULT_ARCHIVE_DOWNLOAD_TASK_TTL_SECONDS = 24 * 60 * 60
ARCHIVE_DOWNLOAD_TASK_LOCK_TIMEOUT_SECONDS = 30
_CACHE_KEY_PREFIX = "workflow_run_archive_download"
class WorkflowRunArchiveDownloadStatus(StrEnum):
"""Lifecycle state for an asynchronous archive download request."""
PENDING = "pending"
PROCESSING = "processing"
READY = "ready"
FAILED = "failed"
class WorkflowRunArchiveBundleRef(BaseModel):
"""Immutable object-store identity for one bundle included in a download task."""
model_config = ConfigDict(extra="forbid")
shard: str
bundle_id: str
class WorkflowRunArchiveDownloadTask(BaseModel):
"""Temporary Redis payload for a monthly archive download request."""
model_config = ConfigDict(extra="forbid")
download_id: str
tenant_id: str
requested_by: str
year: int = Field(ge=1)
month: int = Field(ge=1, le=12)
bundle_ids: list[str]
bundle_refs: list[WorkflowRunArchiveBundleRef] = Field(default_factory=list)
bundle_count: int = Field(ge=0)
archive_bytes: int = Field(ge=0)
status: WorkflowRunArchiveDownloadStatus
file_name: str | None = None
storage_key: str | None = None
file_size_bytes: int | None = Field(default=None, ge=0)
celery_task_id: str | None = None
error: str | None = None
created_at: datetime.datetime
updated_at: datetime.datetime
expires_at: datetime.datetime
started_at: datetime.datetime | None = None
finished_at: datetime.datetime | None = None
class WorkflowRunArchiveDownloadTaskCache:
"""Store ephemeral archive download task state in Redis with a TTL."""
_redis: RedisClientWrapper
def __init__(self, redis: RedisClientWrapper = redis_client) -> None:
def __init__(self, *, redis: RedisClientWrapper) -> None:
self._redis = redis
def get(self, *, tenant_id: str, download_id: str) -> WorkflowRunArchiveDownloadTask | None:
@@ -111,67 +59,3 @@ class WorkflowRunArchiveDownloadTaskCache:
expires_at_utc = expires_at if expires_at.tzinfo else expires_at.replace(tzinfo=datetime.UTC)
remaining = expires_at_utc - datetime.datetime.now(datetime.UTC)
return max(int(remaining.total_seconds()), 1)
def build_pending_archive_download_task(
*,
tenant_id: str,
requested_by: str,
year: int,
month: int,
bundle_ids: Sequence[str],
bundle_refs: Sequence[tuple[str, str]] = (),
archive_bytes: int,
download_id: str,
ttl_seconds: int = DEFAULT_ARCHIVE_DOWNLOAD_TASK_TTL_SECONDS,
now: datetime.datetime | None = None,
) -> WorkflowRunArchiveDownloadTask:
"""Create the Redis payload stored when the console starts an archive download."""
created_at = now or datetime.datetime.now(datetime.UTC)
if created_at.tzinfo is None:
created_at = created_at.replace(tzinfo=datetime.UTC)
normalized_bundle_ids = list(bundle_ids)
normalized_bundle_refs = [
WorkflowRunArchiveBundleRef(shard=shard, bundle_id=bundle_id) for shard, bundle_id in bundle_refs
]
return WorkflowRunArchiveDownloadTask(
download_id=download_id,
tenant_id=tenant_id,
requested_by=requested_by,
year=year,
month=month,
bundle_ids=normalized_bundle_ids,
bundle_refs=normalized_bundle_refs,
bundle_count=len(normalized_bundle_ids),
archive_bytes=archive_bytes,
status=WorkflowRunArchiveDownloadStatus.PENDING,
created_at=created_at,
updated_at=created_at,
expires_at=created_at + datetime.timedelta(seconds=ttl_seconds),
)
def build_archive_download_id(
*,
tenant_id: str,
year: int,
month: int,
bundle_refs: Sequence[tuple[str, str]],
download_format_version: str = ARCHIVE_DOWNLOAD_FORMAT_VERSION,
) -> str:
"""Build a stable id for the exact archive download content."""
if not bundle_refs:
raise ValueError("bundle_refs must not be empty")
normalized_refs = sorted(f"{shard}:{bundle_id}" for shard, bundle_id in bundle_refs)
payload = json.dumps(
{
"tenant_id": tenant_id,
"year": year,
"month": month,
"bundle_refs": normalized_refs,
"download_format_version": download_format_version,
},
sort_keys=True,
separators=(",", ":"),
)
return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:32]

413
api/services/retention/workflow_run/archive_log_service.py
View File

@@ -1,34 +1,69 @@
"""
Console-facing workflow-run archive queries.
The object store remains the recoverable archive source of truth. This module only reads the DB bundle index and writes
temporary Redis download-task state, so console requests never list R2 online.
"""
"""Application service for Console workflow-run archive downloads."""
import datetime
import logging
import uuid
from collections.abc import Callable
from collections.abc import Sequence
from contextlib import AbstractContextManager
from dataclasses import dataclass
from typing import Any, NamedTuple, Protocol
from sqlalchemy import select
from sqlalchemy.orm import Session
from models.workflow import WorkflowRunArchiveBundle
from services.retention.workflow_run.archive_download_task_cache import (
from machinery.context import RequestContext
from services.retention.workflow_run.archive_download_task import (
WorkflowRunArchiveDownloadStatus,
WorkflowRunArchiveDownloadTask,
WorkflowRunArchiveDownloadTaskCache,
build_archive_download_id,
build_pending_archive_download_task,
)
logger = logging.getLogger(__name__)
ArchiveDownloadTaskDispatcher = Callable[
[WorkflowRunArchiveDownloadTask, WorkflowRunArchiveDownloadTaskCache],
WorkflowRunArchiveDownloadTask,
]
class WorkflowRunArchiveBundleRecord(NamedTuple):
"""Persistence-neutral metadata for one immutable archive bundle."""
year: int
month: int
shard: str
bundle_id: str
workflow_run_count: int
row_count: int
archive_bytes: int
archived_at: datetime.datetime
class WorkflowRunArchiveBundleQuery(Protocol):
def list_for_tenant(self, tenant_id: str) -> Sequence[WorkflowRunArchiveBundleRecord]: ...
def list_for_tenant_month(
self,
tenant_id: str,
*,
year: int,
month: int,
) -> Sequence[WorkflowRunArchiveBundleRecord]: ...
class WorkflowRunArchiveDownloadTaskStore(Protocol):
def get(self, *, tenant_id: str, download_id: str) -> WorkflowRunArchiveDownloadTask | None: ...
def save(self, task: WorkflowRunArchiveDownloadTask) -> None: ...
def lock(self, *, tenant_id: str, download_id: str) -> AbstractContextManager[Any]: ...
class WorkflowRunArchiveDownloadTaskDispatcher(Protocol):
def __call__(self, task: WorkflowRunArchiveDownloadTask) -> None: ...
class WorkflowRunArchiveDownloadUrlSigner(Protocol):
def __call__(
self,
storage_key: str,
*,
expires_in: int,
filename: str,
) -> str: ...
@dataclass(frozen=True)
@@ -47,7 +82,7 @@ class WorkflowRunArchiveMonth:
@dataclass(frozen=True)
class WorkflowRunArchiveSummary:
"""Top-level archive totals shown on the console page."""
"""Top-level archive totals shown on the Console page."""
archived_month_count: int
workflow_run_count: int
@@ -68,212 +103,161 @@ class WorkflowRunArchiveNotFoundError(Exception):
class WorkflowRunArchiveDownloadTaskNotFoundError(Exception):
"""Raised when the temporary Redis task has expired or never existed."""
"""Raised when the temporary download task has expired or never existed."""
class WorkflowRunArchiveDownloadNotReadyError(Exception):
"""Raised when a cached download task has not produced a file yet."""
def list_workflow_run_archives(
session: Session,
tenant_id: str,
*,
cache: WorkflowRunArchiveDownloadTaskCache | None = None,
) -> WorkflowRunArchiveList:
"""Return monthly archive metadata for one tenant from the DB bundle index."""
stmt = (
select(WorkflowRunArchiveBundle)
.where(WorkflowRunArchiveBundle.tenant_id == tenant_id)
.order_by(
WorkflowRunArchiveBundle.year.desc(),
WorkflowRunArchiveBundle.month.desc(),
WorkflowRunArchiveBundle.shard,
WorkflowRunArchiveBundle.bundle_id,
)
)
month_bundles: dict[tuple[int, int], list[WorkflowRunArchiveBundle]] = {}
for bundle in session.scalars(stmt):
month_bundles.setdefault((bundle.year, bundle.month), []).append(bundle)
class WorkflowRunArchiveService:
def __init__(
self,
*,
bundles: WorkflowRunArchiveBundleQuery,
tasks: WorkflowRunArchiveDownloadTaskStore,
dispatcher: WorkflowRunArchiveDownloadTaskDispatcher,
sign_download_url: WorkflowRunArchiveDownloadUrlSigner,
) -> None:
self._bundles = bundles
self._tasks = tasks
self._dispatcher = dispatcher
self._sign_download_url = sign_download_url
task_cache = cache or WorkflowRunArchiveDownloadTaskCache()
months: list[WorkflowRunArchiveMonth] = []
for (year, month), bundles in month_bundles.items():
bundle_refs = [(bundle.shard, bundle.bundle_id) for bundle in bundles]
months.append(
WorkflowRunArchiveMonth(
year=year,
month=month,
bundle_count=len(bundles),
workflow_run_count=sum(bundle.workflow_run_count for bundle in bundles),
row_count=sum(bundle.row_count for bundle in bundles),
archive_bytes=sum(bundle.archive_bytes for bundle in bundles),
latest_archived_at=max(bundle.archived_at for bundle in bundles),
download_task=_get_cached_month_download_task(
task_cache,
tenant_id=tenant_id,
def list_archives(self, context: RequestContext) -> WorkflowRunArchiveList:
"""Return monthly archive metadata for the active workspace."""
tenant_id = self._active_workspace_id(context)
records_by_month: dict[tuple[int, int], list[WorkflowRunArchiveBundleRecord]] = {}
for record in self._bundles.list_for_tenant(tenant_id):
records_by_month.setdefault((record.year, record.month), []).append(record)
months: list[WorkflowRunArchiveMonth] = []
for year, month in sorted(records_by_month, reverse=True):
records = records_by_month[(year, month)]
bundle_refs = [(record.shard, record.bundle_id) for record in records]
months.append(
WorkflowRunArchiveMonth(
year=year,
month=month,
bundle_refs=bundle_refs,
),
bundle_count=len(records),
workflow_run_count=sum(record.workflow_run_count for record in records),
row_count=sum(record.row_count for record in records),
archive_bytes=sum(record.archive_bytes for record in records),
latest_archived_at=max(record.archived_at for record in records),
download_task=self._get_cached_month_download_task(
tenant_id=tenant_id,
year=year,
month=month,
bundle_refs=bundle_refs,
),
)
)
latest_archived_at = max((archive.latest_archived_at for archive in months), default=None)
return WorkflowRunArchiveList(
summary=WorkflowRunArchiveSummary(
archived_month_count=len(months),
workflow_run_count=sum(archive.workflow_run_count for archive in months),
archive_bytes=sum(archive.archive_bytes for archive in months),
latest_archived_at=latest_archived_at,
),
months=months,
)
latest_archived_at = max((month.latest_archived_at for month in months), default=None)
return WorkflowRunArchiveList(
summary=WorkflowRunArchiveSummary(
archived_month_count=len(months),
workflow_run_count=sum(month.workflow_run_count for month in months),
archive_bytes=sum(month.archive_bytes for month in months),
latest_archived_at=latest_archived_at,
),
months=months,
)
def create_download(
self,
context: RequestContext,
*,
year: int,
month: int,
) -> WorkflowRunArchiveDownloadTask:
"""Create or return the idempotent download task for one workspace/month."""
tenant_id = self._active_workspace_id(context)
bundles = self._bundles.list_for_tenant_month(tenant_id, year=year, month=month)
if not bundles:
raise WorkflowRunArchiveNotFoundError(f"Workflow run archive not found: {year:04d}-{month:02d}")
def _get_cached_month_download_task(
cache: WorkflowRunArchiveDownloadTaskCache,
*,
tenant_id: str,
year: int,
month: int,
bundle_refs: list[tuple[str, str]],
) -> WorkflowRunArchiveDownloadTask | None:
if not bundle_refs:
return None
download_id = build_archive_download_id(
tenant_id=tenant_id,
year=year,
month=month,
bundle_refs=bundle_refs,
)
try:
return cache.get(tenant_id=tenant_id, download_id=download_id)
except Exception:
logger.warning("Failed to read cached workflow run archive download task: %s", download_id, exc_info=True)
return None
def create_workflow_run_archive_download_task(
session: Session,
*,
tenant_id: str,
requested_by: str,
year: int,
month: int,
cache: WorkflowRunArchiveDownloadTaskCache | None = None,
dispatcher: ArchiveDownloadTaskDispatcher | None = None,
) -> WorkflowRunArchiveDownloadTask:
"""
Create or return the idempotent Redis task for downloading one tenant/month archive.
The task identity is based on the exact ordered bundle set currently indexed for the month. If the month receives a
new bundle later, the next request gets a different download id and prepares a fresh file.
"""
bundles = _list_archive_bundles(session, tenant_id=tenant_id, year=year, month=month)
if not bundles:
raise WorkflowRunArchiveNotFoundError(f"Workflow run archive not found: {year:04d}-{month:02d}")
bundle_refs = [(bundle.shard, bundle.bundle_id) for bundle in bundles]
download_id = build_archive_download_id(
tenant_id=tenant_id,
year=year,
month=month,
bundle_refs=bundle_refs,
)
task = build_pending_archive_download_task(
tenant_id=tenant_id,
requested_by=requested_by,
year=year,
month=month,
bundle_ids=[bundle.bundle_id for bundle in bundles],
bundle_refs=bundle_refs,
archive_bytes=sum(bundle.archive_bytes for bundle in bundles),
download_id=download_id,
)
task_cache = cache or WorkflowRunArchiveDownloadTaskCache()
dispatch = dispatcher or _dispatch_workflow_run_archive_download_task
with task_cache.lock(tenant_id=tenant_id, download_id=download_id):
existing = task_cache.get(tenant_id=tenant_id, download_id=download_id)
if existing is None or existing.status == WorkflowRunArchiveDownloadStatus.FAILED:
task_to_queue = task
elif existing.status == WorkflowRunArchiveDownloadStatus.PENDING and not existing.celery_task_id:
task_to_queue = existing
else:
return existing
queued_task = task_to_queue.model_copy(
update={"celery_task_id": uuid.uuid4().hex, "updated_at": datetime.datetime.now(datetime.UTC)}
bundle_refs = [(bundle.shard, bundle.bundle_id) for bundle in bundles]
download_id = build_archive_download_id(
tenant_id=tenant_id,
year=year,
month=month,
bundle_refs=bundle_refs,
)
task_cache.save(queued_task)
return dispatch(queued_task, task_cache)
def get_workflow_run_archive_download_task(
*,
tenant_id: str,
download_id: str,
cache: WorkflowRunArchiveDownloadTaskCache | None = None,
) -> WorkflowRunArchiveDownloadTask:
"""Return a cached archive download task or raise when the TTL has expired."""
task_cache = cache or WorkflowRunArchiveDownloadTaskCache()
task = task_cache.get(tenant_id=tenant_id, download_id=download_id)
if task is None:
raise WorkflowRunArchiveDownloadTaskNotFoundError(f"Workflow run archive download not found: {download_id}")
return task
def get_ready_workflow_run_archive_download_task(
*,
tenant_id: str,
download_id: str,
cache: WorkflowRunArchiveDownloadTaskCache | None = None,
) -> WorkflowRunArchiveDownloadTask:
"""Return a ready cached archive download task or raise when the file is not available."""
task = get_workflow_run_archive_download_task(tenant_id=tenant_id, download_id=download_id, cache=cache)
if task.status != WorkflowRunArchiveDownloadStatus.READY or not task.storage_key or not task.file_name:
raise WorkflowRunArchiveDownloadNotReadyError(f"Workflow run archive download is not ready: {download_id}")
return task
def _list_archive_bundles(
session: Session,
*,
tenant_id: str,
year: int,
month: int,
) -> list[WorkflowRunArchiveBundle]:
stmt = (
select(WorkflowRunArchiveBundle)
.where(
WorkflowRunArchiveBundle.tenant_id == tenant_id,
WorkflowRunArchiveBundle.year == year,
WorkflowRunArchiveBundle.month == month,
task = build_pending_archive_download_task(
tenant_id=tenant_id,
requested_by=context.account_id,
year=year,
month=month,
bundle_ids=[bundle.bundle_id for bundle in bundles],
bundle_refs=bundle_refs,
archive_bytes=sum(bundle.archive_bytes for bundle in bundles),
download_id=download_id,
)
.order_by(WorkflowRunArchiveBundle.shard, WorkflowRunArchiveBundle.bundle_id)
)
return list(session.scalars(stmt))
with self._tasks.lock(tenant_id=tenant_id, download_id=download_id):
existing = self._tasks.get(tenant_id=tenant_id, download_id=download_id)
if existing is None or existing.status == WorkflowRunArchiveDownloadStatus.FAILED:
task_to_queue = task
elif existing.status == WorkflowRunArchiveDownloadStatus.PENDING and not existing.celery_task_id:
task_to_queue = existing
else:
return existing
def _dispatch_workflow_run_archive_download_task(
task: WorkflowRunArchiveDownloadTask,
cache: WorkflowRunArchiveDownloadTaskCache,
) -> WorkflowRunArchiveDownloadTask:
"""
Enqueue background ZIP preparation after the caller atomically claimed and cached the Celery id.
"""
from tasks.workflow_run_archive_download_tasks import prepare_workflow_run_archive_download_task
queued_task = task_to_queue.model_copy(
update={"celery_task_id": uuid.uuid4().hex, "updated_at": datetime.datetime.now(datetime.UTC)}
)
self._tasks.save(queued_task)
celery_task_id = task.celery_task_id
if celery_task_id is None:
raise ValueError("celery_task_id is required before dispatch")
try:
self._dispatcher(queued_task)
except Exception:
return self._record_dispatch_failure(queued_task)
return queued_task
try:
prepare_workflow_run_archive_download_task.apply_async(
args=(task.tenant_id, task.download_id),
task_id=celery_task_id,
def get_download(self, context: RequestContext, *, download_id: str) -> WorkflowRunArchiveDownloadTask:
"""Return a cached download task or raise after its TTL expires."""
tenant_id = self._active_workspace_id(context)
task = self._tasks.get(tenant_id=tenant_id, download_id=download_id)
if task is None:
raise WorkflowRunArchiveDownloadTaskNotFoundError(f"Workflow run archive download not found: {download_id}")
return task
def get_download_url(self, context: RequestContext, *, download_id: str) -> str:
"""Return a short-lived URL for a ready archive download."""
task = self.get_download(context, download_id=download_id)
if task.status != WorkflowRunArchiveDownloadStatus.READY or not task.storage_key or not task.file_name:
raise WorkflowRunArchiveDownloadNotReadyError(f"Workflow run archive download is not ready: {download_id}")
return self._sign_download_url(
task.storage_key,
expires_in=self._presigned_url_expires_in(task.expires_at),
filename=task.file_name,
)
except Exception:
def _get_cached_month_download_task(
self,
*,
tenant_id: str,
year: int,
month: int,
bundle_refs: Sequence[tuple[str, str]],
) -> WorkflowRunArchiveDownloadTask | None:
download_id = build_archive_download_id(
tenant_id=tenant_id,
year=year,
month=month,
bundle_refs=bundle_refs,
)
try:
return self._tasks.get(tenant_id=tenant_id, download_id=download_id)
except Exception:
logger.warning("Failed to read cached workflow run archive download task: %s", download_id, exc_info=True)
return None
def _record_dispatch_failure(
self,
task: WorkflowRunArchiveDownloadTask,
) -> WorkflowRunArchiveDownloadTask:
failure_time = datetime.datetime.now(datetime.UTC)
failed_task = task.model_copy(
update={
@@ -283,16 +267,27 @@ def _dispatch_workflow_run_archive_download_task(
"finished_at": failure_time,
}
)
with cache.lock(tenant_id=task.tenant_id, download_id=task.download_id):
current = cache.get(tenant_id=task.tenant_id, download_id=task.download_id)
with self._tasks.lock(tenant_id=task.tenant_id, download_id=task.download_id):
current = self._tasks.get(tenant_id=task.tenant_id, download_id=task.download_id)
if (
current is not None
and current.status == WorkflowRunArchiveDownloadStatus.PENDING
and current.celery_task_id == celery_task_id
and current.celery_task_id == task.celery_task_id
):
cache.save(failed_task)
self._tasks.save(failed_task)
current = failed_task
logger.exception("Failed to enqueue workflow run archive download task %s", task.download_id)
return current or failed_task
return task
@staticmethod
def _presigned_url_expires_in(expires_at: datetime.datetime) -> int:
expires_at_utc = expires_at if expires_at.tzinfo else expires_at.replace(tzinfo=datetime.UTC)
remaining_seconds = int((expires_at_utc - datetime.datetime.now(datetime.UTC)).total_seconds())
return max(1, min(3600, remaining_seconds))
@staticmethod
def _active_workspace_id(context: RequestContext) -> str:
workspace_id = context.active_workspace_id
if workspace_id is None:
raise RuntimeError("Console account admission did not resolve an active workspace")
return workspace_id

12
api/tasks/workflow_run_archive_download_tasks.py
View File

@@ -4,7 +4,12 @@ import logging
from celery import shared_task
from core.db.session_factory import get_session_maker
from extensions.ext_redis import redis_client
from libs.archive_storage import get_archive_storage, get_export_storage
from repositories.workflow_run_archive_repository import WorkflowRunArchiveBundleQueryRepository
from services.retention.workflow_run.archive_download_preparation import WorkflowRunArchiveDownloadPreparer
from services.retention.workflow_run.archive_download_task_cache import WorkflowRunArchiveDownloadTaskCache
logger = logging.getLogger(__name__)
@@ -15,4 +20,9 @@ WORKFLOW_RUN_ARCHIVE_DOWNLOAD_QUEUE = "workflow_archive"
def prepare_workflow_run_archive_download_task(tenant_id: str, download_id: str) -> None:
"""Prepare a cached workflow-run archive download in the background."""
logger.info("Preparing workflow run archive download: tenant=%s download_id=%s", tenant_id, download_id)
WorkflowRunArchiveDownloadPreparer().prepare(tenant_id=tenant_id, download_id=download_id)
WorkflowRunArchiveDownloadPreparer(
archive_storage_provider=get_archive_storage,
bundles=WorkflowRunArchiveBundleQueryRepository(session_factory=get_session_maker()),
cache=WorkflowRunArchiveDownloadTaskCache(redis=redis_client),
download_storage_provider=get_export_storage,
).prepare(tenant_id=tenant_id, download_id=download_id)

267
api/tests/unit_tests/controllers/console/test_workflow_run_archive.py
View File

@@ -1,15 +1,42 @@
import inspect
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from werkzeug.exceptions import Forbidden
from flask import Flask
from werkzeug.exceptions import Conflict, Forbidden, NotFound
from configs import dify_config
from controllers.console import workflow_run_archive
from controllers.console import flask_admission, workflow_run_archive
from controllers.console.workflow_run_archive import (
WorkflowRunArchiveDownloadApi,
WorkflowRunArchiveDownloadFileApi,
WorkflowRunArchiveDownloadPayload,
WorkflowRunArchiveDownloadsApi,
WorkflowRunArchivesApi,
)
from enums import CloudPlan, DeploymentEdition
from libs.login import AccountWithTenant
from machinery.context import RequestContext
from models import Account, TenantAccountRole
from services.retention.workflow_run.archive_log_service import (
WorkflowRunArchiveDownloadNotReadyError,
WorkflowRunArchiveDownloadTaskNotFoundError,
WorkflowRunArchiveNotFoundError,
)
_CONTEXT = RequestContext(
request_id="request-1",
trace_id="trace-1",
account_id="account-1",
active_workspace_id="tenant-1",
)
_ENDPOINTS = [
WorkflowRunArchivesApi.get,
WorkflowRunArchiveDownloadsApi.post,
WorkflowRunArchiveDownloadApi.get,
WorkflowRunArchiveDownloadFileApi.get,
]
def _account(role: TenantAccountRole) -> Account:
@@ -19,83 +46,209 @@ def _account(role: TenantAccountRole) -> Account:
return account
@pytest.mark.parametrize(
"role",
[TenantAccountRole.EDITOR, TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR],
)
@pytest.mark.parametrize("rbac_enabled", [False, True])
def test_current_owner_or_admin_ids_rejects_non_manager(
monkeypatch: pytest.MonkeyPatch, role: TenantAccountRole, rbac_enabled: bool
def _original(method):
while hasattr(method, "__wrapped__"):
method = method.__wrapped__
return method
def _admission_injector(method):
method = inspect.getclosurevars(method).nonlocals["admitted"]
while "inject_request_context" not in method.__code__.co_qualname:
method = method.__wrapped__
return method
@pytest.mark.parametrize("method", _ENDPOINTS)
def test_workflow_run_archive_endpoints_declare_cloud_paid_plan_admission(method) -> None:
decorator_names = set()
current = method
while hasattr(current, "__wrapped__"):
decorator_names.add(current.__code__.co_qualname.partition(".<locals>")[0])
current = current.__wrapped__
admission = _admission_injector(method)
allowed_roles = inspect.getclosurevars(admission).nonlocals["allowed_roles"]
assert "console_account_admission" in decorator_names
assert "cloud_edition_billing_paid_plan_required" in decorator_names
assert "only_edition_cloud" not in decorator_names
assert allowed_roles == frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN})
@pytest.mark.parametrize("method", _ENDPOINTS)
def test_workflow_run_archive_endpoints_reject_non_manager_when_rbac_is_disabled(
monkeypatch: pytest.MonkeyPatch,
method,
) -> None:
current_user = _account(role)
monkeypatch.setattr(dify_config, "RBAC_ENABLED", rbac_enabled)
account = _account(TenantAccountRole.NORMAL)
monkeypatch.setattr(dify_config, "RBAC_ENABLED", False)
monkeypatch.setattr(
workflow_run_archive,
flask_admission,
"current_account_with_tenant",
lambda: (current_user, "tenant-1"),
lambda: AccountWithTenant(account=account, tenant_id="tenant-1"),
)
app = Flask(__name__)
with pytest.raises(Forbidden):
workflow_run_archive._current_owner_or_admin_ids()
@pytest.mark.parametrize("role", [TenantAccountRole.OWNER, TenantAccountRole.ADMIN])
@pytest.mark.parametrize("rbac_enabled", [False, True])
def test_current_owner_or_admin_ids_returns_current_ids_for_manager(
monkeypatch: pytest.MonkeyPatch, role: TenantAccountRole, rbac_enabled: bool
) -> None:
current_user = _account(role)
monkeypatch.setattr(dify_config, "RBAC_ENABLED", rbac_enabled)
monkeypatch.setattr(
workflow_run_archive,
"current_account_with_tenant",
lambda: (current_user, "tenant-1"),
)
assert workflow_run_archive._current_owner_or_admin_ids() == ("tenant-1", "account-1")
with app.test_request_context(), pytest.raises(Forbidden):
_admission_injector(method)(None)
@pytest.mark.parametrize(
("method", "args"),
[
(WorkflowRunArchivesApi.get, ()),
(WorkflowRunArchiveDownloadsApi.post, (None,)),
(WorkflowRunArchiveDownloadsApi.post, ()),
(WorkflowRunArchiveDownloadApi.get, ("download-1",)),
(WorkflowRunArchiveDownloadFileApi.get, ("download-1",)),
],
)
def test_workflow_run_archive_endpoints_enforce_fixed_workspace_roles(
monkeypatch: pytest.MonkeyPatch, method, args: tuple[str, ...]
def test_workflow_run_archive_endpoints_are_hidden_outside_cloud(
monkeypatch: pytest.MonkeyPatch,
method,
args: tuple[object, ...],
) -> None:
while hasattr(method, "__wrapped__"):
method = method.__wrapped__
monkeypatch.setattr(dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY)
app = Flask(__name__)
def reject_non_manager() -> tuple[str, str]:
raise Forbidden()
monkeypatch.setattr(workflow_run_archive, "_current_owner_or_admin_ids", reject_non_manager)
with pytest.raises(Forbidden):
with app.test_request_context(), pytest.raises(NotFound):
method(None, *args)
def test_workflow_run_archive_endpoint_allows_admitted_role_when_rbac_is_enabled(
monkeypatch: pytest.MonkeyPatch,
) -> None:
account = _account(TenantAccountRole.NORMAL)
monkeypatch.setattr(dify_config, "RBAC_ENABLED", True)
account_with_tenant = AccountWithTenant(account=account, tenant_id="tenant-1")
monkeypatch.setattr(flask_admission, "current_account_with_tenant", lambda: account_with_tenant)
monkeypatch.setattr("controllers.console.wraps.current_account_with_tenant", lambda: account_with_tenant)
def get_billing_info(tenant_id: str, *, exclude_vector_space: bool):
assert tenant_id == "tenant-1"
assert exclude_vector_space
return {
"enabled": True,
"subscription": {"plan": CloudPlan.TEAM},
}
monkeypatch.setattr(
"controllers.console.wraps.BillingService.get_info",
get_billing_info,
)
service = MagicMock()
service.list_archives.return_value = object()
monkeypatch.setattr(
workflow_run_archive,
"application_services",
lambda: SimpleNamespace(workflow_run_archives=service),
)
monkeypatch.setattr(workflow_run_archive, "dump_response", lambda _model, value: value)
app = Flask(__name__)
with app.test_request_context():
_admission_injector(WorkflowRunArchivesApi.get)(None)
context = service.list_archives.call_args.args[0]
assert context.account_id == "account-1"
assert context.active_workspace_id == "tenant-1"
def test_workflow_run_archive_endpoints_delegate_to_application_service(
monkeypatch: pytest.MonkeyPatch,
) -> None:
service = MagicMock()
archives = object()
task = object()
service.list_archives.return_value = archives
service.create_download.return_value = task
service.get_download.return_value = task
monkeypatch.setattr(
workflow_run_archive,
"application_services",
lambda: SimpleNamespace(workflow_run_archives=service),
)
dump_response = MagicMock(side_effect=lambda _model, value: value)
monkeypatch.setattr(workflow_run_archive, "dump_response", dump_response)
assert _original(WorkflowRunArchivesApi.get)(None, _CONTEXT) is archives
assert _original(WorkflowRunArchiveDownloadsApi.post)(
None,
WorkflowRunArchiveDownloadPayload(year=2025, month=3),
_CONTEXT,
) == (task, 202)
assert _original(WorkflowRunArchiveDownloadApi.get)(None, _CONTEXT, "download-1") is task
service.list_archives.assert_called_once_with(_CONTEXT)
service.create_download.assert_called_once_with(_CONTEXT, year=2025, month=3)
service.get_download.assert_called_once_with(_CONTEXT, download_id="download-1")
@pytest.mark.parametrize(
"method",
("method", "service_method", "service_error", "http_error", "args"),
[
WorkflowRunArchivesApi.get,
WorkflowRunArchiveDownloadsApi.post,
WorkflowRunArchiveDownloadApi.get,
WorkflowRunArchiveDownloadFileApi.get,
(
WorkflowRunArchiveDownloadsApi.post,
"create_download",
WorkflowRunArchiveNotFoundError("archive missing"),
NotFound,
(WorkflowRunArchiveDownloadPayload(year=2025, month=3), _CONTEXT),
),
(
WorkflowRunArchiveDownloadApi.get,
"get_download",
WorkflowRunArchiveDownloadTaskNotFoundError("task missing"),
NotFound,
(_CONTEXT, "download-1"),
),
(
WorkflowRunArchiveDownloadFileApi.get,
"get_download_url",
WorkflowRunArchiveDownloadTaskNotFoundError("task missing"),
NotFound,
(_CONTEXT, "download-1"),
),
(
WorkflowRunArchiveDownloadFileApi.get,
"get_download_url",
WorkflowRunArchiveDownloadNotReadyError("task pending"),
Conflict,
(_CONTEXT, "download-1"),
),
],
)
def test_workflow_run_archive_endpoints_require_cloud_paid_plan(method) -> None:
decorator_names = set()
while hasattr(method, "__wrapped__"):
decorator_names.add(method.__code__.co_qualname.partition(".<locals>")[0])
method = method.__wrapped__
def test_workflow_run_archive_endpoints_translate_service_errors(
monkeypatch: pytest.MonkeyPatch,
method,
service_method: str,
service_error: Exception,
http_error: type[Exception],
args: tuple[object, ...],
) -> None:
service = MagicMock()
service.configure_mock(**{f"{service_method}.side_effect": service_error})
monkeypatch.setattr(
workflow_run_archive,
"application_services",
lambda: SimpleNamespace(workflow_run_archives=service),
)
assert {
"only_edition_cloud",
"cloud_edition_billing_paid_plan_required",
} <= decorator_names
assert "rbac_permission_required" not in decorator_names
with pytest.raises(http_error):
_original(method)(None, *args)
def test_workflow_run_archive_file_redirects_to_signed_url(monkeypatch: pytest.MonkeyPatch) -> None:
service = MagicMock()
service.get_download_url.return_value = "https://storage.example.com/archive.zip"
monkeypatch.setattr(
workflow_run_archive,
"application_services",
lambda: SimpleNamespace(workflow_run_archives=service),
)
app = Flask(__name__)
with app.test_request_context():
response = _original(WorkflowRunArchiveDownloadFileApi.get)(None, _CONTEXT, "download-1")
assert response.status_code == 302
assert response.location == "https://storage.example.com/archive.zip"

25
api/tests/unit_tests/extensions/test_ext_application_services.py
View File

@@ -22,6 +22,7 @@ from repositories.account_activation_repository import SQLAlchemyAccountActivati
from repositories.account_integration_repository import SQLAlchemyAccountIntegrationRepository
from repositories.account_repository import SQLAlchemyAccountRepository
from repositories.app_site_command_repository import AppSiteCommandRepository
from repositories.workflow_run_archive_repository import WorkflowRunArchiveBundleQueryRepository
from services import recommended_app_catalog_gateway
from services.account_activation_adapters import (
BillingAccountActivationEligibility,
@@ -39,6 +40,8 @@ from services.enterprise.enterprise_service import WebAppSettings
from services.errors.enterprise import EnterpriseAPIError, EnterpriseAPINotFoundError
from services.init_validation_service import InvalidInitializationPasswordError
from services.partner_tenant_binding_service import PartnerTenantBindingService
from services.retention.workflow_run.archive_download_task_cache import WorkflowRunArchiveDownloadTaskCache
from services.retention.workflow_run.archive_log_service import WorkflowRunArchiveService
from services.tag_application_service import TagApplicationService
from services.webapp_access_query_service import WebAppAccessUnavailableError
from services.workflow_statistic_query_service import WorkflowStatisticQueryService
@@ -182,6 +185,28 @@ def test_build_application_services_wires_tag_boundary(
assert isinstance(services.tags, TagApplicationService)
def test_build_application_services_wires_workflow_run_archives(
sqlite_session_factory: sessionmaker[Session],
) -> None:
redis = MagicMock(spec=RedisClientWrapper)
services = ext_application_services.build_application_services(
database_client=sqlite_session_factory,
deployment_edition=DeploymentEdition.COMMUNITY,
initialization_password="",
redis=redis,
)
workflow_run_archives = services.workflow_run_archives
assert isinstance(workflow_run_archives, WorkflowRunArchiveService)
assert isinstance(workflow_run_archives._bundles, WorkflowRunArchiveBundleQueryRepository)
assert workflow_run_archives._bundles._session_factory is sqlite_session_factory
assert isinstance(workflow_run_archives._tasks, WorkflowRunArchiveDownloadTaskCache)
assert workflow_run_archives._tasks._redis is redis
assert workflow_run_archives._dispatcher is ext_application_services.dispatch_workflow_run_archive_download_task
assert workflow_run_archives._sign_download_url is ext_application_services.sign_workflow_run_archive_download_url
def test_build_application_services_wires_app_site_boundary(
sqlite_session_factory: sessionmaker[Session],
) -> None:

138
api/tests/unit_tests/repositories/test_workflow_run_archive_repository.py Normal file
View File

@@ -0,0 +1,138 @@
import datetime
from sqlalchemy.orm import Session, sessionmaker
from models.workflow import WorkflowRunArchiveBundle
from repositories.workflow_run_archive_repository import WorkflowRunArchiveBundleQueryRepository
from services.retention.workflow_run.archive_log_service import WorkflowRunArchiveBundleRecord
def _bundle(
*,
tenant_id: str,
year: int,
month: int,
shard: str,
bundle_id: str,
workflow_run_count: int = 1,
row_count: int = 9,
archive_bytes: int = 1024,
archived_at: datetime.datetime | None = None,
) -> WorkflowRunArchiveBundle:
return WorkflowRunArchiveBundle(
tenant_id=tenant_id,
year=year,
month=month,
shard=shard,
bundle_id=bundle_id,
workflow_run_count=workflow_run_count,
row_count=row_count,
archive_bytes=archive_bytes,
archived_at=archived_at or datetime.datetime(2026, 6, 25, 8, 0),
)
def test_list_for_tenant_returns_ordered_pure_records(
sqlite_session_factory: sessionmaker[Session],
) -> None:
latest = datetime.datetime(2026, 6, 26, 8, 0)
previous = datetime.datetime(2026, 6, 25, 8, 0)
with sqlite_session_factory.begin() as session:
session.add_all(
[
_bundle(
tenant_id="tenant-1",
year=2025,
month=2,
shard="00-of-01",
bundle_id="bundle-c",
),
_bundle(
tenant_id="tenant-1",
year=2025,
month=3,
shard="01-of-02",
bundle_id="bundle-b",
workflow_run_count=60,
row_count=540,
archive_bytes=3072,
archived_at=latest,
),
_bundle(
tenant_id="tenant-1",
year=2025,
month=3,
shard="00-of-02",
bundle_id="bundle-a",
workflow_run_count=40,
row_count=360,
archive_bytes=1024,
archived_at=previous,
),
_bundle(
tenant_id="tenant-2",
year=2026,
month=1,
shard="00-of-01",
bundle_id="other-tenant-bundle",
),
]
)
records = WorkflowRunArchiveBundleQueryRepository(session_factory=sqlite_session_factory).list_for_tenant(
"tenant-1"
)
assert records == (
WorkflowRunArchiveBundleRecord(2025, 3, "00-of-02", "bundle-a", 40, 360, 1024, previous),
WorkflowRunArchiveBundleRecord(2025, 3, "01-of-02", "bundle-b", 60, 540, 3072, latest),
WorkflowRunArchiveBundleRecord(2025, 2, "00-of-01", "bundle-c", 1, 9, 1024, previous),
)
def test_list_for_tenant_month_filters_and_orders_bundles(
sqlite_session_factory: sessionmaker[Session],
) -> None:
archived_at = datetime.datetime(2026, 6, 25, 8, 0)
with sqlite_session_factory.begin() as session:
session.add_all(
[
_bundle(
tenant_id="tenant-1",
year=2025,
month=3,
shard="00-of-01",
bundle_id="bundle-b",
),
_bundle(
tenant_id="tenant-1",
year=2025,
month=3,
shard="00-of-01",
bundle_id="bundle-a",
),
_bundle(
tenant_id="tenant-1",
year=2025,
month=2,
shard="00-of-01",
bundle_id="previous-month",
),
_bundle(
tenant_id="tenant-2",
year=2025,
month=3,
shard="00-of-01",
bundle_id="other-tenant-bundle",
),
]
)
records = WorkflowRunArchiveBundleQueryRepository(session_factory=sqlite_session_factory).list_for_tenant_month(
"tenant-1", year=2025, month=3
)
assert records == (
WorkflowRunArchiveBundleRecord(2025, 3, "00-of-01", "bundle-a", 1, 9, 1024, archived_at),
WorkflowRunArchiveBundleRecord(2025, 3, "00-of-01", "bundle-b", 1, 9, 1024, archived_at),
)

100
api/tests/unit_tests/services/retention/workflow_run/test_archive_download_adapters.py Normal file
View File

@@ -0,0 +1,100 @@
import sys
from types import ModuleType, SimpleNamespace
from unittest.mock import MagicMock
import pytest
from services.retention.workflow_run import archive_download_adapters
from services.retention.workflow_run.archive_download_adapters import (
dispatch_workflow_run_archive_download_task,
sign_workflow_run_archive_download_url,
)
from services.retention.workflow_run.archive_download_task import (
WorkflowRunArchiveDownloadTask,
build_pending_archive_download_task,
)
def _task(*, celery_task_id: str | None = "celery-task-1") -> WorkflowRunArchiveDownloadTask:
return build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
bundle_ids=["bundle-1"],
bundle_refs=[("00-of-01", "bundle-1")],
archive_bytes=1024,
download_id="download-1",
).model_copy(update={"celery_task_id": celery_task_id})
def _patch_archive_download_task(monkeypatch: pytest.MonkeyPatch, *, apply_async: MagicMock) -> None:
task_module = ModuleType("tasks.workflow_run_archive_download_tasks")
task_module.__dict__["prepare_workflow_run_archive_download_task"] = SimpleNamespace(apply_async=apply_async)
monkeypatch.setitem(sys.modules, task_module.__name__, task_module)
def test_dispatch_workflow_run_archive_download_task_enqueues_claimed_task(
monkeypatch: pytest.MonkeyPatch,
) -> None:
apply_async = MagicMock()
_patch_archive_download_task(monkeypatch, apply_async=apply_async)
result = dispatch_workflow_run_archive_download_task(_task())
assert result is None
apply_async.assert_called_once_with(
args=("tenant-1", "download-1"),
task_id="celery-task-1",
)
def test_dispatch_workflow_run_archive_download_task_requires_claim_id(
monkeypatch: pytest.MonkeyPatch,
) -> None:
apply_async = MagicMock()
_patch_archive_download_task(monkeypatch, apply_async=apply_async)
with pytest.raises(ValueError, match="celery_task_id is required before dispatch"):
dispatch_workflow_run_archive_download_task(_task(celery_task_id=None))
apply_async.assert_not_called()
def test_dispatch_workflow_run_archive_download_task_propagates_enqueue_error(
monkeypatch: pytest.MonkeyPatch,
) -> None:
failure = RuntimeError("broker unavailable")
apply_async = MagicMock(side_effect=failure)
_patch_archive_download_task(monkeypatch, apply_async=apply_async)
with pytest.raises(RuntimeError) as raised:
dispatch_workflow_run_archive_download_task(_task())
assert raised.value is failure
def test_sign_workflow_run_archive_download_url_resolves_export_storage_lazily(
monkeypatch: pytest.MonkeyPatch,
) -> None:
storage = MagicMock()
storage.generate_presigned_url.return_value = "https://download.example/archive.zip"
get_export_storage = MagicMock(return_value=storage)
monkeypatch.setattr(archive_download_adapters, "get_export_storage", get_export_storage)
assert get_export_storage.call_count == 0
result = sign_workflow_run_archive_download_url(
"downloads/archive.zip",
expires_in=900,
filename="workflow-run-logs-2025-03.zip",
)
assert result == "https://download.example/archive.zip"
get_export_storage.assert_called_once_with()
storage.generate_presigned_url.assert_called_once_with(
"downloads/archive.zip",
expires_in=900,
filename="workflow-run-logs-2025-03.zip",
content_type="application/zip",
)

86
api/tests/unit_tests/services/retention/workflow_run/test_archive_download_preparation.py
View File

@@ -9,22 +9,20 @@ from unittest.mock import MagicMock
import pyarrow as pa
import pyarrow.parquet as pq
from sqlalchemy.orm import Session, sessionmaker
from libs.archive_storage import ArchiveStorage
from models.workflow import WorkflowRunArchiveBundle
from services.retention.workflow_run.archive_bundle_index import ARCHIVE_BUNDLE_ROOT_PREFIX, ArchiveBundleManifest
from services.retention.workflow_run.archive_download_preparation import (
WorkflowRunArchiveDownloadPreparer,
build_archive_download_storage_key,
)
from services.retention.workflow_run.archive_download_task_cache import (
from services.retention.workflow_run.archive_download_task import (
WorkflowRunArchiveDownloadStatus,
WorkflowRunArchiveDownloadTask,
WorkflowRunArchiveDownloadTaskCache,
build_archive_download_id,
build_pending_archive_download_task,
)
from services.retention.workflow_run.archive_log_service import WorkflowRunArchiveBundleRecord
from services.retention.workflow_run.constants import ARCHIVE_BUNDLE_FORMAT, ARCHIVE_BUNDLE_SCHEMA_VERSION
TENANT_ID = "1251fe32-c0c7-4fe2-a7bd-a8105267faf5"
@@ -74,6 +72,25 @@ class FakeTaskCache:
self.saved_tasks.append(task)
class FakeBundleQuery:
def __init__(self, records: tuple[WorkflowRunArchiveBundleRecord, ...]) -> None:
self.records = records
def list_for_tenant(self, tenant_id: str) -> tuple[WorkflowRunArchiveBundleRecord, ...]:
assert tenant_id == TENANT_ID
return self.records
def list_for_tenant_month(
self,
tenant_id: str,
*,
year: int,
month: int,
) -> tuple[WorkflowRunArchiveBundleRecord, ...]:
assert tenant_id == TENANT_ID
return tuple(record for record in self.records if record.year == year and record.month == month)
def _object_prefix(bundle_id: str = BUNDLE_ID) -> str:
return (
f"{ARCHIVE_BUNDLE_ROOT_PREFIX}tenant_prefix=1/tenant_id={TENANT_ID}/"
@@ -81,9 +98,8 @@ def _object_prefix(bundle_id: str = BUNDLE_ID) -> str:
)
def _bundle(bundle_id: str = BUNDLE_ID) -> WorkflowRunArchiveBundle:
return WorkflowRunArchiveBundle(
tenant_id=TENANT_ID,
def _bundle(bundle_id: str = BUNDLE_ID) -> WorkflowRunArchiveBundleRecord:
return WorkflowRunArchiveBundleRecord(
year=2025,
month=3,
shard=SHARD,
@@ -154,21 +170,17 @@ def _preparer(
archive_storage: FakeArchiveStorage | None = None,
download_storage: FakeArchiveStorage | None = None,
cache: FakeTaskCache,
session_factory: sessionmaker[Session],
bundles: list[WorkflowRunArchiveBundle] | None = None,
bundles: list[WorkflowRunArchiveBundleRecord] | None = None,
) -> WorkflowRunArchiveDownloadPreparer:
source_storage = archive_storage or storage
target_storage = download_storage or storage
assert source_storage is not None
assert target_storage is not None
with session_factory() as session:
session.add_all([_bundle()] if bundles is None else bundles)
session.commit()
return WorkflowRunArchiveDownloadPreparer(
archive_storage=cast(ArchiveStorage, source_storage),
download_storage=cast(ArchiveStorage, target_storage),
cache=cast(WorkflowRunArchiveDownloadTaskCache, cache),
session_factory=cast(sessionmaker[Session], session_factory),
archive_storage_provider=lambda: cast(ArchiveStorage, source_storage),
bundles=FakeBundleQuery(tuple([_bundle()] if bundles is None else bundles)),
cache=cache,
download_storage_provider=lambda: cast(ArchiveStorage, target_storage),
)
@@ -178,9 +190,7 @@ def _parquet_bytes(records: list[dict[str, object]]) -> bytes:
return buffer.getvalue()
def test_prepare_workflow_run_archive_download_builds_csv_zip_and_marks_ready(
sqlite_session_factory: sessionmaker[Session],
) -> None:
def test_prepare_workflow_run_archive_download_builds_csv_zip_and_marks_ready() -> None:
bundle_refs = [(SHARD, "bundle-a"), (SHARD, "bundle-b")]
task = _task(bundle_refs)
first_bundle_payloads = {
@@ -217,7 +227,6 @@ def test_prepare_workflow_run_archive_download_builds_csv_zip_and_marks_ready(
archive_storage=archive_storage,
download_storage=download_storage,
cache=cache,
session_factory=sqlite_session_factory,
bundles=[_bundle("bundle-a"), _bundle("bundle-b")],
)
@@ -244,9 +253,7 @@ def test_prepare_workflow_run_archive_download_builds_csv_zip_and_marks_ready(
assert '"run-b","failed","safe"' in workflow_runs_csv
def test_prepare_workflow_run_archive_download_marks_failed_on_checksum_mismatch(
sqlite_session_factory: sessionmaker[Session],
) -> None:
def test_prepare_workflow_run_archive_download_marks_failed_on_checksum_mismatch() -> None:
task = _task()
table_payloads = {"workflow_runs": _parquet_bytes([{"id": "run-a", "status": "succeeded"}])}
manifest_data = json.loads(_manifest_bytes(table_payloads).decode("utf-8"))
@@ -258,7 +265,7 @@ def test_prepare_workflow_run_archive_download_marks_failed_on_checksum_mismatch
}
)
cache = FakeTaskCache(task)
preparer = _preparer(storage=storage, cache=cache, session_factory=sqlite_session_factory)
preparer = _preparer(storage=storage, cache=cache)
result = preparer.prepare(tenant_id=TENANT_ID, download_id=task.download_id)
@@ -268,16 +275,33 @@ def test_prepare_workflow_run_archive_download_marks_failed_on_checksum_mismatch
assert storage.put_objects == {}
def test_prepare_workflow_run_archive_download_skips_duplicate_worker(
sqlite_session_factory: sessionmaker[Session],
) -> None:
def test_prepare_workflow_run_archive_download_marks_failed_when_storage_is_unavailable() -> None:
task = _task()
cache = FakeTaskCache(task)
archive_storage_provider = MagicMock(side_effect=RuntimeError("archive storage unavailable"))
download_storage_provider = MagicMock()
preparer = WorkflowRunArchiveDownloadPreparer(
archive_storage_provider=archive_storage_provider,
bundles=FakeBundleQuery((_bundle(),)),
cache=cache,
download_storage_provider=download_storage_provider,
)
result = preparer.prepare(tenant_id=TENANT_ID, download_id=task.download_id)
assert result is not None
assert result.status == WorkflowRunArchiveDownloadStatus.FAILED
assert result.error == "archive storage unavailable"
download_storage_provider.assert_not_called()
def test_prepare_workflow_run_archive_download_skips_duplicate_worker() -> None:
task = _task().model_copy(update={"celery_task_id": "celery-task-1"})
storage = FakeArchiveStorage({})
cache = FakeTaskCache(task)
preparer = _preparer(
storage=storage,
cache=cache,
session_factory=sqlite_session_factory,
bundles=[],
)
nested_results: list[WorkflowRunArchiveDownloadTask | None] = []
@@ -298,15 +322,13 @@ def test_prepare_workflow_run_archive_download_skips_duplicate_worker(
preparer._build_zip_payload.assert_called_once()
def test_failed_worker_cannot_overwrite_ready_task(
sqlite_session_factory: sessionmaker[Session],
) -> None:
def test_failed_worker_cannot_overwrite_ready_task() -> None:
processing_task = _task().model_copy(
update={"status": WorkflowRunArchiveDownloadStatus.PROCESSING, "celery_task_id": "celery-task-1"}
)
ready_task = processing_task.model_copy(update={"status": WorkflowRunArchiveDownloadStatus.READY})
cache = FakeTaskCache(ready_task)
preparer = _preparer(storage=FakeArchiveStorage({}), cache=cache, session_factory=sqlite_session_factory)
preparer = _preparer(storage=FakeArchiveStorage({}), cache=cache)
result = preparer._mark_failed(processing_task, error="late failure")

16
api/tests/unit_tests/services/retention/workflow_run/test_archive_download_task_cache.py
View File

@@ -3,15 +3,17 @@ from contextlib import nullcontext
import pytest
from services.retention.workflow_run.archive_download_task_cache import (
from services.retention.workflow_run.archive_download_task import (
ARCHIVE_DOWNLOAD_FORMAT_VERSION,
ARCHIVE_DOWNLOAD_TASK_LOCK_TIMEOUT_SECONDS,
WorkflowRunArchiveDownloadStatus,
WorkflowRunArchiveDownloadTask,
WorkflowRunArchiveDownloadTaskCache,
build_archive_download_id,
build_pending_archive_download_task,
)
from services.retention.workflow_run.archive_download_task_cache import (
ARCHIVE_DOWNLOAD_TASK_LOCK_TIMEOUT_SECONDS,
WorkflowRunArchiveDownloadTaskCache,
)
class FakeRedis:
@@ -119,7 +121,7 @@ def test_build_archive_download_id_rejects_empty_bundle_refs() -> None:
def test_archive_download_task_cache_round_trips_with_ttl() -> None:
redis = FakeRedis()
cache = WorkflowRunArchiveDownloadTaskCache(redis)
cache = WorkflowRunArchiveDownloadTaskCache(redis=redis)
task = build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
@@ -142,7 +144,7 @@ def test_archive_download_task_cache_round_trips_with_ttl() -> None:
def test_archive_download_task_cache_uses_per_download_lock() -> None:
redis = FakeRedis()
cache = WorkflowRunArchiveDownloadTaskCache(redis)
cache = WorkflowRunArchiveDownloadTaskCache(redis=redis)
with cache.lock(tenant_id="tenant-1", download_id="download-1"):
pass
@@ -158,7 +160,7 @@ def test_archive_download_task_cache_uses_per_download_lock() -> None:
def test_archive_download_task_cache_delete_removes_entry() -> None:
redis = FakeRedis()
cache = WorkflowRunArchiveDownloadTaskCache(redis)
cache = WorkflowRunArchiveDownloadTaskCache(redis=redis)
task = WorkflowRunArchiveDownloadTask(
download_id="download-1",
tenant_id="tenant-1",
@@ -183,7 +185,7 @@ def test_archive_download_task_cache_delete_removes_entry() -> None:
def test_archive_download_task_cache_ignores_malformed_json() -> None:
redis = FakeRedis()
cache = WorkflowRunArchiveDownloadTaskCache(redis)
cache = WorkflowRunArchiveDownloadTaskCache(redis=redis)
redis.setex("workflow_run_archive_download:tenant-1:download-1", 3600, "{")
assert cache.get(tenant_id="tenant-1", download_id="download-1") is None

522
api/tests/unit_tests/services/retention/workflow_run/test_archive_log_service.py
View File

@@ -1,54 +1,75 @@
import datetime
from contextlib import nullcontext
from typing import cast
import pytest
from sqlalchemy.orm import Session
from models.workflow import WorkflowRunArchiveBundle
from services.retention.workflow_run.archive_download_task_cache import (
from machinery.context import RequestContext
from services.retention.workflow_run.archive_download_task import (
WorkflowRunArchiveDownloadStatus,
WorkflowRunArchiveDownloadTask,
WorkflowRunArchiveDownloadTaskCache,
build_archive_download_id,
build_pending_archive_download_task,
)
from services.retention.workflow_run.archive_log_service import (
ArchiveDownloadTaskDispatcher,
WorkflowRunArchiveBundleRecord,
WorkflowRunArchiveDownloadNotReadyError,
WorkflowRunArchiveDownloadTaskDispatcher,
WorkflowRunArchiveDownloadTaskNotFoundError,
WorkflowRunArchiveNotFoundError,
create_workflow_run_archive_download_task,
get_ready_workflow_run_archive_download_task,
list_workflow_run_archives,
WorkflowRunArchiveService,
)
_CONTEXT = RequestContext(
request_id="request-1",
trace_id="trace-1",
account_id="account-1",
active_workspace_id="tenant-1",
)
class FakeTaskCache:
saved_task: WorkflowRunArchiveDownloadTask | None
existing_task: WorkflowRunArchiveDownloadTask | None
tasks_by_download_id: dict[str, WorkflowRunArchiveDownloadTask]
class FakeBundleQuery:
def __init__(self, records: tuple[WorkflowRunArchiveBundleRecord, ...] = ()) -> None:
self.records = records
def list_for_tenant(self, tenant_id: str) -> tuple[WorkflowRunArchiveBundleRecord, ...]:
assert tenant_id == "tenant-1"
return self.records
def list_for_tenant_month(
self,
tenant_id: str,
*,
year: int,
month: int,
) -> tuple[WorkflowRunArchiveBundleRecord, ...]:
assert tenant_id == "tenant-1"
return tuple(record for record in self.records if record.year == year and record.month == month)
class FakeTaskStore:
def __init__(
self,
tasks: dict[str, WorkflowRunArchiveDownloadTask] | None = None,
*,
existing_task: WorkflowRunArchiveDownloadTask | None = None,
tasks_by_download_id: dict[str, WorkflowRunArchiveDownloadTask] | None = None,
fail_reads: bool = False,
) -> None:
self.saved_task = None
self.existing_task = existing_task
self.tasks_by_download_id = tasks_by_download_id or {}
self.tasks = tasks or {}
self.saved_tasks: list[WorkflowRunArchiveDownloadTask] = []
self.fail_reads = fail_reads
self.lock_calls: list[tuple[str, str]] = []
def lock(self, *, tenant_id: str, download_id: str):
self.lock_calls.append((tenant_id, download_id))
return nullcontext()
def get(self, *, tenant_id: str, download_id: str) -> WorkflowRunArchiveDownloadTask | None:
if self.tasks_by_download_id:
return self.tasks_by_download_id.get(download_id)
return self.existing_task
if self.fail_reads:
raise RuntimeError("cache unavailable")
return self.tasks.get(download_id)
def save(self, task: WorkflowRunArchiveDownloadTask) -> None:
self.saved_task = task
self.existing_task = task
self.saved_tasks.append(task)
self.tasks[task.download_id] = task
def _bundle(
@@ -61,9 +82,8 @@ def _bundle(
workflow_run_count: int = 1,
row_count: int = 9,
archived_at: datetime.datetime | None = None,
) -> WorkflowRunArchiveBundle:
return WorkflowRunArchiveBundle(
tenant_id="tenant-1",
) -> WorkflowRunArchiveBundleRecord:
return WorkflowRunArchiveBundleRecord(
year=year,
month=month,
shard=shard,
@@ -75,20 +95,68 @@ def _bundle(
)
def _fake_dispatcher(dispatched_tasks: list[WorkflowRunArchiveDownloadTask]) -> ArchiveDownloadTaskDispatcher:
def dispatch(
task: WorkflowRunArchiveDownloadTask,
cache: WorkflowRunArchiveDownloadTaskCache,
) -> WorkflowRunArchiveDownloadTask:
dispatched_tasks.append(task)
return task
def _service(
*,
records: tuple[WorkflowRunArchiveBundleRecord, ...] = (),
tasks: FakeTaskStore | None = None,
dispatched_tasks: list[WorkflowRunArchiveDownloadTask] | None = None,
signed_urls: list[tuple[str, int, str]] | None = None,
dispatcher: WorkflowRunArchiveDownloadTaskDispatcher | None = None,
) -> WorkflowRunArchiveService:
task_store = tasks or FakeTaskStore()
return dispatch
def dispatch(task: WorkflowRunArchiveDownloadTask) -> None:
if dispatched_tasks is not None:
dispatched_tasks.append(task)
def sign(storage_key: str, *, expires_in: int, filename: str) -> str:
if signed_urls is not None:
signed_urls.append((storage_key, expires_in, filename))
return "https://storage.example.com/archive.zip"
return WorkflowRunArchiveService(
bundles=FakeBundleQuery(records),
tasks=task_store,
dispatcher=dispatcher or dispatch,
sign_download_url=sign,
)
def test_list_workflow_run_archives_aggregates_month_rows(sqlite_session: Session) -> None:
def test_list_archives_aggregates_month_rows_and_includes_cached_task() -> None:
latest = datetime.datetime(2026, 6, 25, 8, 0)
previous = datetime.datetime(2026, 6, 24, 8, 0)
records = (
_bundle(
year=2025,
month=3,
shard="00-of-01",
bundle_id="bundle-a",
workflow_run_count=40,
row_count=360,
archive_bytes=1024,
archived_at=previous,
),
_bundle(
year=2025,
month=3,
shard="00-of-01",
bundle_id="bundle-b",
workflow_run_count=60,
row_count=540,
archive_bytes=3072,
archived_at=latest,
),
_bundle(
year=2025,
month=2,
shard="00-of-01",
bundle_id="bundle-c",
workflow_run_count=20,
row_count=180,
archive_bytes=1024,
archived_at=previous,
),
)
march_download_id = build_archive_download_id(
tenant_id="tenant-1",
year=2025,
@@ -104,54 +172,9 @@ def test_list_workflow_run_archives_aggregates_month_rows(sqlite_session: Sessio
bundle_refs=[("00-of-01", "bundle-a"), ("00-of-01", "bundle-b")],
archive_bytes=4096,
download_id=march_download_id,
).model_copy(
update={
"status": WorkflowRunArchiveDownloadStatus.READY,
"file_name": "workflow-run-logs-2025-03.zip",
"storage_key": "workflow-run-archive-downloads/tenant-1/2025/03/download.zip",
"file_size_bytes": 8192,
}
)
cache = FakeTaskCache(tasks_by_download_id={march_download_id: ready_task})
sqlite_session.add_all(
[
_bundle(
year=2025,
month=3,
shard="00-of-01",
bundle_id="bundle-a",
workflow_run_count=40,
row_count=360,
archive_bytes=1024,
archived_at=previous,
),
_bundle(
year=2025,
month=3,
shard="00-of-01",
bundle_id="bundle-b",
workflow_run_count=60,
row_count=540,
archive_bytes=3072,
archived_at=latest,
),
_bundle(
year=2025,
month=2,
shard="00-of-01",
bundle_id="bundle-c",
workflow_run_count=20,
row_count=180,
archive_bytes=1024,
archived_at=previous,
),
]
)
sqlite_session.flush()
).model_copy(update={"status": WorkflowRunArchiveDownloadStatus.READY})
result = list_workflow_run_archives(
sqlite_session, "tenant-1", cache=cast(WorkflowRunArchiveDownloadTaskCache, cache)
)
result = _service(records=records, tasks=FakeTaskStore({march_download_id: ready_task})).list_archives(_CONTEXT)
assert result.summary.archived_month_count == 2
assert result.summary.workflow_run_count == 120
@@ -166,25 +189,27 @@ def test_list_workflow_run_archives_aggregates_month_rows(sqlite_session: Sessio
assert result.months[1].download_task is None
def test_create_workflow_run_archive_download_task_creates_stable_pending_task(sqlite_session: Session) -> None:
sqlite_session.add_all(
[
_bundle(shard="01-of-02", bundle_id="bundle-b", archive_bytes=2048),
_bundle(shard="00-of-02", bundle_id="bundle-a", archive_bytes=1024),
]
def test_list_archives_tolerates_unavailable_task_cache() -> None:
result = _service(
records=(_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024),),
tasks=FakeTaskStore(fail_reads=True),
).list_archives(_CONTEXT)
assert result.months[0].download_task is None
def test_create_download_creates_stable_pending_task() -> None:
records = (
_bundle(shard="00-of-02", bundle_id="bundle-a", archive_bytes=1024),
_bundle(shard="01-of-02", bundle_id="bundle-b", archive_bytes=2048),
)
sqlite_session.flush()
cache = FakeTaskCache()
task_store = FakeTaskStore()
dispatched_tasks: list[WorkflowRunArchiveDownloadTask] = []
task = create_workflow_run_archive_download_task(
sqlite_session,
tenant_id="tenant-1",
requested_by="account-1",
task = _service(records=records, tasks=task_store, dispatched_tasks=dispatched_tasks).create_download(
_CONTEXT,
year=2025,
month=3,
cache=cast(WorkflowRunArchiveDownloadTaskCache, cache),
dispatcher=_fake_dispatcher(dispatched_tasks),
)
assert task.download_id == build_archive_download_id(
@@ -200,156 +225,167 @@ def test_create_workflow_run_archive_download_task_creates_stable_pending_task(s
("01-of-02", "bundle-b"),
]
assert task.archive_bytes == 3072
assert cache.saved_task == dispatched_tasks[0]
assert task_store.saved_tasks == dispatched_tasks
assert task.celery_task_id is not None
def test_create_workflow_run_archive_download_task_returns_existing_task_when_cache_key_exists(
sqlite_session: Session,
) -> None:
sqlite_session.add(_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024))
sqlite_session.flush()
existing_task = build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
bundle_ids=["bundle-a"],
archive_bytes=1024,
download_id="existing-download",
).model_copy(update={"celery_task_id": "celery-task-1"})
cache = FakeTaskCache(existing_task=existing_task)
dispatched_tasks: list[WorkflowRunArchiveDownloadTask] = []
task = create_workflow_run_archive_download_task(
sqlite_session,
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
cache=cast(WorkflowRunArchiveDownloadTaskCache, cache),
dispatcher=_fake_dispatcher(dispatched_tasks),
)
assert task == existing_task
assert cache.saved_task is None
assert dispatched_tasks == []
def test_create_workflow_run_archive_download_task_retries_failed_cached_task(
sqlite_session: Session,
) -> None:
sqlite_session.add(_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024))
sqlite_session.flush()
existing_task = build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
bundle_ids=["bundle-a"],
bundle_refs=[("00-of-01", "bundle-a")],
archive_bytes=1024,
download_id=build_archive_download_id(
tenant_id="tenant-1",
year=2025,
month=3,
bundle_refs=[("00-of-01", "bundle-a")],
),
).model_copy(update={"status": WorkflowRunArchiveDownloadStatus.FAILED, "error": "failed"})
cache = FakeTaskCache(existing_task=existing_task)
dispatched_tasks: list[WorkflowRunArchiveDownloadTask] = []
task = create_workflow_run_archive_download_task(
sqlite_session,
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
cache=cast(WorkflowRunArchiveDownloadTaskCache, cache),
dispatcher=_fake_dispatcher(dispatched_tasks),
)
assert task.status == WorkflowRunArchiveDownloadStatus.PENDING
assert task.error is None
assert task.celery_task_id is not None
assert cache.saved_task == dispatched_tasks[0]
@pytest.mark.parametrize("retry_failed_task", [False, True])
def test_create_workflow_run_archive_download_task_claims_dispatch_once(
retry_failed_task: bool, sqlite_session: Session
) -> None:
sqlite_session.add(_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024))
sqlite_session.flush()
def test_create_download_returns_existing_queued_task() -> None:
records = (_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024),)
download_id = build_archive_download_id(
tenant_id="tenant-1",
year=2025,
month=3,
bundle_refs=[("00-of-01", "bundle-a")],
)
existing_task = None
if retry_failed_task:
existing_task = build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
bundle_ids=["bundle-a"],
bundle_refs=[("00-of-01", "bundle-a")],
archive_bytes=1024,
download_id=download_id,
).model_copy(update={"status": WorkflowRunArchiveDownloadStatus.FAILED})
cache = FakeTaskCache(existing_task=existing_task)
dispatched_tasks: list[WorkflowRunArchiveDownloadTask] = []
concurrent_results: list[WorkflowRunArchiveDownloadTask] = []
def dispatch(
task: WorkflowRunArchiveDownloadTask,
task_cache: WorkflowRunArchiveDownloadTaskCache,
) -> WorkflowRunArchiveDownloadTask:
dispatched_tasks.append(task)
concurrent_results.append(
create_workflow_run_archive_download_task(
sqlite_session,
tenant_id="tenant-1",
requested_by="account-2",
year=2025,
month=3,
cache=task_cache,
dispatcher=dispatch,
)
)
return task
result = create_workflow_run_archive_download_task(
sqlite_session,
existing = build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
cache=cast(WorkflowRunArchiveDownloadTaskCache, cache),
dispatcher=dispatch,
bundle_ids=["bundle-a"],
archive_bytes=1024,
download_id=download_id,
).model_copy(update={"celery_task_id": "celery-task-1"})
task_store = FakeTaskStore({download_id: existing})
dispatched_tasks: list[WorkflowRunArchiveDownloadTask] = []
result = _service(records=records, tasks=task_store, dispatched_tasks=dispatched_tasks).create_download(
_CONTEXT,
year=2025,
month=3,
)
assert result == existing
assert task_store.saved_tasks == []
assert dispatched_tasks == []
def test_create_download_retries_failed_task() -> None:
records = (_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024),)
download_id = build_archive_download_id(
tenant_id="tenant-1",
year=2025,
month=3,
bundle_refs=[("00-of-01", "bundle-a")],
)
failed = build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
bundle_ids=["bundle-a"],
bundle_refs=[("00-of-01", "bundle-a")],
archive_bytes=1024,
download_id=download_id,
).model_copy(update={"status": WorkflowRunArchiveDownloadStatus.FAILED, "error": "failed"})
task_store = FakeTaskStore({download_id: failed})
dispatched_tasks: list[WorkflowRunArchiveDownloadTask] = []
task = _service(records=records, tasks=task_store, dispatched_tasks=dispatched_tasks).create_download(
_CONTEXT,
year=2025,
month=3,
)
assert task.status == WorkflowRunArchiveDownloadStatus.PENDING
assert task.error is None
assert task.celery_task_id is not None
assert task_store.saved_tasks == dispatched_tasks
def test_create_download_marks_matching_claim_failed_when_dispatch_fails() -> None:
records = (_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024),)
task_store = FakeTaskStore()
def dispatch(_task: WorkflowRunArchiveDownloadTask) -> None:
raise RuntimeError("broker unavailable")
result = _service(records=records, tasks=task_store, dispatcher=dispatch).create_download(
_CONTEXT,
year=2025,
month=3,
)
assert result.status == WorkflowRunArchiveDownloadStatus.FAILED
assert result.error == "Failed to enqueue archive download task."
assert result.finished_at is not None
assert len(task_store.saved_tasks) == 2
assert task_store.saved_tasks[0].status == WorkflowRunArchiveDownloadStatus.PENDING
assert task_store.saved_tasks[1] == result
assert task_store.lock_calls == [
("tenant-1", result.download_id),
("tenant-1", result.download_id),
]
def test_create_download_does_not_overwrite_newer_state_when_dispatch_fails() -> None:
records = (_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024),)
task_store = FakeTaskStore()
def dispatch(task: WorkflowRunArchiveDownloadTask) -> None:
task_store.tasks[task.download_id] = task.model_copy(
update={
"status": WorkflowRunArchiveDownloadStatus.READY,
"storage_key": "downloads/archive.zip",
"file_name": "archive.zip",
}
)
raise RuntimeError("broker response lost")
result = _service(records=records, tasks=task_store, dispatcher=dispatch).create_download(
_CONTEXT,
year=2025,
month=3,
)
assert result.status == WorkflowRunArchiveDownloadStatus.READY
assert len(task_store.saved_tasks) == 1
assert task_store.lock_calls == [
("tenant-1", result.download_id),
("tenant-1", result.download_id),
]
def test_create_download_claims_dispatch_once() -> None:
records = (_bundle(shard="00-of-01", bundle_id="bundle-a", archive_bytes=1024),)
task_store = FakeTaskStore()
dispatched_tasks: list[WorkflowRunArchiveDownloadTask] = []
concurrent_results: list[WorkflowRunArchiveDownloadTask] = []
def dispatch(task: WorkflowRunArchiveDownloadTask) -> None:
dispatched_tasks.append(task)
concurrent_results.append(service.create_download(_CONTEXT, year=2025, month=3))
def sign_download_url(_storage_key: str, *, expires_in: int, filename: str) -> str:
assert expires_in > 0
assert filename
return "unused"
service = WorkflowRunArchiveService(
bundles=FakeBundleQuery(records),
tasks=task_store,
dispatcher=dispatch,
sign_download_url=sign_download_url,
)
result = service.create_download(_CONTEXT, year=2025, month=3)
assert dispatched_tasks == [result]
assert concurrent_results == [result]
def test_create_workflow_run_archive_download_task_rejects_missing_month(sqlite_session: Session) -> None:
def test_create_download_rejects_missing_month() -> None:
with pytest.raises(WorkflowRunArchiveNotFoundError):
create_workflow_run_archive_download_task(
sqlite_session,
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
cache=cast(WorkflowRunArchiveDownloadTaskCache, FakeTaskCache()),
dispatcher=_fake_dispatcher([]),
)
_service().create_download(_CONTEXT, year=2025, month=3)
def test_get_ready_workflow_run_archive_download_task_requires_ready_file() -> None:
pending_task = build_pending_archive_download_task(
def test_get_download_rejects_missing_task() -> None:
with pytest.raises(WorkflowRunArchiveDownloadTaskNotFoundError):
_service().get_download(_CONTEXT, download_id="missing")
def test_get_download_url_requires_ready_file() -> None:
pending = build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
@@ -358,29 +394,37 @@ def test_get_ready_workflow_run_archive_download_task_requires_ready_file() -> N
archive_bytes=1024,
download_id="download-1",
)
cache = FakeTaskCache(existing_task=pending_task)
with pytest.raises(WorkflowRunArchiveDownloadNotReadyError):
get_ready_workflow_run_archive_download_task(
tenant_id="tenant-1",
download_id="download-1",
cache=cast(WorkflowRunArchiveDownloadTaskCache, cache),
_service(tasks=FakeTaskStore({pending.download_id: pending})).get_download_url(
_CONTEXT,
download_id=pending.download_id,
)
ready_task = pending_task.model_copy(
def test_get_download_url_signs_ready_file_for_at_most_one_hour() -> None:
ready = build_pending_archive_download_task(
tenant_id="tenant-1",
requested_by="account-1",
year=2025,
month=3,
bundle_ids=["bundle-a"],
archive_bytes=1024,
download_id="download-1",
ttl_seconds=7200,
).model_copy(
update={
"status": WorkflowRunArchiveDownloadStatus.READY,
"storage_key": "downloads/download-1.zip",
"file_name": "workflow-run-logs-2025-03.zip",
}
)
cache = FakeTaskCache(existing_task=ready_task)
signed_urls: list[tuple[str, int, str]] = []
assert (
get_ready_workflow_run_archive_download_task(
tenant_id="tenant-1",
download_id="download-1",
cache=cast(WorkflowRunArchiveDownloadTaskCache, cache),
)
== ready_task
)
url = _service(
tasks=FakeTaskStore({ready.download_id: ready}),
signed_urls=signed_urls,
).get_download_url(_CONTEXT, download_id=ready.download_id)
assert url == "https://storage.example.com/archive.zip"
assert signed_urls == [("downloads/download-1.zip", 3600, "workflow-run-logs-2025-03.zip")]