mirror of
https://github.com/langgenius/dify.git
synced 2026-07-21 02:28:30 +08:00
348 lines
12 KiB
Python
348 lines
12 KiB
Python
"""
|
|
Workflow-run archive bundle index helpers.
|
|
|
|
Archive manifests in object storage remain the recoverable source of truth. This module mirrors their small query
|
|
surface into `workflow_run_archive_bundles` so console listing and download jobs can avoid listing R2 on request.
|
|
The backfill path is intentionally idempotent: every manifest is decoded, checked against the V2 schema markers, and
|
|
upserted by immutable bundle identity.
|
|
"""
|
|
|
|
import datetime
|
|
import json
|
|
import logging
|
|
import time
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass, field
|
|
from typing import TypedDict, cast
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
from extensions.ext_database import db
|
|
from libs.archive_storage import ArchiveStorage, get_archive_storage
|
|
from models.workflow import WorkflowRunArchiveBundle
|
|
from services.retention.workflow_run.constants import (
|
|
ARCHIVE_BUNDLE_FORMAT,
|
|
ARCHIVE_BUNDLE_MANIFEST_NAME,
|
|
ARCHIVE_BUNDLE_SCHEMA_VERSION,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ARCHIVE_BUNDLE_ROOT_PREFIX = "workflow-runs/v2/"
|
|
|
|
|
|
class ArchiveBundleTableManifestEntry(TypedDict):
|
|
row_count: int
|
|
checksum: str
|
|
size_bytes: int
|
|
object_key: str
|
|
|
|
|
|
class ArchiveBundleManifest(TypedDict):
|
|
schema_version: str
|
|
archive_format: str
|
|
tenant_id: str
|
|
tenant_prefix: str
|
|
year: int
|
|
month: int
|
|
shard: str
|
|
bundle_id: str
|
|
object_prefix: str
|
|
workflow_run_count: int
|
|
workflow_node_execution_count: int
|
|
min_created_at: str
|
|
max_created_at: str
|
|
min_run_id: str
|
|
max_run_id: str
|
|
archived_at: str
|
|
tables: dict[str, ArchiveBundleTableManifestEntry]
|
|
run_ids: list[str]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArchiveBundleIndexValues:
|
|
"""Computed DB-index values derived from one manifest."""
|
|
|
|
row_count: int
|
|
archive_bytes: int
|
|
archived_at: datetime.datetime
|
|
|
|
|
|
@dataclass
|
|
class ArchiveBundleIndexBackfillSummary:
|
|
"""Aggregate result for a manifest-to-DB-index reconciliation run."""
|
|
|
|
manifests_found: int = 0
|
|
bundles_processed: int = 0
|
|
bundles_upserted: int = 0
|
|
bundles_failed: int = 0
|
|
workflow_run_count: int = 0
|
|
row_count: int = 0
|
|
archive_bytes: int = 0
|
|
elapsed_time: float = 0.0
|
|
errors: list[str] = field(default_factory=list)
|
|
|
|
|
|
def decode_archive_bundle_manifest(manifest_data: bytes) -> ArchiveBundleManifest:
|
|
"""Decode raw manifest bytes into the V2 archive manifest shape."""
|
|
return cast(ArchiveBundleManifest, json.loads(manifest_data.decode("utf-8")))
|
|
|
|
|
|
def parse_archive_manifest_datetime(value: str) -> datetime.datetime:
|
|
"""Parse manifest datetimes and normalize timezone-aware values to naive UTC for DB storage."""
|
|
parsed = datetime.datetime.fromisoformat(value)
|
|
if parsed.tzinfo is None:
|
|
return parsed
|
|
return parsed.astimezone(datetime.UTC).replace(tzinfo=None)
|
|
|
|
|
|
def calculate_archive_bundle_index_values(
|
|
manifest: ArchiveBundleManifest,
|
|
manifest_size_bytes: int,
|
|
) -> ArchiveBundleIndexValues:
|
|
"""Calculate row count, stored bytes, and archived timestamp for the DB index."""
|
|
_validate_archive_bundle_manifest(manifest)
|
|
row_count = sum(entry["row_count"] for entry in manifest["tables"].values())
|
|
archive_bytes = manifest_size_bytes + sum(entry["size_bytes"] for entry in manifest["tables"].values())
|
|
return ArchiveBundleIndexValues(
|
|
row_count=row_count,
|
|
archive_bytes=archive_bytes,
|
|
archived_at=parse_archive_manifest_datetime(manifest["archived_at"]),
|
|
)
|
|
|
|
|
|
def upsert_archive_bundle_index_from_manifest(
|
|
session: Session,
|
|
manifest: ArchiveBundleManifest,
|
|
manifest_size_bytes: int,
|
|
) -> WorkflowRunArchiveBundle:
|
|
"""
|
|
Persist one archive manifest into `workflow_run_archive_bundles`.
|
|
|
|
The caller owns transaction boundaries. Re-running this function for the same manifest is safe and refreshes the
|
|
mutable metrics derived from object sizes and row counts.
|
|
"""
|
|
values = calculate_archive_bundle_index_values(manifest, manifest_size_bytes)
|
|
existing = session.scalar(
|
|
select(WorkflowRunArchiveBundle).where(
|
|
WorkflowRunArchiveBundle.tenant_id == manifest["tenant_id"],
|
|
WorkflowRunArchiveBundle.year == manifest["year"],
|
|
WorkflowRunArchiveBundle.month == manifest["month"],
|
|
WorkflowRunArchiveBundle.shard == manifest["shard"],
|
|
WorkflowRunArchiveBundle.bundle_id == manifest["bundle_id"],
|
|
)
|
|
)
|
|
if existing is None:
|
|
bundle = WorkflowRunArchiveBundle(
|
|
tenant_id=manifest["tenant_id"],
|
|
year=manifest["year"],
|
|
month=manifest["month"],
|
|
shard=manifest["shard"],
|
|
bundle_id=manifest["bundle_id"],
|
|
workflow_run_count=manifest["workflow_run_count"],
|
|
row_count=values.row_count,
|
|
archive_bytes=values.archive_bytes,
|
|
archived_at=values.archived_at,
|
|
)
|
|
session.add(bundle)
|
|
return bundle
|
|
|
|
existing.workflow_run_count = manifest["workflow_run_count"]
|
|
existing.row_count = values.row_count
|
|
existing.archive_bytes = values.archive_bytes
|
|
existing.archived_at = values.archived_at
|
|
return existing
|
|
|
|
|
|
class WorkflowRunArchiveBundleIndexBackfill:
|
|
"""
|
|
Rebuild the DB bundle index by scanning object-store manifests.
|
|
|
|
Tenant IDs are the cheapest scope because they map directly to the object prefix. Tenant prefixes are supported for
|
|
rollout reconciliation, but they still require listing all tenants under that prefix and filtering keys locally.
|
|
"""
|
|
|
|
storage: ArchiveStorage | None
|
|
session_factory: sessionmaker[Session]
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
storage: ArchiveStorage | None = None,
|
|
session_factory: sessionmaker[Session] | None = None,
|
|
) -> None:
|
|
self.storage = storage
|
|
self.session_factory = session_factory or sessionmaker(bind=db.engine, expire_on_commit=False)
|
|
|
|
def run(
|
|
self,
|
|
*,
|
|
tenant_ids: Sequence[str] | None = None,
|
|
tenant_prefixes: Sequence[str] | None = None,
|
|
year: int | None = None,
|
|
month: int | None = None,
|
|
limit: int | None = None,
|
|
dry_run: bool = False,
|
|
) -> ArchiveBundleIndexBackfillSummary:
|
|
"""Scan matching manifest objects and idempotently upsert their DB index rows."""
|
|
start_time = time.time()
|
|
summary = ArchiveBundleIndexBackfillSummary()
|
|
storage = self.storage or get_archive_storage()
|
|
manifest_keys = self._list_manifest_keys(
|
|
storage,
|
|
tenant_ids=tenant_ids,
|
|
tenant_prefixes=tenant_prefixes,
|
|
year=year,
|
|
month=month,
|
|
)
|
|
summary.manifests_found = len(manifest_keys)
|
|
|
|
if limit is not None:
|
|
manifest_keys = manifest_keys[:limit]
|
|
|
|
for manifest_key in manifest_keys:
|
|
try:
|
|
manifest_data = storage.get_object(manifest_key)
|
|
manifest = decode_archive_bundle_manifest(manifest_data)
|
|
self._validate_manifest_scope(
|
|
manifest,
|
|
manifest_key=manifest_key,
|
|
tenant_ids=tenant_ids,
|
|
tenant_prefixes=tenant_prefixes,
|
|
year=year,
|
|
month=month,
|
|
)
|
|
values = calculate_archive_bundle_index_values(manifest, len(manifest_data))
|
|
summary.bundles_processed += 1
|
|
summary.workflow_run_count += manifest["workflow_run_count"]
|
|
summary.row_count += values.row_count
|
|
summary.archive_bytes += values.archive_bytes
|
|
if dry_run:
|
|
continue
|
|
|
|
with self.session_factory() as session:
|
|
upsert_archive_bundle_index_from_manifest(session, manifest, len(manifest_data))
|
|
session.commit()
|
|
summary.bundles_upserted += 1
|
|
except Exception as exc:
|
|
logger.warning("Failed to backfill workflow archive bundle index from %s", manifest_key, exc_info=True)
|
|
summary.bundles_failed += 1
|
|
summary.errors.append(f"{manifest_key}: {exc}")
|
|
|
|
summary.elapsed_time = time.time() - start_time
|
|
return summary
|
|
|
|
@classmethod
|
|
def _list_manifest_keys(
|
|
cls,
|
|
storage: ArchiveStorage,
|
|
*,
|
|
tenant_ids: Sequence[str] | None,
|
|
tenant_prefixes: Sequence[str] | None,
|
|
year: int | None,
|
|
month: int | None,
|
|
) -> list[str]:
|
|
prefixes = cls._list_prefixes(tenant_ids=tenant_ids, tenant_prefixes=tenant_prefixes, year=year, month=month)
|
|
keys: list[str] = []
|
|
for prefix in prefixes:
|
|
keys.extend(storage.list_objects(prefix))
|
|
return sorted(
|
|
key
|
|
for key in keys
|
|
if key.endswith(f"/{ARCHIVE_BUNDLE_MANIFEST_NAME}")
|
|
and cls._manifest_key_matches_scope(
|
|
key,
|
|
tenant_ids=tenant_ids,
|
|
tenant_prefixes=tenant_prefixes,
|
|
year=year,
|
|
month=month,
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _list_prefixes(
|
|
*,
|
|
tenant_ids: Sequence[str] | None,
|
|
tenant_prefixes: Sequence[str] | None,
|
|
year: int | None,
|
|
month: int | None,
|
|
) -> list[str]:
|
|
if tenant_ids:
|
|
prefixes = []
|
|
for tenant_id in sorted(set(tenant_ids)):
|
|
prefix = f"{ARCHIVE_BUNDLE_ROOT_PREFIX}tenant_prefix={tenant_id[0].lower()}/tenant_id={tenant_id}/"
|
|
if year is not None:
|
|
prefix += f"year={year:04d}/"
|
|
if month is not None:
|
|
prefix += f"month={month:02d}/"
|
|
prefixes.append(prefix)
|
|
return prefixes
|
|
|
|
if tenant_prefixes:
|
|
return [
|
|
f"{ARCHIVE_BUNDLE_ROOT_PREFIX}tenant_prefix={tenant_prefix}/"
|
|
for tenant_prefix in sorted(set(tenant_prefixes))
|
|
]
|
|
|
|
return [ARCHIVE_BUNDLE_ROOT_PREFIX]
|
|
|
|
@staticmethod
|
|
def _manifest_key_matches_scope(
|
|
key: str,
|
|
*,
|
|
tenant_ids: Sequence[str] | None,
|
|
tenant_prefixes: Sequence[str] | None,
|
|
year: int | None,
|
|
month: int | None,
|
|
) -> bool:
|
|
if tenant_ids and _extract_key_part(key, "tenant_id") not in set(tenant_ids):
|
|
return False
|
|
if tenant_prefixes and _extract_key_part(key, "tenant_prefix") not in set(tenant_prefixes):
|
|
return False
|
|
if year is not None and _extract_key_part(key, "year") != f"{year:04d}":
|
|
return False
|
|
if month is not None and _extract_key_part(key, "month") != f"{month:02d}":
|
|
return False
|
|
return True
|
|
|
|
@staticmethod
|
|
def _validate_manifest_scope(
|
|
manifest: ArchiveBundleManifest,
|
|
*,
|
|
manifest_key: str,
|
|
tenant_ids: Sequence[str] | None,
|
|
tenant_prefixes: Sequence[str] | None,
|
|
year: int | None,
|
|
month: int | None,
|
|
) -> None:
|
|
expected_object_prefix = manifest_key.removesuffix(f"/{ARCHIVE_BUNDLE_MANIFEST_NAME}")
|
|
if manifest["object_prefix"] != expected_object_prefix:
|
|
raise ValueError(
|
|
f"manifest object_prefix mismatch: expected={expected_object_prefix}, "
|
|
f"actual={manifest['object_prefix']}"
|
|
)
|
|
if tenant_ids and manifest["tenant_id"] not in tenant_ids:
|
|
raise ValueError(f"manifest tenant_id is outside requested scope: {manifest['tenant_id']}")
|
|
if tenant_prefixes and manifest["tenant_prefix"] not in tenant_prefixes:
|
|
raise ValueError(f"manifest tenant_prefix is outside requested scope: {manifest['tenant_prefix']}")
|
|
if year is not None and manifest["year"] != year:
|
|
raise ValueError(f"manifest year is outside requested scope: {manifest['year']}")
|
|
if month is not None and manifest["month"] != month:
|
|
raise ValueError(f"manifest month is outside requested scope: {manifest['month']}")
|
|
|
|
|
|
def _validate_archive_bundle_manifest(manifest: ArchiveBundleManifest) -> None:
|
|
if manifest["schema_version"] != ARCHIVE_BUNDLE_SCHEMA_VERSION:
|
|
raise ValueError(f"unsupported archive bundle schema version: {manifest['schema_version']}")
|
|
if manifest["archive_format"] != ARCHIVE_BUNDLE_FORMAT:
|
|
raise ValueError(f"unsupported archive bundle format: {manifest['archive_format']}")
|
|
|
|
|
|
def _extract_key_part(key: str, name: str) -> str | None:
|
|
prefix = f"{name}="
|
|
for part in key.split("/"):
|
|
if part.startswith(prefix):
|
|
return part[len(prefix) :]
|
|
return None
|