Aether/_deprecated_py_src/services/provider/delete_task.py

"""Provider 异步删除任务。

接口提交后立即返回 task_id，后台分阶段删除 provider 及其子资源。
任务状态存储在 Redis 中，支持多 worker 进程共享。
"""

from __future__ import annotations

import asyncio
import json
import time
import uuid
from collections.abc import Callable, Sequence
from concurrent.futures import Future
from typing import Any

import redis.asyncio as aioredis
from sqlalchemy import text

from src.clients.redis_client import get_redis_client
from src.core.logger import logger
from src.database import create_session
from src.models.database import (
    ApiKey,
    Model,
    Provider,
    ProviderAPIKey,
    ProviderEndpoint,
    RequestCandidate,
    Usage,
    User,
    UserPreference,
    VideoTask,
)
from src.models.database_extensions import ApiKeyProviderMapping, ProviderUsageTracking
from src.services.cache.model_cache import ModelCacheService
from src.services.cache.model_list_cache import invalidate_models_list_cache
from src.services.cache.provider_cache import ProviderCacheService
from src.services.provider.delete_cleanup import prune_allowed_provider_refs
from src.services.provider_keys.key_side_effects import cleanup_key_references

STATUS_PENDING = "pending"
STATUS_RUNNING = "running"
STATUS_COMPLETED = "completed"
STATUS_FAILED = "failed"

_TASK_RETAIN_SECONDS = 600
_KEY_BATCH_SIZE = 50
_ENDPOINT_BATCH_SIZE = 200
_BATCH_STATEMENT_TIMEOUT_S = 30
_BATCH_LOCK_TIMEOUT_S = 5
_TASK_TIMEOUT_S = 1800
_REDIS_KEY_PREFIX = "provider_delete_task"

_running_tasks: set[asyncio.Task[None]] = set()


def _task_key(task_id: str) -> str:
    return f"{_REDIS_KEY_PREFIX}:{task_id}"


def _provider_lock_key(provider_id: str) -> str:
    return f"{_REDIS_KEY_PREFIX}:provider:{provider_id}"


