Aether/_deprecated_py_src/services/provider/delete_cleanup.py

from __future__ import annotations

from collections.abc import Iterable, Sequence
from typing import Any

from sqlalchemy.orm import Session

from src.core.logger import logger
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.provider_keys.key_side_effects import cleanup_key_references

_BATCH_SIZE = 2000


def _empty_cleanup_stats() -> dict[str, int]:
    return {
        "users": 0,
        "api_keys": 0,
        "user_preferences": 0,
        "usage_provider": 0,
        "usage_endpoint": 0,
        "video_tasks_provider": 0,
        "video_tasks_endpoint": 0,
        "request_candidates_provider": 0,
        "request_candidates_endpoint": 0,
    }


def _empty_delete_stats() -> dict[str, int]:
    return {
        "api_key_mappings": 0,
        "usage_tracking": 0,
        "models": 0,
        "api_keys": 0,
        "endpoints": 0,
        "providers": 0,
    }


def _iter_batches(items: Sequence[str], batch_size: int = _BATCH_SIZE) -> 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 _collect_provider_child_ids(db: Session, 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 prune_allowed_provider_list(
    allowed_providers: Any, provider_id: str
) -> tuple[list[str] | None | Any, bool]:
    """从访问限制列表中移除指定 Provider ID。"""
    if not isinstance(allowed_providers, list):
        return allowed_providers, False
    if provider_id not in allowed_providers:
        return allowed_providers, False

    next_allowed = [value for value in allowed_providers if value != provider_id]
    return next_allowed, True


def prune_allowed_provider_refs(records: Iterable[Any], provider_id: str) -> int:
    """批量移除记录中的 allowed_providers 引用。"""
    updated = 0
    for record in records:
        next_allowed, changed = prune_allowed_provider_list(
            getattr(record, "allowed_providers", None),
            provider_id,
        )
        if not changed:
            continue
        record.allowed_providers = next_allowed
        updated += 1
    return updated


def cleanup_deleted_provider_references(
    db: Session,
    provider_id: str,
    *,
    endpoint_ids: Sequence[str] | None = None,
    key_ids: Sequence[str] | None = None,
) -> dict[str, int]:
    """清理 Provider 删除时的大扇出引用，避免依赖数据库级联导致慢删。"""
    if not provider_id:
        return _empty_cleanup_stats()

    if endpoint_ids is None or key_ids is None:
        resolved_endpoint_ids, resolved_key_ids = _collect_provider_child_ids(db, provider_id)
        endpoint_ids = resolved_endpoint_ids if endpoint_ids is None else list(endpoint_ids)
        key_ids = resolved_key_ids if key_ids is None else list(key_ids)
    else:
        endpoint_ids = list(endpoint_ids)
        key_ids = list(key_ids)

    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,
    )

    cleared_preferences = int(
        db.query(UserPreference)
        .filter(UserPreference.default_provider_id == provider_id)
        .update({UserPreference.default_provider_id: None}, synchronize_session=False)
        or 0
    )
    cleared_usage_providers = int(
        db.query(Usage)
        .filter(Usage.provider_id == provider_id)
        .update({Usage.provider_id: None}, synchronize_session=False)
        or 0
    )
    cleared_video_task_providers = int(
        db.query(VideoTask)
        .filter(VideoTask.provider_id == provider_id)
        .update({VideoTask.provider_id: None}, synchronize_session=False)
        or 0
    )

    if key_ids:
        cleanup_key_references(db, list(key_ids))

    cleared_usage_endpoints = 0
    cleared_video_task_endpoints = 0
    deleted_request_candidates_endpoints = 0
    for batch in _iter_batches(endpoint_ids):
        cleared_usage_endpoints += int(
            db.query(Usage)
            .filter(Usage.provider_endpoint_id.in_(batch))
            .update({Usage.provider_endpoint_id: None}, synchronize_session=False)
            or 0
        )
        cleared_video_task_endpoints += int(
            db.query(VideoTask)
            .filter(VideoTask.endpoint_id.in_(batch))
            .update({VideoTask.endpoint_id: None}, synchronize_session=False)
            or 0
        )
        deleted_request_candidates_endpoints += int(
            db.query(RequestCandidate)
            .filter(RequestCandidate.endpoint_id.in_(batch))
            .delete(synchronize_session=False)
            or 0
        )

    deleted_request_candidates_provider = int(
        db.query(RequestCandidate)
        .filter(RequestCandidate.provider_id == provider_id)
        .delete(synchronize_session=False)
        or 0
    )

    stats = {
        "users": updated_users,
        "api_keys": updated_api_keys,
        "user_preferences": cleared_preferences,
        "usage_provider": cleared_usage_providers,
        "usage_endpoint": cleared_usage_endpoints,
        "video_tasks_provider": cleared_video_task_providers,
        "video_tasks_endpoint": cleared_video_task_endpoints,
        "request_candidates_provider": deleted_request_candidates_provider,
        "request_candidates_endpoint": deleted_request_candidates_endpoints,
    }

    if any(stats.values()) or key_ids:
        logger.info(
            "Provider 删除引用清理: provider_id={}, key_refs={}, users={}, api_keys={}, "
            "user_preferences={}, usage_provider={}, usage_endpoint={}, "
            "video_tasks_provider={}, video_tasks_endpoint={}, "
            "request_candidates_provider={}, request_candidates_endpoint={}",
            provider_id,
            len(key_ids),
            stats["users"],
            stats["api_keys"],
            stats["user_preferences"],
            stats["usage_provider"],
            stats["usage_endpoint"],
            stats["video_tasks_provider"],
            stats["video_tasks_endpoint"],
            stats["request_candidates_provider"],
            stats["request_candidates_endpoint"],
        )

    return stats


