i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-10 13:10:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(retention): 删除用户/Key 时保留历史记录，外键改 SET NULL 并添加名称快照

Browse Source 
- Usage/RequestCandidate/VideoTask/Stats 等表的 user_id/api_key_id 外键从
  CASCADE 改为 SET NULL，删除用户或 Key 后历史记录不再丢失
- 各表添加 username/api_key_name 快照字段，删除后仍可追溯归属
- 新增 bulk_cleanup 模块，分批置空大表外键避免长事务锁
- 删除用户/Key 流程集成预清理步骤，先置空再删除
- 精简 candidate_builder 冗余 debug 日志
- 修复 proxy_nodes 启动日志 format 占位符错误（{} -> %s）
- 前端批量操作请求增加 5 分钟超时配置
This commit is contained in:
fawney19
2026-03-08 14:31:15 +08:00
parent 25c33846be
commit bd3f73c2fc
26 changed files with 799 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
alembic/versions/20260308_0348_45b118150a78_usage_stats_retention_set_null_on_.py
+166
View File
@@ -0,0 +1,166 @@
"""usage stats retention: SET NULL on delete and add name snapshots
Revision ID: 45b118150a78
Revises: 2d932114930d
Create Date: 2026-03-08 03:48:49.622091+00:00
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = '45b118150a78'
down_revision = '2d932114930d'
branch_labels = None
depends_on = None
# ---------------------------------------------------------------------------
# Idempotent helpers
# ---------------------------------------------------------------------------
def _column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.columns "
"WHERE table_name = :table AND column_name = :col"
),
{"table": table_name, "col": column_name},
)
return result.scalar() is not None
def _constraint_exists(table_name: str, constraint_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.table_constraints "
"WHERE table_name = :table AND constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
return result.scalar() is not None
def _fk_ondelete(table_name: str, constraint_name: str) -> str | None:
"""Return the ON DELETE action for a foreign key, or None if it doesn't exist."""
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT rc.delete_rule "
"FROM information_schema.referential_constraints rc "
"JOIN information_schema.table_constraints tc "
" ON rc.constraint_name = tc.constraint_name "
"WHERE tc.table_name = :table AND tc.constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
row = result.first()
return row[0] if row else None
def _replace_fk_if_needed(
constraint_name: str,
table_name: str,
ref_table: str,
local_cols: list[str],
remote_cols: list[str],
desired_ondelete: str,
) -> None:
"""Drop and recreate a FK only if the current ON DELETE rule differs."""
current = _fk_ondelete(table_name, constraint_name)
if current and current.upper() == desired_ondelete.upper():
return # already correct
if current:
op.drop_constraint(constraint_name, table_name, type_='foreignkey')
op.create_foreign_key(
constraint_name, table_name, ref_table,
local_cols, remote_cols, ondelete=desired_ondelete,
)
def upgrade() -> None:
# --- Usage: add name snapshot columns ---
if not _column_exists('usage', 'username'):
op.add_column('usage', sa.Column('username', sa.String(100), nullable=True,
comment='用户名快照'))
if not _column_exists('usage', 'api_key_name'):
op.add_column('usage', sa.Column('api_key_name', sa.String(200), nullable=True,
comment='API Key 名称快照'))
# --- StatsUserDaily: CASCADE -> SET NULL, add username snapshot ---
_replace_fk_if_needed(
'stats_user_daily_user_id_fkey', 'stats_user_daily',
'users', ['user_id'], ['id'], 'SET NULL',
)
op.alter_column('stats_user_daily', 'user_id', existing_type=sa.String(36), nullable=True)
if not _column_exists('stats_user_daily', 'username'):
op.add_column('stats_user_daily', sa.Column('username', sa.String(100), nullable=True,
comment='用户名快照（删除用户后仍可追溯）'))
# --- StatsDailyApiKey: CASCADE -> SET NULL, add api_key_name snapshot ---
_replace_fk_if_needed(
'stats_daily_api_key_api_key_id_fkey', 'stats_daily_api_key',
'api_keys', ['api_key_id'], ['id'], 'SET NULL',
)
op.alter_column('stats_daily_api_key', 'api_key_id', existing_type=sa.String(36),
nullable=True)
if not _column_exists('stats_daily_api_key', 'api_key_name'):
op.add_column('stats_daily_api_key', sa.Column('api_key_name', sa.String(200),
nullable=True,
comment='API Key 名称快照（删除 Key 后仍可追溯）'))
# --- Backfill: populate snapshots from existing FK joins ---
# (WHERE ... IS NULL makes these inherently idempotent)
op.execute("""
UPDATE usage u
SET username = usr.username
FROM users usr
WHERE u.user_id = usr.id AND u.username IS NULL
""")
op.execute("""
UPDATE usage u
SET api_key_name = ak.name
FROM api_keys ak
WHERE u.api_key_id = ak.id AND u.api_key_name IS NULL
""")
op.execute("""
UPDATE stats_user_daily s
SET username = usr.username
FROM users usr
WHERE s.user_id = usr.id AND s.username IS NULL
""")
op.execute("""
UPDATE stats_daily_api_key s
SET api_key_name = ak.name
FROM api_keys ak
WHERE s.api_key_id = ak.id AND s.api_key_name IS NULL
""")
def downgrade() -> None:
# --- Remove snapshot columns ---
if _column_exists('stats_daily_api_key', 'api_key_name'):
op.drop_column('stats_daily_api_key', 'api_key_name')
if _column_exists('stats_user_daily', 'username'):
op.drop_column('stats_user_daily', 'username')
if _column_exists('usage', 'api_key_name'):
op.drop_column('usage', 'api_key_name')
if _column_exists('usage', 'username'):
op.drop_column('usage', 'username')
# --- StatsDailyApiKey: SET NULL -> CASCADE ---
_replace_fk_if_needed(
'stats_daily_api_key_api_key_id_fkey', 'stats_daily_api_key',
'api_keys', ['api_key_id'], ['id'], 'CASCADE',
)
op.alter_column('stats_daily_api_key', 'api_key_id', existing_type=sa.String(36),
nullable=False)
# --- StatsUserDaily: SET NULL -> CASCADE ---
_replace_fk_if_needed(
'stats_user_daily_user_id_fkey', 'stats_user_daily',
'users', ['user_id'], ['id'], 'CASCADE',
)
op.alter_column('stats_user_daily', 'user_id', existing_type=sa.String(36), nullable=False)
alembic/versions/20260308_1215_13a4c8f6d9e0_request_candidate_video_task_retention.py
+208
View File
@@ -0,0 +1,208 @@
"""request_candidates/video_tasks retention: SET NULL and add snapshots
Revision ID: 13a4c8f6d9e0
Revises: 45b118150a78
Create Date: 2026-03-08 12:15:00.000000+00:00
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = "13a4c8f6d9e0"
down_revision = "45b118150a78"
branch_labels = None
depends_on = None
# ---------------------------------------------------------------------------
# Idempotent helpers
# ---------------------------------------------------------------------------
def _column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.columns "
"WHERE table_name = :table AND column_name = :col"
),
{"table": table_name, "col": column_name},
)
return result.scalar() is not None
def _constraint_exists(table_name: str, constraint_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.table_constraints "
"WHERE table_name = :table AND constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
return result.scalar() is not None
def _fk_ondelete(table_name: str, constraint_name: str) -> str | None:
"""Return the ON DELETE action for a foreign key, or None if it doesn't exist."""
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT rc.delete_rule "
"FROM information_schema.referential_constraints rc "
"JOIN information_schema.table_constraints tc "
" ON rc.constraint_name = tc.constraint_name "
"WHERE tc.table_name = :table AND tc.constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
row = result.first()
return row[0] if row else None
def _replace_fk_if_needed(
constraint_name: str,
table_name: str,
ref_table: str,
local_cols: list[str],
remote_cols: list[str],
desired_ondelete: str,
) -> None:
"""Drop and recreate a FK only if the current ON DELETE rule differs."""
current = _fk_ondelete(table_name, constraint_name)
if current and current.upper() == desired_ondelete.upper():
return # already correct
if current:
op.drop_constraint(constraint_name, table_name, type_="foreignkey")
op.create_foreign_key(
constraint_name, table_name, ref_table,
local_cols, remote_cols, ondelete=desired_ondelete,
)
def upgrade() -> None:
# --- request_candidates: add snapshot columns ---
if not _column_exists("request_candidates", "username"):
op.add_column(
"request_candidates",
sa.Column("username", sa.String(length=100), nullable=True, comment="用户名快照"),
)
if not _column_exists("request_candidates", "api_key_name"):
op.add_column(
"request_candidates",
sa.Column(
"api_key_name",
sa.String(length=200),
nullable=True,
comment="API Key 名称快照",
),
)
# --- request_candidates: CASCADE -> SET NULL ---
_replace_fk_if_needed(
"request_candidates_user_id_fkey", "request_candidates",
"users", ["user_id"], ["id"], "SET NULL",
)
_replace_fk_if_needed(
"request_candidates_api_key_id_fkey", "request_candidates",
"api_keys", ["api_key_id"], ["id"], "SET NULL",
)
# --- video_tasks: add snapshot columns ---
if not _column_exists("video_tasks", "username"):
op.add_column(
"video_tasks",
sa.Column("username", sa.String(length=100), nullable=True, comment="用户名快照"),
)
if not _column_exists("video_tasks", "api_key_name"):
op.add_column(
"video_tasks",
sa.Column(
"api_key_name",
sa.String(length=200),
nullable=True,
comment="API Key 名称快照",
),
)
# --- video_tasks: CASCADE -> SET NULL, user_id nullable ---
op.alter_column("video_tasks", "user_id", existing_type=sa.String(length=36), nullable=True)
_replace_fk_if_needed(
"video_tasks_user_id_fkey", "video_tasks",
"users", ["user_id"], ["id"], "SET NULL",
)
_replace_fk_if_needed(
"video_tasks_api_key_id_fkey", "video_tasks",
"api_keys", ["api_key_id"], ["id"], "SET NULL",
)
# --- Backfill: populate snapshots from existing FK joins ---
# (WHERE ... IS NULL makes these inherently idempotent)
op.execute(
"""
UPDATE request_candidates
SET username = (
SELECT users.username FROM users WHERE users.id = request_candidates.user_id
)
WHERE username IS NULL AND user_id IS NOT NULL
"""
)
op.execute(
"""
UPDATE request_candidates
SET api_key_name = (
SELECT api_keys.name FROM api_keys WHERE api_keys.id = request_candidates.api_key_id
)
WHERE api_key_name IS NULL AND api_key_id IS NOT NULL
"""
)
op.execute(
"""
UPDATE video_tasks
SET username = (
SELECT users.username FROM users WHERE users.id = video_tasks.user_id
)
WHERE username IS NULL AND user_id IS NOT NULL
"""
)
op.execute(
"""
UPDATE video_tasks
SET api_key_name = (
SELECT api_keys.name FROM api_keys WHERE api_keys.id = video_tasks.api_key_id
)
WHERE api_key_name IS NULL AND api_key_id IS NOT NULL
"""
)
def downgrade() -> None:
# --- video_tasks: SET NULL -> default (no action), restore NOT NULL ---
_replace_fk_if_needed(
"video_tasks_api_key_id_fkey", "video_tasks",
"api_keys", ["api_key_id"], ["id"], "NO ACTION",
)
_replace_fk_if_needed(
"video_tasks_user_id_fkey", "video_tasks",
"users", ["user_id"], ["id"], "NO ACTION",
)
op.alter_column("video_tasks", "user_id", existing_type=sa.String(length=36), nullable=False)
if _column_exists("video_tasks", "api_key_name"):
op.drop_column("video_tasks", "api_key_name")
if _column_exists("video_tasks", "username"):
op.drop_column("video_tasks", "username")
# --- request_candidates: SET NULL -> CASCADE ---
_replace_fk_if_needed(
"request_candidates_api_key_id_fkey", "request_candidates",
"api_keys", ["api_key_id"], ["id"], "CASCADE",
)
_replace_fk_if_needed(
"request_candidates_user_id_fkey", "request_candidates",
"users", ["user_id"], ["id"], "CASCADE",
)
if _column_exists("request_candidates", "api_key_name"):
op.drop_column("request_candidates", "api_key_name")
if _column_exists("request_candidates", "username"):
op.drop_column("request_candidates", "username")
frontend/src/api/endpoints/pool.ts
+12 -2
View File
@@ -2,6 +2,8 @@ import client from '../client'
import { dedupedRequest } from '@/utils/cache'
import type { AllowedModels, ProxyConfig } from './types/provider'
const POOL_BATCH_ACTION_TIMEOUT_MS = 5 * 60 * 1000
export interface PoolKeyStatus {
key_id: string
key_name: string
@@ -212,13 +214,21 @@ export async function batchActionPoolKeys(
providerId: string,
body: PoolBatchAction,
): Promise<{ affected: number; message: string }> {
const response = await client.post(`/api/admin/pool/${providerId}/keys/batch-action`, body)
const response = await client.post(
`/api/admin/pool/${providerId}/keys/batch-action`,
body,
{ timeout: POOL_BATCH_ACTION_TIMEOUT_MS },
)
return response.data
}
export async function cleanupBannedPoolKeys(
providerId: string,
): Promise<{ affected: number; message: string }> {
const response = await client.post(`/api/admin/pool/${providerId}/keys/cleanup-banned`)
const response = await client.post(
`/api/admin/pool/${providerId}/keys/cleanup-banned`,
undefined,
{ timeout: POOL_BATCH_ACTION_TIMEOUT_MS },
)
return response.data
}
src/api/admin/api_keys/routes.py
+3 -4
View File
@@ -23,6 +23,7 @@ from src.database import get_db
from src.models.api import CreateApiKeyRequest
from src.models.database import ApiKey, Wallet
from src.services.user.apikey import ApiKeyService
from src.services.user.bulk_cleanup import pre_clean_api_key
from src.services.wallet import WalletService
# 应用时区配置，默认为 Asia/Shanghai
@@ -349,10 +350,7 @@ class AdminCreateStandaloneKeyAdapter(AdminApiAdapter):
db = context.db
# 独立Key支持无限制额度（initial_balance_usd = null）
if (
self.key_data.initial_balance_usd is not None
and self.key_data.initial_balance_usd <= 0
):
if self.key_data.initial_balance_usd is not None and self.key_data.initial_balance_usd <= 0:
raise HTTPException(
status_code=400,
detail="创建独立余额Key时，初始余额必须大于 0（或设置为 null 表示无限制）",
@@ -568,6 +566,7 @@ class AdminDeleteApiKeyAdapter(AdminApiAdapter):
raise InvalidRequestException("仅支持删除独立密钥")
user = api_key.user
pre_clean_api_key(db, api_key.id)
db.delete(api_key)
db.commit()
src/api/admin/users/routes.py
+2
View File
@@ -21,6 +21,7 @@ from src.models.api import CreateApiKeyRequest, CreateUserRequest
from src.models.database import ApiKey, User, UserRole
from src.services.system.config import SystemConfigService
from src.services.user.apikey import ApiKeyService
from src.services.user.bulk_cleanup import pre_clean_api_key
from src.services.user.service import UserService
from src.services.wallet import WalletService
from src.utils.cache_decorator import cache_result
@@ -584,6 +585,7 @@ class AdminDeleteUserKeyAdapter(AdminApiAdapter):
if not api_key:
raise NotFoundException("API Key不存在或不属于该用户", "api_key")
pre_clean_api_key(db, api_key.id)
db.delete(api_key)
db.commit()
src/api/handlers/gemini/video_handler.py
+2
View File
@@ -634,6 +634,8 @@ class GeminiVeoHandler(VideoHandlerBase):
external_task_id=external_task_id,
user_id=self.user.id,
api_key_id=self.api_key.id,
username=self.user.username,
api_key_name=self.api_key.name,
provider_id=candidate.provider.id,
endpoint_id=candidate.endpoint.id,
key_id=candidate.key.id,
src/api/handlers/openai/video_handler.py
+4
View File
@@ -887,6 +887,8 @@ class OpenAIVideoHandler(VideoHandlerBase):
external_task_id=external_task_id,
user_id=self.user.id,
api_key_id=self.api_key.id,
username=self.user.username,
api_key_name=self.api_key.name,
provider_id=candidate.provider.id,
endpoint_id=candidate.endpoint.id,
key_id=candidate.key.id,
@@ -1049,6 +1051,8 @@ class OpenAIVideoHandler(VideoHandlerBase):
external_task_id=None,
user_id=self.user.id,
api_key_id=self.api_key.id,
username=self.user.username,
api_key_name=self.api_key.name,
provider_id=None,
endpoint_id=None,
key_id=None,
src/api/user_me/routes.py
+2
View File
@@ -46,6 +46,7 @@ from src.models.database import (
from src.services.system.time_range import TimeRangeParams
from src.services.usage.service import UsageService
from src.services.user.apikey import ApiKeyService
from src.services.user.bulk_cleanup import pre_clean_api_key
from src.services.user.preference import PreferenceService
from src.services.wallet import WalletService
from src.utils.cache_decorator import cache_result
@@ -731,6 +732,7 @@ class DeleteMyApiKeyAdapter(AuthenticatedApiAdapter):
raise NotFoundException("API密钥不存在", "api_key")
if api_key.is_locked:
raise ForbiddenException("该密钥已被管理员锁定，无法删除")
pre_clean_api_key(context.db, api_key.id)
context.db.delete(api_key)
context.db.commit()
return {"message": "API密钥已删除"}
src/models/database.py
+27 -9
View File
@@ -32,7 +32,7 @@ from sqlalchemy import (
text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import declarative_base, relationship
from sqlalchemy.orm import backref, declarative_base, relationship
from ..config import config
from ..core.enums import AuthSource, ProviderBillingType, UserRole
@@ -235,7 +235,7 @@ class ApiKey(Base):
# 关系
user = relationship("User", back_populates="api_keys")
usage_records = relationship("Usage", back_populates="api_key")
usage_records = relationship("Usage", back_populates="api_key", passive_deletes=True)
wallet = relationship("Wallet", back_populates="api_key", uselist=False, passive_deletes=True)
provider_mappings = relationship(
"ApiKeyProviderMapping", back_populates="api_key", cascade="all, delete-orphan"
@@ -325,6 +325,10 @@ class Usage(Base):
String(36), ForeignKey("wallets.id", ondelete="SET NULL"), nullable=True, index=True
)
# 归属快照（删除用户/Key 后仍可追溯）
username = Column(String(100), nullable=True, comment="用户名快照")
api_key_name = Column(String(200), nullable=True, comment="API Key 名称快照")
# 请求信息
request_id = Column(String(100), unique=True, index=True, nullable=False)
provider_name = Column(String(100), nullable=False) # Provider 名称（非外键）
@@ -1932,8 +1936,14 @@ class VideoTask(Base):
external_task_id = Column(String(200))
# 关联
user_id = Column(String(36), ForeignKey("users.id"), nullable=False)
api_key_id = Column(String(36), ForeignKey("api_keys.id"), index=True)
user_id = Column(String(36), ForeignKey("users.id", ondelete="SET NULL"), nullable=True)
api_key_id = Column(
String(36), ForeignKey("api_keys.id", ondelete="SET NULL"), nullable=True, index=True
)
# 归属快照（删除用户/Key 后仍可追溯）
username = Column(String(100), nullable=True, comment="用户名快照")
api_key_name = Column(String(200), nullable=True, comment="API Key 名称快照")
provider_id = Column(String(36), ForeignKey("providers.id"), index=True)
endpoint_id = Column(String(36), ForeignKey("provider_endpoints.id"), index=True)
key_id = Column(String(36), ForeignKey("provider_api_keys.id"), index=True)
@@ -2013,7 +2023,7 @@ class VideoTask(Base):
)
# 关系
user = relationship("User", backref="video_tasks")
user = relationship("User", backref=backref("video_tasks", passive_deletes=True))
remixed_from = relationship("VideoTask", remote_side=[id], backref="remixes")
# 复合索引和唯一约束
@@ -2395,12 +2405,16 @@ class RequestCandidate(Base):
# 关联字段
request_id = Column(String(100), nullable=False, index=True)
user_id = Column(
String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=True, index=True
String(36), ForeignKey("users.id", ondelete="SET NULL"), nullable=True, index=True
)
api_key_id = Column(
String(36), ForeignKey("api_keys.id", ondelete="CASCADE"), nullable=True, index=True
String(36), ForeignKey("api_keys.id", ondelete="SET NULL"), nullable=True, index=True
)
# 归属快照（删除用户/Key 后仍可追溯）
username = Column(String(100), nullable=True, comment="用户名快照")
api_key_name = Column(String(200), nullable=True, comment="API Key 名称快照")
# 候选信息
candidate_index = Column(Integer, nullable=False) # 候选序号（从0开始）
retry_index = Column(Integer, nullable=False, default=0) # 重试序号（从0开始）
@@ -2766,7 +2780,10 @@ class StatsDailyApiKey(Base):
__tablename__ = "stats_daily_api_key"
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
api_key_id = Column(String(36), ForeignKey("api_keys.id", ondelete="CASCADE"), nullable=False)
api_key_id = Column(String(36), ForeignKey("api_keys.id", ondelete="SET NULL"), nullable=True)
api_key_name = Column(
String(200), nullable=True, comment="API Key 名称快照（删除 Key 后仍可追溯）"
)
date = Column(DateTime(timezone=True), nullable=False, index=True)
total_requests = Column(Integer, default=0, nullable=False)
@@ -2887,7 +2904,8 @@ class StatsUserDaily(Base):
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
# 用户关联
user_id = Column(String(36), ForeignKey("users.id", ondelete="CASCADE"), nullable=False)
user_id = Column(String(36), ForeignKey("users.id", ondelete="SET NULL"), nullable=True)
username = Column(String(100), nullable=True, comment="用户名快照（删除用户后仍可追溯）")
# 统计日期 (UTC)
date = Column(DateTime(timezone=True), nullable=False, index=True)
src/modules/proxy_nodes/__init__.py
+1 -1
View File
@@ -109,7 +109,7 @@ async def _on_startup() -> None:
logger.info("Hub worker channel initialized on startup")
except Exception as e:
# ensure_connected 失败时内部会启动重连循环，这里仅记录告警不阻塞启动
logger.warning("Hub worker channel init failed, reconnecting in background: {}", e)
logger.warning("Hub worker channel init failed, reconnecting in background: %s", e)
if active:
logger.info("启动 ProxyNode 心跳检测调度器...")
src/services/candidate/failover.py
+22
View File
@@ -332,6 +332,8 @@ class FailoverEngine:
request_id: str | None = None,
user_id: str | None = None,
api_key_id: str | None = None,
username: str | None = None,
api_key_name: str | None = None,
candidate_record_map: dict[tuple[int, int], str] | None = None,
max_candidates: int | None = None,
max_attempts: int | None = None,
@@ -405,6 +407,8 @@ class FailoverEngine:
candidate_index=candidate_index,
user_id=user_id,
api_key_id=api_key_id,
username=username,
api_key_name=api_key_name,
skip_reason=skip_reason,
)
candidate_keys_fallback.append(
@@ -428,6 +432,8 @@ class FailoverEngine:
request_id=request_id,
user_id=user_id,
api_key_id=api_key_id,
username=username,
api_key_name=api_key_name,
candidate_record_map=candidate_record_map,
candidate_keys_fallback=candidate_keys_fallback,
candidates=candidates,
@@ -476,6 +482,8 @@ class FailoverEngine:
retry_index=retry_index,
user_id=user_id,
api_key_id=api_key_id,
username=username,
api_key_name=api_key_name,
)
self._attach_attempt_context(
@@ -620,6 +628,8 @@ class FailoverEngine:
request_id: str | None,
user_id: str | None,
api_key_id: str | None,
username: str | None,
api_key_name: str | None,
candidate_record_map: dict[tuple[int, int], str] | None,
candidate_keys_fallback: list[CandidateKey],
candidates: list[ProviderCandidate],
@@ -677,6 +687,8 @@ class FailoverEngine:
retry_index=base_retry_index,
user_id=user_id,
api_key_id=api_key_id,
username=username,
api_key_name=api_key_name,
skip_reason=skip_reason,
)
candidate_keys_fallback.append(
@@ -724,6 +736,8 @@ class FailoverEngine:
retry_index=composite_retry_index,
user_id=user_id,
api_key_id=api_key_id,
username=username,
api_key_name=api_key_name,
)
self._attach_attempt_context(
@@ -1202,6 +1216,8 @@ class FailoverEngine:
retry_index: int,
user_id: str | None,
api_key_id: str | None,
username: str | None,
api_key_name: str | None,
) -> str:
# Create "available" record, then caller will mark pending.
extra = self._build_pool_extra_data(candidate)
@@ -1212,6 +1228,8 @@ class FailoverEngine:
retry_index=retry_index,
user_id=user_id,
api_key_id=api_key_id,
username=username,
api_key_name=api_key_name,
provider_id=str(candidate.provider.id),
endpoint_id=str(candidate.endpoint.id),
key_id=str(candidate.key.id),
@@ -1230,6 +1248,8 @@ class FailoverEngine:
retry_index: int = 0,
user_id: str | None,
api_key_id: str | None,
username: str | None,
api_key_name: str | None,
skip_reason: str | None,
) -> str:
extra = self._build_pool_extra_data(candidate)
@@ -1240,6 +1260,8 @@ class FailoverEngine:
retry_index=retry_index,
user_id=user_id,
api_key_id=api_key_id,
username=username,
api_key_name=api_key_name,
provider_id=str(candidate.provider.id),
endpoint_id=str(candidate.endpoint.id),
key_id=str(candidate.key.id),
src/services/orchestration/candidate_resolver.py
+15
View File
@@ -216,6 +216,15 @@ class CandidateResolver:
candidate_records_to_insert: list[dict[str, Any]] = []
candidate_record_map: dict[tuple[int, int], str] = {}
username = None
api_key_name = getattr(user_api_key, "name", None) if user_api_key else None
if user_api_key is not None:
try:
user = getattr(user_api_key, "user", None)
except Exception:
user = None
username = getattr(user, "username", None) if user is not None else None
# 只保存启用的能力（值为 True 的）
active_capabilities = None
@@ -290,6 +299,8 @@ class CandidateResolver:
"retry_index": retry_index,
"user_id": user_id,
"api_key_id": user_api_key.id if user_api_key else None,
"username": username,
"api_key_name": api_key_name,
"provider_id": provider.id,
"endpoint_id": endpoint.id,
"key_id": key_id,
@@ -314,6 +325,8 @@ class CandidateResolver:
"retry_index": 0,
"user_id": user_id,
"api_key_id": user_api_key.id if user_api_key else None,
"username": username,
"api_key_name": api_key_name,
"provider_id": provider.id,
"endpoint_id": endpoint.id,
"key_id": key.id,
@@ -339,6 +352,8 @@ class CandidateResolver:
"retry_index": retry_index,
"user_id": user_id,
"api_key_id": user_api_key.id if user_api_key else None,
"username": username,
"api_key_name": api_key_name,
"provider_id": provider.id,
"endpoint_id": endpoint.id,
"key_id": key.id,
src/services/request/candidate.py
+6
View File
@@ -23,6 +23,8 @@ class RequestCandidateService:
retry_index: int = 0, # 新增：重试序号
user_id: str | None = None,
api_key_id: str | None = None,
username: str | None = None,
api_key_name: str | None = None,
provider_id: str | None = None,
endpoint_id: str | None = None,
key_id: str | None = None,
@@ -42,6 +44,8 @@ class RequestCandidateService:
retry_index: 重试序号（从0开始）
user_id: 用户ID
api_key_id: API Key ID
username: 用户名快照
api_key_name: API Key 名称快照
provider_id: Provider ID
endpoint_id: Endpoint ID
key_id: API Key ID
@@ -58,6 +62,8 @@ class RequestCandidateService:
retry_index=retry_index, # 新增
user_id=user_id,
api_key_id=api_key_id,
username=username,
api_key_name=api_key_name,
provider_id=provider_id,
endpoint_id=endpoint_id,
key_id=key_id,
src/services/scheduling/aware_scheduler.py
+2 -10
View File
@@ -476,18 +476,10 @@ class CacheAwareScheduler:
release_db_connection_before_await(db)
logger.debug(
"[Scheduler] Found {} active providers",
"[Scheduler] Found {} active providers: {}",
len(providers),
", ".join(p.name for p in providers),
)
for p in providers:
logger.debug(
"[Scheduler] Provider: id={}, name={}, is_active={}, endpoints={}, models={}",
p.id[:8] if p.id else "N/A",
p.name,
p.is_active,
len(p.endpoints) if p.endpoints else 0,
len(p.models) if p.models else 0,
)
if not providers:
return [], global_model_id, queried_provider_count
src/services/scheduling/candidate_builder.py
-38
View File
@@ -425,11 +425,6 @@ class CandidateBuilder:
allowed_kinds = {client_kind}
for provider in providers:
logger.debug(
"[Scheduler] Checking provider: {}, endpoints={}",
provider.name,
len(provider.endpoints) if provider.endpoints else 0,
)
# 按端点格式分别判断兼容性与模型/Key 可用性：
# - 同格式端点优先（needs_conversion=False）
# - 跨格式端点次之（needs_conversion=True）
@@ -487,15 +482,7 @@ class CandidateBuilder:
)
for endpoint in endpoints:
logger.debug(
"[Scheduler] Checking endpoint: family={}, kind={}, is_active={}, base_url={}",
getattr(endpoint, "api_family", None),
getattr(endpoint, "endpoint_kind", None),
getattr(endpoint, "is_active", None),
(endpoint.base_url[:50] if endpoint.base_url else "N/A"),
)
if not endpoint.is_active:
logger.debug("[Scheduler] Endpoint skipped: not active")
continue
endpoint_format_str = make_signature_key(
@@ -520,17 +507,6 @@ class CandidateBuilder:
global_conversion_enabled,
skip_endpoint_check=skip_endpoint_check,
)
logger.debug(
"[Scheduler] Format compatibility: client={}, endpoint={}, compatible={}, "
"global={}, provider={}, skip_endpoint={}, reason={}",
client_format_str,
endpoint_format_str,
is_compatible,
global_conversion_enabled,
provider_conversion_enabled,
skip_endpoint_check,
_compat_reason,
)
if not is_compatible:
continue
@@ -547,21 +523,7 @@ class CandidateBuilder:
supports_model, skip_reason, _model_caps, provider_model_names = (
model_support_cache[endpoint_format_str]
)
logger.debug(
"[Scheduler] Model support: provider={}, model={}, supports={}, reason={}",
provider.name,
model_name,
supports_model,
skip_reason,
)
if not supports_model:
logger.debug(
"Provider {} 端点 {} 不支持模型 {}: {}",
provider.name,
endpoint_format_str,
model_name,
skip_reason,
)
continue
# Key 直属 Provider，通过 api_formats 按端点格式筛选
src/services/system/stats_aggregator.py
+12
View File
@@ -426,6 +426,7 @@ class StatsAggregatorService:
stats = (
db.query(
Usage.api_key_id,
func.max(Usage.api_key_name).label("api_key_name"),
func.count(Usage.id).label("total_requests"),
func.sum(case((error_cond, 1), else_=0)).label("error_requests"),
func.sum(Usage.input_tokens).label("input_tokens"),
@@ -460,6 +461,10 @@ class StatsAggregatorService:
id=str(uuid.uuid4()), date=day_start, api_key_id=stat.api_key_id
)
# 填充 api_key_name 快照（优先用已有值，新数据从 usage 聚合获取）
if not record.api_key_name and stat.api_key_name:
record.api_key_name = stat.api_key_name
error_requests = int(stat.error_requests or 0)
total_requests = int(stat.total_requests or 0)
record.total_requests = total_requests
@@ -598,6 +603,7 @@ class StatsAggregatorService:
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
func.sum(Usage.total_cost_usd).label("total_cost"),
func.max(Usage.username).label("username"),
)
.filter(
and_(
@@ -609,6 +615,12 @@ class StatsAggregatorService:
.first()
)
# 填充 username 快照：从 Usage 聚合获取，用户删除后仍可追溯
if not stats.username:
username = getattr(aggregated, "username", None)
if username:
stats.username = username
total_requests = int(getattr(aggregated, "total_requests", 0) or 0)
if total_requests == 0:
stats.total_requests = 0
src/services/task/execute/sync_execute.py
+15 -1
View File
@@ -128,10 +128,13 @@ class SyncTaskExecutionService:
affinity_key = str(user_api_key.id)
user_id = str(user_api_key.user_id)
api_format_norm = normalize_endpoint_signature(api_format)
username_snapshot = None
api_key_name_snapshot = getattr(user_api_key, "name", None)
# Keep pending usage creation behavior consistent with previous behavior
try:
user = self.db.query(User).filter(User.id == user_api_key.user_id).first()
username_snapshot = getattr(user, "username", None) if user else None
UsageService.create_pending_usage(
db=self.db,
request_id=request_id,
@@ -216,6 +219,8 @@ class SyncTaskExecutionService:
retry_index=retry_index,
user_id=user_id,
api_key_id=str(user_api_key.id),
username=username_snapshot,
api_key_name=api_key_name_snapshot,
provider_id=str(candidate.provider.id),
endpoint_id=str(candidate.endpoint.id),
key_id=str(candidate.key.id),
@@ -253,7 +258,14 @@ class SyncTaskExecutionService:
max_attempts=max_attempts_local,
is_stream=is_stream,
)
_ = (attempt_id, _provider_name, _provider_id, _endpoint_id, _key_id, _first_byte_time_ms)
_ = (
attempt_id,
_provider_name,
_provider_id,
_endpoint_id,
_key_id,
_first_byte_time_ms,
)
# Account Pool: on success, update sticky binding + LRU.
await self._pool_ops.pool_on_success(candidate, request_body)
@@ -335,6 +347,8 @@ class SyncTaskExecutionService:
request_id=request_id,
user_id=user_id,
api_key_id=str(user_api_key.id),
username=username_snapshot,
api_key_name=api_key_name_snapshot,
candidate_record_map=candidate_record_map,
max_attempts=max_attempts,
execution_error_handler=_handle_exec_err,
src/services/task/service.py
+7 -1
View File
@@ -363,6 +363,9 @@ class TaskService:
elif user_api_key is not None and getattr(user_api_key, "user_id", None):
user_id = str(user_api_key.user_id)
username_snapshot = getattr(resolved_user, "username", None) if resolved_user else None
api_key_name_snapshot = getattr(user_api_key, "name", None) if user_api_key else None
resolved_affinity_key = affinity_key
if not resolved_affinity_key:
api_key_id = getattr(user_api_key, "id", None) if user_api_key is not None else None
@@ -443,6 +446,8 @@ class TaskService:
retry_index=retry_index,
user_id=user_id,
api_key_id=(getattr(user_api_key, "id", None) if user_api_key else None),
username=username_snapshot,
api_key_name=api_key_name_snapshot,
provider_id=str(candidate.provider.id),
endpoint_id=str(candidate.endpoint.id),
key_id=str(candidate.key.id),
@@ -564,6 +569,8 @@ class TaskService:
api_key_id=(
str(user_api_key.id) if user_api_key and getattr(user_api_key, "id", None) else None
),
username=username_snapshot,
api_key_name=api_key_name_snapshot,
candidate_record_map=candidate_record_map,
max_attempts=max_attempts,
execution_error_handler=_handle_exec_err,
@@ -605,7 +612,6 @@ class TaskService:
execution_state.last_error,
)
async def submit_with_failover(
self,
*,
src/services/usage/_recording_helpers.py
+2
View File
@@ -193,6 +193,8 @@ def build_usage_params(
return {
"user_id": user.id if user else None,
"api_key_id": api_key.id if api_key else None,
"username": user.username if user else None,
"api_key_name": api_key.name if api_key else None,
"request_id": request_id,
"provider_name": provider,
"model": model,
src/services/usage/lifecycle.py
+2
View File
@@ -65,6 +65,8 @@ class UsageLifecycleMixin:
usage = Usage(
user_id=user.id if user else None,
api_key_id=api_key.id if api_key else None,
username=user.username if user else None,
api_key_name=api_key.name if api_key else None,
request_id=request_id,
provider_name="pending", # 尚未确定 provider
model=model,
src/services/user/apikey.py
+4 -3
View File
@@ -13,6 +13,7 @@ from sqlalchemy.orm import Session
from src.core.crypto import crypto_service
from src.core.logger import logger
from src.models.database import ApiKey, Usage
from src.services.user.bulk_cleanup import pre_clean_api_key
class ApiKeyService:
@@ -82,8 +83,7 @@ class ApiKeyService:
db.refresh(api_key)
logger.info(
f"创建API密钥: 用户ID {user_id}, 密钥名 {api_key.name}, "
f"独立Key={is_standalone}"
f"创建API密钥: 用户ID {user_id}, 密钥名 {api_key.name}, " f"独立Key={is_standalone}"
)
return api_key, key # 返回密钥对象和明文密钥
@@ -245,7 +245,8 @@ class ApiKeyService:
)
if should_delete:
# 物理删除（Usage记录会保留，因为是 SET NULL）
# 物理删除（Usage / RequestCandidate / VideoTask 等记录保留）
pre_clean_api_key(db, api_key.id)
db.delete(api_key)
logger.info(
f"删除过期API密钥: ID {api_key.id}, 名称 {api_key.name}, "
src/services/user/bulk_cleanup.py
+111
View File
@@ -0,0 +1,111 @@
from __future__ import annotations
import time
from typing import Any
from sqlalchemy.orm import Session
from src.core.logger import logger
from src.models.database import RequestCandidate, Usage
_POSTGRES_BATCH_SIZE = 2000
_SQLITE_BATCH_SIZE = 900
def _resolve_batch_size(db: Session) -> int:
try:
bind = db.get_bind()
dialect_name = str(getattr(getattr(bind, "dialect", None), "name", "") or "").lower()
except Exception:
dialect_name = ""
if dialect_name == "sqlite":
return _SQLITE_BATCH_SIZE
return _POSTGRES_BATCH_SIZE
def batch_nullify_fk(
db: Session,
model: type[Any],
column_name: str,
entity_id: str | None,
) -> int:
"""分批将大表外键置空，避免单个长事务阻塞删除流程。"""
if not entity_id:
return 0
column = getattr(model, column_name)
primary_key_column = next(iter(model.__table__.primary_key.columns))
batch_size = _resolve_batch_size(db)
total_updated = 0
batch_index = 0
started_at = time.monotonic()
while True:
batch_ids = [
row[0]
for row in db.query(primary_key_column)
.filter(column == entity_id)
.limit(batch_size)
.all()
]
if not batch_ids:
break
batch_index += 1
batch_started_at = time.monotonic()
updated = int(
db.query(model)
.filter(primary_key_column.in_(batch_ids))
.update({column: None}, synchronize_session=False)
or 0
)
db.commit()
total_updated += updated
elapsed_ms = int((time.monotonic() - batch_started_at) * 1000)
logger.info(
"批量清理 {}.{}: batch={}, updated={}, entity_id={}, elapsed_ms={}",
model.__tablename__,
column_name,
batch_index,
updated,
entity_id,
elapsed_ms,
)
if len(batch_ids) < batch_size:
break
if total_updated > 0:
total_elapsed_ms = int((time.monotonic() - started_at) * 1000)
logger.info(
"批量清理完成 {}.{}: total_updated={}, entity_id={}, elapsed_ms={}",
model.__tablename__,
column_name,
total_updated,
entity_id,
total_elapsed_ms,
)
return total_updated
def pre_clean_api_key(db: Session, api_key_id: str | None) -> int:
"""预清理 API Key 在大表中的外键引用，减少后续删除锁竞争。"""
if not api_key_id:
return 0
usage_rows = batch_nullify_fk(db, Usage, "api_key_id", api_key_id)
candidate_rows = batch_nullify_fk(db, RequestCandidate, "api_key_id", api_key_id)
total_rows = usage_rows + candidate_rows
if total_rows > 0:
logger.info(
"API Key 预清理完成: api_key_id={}, usage={}, request_candidates={}",
api_key_id,
usage_rows,
candidate_rows,
)
return total_rows
src/services/user/service.py
+36 -23
View File
@@ -15,6 +15,7 @@ from src.core.logger import logger
from src.core.validators import EmailValidator, PasswordValidator, UsernameValidator
from src.models.database import ApiKey, GlobalModel, Model, Provider, Usage, User, UserRole
from src.services.cache.user_cache import UserCacheService
from src.services.user.bulk_cleanup import batch_nullify_fk, pre_clean_api_key
from src.utils.transaction_manager import retry_on_database_error, transactional
@@ -252,15 +253,15 @@ class UserService:
return user
@staticmethod
@transactional()
def delete_user(db: Session, user_id: str) -> bool:
"""删除用户（硬删除）
删除流程：
1. 检查未完结账务，阻止删除
2. 手动删除 ORM cascade 冲突的子记录
3. 删除用户记录
4. 财务记录（Wallet/PaymentOrder/RefundRequest/WalletTransaction）和
2. 预清理 Usage / RequestCandidate 等大表外键
3. 手动删除 ORM cascade 冲突的子记录
4. 删除用户记录
5. 财务记录（Wallet/PaymentOrder/RefundRequest/WalletTransaction）和
Usage 记录保留，外键 SET NULL，由自动清理策略统一回收
"""
from src.models.database import (
@@ -268,6 +269,7 @@ class UserService:
ApiKey,
PaymentOrder,
RefundRequest,
RequestCandidate,
UserPreference,
Wallet,
)
@@ -312,28 +314,39 @@ class UserService:
if pending_order_count > 0:
raise ValueError("用户存在未完结充值订单，禁止删除")
# 手动删除子记录，避免 SQLAlchemy 的 ORM cascade 与数据库 CASCADE 冲突
# （UserPreference/AnnouncementRead 的数据库外键是 ON DELETE CASCADE，
# 但 SQLAlchemy 会先尝试 UPDATE SET NULL 导致冲突）
db.query(UserPreference).filter(UserPreference.user_id == user_id).delete(
synchronize_session=False
)
db.query(AnnouncementRead).filter(AnnouncementRead.user_id == user_id).delete(
synchronize_session=False
)
api_key_ids = [
api_key_id
for (api_key_id,) in db.query(ApiKey.id).filter(ApiKey.user_id == user_id).all()
]
api_key_count = len(api_key_ids)
# 财务记录（Wallet/WalletTransaction/PaymentOrder/RefundRequest/PaymentCallback）
# 和 Usage 记录全部保留，数据库外键 SET NULL 自动断开关联，
# 由自动清理策略统一回收。
# 注意：batch_nullify_fk 内部分批 commit，预清理部分不可回滚。
# 这是预期行为：SET NULL 是幂等操作，即使后续步骤失败，
# 已置空的外键不影响数据完整性，重新执行删除即可。
try:
for api_key_id in api_key_ids:
pre_clean_api_key(db, api_key_id)
api_key_count = int(
db.query(func.count(ApiKey.id)).filter(ApiKey.user_id == user_id).scalar() or 0
)
db.query(ApiKey).filter(ApiKey.user_id == user_id).delete(synchronize_session=False)
batch_nullify_fk(db, Usage, "user_id", user_id)
batch_nullify_fk(db, RequestCandidate, "user_id", user_id)
# 现在删除用户（Usage, AuditLog, RequestAttempt 会通过数据库 SET NULL 保留）
db.delete(user)
db.commit() # 立即提交事务,释放数据库锁
db.query(UserPreference).filter(UserPreference.user_id == user_id).delete(
synchronize_session=False
)
db.query(AnnouncementRead).filter(AnnouncementRead.user_id == user_id).delete(
synchronize_session=False
)
# 财务记录（Wallet/WalletTransaction/PaymentOrder/RefundRequest/PaymentCallback）
# 和 Usage / RequestCandidate / VideoTask 记录全部保留，数据库外键 SET NULL 自动断开关联。
db.query(ApiKey).filter(ApiKey.user_id == user_id).delete(synchronize_session=False)
# 现在删除用户（Usage, AuditLog, RequestAttempt 等会通过数据库 SET NULL 保留）
db.delete(user)
db.commit() # 立即提交事务,释放数据库锁
except Exception:
db.rollback()
raise
# 清除用户缓存
asyncio.create_task(UserCacheService.invalidate_user_cache(user_id, email))
tests/services/test_user_bulk_cleanup.py
+35
View File
@@ -0,0 +1,35 @@
from __future__ import annotations
from sqlalchemy import Column, Integer, String, create_engine
from sqlalchemy.orm import Session, declarative_base, sessionmaker
from src.services.user.bulk_cleanup import batch_nullify_fk
Base = declarative_base()
class DemoRow(Base):
__tablename__ = "demo_rows"
id = Column(Integer, primary_key=True, autoincrement=True)
ref_id = Column(String(36), nullable=True)
def test_batch_nullify_fk_handles_sqlite_batches() -> None:
engine = create_engine("sqlite:///:memory:")
Base.metadata.create_all(engine)
session_factory = sessionmaker(bind=engine)
db: Session = session_factory()
try:
db.add_all([DemoRow(ref_id="entity-1") for _ in range(905)])
db.add_all([DemoRow(ref_id="entity-2") for _ in range(3)])
db.commit()
updated = batch_nullify_fk(db, DemoRow, "ref_id", "entity-1")
assert updated == 905
assert db.query(DemoRow).filter(DemoRow.ref_id.is_(None)).count() == 905
assert db.query(DemoRow).filter(DemoRow.ref_id == "entity-2").count() == 3
finally:
db.close()
tests/services/test_user_service_delete_user.py
+74 -4
View File
@@ -1,11 +1,19 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import MagicMock
from unittest.mock import AsyncMock, MagicMock, call
import pytest
from src.models.database import PaymentOrder, RefundRequest, User, Wallet
from src.models.database import (
ApiKey,
PaymentOrder,
RefundRequest,
RequestCandidate,
Usage,
User,
Wallet,
)
from src.services.user.service import UserService
@@ -40,7 +48,7 @@ def test_delete_user_blocks_when_unfinished_refund_exists() -> None:
db.query.side_effect = _query
with pytest.raises(ValueError, match="未完结退款"):
UserService.delete_user.__wrapped__(db, "user-1")
UserService.delete_user(db, "user-1")
db.delete.assert_not_called()
@@ -82,6 +90,68 @@ def test_delete_user_blocks_when_unfinished_payment_order_exists() -> None:
db.query.side_effect = _query
with pytest.raises(ValueError, match="未完结充值订单"):
UserService.delete_user.__wrapped__(db, "user-2")
UserService.delete_user(db, "user-2")
db.delete.assert_not_called()
def test_delete_user_precleans_large_tables_before_final_delete(
monkeypatch: pytest.MonkeyPatch,
) -> None:
user = SimpleNamespace(id="user-3", email="u3@example.com")
user_query = MagicMock()
user_query.filter.return_value = user_query
user_query.first.return_value = user
wallet_ids_query = MagicMock()
wallet_ids_query.outerjoin.return_value = wallet_ids_query
wallet_ids_query.filter.return_value = wallet_ids_query
wallet_ids_query.all.return_value = []
api_key_ids_query = MagicMock()
api_key_ids_query.filter.return_value = api_key_ids_query
api_key_ids_query.all.return_value = [("key-1",), ("key-2",)]
delete_query = MagicMock()
delete_query.filter.return_value = delete_query
delete_query.delete.return_value = 2
db = MagicMock()
def _query(model: object) -> MagicMock:
if model is User:
return user_query
if model is Wallet.id:
return wallet_ids_query
if model is ApiKey.id:
return api_key_ids_query
return delete_query
db.query.side_effect = _query
pre_clean_api_key = MagicMock()
batch_nullify_fk = MagicMock()
invalidate_user_cache = AsyncMock()
create_task = MagicMock()
monkeypatch.setattr("src.services.user.service.pre_clean_api_key", pre_clean_api_key)
monkeypatch.setattr("src.services.user.service.batch_nullify_fk", batch_nullify_fk)
monkeypatch.setattr(
"src.services.user.service.UserCacheService.invalidate_user_cache",
invalidate_user_cache,
)
monkeypatch.setattr("src.services.user.service.asyncio.create_task", create_task)
assert UserService.delete_user(db, "user-3") is True
assert pre_clean_api_key.call_args_list == [call(db, "key-1"), call(db, "key-2")]
assert batch_nullify_fk.call_args_list == [
call(db, Usage, "user_id", "user-3"),
call(db, RequestCandidate, "user_id", "user-3"),
]
db.delete.assert_called_once_with(user)
db.commit.assert_called_once()
db.rollback.assert_not_called()
invalidate_user_cache.assert_called_once_with("user-3", "u3@example.com")
create_task.assert_called_once()
tests/unit/test_request_candidate_create_candidate.py
+29
View File
@@ -0,0 +1,29 @@
from __future__ import annotations
from unittest.mock import MagicMock
from src.services.request.candidate import RequestCandidateService
def test_create_candidate_persists_snapshot_fields() -> None:
db = MagicMock()
candidate = RequestCandidateService.create_candidate(
db=db,
request_id="req-1",
candidate_index=0,
retry_index=1,
user_id="user-1",
api_key_id="key-1",
username="alice",
api_key_name="Primary Key",
provider_id="provider-1",
endpoint_id="endpoint-1",
key_id="pool-key-1",
status="available",
)
assert candidate.username == "alice"
assert candidate.api_key_name == "Primary Key"
db.add.assert_called_once_with(candidate)
db.flush.assert_called_once()