class ProviderDeleteTaskInfo:
    __slots__ = (
        "task_id",
        "provider_id",
        "status",
        "stage",
        "total_keys",
        "deleted_keys",
        "total_endpoints",
        "deleted_endpoints",
        "message",
    )

    def __init__(
        self,
        task_id: str,
        provider_id: str,
        status: str = STATUS_PENDING,
        stage: str = "queued",
        total_keys: int = 0,
        deleted_keys: int = 0,
        total_endpoints: int = 0,
        deleted_endpoints: int = 0,
        message: str = "",
    ) -> None:
        self.task_id = task_id
        self.provider_id = provider_id
        self.status = status
        self.stage = stage
        self.total_keys = total_keys
        self.deleted_keys = deleted_keys
        self.total_endpoints = total_endpoints
        self.deleted_endpoints = deleted_endpoints
        self.message = message

    def to_dict(self) -> dict[str, Any]:
        return {
            "task_id": self.task_id,
            "provider_id": self.provider_id,
            "status": self.status,
            "stage": self.stage,
            "total_keys": self.total_keys,
            "deleted_keys": self.deleted_keys,
            "total_endpoints": self.total_endpoints,
            "deleted_endpoints": self.deleted_endpoints,
            "message": self.message,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ProviderDeleteTaskInfo":
        return cls(
            task_id=str(data["task_id"]),
            provider_id=str(data["provider_id"]),
            status=str(data.get("status", STATUS_PENDING)),
            stage=str(data.get("stage", "queued")),
            total_keys=int(data.get("total_keys", 0)),
            deleted_keys=int(data.get("deleted_keys", 0)),
            total_endpoints=int(data.get("total_endpoints", 0)),
            deleted_endpoints=int(data.get("deleted_endpoints", 0)),
            message=str(data.get("message", "")),
        )


async def _save_task(
    task: ProviderDeleteTaskInfo,
    ttl: int = _TASK_RETAIN_SECONDS,
    r: aioredis.Redis | None = None,
) -> None:
    if r is None:
        r = await get_redis_client(require_redis=False)
    if not r:
        return
    try:
        await r.setex(_task_key(task.task_id), ttl, json.dumps(task.to_dict()))
        await r.setex(_provider_lock_key(task.provider_id), ttl, task.task_id)
    except Exception as exc:
        logger.warning("Failed to save provider delete task: {}", exc)


async def _load_task(
    task_id: str,
    r: aioredis.Redis | None = None,
) -> ProviderDeleteTaskInfo | None:
    if r is None:
        r = await get_redis_client(require_redis=False)
    if not r:
        return None
    try:
        data = await r.get(_task_key(task_id))
        if data is None:
            return None
        return ProviderDeleteTaskInfo.from_dict(json.loads(data))
    except Exception as exc:
        logger.warning("Failed to load provider delete task: {}", exc)
        return None


async def _update_task_field(
    task_id: str,
    r: aioredis.Redis | None = None,
    **fields: object,
) -> None:
    if r is None:
        r = await get_redis_client(require_redis=False)
    if not r:
        return
    task = await _load_task(task_id, r=r)
    if task is None:
        return
    for key, value in fields.items():
        setattr(task, key, value)
    ttl = (
        _TASK_RETAIN_SECONDS
        if task.status in (STATUS_COMPLETED, STATUS_FAILED)
        else _TASK_RETAIN_SECONDS * 2
    )
    await _save_task(task, ttl=ttl, r=r)


async def submit_provider_delete(provider_id: str) -> str:
    r = await get_redis_client(require_redis=False)
    if not r:
        raise RuntimeError("Redis is required for provider delete tasks but is not available")

    existing_task_id = await r.get(_provider_lock_key(provider_id))
    if isinstance(existing_task_id, bytes):
        existing_task_id = existing_task_id.decode()
    if existing_task_id:
        existing_task = await _load_task(str(existing_task_id), r=r)
        if existing_task and existing_task.status in (STATUS_PENDING, STATUS_RUNNING):
            return existing_task.task_id

    task_id = uuid.uuid4().hex[:16]
    task = ProviderDeleteTaskInfo(
        task_id=task_id,
        provider_id=provider_id,
        message="delete task submitted",
    )
    await _save_task(task, ttl=_TASK_RETAIN_SECONDS * 2, r=r)

    def _on_task_done(task: asyncio.Task[None]) -> None:
        _running_tasks.discard(task)
        if not task.cancelled() and task.exception():
            logger.error("[PROVIDER_DELETE_TASK] unhandled error: {}", task.exception())

    bg = asyncio.create_task(
        _run_provider_delete(task_id, provider_id),
        name=f"provider-delete-{task_id}",
    )
    _running_tasks.add(bg)
    bg.add_done_callback(_on_task_done)
    return task_id


async def get_provider_delete_task(task_id: str) -> ProviderDeleteTaskInfo | None:
    return await _load_task(task_id)


def _iter_batches(items: Sequence[str], batch_size: int) -> list[list[str]]:
    if not items:
        return []
    if batch_size <= 0:
        return [list(items)]
    return [list(items[i : i + batch_size]) for i in range(0, len(items), batch_size)]


def _apply_statement_timeouts(db: Any) -> None:
    db.execute(text(f"SET LOCAL statement_timeout = '{_BATCH_STATEMENT_TIMEOUT_S * 1000}'"))
    db.execute(text(f"SET LOCAL lock_timeout = '{_BATCH_LOCK_TIMEOUT_S * 1000}'"))


def _collect_ids(db: Any, provider_id: str) -> tuple[list[str], list[str]]:
    endpoint_ids = [
        endpoint_id
        for endpoint_id, in db.query(ProviderEndpoint.id)
        .filter(ProviderEndpoint.provider_id == provider_id)
        .all()
    ]
    key_ids = [
        key_id
        for key_id, in db.query(ProviderAPIKey.id)
        .filter(ProviderAPIKey.provider_id == provider_id)
        .all()
    ]
    return endpoint_ids, key_ids


def _sync_delete_provider(
    provider_id: str,
    progress_callback: Callable[[dict[str, object]], None] | None = None,
) -> dict[str, Any]:
    db = create_session()
    try:
        task_start = time.monotonic()
        provider = db.query(Provider).filter(Provider.id == provider_id).first()
        if not provider:
            raise RuntimeError("provider not found")

        _apply_statement_timeouts(db)
        endpoint_ids, key_ids = _collect_ids(db, provider_id)
        if progress_callback is not None:
            progress_callback(
                {
                    "stage": "preparing",
                    "total_keys": len(key_ids),
                    "total_endpoints": len(endpoint_ids),
                    "message": f"preparing delete for {len(key_ids)} keys and {len(endpoint_ids)} endpoints",
                }
            )

        if getattr(provider, "is_active", False):
            provider.is_active = False
            db.commit()
            if progress_callback is not None:
                progress_callback(
                    {
                        "stage": "disabling",
                        "message": "provider disabled; starting cleanup",
                    }
                )

        _apply_statement_timeouts(db)
        updated_users = prune_allowed_provider_refs(
            db.query(User).filter(User.allowed_providers.isnot(None)).all(),
            provider_id,
        )
        updated_api_keys = prune_allowed_provider_refs(
            db.query(ApiKey).filter(ApiKey.allowed_providers.isnot(None)).all(),
            provider_id,
        )
        db.commit()
        if progress_callback is not None:
            progress_callback(
                {
                    "stage": "cleaning_restrictions",
                    "message": f"cleaned access restrictions (users={updated_users}, api_keys={updated_api_keys})",
                }
            )

        _apply_statement_timeouts(db)
        db.query(UserPreference).filter(UserPreference.default_provider_id == provider_id).update(
            {UserPreference.default_provider_id: None},
            synchronize_session=False,
        )
        db.query(Usage).filter(Usage.provider_id == provider_id).update(
            {Usage.provider_id: None},
            synchronize_session=False,
        )
        db.query(VideoTask).filter(VideoTask.provider_id == provider_id).update(
            {VideoTask.provider_id: None},
            synchronize_session=False,
        )
        db.query(RequestCandidate).filter(RequestCandidate.provider_id == provider_id).delete(
            synchronize_session=False,
        )
        db.commit()
        if progress_callback is not None:
            progress_callback(
                {
                    "stage": "cleaning_provider_refs",
                    "message": "cleaned provider-wide history references",
                }
            )

        deleted_keys = 0
        key_batches = _iter_batches(key_ids, _KEY_BATCH_SIZE)
        for index, batch in enumerate(key_batches, start=1):
            if time.monotonic() - task_start > _TASK_TIMEOUT_S:
                raise RuntimeError(f"task timeout after {_TASK_TIMEOUT_S}s")
            _apply_statement_timeouts(db)
            cleanup_key_references(db, batch)
            deleted_batch = int(
                db.query(ProviderAPIKey)
                .filter(ProviderAPIKey.provider_id == provider_id, ProviderAPIKey.id.in_(batch))
                .delete(synchronize_session=False)
                or 0
            )
            db.commit()
            deleted_keys += deleted_batch
            if progress_callback is not None:
                progress_callback(
                    {
                        "stage": "deleting_keys",
                        "deleted_keys": deleted_keys,
                        "message": f"deleted key batch {index}/{max(len(key_batches), 1)}",
                    }
                )

        deleted_endpoints = 0
        endpoint_batches = _iter_batches(endpoint_ids, _ENDPOINT_BATCH_SIZE)
        for index, batch in enumerate(endpoint_batches, start=1):
            if time.monotonic() - task_start > _TASK_TIMEOUT_S:
                raise RuntimeError(f"task timeout after {_TASK_TIMEOUT_S}s")
            _apply_statement_timeouts(db)
            db.query(Usage).filter(Usage.provider_endpoint_id.in_(batch)).update(
                {Usage.provider_endpoint_id: None},
                synchronize_session=False,
            )
            db.query(VideoTask).filter(VideoTask.endpoint_id.in_(batch)).update(
                {VideoTask.endpoint_id: None},
                synchronize_session=False,
            )
            db.query(RequestCandidate).filter(RequestCandidate.endpoint_id.in_(batch)).delete(
                synchronize_session=False,
            )
            deleted_batch = int(
                db.query(ProviderEndpoint)
                .filter(ProviderEndpoint.provider_id == provider_id, ProviderEndpoint.id.in_(batch))
                .delete(synchronize_session=False)
                or 0
            )
            db.commit()
            deleted_endpoints += deleted_batch
            if progress_callback is not None:
                progress_callback(
                    {
                        "stage": "deleting_endpoints",
                        "deleted_endpoints": deleted_endpoints,
                        "message": f"deleted endpoint batch {index}/{max(len(endpoint_batches), 1)}",
                    }
                )

        _apply_statement_timeouts(db)
        deleted_models = int(
            db.query(Model)
            .filter(Model.provider_id == provider_id)
            .delete(synchronize_session=False)
            or 0
        )
        deleted_mappings = int(
            db.query(ApiKeyProviderMapping)
            .filter(ApiKeyProviderMapping.provider_id == provider_id)
            .delete(synchronize_session=False)
            or 0
        )
        deleted_usage_tracking = int(
            db.query(ProviderUsageTracking)
            .filter(ProviderUsageTracking.provider_id == provider_id)
            .delete(synchronize_session=False)
            or 0
        )
        deleted_provider = int(
            db.query(Provider).filter(Provider.id == provider_id).delete(synchronize_session=False)
            or 0
        )
        db.commit()

        return {
            "provider_id": provider_id,
            "total_keys": len(key_ids),
            "deleted_keys": deleted_keys,
            "total_endpoints": len(endpoint_ids),
            "deleted_endpoints": deleted_endpoints,
            "deleted_models": deleted_models,
            "deleted_mappings": deleted_mappings,
            "deleted_usage_tracking": deleted_usage_tracking,
            "deleted_provider": deleted_provider,
            "elapsed_seconds": time.monotonic() - task_start,
        }
    except Exception:
        try:
            db.rollback()
        except Exception:
            pass
        raise
    finally:
        try:
            db.close()
        except Exception:
            pass


async def _run_provider_delete(task_id: str, provider_id: str) -> None:
    r = await get_redis_client(require_redis=False)
    await _update_task_field(
        task_id,
        r=r,
        status=STATUS_RUNNING,
        stage="queued",
        message="delete task started",
    )

    loop = asyncio.get_running_loop()
    progress_futures: list[Future[object]] = []

    def on_progress(fields: dict[str, object]) -> None:
        try:
            future = asyncio.run_coroutine_threadsafe(
                _update_task_field(task_id, r=r, **fields),
                loop,
            )
            progress_futures.append(future)
        except RuntimeError:
            pass

    async def _drain_progress() -> None:
        if progress_futures:
            await asyncio.gather(
                *(asyncio.wrap_future(f) for f in progress_futures),
                return_exceptions=True,
            )

    try:
        summary = await asyncio.wait_for(
            asyncio.to_thread(_sync_delete_provider, provider_id, on_progress),
            timeout=_TASK_TIMEOUT_S + 60,
        )

        await _drain_progress()

        try:
            await invalidate_models_list_cache()
            await ModelCacheService.invalidate_all_resolve_cache()
            await ProviderCacheService.invalidate_provider_cache(provider_id)
        except Exception as exc:
            logger.error("provider delete cache invalidation failed: {}", exc)

        await _update_task_field(
            task_id,
            r=r,
            status=STATUS_COMPLETED,
            stage="completed",
            total_keys=int(summary.get("total_keys", 0)),
            deleted_keys=int(summary.get("deleted_keys", 0)),
            total_endpoints=int(summary.get("total_endpoints", 0)),
            deleted_endpoints=int(summary.get("deleted_endpoints", 0)),
            message=(
                f"provider deleted: keys={summary.get('deleted_keys', 0)}, "
                f"endpoints={summary.get('deleted_endpoints', 0)}"
            ),
        )
        logger.info(
            "[PROVIDER_DELETE_TASK] completed task={} provider={} keys={}/{} endpoints={}/{} elapsed={:.1f}s",
            task_id,
            provider_id[:8],
            summary.get("deleted_keys", 0),
            summary.get("total_keys", 0),
            summary.get("deleted_endpoints", 0),
            summary.get("total_endpoints", 0),
            float(summary.get("elapsed_seconds", 0.0) or 0.0),
        )
    except asyncio.TimeoutError:
        await _drain_progress()
        msg = f"task timeout after {_TASK_TIMEOUT_S + 60}s"
        await _update_task_field(task_id, r=r, status=STATUS_FAILED, stage="failed", message=msg)
        logger.error("[PROVIDER_DELETE_TASK] {} task={} provider={}", msg, task_id, provider_id[:8])
    except asyncio.CancelledError:
        await _drain_progress()
        await _update_task_field(
            task_id,
            r=r,
            status=STATUS_FAILED,
            stage="failed",
            message="task cancelled (shutdown)",
        )
        logger.warning(
            "[PROVIDER_DELETE_TASK] cancelled task={} provider={}",
            task_id,
            provider_id[:8],
        )
    except Exception as exc:
        await _drain_progress()
        await _update_task_field(
            task_id,
            r=r,
            status=STATUS_FAILED,
            stage="failed",
            message=str(exc),
        )
        logger.error(
            "[PROVIDER_DELETE_TASK] failed task={} provider={} error={}",
            task_id,
            provider_id[:8],
            exc,
        )