def delete_provider_tree(db: Session, provider_id: str) -> dict[str, Any]:
    """分阶段删除 Provider 及其子资源，降低 ORM/FK 级联导致的超时风险。"""
    if not provider_id:
        return {
            "cleanup": _empty_cleanup_stats(),
            "deleted": _empty_delete_stats(),
            "key_count": 0,
            "endpoint_count": 0,
        }

    endpoint_ids, key_ids = _collect_provider_child_ids(db, provider_id)
    cleanup_stats = cleanup_deleted_provider_references(
        db,
        provider_id,
        endpoint_ids=endpoint_ids,
        key_ids=key_ids,
    )

    deleted_stats = {
        "api_key_mappings": int(
            db.query(ApiKeyProviderMapping)
            .filter(ApiKeyProviderMapping.provider_id == provider_id)
            .delete(synchronize_session=False)
            or 0
        ),
        "usage_tracking": int(
            db.query(ProviderUsageTracking)
            .filter(ProviderUsageTracking.provider_id == provider_id)
            .delete(synchronize_session=False)
            or 0
        ),
        "models": int(
            db.query(Model)
            .filter(Model.provider_id == provider_id)
            .delete(synchronize_session=False)
            or 0
        ),
        "api_keys": int(
            db.query(ProviderAPIKey)
            .filter(ProviderAPIKey.provider_id == provider_id)
            .delete(synchronize_session=False)
            or 0
        ),
        "endpoints": int(
            db.query(ProviderEndpoint)
            .filter(ProviderEndpoint.provider_id == provider_id)
            .delete(synchronize_session=False)
            or 0
        ),
        "providers": int(
            db.query(Provider).filter(Provider.id == provider_id).delete(synchronize_session=False)
            or 0
        ),
    }

    logger.info(
        "Provider 分阶段删除: provider_id={}, key_count={}, endpoint_count={}, "
        "deleted_mappings={}, deleted_usage_tracking={}, deleted_models={}, "
        "deleted_api_keys={}, deleted_endpoints={}, deleted_providers={}",
        provider_id,
        len(key_ids),
        len(endpoint_ids),
        deleted_stats["api_key_mappings"],
        deleted_stats["usage_tracking"],
        deleted_stats["models"],
        deleted_stats["api_keys"],
        deleted_stats["endpoints"],
        deleted_stats["providers"],
    )

    return {
        "cleanup": cleanup_stats,
        "deleted": deleted_stats,
        "key_count": len(key_ids),
        "endpoint_count": len(endpoint_ids),
    }