mirror of
https://github.com/langgenius/dify.git
synced 2026-09-02 22:36:15 +08:00
94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
"""Dedicated-queue execution for legacy Dataset upgrades."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from celery import shared_task
|
|
|
|
from core.db.session_factory import session_factory
|
|
from services.dataset_knowledge_fs_upgrade_file_lease import cleanup_deferred_upgrade_files
|
|
from services.dataset_knowledge_fs_upgrade_service import (
|
|
KnowledgeFSUpgradeDocumentReconciler,
|
|
KnowledgeFSUpgradeNotReadyError,
|
|
KnowledgeFSUpgradeRunner,
|
|
)
|
|
|
|
KNOWLEDGE_FS_UPGRADE_QUEUE = "knowledge_fs_upgrade"
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@shared_task(
|
|
bind=True,
|
|
queue=KNOWLEDGE_FS_UPGRADE_QUEUE,
|
|
max_retries=360,
|
|
default_retry_delay=5,
|
|
)
|
|
def run_knowledge_fs_upgrade(self, *, job_id: str) -> None:
|
|
"""Run one checkpoint and enqueue the next only after it is committed."""
|
|
|
|
runner = KnowledgeFSUpgradeRunner(session_factory.get_session_maker())
|
|
try:
|
|
has_more = runner.run_next(job_id=job_id, celery_task_id=self.request.id)
|
|
except KnowledgeFSUpgradeNotReadyError as error:
|
|
if self.request.retries >= self.max_retries:
|
|
runner.fail(job_id=job_id, error=error)
|
|
logger.exception(
|
|
"KnowledgeFS Dataset upgrade exhausted provisioning retries",
|
|
extra={"upgrade_job_id": job_id},
|
|
)
|
|
raise
|
|
raise self.retry(exc=error)
|
|
except Exception as error:
|
|
runner.fail(job_id=job_id, error=error)
|
|
logger.exception("KnowledgeFS Dataset upgrade failed", extra={"upgrade_job_id": job_id})
|
|
raise
|
|
if has_more:
|
|
run_knowledge_fs_upgrade.apply_async(kwargs={"job_id": job_id})
|
|
else:
|
|
reconcile_knowledge_fs_upgrade_documents.apply_async(kwargs={"job_id": job_id})
|
|
|
|
|
|
@shared_task(
|
|
bind=True,
|
|
queue=KNOWLEDGE_FS_UPGRADE_QUEUE,
|
|
max_retries=10_080,
|
|
default_retry_delay=60,
|
|
)
|
|
def reconcile_knowledge_fs_upgrade_documents(self, *, job_id: str) -> None:
|
|
"""Eventually apply click-time metadata and availability without gating migration success."""
|
|
|
|
reconciler = KnowledgeFSUpgradeDocumentReconciler(session_factory.get_session_maker())
|
|
try:
|
|
remaining = reconciler.reconcile(job_id=job_id)
|
|
except Exception as error:
|
|
logger.warning(
|
|
"KnowledgeFS Dataset upgrade document reconciliation is not ready",
|
|
extra={"upgrade_job_id": job_id},
|
|
exc_info=True,
|
|
)
|
|
raise self.retry(exc=error)
|
|
if remaining:
|
|
raise self.retry(exc=KnowledgeFSUpgradeNotReadyError(f"{remaining} migrated documents are not visible yet"))
|
|
|
|
|
|
@shared_task(queue=KNOWLEDGE_FS_UPGRADE_QUEUE)
|
|
def cleanup_deferred_knowledge_fs_upgrade_files() -> int:
|
|
"""Delete orphaned source files after abandoned upgrade leases expire."""
|
|
|
|
return cleanup_deferred_upgrade_files(session_factory.get_session_maker())
|
|
|
|
|
|
def enqueue_knowledge_fs_upgrade(*, job_id: str) -> str:
|
|
result = run_knowledge_fs_upgrade.apply_async(kwargs={"job_id": job_id})
|
|
return str(result.id)
|
|
|
|
|
|
__all__ = [
|
|
"KNOWLEDGE_FS_UPGRADE_QUEUE",
|
|
"cleanup_deferred_knowledge_fs_upgrade_files",
|
|
"enqueue_knowledge_fs_upgrade",
|
|
"reconcile_knowledge_fs_upgrade_documents",
|
|
"run_knowledge_fs_upgrade",
|
|
]
|