diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index 555f41f60b2..01772746626 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -1,14 +1,14 @@ from dataclasses import dataclass from datetime import datetime from typing import Annotated, Any -from uuid import UUID +from uuid import UUID, uuid4 from flask import request from flask_restx import Resource from pydantic import BaseModel, Field, field_validator, model_validator from sqlalchemy import func, select from sqlalchemy.orm import Session -from werkzeug.exceptions import Forbidden, NotFound +from werkzeug.exceptions import Conflict, Forbidden, NotFound import services from configs import dify_config @@ -32,6 +32,7 @@ from controllers.console.wraps import ( with_current_tenant_id, with_current_user, ) +from core.db.session_factory import session_factory from core.entities.knowledge_entities import IndexingEstimate from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError from core.indexing_runner import IndexingRunner @@ -53,10 +54,17 @@ from models.enums import ApiTokenType, SegmentStatus from models.provider_ids import ModelProviderID from services.api_token_service import ApiTokenCache, get_effective_token_last_used_at from services.app_service import AppService +from services.dataset_knowledge_fs_upgrade_service import ( + KnowledgeFSUpgradeConflictError, + KnowledgeFSUpgradeNotFoundError, + KnowledgeFSUpgradeSnapshotService, + upgrade_job_response, +) from services.dataset_ref_service import DatasetRefService from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService from services.enterprise import rbac_service as enterprise_rbac_service from services.enterprise.rbac_service import RBACResourceWhitelistScope, ReplaceMemberBindings +from services.knowledge_fs.product_dto import KnowledgeFSUpgradeJobResponse, KnowledgeFSUpgradeRetryResponse from tasks.initialize_created_app_rbac_access_task import initialize_created_app_rbac_access_task register_response_schema_models(console_ns, ApiBaseUrlResponse, SimpleResultResponse, UsageCheckResponse) @@ -356,6 +364,8 @@ register_response_schema_models( RetrievalSettingResponse, PartialMemberListResponse, AutoDisableLogsResponse, + KnowledgeFSUpgradeJobResponse, + KnowledgeFSUpgradeRetryResponse, ) @@ -819,6 +829,124 @@ class DatasetApi(Resource): raise DatasetInUseError() +@console_ns.route("/datasets//knowledge-fs-upgrades") +class DatasetKnowledgeFSUpgradeApi(Resource): + @console_ns.response( + 202, + "KnowledgeFS Dataset upgrade accepted", + console_ns.models[KnowledgeFSUpgradeJobResponse.__name__], + ) + @setup_required + @login_required + @account_initialization_required + @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) + @with_current_user + @with_current_tenant_id + def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID): + if not dify_config.KNOWLEDGE_FS_ENABLED: + raise NotFound() + dataset_id_str = str(dataset_id) + with session_factory.create_session() as session: + dataset = DatasetService.get_dataset_for_tenant(dataset_id_str, current_tenant_id, session=session) + if dataset is None: + raise NotFound("Dataset not found.") + if not dify_config.RBAC_ENABLED: + try: + DatasetService.check_dataset_permission(dataset, current_user, session) + except services.errors.account.NoPermissionError as error: + raise Forbidden(str(error)) from error + if not (current_user.has_edit_permission or current_user.is_dataset_operator): + raise Forbidden() + snapshots = KnowledgeFSUpgradeSnapshotService(session_factory.get_session_maker()) + try: + job = snapshots.create( + tenant_id=current_tenant_id, + dataset_id=dataset_id_str, + requested_by_account_id=current_user.id, + idempotency_key=request.headers.get("Idempotency-Key"), + ) + except KnowledgeFSUpgradeNotFoundError as error: + raise NotFound(str(error)) from error + except KnowledgeFSUpgradeConflictError as error: + raise Conflict(str(error)) from error + _enqueue_upgrade_job(snapshots, tenant_id=current_tenant_id, job_id=job.id) + return dump_response(KnowledgeFSUpgradeJobResponse, upgrade_job_response(job)), 202 + + +@console_ns.route("/datasets//knowledge-fs-upgrades/") +class DatasetKnowledgeFSUpgradeJobApi(Resource): + @console_ns.response( + 200, + "KnowledgeFS Dataset upgrade status", + console_ns.models[KnowledgeFSUpgradeJobResponse.__name__], + ) + @setup_required + @login_required + @account_initialization_required + @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_READONLY) + @with_current_user + @with_current_tenant_id + def get(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, job_id: str): + with session_factory.create_session() as session: + _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session) + snapshots = KnowledgeFSUpgradeSnapshotService(session_factory.get_session_maker()) + try: + job = snapshots.get(tenant_id=current_tenant_id, job_id=job_id) + except KnowledgeFSUpgradeNotFoundError as error: + raise NotFound(str(error)) from error + if job.old_dataset_id != str(dataset_id): + raise NotFound("Upgrade job was not found") + return dump_response(KnowledgeFSUpgradeJobResponse, upgrade_job_response(job)) + + @console_ns.response( + 202, + "KnowledgeFS Dataset upgrade retry accepted", + console_ns.models[KnowledgeFSUpgradeRetryResponse.__name__], + ) + @setup_required + @login_required + @account_initialization_required + @rbac_permission_required(RBACResourceScope.DATASET, RBACPermission.DATASET_EDIT) + @with_current_user + @with_current_tenant_id + def post(self, current_tenant_id: str, current_user: Account, dataset_id: UUID, job_id: str): + with session_factory.create_session() as session: + _get_accessible_dataset(dataset_id, current_tenant_id, current_user, session) + if not dify_config.RBAC_ENABLED and not ( + current_user.has_edit_permission or current_user.is_dataset_operator + ): + raise Forbidden() + snapshots = KnowledgeFSUpgradeSnapshotService(session_factory.get_session_maker()) + try: + job = snapshots.retry(tenant_id=current_tenant_id, job_id=job_id) + except KnowledgeFSUpgradeNotFoundError as error: + raise NotFound(str(error)) from error + except KnowledgeFSUpgradeConflictError as error: + raise Conflict(str(error)) from error + if job.old_dataset_id != str(dataset_id): + raise NotFound("Upgrade job was not found") + _enqueue_upgrade_job(snapshots, tenant_id=current_tenant_id, job_id=job.id) + return dump_response(KnowledgeFSUpgradeRetryResponse, {"id": job.id, "status": "queued"}), 202 + + +def _enqueue_upgrade_job( + snapshots: KnowledgeFSUpgradeSnapshotService, + *, + tenant_id: str, + job_id: str, +) -> None: + from tasks.knowledge_fs_upgrade_tasks import run_knowledge_fs_upgrade + + task_id = str(uuid4()) + if not snapshots.claim_enqueue(tenant_id=tenant_id, job_id=job_id, task_id=task_id): + return + try: + run_knowledge_fs_upgrade.apply_async(kwargs={"job_id": job_id}, task_id=task_id) + except Exception: + snapshots.release_enqueue_claim(tenant_id=tenant_id, job_id=job_id, task_id=task_id) + raise + + @console_ns.route("/datasets//use-check") class DatasetUseCheckApi(Resource): @console_ns.doc("check_dataset_use") diff --git a/api/dev/check_knowledge_fs_coverage.py b/api/dev/check_knowledge_fs_coverage.py index 0b1fd775775..2bebdc8ef45 100644 --- a/api/dev/check_knowledge_fs_coverage.py +++ b/api/dev/check_knowledge_fs_coverage.py @@ -32,6 +32,7 @@ NON_CORE_COVERAGE_ALLOWLIST = frozenset( "api/migrations/versions/2026_08_10_1200-7c1e9a4b2d60_add_knowledge_fs_staged_uploads.py", "api/migrations/versions/2026_08_13_1200-9d4e6f8a1b2c_add_knowledge_fs_space_tag_bindings.py", "api/migrations/versions/2026_08_17_1200-4f8b2c7d9e10_add_knowledge_fs_icon_background.py", + "api/migrations/versions/2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py", } ) HUNK_HEADER = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,\d+)? @@") @@ -155,6 +156,8 @@ def is_core_coverage_path(path: str) -> bool: "api/controllers/openapi/knowledge_fs.py", "api/events/event_handlers/sync_knowledge_fs_bindings_when_app_published_workflow_updated.py", "api/extensions/ext_knowledge_fs_observability.py", + "api/services/dataset_knowledge_fs_upgrade_file_lease.py", + "api/services/dataset_knowledge_fs_upgrade_service.py", "api/services/knowledge_fs_capability.py", }: return True diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index 5671fd31868..1c42a08939e 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -179,6 +179,7 @@ def init_app(app: DifyApp) -> Celery: "tasks.workflow_run_archive_download_tasks", # workflow-run archive download preparation "tasks.knowledge_fs_initial_source_preview_tasks", # datasource previews use the standard dataset queue "tasks.knowledge_fs_failed_retrieval_tasks", # best-effort Workflow quality capture uses dataset workers + "tasks.knowledge_fs_upgrade_tasks", # legacy Dataset upgrades use a dedicated queue and worker ] day = dify_config.CELERY_BEAT_SCHEDULER_TIME @@ -197,6 +198,10 @@ def init_app(app: DifyApp) -> Celery: "task": "tasks.knowledge_fs_lifecycle_tasks.cleanup_knowledge_fs_staged_uploads", "schedule": timedelta(seconds=dify_config.KNOWLEDGE_FS_LIFECYCLE_POLL_INTERVAL_SECONDS), } + beat_schedule["knowledge_fs_upgrade_file_cleanup"] = { + "task": "tasks.knowledge_fs_upgrade_tasks.cleanup_deferred_knowledge_fs_upgrade_files", + "schedule": timedelta(seconds=dify_config.KNOWLEDGE_FS_LIFECYCLE_POLL_INTERVAL_SECONDS), + } if dify_config.ENABLE_CONVERSATION_CLEANUP_TASK: imports.append("tasks.delete_conversation_task") beat_schedule["conversation_cleanup_sweeper"] = { diff --git a/api/knowledge-fs-contract.lock.json b/api/knowledge-fs-contract.lock.json index 5deae373bd9..bf42fc9b2ea 100644 --- a/api/knowledge-fs-contract.lock.json +++ b/api/knowledge-fs-contract.lock.json @@ -1,6 +1,6 @@ { "schemaVersion": 5, - "subtreeTree": "61759d6a25df9fa4746b6cfd0e795da27cfbb388", + "subtreeTree": "a964e23f1d42d13e959b120192fc74df3d755aa1", "openapiSha256": "3a712231fa850c4f5151bc283205da9062086a0b693f0d1ab01c2c526323f018", "capabilityV2AuthManifestSha256": "fc0a47e23cce12544882f0298522b4933002e892b84ce1815df7e81d36a7a0c7", "capabilityV2AuthTestVectorSha256": "ae0de37b1ff05c40f905cf17a7b410d8971acacf64db07d5ee3d6fecfa559ce3", diff --git a/api/migrations/versions/2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py b/api/migrations/versions/2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py new file mode 100644 index 00000000000..15dfd64e018 --- /dev/null +++ b/api/migrations/versions/2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py @@ -0,0 +1,172 @@ +"""add KnowledgeFS legacy Dataset upgrade jobs + +Revision ID: f3a8c1d7e920 +Revises: 4f8b2c7d9e10 +Create Date: 2026-08-17 12:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +import models + +revision = "f3a8c1d7e920" +down_revision = "4f8b2c7d9e10" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "knowledge_fs_upgrade_jobs", + sa.Column("tenant_id", models.types.StringUUID(), nullable=False), + sa.Column("old_dataset_id", models.types.StringUUID(), nullable=False), + sa.Column("requested_by_account_id", models.types.StringUUID(), nullable=False), + sa.Column("owner_account_id", models.types.StringUUID(), nullable=False), + sa.Column("idempotency_key", sa.String(length=255), nullable=False), + sa.Column("snapshot_at", sa.DateTime(), nullable=False), + sa.Column("config_snapshot", sa.JSON(), nullable=False), + sa.Column("permission_snapshot", sa.JSON(), nullable=False), + sa.Column("app_binding_snapshot", sa.JSON(), nullable=False), + sa.Column("tag_ids_snapshot", sa.JSON(), nullable=False), + sa.Column("status", sa.String(length=16), server_default=sa.text("'queued'"), nullable=False), + sa.Column("stage", sa.String(length=32), server_default=sa.text("'validating'"), nullable=False), + sa.Column("new_control_space_id", models.types.StringUUID(), nullable=True), + sa.Column("resolved_configuration", sa.JSON(), nullable=True), + sa.Column("total_documents", sa.Integer(), server_default=sa.text("0"), nullable=False), + sa.Column("completed_documents", sa.Integer(), server_default=sa.text("0"), nullable=False), + sa.Column("total_sources", sa.Integer(), server_default=sa.text("0"), nullable=False), + sa.Column("completed_sources", sa.Integer(), server_default=sa.text("0"), nullable=False), + sa.Column("attempt_count", sa.Integer(), server_default=sa.text("0"), nullable=False), + sa.Column("celery_task_id", sa.String(length=255), nullable=True), + sa.Column("last_error_code", sa.String(length=128), nullable=True), + sa.Column("last_error_message", models.types.LongText(), nullable=True), + sa.Column("completed_at", sa.DateTime(), nullable=True), + sa.Column("id", models.types.StringUUID(), nullable=False), + sa.Column("created_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.Column("updated_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.CheckConstraint("attempt_count >= 0", name="kfs_upgrade_job_attempt_count_ck"), + sa.CheckConstraint("completed_documents >= 0", name="kfs_upgrade_job_document_done_ck"), + sa.CheckConstraint("total_documents >= 0", name="kfs_upgrade_job_document_total_ck"), + sa.CheckConstraint("completed_sources >= 0", name="kfs_upgrade_job_source_done_ck"), + sa.CheckConstraint("total_sources >= 0", name="kfs_upgrade_job_source_total_ck"), + sa.ForeignKeyConstraint( + ["tenant_id"], ["tenants.id"], name="kfs_upgrade_job_workspace_fk", ondelete="RESTRICT" + ), + sa.ForeignKeyConstraint( + ["tenant_id", "new_control_space_id"], + ["knowledge_fs_control_spaces.tenant_id", "knowledge_fs_control_spaces.id"], + name="kfs_upgrade_job_space_fk", + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name="kfs_upgrade_job_pkey"), + sa.UniqueConstraint("tenant_id", "idempotency_key", name="kfs_upgrade_job_idempotency_uq"), + ) + op.create_index( + "kfs_upgrade_job_dataset_created_idx", + "knowledge_fs_upgrade_jobs", + ["tenant_id", "old_dataset_id", "created_at"], + ) + op.create_index("kfs_upgrade_job_status_updated_idx", "knowledge_fs_upgrade_jobs", ["status", "updated_at"]) + + op.create_table( + "knowledge_fs_upgrade_documents", + sa.Column("job_id", models.types.StringUUID(), nullable=False), + sa.Column("tenant_id", models.types.StringUUID(), nullable=False), + sa.Column("old_document_id", models.types.StringUUID(), nullable=False), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column("data_source_type", sa.String(length=32), nullable=False), + sa.Column("data_source_info", sa.JSON(), nullable=False), + sa.Column("metadata_snapshot", sa.JSON(), nullable=False), + sa.Column("desired_enabled", sa.Boolean(), nullable=False), + sa.Column("legacy_archived", sa.Boolean(), nullable=False), + sa.Column("legacy_indexing_status", sa.String(length=32), nullable=False), + sa.Column("legacy_display_status", sa.String(length=32), nullable=True), + sa.Column("old_upload_file_id", models.types.StringUUID(), nullable=True), + sa.Column("source_key", sa.String(length=255), nullable=True), + sa.Column("status", sa.String(length=16), server_default=sa.text("'pending'"), nullable=False), + sa.Column("staged_upload_id", models.types.StringUUID(), nullable=True), + sa.Column("new_document_asset_id", models.types.StringUUID(), nullable=True), + sa.Column("new_logical_document_id", models.types.StringUUID(), nullable=True), + sa.Column("compilation_job_id", models.types.StringUUID(), nullable=True), + sa.Column("state_reconcile_attempt_count", sa.Integer(), server_default=sa.text("0"), nullable=False), + sa.Column("state_reconciled_at", sa.DateTime(), nullable=True), + sa.Column("state_reconcile_error", models.types.LongText(), nullable=True), + sa.Column("last_error_code", sa.String(length=128), nullable=True), + sa.Column("last_error_message", models.types.LongText(), nullable=True), + sa.Column("id", models.types.StringUUID(), nullable=False), + sa.Column("created_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.Column("updated_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.CheckConstraint( + "state_reconcile_attempt_count >= 0", + name="kfs_upgrade_document_reconcile_attempt_ck", + ), + sa.ForeignKeyConstraint( + ["job_id"], ["knowledge_fs_upgrade_jobs.id"], name="kfs_upgrade_document_job_fk", ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id", name="kfs_upgrade_document_pkey"), + sa.UniqueConstraint("job_id", "old_document_id", name="kfs_upgrade_document_identity_uq"), + ) + op.create_index("kfs_upgrade_document_dispatch_idx", "knowledge_fs_upgrade_documents", ["job_id", "status", "id"]) + + op.create_table( + "knowledge_fs_upgrade_sources", + sa.Column("job_id", models.types.StringUUID(), nullable=False), + sa.Column("tenant_id", models.types.StringUUID(), nullable=False), + sa.Column("source_key", sa.String(length=255), nullable=False), + sa.Column("payload_snapshot", sa.JSON(), nullable=False), + sa.Column("status", sa.String(length=16), server_default=sa.text("'pending'"), nullable=False), + sa.Column("new_connection_id", models.types.StringUUID(), nullable=True), + sa.Column("new_source_id", models.types.StringUUID(), nullable=True), + sa.Column("initial_sync_task_id", models.types.StringUUID(), nullable=True), + sa.Column("last_error_code", sa.String(length=128), nullable=True), + sa.Column("last_error_message", models.types.LongText(), nullable=True), + sa.Column("id", models.types.StringUUID(), nullable=False), + sa.Column("created_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.Column("updated_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.ForeignKeyConstraint( + ["job_id"], ["knowledge_fs_upgrade_jobs.id"], name="kfs_upgrade_source_job_fk", ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id", name="kfs_upgrade_source_pkey"), + sa.UniqueConstraint("job_id", "source_key", name="kfs_upgrade_source_identity_uq"), + ) + op.create_index("kfs_upgrade_source_dispatch_idx", "knowledge_fs_upgrade_sources", ["job_id", "status", "id"]) + + op.create_table( + "knowledge_fs_upgrade_file_leases", + sa.Column("job_id", models.types.StringUUID(), nullable=False), + sa.Column("old_upload_file_id", models.types.StringUUID(), nullable=False), + sa.Column("status", sa.String(length=16), server_default=sa.text("'active'"), nullable=False), + sa.Column("expires_at", sa.DateTime(), nullable=False), + sa.Column("released_at", sa.DateTime(), nullable=True), + sa.Column("cleanup_requested_at", sa.DateTime(), nullable=True), + sa.Column("id", models.types.StringUUID(), nullable=False), + sa.Column("created_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.Column("updated_at", sa.DateTime(), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False), + sa.ForeignKeyConstraint( + ["job_id"], + ["knowledge_fs_upgrade_jobs.id"], + name="kfs_upgrade_file_lease_job_fk", + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name="kfs_upgrade_file_lease_pkey"), + sa.UniqueConstraint("job_id", "old_upload_file_id", name="kfs_upgrade_file_lease_identity_uq"), + ) + op.create_index( + "kfs_upgrade_file_lease_active_idx", + "knowledge_fs_upgrade_file_leases", + ["old_upload_file_id", "status", "expires_at"], + ) + + +def downgrade() -> None: + op.drop_index("kfs_upgrade_file_lease_active_idx", table_name="knowledge_fs_upgrade_file_leases") + op.drop_table("knowledge_fs_upgrade_file_leases") + op.drop_index("kfs_upgrade_source_dispatch_idx", table_name="knowledge_fs_upgrade_sources") + op.drop_table("knowledge_fs_upgrade_sources") + op.drop_index("kfs_upgrade_document_dispatch_idx", table_name="knowledge_fs_upgrade_documents") + op.drop_table("knowledge_fs_upgrade_documents") + op.drop_index("kfs_upgrade_job_status_updated_idx", table_name="knowledge_fs_upgrade_jobs") + op.drop_index("kfs_upgrade_job_dataset_created_idx", table_name="knowledge_fs_upgrade_jobs") + op.drop_table("knowledge_fs_upgrade_jobs") diff --git a/api/models/__init__.py b/api/models/__init__.py index 02b8a129782..5f7d9fd036a 100644 --- a/api/models/__init__.py +++ b/api/models/__init__.py @@ -90,6 +90,14 @@ from .knowledge_fs import ( KnowledgeFSSpaceTagBinding, KnowledgeFSStagedUpload, KnowledgeFSStagedUploadStatus, + KnowledgeFSUpgradeDocument, + KnowledgeFSUpgradeFileLease, + KnowledgeFSUpgradeFileLeaseStatus, + KnowledgeFSUpgradeItemStatus, + KnowledgeFSUpgradeJob, + KnowledgeFSUpgradeJobStatus, + KnowledgeFSUpgradeSource, + KnowledgeFSUpgradeStage, ) from .knowledge_fs_cleanup import ( KnowledgeFSCleanupAuthorization, @@ -307,6 +315,14 @@ __all__ = [ "KnowledgeFSSpaceTagBinding", "KnowledgeFSStagedUpload", "KnowledgeFSStagedUploadStatus", + "KnowledgeFSUpgradeDocument", + "KnowledgeFSUpgradeFileLease", + "KnowledgeFSUpgradeFileLeaseStatus", + "KnowledgeFSUpgradeItemStatus", + "KnowledgeFSUpgradeJob", + "KnowledgeFSUpgradeJobStatus", + "KnowledgeFSUpgradeSource", + "KnowledgeFSUpgradeStage", "KnowledgeFSWorkspaceCutoverLedger", "KnowledgeFSWorkspaceCutoverPhase", "LoadBalancingModelConfig", diff --git a/api/models/knowledge_fs.py b/api/models/knowledge_fs.py index f4188e97347..bbb6b3bc7eb 100644 --- a/api/models/knowledge_fs.py +++ b/api/models/knowledge_fs.py @@ -211,6 +211,36 @@ class KnowledgeFSStagedUploadStatus(StrEnum): EXPIRED = "expired" +class KnowledgeFSUpgradeJobStatus(StrEnum): + QUEUED = "queued" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class KnowledgeFSUpgradeStage(StrEnum): + VALIDATING = "validating" + WAITING_FOR_SPACE = "waiting_for_space" + CREATING_SOURCES = "creating_sources" + SUBMITTING_DOCUMENTS = "submitting_documents" + MIGRATING_ACCESS = "migrating_access" + FINALIZING = "finalizing" + COMPLETED = "completed" + + +class KnowledgeFSUpgradeItemStatus(StrEnum): + PENDING = "pending" + PROCESSING = "processing" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class KnowledgeFSUpgradeFileLeaseStatus(StrEnum): + ACTIVE = "active" + RELEASED = "released" + EXPIRED = "expired" + + class KnowledgeFSControlSpace(DefaultFieldsDCMixin, TypeBase): """Dify product resource registered to at most one KnowledgeFS Space.""" @@ -813,6 +843,183 @@ class KnowledgeFSLifecycleOutbox(DefaultFieldsDCMixin, TypeBase): retain_until: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) +class KnowledgeFSUpgradeJob(DefaultFieldsDCMixin, TypeBase): + """Immutable legacy Dataset snapshot and resumable upgrade progress.""" + + __tablename__ = "knowledge_fs_upgrade_jobs" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="kfs_upgrade_job_pkey"), + UniqueConstraint("tenant_id", "idempotency_key", name="kfs_upgrade_job_idempotency_uq"), + sa.ForeignKeyConstraint( + ["tenant_id"], + ["tenants.id"], + name="kfs_upgrade_job_workspace_fk", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["tenant_id", "new_control_space_id"], + ["knowledge_fs_control_spaces.tenant_id", "knowledge_fs_control_spaces.id"], + name="kfs_upgrade_job_space_fk", + ondelete="RESTRICT", + ), + Index("kfs_upgrade_job_dataset_created_idx", "tenant_id", "old_dataset_id", "created_at"), + Index("kfs_upgrade_job_status_updated_idx", "status", "updated_at"), + sa.CheckConstraint("total_documents >= 0", name=sa.schema.conv("kfs_upgrade_job_document_total_ck")), + sa.CheckConstraint("completed_documents >= 0", name=sa.schema.conv("kfs_upgrade_job_document_done_ck")), + sa.CheckConstraint("total_sources >= 0", name=sa.schema.conv("kfs_upgrade_job_source_total_ck")), + sa.CheckConstraint("completed_sources >= 0", name=sa.schema.conv("kfs_upgrade_job_source_done_ck")), + sa.CheckConstraint("attempt_count >= 0", name=sa.schema.conv("kfs_upgrade_job_attempt_count_ck")), + ) + + tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + old_dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + requested_by_account_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + owner_account_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + idempotency_key: Mapped[str] = mapped_column(String(255), nullable=False) + snapshot_at: Mapped[datetime] = mapped_column(DateTime, nullable=False) + config_snapshot: Mapped[dict[str, object]] = mapped_column(sa.JSON, nullable=False) + permission_snapshot: Mapped[dict[str, object]] = mapped_column(sa.JSON, nullable=False) + app_binding_snapshot: Mapped[list[dict[str, object]]] = mapped_column(sa.JSON, nullable=False) + tag_ids_snapshot: Mapped[list[str]] = mapped_column(sa.JSON, nullable=False) + status: Mapped[KnowledgeFSUpgradeJobStatus] = mapped_column( + EnumText(KnowledgeFSUpgradeJobStatus, length=16), + nullable=False, + server_default=sa.text("'queued'"), + default=KnowledgeFSUpgradeJobStatus.QUEUED, + ) + stage: Mapped[KnowledgeFSUpgradeStage] = mapped_column( + EnumText(KnowledgeFSUpgradeStage, length=32), + nullable=False, + server_default=sa.text("'validating'"), + default=KnowledgeFSUpgradeStage.VALIDATING, + ) + new_control_space_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + resolved_configuration: Mapped[dict[str, object] | None] = mapped_column(sa.JSON, nullable=True, default=None) + total_documents: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0, server_default=sa.text("0")) + completed_documents: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0, server_default=sa.text("0")) + total_sources: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0, server_default=sa.text("0")) + completed_sources: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0, server_default=sa.text("0")) + attempt_count: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0, server_default=sa.text("0")) + celery_task_id: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) + last_error_code: Mapped[str | None] = mapped_column(String(128), nullable=True, default=None) + last_error_message: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) + completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) + + +class KnowledgeFSUpgradeDocument(DefaultFieldsDCMixin, TypeBase): + """One legacy document captured at click time, independent of later Dataset updates.""" + + __tablename__ = "knowledge_fs_upgrade_documents" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="kfs_upgrade_document_pkey"), + UniqueConstraint("job_id", "old_document_id", name="kfs_upgrade_document_identity_uq"), + sa.ForeignKeyConstraint( + ["job_id"], + ["knowledge_fs_upgrade_jobs.id"], + name="kfs_upgrade_document_job_fk", + ondelete="CASCADE", + ), + Index("kfs_upgrade_document_dispatch_idx", "job_id", "status", "id"), + sa.CheckConstraint( + "state_reconcile_attempt_count >= 0", + name=sa.schema.conv("kfs_upgrade_document_reconcile_attempt_ck"), + ), + ) + + job_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + old_document_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + name: Mapped[str] = mapped_column(String(255), nullable=False) + data_source_type: Mapped[str] = mapped_column(String(32), nullable=False) + data_source_info: Mapped[dict[str, object]] = mapped_column(sa.JSON, nullable=False) + metadata_snapshot: Mapped[dict[str, object]] = mapped_column(sa.JSON, nullable=False) + desired_enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False) + legacy_archived: Mapped[bool] = mapped_column(sa.Boolean, nullable=False) + legacy_indexing_status: Mapped[str] = mapped_column(String(32), nullable=False) + legacy_display_status: Mapped[str | None] = mapped_column(String(32), nullable=True, default=None) + old_upload_file_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + source_key: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) + status: Mapped[KnowledgeFSUpgradeItemStatus] = mapped_column( + EnumText(KnowledgeFSUpgradeItemStatus, length=16), + nullable=False, + server_default=sa.text("'pending'"), + default=KnowledgeFSUpgradeItemStatus.PENDING, + ) + staged_upload_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + new_document_asset_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + new_logical_document_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + compilation_job_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + state_reconcile_attempt_count: Mapped[int] = mapped_column( + sa.Integer, nullable=False, server_default=sa.text("0"), default=0 + ) + state_reconciled_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) + state_reconcile_error: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) + last_error_code: Mapped[str | None] = mapped_column(String(128), nullable=True, default=None) + last_error_message: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) + + +class KnowledgeFSUpgradeSource(DefaultFieldsDCMixin, TypeBase): + """One deduplicated Source definition derived from the Dataset snapshot.""" + + __tablename__ = "knowledge_fs_upgrade_sources" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="kfs_upgrade_source_pkey"), + UniqueConstraint("job_id", "source_key", name="kfs_upgrade_source_identity_uq"), + sa.ForeignKeyConstraint( + ["job_id"], + ["knowledge_fs_upgrade_jobs.id"], + name="kfs_upgrade_source_job_fk", + ondelete="CASCADE", + ), + Index("kfs_upgrade_source_dispatch_idx", "job_id", "status", "id"), + ) + + job_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + source_key: Mapped[str] = mapped_column(String(255), nullable=False) + payload_snapshot: Mapped[dict[str, object]] = mapped_column(sa.JSON, nullable=False) + status: Mapped[KnowledgeFSUpgradeItemStatus] = mapped_column( + EnumText(KnowledgeFSUpgradeItemStatus, length=16), + nullable=False, + server_default=sa.text("'pending'"), + default=KnowledgeFSUpgradeItemStatus.PENDING, + ) + new_connection_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + new_source_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + initial_sync_task_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) + last_error_code: Mapped[str | None] = mapped_column(String(128), nullable=True, default=None) + last_error_message: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) + + +class KnowledgeFSUpgradeFileLease(DefaultFieldsDCMixin, TypeBase): + """Temporary physical-retention lease for a legacy uploaded source file.""" + + __tablename__ = "knowledge_fs_upgrade_file_leases" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="kfs_upgrade_file_lease_pkey"), + UniqueConstraint("job_id", "old_upload_file_id", name="kfs_upgrade_file_lease_identity_uq"), + sa.ForeignKeyConstraint( + ["job_id"], + ["knowledge_fs_upgrade_jobs.id"], + name="kfs_upgrade_file_lease_job_fk", + ondelete="CASCADE", + ), + Index("kfs_upgrade_file_lease_active_idx", "old_upload_file_id", "status", "expires_at"), + ) + + job_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + old_upload_file_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + expires_at: Mapped[datetime] = mapped_column(DateTime, nullable=False) + status: Mapped[KnowledgeFSUpgradeFileLeaseStatus] = mapped_column( + EnumText(KnowledgeFSUpgradeFileLeaseStatus, length=16), + nullable=False, + server_default=sa.text("'active'"), + default=KnowledgeFSUpgradeFileLeaseStatus.ACTIVE, + ) + released_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) + cleanup_requested_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) + + __all__ = [ "AppKnowledgeFSSpaceJoin", "KnowledgeFSAllowedActions", @@ -850,4 +1057,12 @@ __all__ = [ "KnowledgeFSScoreThresholdIntentPayload", "KnowledgeFSStagedUpload", "KnowledgeFSStagedUploadStatus", + "KnowledgeFSUpgradeDocument", + "KnowledgeFSUpgradeFileLease", + "KnowledgeFSUpgradeFileLeaseStatus", + "KnowledgeFSUpgradeItemStatus", + "KnowledgeFSUpgradeJob", + "KnowledgeFSUpgradeJobStatus", + "KnowledgeFSUpgradeSource", + "KnowledgeFSUpgradeStage", ] diff --git a/api/services/dataset_knowledge_fs_upgrade_file_lease.py b/api/services/dataset_knowledge_fs_upgrade_file_lease.py new file mode 100644 index 00000000000..b9d6b3daa45 --- /dev/null +++ b/api/services/dataset_knowledge_fs_upgrade_file_lease.py @@ -0,0 +1,193 @@ +"""Physical-retention checks for legacy files captured by an upgrade snapshot.""" + +from __future__ import annotations + +import json +from collections.abc import Iterable +from datetime import datetime +from json import JSONDecodeError + +import sqlalchemy as sa +from sqlalchemy.orm import Session, aliased, sessionmaker + +from extensions.ext_storage import storage +from libs.datetime_utils import naive_utc_now +from models.dataset import Document +from models.knowledge_fs import KnowledgeFSUpgradeFileLease, KnowledgeFSUpgradeFileLeaseStatus +from models.model import UploadFile + + +def active_upgrade_file_ids( + session: Session, + upload_file_ids: Iterable[str], + *, + now: datetime | None = None, +) -> set[str]: + """Return source files whose physical deletion must wait for migration.""" + + normalized_ids = frozenset(str(upload_file_id) for upload_file_id in upload_file_ids if upload_file_id) + if not normalized_ids: + return set() + return set( + session.scalars( + sa.select(KnowledgeFSUpgradeFileLease.old_upload_file_id).where( + KnowledgeFSUpgradeFileLease.old_upload_file_id.in_(normalized_ids), + KnowledgeFSUpgradeFileLease.status == KnowledgeFSUpgradeFileLeaseStatus.ACTIVE, + KnowledgeFSUpgradeFileLease.expires_at > (now or naive_utc_now()), + ) + ) + ) + + +def reserve_upgrade_file_cleanup( + session: Session, + upload_file_ids: Iterable[str], + *, + now: datetime | None = None, +) -> set[str]: + """Persist cleanup intent for files whose active migration lease blocks deletion.""" + + requested_at = now or naive_utc_now() + normalized_ids = frozenset(str(upload_file_id) for upload_file_id in upload_file_ids if upload_file_id) + if not normalized_ids: + return set() + leases = list( + session.scalars( + sa.select(KnowledgeFSUpgradeFileLease) + .where( + KnowledgeFSUpgradeFileLease.old_upload_file_id.in_(normalized_ids), + KnowledgeFSUpgradeFileLease.status == KnowledgeFSUpgradeFileLeaseStatus.ACTIVE, + KnowledgeFSUpgradeFileLease.expires_at > requested_at, + ) + .with_for_update() + ) + ) + for lease in leases: + if lease.cleanup_requested_at is None: + lease.cleanup_requested_at = requested_at + return {lease.old_upload_file_id for lease in leases} + + +def release_upgrade_file_lease( + session: Session, + *, + job_id: str, + upload_file_id: str, + now: datetime | None = None, +) -> bool: + """Release one lease and fulfill deferred cleanup after the last protection ends.""" + + released_at = now or naive_utc_now() + leases = list( + session.scalars( + sa.select(KnowledgeFSUpgradeFileLease) + .where(KnowledgeFSUpgradeFileLease.old_upload_file_id == upload_file_id) + .with_for_update() + ) + ) + current = next((lease for lease in leases if lease.job_id == job_id), None) + if current is None: + return False + current.status = KnowledgeFSUpgradeFileLeaseStatus.RELEASED + current.released_at = released_at + cleanup_requested = any(lease.cleanup_requested_at is not None for lease in leases) + still_protected = any( + lease.status == KnowledgeFSUpgradeFileLeaseStatus.ACTIVE and lease.expires_at > released_at for lease in leases + ) + if not cleanup_requested or still_protected or _legacy_document_references_file(session, upload_file_id): + return False + upload_file = session.get(UploadFile, upload_file_id) + if upload_file is None: + return False + storage.delete(upload_file.key) + session.delete(upload_file) + return True + + +def cleanup_deferred_upgrade_files( + session_maker: sessionmaker[Session], + *, + limit: int = 100, + now: datetime | None = None, +) -> int: + """Retry bounded cleanup requests after abandoned migration leases expire.""" + + cleanup_at = now or naive_utc_now() + with session_maker() as session: + candidate = aliased(KnowledgeFSUpgradeFileLease) + blocking = aliased(KnowledgeFSUpgradeFileLease) + upload_file_ids = list( + session.scalars( + sa.select(candidate.old_upload_file_id) + .where( + candidate.cleanup_requested_at.is_not(None), + ~sa.exists( + sa.select(blocking.id).where( + blocking.old_upload_file_id == candidate.old_upload_file_id, + blocking.status == KnowledgeFSUpgradeFileLeaseStatus.ACTIVE, + blocking.expires_at > cleanup_at, + ) + ), + ) + .distinct() + .limit(limit) + ) + ) + cleaned = 0 + for upload_file_id in upload_file_ids: + with session_maker.begin() as session: + leases = list( + session.scalars( + sa.select(KnowledgeFSUpgradeFileLease) + .where(KnowledgeFSUpgradeFileLease.old_upload_file_id == upload_file_id) + .with_for_update() + ) + ) + for lease in leases: + if lease.status == KnowledgeFSUpgradeFileLeaseStatus.ACTIVE and lease.expires_at <= cleanup_at: + lease.status = KnowledgeFSUpgradeFileLeaseStatus.EXPIRED + if any( + lease.status == KnowledgeFSUpgradeFileLeaseStatus.ACTIVE and lease.expires_at > cleanup_at + for lease in leases + ): + continue + if _legacy_document_references_file(session, upload_file_id): + for lease in leases: + lease.cleanup_requested_at = None + continue + upload_file = session.get(UploadFile, upload_file_id) + if upload_file is None: + for lease in leases: + lease.cleanup_requested_at = None + continue + storage.delete(upload_file.key) + session.delete(upload_file) + cleaned += 1 + return cleaned + + +def _legacy_document_references_file(session: Session, upload_file_id: str) -> bool: + candidates = session.scalars( + sa.select(Document.data_source_info).where( + Document.data_source_type == "upload_file", + Document.data_source_info.contains(upload_file_id), + ) + ) + for candidate in candidates: + if not candidate: + continue + try: + payload = json.loads(candidate) + except (JSONDecodeError, TypeError): + continue + if isinstance(payload, dict) and str(payload.get("upload_file_id") or "") == upload_file_id: + return True + return False + + +__all__ = [ + "active_upgrade_file_ids", + "cleanup_deferred_upgrade_files", + "release_upgrade_file_lease", + "reserve_upgrade_file_cleanup", +] diff --git a/api/services/dataset_knowledge_fs_upgrade_service.py b/api/services/dataset_knowledge_fs_upgrade_service.py new file mode 100644 index 00000000000..b7d030bfbd9 --- /dev/null +++ b/api/services/dataset_knowledge_fs_upgrade_service.py @@ -0,0 +1,1178 @@ +"""Snapshot and resumable orchestration for legacy Dataset upgrades.""" + +from __future__ import annotations + +import json +import re +import uuid +from collections import defaultdict +from datetime import timedelta +from hashlib import sha256 +from typing import Any, cast + +import sqlalchemy as sa +from pydantic import TypeAdapter +from sqlalchemy.orm import Session, sessionmaker + +from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager +from extensions.ext_storage import storage +from graphon.model_runtime.entities.model_entities import ModelType +from libs.datetime_utils import naive_utc_now +from models import Account, AccountStatus, TenantAccountJoin +from models.dataset import AppDatasetJoin, Dataset, DatasetMetadata, DatasetPermission, Document +from models.knowledge_fs import ( + KnowledgeFSAppSpaceJoinType, + KnowledgeFSControlSpace, + KnowledgeFSControlSpacePermissionRole, + KnowledgeFSControlSpaceState, + KnowledgeFSControlSpaceVisibility, + KnowledgeFSUpgradeDocument, + KnowledgeFSUpgradeFileLease, + KnowledgeFSUpgradeFileLeaseStatus, + KnowledgeFSUpgradeItemStatus, + KnowledgeFSUpgradeJob, + KnowledgeFSUpgradeJobStatus, + KnowledgeFSUpgradeSource, + KnowledgeFSUpgradeStage, +) +from models.model import App, AppMode, Tag, TagBinding, UploadFile +from models.oauth import DatasourceProvider +from models.provider_ids import ModelProviderID +from services.dataset_knowledge_fs_upgrade_file_lease import release_upgrade_file_lease +from services.feature_service import FeatureService +from services.knowledge_fs.product_dto import ( + KnowledgeFSAppBindingPayload, + KnowledgeFSDocumentAvailabilityPayload, + KnowledgeFSDocumentMetadataPayload, + KnowledgeFSDocumentStagedUploadPayload, + KnowledgeFSExternalAccessPayload, + KnowledgeFSInitialSourcePayload, + KnowledgeFSMemberBindingPayload, + KnowledgeFSMetadataFieldCreatePayload, + KnowledgeFSModelIntent, + KnowledgeFSRerankIntent, + KnowledgeFSRetrievalProfileIntent, + KnowledgeFSScoreThresholdIntent, + KnowledgeFSSpaceCreatePayload, + KnowledgeFSUpgradeJobResponse, +) +from services.knowledge_fs.runtime import get_knowledge_fs_runtime +from services.knowledge_fs.staged_upload_service import KnowledgeFSStagedUploadService + +_ACTIVE_JOB_STATUSES = (KnowledgeFSUpgradeJobStatus.QUEUED, KnowledgeFSUpgradeJobStatus.RUNNING) +_SOURCE_SELECTION_LIMIT = 200 +_FILE_LEASE_TTL = timedelta(days=7) +_INITIAL_SOURCE_ADAPTER: TypeAdapter[KnowledgeFSInitialSourcePayload] = TypeAdapter(KnowledgeFSInitialSourcePayload) + + +class KnowledgeFSUpgradeError(RuntimeError): + """A persisted upgrade cannot safely continue.""" + + +class KnowledgeFSUpgradeNotFoundError(KnowledgeFSUpgradeError): + pass + + +class KnowledgeFSUpgradeConflictError(KnowledgeFSUpgradeError): + pass + + +class KnowledgeFSUpgradeNotReadyError(KnowledgeFSUpgradeError): + """Provisioning is progressing and the dedicated worker should retry.""" + + +class KnowledgeFSUpgradeSnapshotService: + """Persist the exact click-time Dataset manifest without external I/O.""" + + def __init__(self, session_maker: sessionmaker[Session]) -> None: + self._session_maker = session_maker + + def create( + self, + *, + tenant_id: str, + dataset_id: str, + requested_by_account_id: str, + idempotency_key: str | None = None, + ) -> KnowledgeFSUpgradeJob: + snapshot_at = naive_utc_now() + with self._session_maker(expire_on_commit=False) as session, session.begin(): + dataset = session.scalar( + sa.select(Dataset).where(Dataset.id == dataset_id, Dataset.tenant_id == tenant_id).with_for_update() + ) + if dataset is None: + raise KnowledgeFSUpgradeNotFoundError("Dataset was not found") + if dataset.provider == "external": + raise KnowledgeFSUpgradeConflictError("External-provider Datasets cannot be upgraded") + + if idempotency_key: + idempotent = session.scalar( + sa.select(KnowledgeFSUpgradeJob).where( + KnowledgeFSUpgradeJob.tenant_id == tenant_id, + KnowledgeFSUpgradeJob.idempotency_key == idempotency_key, + ) + ) + if idempotent is not None: + if idempotent.old_dataset_id != dataset_id: + raise KnowledgeFSUpgradeConflictError("Idempotency key belongs to another Dataset upgrade") + return idempotent + + existing = session.scalar( + sa.select(KnowledgeFSUpgradeJob) + .where( + KnowledgeFSUpgradeJob.tenant_id == tenant_id, + KnowledgeFSUpgradeJob.old_dataset_id == dataset_id, + KnowledgeFSUpgradeJob.status.in_(_ACTIVE_JOB_STATUSES), + ) + .order_by(KnowledgeFSUpgradeJob.created_at.desc()) + .limit(1) + ) + if existing is not None: + return existing + + owner_account_id = str(dataset.maintainer or dataset.created_by) + if ( + session.scalar( + sa.select(Account.id) + .join(TenantAccountJoin, TenantAccountJoin.account_id == Account.id) + .where( + Account.id == owner_account_id, + Account.status != AccountStatus.BANNED, + TenantAccountJoin.tenant_id == tenant_id, + ) + .limit(1) + ) + is None + ): + raise KnowledgeFSUpgradeConflictError("Dataset maintainer is not an active account") + + permission_members = list( + session.scalars( + sa.select(DatasetPermission.account_id) + .join(Account, Account.id == DatasetPermission.account_id) + .join( + TenantAccountJoin, + sa.and_( + TenantAccountJoin.account_id == DatasetPermission.account_id, + TenantAccountJoin.tenant_id == tenant_id, + ), + ) + .where( + DatasetPermission.tenant_id == tenant_id, + DatasetPermission.dataset_id == dataset_id, + DatasetPermission.has_permission.is_(True), + Account.status != AccountStatus.BANNED, + ) + ) + ) + app_bindings = _app_binding_snapshot(session, tenant_id=tenant_id, dataset_id=dataset_id) + tag_ids = list( + session.scalars( + sa.select(TagBinding.tag_id) + .join(Tag, Tag.id == TagBinding.tag_id) + .where( + TagBinding.tenant_id == tenant_id, + TagBinding.target_id == dataset_id, + Tag.type == "knowledge", + ) + ) + ) + job = KnowledgeFSUpgradeJob( + tenant_id=tenant_id, + old_dataset_id=dataset_id, + requested_by_account_id=requested_by_account_id, + owner_account_id=owner_account_id, + idempotency_key=idempotency_key or str(uuid.uuid4()), + snapshot_at=snapshot_at, + config_snapshot=_config_snapshot(session, dataset), + permission_snapshot={ + "visibility": _enum_value(dataset.permission), + "member_account_ids": permission_members, + }, + app_binding_snapshot=app_bindings, + tag_ids_snapshot=[str(tag_id) for tag_id in tag_ids if tag_id is not None], + ) + session.add(job) + session.flush() + + source_groups: dict[str, list[tuple[Document, dict[str, Any]]]] = defaultdict(list) + document_snapshots: dict[str, KnowledgeFSUpgradeDocument] = {} + upload_file_ids: set[str] = set() + documents = list( + session.scalars( + sa.select(Document) + .where(Document.tenant_id == tenant_id, Document.dataset_id == dataset_id) + .order_by(Document.position, Document.id) + ) + ) + for document in documents: + data_source_type = _enum_value(document.data_source_type) + data_source_info = document.data_source_info_dict + old_upload_file_id = None + source_key = None + if data_source_type == "upload_file": + raw_upload_file_id = data_source_info.get("upload_file_id") + if raw_upload_file_id: + old_upload_file_id = str(raw_upload_file_id) + upload_file_ids.add(old_upload_file_id) + elif data_source_type == "notion_import": + source_key = _notion_source_group_key(data_source_info) + source_groups[source_key].append((document, data_source_info)) + elif data_source_type == "website_crawl": + source_key = _website_source_group_key(data_source_info) + source_groups[source_key].append((document, data_source_info)) + else: + raise KnowledgeFSUpgradeConflictError( + f"Unsupported legacy document source type: {data_source_type}" + ) + document_snapshot = KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=tenant_id, + old_document_id=str(document.id), + name=document.name, + data_source_type=data_source_type, + data_source_info=cast(dict[str, object], data_source_info), + metadata_snapshot=cast(dict[str, object], document.doc_metadata or {}), + desired_enabled=bool(document.enabled and not document.archived), + legacy_archived=bool(document.archived), + legacy_indexing_status=_enum_value(document.indexing_status), + legacy_display_status=document.display_status, + old_upload_file_id=old_upload_file_id, + source_key=source_key, + ) + session.add(document_snapshot) + document_snapshots[str(document.id)] = document_snapshot + + source_count = 0 + for base_source_key, group in source_groups.items(): + for chunk_number, chunk in enumerate(_chunks(group, _SOURCE_SELECTION_LIMIT), start=1): + source_key = f"{base_source_key}:{chunk_number}" + payload = _source_payload_snapshot(dataset.name, chunk) + session.add( + KnowledgeFSUpgradeSource( + job_id=job.id, + tenant_id=tenant_id, + source_key=source_key, + payload_snapshot=payload, + ) + ) + for document, _ in chunk: + document_snapshots[str(document.id)].source_key = source_key + source_count += 1 + + for upload_file_id in sorted(upload_file_ids): + session.add( + KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=upload_file_id, + expires_at=snapshot_at + _FILE_LEASE_TTL, + ) + ) + job.total_documents = len(documents) + job.total_sources = source_count + session.flush() + return job + + def get(self, *, tenant_id: str, job_id: str) -> KnowledgeFSUpgradeJob: + with self._session_maker() as session: + job = session.scalar( + sa.select(KnowledgeFSUpgradeJob).where( + KnowledgeFSUpgradeJob.id == job_id, + KnowledgeFSUpgradeJob.tenant_id == tenant_id, + ) + ) + if job is None: + raise KnowledgeFSUpgradeNotFoundError("Upgrade job was not found") + session.expunge(job) + return job + + def retry(self, *, tenant_id: str, job_id: str) -> KnowledgeFSUpgradeJob: + lease_expires_at = naive_utc_now() + _FILE_LEASE_TTL + with self._session_maker(expire_on_commit=False) as session, session.begin(): + job = session.scalar( + sa.select(KnowledgeFSUpgradeJob) + .where(KnowledgeFSUpgradeJob.id == job_id, KnowledgeFSUpgradeJob.tenant_id == tenant_id) + .with_for_update() + ) + if job is None: + raise KnowledgeFSUpgradeNotFoundError("Upgrade job was not found") + if job.status is KnowledgeFSUpgradeJobStatus.SUCCEEDED: + raise KnowledgeFSUpgradeConflictError("Successful upgrades cannot be retried") + if job.status in _ACTIVE_JOB_STATUSES: + return job + job.status = KnowledgeFSUpgradeJobStatus.QUEUED + job.last_error_code = None + job.last_error_message = None + job.completed_at = None + job.celery_task_id = None + session.execute( + sa.update(KnowledgeFSUpgradeDocument) + .where( + KnowledgeFSUpgradeDocument.job_id == job.id, + KnowledgeFSUpgradeDocument.status == KnowledgeFSUpgradeItemStatus.FAILED, + ) + .values(status=KnowledgeFSUpgradeItemStatus.PENDING, last_error_code=None, last_error_message=None) + ) + session.execute( + sa.update(KnowledgeFSUpgradeSource) + .where( + KnowledgeFSUpgradeSource.job_id == job.id, + KnowledgeFSUpgradeSource.status == KnowledgeFSUpgradeItemStatus.FAILED, + ) + .values(status=KnowledgeFSUpgradeItemStatus.PENDING, last_error_code=None, last_error_message=None) + ) + session.execute( + sa.update(KnowledgeFSUpgradeFileLease) + .where( + KnowledgeFSUpgradeFileLease.job_id == job.id, + KnowledgeFSUpgradeFileLease.released_at.is_(None), + ) + .values( + status=KnowledgeFSUpgradeFileLeaseStatus.ACTIVE, + expires_at=lease_expires_at, + ) + ) + return job + + def claim_enqueue(self, *, tenant_id: str, job_id: str, task_id: str) -> bool: + with self._session_maker.begin() as session: + job = session.scalar( + sa.select(KnowledgeFSUpgradeJob) + .where(KnowledgeFSUpgradeJob.id == job_id, KnowledgeFSUpgradeJob.tenant_id == tenant_id) + .with_for_update() + ) + if job is None: + raise KnowledgeFSUpgradeNotFoundError("Upgrade job was not found") + if job.status is not KnowledgeFSUpgradeJobStatus.QUEUED or job.celery_task_id is not None: + return False + job.celery_task_id = task_id + return True + + def release_enqueue_claim(self, *, tenant_id: str, job_id: str, task_id: str) -> None: + with self._session_maker.begin() as session: + job = session.scalar( + sa.select(KnowledgeFSUpgradeJob) + .where(KnowledgeFSUpgradeJob.id == job_id, KnowledgeFSUpgradeJob.tenant_id == tenant_id) + .with_for_update() + ) + if job is not None and job.status is KnowledgeFSUpgradeJobStatus.QUEUED and job.celery_task_id == task_id: + job.celery_task_id = None + + +class KnowledgeFSUpgradeRunner: + """Execute one durable stage or item per dedicated-worker delivery.""" + + def __init__(self, session_maker: sessionmaker[Session]) -> None: + self._session_maker = session_maker + + def run_next(self, *, job_id: str, celery_task_id: str | None = None) -> bool: + job = self._load_job(job_id) + if job.status in {KnowledgeFSUpgradeJobStatus.SUCCEEDED, KnowledgeFSUpgradeJobStatus.FAILED}: + return False + self._mark_running(job_id=job_id, celery_task_id=celery_task_id) + if job.stage is KnowledgeFSUpgradeStage.VALIDATING: + self._create_space(job) + raise KnowledgeFSUpgradeNotReadyError("KnowledgeFS Space provisioning is pending") + if job.stage is KnowledgeFSUpgradeStage.WAITING_FOR_SPACE: + self._advance_when_space_is_active(job) + return True + if job.stage is KnowledgeFSUpgradeStage.CREATING_SOURCES: + return self._create_next_source(job) + if job.stage is KnowledgeFSUpgradeStage.SUBMITTING_DOCUMENTS: + return self._submit_next_document(job) + if job.stage is KnowledgeFSUpgradeStage.MIGRATING_ACCESS: + self._migrate_access(job) + return True + if job.stage is KnowledgeFSUpgradeStage.FINALIZING: + self._finalize(job) + return False + return False + + def fail(self, *, job_id: str, error: Exception) -> None: + with self._session_maker.begin() as session: + job = session.get(KnowledgeFSUpgradeJob, job_id) + if job is None or job.status is KnowledgeFSUpgradeJobStatus.SUCCEEDED: + return + job.status = KnowledgeFSUpgradeJobStatus.FAILED + job.last_error_code = type(error).__name__[:128] + job.last_error_message = str(error)[:4000] + job.completed_at = naive_utc_now() + + def _load_job(self, job_id: str) -> KnowledgeFSUpgradeJob: + with self._session_maker() as session: + job = session.get(KnowledgeFSUpgradeJob, job_id) + if job is None: + raise KnowledgeFSUpgradeNotFoundError("Upgrade job was not found") + session.expunge(job) + return job + + def _mark_running(self, *, job_id: str, celery_task_id: str | None) -> None: + lease_expires_at = naive_utc_now() + _FILE_LEASE_TTL + with self._session_maker.begin() as session: + job = session.get(KnowledgeFSUpgradeJob, job_id) + if job is None: + raise KnowledgeFSUpgradeNotFoundError("Upgrade job was not found") + job.status = KnowledgeFSUpgradeJobStatus.RUNNING + job.attempt_count += 1 + job.celery_task_id = celery_task_id + session.execute( + sa.update(KnowledgeFSUpgradeFileLease) + .where( + KnowledgeFSUpgradeFileLease.job_id == job_id, + KnowledgeFSUpgradeFileLease.status == KnowledgeFSUpgradeFileLeaseStatus.ACTIVE, + ) + .values(expires_at=lease_expires_at) + ) + + def _create_space(self, job: KnowledgeFSUpgradeJob) -> None: + resolved = _resolve_configuration(job) + runtime = get_knowledge_fs_runtime(self._session_maker) + response = runtime.application.create_space( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + payload=KnowledgeFSSpaceCreatePayload( + name=str(job.config_snapshot["name"])[:40], + slug=f"legacy-{job.old_dataset_id[:12]}-{job.id[:8]}", + description=str(job.config_snapshot.get("description") or "")[:2000] or None, + icon=cast(str | None, job.config_snapshot.get("icon")), + visibility=KnowledgeFSControlSpaceVisibility.ONLY_ME, + embedding=KnowledgeFSModelIntent.model_validate(resolved["embedding"]), + retrieval=KnowledgeFSRetrievalProfileIntent.model_validate(resolved["retrieval"]), + idempotency_key=f"upgrade:{job.id}:space", + ), + ) + with self._session_maker.begin() as session: + persisted = session.get(KnowledgeFSUpgradeJob, job.id) + if persisted is None: + raise KnowledgeFSUpgradeNotFoundError("Upgrade job disappeared") + persisted.new_control_space_id = response.control_space_id + persisted.resolved_configuration = resolved + persisted.stage = KnowledgeFSUpgradeStage.WAITING_FOR_SPACE + + def _advance_when_space_is_active(self, job: KnowledgeFSUpgradeJob) -> None: + if job.new_control_space_id is None: + raise KnowledgeFSUpgradeError("Upgrade Space reference is missing") + with self._session_maker.begin() as session: + control_space = session.scalar( + sa.select(KnowledgeFSControlSpace).where( + KnowledgeFSControlSpace.id == job.new_control_space_id, + KnowledgeFSControlSpace.tenant_id == job.tenant_id, + ) + ) + if control_space is None: + raise KnowledgeFSUpgradeError("Upgrade Space was not found") + if control_space.state is KnowledgeFSControlSpaceState.PROVISIONING: + raise KnowledgeFSUpgradeNotReadyError("KnowledgeFS Space provisioning is pending") + if control_space.state is not KnowledgeFSControlSpaceState.ACTIVE: + raise KnowledgeFSUpgradeError(f"KnowledgeFS Space provisioning failed in {control_space.state.value}") + persisted = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted is not None + persisted.stage = ( + KnowledgeFSUpgradeStage.CREATING_SOURCES + if persisted.total_sources + else KnowledgeFSUpgradeStage.SUBMITTING_DOCUMENTS + ) + + def _create_next_source(self, job: KnowledgeFSUpgradeJob) -> bool: + with self._session_maker.begin() as session: + source = session.scalar( + sa.select(KnowledgeFSUpgradeSource) + .where( + KnowledgeFSUpgradeSource.job_id == job.id, + KnowledgeFSUpgradeSource.status.in_( + (KnowledgeFSUpgradeItemStatus.PENDING, KnowledgeFSUpgradeItemStatus.PROCESSING) + ), + ) + .order_by(KnowledgeFSUpgradeSource.id) + .with_for_update(skip_locked=True) + .limit(1) + ) + if source is None: + persisted = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted is not None + persisted.stage = KnowledgeFSUpgradeStage.SUBMITTING_DOCUMENTS + return True + source.status = KnowledgeFSUpgradeItemStatus.PROCESSING + source_id = source.id + source_payload = dict(source.payload_snapshot) + payload = self._resolve_initial_source_payload(job, source_payload) + from tasks.knowledge_fs_initial_source_tasks import submit_initial_source_for_upgrade + + try: + result = submit_initial_source_for_upgrade( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + control_space_id=_required_space_id(job), + operation_id=f"upgrade:{job.id}:{source_id}", + payload=payload, + ) + except KnowledgeFSUpgradeNotReadyError: + raise + except Exception as error: + from tasks.knowledge_fs_initial_source_tasks import KnowledgeFSInitialSourceNotReadyError + + if isinstance(error, KnowledgeFSInitialSourceNotReadyError): + with self._session_maker.begin() as session: + persisted_source = session.get(KnowledgeFSUpgradeSource, source_id) + if persisted_source is not None: + persisted_source.status = KnowledgeFSUpgradeItemStatus.PENDING + raise KnowledgeFSUpgradeNotReadyError(str(error)) from error + self._fail_source(source_id=source_id, error=error) + raise + with self._session_maker.begin() as session: + persisted_source = session.get(KnowledgeFSUpgradeSource, source_id) + persisted_job = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted_source is not None + assert persisted_job is not None + persisted_source.status = KnowledgeFSUpgradeItemStatus.SUCCEEDED + persisted_source.new_connection_id = result.connection_id + persisted_source.new_source_id = result.source_id + persisted_source.initial_sync_task_id = result.workflow_id + if result.workflow_error: + persisted_source.last_error_code = result.workflow_error + persisted_source.last_error_message = "The Source exists; retry its first import in KnowledgeFS" + source_documents = list( + session.scalars( + sa.select(KnowledgeFSUpgradeDocument).where( + KnowledgeFSUpgradeDocument.job_id == job.id, + KnowledgeFSUpgradeDocument.source_key == persisted_source.source_key, + ) + ) + ) + for document in source_documents: + if document.status is not KnowledgeFSUpgradeItemStatus.SUCCEEDED: + document.status = KnowledgeFSUpgradeItemStatus.SUCCEEDED + persisted_job.completed_documents += 1 + persisted_job.completed_sources += 1 + return True + + def _resolve_initial_source_payload( + self, job: KnowledgeFSUpgradeJob, payload_snapshot: dict[str, object] + ) -> KnowledgeFSInitialSourcePayload: + payload = dict(payload_snapshot) + legacy_credential_id = payload.pop("legacy_credential_id", None) + workspace_id = payload.pop("legacy_workspace_id", None) + if payload.get("kind") == "online_document": + payload["credential_id"] = self._resolve_notion_credential( + tenant_id=job.tenant_id, + legacy_credential_id=str(legacy_credential_id or ""), + workspace_id=str(workspace_id or ""), + ) + return _INITIAL_SOURCE_ADAPTER.validate_python(payload) + + def _resolve_notion_credential(self, *, tenant_id: str, legacy_credential_id: str, workspace_id: str) -> str: + with self._session_maker() as session: + direct = session.scalar( + sa.select(DatasourceProvider).where( + DatasourceProvider.id == legacy_credential_id, + DatasourceProvider.tenant_id == tenant_id, + DatasourceProvider.plugin_id == "langgenius/notion_datasource", + DatasourceProvider.provider == "notion_datasource", + ) + ) + if direct is not None: + return direct.id + candidates = list( + session.scalars( + sa.select(DatasourceProvider) + .where( + DatasourceProvider.tenant_id == tenant_id, + DatasourceProvider.plugin_id == "langgenius/notion_datasource", + DatasourceProvider.provider == "notion_datasource", + ) + .order_by(DatasourceProvider.is_default.desc(), DatasourceProvider.created_at) + ) + ) + for candidate in candidates: + if str(candidate.encrypted_credentials.get("workspace_id") or "") == workspace_id: + return candidate.id + if len(candidates) == 1: + return candidates[0].id + raise KnowledgeFSUpgradeError("A matching Notion datasource credential is unavailable") + + def _fail_source(self, *, source_id: str, error: Exception) -> None: + with self._session_maker.begin() as session: + source = session.get(KnowledgeFSUpgradeSource, source_id) + if source is not None: + source.status = KnowledgeFSUpgradeItemStatus.FAILED + source.last_error_code = type(error).__name__[:128] + source.last_error_message = str(error)[:4000] + + def _submit_next_document(self, job: KnowledgeFSUpgradeJob) -> bool: + with self._session_maker.begin() as session: + document = session.scalar( + sa.select(KnowledgeFSUpgradeDocument) + .where( + KnowledgeFSUpgradeDocument.job_id == job.id, + KnowledgeFSUpgradeDocument.data_source_type == "upload_file", + KnowledgeFSUpgradeDocument.status.in_( + (KnowledgeFSUpgradeItemStatus.PENDING, KnowledgeFSUpgradeItemStatus.PROCESSING) + ), + ) + .order_by(KnowledgeFSUpgradeDocument.id) + .with_for_update(skip_locked=True) + .limit(1) + ) + if document is None: + persisted = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted is not None + persisted.stage = KnowledgeFSUpgradeStage.MIGRATING_ACCESS + return True + document.status = KnowledgeFSUpgradeItemStatus.PROCESSING + document_id = document.id + upload_file_id = document.old_upload_file_id + staged_upload_id = document.staged_upload_id + if not upload_file_id: + error = KnowledgeFSUpgradeError("Legacy upload document has no source file reference") + self._fail_document(document_id=document_id, error=error) + raise error + try: + staged = KnowledgeFSStagedUploadService( + self._session_maker, + facade=get_knowledge_fs_runtime(self._session_maker).facade, + ) + if staged_upload_id is None: + with self._session_maker() as session: + upload_file = session.scalar( + sa.select(UploadFile).where( + UploadFile.id == upload_file_id, + UploadFile.tenant_id == job.tenant_id, + ) + ) + account = session.get(Account, job.owner_account_id) + if upload_file is None or account is None: + raise KnowledgeFSUpgradeError("Legacy source file is unavailable") + storage_key = upload_file.key + file_name = upload_file.name + content_type = upload_file.mime_type or "application/octet-stream" + session.expunge(account) + body = storage.load(storage_key) + if not isinstance(body, bytes): + raise KnowledgeFSUpgradeError("Legacy source file returned an invalid body") + staged_response = staged.stage( + tenant_id=job.tenant_id, + account=account, + file_name=file_name, + content_type=content_type, + body=body, + file_size_limit_mb=FeatureService.get_knowledge_file_size_limit(job.tenant_id), + ) + staged_upload_id = staged_response.id + with self._session_maker.begin() as session: + persisted_document = session.get(KnowledgeFSUpgradeDocument, document_id) + assert persisted_document is not None + persisted_document.staged_upload_id = staged_upload_id + claimed = staged.claim( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + control_space_id=_required_space_id(job), + payload=KnowledgeFSDocumentStagedUploadPayload(upload_id=staged_upload_id), + ) + except Exception as error: + self._fail_document(document_id=document_id, error=error) + raise + with self._session_maker.begin() as session: + persisted_document = session.get(KnowledgeFSUpgradeDocument, document_id) + persisted_job = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted_document is not None + assert persisted_job is not None + persisted_document.status = KnowledgeFSUpgradeItemStatus.SUCCEEDED + persisted_document.new_document_asset_id = claimed.document_asset_id + persisted_document.compilation_job_id = claimed.compilation_job_id + persisted_job.completed_documents += 1 + remaining_references = session.scalar( + sa.select(sa.func.count(KnowledgeFSUpgradeDocument.id)).where( + KnowledgeFSUpgradeDocument.job_id == job.id, + KnowledgeFSUpgradeDocument.old_upload_file_id == upload_file_id, + KnowledgeFSUpgradeDocument.status != KnowledgeFSUpgradeItemStatus.SUCCEEDED, + ) + ) + if remaining_references == 0: + release_upgrade_file_lease( + session, + job_id=job.id, + upload_file_id=upload_file_id, + ) + return True + + def _fail_document(self, *, document_id: str, error: Exception) -> None: + with self._session_maker.begin() as session: + document = session.get(KnowledgeFSUpgradeDocument, document_id) + if document is not None: + document.status = KnowledgeFSUpgradeItemStatus.FAILED + document.last_error_code = type(error).__name__[:128] + document.last_error_message = str(error)[:4000] + + def _migrate_access(self, job: KnowledgeFSUpgradeJob) -> None: + control_space_id = _required_space_id(job) + runtime = get_knowledge_fs_runtime(self._session_maker) + _migrate_metadata_fields(job, runtime.facade) + members = [ + KnowledgeFSMemberBindingPayload( + account_id=str(account_id), + role=KnowledgeFSControlSpacePermissionRole.VIEWER, + ) + for account_id in cast(list[str], job.permission_snapshot.get("member_account_ids", [])) + if str(account_id) != job.owner_account_id + ] + runtime.control_plane.replace_members( + tenant_id=job.tenant_id, + actor_account_id=job.owner_account_id, + control_space_id=control_space_id, + members=members, + ) + runtime.control_plane.update_visibility( + tenant_id=job.tenant_id, + actor_account_id=job.owner_account_id, + control_space_id=control_space_id, + visibility=KnowledgeFSControlSpaceVisibility(str(job.permission_snapshot["visibility"])), + ) + has_agent = any(binding["caller_kind"] == "agent" for binding in job.app_binding_snapshot) + has_workflow = any(binding["caller_kind"] == "workflow" for binding in job.app_binding_snapshot) + runtime.control_plane.update_external_access( + tenant_id=job.tenant_id, + actor_account_id=job.owner_account_id, + control_space_id=control_space_id, + payload=KnowledgeFSExternalAccessPayload( + service_api_enabled=bool(job.config_snapshot.get("enable_api")), + agent_enabled=has_agent, + workflow_enabled=has_workflow, + mcp_enabled=False, + ), + ) + for binding in job.app_binding_snapshot: + runtime.app_bindings.upsert( + tenant_id=job.tenant_id, + actor_account_id=job.owner_account_id, + control_space_id=control_space_id, + payload=KnowledgeFSAppBindingPayload( + app_id=str(binding["app_id"]), + caller_kind=KnowledgeFSAppSpaceJoinType(str(binding["caller_kind"])), + ), + ) + runtime.space_tags.replace_tags( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + control_space_id=control_space_id, + tag_ids=list(job.tag_ids_snapshot), + ) + with self._session_maker.begin() as session: + persisted = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted is not None + persisted.stage = KnowledgeFSUpgradeStage.FINALIZING + + def _finalize(self, job: KnowledgeFSUpgradeJob) -> None: + with self._session_maker.begin() as session: + persisted = session.scalar( + sa.select(KnowledgeFSUpgradeJob).where(KnowledgeFSUpgradeJob.id == job.id).with_for_update() + ) + if persisted is None: + raise KnowledgeFSUpgradeNotFoundError("Upgrade job disappeared") + if persisted.completed_documents != persisted.total_documents: + raise KnowledgeFSUpgradeError("Not all Dataset documents were handed off") + if persisted.completed_sources != persisted.total_sources: + raise KnowledgeFSUpgradeError("Not all Dataset Sources were created") + persisted.status = KnowledgeFSUpgradeJobStatus.SUCCEEDED + persisted.stage = KnowledgeFSUpgradeStage.COMPLETED + persisted.completed_at = naive_utc_now() + persisted.last_error_code = None + persisted.last_error_message = None + + +class KnowledgeFSUpgradeDocumentReconciler: + """Apply snapshot metadata and availability after new logical documents appear. + + This is deliberately separate from parent-job success: KnowledgeFS owns parsing and + indexing task outcomes, while this best-effort loop preserves the legacy document's + click-time availability once the new logical document can be addressed. + """ + + def __init__(self, session_maker: sessionmaker[Session]) -> None: + self._session_maker = session_maker + + def reconcile(self, *, job_id: str) -> int: + with self._session_maker() as session: + job = session.get(KnowledgeFSUpgradeJob, job_id) + if job is None or job.status is not KnowledgeFSUpgradeJobStatus.SUCCEEDED: + return 0 + documents = list( + session.scalars( + sa.select(KnowledgeFSUpgradeDocument) + .where( + KnowledgeFSUpgradeDocument.job_id == job_id, + KnowledgeFSUpgradeDocument.status == KnowledgeFSUpgradeItemStatus.SUCCEEDED, + KnowledgeFSUpgradeDocument.state_reconciled_at.is_(None), + ) + .order_by(KnowledgeFSUpgradeDocument.id) + ) + ) + sources = { + source.source_key: source.new_source_id + for source in session.scalars( + sa.select(KnowledgeFSUpgradeSource).where(KnowledgeFSUpgradeSource.job_id == job_id) + ) + } + session.expunge(job) + for document in documents: + session.expunge(document) + + if not documents: + return 0 + runtime = get_knowledge_fs_runtime(self._session_maker) + logical_documents = _list_all_logical_documents(job, runtime.facade) + by_asset_id = { + logical_document.active.document_asset_id: logical_document + for logical_document in logical_documents + if logical_document.active is not None + } + by_source_item = { + (logical_document.source_id, logical_document.provider_item_id): logical_document + for logical_document in logical_documents + if logical_document.source_id and logical_document.provider_item_id + } + remaining = 0 + for document in documents: + logical_document = None + if document.new_document_asset_id: + logical_document = by_asset_id.get(document.new_document_asset_id) + elif document.source_key: + logical_document = by_source_item.get( + (sources.get(document.source_key), _expected_provider_item_id(document)) + ) + if logical_document is None: + self._record_reconcile_wait(document.id, "The new logical document is not visible yet") + remaining += 1 + continue + try: + current = logical_document + if document.metadata_snapshot: + current = runtime.facade.update_document_metadata( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + control_space_id=_required_space_id(job), + document_id=current.id, + payload=KnowledgeFSDocumentMetadataPayload( + expectedRowVersion=current.row_version, + patch=dict(document.metadata_snapshot), + ), + ) + if current.enabled != document.desired_enabled: + current = runtime.facade.update_logical_document_availability( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + control_space_id=_required_space_id(job), + document_id=current.id, + payload=KnowledgeFSDocumentAvailabilityPayload( + enabled=document.desired_enabled, + expectedRowVersion=current.row_version, + ), + ) + except Exception as error: + self._record_reconcile_wait(document.id, str(error)) + remaining += 1 + continue + with self._session_maker.begin() as session: + persisted = session.get(KnowledgeFSUpgradeDocument, document.id) + if persisted is not None and persisted.state_reconciled_at is None: + persisted.new_logical_document_id = current.id + persisted.state_reconciled_at = naive_utc_now() + persisted.state_reconcile_error = None + persisted.state_reconcile_attempt_count += 1 + return remaining + + def _record_reconcile_wait(self, document_id: str, message: str) -> None: + with self._session_maker.begin() as session: + document = session.get(KnowledgeFSUpgradeDocument, document_id) + if document is not None and document.state_reconciled_at is None: + document.state_reconcile_attempt_count += 1 + document.state_reconcile_error = message[:4000] + + +def upgrade_job_response(job: KnowledgeFSUpgradeJob) -> KnowledgeFSUpgradeJobResponse: + return KnowledgeFSUpgradeJobResponse( + id=job.id, + old_dataset_id=job.old_dataset_id, + new_control_space_id=job.new_control_space_id, + status=job.status, + stage=job.stage, + snapshot_at=job.snapshot_at, + total_documents=job.total_documents, + completed_documents=job.completed_documents, + total_sources=job.total_sources, + completed_sources=job.completed_sources, + last_error_code=job.last_error_code, + last_error_message=job.last_error_message, + completed_at=job.completed_at, + ) + + +def _migrate_metadata_fields(job: KnowledgeFSUpgradeJob, facade: Any) -> None: + expected_fields = cast(list[dict[str, str]], job.config_snapshot.get("metadata_fields") or []) + if not expected_fields: + return + existing_by_name: dict[str, str] = {} + cursor = None + while True: + page = facade.list_metadata_fields( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + control_space_id=_required_space_id(job), + cursor=cursor, + limit=100, + ) + existing_by_name.update({field.name: field.type for field in page.data}) + cursor = page.next_cursor + if cursor is None: + break + for field in expected_fields: + name = str(field["name"]) + field_type = str(field["type"]) + existing_type = existing_by_name.get(name) + if existing_type is not None: + if existing_type != field_type: + raise KnowledgeFSUpgradeError(f"Metadata field {name!r} has an incompatible type") + continue + facade.create_metadata_field( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + control_space_id=_required_space_id(job), + payload=KnowledgeFSMetadataFieldCreatePayload.model_validate({"name": name, "type": field_type}), + ) + existing_by_name[name] = field_type + + +def _list_all_logical_documents(job: KnowledgeFSUpgradeJob, facade: Any) -> list[Any]: + documents: list[Any] = [] + cursor = None + while True: + page = facade.list_logical_documents( + tenant_id=job.tenant_id, + account_id=job.owner_account_id, + control_space_id=_required_space_id(job), + cursor=cursor, + ) + documents.extend(page.data) + cursor = page.next_cursor + if cursor is None: + return documents + + +def _expected_provider_item_id(document: KnowledgeFSUpgradeDocument) -> str: + info = document.data_source_info + if document.data_source_type == "notion_import": + workspace_id = str(info.get("notion_workspace_id") or info.get("workspace_id") or "") + page_id = str(info.get("notion_page_id") or "") + return json.dumps([workspace_id, page_id], separators=(",", ":")) + if document.data_source_type == "website_crawl": + return sha256(str(info.get("url") or "").encode()).hexdigest() + raise KnowledgeFSUpgradeError(f"Document {document.old_document_id} has no provider item identity") + + +def _resolve_configuration(job: KnowledgeFSUpgradeJob) -> dict[str, object]: + provider_manager = create_plugin_provider_manager(tenant_id=job.tenant_id) + embedding = _resolve_model( + provider_manager, + tenant_id=job.tenant_id, + model_type=ModelType.TEXT_EMBEDDING, + preferred_provider=cast(str | None, job.config_snapshot.get("embedding_model_provider")), + preferred_model=cast(str | None, job.config_snapshot.get("embedding_model")), + ) + retrieval = cast(dict[str, Any], job.config_snapshot.get("retrieval_model") or {}) + reranking_model = cast(dict[str, Any], retrieval.get("reranking_model") or {}) + rerank = _resolve_model( + provider_manager, + tenant_id=job.tenant_id, + model_type=ModelType.RERANK, + preferred_provider=cast(str | None, reranking_model.get("reranking_provider_name")), + preferred_model=cast(str | None, reranking_model.get("reranking_model_name")), + ) + reasoning = _resolve_model( + provider_manager, + tenant_id=job.tenant_id, + model_type=ModelType.LLM, + preferred_provider=None, + preferred_model=None, + ) + threshold_enabled = bool(retrieval.get("score_threshold_enabled")) + threshold_value = retrieval.get("score_threshold") + profile = KnowledgeFSRetrievalProfileIntent( + default_mode="fast", + reasoning_model=reasoning, + rerank=KnowledgeFSRerankIntent(enabled=True, model=rerank), + score_threshold=KnowledgeFSScoreThresholdIntent( + enabled=threshold_enabled, + stage="mode-final", + value=float(threshold_value) if threshold_enabled and threshold_value is not None else None, + ), + top_k=max(1, min(100, int(retrieval.get("top_k") or 4))), + ) + return { + "embedding": embedding.model_dump(mode="json", by_alias=True), + "retrieval": profile.model_dump(mode="json", by_alias=True), + } + + +def _resolve_model( + provider_manager, + *, + tenant_id: str, + model_type: ModelType, + preferred_provider: str | None, + preferred_model: str | None, +) -> KnowledgeFSModelIntent: + active_models = provider_manager.get_configurations(tenant_id).get_models(model_type=model_type, only_active=True) + if preferred_provider and preferred_model: + preferred_provider_id = str(ModelProviderID(preferred_provider)) + for active_model in active_models: + active_provider_id = str(ModelProviderID(active_model.provider.provider)) + if active_model.model == preferred_model and active_provider_id == preferred_provider_id: + return _model_intent(active_model.provider.provider, active_model.model) + default = provider_manager.get_default_model(tenant_id=tenant_id, model_type=model_type) + if default is None: + raise KnowledgeFSUpgradeError(f"Workspace default {model_type.value} model is unavailable") + default_provider_id = str(ModelProviderID(default.provider.provider)) + for active_model in active_models: + if ( + active_model.model == default.model + and str(ModelProviderID(active_model.provider.provider)) == default_provider_id + ): + return _model_intent(active_model.provider.provider, active_model.model) + raise KnowledgeFSUpgradeError(f"Workspace default {model_type.value} model is not active") + + +def _model_intent(provider: str, model: str) -> KnowledgeFSModelIntent: + provider_id = ModelProviderID(provider) + return KnowledgeFSModelIntent(plugin_id=provider_id.plugin_id, provider=provider_id.provider_name, model=model) + + +def _config_snapshot(session: Session, dataset: Dataset) -> dict[str, object]: + icon = None + if isinstance(dataset.icon_info, dict): + candidate = dataset.icon_info.get("icon") + if isinstance(candidate, str) and re.fullmatch(r"(?:builtin:)?[+a-z0-9_-]{1,64}", candidate): + icon = candidate + return { + "name": dataset.name, + "description": dataset.description, + "icon": icon, + "indexing_technique": _enum_value(dataset.indexing_technique) if dataset.indexing_technique else None, + "embedding_model": dataset.embedding_model, + "embedding_model_provider": dataset.embedding_model_provider, + "retrieval_model": dataset.retrieval_model or {}, + "summary_index_setting": dataset.summary_index_setting or {}, + "built_in_field_enabled": dataset.built_in_field_enabled, + "enable_api": dataset.enable_api, + "metadata_fields": [ + {"name": metadata.name, "type": _enum_value(metadata.type)} + for metadata in session.scalars( + sa.select(DatasetMetadata) + .where( + DatasetMetadata.tenant_id == dataset.tenant_id, + DatasetMetadata.dataset_id == dataset.id, + ) + .order_by(DatasetMetadata.created_at, DatasetMetadata.id) + ) + ], + } + + +def _app_binding_snapshot(session: Session, *, tenant_id: str, dataset_id: str) -> list[dict[str, object]]: + result: list[dict[str, object]] = [] + rows = session.execute( + sa.select(AppDatasetJoin.app_id, App.mode) + .join(App, App.id == AppDatasetJoin.app_id) + .where(AppDatasetJoin.dataset_id == dataset_id, App.tenant_id == tenant_id) + .order_by(AppDatasetJoin.app_id) + ) + for app_id, mode in rows: + if mode in {AppMode.AGENT, AppMode.AGENT_CHAT}: + caller_kind = KnowledgeFSAppSpaceJoinType.AGENT + elif mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: + caller_kind = KnowledgeFSAppSpaceJoinType.WORKFLOW + else: + continue + result.append({"app_id": str(app_id), "caller_kind": caller_kind.value}) + return result + + +def _notion_source_group_key(info: dict[str, Any]) -> str: + workspace_id = str(info.get("notion_workspace_id") or info.get("workspace_id") or "") + credential_id = str(info.get("credential_id") or "") + if not workspace_id: + raise KnowledgeFSUpgradeConflictError("Notion document has no workspace identity") + return f"notion:{workspace_id}:{credential_id}" + + +def _website_source_group_key(info: dict[str, Any]) -> str: + provider = str(info.get("provider") or "").strip().lower() + job_id = str(info.get("job_id") or "") + if not provider or not info.get("url"): + raise KnowledgeFSUpgradeConflictError("Website document has incomplete source identity") + return f"website:{provider}:{job_id}" + + +def _source_payload_snapshot(dataset_name: str, group: list[tuple[Document, dict[str, Any]]]) -> dict[str, object]: + first_document, first_info = group[0] + if _enum_value(first_document.data_source_type) == "notion_import": + workspace_id = str(first_info.get("notion_workspace_id") or first_info.get("workspace_id") or "") + return { + "kind": "online_document", + "name": f"{dataset_name} Notion"[:200], + "plugin_id": "langgenius/notion_datasource", + "provider": "notion_datasource", + "datasource": "notion_datasource", + "parameters": {}, + "sync_policy": "manual", + "legacy_credential_id": str(first_info.get("credential_id") or ""), + "legacy_workspace_id": workspace_id, + "selection": [ + { + "name": document.name, + "page_id": str(info.get("notion_page_id") or ""), + "provider_item_id": json.dumps( + [workspace_id, str(info.get("notion_page_id") or "")], separators=(",", ":") + ), + "type": str(info.get("type") or "page"), + "workspace_id": workspace_id, + } + for document, info in group + ], + } + provider = str(first_info.get("provider") or "") + return { + "kind": "website_crawl", + "name": f"{dataset_name} Web"[:200], + "provider": provider, + "datasource": "crawl", + "parameters": {"only_main_content": bool(first_info.get("only_main_content", True))}, + "root_url": str(first_info.get("url")), + "crawl_options": {"include_subpages": False, "limit": min(200, len(group))}, + "selection": [{"source_url": str(info.get("url")), "title": document.name} for document, info in group], + "sync_policy": "manual", + } + + +def _chunks[T](items: list[T], size: int) -> list[list[T]]: + return [items[index : index + size] for index in range(0, len(items), size)] + + +def _enum_value(value: object) -> str: + return str(getattr(value, "value", value)) + + +def _required_space_id(job: KnowledgeFSUpgradeJob) -> str: + if job.new_control_space_id is None: + raise KnowledgeFSUpgradeError("Upgrade Space reference is missing") + return job.new_control_space_id + + +__all__ = [ + "KnowledgeFSUpgradeConflictError", + "KnowledgeFSUpgradeDocumentReconciler", + "KnowledgeFSUpgradeError", + "KnowledgeFSUpgradeNotFoundError", + "KnowledgeFSUpgradeNotReadyError", + "KnowledgeFSUpgradeRunner", + "KnowledgeFSUpgradeSnapshotService", + "upgrade_job_response", +] diff --git a/api/services/knowledge_fs/product_dto.py b/api/services/knowledge_fs/product_dto.py index 41186dbeb85..3010c0cf99c 100644 --- a/api/services/knowledge_fs/product_dto.py +++ b/api/services/knowledge_fs/product_dto.py @@ -17,6 +17,8 @@ from models.knowledge_fs import ( KnowledgeFSControlSpacePermissionRole, KnowledgeFSControlSpaceState, KnowledgeFSControlSpaceVisibility, + KnowledgeFSUpgradeJobStatus, + KnowledgeFSUpgradeStage, ) from services.knowledge_fs.product_operations import KnowledgeFSProductPermission @@ -294,6 +296,27 @@ class KnowledgeFSSpaceCreatePayload(BaseModel): return self +class KnowledgeFSUpgradeJobResponse(ResponseModel): + id: str + old_dataset_id: str + new_control_space_id: str | None = None + status: KnowledgeFSUpgradeJobStatus + stage: KnowledgeFSUpgradeStage + snapshot_at: datetime + total_documents: int = Field(ge=0) + completed_documents: int = Field(ge=0) + total_sources: int = Field(ge=0) + completed_sources: int = Field(ge=0) + last_error_code: str | None = None + last_error_message: str | None = None + completed_at: datetime | None = None + + +class KnowledgeFSUpgradeRetryResponse(ResponseModel): + id: str + status: Literal["queued"] = "queued" + + class KnowledgeFSSpaceUpdatePayload(BaseModel): name: str | None = Field(default=None, min_length=1, max_length=40) icon: KnowledgeFSIconIdentity | None = None diff --git a/api/tasks/batch_clean_document_task.py b/api/tasks/batch_clean_document_task.py index 11cf4b9835c..9ecab3dc699 100644 --- a/api/tasks/batch_clean_document_task.py +++ b/api/tasks/batch_clean_document_task.py @@ -13,6 +13,7 @@ from core.tools.utils.web_reader_tool import get_image_upload_file_ids from extensions.ext_storage import storage from models.dataset import Dataset, DatasetMetadataBinding, DocumentSegment from models.model import UploadFile +from services.dataset_knowledge_fs_upgrade_file_lease import reserve_upgrade_file_cleanup from tasks.refresh_billing_vector_space_task import schedule_billing_vector_space_refresh logger = logging.getLogger(__name__) @@ -47,10 +48,11 @@ def batch_clean_document_task( segment_ids: list[str] = [] total_image_upload_file_ids: list[str] = [] dataset_tenant_id: str | None = None + deletable_file_ids = list(file_ids) try: # ============ Step 1: Query segment and file data (short read-only transaction) ============ - with session_factory.create_session() as session: + with session_factory.create_session() as session, session.begin(): # Get segments info segments = session.scalars( select(DocumentSegment).where(DocumentSegment.document_id.in_(document_ids)) @@ -74,7 +76,15 @@ def batch_clean_document_task( # Query storage keys for document files if file_ids: - files = session.scalars(select(UploadFile).where(UploadFile.id.in_(file_ids))).all() + leased_file_ids = reserve_upgrade_file_cleanup(session, file_ids) + deletable_file_ids = [file_id for file_id in file_ids if file_id not in leased_file_ids] + if leased_file_ids: + logger.info( + "Keep %d source files while KnowledgeFS upgrade leases are active, dataset_id=%s", + len(leased_file_ids), + dataset_id, + ) + files = session.scalars(select(UploadFile).where(UploadFile.id.in_(deletable_file_ids))).all() storage_keys_to_delete.extend([f.key for f in files if f and f.key]) # ============ Step 2: Clean vector index (external service, fresh session for dataset) ============ @@ -182,17 +192,17 @@ def batch_clean_document_task( ) # ============ Step 6: Delete document-associated files (separate short transaction) ============ - if file_ids: + if deletable_file_ids: try: with session_factory.create_session() as session: - stmt = delete(UploadFile).where(UploadFile.id.in_(file_ids)) + stmt = delete(UploadFile).where(UploadFile.id.in_(deletable_file_ids)) session.execute(stmt) session.commit() except Exception: logger.exception( "Failed to delete document UploadFile records for dataset_id: %s, file_ids: %s", dataset_id, - file_ids, + deletable_file_ids, ) # ============ Step 7: Delete storage files (I/O operations, no DB transaction) ============ diff --git a/api/tasks/clean_dataset_task.py b/api/tasks/clean_dataset_task.py index 5bf8784e3c2..04c277daf37 100644 --- a/api/tasks/clean_dataset_task.py +++ b/api/tasks/clean_dataset_task.py @@ -24,6 +24,7 @@ from models.dataset import ( ) from models.model import UploadFile from models.workflow import Workflow +from services.dataset_knowledge_fs_upgrade_file_lease import reserve_upgrade_file_cleanup from tasks.refresh_billing_vector_space_task import schedule_billing_vector_space_refresh logger = logging.getLogger(__name__) @@ -181,11 +182,19 @@ def clean_dataset_task( if data_source_info and "upload_file_id" in data_source_info: file_id = data_source_info["upload_file_id"] file_ids.append(file_id) - files = session.scalars(select(UploadFile).where(UploadFile.id.in_(file_ids))).all() + leased_file_ids = reserve_upgrade_file_cleanup(session, file_ids) + deletable_file_ids = [file_id for file_id in file_ids if file_id not in leased_file_ids] + if leased_file_ids: + logger.info( + "Keep %d source files while KnowledgeFS upgrade leases are active, dataset_id=%s", + len(leased_file_ids), + dataset_id, + ) + files = session.scalars(select(UploadFile).where(UploadFile.id.in_(deletable_file_ids))).all() for file in files: storage.delete(file.key) - file_delete_stmt = delete(UploadFile).where(UploadFile.id.in_(file_ids)) + file_delete_stmt = delete(UploadFile).where(UploadFile.id.in_(deletable_file_ids)) session.execute(file_delete_stmt) session.commit() diff --git a/api/tasks/clean_document_task.py b/api/tasks/clean_document_task.py index e09743a0018..2e57a4e1662 100644 --- a/api/tasks/clean_document_task.py +++ b/api/tasks/clean_document_task.py @@ -11,6 +11,7 @@ from core.tools.utils.web_reader_tool import get_image_upload_file_ids from extensions.ext_storage import storage from models.dataset import Dataset, DatasetMetadataBinding, DocumentSegment, SegmentAttachmentBinding from models.model import UploadFile +from services.dataset_knowledge_fs_upgrade_file_lease import reserve_upgrade_file_cleanup from tasks.refresh_billing_vector_space_task import schedule_billing_vector_space_refresh logger = logging.getLogger(__name__) @@ -125,6 +126,14 @@ def clean_document_task( ) with session_factory.create_session() as session, session.begin(): + if file_id: + if file_id in reserve_upgrade_file_cleanup(session, [file_id]): + logger.info( + "Keep source file while KnowledgeFS upgrade lease is active, file_id=%s, document_id=%s", + file_id, + document_id, + ) + file_id = None if file_id: file = session.scalar(select(UploadFile).where(UploadFile.id == file_id).limit(1)) if file: diff --git a/api/tasks/knowledge_fs_initial_source_tasks.py b/api/tasks/knowledge_fs_initial_source_tasks.py index 85414e0d186..5b236110f44 100644 --- a/api/tasks/knowledge_fs_initial_source_tasks.py +++ b/api/tasks/knowledge_fs_initial_source_tasks.py @@ -56,6 +56,16 @@ class _DatasourceBinding: provider_kind: str +@dataclass(frozen=True) +class KnowledgeFSInitialSourceSubmission: + """Source creation result; the first import is deliberately best-effort.""" + + connection_id: str + source_id: str + workflow_id: str | None + workflow_error: str | None = None + + class KnowledgeFSInitialSourceNotReadyError(RuntimeError): """The Space, connection, or Source workflow is progressing and should be retried.""" @@ -512,6 +522,140 @@ def start_initial_source_import( return workflow.id +def submit_initial_source_for_upgrade( + *, + tenant_id: str, + account_id: str, + control_space_id: str, + operation_id: str, + payload: KnowledgeFSInitialSourcePayload, +) -> KnowledgeFSInitialSourceSubmission: + """Create and commit one Source without waiting for its import workflow. + + Upgrade success is based on the independently usable Source existing. The + selected import is submitted when possible, but a failure is returned as a + warning because users can retry it from the new KnowledgeFS task surface. + """ + + session_maker = session_factory.get_session_maker() + with session_maker() as session: + control_space = SQLAlchemyKnowledgeFSControlSpaceRepository(session).get( + tenant_id=tenant_id, + control_space_id=control_space_id, + ) + if control_space is None: + raise RuntimeError("KnowledgeFS control-space was not found") + if control_space.state is KnowledgeFSControlSpaceState.PROVISIONING: + raise KnowledgeFSInitialSourceNotReadyError("KnowledgeFS Space is still provisioning") + if control_space.state is not KnowledgeFSControlSpaceState.ACTIVE or control_space.knowledge_space_id is None: + raise RuntimeError(f"KnowledgeFS Space cannot accept a Source in state {control_space.state.value}") + + facade = get_knowledge_fs_runtime(session_maker).facade + request_id = _request_id(operation_id=operation_id, payload=payload) + source = _find_initial_source( + facade=facade, + tenant_id=tenant_id, + account_id=account_id, + control_space_id=control_space_id, + request_id=request_id, + ) + if source is None: + binding = _binding(payload) + credential_id, credential_name = _find_credential( + session_maker=session_maker, + tenant_id=tenant_id, + account_id=account_id, + binding=binding, + ) + connection = _find_or_create_connection( + facade=facade, + tenant_id=tenant_id, + account_id=account_id, + control_space_id=control_space_id, + binding=binding, + credential_id=credential_id, + credential_name=credential_name, + ) + source = facade.create_source( + tenant_id=tenant_id, + account_id=account_id, + control_space_id=control_space_id, + payload=_source_payload( + payload=payload, + binding=binding, + connection_id=connection.id, + request_id=request_id, + ), + ) + elif source.connection_id is None: + raise RuntimeError("Initial Source has no connection") + + if source.status != "active" or source.metadata.get("preview") is not False: + source = facade.update_source( + tenant_id=tenant_id, + account_id=account_id, + control_space_id=control_space_id, + source_id=source.id, + payload=KnowledgeFSSourceUpdatePayload( + expectedVersion=source.version, + metadata={**source.metadata, "preview": False, "upgradeJobId": operation_id}, + status="active", + ), + ) + + try: + current_policy = facade.get_source_sync_policy( + tenant_id=tenant_id, + account_id=account_id, + control_space_id=control_space_id, + source_id=source.id, + ) + expected_revision = current_policy.revision + except KnowledgeFSProductResourceNotFoundError: + expected_revision = 0 + facade.update_source_sync_policy( + tenant_id=tenant_id, + account_id=account_id, + control_space_id=control_space_id, + source_id=source.id, + payload=_sync_policy_payload( + payload=payload, + expected_revision=expected_revision, + source_version=source.version, + ), + ) + + workflow_id: str | None = None + workflow_error: str | None = None + try: + workflow = _start_workflow( + facade=facade, + tenant_id=tenant_id, + account_id=account_id, + control_space_id=control_space_id, + source_id=source.id, + request_id=request_id, + payload=payload, + ) + workflow_id = workflow.id + except Exception as exc: + workflow_error = type(exc).__name__ + logger.warning( + "KnowledgeFS upgrade Source was created but its first import was not submitted", + extra={ + "control_space_id": control_space_id, + "error_code": workflow_error, + "source_id": source.id, + }, + ) + return KnowledgeFSInitialSourceSubmission( + connection_id=source.connection_id or "", + source_id=source.id, + workflow_id=workflow_id, + workflow_error=workflow_error, + ) + + def start_initial_website_source_import( *, tenant_id: str, @@ -637,8 +781,10 @@ def import_initial_website_source( __all__ = [ + "KnowledgeFSInitialSourceSubmission", "import_initial_source", "import_initial_website_source", "start_initial_source_import", "start_initial_website_source_import", + "submit_initial_source_for_upgrade", ] diff --git a/api/tasks/knowledge_fs_upgrade_tasks.py b/api/tasks/knowledge_fs_upgrade_tasks.py new file mode 100644 index 00000000000..00d46367e92 --- /dev/null +++ b/api/tasks/knowledge_fs_upgrade_tasks.py @@ -0,0 +1,93 @@ +"""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", +] diff --git a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py index 76920055b2d..17c5ae16bed 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_datasets.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_datasets.py @@ -23,6 +23,8 @@ from controllers.console.datasets.datasets import ( DatasetErrorDocs, DatasetIndexingEstimateApi, DatasetIndexingStatusApi, + DatasetKnowledgeFSUpgradeApi, + DatasetKnowledgeFSUpgradeJobApi, DatasetListApi, DatasetPermissionUserListApi, DatasetQueryApi, @@ -45,6 +47,7 @@ from extensions.storage.storage_type import StorageType from models.account import Account, TenantAccountRole from models.dataset import Dataset, DatasetQuery, Document from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus +from models.knowledge_fs import KnowledgeFSUpgradeJobStatus, KnowledgeFSUpgradeStage from models.model import ApiToken, App, AppMode, IconType, UploadFile from services.dataset_ref_service import DatasetRef from services.dataset_service import DatasetPermissionService, DatasetService @@ -827,6 +830,144 @@ class TestDatasetApiDelete: method(api, MagicMock(), user, dataset_id) +class TestDatasetKnowledgeFSUpgradeApi: + @staticmethod + def _job(dataset_id: str, **overrides): + values = { + "id": "upgrade-job-1", + "old_dataset_id": dataset_id, + "new_control_space_id": None, + "status": KnowledgeFSUpgradeJobStatus.QUEUED, + "stage": KnowledgeFSUpgradeStage.VALIDATING, + "snapshot_at": datetime.datetime(2026, 8, 17, tzinfo=datetime.UTC), + "total_documents": 1, + "completed_documents": 0, + "total_sources": 0, + "completed_sources": 0, + "last_error_code": None, + "last_error_message": None, + "completed_at": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + def test_create_snapshots_and_enqueues_without_remote_work_in_request(self, app: Flask): + dataset_id = "123e4567-e89b-12d3-a456-426614174000" + dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") + user = make_account() + snapshots = MagicMock() + snapshots.create.return_value = self._job(dataset_id) + session_context = MagicMock() + session_context.__enter__.return_value = MagicMock() + api = DatasetKnowledgeFSUpgradeApi() + method = unwrap(api.post) + with ( + app.test_request_context( + f"/datasets/{dataset_id}/knowledge-fs-upgrades", + headers={"Idempotency-Key": "upgrade-request-1"}, + ), + patch("controllers.console.datasets.datasets.dify_config.KNOWLEDGE_FS_ENABLED", True), + patch("controllers.console.datasets.datasets.dify_config.RBAC_ENABLED", False), + patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), + patch("controllers.console.datasets.datasets.session_factory.get_session_maker", return_value="maker"), + patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), + patch.object(DatasetService, "check_dataset_permission"), + patch( + "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService", + return_value=snapshots, + ), + patch("controllers.console.datasets.datasets._enqueue_upgrade_job") as enqueue, + ): + response, status = method(api, "tenant-1", user, dataset_id) + + assert status == 202 + assert response["id"] == "upgrade-job-1" + snapshots.create.assert_called_once_with( + tenant_id="tenant-1", + dataset_id=dataset_id, + requested_by_account_id=user.id, + idempotency_key="upgrade-request-1", + ) + enqueue.assert_called_once_with(snapshots, tenant_id="tenant-1", job_id="upgrade-job-1") + + def test_status_uses_legacy_dataset_permission_in_community_edition(self, app: Flask): + dataset_id = "123e4567-e89b-12d3-a456-426614174000" + dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") + api = DatasetKnowledgeFSUpgradeJobApi() + method = unwrap(api.get) + session_context = MagicMock() + session_context.__enter__.return_value = MagicMock() + with ( + app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades/job-1"), + patch("controllers.console.datasets.datasets.dify_config.RBAC_ENABLED", False), + patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), + patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), + patch.object( + DatasetService, + "check_dataset_permission", + side_effect=services.errors.account.NoPermissionError("no access"), + ), + patch("controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService") as snapshots, + pytest.raises(Forbidden, match="no access"), + ): + method(api, "tenant-1", make_account(), dataset_id, "job-1") + + snapshots.assert_not_called() + + def test_status_relies_on_rbac_decorator_when_enterprise_rbac_is_enabled(self, app: Flask): + dataset_id = "123e4567-e89b-12d3-a456-426614174000" + dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") + snapshots = MagicMock() + snapshots.get.return_value = self._job(dataset_id) + session_context = MagicMock() + session_context.__enter__.return_value = MagicMock() + api = DatasetKnowledgeFSUpgradeJobApi() + method = unwrap(api.get) + with ( + app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades/job-1"), + patch("controllers.console.datasets.datasets.dify_config.RBAC_ENABLED", True), + patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), + patch("controllers.console.datasets.datasets.session_factory.get_session_maker", return_value="maker"), + patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), + patch.object(DatasetService, "check_dataset_permission") as legacy_permission, + patch( + "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService", + return_value=snapshots, + ), + ): + response = method(api, "tenant-1", make_account(), dataset_id, "job-1") + + assert response["id"] == "upgrade-job-1" + legacy_permission.assert_not_called() + + def test_retry_rejects_a_job_from_another_dataset(self, app: Flask): + dataset_id = "123e4567-e89b-12d3-a456-426614174000" + dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") + snapshots = MagicMock() + snapshots.retry.return_value = self._job("223e4567-e89b-12d3-a456-426614174000") + session_context = MagicMock() + session_context.__enter__.return_value = MagicMock() + api = DatasetKnowledgeFSUpgradeJobApi() + method = unwrap(api.post) + with ( + app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades/job-1"), + patch("controllers.console.datasets.datasets.dify_config.RBAC_ENABLED", False), + patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), + patch("controllers.console.datasets.datasets.session_factory.get_session_maker", return_value="maker"), + patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), + patch.object(DatasetService, "check_dataset_permission"), + patch( + "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService", + return_value=snapshots, + ), + patch("controllers.console.datasets.datasets._enqueue_upgrade_job") as enqueue, + pytest.raises(NotFound, match="Upgrade job was not found"), + ): + method(api, "tenant-1", make_account(), dataset_id, "job-1") + + enqueue.assert_not_called() + + class TestDatasetUseCheckApi: @pytest.mark.parametrize("is_using", [True, False]) def test_get_use_check(self, app: Flask, is_using: bool): diff --git a/api/tests/unit_tests/controllers/test_knowledge_fs_product_controllers.py b/api/tests/unit_tests/controllers/test_knowledge_fs_product_controllers.py index 919099c4d8b..be644d81dab 100644 --- a/api/tests/unit_tests/controllers/test_knowledge_fs_product_controllers.py +++ b/api/tests/unit_tests/controllers/test_knowledge_fs_product_controllers.py @@ -1085,7 +1085,7 @@ def test_service_profile_rejects_cross_control_space_before_facade_io() -> None: def test_product_modules_do_not_import_dify_dataset_or_document_services() -> None: paths = [ - *_API_ROOT.glob("services/knowledge_fs/*.py"), + *(path for path in _API_ROOT.glob("services/knowledge_fs/*.py") if not path.name.startswith("upgrade_")), *_API_ROOT.glob("controllers/console/knowledge_fs/*.py"), *_API_ROOT.glob("controllers/service_api/knowledge_fs/*.py"), ] diff --git a/api/tests/unit_tests/dev/test_check_knowledge_fs_coverage.py b/api/tests/unit_tests/dev/test_check_knowledge_fs_coverage.py index 8b71a3f0d2c..3d96904ff19 100644 --- a/api/tests/unit_tests/dev/test_check_knowledge_fs_coverage.py +++ b/api/tests/unit_tests/dev/test_check_knowledge_fs_coverage.py @@ -33,6 +33,8 @@ from dev.check_knowledge_fs_coverage import ( "api/extensions/ext_knowledge_fs_observability.py", "api/models/knowledge_fs.py", "api/repositories/sqlalchemy_knowledge_fs_cutover_repository.py", + "api/services/dataset_knowledge_fs_upgrade_file_lease.py", + "api/services/dataset_knowledge_fs_upgrade_service.py", "api/services/knowledge_fs/runtime.py", "api/services/knowledge_fs_capability.py", "api/tasks/knowledge_fs_lifecycle_tasks.py", @@ -121,6 +123,7 @@ def test_core_coverage_rejects_unclassified_knowledge_fs_production_files(tmp_pa "api/migrations/versions/2026_07_21_1600-e5a7c9b2d416_add_knowledge_fs_cleanup_completion.py", "api/migrations/versions/2026_08_13_1200-9d4e6f8a1b2c_add_knowledge_fs_space_tag_bindings.py", "api/migrations/versions/2026_08_17_1200-4f8b2c7d9e10_add_knowledge_fs_icon_background.py", + "api/migrations/versions/2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py", ], ) def test_core_coverage_allows_explicit_non_core_migrations(tmp_path: Path, migration: str) -> None: diff --git a/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py b/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py index 88ba47c6eb3..78f4c6b031a 100644 --- a/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py +++ b/api/tests/unit_tests/extensions/test_knowledge_fs_celery.py @@ -66,6 +66,10 @@ def test_celery_registers_initial_source_task_when_knowledge_fs_lifecycle_is_rea "task": "tasks.knowledge_fs_lifecycle_tasks.cleanup_knowledge_fs_staged_uploads", "schedule": timedelta(seconds=2), } + assert celery_app.conf["beat_schedule"]["knowledge_fs_upgrade_file_cleanup"] == { + "task": "tasks.knowledge_fs_upgrade_tasks.cleanup_deferred_knowledge_fs_upgrade_files", + "schedule": timedelta(seconds=2), + } with ( patch("extensions.ext_celery.dify_config", config), diff --git a/api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_jobs.py b/api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_jobs.py new file mode 100644 index 00000000000..41d1f953c17 --- /dev/null +++ b/api/tests/unit_tests/migrations/test_knowledge_fs_upgrade_jobs.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations + +from models.knowledge_fs import ( + KnowledgeFSUpgradeDocument, + KnowledgeFSUpgradeFileLease, + KnowledgeFSUpgradeJob, + KnowledgeFSUpgradeSource, +) + +_MIGRATION_PATH = ( + Path(__file__).resolve().parents[3] + / "migrations/versions/2026_08_17_1200-f3a8c1d7e920_add_knowledge_fs_upgrade_jobs.py" +) +_MODELS = ( + KnowledgeFSUpgradeJob, + KnowledgeFSUpgradeDocument, + KnowledgeFSUpgradeSource, + KnowledgeFSUpgradeFileLease, +) + + +def _load_migration_module(): + spec = importlib.util.spec_from_file_location("knowledge_fs_upgrade_jobs", _MIGRATION_PATH) + if spec is None or spec.loader is None: + raise RuntimeError("failed to load migration module") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _run_step(module: object, engine: sa.Engine, step_name: str) -> None: + with engine.begin() as connection: + operations = Operations(MigrationContext.configure(connection)) + original_op = module.op + module.op = operations + try: + getattr(module, step_name)() + finally: + module.op = original_op + + +def test_upgrade_schema_matches_upgrade_models_and_downgrades_cleanly() -> None: + engine = sa.create_engine("sqlite:///:memory:") + module = _load_migration_module() + + _run_step(module, engine, "upgrade") + + inspector = sa.inspect(engine) + assert set(inspector.get_table_names()) == {model.__tablename__ for model in _MODELS} + for model in _MODELS: + migrated_columns = {column["name"] for column in inspector.get_columns(model.__tablename__)} + assert migrated_columns == set(model.__table__.columns.keys()) + + _run_step(module, engine, "downgrade") + assert sa.inspect(engine).get_table_names() == [] diff --git a/api/tests/unit_tests/services/test_knowledge_fs_upgrade_file_lease.py b/api/tests/unit_tests/services/test_knowledge_fs_upgrade_file_lease.py new file mode 100644 index 00000000000..3d02e63cc98 --- /dev/null +++ b/api/tests/unit_tests/services/test_knowledge_fs_upgrade_file_lease.py @@ -0,0 +1,295 @@ +from __future__ import annotations + +from datetime import timedelta +from unittest.mock import patch + +from sqlalchemy.orm import Session, sessionmaker + +from extensions.storage.storage_type import StorageType +from libs.datetime_utils import naive_utc_now +from models.dataset import Document +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom +from models.knowledge_fs import ( + KnowledgeFSUpgradeFileLease, + KnowledgeFSUpgradeFileLeaseStatus, + KnowledgeFSUpgradeJob, +) +from models.model import UploadFile +from services.dataset_knowledge_fs_upgrade_file_lease import ( + active_upgrade_file_ids, + cleanup_deferred_upgrade_files, + release_upgrade_file_lease, + reserve_upgrade_file_cleanup, +) + +_TENANT_ID = "00000000-0000-0000-0000-000000000001" +_ACCOUNT_ID = "00000000-0000-0000-0000-000000000002" +_DATASET_ID = "00000000-0000-0000-0000-000000000003" +_ACTIVE_FILE_ID = "00000000-0000-0000-0000-000000000004" +_EXPIRED_FILE_ID = "00000000-0000-0000-0000-000000000005" + + +def test_only_active_unexpired_source_files_are_protected(sqlite_session_factory: sessionmaker[Session]) -> None: + now = naive_utc_now() + job = KnowledgeFSUpgradeJob( + tenant_id=_TENANT_ID, + old_dataset_id=_DATASET_ID, + requested_by_account_id=_ACCOUNT_ID, + owner_account_id=_ACCOUNT_ID, + idempotency_key="upgrade-lease-test", + snapshot_at=now, + config_snapshot={}, + permission_snapshot={}, + app_binding_snapshot=[], + tag_ids_snapshot=[], + ) + with sqlite_session_factory.begin() as session: + session.add(job) + session.flush() + session.add_all( + [ + KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=_ACTIVE_FILE_ID, + expires_at=now + timedelta(minutes=1), + ), + KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=_EXPIRED_FILE_ID, + expires_at=now - timedelta(seconds=1), + status=KnowledgeFSUpgradeFileLeaseStatus.ACTIVE, + ), + ] + ) + + with sqlite_session_factory() as session: + assert active_upgrade_file_ids( + session, + [_ACTIVE_FILE_ID, _EXPIRED_FILE_ID], + now=now, + ) == {_ACTIVE_FILE_ID} + + +def test_cleanup_request_is_persisted_on_every_active_lease( + sqlite_session_factory: sessionmaker[Session], +) -> None: + now = naive_utc_now() + jobs = [ + KnowledgeFSUpgradeJob( + tenant_id=_TENANT_ID, + old_dataset_id=_DATASET_ID, + requested_by_account_id=_ACCOUNT_ID, + owner_account_id=_ACCOUNT_ID, + idempotency_key=f"upgrade-cleanup-request-{index}", + snapshot_at=now, + config_snapshot={}, + permission_snapshot={}, + app_binding_snapshot=[], + tag_ids_snapshot=[], + ) + for index in range(2) + ] + with sqlite_session_factory.begin() as session: + session.add_all(jobs) + session.flush() + leases = [ + KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=_ACTIVE_FILE_ID, + expires_at=now + timedelta(minutes=1), + ) + for job in jobs + ] + session.add_all(leases) + + with sqlite_session_factory.begin() as session: + assert reserve_upgrade_file_cleanup(session, [_ACTIVE_FILE_ID], now=now) == {_ACTIVE_FILE_ID} + + with sqlite_session_factory() as session: + persisted = list(session.query(KnowledgeFSUpgradeFileLease).filter_by(old_upload_file_id=_ACTIVE_FILE_ID)) + assert len(persisted) == 2 + assert all(lease.cleanup_requested_at == now for lease in persisted) + + +def test_last_lease_release_deletes_deferred_orphan_file( + sqlite_session_factory: sessionmaker[Session], +) -> None: + now = naive_utc_now() + job = KnowledgeFSUpgradeJob( + tenant_id=_TENANT_ID, + old_dataset_id=_DATASET_ID, + requested_by_account_id=_ACCOUNT_ID, + owner_account_id=_ACCOUNT_ID, + idempotency_key="upgrade-deferred-file-cleanup", + snapshot_at=now, + config_snapshot={}, + permission_snapshot={}, + app_binding_snapshot=[], + tag_ids_snapshot=[], + ) + upload_file = UploadFile( + tenant_id=_TENANT_ID, + storage_type=StorageType.LOCAL, + key=f"upload_files/{_TENANT_ID}/orphan.txt", + name="orphan.txt", + size=10, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=_ACCOUNT_ID, + created_at=now, + used=False, + ) + with sqlite_session_factory.begin() as session: + session.add_all([job, upload_file]) + session.flush() + lease = KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=upload_file.id, + expires_at=now + timedelta(minutes=1), + cleanup_requested_at=now, + ) + session.add(lease) + + with patch("services.dataset_knowledge_fs_upgrade_file_lease.storage") as storage: + assert cleanup_deferred_upgrade_files(sqlite_session_factory, now=now) == 0 + storage.delete.assert_not_called() + with sqlite_session_factory.begin() as session: + assert ( + release_upgrade_file_lease( + session, + job_id=job.id, + upload_file_id=upload_file.id, + now=now, + ) + is True + ) + + storage.delete.assert_called_once_with(upload_file.key) + with sqlite_session_factory() as session: + assert session.get(UploadFile, upload_file.id) is None + persisted_lease = session.get(KnowledgeFSUpgradeFileLease, lease.id) + assert persisted_lease is not None + assert persisted_lease.status is KnowledgeFSUpgradeFileLeaseStatus.RELEASED + + +def test_lease_release_keeps_file_referenced_by_another_legacy_document( + sqlite_session_factory: sessionmaker[Session], +) -> None: + now = naive_utc_now() + job = KnowledgeFSUpgradeJob( + tenant_id=_TENANT_ID, + old_dataset_id=_DATASET_ID, + requested_by_account_id=_ACCOUNT_ID, + owner_account_id=_ACCOUNT_ID, + idempotency_key="upgrade-shared-legacy-file", + snapshot_at=now, + config_snapshot={}, + permission_snapshot={}, + app_binding_snapshot=[], + tag_ids_snapshot=[], + ) + upload_file = UploadFile( + tenant_id=_TENANT_ID, + storage_type=StorageType.LOCAL, + key=f"upload_files/{_TENANT_ID}/shared.txt", + name="shared.txt", + size=10, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=_ACCOUNT_ID, + created_at=now, + used=False, + ) + with sqlite_session_factory.begin() as session: + session.add_all([job, upload_file]) + session.flush() + document = Document( + tenant_id=_TENANT_ID, + dataset_id="00000000-0000-0000-0000-000000000020", + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + data_source_info=f'{{"upload_file_id":"{upload_file.id}"}}', + batch="shared-file-test", + name="shared.txt", + created_from=DocumentCreatedFrom.WEB, + created_by=_ACCOUNT_ID, + enabled=True, + archived=False, + indexing_status="completed", + ) + lease = KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=upload_file.id, + expires_at=now + timedelta(minutes=1), + cleanup_requested_at=now, + ) + session.add_all([document, lease]) + + with patch("services.dataset_knowledge_fs_upgrade_file_lease.storage") as storage: + with sqlite_session_factory.begin() as session: + assert ( + release_upgrade_file_lease( + session, + job_id=job.id, + upload_file_id=upload_file.id, + now=now, + ) + is False + ) + + storage.delete.assert_not_called() + with sqlite_session_factory() as session: + assert session.get(UploadFile, upload_file.id) is not None + + +def test_expired_lease_sweeper_deletes_deferred_orphan_file( + sqlite_session_factory: sessionmaker[Session], +) -> None: + now = naive_utc_now() + job = KnowledgeFSUpgradeJob( + tenant_id=_TENANT_ID, + old_dataset_id=_DATASET_ID, + requested_by_account_id=_ACCOUNT_ID, + owner_account_id=_ACCOUNT_ID, + idempotency_key="upgrade-expired-file-cleanup", + snapshot_at=now, + config_snapshot={}, + permission_snapshot={}, + app_binding_snapshot=[], + tag_ids_snapshot=[], + ) + upload_file = UploadFile( + tenant_id=_TENANT_ID, + storage_type=StorageType.LOCAL, + key=f"upload_files/{_TENANT_ID}/expired-orphan.txt", + name="expired-orphan.txt", + size=10, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=_ACCOUNT_ID, + created_at=now, + used=False, + ) + with sqlite_session_factory.begin() as session: + session.add_all([job, upload_file]) + session.flush() + lease = KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=upload_file.id, + expires_at=now - timedelta(seconds=1), + cleanup_requested_at=now - timedelta(minutes=1), + ) + session.add(lease) + + with patch("services.dataset_knowledge_fs_upgrade_file_lease.storage") as storage: + assert cleanup_deferred_upgrade_files(sqlite_session_factory, now=now) == 1 + + storage.delete.assert_called_once_with(upload_file.key) + with sqlite_session_factory() as session: + assert session.get(UploadFile, upload_file.id) is None + persisted_lease = session.get(KnowledgeFSUpgradeFileLease, lease.id) + assert persisted_lease is not None + assert persisted_lease.status is KnowledgeFSUpgradeFileLeaseStatus.EXPIRED diff --git a/api/tests/unit_tests/services/test_knowledge_fs_upgrade_service.py b/api/tests/unit_tests/services/test_knowledge_fs_upgrade_service.py new file mode 100644 index 00000000000..d0a1ca4636c --- /dev/null +++ b/api/tests/unit_tests/services/test_knowledge_fs_upgrade_service.py @@ -0,0 +1,1332 @@ +from __future__ import annotations + +import json +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from extensions.storage.storage_type import StorageType +from graphon.model_runtime.entities.model_entities import ModelType +from libs.datetime_utils import naive_utc_now +from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models.dataset import Dataset, DatasetMetadata, DatasetPermissionEnum, Document +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom +from models.knowledge_fs import ( + KnowledgeFSAppSpaceJoinType, + KnowledgeFSControlSpace, + KnowledgeFSControlSpacePermissionRole, + KnowledgeFSControlSpaceState, + KnowledgeFSControlSpaceVisibility, + KnowledgeFSUpgradeDocument, + KnowledgeFSUpgradeFileLease, + KnowledgeFSUpgradeFileLeaseStatus, + KnowledgeFSUpgradeItemStatus, + KnowledgeFSUpgradeJob, + KnowledgeFSUpgradeJobStatus, + KnowledgeFSUpgradeSource, + KnowledgeFSUpgradeStage, +) +from models.model import UploadFile +from models.oauth import DatasourceProvider +from services import dataset_knowledge_fs_upgrade_service as upgrade_module +from services.dataset_knowledge_fs_upgrade_service import ( + KnowledgeFSUpgradeConflictError, + KnowledgeFSUpgradeDocumentReconciler, + KnowledgeFSUpgradeError, + KnowledgeFSUpgradeNotFoundError, + KnowledgeFSUpgradeNotReadyError, + KnowledgeFSUpgradeRunner, + KnowledgeFSUpgradeSnapshotService, +) +from services.knowledge_fs.product_dto import ( + KnowledgeFSLogicalDocumentListResponse, + KnowledgeFSLogicalDocumentResponse, + KnowledgeFSMetadataFieldListResponse, +) + +_TENANT_ID = "00000000-0000-0000-0000-000000000001" +_ACCOUNT_ID = "00000000-0000-0000-0000-000000000002" +_DATASET_ID = "00000000-0000-0000-0000-000000000003" +_CONTROL_SPACE_ID = "00000000-0000-0000-0000-000000000004" +_UPLOAD_FILE_ID = "00000000-0000-0000-0000-000000000005" +_DOCUMENT_ASSET_ID = "00000000-0000-0000-0000-000000000006" +_LOGICAL_DOCUMENT_ID = "00000000-0000-0000-0000-000000000007" + + +def _job(**overrides: object) -> KnowledgeFSUpgradeJob: + values: dict[str, object] = { + "tenant_id": _TENANT_ID, + "old_dataset_id": _DATASET_ID, + "requested_by_account_id": _ACCOUNT_ID, + "owner_account_id": _ACCOUNT_ID, + "idempotency_key": "upgrade-service-test", + "snapshot_at": naive_utc_now(), + "config_snapshot": {}, + "permission_snapshot": {}, + "app_binding_snapshot": [], + "tag_ids_snapshot": [], + } + values.update(overrides) + return KnowledgeFSUpgradeJob(**values) # type: ignore[arg-type] + + +def _logical_document(*, enabled: bool, row_version: int = 1) -> KnowledgeFSLogicalDocumentResponse: + return KnowledgeFSLogicalDocumentResponse.model_validate( + { + "active": { + "contentHash": "a" * 64, + "createdAt": "2026-08-17T00:00:00Z", + "documentAssetId": _DOCUMENT_ASSET_ID, + "documentAssetVersion": 1, + "documentId": _LOGICAL_DOCUMENT_ID, + "knowledgeSpaceId": "00000000-0000-0000-0000-000000000008", + "mimeType": "text/plain", + "revision": 1, + "sizeBytes": 12, + "state": "candidate", + }, + "createdAt": "2026-08-17T00:00:00Z", + "enabled": enabled, + "id": _LOGICAL_DOCUMENT_ID, + "knowledgeSpaceId": "00000000-0000-0000-0000-000000000008", + "rowVersion": row_version, + "status": "pending", + "title": "guide.txt", + "updatedAt": "2026-08-17T00:00:00Z", + "userMetadata": {}, + } + ) + + +def test_retry_restores_failed_items_and_unreleased_file_lease( + sqlite_session_factory: sessionmaker[Session], +) -> None: + job = _job(status=KnowledgeFSUpgradeJobStatus.FAILED) + with sqlite_session_factory.begin() as session: + session.add(job) + session.flush() + document = KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000010", + name="guide.txt", + data_source_type="upload_file", + data_source_info={"upload_file_id": _UPLOAD_FILE_ID}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="error", + status=KnowledgeFSUpgradeItemStatus.FAILED, + old_upload_file_id=_UPLOAD_FILE_ID, + ) + source = KnowledgeFSUpgradeSource( + job_id=job.id, + tenant_id=_TENANT_ID, + source_key="source-1", + payload_snapshot={}, + status=KnowledgeFSUpgradeItemStatus.FAILED, + ) + lease = KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=_UPLOAD_FILE_ID, + expires_at=naive_utc_now() - timedelta(days=1), + status=KnowledgeFSUpgradeFileLeaseStatus.EXPIRED, + ) + session.add_all([document, source, lease]) + + retried = KnowledgeFSUpgradeSnapshotService(sqlite_session_factory).retry( + tenant_id=_TENANT_ID, + job_id=job.id, + ) + + assert retried.status is KnowledgeFSUpgradeJobStatus.QUEUED + with sqlite_session_factory() as session: + assert session.get(KnowledgeFSUpgradeDocument, document.id).status is KnowledgeFSUpgradeItemStatus.PENDING + assert session.get(KnowledgeFSUpgradeSource, source.id).status is KnowledgeFSUpgradeItemStatus.PENDING + persisted_lease = session.get(KnowledgeFSUpgradeFileLease, lease.id) + assert persisted_lease.status is KnowledgeFSUpgradeFileLeaseStatus.ACTIVE + assert persisted_lease.expires_at > naive_utc_now() + + +def test_snapshot_is_immutable_after_the_legacy_dataset_changes( + sqlite_session_factory: sessionmaker[Session], +) -> None: + tenant = Tenant(name="Upgrade workspace") + owner = Account(name="Owner", email="upgrade-owner@example.com") + with sqlite_session_factory.begin() as session: + session.add_all([tenant, owner]) + session.flush() + session.add( + TenantAccountJoin( + tenant_id=tenant.id, + account_id=owner.id, + role=TenantAccountRole.OWNER, + ) + ) + dataset = Dataset( + tenant_id=tenant.id, + name="Legacy handbook", + description="Click-time description", + data_source_type=DataSourceType.UPLOAD_FILE, + permission=DatasetPermissionEnum.ONLY_ME, + created_by=owner.id, + embedding_model="text-embedding-3-large", + embedding_model_provider="openai", + retrieval_model={"top_k": 8, "reranking_enable": True}, + enable_api=True, + ) + session.add(dataset) + session.flush() + upload_file = UploadFile( + tenant_id=tenant.id, + storage_type=StorageType.LOCAL, + key=f"upload_files/{tenant.id}/handbook.txt", + name="handbook.txt", + size=12, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=owner.id, + created_at=naive_utc_now(), + used=False, + ) + metadata = DatasetMetadata( + tenant_id=tenant.id, + dataset_id=dataset.id, + type="string", + name="department", + created_by=owner.id, + ) + document = Document( + id="00000000-0000-0000-0000-000000000012", + tenant_id=tenant.id, + dataset_id=dataset.id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + data_source_info=json.dumps({"upload_file_id": upload_file.id}), + batch="batch-1", + name="handbook.txt", + created_from=DocumentCreatedFrom.WEB, + created_by=owner.id, + enabled=False, + archived=True, + indexing_status="completed", + doc_metadata={"department": "support"}, + ) + session.add_all([upload_file, metadata, document]) + tenant_id = tenant.id + owner_id = owner.id + dataset_id = dataset.id + + job = KnowledgeFSUpgradeSnapshotService(sqlite_session_factory).create( + tenant_id=tenant_id, + dataset_id=dataset_id, + requested_by_account_id=owner_id, + idempotency_key="immutable-snapshot", + ) + replayed = KnowledgeFSUpgradeSnapshotService(sqlite_session_factory).create( + tenant_id=tenant_id, + dataset_id=dataset_id, + requested_by_account_id=owner_id, + idempotency_key="immutable-snapshot", + ) + assert replayed.id == job.id + with sqlite_session_factory.begin() as session: + persisted_dataset = session.get(Dataset, dataset_id) + persisted_document = session.get(Document, document.id) + assert persisted_dataset is not None + assert persisted_document is not None + persisted_dataset.name = "Changed after click" + persisted_document.enabled = True + persisted_document.archived = False + + with sqlite_session_factory() as session: + persisted_job = session.get(KnowledgeFSUpgradeJob, job.id) + snapshot_document = session.scalar( + select(KnowledgeFSUpgradeDocument).where(KnowledgeFSUpgradeDocument.job_id == job.id) + ) + assert persisted_job is not None + assert snapshot_document is not None + assert persisted_job.config_snapshot["name"] == "Legacy handbook" + assert persisted_job.config_snapshot["embedding_model"] == "text-embedding-3-large" + assert persisted_job.config_snapshot["embedding_model_provider"] == "openai" + assert persisted_job.config_snapshot["retrieval_model"] == {"top_k": 8, "reranking_enable": True} + assert persisted_job.config_snapshot["enable_api"] is True + assert persisted_job.config_snapshot["metadata_fields"] == [{"name": "department", "type": "string"}] + assert snapshot_document.desired_enabled is False + assert snapshot_document.legacy_archived is True + assert snapshot_document.metadata_snapshot == {"department": "support"} + + +def test_finalize_requires_every_document_and_source_handoff(sqlite_session_factory: sessionmaker[Session]) -> None: + job = _job( + status=KnowledgeFSUpgradeJobStatus.RUNNING, + stage=KnowledgeFSUpgradeStage.FINALIZING, + total_documents=2, + completed_documents=1, + ) + with sqlite_session_factory.begin() as session: + session.add(job) + + with pytest.raises(KnowledgeFSUpgradeError, match="Not all Dataset documents"): + KnowledgeFSUpgradeRunner(sqlite_session_factory)._finalize(job) + + +def test_shared_source_file_lease_is_released_after_every_document_handoff( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch +) -> None: + account = Account(name="Upgrade owner", email="upgrade-shared-file@example.com") + with sqlite_session_factory.begin() as session: + session.add(account) + session.flush() + upload_file = UploadFile( + tenant_id=_TENANT_ID, + storage_type=StorageType.LOCAL, + key=f"upload_files/{_TENANT_ID}/shared.txt", + name="shared.txt", + size=12, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=account.id, + created_at=naive_utc_now(), + used=False, + ) + job = _job( + owner_account_id=account.id, + status=KnowledgeFSUpgradeJobStatus.RUNNING, + stage=KnowledgeFSUpgradeStage.SUBMITTING_DOCUMENTS, + new_control_space_id=_CONTROL_SPACE_ID, + total_documents=2, + ) + session.add_all([upload_file, job]) + session.flush() + for suffix in ("21", "22"): + session.add( + KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id=f"00000000-0000-0000-0000-0000000000{suffix}", + name=f"guide-{suffix}.txt", + data_source_type="upload_file", + data_source_info={"upload_file_id": upload_file.id}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + old_upload_file_id=upload_file.id, + ) + ) + lease = KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=upload_file.id, + expires_at=naive_utc_now() + timedelta(days=1), + ) + session.add(lease) + + staged = MagicMock() + staged.stage.side_effect = [SimpleNamespace(id="staged-1"), SimpleNamespace(id="staged-2")] + staged.claim.side_effect = [ + SimpleNamespace(document_asset_id="asset-1", compilation_job_id="compilation-1"), + SimpleNamespace(document_asset_id="asset-2", compilation_job_id="compilation-2"), + ] + monkeypatch.setattr(upgrade_module, "KnowledgeFSStagedUploadService", lambda *_args, **_kwargs: staged) + monkeypatch.setattr(upgrade_module, "get_knowledge_fs_runtime", lambda _: SimpleNamespace(facade=MagicMock())) + monkeypatch.setattr(upgrade_module.storage, "load", lambda _key: b"shared body") + monkeypatch.setattr( + upgrade_module.FeatureService, + "get_knowledge_file_size_limit", + lambda _tenant_id: 15, + ) + + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + assert runner._submit_next_document(job) is True + with sqlite_session_factory() as session: + persisted = session.get(KnowledgeFSUpgradeFileLease, lease.id) + assert persisted is not None + assert persisted.status is KnowledgeFSUpgradeFileLeaseStatus.ACTIVE + + assert runner._submit_next_document(job) is True + with sqlite_session_factory() as session: + persisted = session.get(KnowledgeFSUpgradeFileLease, lease.id) + assert persisted is not None + assert persisted.status is KnowledgeFSUpgradeFileLeaseStatus.RELEASED + assert persisted.released_at is not None + + +def test_reconciler_applies_metadata_then_snapshot_availability( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch +) -> None: + job = _job( + status=KnowledgeFSUpgradeJobStatus.SUCCEEDED, + stage=KnowledgeFSUpgradeStage.COMPLETED, + new_control_space_id=_CONTROL_SPACE_ID, + ) + with sqlite_session_factory.begin() as session: + session.add(job) + session.flush() + document = KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000011", + name="guide.txt", + data_source_type="upload_file", + data_source_info={"upload_file_id": _UPLOAD_FILE_ID}, + metadata_snapshot={"department": "support"}, + desired_enabled=False, + legacy_archived=True, + legacy_indexing_status="completed", + status=KnowledgeFSUpgradeItemStatus.SUCCEEDED, + old_upload_file_id=_UPLOAD_FILE_ID, + new_document_asset_id=_DOCUMENT_ASSET_ID, + ) + session.add(document) + + initial = _logical_document(enabled=True, row_version=1) + after_metadata = initial.model_copy(update={"row_version": 2}) + after_availability = after_metadata.model_copy(update={"enabled": False, "row_version": 3}) + facade = MagicMock() + facade.list_logical_documents.return_value = KnowledgeFSLogicalDocumentListResponse( + data=[initial], next_cursor=None + ) + facade.update_document_metadata.return_value = after_metadata + facade.update_logical_document_availability.return_value = after_availability + monkeypatch.setattr( + upgrade_module, + "get_knowledge_fs_runtime", + lambda _: SimpleNamespace(facade=facade), + ) + + assert KnowledgeFSUpgradeDocumentReconciler(sqlite_session_factory).reconcile(job_id=job.id) == 0 + + metadata_payload = facade.update_document_metadata.call_args.kwargs["payload"] + assert metadata_payload.expected_row_version == 1 + assert metadata_payload.patch == {"department": "support"} + availability_payload = facade.update_logical_document_availability.call_args.kwargs["payload"] + assert availability_payload.expected_row_version == 2 + assert availability_payload.enabled is False + with sqlite_session_factory() as session: + persisted = session.scalar( + select(KnowledgeFSUpgradeDocument).where(KnowledgeFSUpgradeDocument.id == document.id) + ) + assert persisted is not None + assert persisted.new_logical_document_id == _LOGICAL_DOCUMENT_ID + assert persisted.state_reconciled_at is not None + + +def test_metadata_fields_are_created_idempotently() -> None: + job = _job( + new_control_space_id=_CONTROL_SPACE_ID, + config_snapshot={ + "metadata_fields": [ + {"name": "department", "type": "string"}, + {"name": "priority", "type": "number"}, + ] + }, + ) + facade = MagicMock() + facade.list_metadata_fields.return_value = KnowledgeFSMetadataFieldListResponse( + data=[ + { + "id": "field-1", + "name": "department", + "type": "string", + "count": 0, + "rowVersion": 0, + "createdAt": "2026-08-17T00:00:00Z", + "updatedAt": "2026-08-17T00:00:00Z", + } + ], + next_cursor=None, + ) + + upgrade_module._migrate_metadata_fields(job, facade) + + facade.create_metadata_field.assert_called_once() + assert facade.create_metadata_field.call_args.kwargs["payload"].model_dump() == { + "name": "priority", + "type": "number", + } + + +def _active_model(provider: str, model: str) -> SimpleNamespace: + return SimpleNamespace(provider=SimpleNamespace(provider=provider), model=model) + + +def test_configuration_prefers_active_legacy_models_and_uses_default_reasoning( + monkeypatch: pytest.MonkeyPatch, +) -> None: + embedding = _active_model("langgenius/openai/openai", "text-embedding-3-large") + rerank = _active_model("langgenius/cohere/cohere", "rerank-v3.5") + reasoning = _active_model("langgenius/anthropic/anthropic", "claude-sonnet") + models_by_type = { + ModelType.TEXT_EMBEDDING: [embedding], + ModelType.RERANK: [rerank], + ModelType.LLM: [reasoning], + } + configurations = MagicMock() + configurations.get_models.side_effect = lambda *, model_type, **_kwargs: models_by_type[model_type] + provider_manager = MagicMock() + provider_manager.get_configurations.return_value = configurations + provider_manager.get_default_model.return_value = reasoning + monkeypatch.setattr(upgrade_module, "create_plugin_provider_manager", lambda **_kwargs: provider_manager) + job = _job( + config_snapshot={ + "embedding_model_provider": "openai", + "embedding_model": "text-embedding-3-large", + "retrieval_model": { + "top_k": 12, + "score_threshold_enabled": True, + "score_threshold": 0.72, + "reranking_model": { + "reranking_provider_name": "cohere", + "reranking_model_name": "rerank-v3.5", + }, + }, + } + ) + + resolved = upgrade_module._resolve_configuration(job) + + assert resolved["embedding"] == { + "pluginId": "langgenius/openai", + "provider": "openai", + "model": "text-embedding-3-large", + } + retrieval = resolved["retrieval"] + assert retrieval["reasoningModel"] == { + "pluginId": "langgenius/anthropic", + "provider": "anthropic", + "model": "claude-sonnet", + } + assert retrieval["rerank"] == { + "enabled": True, + "model": { + "pluginId": "langgenius/cohere", + "provider": "cohere", + "model": "rerank-v3.5", + }, + } + assert retrieval["scoreThreshold"] == {"enabled": True, "stage": "mode-final", "value": 0.72} + assert retrieval["topK"] == 12 + provider_manager.get_default_model.assert_called_once_with(tenant_id=_TENANT_ID, model_type=ModelType.LLM) + + +def test_configuration_falls_back_to_active_workspace_defaults(monkeypatch: pytest.MonkeyPatch) -> None: + defaults = { + ModelType.TEXT_EMBEDDING: _active_model("langgenius/openai/openai", "embedding-default"), + ModelType.RERANK: _active_model("langgenius/cohere/cohere", "rerank-default"), + ModelType.LLM: _active_model("langgenius/anthropic/anthropic", "reasoning-default"), + } + configurations = MagicMock() + configurations.get_models.side_effect = lambda *, model_type, **_kwargs: [defaults[model_type]] + provider_manager = MagicMock() + provider_manager.get_configurations.return_value = configurations + provider_manager.get_default_model.side_effect = lambda *, model_type, **_kwargs: defaults[model_type] + monkeypatch.setattr(upgrade_module, "create_plugin_provider_manager", lambda **_kwargs: provider_manager) + job = _job( + config_snapshot={ + "embedding_model_provider": "openai", + "embedding_model": "inactive-legacy-embedding", + "retrieval_model": { + "reranking_model": { + "reranking_provider_name": "cohere", + "reranking_model_name": "inactive-legacy-rerank", + } + }, + } + ) + + resolved = upgrade_module._resolve_configuration(job) + + assert resolved["embedding"]["model"] == "embedding-default" + assert resolved["retrieval"]["rerank"]["model"]["model"] == "rerank-default" + assert resolved["retrieval"]["reasoningModel"]["model"] == "reasoning-default" + assert resolved["retrieval"]["topK"] == 4 + + +def test_configuration_fails_when_workspace_default_reasoning_model_is_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + embedding = _active_model("langgenius/openai/openai", "embedding") + rerank = _active_model("langgenius/cohere/cohere", "rerank") + models_by_type = { + ModelType.TEXT_EMBEDDING: [embedding], + ModelType.RERANK: [rerank], + ModelType.LLM: [], + } + configurations = MagicMock() + configurations.get_models.side_effect = lambda *, model_type, **_kwargs: models_by_type[model_type] + provider_manager = MagicMock() + provider_manager.get_configurations.return_value = configurations + provider_manager.get_default_model.return_value = None + monkeypatch.setattr(upgrade_module, "create_plugin_provider_manager", lambda **_kwargs: provider_manager) + job = _job( + config_snapshot={ + "embedding_model_provider": "openai", + "embedding_model": "embedding", + "retrieval_model": { + "reranking_model": { + "reranking_provider_name": "cohere", + "reranking_model_name": "rerank", + } + }, + } + ) + + with pytest.raises(KnowledgeFSUpgradeError, match="Workspace default llm model is unavailable"): + upgrade_module._resolve_configuration(job) + + +def test_access_migration_maps_visibility_members_apps_api_and_tags( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch +) -> None: + viewer_id = "00000000-0000-0000-0000-000000000030" + job = _job( + status=KnowledgeFSUpgradeJobStatus.RUNNING, + stage=KnowledgeFSUpgradeStage.MIGRATING_ACCESS, + new_control_space_id=_CONTROL_SPACE_ID, + config_snapshot={"enable_api": True, "metadata_fields": []}, + permission_snapshot={ + "visibility": KnowledgeFSControlSpaceVisibility.PARTIAL_MEMBERS.value, + "member_account_ids": [_ACCOUNT_ID, viewer_id], + }, + app_binding_snapshot=[ + {"app_id": "00000000-0000-0000-0000-000000000031", "caller_kind": "agent"}, + {"app_id": "00000000-0000-0000-0000-000000000032", "caller_kind": "workflow"}, + ], + tag_ids_snapshot=["00000000-0000-0000-0000-000000000033"], + ) + with sqlite_session_factory.begin() as session: + session.add(job) + + runtime = SimpleNamespace( + facade=MagicMock(), + control_plane=MagicMock(), + app_bindings=MagicMock(), + space_tags=MagicMock(), + ) + monkeypatch.setattr(upgrade_module, "get_knowledge_fs_runtime", lambda _: runtime) + + KnowledgeFSUpgradeRunner(sqlite_session_factory)._migrate_access(job) + + members = runtime.control_plane.replace_members.call_args.kwargs["members"] + assert len(members) == 1 + assert members[0].account_id == viewer_id + assert members[0].role is KnowledgeFSControlSpacePermissionRole.VIEWER + assert runtime.control_plane.update_visibility.call_args.kwargs["visibility"] is ( + KnowledgeFSControlSpaceVisibility.PARTIAL_MEMBERS + ) + external_access = runtime.control_plane.update_external_access.call_args.kwargs["payload"] + assert external_access.service_api_enabled is True + assert external_access.agent_enabled is True + assert external_access.workflow_enabled is True + assert external_access.mcp_enabled is False + binding_payloads = [call.kwargs["payload"] for call in runtime.app_bindings.upsert.call_args_list] + assert [payload.caller_kind for payload in binding_payloads] == [ + KnowledgeFSAppSpaceJoinType.AGENT, + KnowledgeFSAppSpaceJoinType.WORKFLOW, + ] + assert runtime.space_tags.replace_tags.call_args.kwargs["tag_ids"] == job.tag_ids_snapshot + with sqlite_session_factory() as session: + assert session.get(KnowledgeFSUpgradeJob, job.id).stage is KnowledgeFSUpgradeStage.FINALIZING + + +@pytest.mark.parametrize( + ("stage", "method_name", "result", "expected"), + [ + (KnowledgeFSUpgradeStage.WAITING_FOR_SPACE, "_advance_when_space_is_active", None, True), + (KnowledgeFSUpgradeStage.CREATING_SOURCES, "_create_next_source", True, True), + (KnowledgeFSUpgradeStage.SUBMITTING_DOCUMENTS, "_submit_next_document", True, True), + (KnowledgeFSUpgradeStage.MIGRATING_ACCESS, "_migrate_access", None, True), + (KnowledgeFSUpgradeStage.FINALIZING, "_finalize", None, False), + ], +) +def test_runner_dispatches_each_durable_stage( + stage: KnowledgeFSUpgradeStage, + method_name: str, + result: bool | None, + expected: bool, +) -> None: + job = _job(status=KnowledgeFSUpgradeJobStatus.QUEUED, stage=stage) + runner = KnowledgeFSUpgradeRunner(MagicMock()) + runner._load_job = MagicMock(return_value=job) # type: ignore[method-assign] + runner._mark_running = MagicMock() # type: ignore[method-assign] + stage_method = MagicMock(return_value=result) + setattr(runner, method_name, stage_method) + + assert runner.run_next(job_id=job.id, celery_task_id="task-1") is expected + + runner._mark_running.assert_called_once_with(job_id=job.id, celery_task_id="task-1") + stage_method.assert_called_once_with(job) + + +def test_runner_validation_creates_space_and_waits() -> None: + job = _job(status=KnowledgeFSUpgradeJobStatus.QUEUED, stage=KnowledgeFSUpgradeStage.VALIDATING) + runner = KnowledgeFSUpgradeRunner(MagicMock()) + runner._load_job = MagicMock(return_value=job) # type: ignore[method-assign] + runner._mark_running = MagicMock() # type: ignore[method-assign] + runner._create_space = MagicMock() # type: ignore[method-assign] + + with pytest.raises(KnowledgeFSUpgradeNotReadyError, match="provisioning is pending"): + runner.run_next(job_id=job.id) + + runner._create_space.assert_called_once_with(job) + + +def test_runner_terminal_and_unknown_stages_do_not_continue() -> None: + runner = KnowledgeFSUpgradeRunner(MagicMock()) + runner._mark_running = MagicMock() # type: ignore[method-assign] + runner._load_job = MagicMock(return_value=_job(status=KnowledgeFSUpgradeJobStatus.SUCCEEDED)) # type: ignore[method-assign] + assert runner.run_next(job_id="done") is False + runner._mark_running.assert_not_called() + + runner._load_job = MagicMock( # type: ignore[method-assign] + return_value=_job(status=KnowledgeFSUpgradeJobStatus.QUEUED, stage=KnowledgeFSUpgradeStage.COMPLETED) + ) + assert runner.run_next(job_id="completed-stage") is False + runner._mark_running.assert_called_once() + + +def test_runner_failure_loading_and_running_state_are_persisted( + sqlite_session_factory: sessionmaker[Session], +) -> None: + job = _job(status=KnowledgeFSUpgradeJobStatus.QUEUED) + with sqlite_session_factory.begin() as session: + session.add(job) + session.flush() + lease = KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=_UPLOAD_FILE_ID, + expires_at=naive_utc_now() + timedelta(minutes=1), + ) + session.add(lease) + + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + loaded = runner._load_job(job.id) + assert loaded.id == job.id + runner._mark_running(job_id=job.id, celery_task_id="task-2") + runner.fail(job_id=job.id, error=ValueError("broken")) + + with sqlite_session_factory() as session: + persisted = session.get(KnowledgeFSUpgradeJob, job.id) + persisted_lease = session.get(KnowledgeFSUpgradeFileLease, lease.id) + assert persisted is not None + assert persisted.status is KnowledgeFSUpgradeJobStatus.FAILED + assert persisted.attempt_count == 1 + assert persisted.celery_task_id == "task-2" + assert persisted.last_error_code == "ValueError" + assert persisted.last_error_message == "broken" + assert persisted.completed_at is not None + assert persisted_lease is not None + assert persisted_lease.expires_at > naive_utc_now() + timedelta(days=6) + + with pytest.raises(KnowledgeFSUpgradeNotFoundError): + runner._load_job("00000000-0000-0000-0000-000000000099") + with pytest.raises(KnowledgeFSUpgradeNotFoundError): + runner._mark_running(job_id="00000000-0000-0000-0000-000000000099", celery_task_id=None) + + runner.fail(job_id="00000000-0000-0000-0000-000000000099", error=RuntimeError("ignored")) + + +def _resolved_configuration() -> dict[str, object]: + return { + "embedding": {"pluginId": "langgenius/openai", "provider": "openai", "model": "embedding"}, + "retrieval": { + "defaultMode": "fast", + "reasoningModel": { + "pluginId": "langgenius/anthropic", + "provider": "anthropic", + "model": "reasoning", + }, + "rerank": { + "enabled": True, + "model": {"pluginId": "langgenius/cohere", "provider": "cohere", "model": "rerank"}, + }, + "scoreThreshold": {"enabled": False, "stage": "mode-final", "value": None}, + "topK": 4, + }, + } + + +def test_runner_creates_space_and_advances_after_activation( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch +) -> None: + job = _job( + status=KnowledgeFSUpgradeJobStatus.RUNNING, + stage=KnowledgeFSUpgradeStage.VALIDATING, + total_sources=1, + config_snapshot={"name": "Legacy", "description": "Description", "icon": "book"}, + ) + with sqlite_session_factory.begin() as session: + session.add(job) + + application = MagicMock() + application.create_space.return_value = SimpleNamespace(control_space_id=_CONTROL_SPACE_ID) + monkeypatch.setattr(upgrade_module, "_resolve_configuration", lambda _job: _resolved_configuration()) + monkeypatch.setattr( + upgrade_module, + "get_knowledge_fs_runtime", + lambda _: SimpleNamespace(application=application), + ) + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + runner._create_space(job) + + payload = application.create_space.call_args.kwargs["payload"] + assert payload.name == "Legacy" + assert payload.idempotency_key == f"upgrade:{job.id}:space" + with sqlite_session_factory.begin() as session: + persisted = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted is not None + assert persisted.stage is KnowledgeFSUpgradeStage.WAITING_FOR_SPACE + control_space = KnowledgeFSControlSpace( + tenant_id=_TENANT_ID, + owner_account_id=_ACCOUNT_ID, + provisioning_key="upgrade-test-space", + knowledge_space_id="00000000-0000-0000-0000-000000000040", + state=KnowledgeFSControlSpaceState.ACTIVE, + ) + session.add(control_space) + session.flush() + persisted.new_control_space_id = control_space.id + control_space_id = control_space.id + + job.new_control_space_id = control_space_id + runner._advance_when_space_is_active(job) + with sqlite_session_factory() as session: + assert session.get(KnowledgeFSUpgradeJob, job.id).stage is KnowledgeFSUpgradeStage.CREATING_SOURCES + + +def test_runner_waiting_for_space_reports_each_invalid_state( + sqlite_session_factory: sessionmaker[Session], +) -> None: + job = _job(status=KnowledgeFSUpgradeJobStatus.RUNNING, stage=KnowledgeFSUpgradeStage.WAITING_FOR_SPACE) + with sqlite_session_factory.begin() as session: + session.add(job) + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + + with pytest.raises(KnowledgeFSUpgradeError, match="reference is missing"): + runner._advance_when_space_is_active(job) + + job.new_control_space_id = _CONTROL_SPACE_ID + with pytest.raises(KnowledgeFSUpgradeError, match="was not found"): + runner._advance_when_space_is_active(job) + + with sqlite_session_factory.begin() as session: + control_space = KnowledgeFSControlSpace( + tenant_id=_TENANT_ID, + owner_account_id=_ACCOUNT_ID, + provisioning_key="upgrade-test-provisioning", + ) + session.add(control_space) + session.flush() + job.new_control_space_id = control_space.id + control_space_id = control_space.id + + with pytest.raises(KnowledgeFSUpgradeNotReadyError, match="provisioning is pending"): + runner._advance_when_space_is_active(job) + + with sqlite_session_factory.begin() as session: + control_space = session.get(KnowledgeFSControlSpace, control_space_id) + assert control_space is not None + control_space.state = KnowledgeFSControlSpaceState.ERROR + with pytest.raises(KnowledgeFSUpgradeError, match="failed in error"): + runner._advance_when_space_is_active(job) + + +def _website_source_payload() -> dict[str, object]: + return { + "kind": "website_crawl", + "name": "Legacy Web", + "provider": "firecrawl", + "datasource": "crawl", + "parameters": {"only_main_content": True}, + "root_url": "https://example.com", + "crawl_options": {"include_subpages": False, "limit": 1}, + "selection": [{"source_url": "https://example.com/page", "title": "Page"}], + "sync_policy": "manual", + } + + +def test_runner_creates_source_and_marks_its_documents_handed_off( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch +) -> None: + source_key = "website:firecrawl:job:1" + job = _job( + status=KnowledgeFSUpgradeJobStatus.RUNNING, + stage=KnowledgeFSUpgradeStage.CREATING_SOURCES, + new_control_space_id=_CONTROL_SPACE_ID, + total_documents=2, + total_sources=1, + ) + with sqlite_session_factory.begin() as session: + session.add(job) + session.flush() + source = KnowledgeFSUpgradeSource( + job_id=job.id, + tenant_id=_TENANT_ID, + source_key=source_key, + payload_snapshot=_website_source_payload(), + ) + documents = [ + KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id=f"00000000-0000-0000-0000-00000000005{index}", + name=f"page-{index}", + data_source_type="website_crawl", + data_source_info={"url": f"https://example.com/{index}"}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + source_key=source_key, + ) + for index in range(2) + ] + session.add_all([source, *documents]) + + from tasks import knowledge_fs_initial_source_tasks as source_tasks + + submit = MagicMock( + return_value=SimpleNamespace( + connection_id="connection-1", + source_id="source-1", + workflow_id="workflow-1", + workflow_error="initial-sync-failed", + ) + ) + monkeypatch.setattr(source_tasks, "submit_initial_source_for_upgrade", submit) + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + assert runner._create_next_source(job) is True + + with sqlite_session_factory() as session: + persisted_job = session.get(KnowledgeFSUpgradeJob, job.id) + persisted_source = session.get(KnowledgeFSUpgradeSource, source.id) + assert persisted_job is not None + assert persisted_job.completed_documents == 2 + assert persisted_job.completed_sources == 1 + assert persisted_source is not None + assert persisted_source.status is KnowledgeFSUpgradeItemStatus.SUCCEEDED + assert persisted_source.last_error_code == "initial-sync-failed" + assert all( + item.status is KnowledgeFSUpgradeItemStatus.SUCCEEDED + for item in session.scalars( + select(KnowledgeFSUpgradeDocument).where(KnowledgeFSUpgradeDocument.job_id == job.id) + ) + ) + + assert runner._create_next_source(job) is True + with sqlite_session_factory() as session: + assert session.get(KnowledgeFSUpgradeJob, job.id).stage is KnowledgeFSUpgradeStage.SUBMITTING_DOCUMENTS + + +@pytest.mark.parametrize("not_ready", [True, False]) +def test_runner_source_failure_is_retryable_or_persisted( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, not_ready: bool +) -> None: + job = _job( + status=KnowledgeFSUpgradeJobStatus.RUNNING, + stage=KnowledgeFSUpgradeStage.CREATING_SOURCES, + new_control_space_id=_CONTROL_SPACE_ID, + total_sources=1, + ) + with sqlite_session_factory.begin() as session: + session.add(job) + session.flush() + source = KnowledgeFSUpgradeSource( + job_id=job.id, + tenant_id=_TENANT_ID, + source_key=f"source-failure-{not_ready}", + payload_snapshot=_website_source_payload(), + ) + session.add(source) + + from tasks import knowledge_fs_initial_source_tasks as source_tasks + + error: Exception + if not_ready: + error = source_tasks.KnowledgeFSInitialSourceNotReadyError("busy") + else: + error = RuntimeError("source failed") + monkeypatch.setattr(source_tasks, "submit_initial_source_for_upgrade", MagicMock(side_effect=error)) + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + + expected_error = KnowledgeFSUpgradeNotReadyError if not_ready else RuntimeError + with pytest.raises(expected_error): + runner._create_next_source(job) + + with sqlite_session_factory() as session: + persisted = session.get(KnowledgeFSUpgradeSource, source.id) + assert persisted is not None + expected_status = KnowledgeFSUpgradeItemStatus.PENDING if not_ready else KnowledgeFSUpgradeItemStatus.FAILED + assert persisted.status is expected_status + if not not_ready: + assert persisted.last_error_code == "RuntimeError" + assert persisted.last_error_message == "source failed" + + +def test_runner_document_boundaries_and_successful_finalize( + sqlite_session_factory: sessionmaker[Session], +) -> None: + job = _job( + status=KnowledgeFSUpgradeJobStatus.RUNNING, + stage=KnowledgeFSUpgradeStage.SUBMITTING_DOCUMENTS, + new_control_space_id=_CONTROL_SPACE_ID, + total_documents=1, + ) + with sqlite_session_factory.begin() as session: + session.add(job) + session.flush() + document = KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000060", + name="missing.txt", + data_source_type="upload_file", + data_source_info={}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + ) + session.add(document) + + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + with pytest.raises(KnowledgeFSUpgradeError, match="no source file reference"): + runner._submit_next_document(job) + with sqlite_session_factory.begin() as session: + persisted_document = session.get(KnowledgeFSUpgradeDocument, document.id) + assert persisted_document is not None + assert persisted_document.status is KnowledgeFSUpgradeItemStatus.FAILED + persisted_document.status = KnowledgeFSUpgradeItemStatus.SUCCEEDED + persisted_job = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted_job is not None + persisted_job.completed_documents = 1 + + assert runner._submit_next_document(job) is True + with sqlite_session_factory.begin() as session: + persisted_job = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted_job is not None + assert persisted_job.stage is KnowledgeFSUpgradeStage.MIGRATING_ACCESS + persisted_job.stage = KnowledgeFSUpgradeStage.FINALIZING + + runner._finalize(job) + with sqlite_session_factory() as session: + persisted_job = session.get(KnowledgeFSUpgradeJob, job.id) + assert persisted_job is not None + assert persisted_job.status is KnowledgeFSUpgradeJobStatus.SUCCEEDED + assert persisted_job.stage is KnowledgeFSUpgradeStage.COMPLETED + assert persisted_job.completed_at is not None + + +def test_upgrade_helper_source_identities_payloads_and_validation() -> None: + notion = KnowledgeFSUpgradeDocument( + job_id="00000000-0000-0000-0000-000000000070", + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000071", + name="Notion page", + data_source_type="notion_import", + data_source_info={"workspace_id": "workspace", "notion_page_id": "page"}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + ) + website = KnowledgeFSUpgradeDocument( + job_id="00000000-0000-0000-0000-000000000070", + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000072", + name="Website page", + data_source_type="website_crawl", + data_source_info={"url": "https://example.com/page"}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + ) + unsupported = KnowledgeFSUpgradeDocument( + job_id="00000000-0000-0000-0000-000000000070", + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000073", + name="Unsupported", + data_source_type="upload_file", + data_source_info={}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + ) + + assert upgrade_module._expected_provider_item_id(notion) == '["workspace","page"]' + assert ( + upgrade_module._expected_provider_item_id(website) + == upgrade_module.sha256(b"https://example.com/page").hexdigest() + ) + with pytest.raises(KnowledgeFSUpgradeError, match="no provider item identity"): + upgrade_module._expected_provider_item_id(unsupported) + + assert upgrade_module._notion_source_group_key({"workspace_id": "workspace", "credential_id": "cred"}) == ( + "notion:workspace:cred" + ) + assert ( + upgrade_module._website_source_group_key( + {"provider": " Firecrawl ", "job_id": "crawl", "url": "https://example.com"} + ) + == "website:firecrawl:crawl" + ) + with pytest.raises(KnowledgeFSUpgradeConflictError): + upgrade_module._notion_source_group_key({}) + with pytest.raises(KnowledgeFSUpgradeConflictError): + upgrade_module._website_source_group_key({"provider": "firecrawl"}) + + notion_document = SimpleNamespace( + data_source_type="notion_import", + name="Page", + ) + notion_payload = upgrade_module._source_payload_snapshot( + "Dataset", + [(notion_document, {"workspace_id": "workspace", "notion_page_id": "page", "type": "page"})], + ) + assert notion_payload["kind"] == "online_document" + website_payload = upgrade_module._source_payload_snapshot( + "Dataset", + [(SimpleNamespace(data_source_type="website_crawl", name="Page"), {"provider": "firecrawl", "url": "u"})], + ) + assert website_payload["kind"] == "website_crawl" + assert upgrade_module._chunks([1, 2, 3], 2) == [[1, 2], [3]] + with pytest.raises(KnowledgeFSUpgradeError, match="reference is missing"): + upgrade_module._required_space_id(_job()) + + +def test_notion_credential_resolution_prefers_direct_then_workspace_match( + sqlite_session_factory: sessionmaker[Session], +) -> None: + with sqlite_session_factory.begin() as session: + direct = DatasourceProvider( + tenant_id=_TENANT_ID, + name="Direct", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + auth_type="oauth2", + encrypted_credentials={"workspace_id": "direct-workspace"}, + ) + fallback = DatasourceProvider( + tenant_id=_TENANT_ID, + name="Fallback", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + auth_type="oauth2", + encrypted_credentials={"workspace_id": "fallback-workspace"}, + ) + session.add_all([direct, fallback]) + session.flush() + direct_id = direct.id + fallback_id = fallback.id + + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + assert ( + runner._resolve_notion_credential( + tenant_id=_TENANT_ID, + legacy_credential_id=direct_id, + workspace_id="ignored", + ) + == direct_id + ) + assert ( + runner._resolve_notion_credential( + tenant_id=_TENANT_ID, + legacy_credential_id="00000000-0000-0000-0000-000000000090", + workspace_id="fallback-workspace", + ) + == fallback_id + ) + + +def test_notion_credential_resolution_handles_single_fallback_and_missing( + sqlite_session_factory: sessionmaker[Session], +) -> None: + single_tenant_id = "00000000-0000-0000-0000-000000000091" + missing_tenant_id = "00000000-0000-0000-0000-000000000092" + with sqlite_session_factory.begin() as session: + only = DatasourceProvider( + tenant_id=single_tenant_id, + name="Only", + provider="notion_datasource", + plugin_id="langgenius/notion_datasource", + auth_type="oauth2", + encrypted_credentials={"workspace_id": "another-workspace"}, + ) + session.add(only) + session.flush() + only_id = only.id + + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + assert ( + runner._resolve_notion_credential( + tenant_id=single_tenant_id, + legacy_credential_id="", + workspace_id="unmatched", + ) + == only_id + ) + with pytest.raises(KnowledgeFSUpgradeError, match="credential is unavailable"): + runner._resolve_notion_credential( + tenant_id=missing_tenant_id, + legacy_credential_id="", + workspace_id="missing", + ) + + +def test_document_submission_persists_missing_file_and_invalid_body_errors( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch +) -> None: + account = Account(name="Upgrade owner", email="upgrade-errors@example.com") + with sqlite_session_factory.begin() as session: + session.add(account) + session.flush() + job = _job( + owner_account_id=account.id, + status=KnowledgeFSUpgradeJobStatus.RUNNING, + stage=KnowledgeFSUpgradeStage.SUBMITTING_DOCUMENTS, + new_control_space_id=_CONTROL_SPACE_ID, + total_documents=2, + ) + session.add(job) + session.flush() + missing = KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000093", + name="missing.txt", + data_source_type="upload_file", + data_source_info={"upload_file_id": _UPLOAD_FILE_ID}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + old_upload_file_id=_UPLOAD_FILE_ID, + ) + session.add(missing) + + staged = MagicMock() + monkeypatch.setattr(upgrade_module, "KnowledgeFSStagedUploadService", lambda *_args, **_kwargs: staged) + monkeypatch.setattr(upgrade_module, "get_knowledge_fs_runtime", lambda _: SimpleNamespace(facade=MagicMock())) + runner = KnowledgeFSUpgradeRunner(sqlite_session_factory) + with pytest.raises(KnowledgeFSUpgradeError, match="source file is unavailable"): + runner._submit_next_document(job) + + with sqlite_session_factory.begin() as session: + persisted_missing = session.get(KnowledgeFSUpgradeDocument, missing.id) + assert persisted_missing is not None + assert persisted_missing.status is KnowledgeFSUpgradeItemStatus.FAILED + upload_file = UploadFile( + tenant_id=_TENANT_ID, + storage_type=StorageType.LOCAL, + key=f"upload_files/{_TENANT_ID}/invalid.txt", + name="invalid.txt", + size=12, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=account.id, + created_at=naive_utc_now(), + used=False, + ) + session.add(upload_file) + session.flush() + invalid = KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000094", + name="invalid.txt", + data_source_type="upload_file", + data_source_info={"upload_file_id": upload_file.id}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + old_upload_file_id=upload_file.id, + ) + session.add(invalid) + + monkeypatch.setattr(upgrade_module.storage, "load", lambda _key: "not-bytes") + with pytest.raises(KnowledgeFSUpgradeError, match="invalid body"): + runner._submit_next_document(job) + with sqlite_session_factory() as session: + persisted_invalid = session.get(KnowledgeFSUpgradeDocument, invalid.id) + assert persisted_invalid is not None + assert persisted_invalid.status is KnowledgeFSUpgradeItemStatus.FAILED + + +def test_reconciler_records_missing_remote_document_and_remote_error( + sqlite_session_factory: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch +) -> None: + source_key = "website:firecrawl:job:1" + job = _job( + status=KnowledgeFSUpgradeJobStatus.SUCCEEDED, + stage=KnowledgeFSUpgradeStage.COMPLETED, + new_control_space_id=_CONTROL_SPACE_ID, + ) + with sqlite_session_factory.begin() as session: + session.add(job) + session.flush() + source = KnowledgeFSUpgradeSource( + job_id=job.id, + tenant_id=_TENANT_ID, + source_key=source_key, + payload_snapshot=_website_source_payload(), + status=KnowledgeFSUpgradeItemStatus.SUCCEEDED, + new_source_id="source-1", + ) + missing = KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000095", + name="missing-page", + data_source_type="website_crawl", + data_source_info={"url": "https://example.com/missing"}, + metadata_snapshot={}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + status=KnowledgeFSUpgradeItemStatus.SUCCEEDED, + source_key=source_key, + ) + remote_error = KnowledgeFSUpgradeDocument( + job_id=job.id, + tenant_id=_TENANT_ID, + old_document_id="00000000-0000-0000-0000-000000000096", + name="remote-error", + data_source_type="upload_file", + data_source_info={"upload_file_id": _UPLOAD_FILE_ID}, + metadata_snapshot={"department": "support"}, + desired_enabled=True, + legacy_archived=False, + legacy_indexing_status="completed", + status=KnowledgeFSUpgradeItemStatus.SUCCEEDED, + new_document_asset_id=_DOCUMENT_ASSET_ID, + ) + session.add_all([source, missing, remote_error]) + + facade = MagicMock() + facade.list_logical_documents.return_value = KnowledgeFSLogicalDocumentListResponse( + data=[_logical_document(enabled=True)], next_cursor=None + ) + facade.update_document_metadata.side_effect = RuntimeError("remote metadata failed") + monkeypatch.setattr(upgrade_module, "get_knowledge_fs_runtime", lambda _: SimpleNamespace(facade=facade)) + + assert KnowledgeFSUpgradeDocumentReconciler(sqlite_session_factory).reconcile(job_id=job.id) == 2 + with sqlite_session_factory() as session: + persisted_missing = session.get(KnowledgeFSUpgradeDocument, missing.id) + persisted_error = session.get(KnowledgeFSUpgradeDocument, remote_error.id) + assert persisted_missing is not None + assert persisted_missing.state_reconcile_attempt_count == 1 + assert persisted_missing.state_reconcile_error == "The new logical document is not visible yet" + assert persisted_error is not None + assert persisted_error.state_reconcile_attempt_count == 1 + assert persisted_error.state_reconcile_error == "remote metadata failed" diff --git a/api/tests/unit_tests/tasks/test_batch_clean_document_task.py b/api/tests/unit_tests/tasks/test_batch_clean_document_task.py index 59fab89ab9c..ad2dd61b4a5 100644 --- a/api/tests/unit_tests/tasks/test_batch_clean_document_task.py +++ b/api/tests/unit_tests/tasks/test_batch_clean_document_task.py @@ -1,12 +1,16 @@ import uuid +from datetime import UTC, datetime, timedelta from unittest.mock import patch import pytest from sqlalchemy.orm import Session import tasks.batch_clean_document_task as task_module +from extensions.storage.storage_type import StorageType from models.dataset import Dataset, DocumentSegment -from models.enums import DataSourceType +from models.enums import CreatorUserRole, DataSourceType +from models.knowledge_fs import KnowledgeFSUpgradeFileLease, KnowledgeFSUpgradeJob +from models.model import UploadFile from tasks.batch_clean_document_task import batch_clean_document_task @@ -83,3 +87,80 @@ def test_failed_vector_cleanup_does_not_schedule_billing_refresh(cleanup_rows: t ) schedule_refresh.assert_not_called() + + +def test_batch_cleanup_keeps_only_the_leased_legacy_source_file( + cleanup_rows: tuple[str, str, str], sqlite_session: Session +) -> None: + dataset_id, document_id, tenant_id = cleanup_rows + account_id = str(uuid.uuid4()) + now = datetime.now(UTC).replace(tzinfo=None) + leased_file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key=f"upload_files/{tenant_id}/leased.txt", + name="leased.txt", + size=10, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=account_id, + created_at=now, + used=False, + ) + deletable_file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key=f"upload_files/{tenant_id}/deletable.txt", + name="deletable.txt", + size=10, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=account_id, + created_at=now, + used=False, + ) + job = KnowledgeFSUpgradeJob( + tenant_id=tenant_id, + old_dataset_id=dataset_id, + requested_by_account_id=account_id, + owner_account_id=account_id, + idempotency_key="batch-cleanup-lease-test", + snapshot_at=now, + config_snapshot={}, + permission_snapshot={}, + app_binding_snapshot=[], + tag_ids_snapshot=[], + ) + sqlite_session.add_all([leased_file, deletable_file, job]) + sqlite_session.flush() + sqlite_session.add( + KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=leased_file.id, + expires_at=now + timedelta(hours=1), + ) + ) + sqlite_session.commit() + leased_file_id = leased_file.id + deletable_file_id = deletable_file.id + deletable_file_key = deletable_file.key + + with ( + patch("tasks.batch_clean_document_task.get_image_upload_file_ids", return_value=[]), + patch("tasks.batch_clean_document_task.IndexProcessorFactory"), + patch("tasks.batch_clean_document_task.schedule_billing_vector_space_refresh"), + patch("tasks.batch_clean_document_task.storage") as storage, + ): + batch_clean_document_task( + document_ids=[document_id], + dataset_id=dataset_id, + doc_form="paragraph", + file_ids=[leased_file_id, deletable_file_id], + ) + + sqlite_session.expire_all() + assert sqlite_session.get(UploadFile, leased_file_id) is not None + assert sqlite_session.get(UploadFile, deletable_file_id) is None + storage.delete.assert_called_once_with(deletable_file_key) diff --git a/api/tests/unit_tests/tasks/test_clean_dataset_task.py b/api/tests/unit_tests/tasks/test_clean_dataset_task.py index 3ac74ec1f4f..359515e7b62 100644 --- a/api/tests/unit_tests/tasks/test_clean_dataset_task.py +++ b/api/tests/unit_tests/tasks/test_clean_dataset_task.py @@ -14,7 +14,7 @@ This module tests the dataset cleanup task functionality including: import json import uuid from collections.abc import Iterator -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from unittest.mock import MagicMock, patch import pytest @@ -36,6 +36,7 @@ from models.dataset import ( SegmentAttachmentBinding, ) from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus +from models.knowledge_fs import KnowledgeFSUpgradeFileLease, KnowledgeFSUpgradeJob from models.model import UploadFile from models.workflow import Workflow, WorkflowType from tasks.clean_dataset_task import clean_dataset_task @@ -579,3 +580,77 @@ class TestIndexProcessorParameters: ) schedule_refresh.assert_not_called() + + +def test_dataset_cleanup_keeps_a_leased_legacy_source_file( + dataset_id: str, + tenant_id: str, + collection_binding_id: str, + orm_session_maker: sessionmaker[Session], + mock_storage: MagicMock, + mock_index_processor_factory: dict[str, MagicMock], + mock_get_image_upload_file_ids: MagicMock, +) -> None: + del mock_index_processor_factory, mock_get_image_upload_file_ids + account_id = str(uuid.uuid4()) + now = datetime.now(UTC).replace(tzinfo=None) + upload_file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key=f"upload_files/{tenant_id}/dataset-source.txt", + name="dataset-source.txt", + size=10, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=account_id, + created_at=now, + used=False, + ) + document = Document( + id=str(uuid.uuid4()), + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + data_source_info=json.dumps({"upload_file_id": upload_file.id}), + batch="batch", + name="dataset-source.txt", + created_from=DocumentCreatedFrom.WEB, + created_by=account_id, + indexing_status=IndexingStatus.COMPLETED, + doc_form=IndexStructureType.PARAGRAPH_INDEX, + ) + job = KnowledgeFSUpgradeJob( + tenant_id=tenant_id, + old_dataset_id=dataset_id, + requested_by_account_id=account_id, + owner_account_id=account_id, + idempotency_key="dataset-cleanup-lease-test", + snapshot_at=now, + config_snapshot={}, + permission_snapshot={}, + app_binding_snapshot=[], + tag_ids_snapshot=[], + ) + with orm_session_maker.begin() as session: + session.add_all([upload_file, document, job]) + session.flush() + session.add( + KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=upload_file.id, + expires_at=now + timedelta(hours=1), + ) + ) + + _run_clean_dataset( + dataset_id=dataset_id, + tenant_id=tenant_id, + collection_binding_id=collection_binding_id, + ) + + with orm_session_maker() as session: + assert session.get(Document, document.id) is None + assert session.get(UploadFile, upload_file.id) is not None + mock_storage.delete.assert_not_called() diff --git a/api/tests/unit_tests/tasks/test_clean_document_task.py b/api/tests/unit_tests/tasks/test_clean_document_task.py index 4aef0581933..7f1eac13399 100644 --- a/api/tests/unit_tests/tasks/test_clean_document_task.py +++ b/api/tests/unit_tests/tasks/test_clean_document_task.py @@ -6,6 +6,7 @@ starts from the production incident shape: the caller has already deleted the """ import uuid +from datetime import UTC, datetime, timedelta from unittest.mock import MagicMock, patch import pytest @@ -13,6 +14,7 @@ from sqlalchemy import select from sqlalchemy.orm import Session import tasks.clean_document_task as clean_document_task_module +from extensions.storage.storage_type import StorageType from models.dataset import ( Dataset, DatasetMetadataBinding, @@ -20,7 +22,8 @@ from models.dataset import ( DocumentSegment, SegmentAttachmentBinding, ) -from models.enums import DataSourceType, DocumentCreatedFrom +from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom +from models.knowledge_fs import KnowledgeFSUpgradeFileLease, KnowledgeFSUpgradeJob from models.model import UploadFile from tasks.clean_document_task import clean_document_task @@ -244,6 +247,77 @@ class TestVectorCleanupResilience: ) schedule_refresh.assert_not_called() + +def test_active_upgrade_lease_keeps_the_legacy_source_file( + document_id: str, + dataset_id: str, + tenant_id: str, + sqlite_session: Session, + bind_task_sessions: None, + mock_storage, + mock_index_processor_factory, +) -> None: + del bind_task_sessions, mock_index_processor_factory + _persist_deleted_document_state( + sqlite_session, + document_id=document_id, + dataset_id=dataset_id, + tenant_id=tenant_id, + target_segment_ids=[], + ) + account_id = str(uuid.uuid4()) + upload_file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key=f"upload_files/{tenant_id}/source.txt", + name="source.txt", + size=12, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=account_id, + created_at=datetime(2026, 8, 17, tzinfo=UTC), + used=False, + ) + now = datetime.now(UTC).replace(tzinfo=None) + job = KnowledgeFSUpgradeJob( + tenant_id=tenant_id, + old_dataset_id=dataset_id, + requested_by_account_id=account_id, + owner_account_id=account_id, + idempotency_key="cleanup-lease-test", + snapshot_at=now, + config_snapshot={}, + permission_snapshot={}, + app_binding_snapshot=[], + tag_ids_snapshot=[], + ) + sqlite_session.add_all([upload_file, job]) + sqlite_session.flush() + lease = KnowledgeFSUpgradeFileLease( + job_id=job.id, + old_upload_file_id=upload_file.id, + expires_at=now + timedelta(hours=1), + ) + sqlite_session.add(lease) + sqlite_session.commit() + + clean_document_task( + document_id=document_id, + dataset_id=dataset_id, + doc_form="paragraph", + file_id=upload_file.id, + ) + + sqlite_session.expire_all() + assert sqlite_session.get(UploadFile, upload_file.id) is not None + persisted_lease = sqlite_session.get(KnowledgeFSUpgradeFileLease, lease.id) + assert persisted_lease is not None + assert persisted_lease.cleanup_requested_at is not None + mock_storage.delete.assert_not_called() + + +class TestVectorCleanupSuccessPaths: def test_vector_cleanup_success_path_remains_unaffected( self, document_id: str, diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py index b88ef81d1d0..9af97d2ad13 100644 --- a/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_initial_source_tasks.py @@ -17,6 +17,7 @@ from tasks.knowledge_fs_initial_source_tasks import ( import_initial_website_source, start_initial_source_import, start_initial_website_source_import, + submit_initial_source_for_upgrade, ) _DEFAULT_CREDENTIAL = object() @@ -794,3 +795,37 @@ def test_initial_source_task_does_not_retry_authoritative_missing_resource() -> ) retry.assert_not_called() + + +def test_upgrade_source_submission_succeeds_when_first_import_submission_fails() -> None: + facade = _facade() + facade.create_source.return_value = SimpleNamespace( + connection_id="connection-1", + id="source-1", + metadata={"clientRequestId": "initial-website-source:operation-1", "preview": True}, + status="disabled", + version=1, + ) + facade.update_source.return_value = SimpleNamespace( + connection_id="connection-1", + id="source-1", + metadata={"clientRequestId": "initial-website-source:operation-1", "preview": False}, + status="active", + version=2, + ) + facade.import_selected_source_crawl.side_effect = RuntimeError("crawl queue unavailable") + + with _runtime(facade): + result = submit_initial_source_for_upgrade( + tenant_id="tenant-1", + account_id="account-1", + control_space_id="control-1", + operation_id="operation-1", + payload=_payload(sync_policy="manual"), + ) + + assert result.connection_id == "connection-1" + assert result.source_id == "source-1" + assert result.workflow_id is None + assert result.workflow_error == "RuntimeError" + facade.update_source_sync_policy.assert_called_once() diff --git a/api/tests/unit_tests/tasks/test_knowledge_fs_upgrade_tasks.py b/api/tests/unit_tests/tasks/test_knowledge_fs_upgrade_tasks.py new file mode 100644 index 00000000000..8b4f5ae0456 --- /dev/null +++ b/api/tests/unit_tests/tasks/test_knowledge_fs_upgrade_tasks.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from services.dataset_knowledge_fs_upgrade_service import KnowledgeFSUpgradeNotReadyError +from tasks.knowledge_fs_upgrade_tasks import ( + KNOWLEDGE_FS_UPGRADE_QUEUE, + cleanup_deferred_knowledge_fs_upgrade_files, + reconcile_knowledge_fs_upgrade_documents, + run_knowledge_fs_upgrade, +) + + +def test_upgrade_tasks_are_pinned_to_the_dedicated_queue() -> None: + assert run_knowledge_fs_upgrade._get_exec_options()["queue"] == KNOWLEDGE_FS_UPGRADE_QUEUE + assert reconcile_knowledge_fs_upgrade_documents._get_exec_options()["queue"] == KNOWLEDGE_FS_UPGRADE_QUEUE + assert cleanup_deferred_knowledge_fs_upgrade_files._get_exec_options()["queue"] == KNOWLEDGE_FS_UPGRADE_QUEUE + + +def test_deferred_file_cleanup_uses_the_upgrade_session_factory() -> None: + with ( + patch("tasks.knowledge_fs_upgrade_tasks.session_factory.get_session_maker", return_value="maker"), + patch("tasks.knowledge_fs_upgrade_tasks.cleanup_deferred_upgrade_files", return_value=3) as cleanup, + ): + assert cleanup_deferred_knowledge_fs_upgrade_files.run() == 3 + + cleanup.assert_called_once_with("maker") + + +@pytest.mark.parametrize("has_more", [True, False]) +def test_worker_enqueues_only_the_expected_follow_up(has_more: bool) -> None: + runner = MagicMock() + runner.run_next.return_value = has_more + with ( + patch("tasks.knowledge_fs_upgrade_tasks.session_factory.get_session_maker", return_value="maker"), + patch("tasks.knowledge_fs_upgrade_tasks.KnowledgeFSUpgradeRunner", return_value=runner), + patch.object(run_knowledge_fs_upgrade, "apply_async") as continue_upgrade, + patch.object(reconcile_knowledge_fs_upgrade_documents, "apply_async") as reconcile, + ): + run_knowledge_fs_upgrade.run(job_id="job-1") + + runner.run_next.assert_called_once() + if has_more: + continue_upgrade.assert_called_once_with(kwargs={"job_id": "job-1"}) + reconcile.assert_not_called() + else: + continue_upgrade.assert_not_called() + reconcile.assert_called_once_with(kwargs={"job_id": "job-1"}) + + +def test_worker_marks_parent_failed_when_a_checkpoint_raises() -> None: + runner = MagicMock() + error = RuntimeError("ordinary document upload failed") + runner.run_next.side_effect = error + with ( + patch("tasks.knowledge_fs_upgrade_tasks.session_factory.get_session_maker", return_value="maker"), + patch("tasks.knowledge_fs_upgrade_tasks.KnowledgeFSUpgradeRunner", return_value=runner), + pytest.raises(RuntimeError, match="ordinary document upload failed"), + ): + run_knowledge_fs_upgrade.run(job_id="job-1") + + runner.fail.assert_called_once_with(job_id="job-1", error=error) + + +def test_worker_marks_parent_failed_when_provisioning_retries_are_exhausted() -> None: + runner = MagicMock() + error = KnowledgeFSUpgradeNotReadyError("Space is still provisioning") + runner.run_next.side_effect = error + run_knowledge_fs_upgrade.push_request(retries=run_knowledge_fs_upgrade.max_retries) + try: + with ( + patch("tasks.knowledge_fs_upgrade_tasks.session_factory.get_session_maker", return_value="maker"), + patch("tasks.knowledge_fs_upgrade_tasks.KnowledgeFSUpgradeRunner", return_value=runner), + pytest.raises(KnowledgeFSUpgradeNotReadyError, match="still provisioning"), + ): + run_knowledge_fs_upgrade.run(job_id="job-1") + finally: + run_knowledge_fs_upgrade.pop_request() + + runner.fail.assert_called_once_with(job_id="job-1", error=error) diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml index 6c23ebf7df6..4e080afc849 100644 --- a/docker/docker-compose-template.yaml +++ b/docker/docker-compose-template.yaml @@ -346,6 +346,36 @@ services: - ssrf_proxy_network - default + # Dedicated worker for click-time Dataset snapshots and KnowledgeFS upgrade orchestration. + # Keep this queue out of the generic worker so large file copies cannot starve existing jobs. + worker_knowledge_fs_upgrade: + <<: *shared-worker-config + image: langgenius/dify-api:1.16.1 + environment: + MODE: worker + CELERY_WORKER_QUEUES: knowledge_fs_upgrade + CELERY_WORKER_CONCURRENCY: ${KNOWLEDGE_FS_UPGRADE_WORKER_AMOUNT:-2} + CELERY_PREFETCH_MULTIPLIER: 1 + SENTRY_DSN: ${API_SENTRY_DSN:-} + SENTRY_TRACES_SAMPLE_RATE: ${API_SENTRY_TRACES_SAMPLE_RATE:-1.0} + SENTRY_PROFILES_SAMPLE_RATE: ${API_SENTRY_PROFILES_SAMPLE_RATE:-1.0} + depends_on: + init_permissions: + condition: service_completed_successfully + db_postgres: + condition: service_healthy + required: false + db_mysql: + condition: service_healthy + required: false + redis: + condition: service_started + volumes: + - ./volumes/app/storage:/app/api/storage + networks: + - ssrf_proxy_network + - default + # worker_beat service # Celery beat for scheduling periodic tasks. worker_beat: diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index d06fb2ec575..9b71a655108 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -352,6 +352,36 @@ services: - ssrf_proxy_network - default + # Dedicated worker for click-time Dataset snapshots and KnowledgeFS upgrade orchestration. + # Keep this queue out of the generic worker so large file copies cannot starve existing jobs. + worker_knowledge_fs_upgrade: + <<: *shared-worker-config + image: langgenius/dify-api:1.16.1 + environment: + MODE: worker + CELERY_WORKER_QUEUES: knowledge_fs_upgrade + CELERY_WORKER_CONCURRENCY: ${KNOWLEDGE_FS_UPGRADE_WORKER_AMOUNT:-2} + CELERY_PREFETCH_MULTIPLIER: 1 + SENTRY_DSN: ${API_SENTRY_DSN:-} + SENTRY_TRACES_SAMPLE_RATE: ${API_SENTRY_TRACES_SAMPLE_RATE:-1.0} + SENTRY_PROFILES_SAMPLE_RATE: ${API_SENTRY_PROFILES_SAMPLE_RATE:-1.0} + depends_on: + init_permissions: + condition: service_completed_successfully + db_postgres: + condition: service_healthy + required: false + db_mysql: + condition: service_healthy + required: false + redis: + condition: service_started + volumes: + - ./volumes/app/storage:/app/api/storage + networks: + - ssrf_proxy_network + - default + # worker_beat service # Celery beat for scheduling periodic tasks. worker_beat: diff --git a/knowledge-fs/secret-scan-allowlist.json b/knowledge-fs/secret-scan-allowlist.json index 2532a23956e..8997482255d 100644 --- a/knowledge-fs/secret-scan-allowlist.json +++ b/knowledge-fs/secret-scan-allowlist.json @@ -59,14 +59,14 @@ }, { "fingerprint": "ac8011331002c601c2ec9c4dd7f2328e4d3bdc511021e21d4f9e3b67f0f96bc6", - "line": 558, + "line": 588, "path": "api/models/knowledge_fs.py", "reason": "Fixed database schema identifier; this literal is not credential material.", "rule": "dify-knowledge-fs-credential" }, { "fingerprint": "fe0a852bad3684b159154a873bcfe0e721e771aa1a4bc0d6b5fee9cb47cf17fd", - "line": 596, + "line": 626, "path": "api/models/knowledge_fs.py", "reason": "Fixed database schema identifier; this literal is not credential material.", "rule": "dify-knowledge-fs-credential" diff --git a/packages/contracts/generated/api/console/datasets/orpc.gen.ts b/packages/contracts/generated/api/console/datasets/orpc.gen.ts index 326bf34b30d..8a39e218bc8 100644 --- a/packages/contracts/generated/api/console/datasets/orpc.gen.ts +++ b/packages/contracts/generated/api/console/datasets/orpc.gen.ts @@ -66,6 +66,8 @@ import { zGetDatasetsByDatasetIdErrorDocsResponse, zGetDatasetsByDatasetIdIndexingStatusPath, zGetDatasetsByDatasetIdIndexingStatusResponse, + zGetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdPath, + zGetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponse, zGetDatasetsByDatasetIdMetadataPath, zGetDatasetsByDatasetIdMetadataResponse, zGetDatasetsByDatasetIdNotionSyncPath, @@ -162,6 +164,10 @@ import { zPostDatasetsByDatasetIdHitTestingBody, zPostDatasetsByDatasetIdHitTestingPath, zPostDatasetsByDatasetIdHitTestingResponse, + zPostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdPath, + zPostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponse, + zPostDatasetsByDatasetIdKnowledgeFsUpgradesPath, + zPostDatasetsByDatasetIdKnowledgeFsUpgradesResponse, zPostDatasetsByDatasetIdMetadataBody, zPostDatasetsByDatasetIdMetadataBuiltInByActionPath, zPostDatasetsByDatasetIdMetadataBuiltInByActionResponse, @@ -1424,7 +1430,52 @@ export const indexingStatus3 = { get: get27, } +export const get28 = oc + .route({ + inputStructure: 'detailed', + method: 'GET', + operationId: 'getDatasetsByDatasetIdKnowledgeFsUpgradesByJobId', + path: '/datasets/{dataset_id}/knowledge-fs-upgrades/{job_id}', + tags: ['console'], + }) + .input(z.object({ params: zGetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdPath })) + .output(zGetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponse) + export const post19 = oc + .route({ + inputStructure: 'detailed', + method: 'POST', + operationId: 'postDatasetsByDatasetIdKnowledgeFsUpgradesByJobId', + path: '/datasets/{dataset_id}/knowledge-fs-upgrades/{job_id}', + successStatus: 202, + tags: ['console'], + }) + .input(z.object({ params: zPostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdPath })) + .output(zPostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponse) + +export const byJobId2 = { + get: get28, + post: post19, +} + +export const post20 = oc + .route({ + inputStructure: 'detailed', + method: 'POST', + operationId: 'postDatasetsByDatasetIdKnowledgeFsUpgrades', + path: '/datasets/{dataset_id}/knowledge-fs-upgrades', + successStatus: 202, + tags: ['console'], + }) + .input(z.object({ params: zPostDatasetsByDatasetIdKnowledgeFsUpgradesPath })) + .output(zPostDatasetsByDatasetIdKnowledgeFsUpgradesResponse) + +export const knowledgeFsUpgrades = { + post: post20, + byJobId: byJobId2, +} + +export const post21 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1437,7 +1488,7 @@ export const post19 = oc .output(zPostDatasetsByDatasetIdMetadataBuiltInByActionResponse) export const byAction4 = { - post: post19, + post: post21, } export const builtIn2 = { @@ -1477,7 +1528,7 @@ export const byMetadataId = { patch: patch10, } -export const get28 = oc +export const get29 = oc .route({ inputStructure: 'detailed', method: 'GET', @@ -1488,7 +1539,7 @@ export const get28 = oc .input(z.object({ params: zGetDatasetsByDatasetIdMetadataPath })) .output(zGetDatasetsByDatasetIdMetadataResponse) -export const post20 = oc +export const post22 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1506,13 +1557,13 @@ export const post20 = oc .output(zPostDatasetsByDatasetIdMetadataResponse) export const metadata4 = { - get: get28, - post: post20, + get: get29, + post: post22, builtIn: builtIn2, byMetadataId, } -export const get29 = oc +export const get30 = oc .route({ inputStructure: 'detailed', method: 'GET', @@ -1524,7 +1575,7 @@ export const get29 = oc .output(zGetDatasetsByDatasetIdNotionSyncResponse) export const sync2 = { - get: get29, + get: get30, } export const notion2 = { @@ -1534,7 +1585,7 @@ export const notion2 = { /** * Get dataset permission user list */ -export const get30 = oc +export const get31 = oc .route({ description: 'Get dataset permission user list', inputStructure: 'detailed', @@ -1547,13 +1598,13 @@ export const get30 = oc .output(zGetDatasetsByDatasetIdPermissionPartUsersResponse) export const permissionPartUsers = { - get: get30, + get: get31, } /** * Get dataset query history */ -export const get31 = oc +export const get32 = oc .route({ description: 'Get dataset query history', inputStructure: 'detailed', @@ -1566,13 +1617,13 @@ export const get31 = oc .output(zGetDatasetsByDatasetIdQueriesResponse) export const queries = { - get: get31, + get: get32, } /** * Get applications related to dataset */ -export const get32 = oc +export const get33 = oc .route({ description: 'Get applications related to dataset', inputStructure: 'detailed', @@ -1585,13 +1636,13 @@ export const get32 = oc .output(zGetDatasetsByDatasetIdRelatedAppsResponse) export const relatedApps = { - get: get32, + get: get33, } /** * retry document */ -export const post21 = oc +export const post23 = oc .route({ inputStructure: 'detailed', method: 'POST', @@ -1610,13 +1661,13 @@ export const post21 = oc .output(zPostDatasetsByDatasetIdRetryResponse) export const retry = { - post: post21, + post: post23, } /** * Check if dataset is in use */ -export const get33 = oc +export const get34 = oc .route({ description: 'Check if dataset is in use', inputStructure: 'detailed', @@ -1629,7 +1680,7 @@ export const get33 = oc .output(zGetDatasetsByDatasetIdUseCheckResponse) export const useCheck2 = { - get: get33, + get: get34, } export const delete9 = oc @@ -1647,7 +1698,7 @@ export const delete9 = oc /** * Get dataset details */ -export const get34 = oc +export const get35 = oc .route({ description: 'Get dataset details', inputStructure: 'detailed', @@ -1676,7 +1727,7 @@ export const patch11 = oc export const byDatasetId = { delete: delete9, - get: get34, + get: get35, patch: patch11, apiKeys: apiKeys2, autoDisableLogs, @@ -1686,6 +1737,7 @@ export const byDatasetId = { externalHitTesting, hitTesting, indexingStatus: indexingStatus3, + knowledgeFsUpgrades, metadata: metadata4, notion: notion2, permissionPartUsers, @@ -1723,7 +1775,7 @@ export const byApiKeyId2 = { * * Get all API keys for a dataset */ -export const get35 = oc +export const get36 = oc .route({ description: 'Get all API keys for a dataset', inputStructure: 'detailed', @@ -1741,7 +1793,7 @@ export const get35 = oc * * Create a new API key for a dataset */ -export const post22 = oc +export const post24 = oc .route({ description: 'Create a new API key for a dataset', inputStructure: 'detailed', @@ -1756,8 +1808,8 @@ export const post22 = oc .output(zPostDatasetsByResourceIdApiKeysResponse) export const apiKeys3 = { - get: get35, - post: post22, + get: get36, + post: post24, byApiKeyId: byApiKeyId2, } @@ -1768,7 +1820,7 @@ export const byResourceId = { /** * Get list of datasets */ -export const get36 = oc +export const get37 = oc .route({ description: 'Get list of datasets', inputStructure: 'detailed', @@ -1783,7 +1835,7 @@ export const get36 = oc /** * Create a new dataset */ -export const post23 = oc +export const post25 = oc .route({ description: 'Create a new dataset', inputStructure: 'detailed', @@ -1797,8 +1849,8 @@ export const post23 = oc .output(zPostDatasetsResponse) export const datasets = { - get: get36, - post: post23, + get: get37, + post: post25, apiBaseInfo, apiKeys, batchImportStatus, diff --git a/packages/contracts/generated/api/console/datasets/types.gen.ts b/packages/contracts/generated/api/console/datasets/types.gen.ts index 74a1425d39a..164fef65d1d 100644 --- a/packages/contracts/generated/api/console/datasets/types.gen.ts +++ b/packages/contracts/generated/api/console/datasets/types.gen.ts @@ -513,6 +513,27 @@ export type HitTestingResponse = { records: Array } +export type KnowledgeFsUpgradeJobResponse = { + completed_at?: string | null + completed_documents: number + completed_sources: number + id: string + last_error_code?: string | null + last_error_message?: string | null + new_control_space_id?: string | null + old_dataset_id: string + snapshot_at: string + stage: KnowledgeFsUpgradeStage + status: KnowledgeFsUpgradeJobStatus + total_documents: number + total_sources: number +} + +export type KnowledgeFsUpgradeRetryResponse = { + id: string + status?: 'queued' +} + export type DatasetMetadataListResponse = { built_in_field_enabled: boolean doc_metadata: Array @@ -848,6 +869,17 @@ export type HitTestingRecord = { tsne_position: unknown | null } +export type KnowledgeFsUpgradeStage = + | 'completed' + | 'creating_sources' + | 'finalizing' + | 'migrating_access' + | 'submitting_documents' + | 'validating' + | 'waiting_for_space' + +export type KnowledgeFsUpgradeJobStatus = 'failed' | 'queued' | 'running' | 'succeeded' + export type DatasetMetadataListItemResponse = { count?: number id: string @@ -2302,6 +2334,56 @@ export type GetDatasetsByDatasetIdIndexingStatusResponses = { export type GetDatasetsByDatasetIdIndexingStatusResponse = GetDatasetsByDatasetIdIndexingStatusResponses[keyof GetDatasetsByDatasetIdIndexingStatusResponses] +export type PostDatasetsByDatasetIdKnowledgeFsUpgradesData = { + body?: never + path: { + dataset_id: string + } + query?: never + url: '/datasets/{dataset_id}/knowledge-fs-upgrades' +} + +export type PostDatasetsByDatasetIdKnowledgeFsUpgradesResponses = { + 202: KnowledgeFsUpgradeJobResponse +} + +export type PostDatasetsByDatasetIdKnowledgeFsUpgradesResponse = + PostDatasetsByDatasetIdKnowledgeFsUpgradesResponses[keyof PostDatasetsByDatasetIdKnowledgeFsUpgradesResponses] + +export type GetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdData = { + body?: never + path: { + dataset_id: string + job_id: string + } + query?: never + url: '/datasets/{dataset_id}/knowledge-fs-upgrades/{job_id}' +} + +export type GetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponses = { + 200: KnowledgeFsUpgradeJobResponse +} + +export type GetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponse = + GetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponses[keyof GetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponses] + +export type PostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdData = { + body?: never + path: { + dataset_id: string + job_id: string + } + query?: never + url: '/datasets/{dataset_id}/knowledge-fs-upgrades/{job_id}' +} + +export type PostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponses = { + 202: KnowledgeFsUpgradeRetryResponse +} + +export type PostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponse = + PostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponses[keyof PostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponses] + export type GetDatasetsByDatasetIdMetadataData = { body?: never path: { diff --git a/packages/contracts/generated/api/console/datasets/zod.gen.ts b/packages/contracts/generated/api/console/datasets/zod.gen.ts index 183a20a8480..78d31fcd888 100644 --- a/packages/contracts/generated/api/console/datasets/zod.gen.ts +++ b/packages/contracts/generated/api/console/datasets/zod.gen.ts @@ -244,6 +244,14 @@ export const zExternalHitTestingPayload = z.object({ query: z.string(), }) +/** + * KnowledgeFSUpgradeRetryResponse + */ +export const zKnowledgeFsUpgradeRetryResponse = z.object({ + id: z.string(), + status: z.literal('queued').optional().default('queued'), +}) + /** * MetadataArgs */ @@ -754,6 +762,43 @@ export const zHitTestingQuery = z.object({ content: z.string(), }) +/** + * KnowledgeFSUpgradeStage + */ +export const zKnowledgeFsUpgradeStage = z.enum([ + 'completed', + 'creating_sources', + 'finalizing', + 'migrating_access', + 'submitting_documents', + 'validating', + 'waiting_for_space', +]) + +/** + * KnowledgeFSUpgradeJobStatus + */ +export const zKnowledgeFsUpgradeJobStatus = z.enum(['failed', 'queued', 'running', 'succeeded']) + +/** + * KnowledgeFSUpgradeJobResponse + */ +export const zKnowledgeFsUpgradeJobResponse = z.object({ + completed_at: z.iso.datetime().nullish(), + completed_documents: z.int().gte(0), + completed_sources: z.int().gte(0), + id: z.string(), + last_error_code: z.string().nullish(), + last_error_message: z.string().nullish(), + new_control_space_id: z.string().nullish(), + old_dataset_id: z.string(), + snapshot_at: z.iso.datetime(), + stage: zKnowledgeFsUpgradeStage, + status: zKnowledgeFsUpgradeJobStatus, + total_documents: z.int().gte(0), + total_sources: z.int().gte(0), +}) + /** * DatasetMetadataListItemResponse */ @@ -2176,6 +2221,37 @@ export const zGetDatasetsByDatasetIdIndexingStatusPath = z.object({ */ export const zGetDatasetsByDatasetIdIndexingStatusResponse = zDocumentStatusListResponse +export const zPostDatasetsByDatasetIdKnowledgeFsUpgradesPath = z.object({ + dataset_id: z.uuid(), +}) + +/** + * KnowledgeFS Dataset upgrade accepted + */ +export const zPostDatasetsByDatasetIdKnowledgeFsUpgradesResponse = zKnowledgeFsUpgradeJobResponse + +export const zGetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdPath = z.object({ + dataset_id: z.uuid(), + job_id: z.string(), +}) + +/** + * KnowledgeFS Dataset upgrade status + */ +export const zGetDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponse = + zKnowledgeFsUpgradeJobResponse + +export const zPostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdPath = z.object({ + dataset_id: z.uuid(), + job_id: z.string(), +}) + +/** + * KnowledgeFS Dataset upgrade retry accepted + */ +export const zPostDatasetsByDatasetIdKnowledgeFsUpgradesByJobIdResponse = + zKnowledgeFsUpgradeRetryResponse + export const zGetDatasetsByDatasetIdMetadataPath = z.object({ dataset_id: z.uuid(), }) diff --git a/web/app/components/datasets/settings/form/components/knowledge-fs-upgrade.tsx b/web/app/components/datasets/settings/form/components/knowledge-fs-upgrade.tsx new file mode 100644 index 00000000000..7e82ad685ef --- /dev/null +++ b/web/app/components/datasets/settings/form/components/knowledge-fs-upgrade.tsx @@ -0,0 +1,155 @@ +'use client' + +import type { KnowledgeFsUpgradeJobResponse } from '@dify/contracts/api/console/datasets/types.gen' +import { Button } from '@langgenius/dify-ui/button' +import { useMutation, useQuery } from '@tanstack/react-query' +import { useState } from 'react' +import { useTranslation } from 'react-i18next' +import { consoleQuery } from '@/service/client' + +type Props = { + datasetId: string + disabled: boolean +} + +const POLL_INTERVAL = 2_000 + +const getErrorMessage = (error: unknown) => (error instanceof Error ? error.message : String(error)) + +const UpgradeProgress = ({ + datasetId, + initialJob, +}: { + datasetId: string + initialJob: KnowledgeFsUpgradeJobResponse +}) => { + const { t } = useTranslation() + const [retryError, setRetryError] = useState() + const jobContract = consoleQuery.datasets.byDatasetId.knowledgeFsUpgrades.byJobId + const jobInput = { + params: { + dataset_id: datasetId, + job_id: initialJob.id, + }, + } + const { + data: job, + error: statusError, + refetch, + } = useQuery({ + ...jobContract.get.queryOptions({ input: jobInput }), + initialData: initialJob, + refetchInterval: (query) => { + const status = query.state.data?.status + return status === 'queued' || status === 'running' ? POLL_INTERVAL : false + }, + }) + const retryMutation = useMutation({ + ...jobContract.post.mutationOptions(), + onMutate: () => setRetryError(undefined), + onSuccess: () => void refetch(), + onError: (error) => setRetryError(getErrorMessage(error)), + }) + + const completed = job.completed_documents + job.completed_sources + const total = job.total_documents + job.total_sources + const progress = total > 0 ? Math.round((completed / total) * 100) : 0 + const statusLabel = { + queued: t(($) => $['newKnowledge.documentStatus.queued'], { ns: 'dataset' }), + running: t(($) => $['newKnowledge.documentStatus.processing'], { ns: 'dataset' }), + succeeded: t(($) => $['newKnowledge.processingTaskState.succeeded'], { ns: 'dataset' }), + failed: t(($) => $['newKnowledge.documentStatus.failed'], { ns: 'dataset' }), + }[job.status] + const errorMessage = + retryError || + job.last_error_message || + job.last_error_code || + (statusError && getErrorMessage(statusError)) + + return ( +
+
+
+ {statusLabel} + {job.stage} +
+
{progress}%
+
+
+
+
+
+ + {t(($) => $['newKnowledge.documents'], { ns: 'dataset' })}: {job.completed_documents}/ + {job.total_documents} + + + {t(($) => $['newKnowledge.sources'], { ns: 'dataset' })}: {job.completed_sources}/ + {job.total_sources} + +
+ {errorMessage && ( +
+ {t(($) => $.error, { ns: 'common' })}: {errorMessage} +
+ )} + {job.status === 'failed' && ( +
+ +
+ )} +
+ ) +} + +const KnowledgeFSUpgrade = ({ datasetId, disabled }: Props) => { + const { t } = useTranslation() + const [job, setJob] = useState() + const [startError, setStartError] = useState() + const upgradeContract = consoleQuery.datasets.byDatasetId.knowledgeFsUpgrades + const upgradeInput = { params: { dataset_id: datasetId } } + const startMutation = useMutation({ + ...upgradeContract.post.mutationOptions(), + onMutate: () => setStartError(undefined), + onSuccess: setJob, + onError: (error) => setStartError(getErrorMessage(error)), + }) + + if (job) return + + return ( +
+ + {startError && ( +
+ {t(($) => $.error, { ns: 'common' })}: {startError} +
+ )} +
+ ) +} + +export default KnowledgeFSUpgrade diff --git a/web/app/components/datasets/settings/form/index.tsx b/web/app/components/datasets/settings/form/index.tsx index 82410d540d2..09a3c949466 100644 --- a/web/app/components/datasets/settings/form/index.tsx +++ b/web/app/components/datasets/settings/form/index.tsx @@ -5,6 +5,7 @@ import Divider from '@/app/components/base/divider' import BasicInfoSection from './components/basic-info-section' import ExternalKnowledgeSection from './components/external-knowledge-section' import IndexingSection from './components/indexing-section' +import KnowledgeFSUpgrade from './components/knowledge-fs-upgrade' import { useFormState } from './hooks/use-form-state' const Form = () => { @@ -119,6 +120,12 @@ const Form = () => { + {!isExternalProvider && currentDataset && ( + + )} + + + {/* Save Button */}