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

refactor(cost,perf): 成本字段 Float 改 Numeric 并优化多处查询性能

Browse Source 
- 数据库所有 cost/price 字段从 Float 改为 Numeric(20,8)，解决浮点精度问题
- API 响应中 Decimal 值统一用 float() 转换确保 JSON 序列化
- 路由中手动 commit 后标记 tx_committed_by_route 防止中间件重复提交
- 多处查询优化：SQL 聚合替代 Python 遍历、load_only/defer 减少字段加载、
  批量 DELETE 替代逐条 ORM 删除、N+1 查询消除、UNION ALL 合并多表日期查询
- 新增 provider_api_keys (provider_id, is_active) 复合索引
- 候选构建热路径 defer 冷字段，钱包扣费合并解析与加锁查询
This commit is contained in:
fawney19
2026-03-08 16:44:16 +08:00
parent bd3f73c2fc
commit 48d13762d9
25 changed files with 474 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files
alembic/versions/20260308_1530_2053ab8ed764_cost_fields_float_to_numeric.py
+165
View File
@@ -0,0 +1,165 @@
"""cost fields: Float -> Numeric(20,8) + provider_api_keys composite index
Revision ID: 2053ab8ed764
Revises: 13a4c8f6d9e0
Create Date: 2026-03-08 15:30:00.000000+00:00
"""
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision = "2053ab8ed764"
down_revision = "13a4c8f6d9e0"
branch_labels = None
depends_on = None
# (table_name, column_name, nullable, server_default)
_COST_COLUMNS: list[tuple[str, str, bool, str | None]] = [
# api_keys
("api_keys", "total_cost_usd", True, "0.0"),
# usage
("usage", "input_cost_usd", True, "0.0"),
("usage", "output_cost_usd", True, "0.0"),
("usage", "cache_cost_usd", True, "0.0"),
("usage", "cache_creation_cost_usd", True, "0.0"),
("usage", "cache_read_cost_usd", True, "0.0"),
("usage", "request_cost_usd", True, "0.0"),
("usage", "total_cost_usd", True, "0.0"),
("usage", "actual_input_cost_usd", True, "0.0"),
("usage", "actual_output_cost_usd", True, "0.0"),
("usage", "actual_cache_creation_cost_usd", True, "0.0"),
("usage", "actual_cache_read_cost_usd", True, "0.0"),
("usage", "actual_request_cost_usd", True, "0.0"),
("usage", "actual_total_cost_usd", True, "0.0"),
("usage", "rate_multiplier", True, "1.0"),
("usage", "input_price_per_1m", True, None),
("usage", "output_price_per_1m", True, None),
("usage", "cache_creation_price_per_1m", True, None),
("usage", "cache_read_price_per_1m", True, None),
("usage", "price_per_request", True, None),
# providers
("providers", "monthly_quota_usd", True, None),
("providers", "monthly_used_usd", True, "0.0"),
# global_models
("global_models", "default_price_per_request", True, None),
# models
("models", "price_per_request", True, None),
# stats_hourly
("stats_hourly", "total_cost", False, "0.0"),
("stats_hourly", "actual_total_cost", False, "0.0"),
# stats_hourly_user
("stats_hourly_user", "total_cost", False, "0.0"),
# stats_hourly_model
("stats_hourly_model", "total_cost", False, "0.0"),
# stats_hourly_provider
("stats_hourly_provider", "total_cost", False, "0.0"),
# stats_daily
("stats_daily", "total_cost", False, "0.0"),
("stats_daily", "actual_total_cost", False, "0.0"),
("stats_daily", "input_cost", False, "0.0"),
("stats_daily", "output_cost", False, "0.0"),
("stats_daily", "cache_creation_cost", False, "0.0"),
("stats_daily", "cache_read_cost", False, "0.0"),
# stats_daily_model
("stats_daily_model", "total_cost", False, "0.0"),
# stats_daily_provider
("stats_daily_provider", "total_cost", False, "0.0"),
# stats_daily_api_key
("stats_daily_api_key", "total_cost", False, "0.0"),
# stats_summary
("stats_summary", "all_time_cost", False, "0.0"),
("stats_summary", "all_time_actual_cost", False, "0.0"),
# stats_user_daily
("stats_user_daily", "total_cost", False, "0.0"),
]
# rate_multiplier uses a smaller precision
_RATE_MULTIPLIER_TYPE = sa.Numeric(10, 6)
_COST_TYPE = sa.Numeric(20, 8)
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 _index_exists(index_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
{"name": index_name},
)
return result.scalar() is not None
def _is_numeric_type(table_name: str, column_name: str) -> bool:
"""Check if a column is already numeric type (not float/double precision)."""
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT data_type FROM information_schema.columns "
"WHERE table_name = :table AND column_name = :col"
),
{"table": table_name, "col": column_name},
)
data_type = result.scalar()
return data_type == "numeric"
def upgrade() -> None:
# -- 1. cost fields: Float -> Numeric
for table, col, _nullable, _default in _COST_COLUMNS:
if not _column_exists(table, col):
continue
if _is_numeric_type(table, col):
continue
target_type = _RATE_MULTIPLIER_TYPE if col == "rate_multiplier" else _COST_TYPE
op.alter_column(
table,
col,
type_=target_type,
existing_type=sa.Float(),
postgresql_using=f"{col}::numeric",
)
# -- 2. provider_api_keys composite index
if not _index_exists("idx_provider_api_keys_provider_active"):
op.create_index(
"idx_provider_api_keys_provider_active",
"provider_api_keys",
["provider_id", "is_active"],
)
def downgrade() -> None:
# -- 2. drop composite index
if _index_exists("idx_provider_api_keys_provider_active"):
op.drop_index(
"idx_provider_api_keys_provider_active",
table_name="provider_api_keys",
)
# -- 1. Numeric -> Float
for table, col, _nullable, _default in _COST_COLUMNS:
if not _column_exists(table, col):
continue
if not _is_numeric_type(table, col):
continue
op.alter_column(
table,
col,
type_=sa.Float(),
existing_type=_COST_TYPE if col != "rate_multiplier" else _RATE_MULTIPLIER_TYPE,
postgresql_using=f"{col}::double precision",
)
src/api/admin/adaptive.py
+29 -11
View File
@@ -16,7 +16,8 @@ from typing import Any
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from pydantic import BaseModel, Field, ValidationError
from sqlalchemy.orm import Session
from sqlalchemy import func
from sqlalchemy.orm import Session, load_only
from src.api.base.admin_adapter import AdminApiAdapter
from src.api.base.context import ApiRequestContext
@@ -360,18 +361,35 @@ class SetRPMLimitAdapter(AdminApiAdapter):
class AdaptiveSummaryAdapter(AdminApiAdapter):
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
# 自适应模式：rpm_limit = NULL
adaptive_keys = (
context.db.query(ProviderAPIKey).filter(ProviderAPIKey.rpm_limit.is_(None)).all()
db = context.db
is_adaptive = ProviderAPIKey.rpm_limit.is_(None)
# SQL 聚合获取 count / sum，避免全表 ORM 加载
total_keys, total_concurrent_429, total_rpm_429 = (
db.query(
func.count(ProviderAPIKey.id),
func.coalesce(func.sum(ProviderAPIKey.concurrent_429_count), 0),
func.coalesce(func.sum(ProviderAPIKey.rpm_429_count), 0),
)
.filter(is_adaptive)
.one()
)
total_keys = len(adaptive_keys)
total_concurrent_429 = sum(key.concurrent_429_count or 0 for key in adaptive_keys)
total_rpm_429 = sum(key.rpm_429_count or 0 for key in adaptive_keys)
total_adjustments = sum(len(key.adjustment_history or []) for key in adaptive_keys)
# adjustment_history 是 JSON 列，长度只能在 Python 侧统计；
# 只加载有历史记录的 key 的必要列
keys_with_history = (
db.query(ProviderAPIKey)
.options(
load_only(ProviderAPIKey.id, ProviderAPIKey.name, ProviderAPIKey.adjustment_history)
)
.filter(is_adaptive, ProviderAPIKey.adjustment_history.isnot(None))
.all()
)
total_adjustments = sum(len(key.adjustment_history or []) for key in keys_with_history)
recent_adjustments = []
for key in adaptive_keys:
for key in keys_with_history:
if key.adjustment_history:
for adj in key.adjustment_history[-3:]:
recent_adjustments.append(
@@ -386,8 +404,8 @@ class AdaptiveSummaryAdapter(AdminApiAdapter):
return {
"total_adaptive_keys": total_keys,
"total_concurrent_429_errors": total_concurrent_429,
"total_rpm_429_errors": total_rpm_429,
"total_concurrent_429_errors": int(total_concurrent_429),
"total_rpm_429_errors": int(total_rpm_429),
"total_adjustments": total_adjustments,
"recent_adjustments": recent_adjustments[:10],
}
src/api/admin/api_keys/routes.py
+4
View File
@@ -298,6 +298,7 @@ class AdminListStandaloneKeysAdapter(AdminApiAdapter):
wallet_initialized = True
if wallet_initialized:
db.commit()
context.request.state.tx_committed_by_route = True
for api_key in api_keys:
db.refresh(api_key)
@@ -390,6 +391,7 @@ class AdminCreateStandaloneKeyAdapter(AdminApiAdapter):
if wallet is None:
raise InvalidRequestException("独立密钥钱包初始化失败")
db.commit()
context.request.state.tx_committed_by_route = True
db.refresh(api_key)
wallet_summary = WalletService.serialize_wallet_summary(wallet)
@@ -533,6 +535,7 @@ class AdminToggleApiKeyAdapter(AdminApiAdapter):
api_key.is_active = not api_key.is_active
api_key.updated_at = datetime.now(timezone.utc)
db.commit()
context.request.state.tx_committed_by_route = True
db.refresh(api_key)
logger.info(
@@ -569,6 +572,7 @@ class AdminDeleteApiKeyAdapter(AdminApiAdapter):
pre_clean_api_key(db, api_key.id)
db.delete(api_key)
db.commit()
context.request.state.tx_committed_by_route = True
logger.info(
f"管理员删除API密钥: Key ID {self.key_id}, 用户 {user.email if user else '未知'}"
src/api/admin/pool/routes.py
+6 -3
View File
@@ -1244,10 +1244,13 @@ class AdminCleanupBannedKeysAdapter(AdminApiAdapter):
return BatchActionResponse(affected=0, message="未发现已知封号账号")
banned_key_ids = [str(key.id) for key in banned_keys]
for key in banned_keys:
db.delete(key)
try:
db.execute(
sa_delete(ProviderAPIKey).where(
ProviderAPIKey.provider_id == pid,
ProviderAPIKey.id.in_(banned_key_ids),
)
)
db.commit()
except Exception as exc:
db.rollback()
src/api/admin/provider_strategy.py
+22 -11
View File
@@ -10,6 +10,7 @@ from typing import Any
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field, ValidationError
from sqlalchemy import func
from sqlalchemy.orm import Session
from src.api.base.admin_adapter import AdminApiAdapter
@@ -228,20 +229,26 @@ class AdminProviderStatsAdapter(AdminApiAdapter):
raise HTTPException(status_code=404, detail="Provider not found")
since = datetime.now(timezone.utc) - timedelta(hours=self.hours)
stats = (
db.query(ProviderUsageTracking)
row = (
db.query(
func.coalesce(func.sum(ProviderUsageTracking.total_requests), 0),
func.coalesce(func.sum(ProviderUsageTracking.successful_requests), 0),
func.coalesce(func.sum(ProviderUsageTracking.failed_requests), 0),
func.coalesce(func.avg(ProviderUsageTracking.avg_response_time_ms), 0),
func.coalesce(func.sum(ProviderUsageTracking.total_cost_usd), 0),
)
.filter(
ProviderUsageTracking.provider_id == self.provider_id,
ProviderUsageTracking.window_start >= since,
)
.all()
.one()
)
total_requests = sum(s.total_requests for s in stats)
total_success = sum(s.successful_requests for s in stats)
total_failures = sum(s.failed_requests for s in stats)
avg_response_time = sum(s.avg_response_time_ms for s in stats) / len(stats) if stats else 0
total_cost = sum(s.total_cost_usd for s in stats)
total_requests = int(row[0])
total_success = int(row[1])
total_failures = int(row[2])
avg_response_time = float(row[3])
total_cost = float(row[4])
return JSONResponse(
{
@@ -250,10 +257,14 @@ class AdminProviderStatsAdapter(AdminApiAdapter):
"period_hours": self.hours,
"billing_info": {
"billing_type": provider.billing_type.value,
"monthly_quota_usd": provider.monthly_quota_usd,
"monthly_used_usd": provider.monthly_used_usd,
"monthly_quota_usd": (
float(provider.monthly_quota_usd)
if provider.monthly_quota_usd is not None
else None
),
"monthly_used_usd": float(provider.monthly_used_usd or 0),
"quota_remaining_usd": (
provider.monthly_quota_usd - provider.monthly_used_usd
float(provider.monthly_quota_usd - provider.monthly_used_usd)
if provider.monthly_quota_usd is not None
else None
),
src/api/admin/stats/quota.py
+2 -2
View File
@@ -41,8 +41,8 @@ class AdminQuotaUsageAdapter(AdminApiAdapter):
result = []
for provider in providers:
quota = provider.monthly_quota_usd or 0.0
used = float(provider.monthly_used_usd or 0.0)
quota = float(provider.monthly_quota_usd or 0)
used = float(provider.monthly_used_usd or 0)
remaining = max(quota - used, 0.0)
usage_percent = round((used / quota) * 100, 2) if quota > 0 else 0.0
src/api/admin/usage/routes.py
+31 -11
View File
@@ -1312,22 +1312,42 @@ class AdminUsageDetailAdapter(AdminApiAdapter):
"total": usage_record.total_tokens,
},
"cost": {
"input": usage_record.input_cost_usd,
"output": usage_record.output_cost_usd,
"total": usage_record.total_cost_usd,
"input": float(usage_record.input_cost_usd or 0),
"output": float(usage_record.output_cost_usd or 0),
"total": float(usage_record.total_cost_usd or 0),
},
"cache_creation_input_tokens": usage_record.cache_creation_input_tokens,
"cache_read_input_tokens": usage_record.cache_read_input_tokens,
"cache_creation_input_tokens_5m": usage_record.cache_creation_input_tokens_5m or 0,
"cache_creation_input_tokens_1h": usage_record.cache_creation_input_tokens_1h or 0,
"cache_creation_cost": getattr(usage_record, "cache_creation_cost_usd", 0.0),
"cache_read_cost": getattr(usage_record, "cache_read_cost_usd", 0.0),
"request_cost": getattr(usage_record, "request_cost_usd", 0.0),
"input_price_per_1m": usage_record.input_price_per_1m,
"output_price_per_1m": usage_record.output_price_per_1m,
"cache_creation_price_per_1m": usage_record.cache_creation_price_per_1m,
"cache_read_price_per_1m": usage_record.cache_read_price_per_1m,
"price_per_request": usage_record.price_per_request,
"cache_creation_cost": float(getattr(usage_record, "cache_creation_cost_usd", 0) or 0),
"cache_read_cost": float(getattr(usage_record, "cache_read_cost_usd", 0) or 0),
"request_cost": float(getattr(usage_record, "request_cost_usd", 0) or 0),
"input_price_per_1m": (
float(usage_record.input_price_per_1m)
if usage_record.input_price_per_1m is not None
else None
),
"output_price_per_1m": (
float(usage_record.output_price_per_1m)
if usage_record.output_price_per_1m is not None
else None
),
"cache_creation_price_per_1m": (
float(usage_record.cache_creation_price_per_1m)
if usage_record.cache_creation_price_per_1m is not None
else None
),
"cache_read_price_per_1m": (
float(usage_record.cache_read_price_per_1m)
if usage_record.cache_read_price_per_1m is not None
else None
),
"price_per_request": (
float(usage_record.price_per_request)
if usage_record.price_per_request is not None
else None
),
"request_type": usage_record.request_type,
"is_stream": usage_record.is_stream,
"status_code": usage_record.status_code,
src/api/admin/users/routes.py
+2
View File
@@ -588,6 +588,7 @@ class AdminDeleteUserKeyAdapter(AdminApiAdapter):
pre_clean_api_key(db, api_key.id)
db.delete(api_key)
db.commit()
context.request.state.tx_committed_by_route = True
logger.info(f"管理员删除用户API Key: 用户ID {self.user_id}, Key ID {self.key_id}")
@@ -624,6 +625,7 @@ class AdminToggleUserKeyLockAdapter(AdminApiAdapter):
api_key.is_locked = not api_key.is_locked
db.commit()
context.request.state.tx_committed_by_route = True
db.refresh(api_key)
logger.info(
src/api/auth/routes.py
+9
View File
@@ -287,6 +287,7 @@ class AuthLoginAdapter(AuthPublicAdapter):
error_reason="邮箱或密码错误",
)
db.commit()
context.request.state.tx_committed_by_route = True
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="邮箱或密码错误")
AuditService.log_login_attempt(
@@ -298,6 +299,7 @@ class AuthLoginAdapter(AuthPublicAdapter):
user_id=user.id,
)
db.commit()
context.request.state.tx_committed_by_route = True
access_token = AuthService.create_access_token(
data={
@@ -472,6 +474,7 @@ class AuthRegisterAdapter(AuthPublicAdapter):
metadata={"username": register_request.username, "reason": "registration_disabled"},
)
db.commit()
context.request.state.tx_committed_by_route = True
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="系统暂不开放注册")
email = register_request.email
@@ -514,6 +517,7 @@ class AuthRegisterAdapter(AuthPublicAdapter):
metadata={"email": email, "reason": "email_suffix_not_allowed"},
)
db.commit()
context.request.state.tx_committed_by_route = True
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=suffix_error,
@@ -550,6 +554,7 @@ class AuthRegisterAdapter(AuthPublicAdapter):
)
db.commit()
context.request.state.tx_committed_by_route = True
# 注册成功后清除验证状态（在 commit 后清理，即使清理失败也不影响注册结果）
if require_verification and email:
@@ -565,6 +570,7 @@ class AuthRegisterAdapter(AuthPublicAdapter):
message="注册成功",
).model_dump()
except ValueError as exc:
db.rollback()
AuditService.log_event(
db=db,
event_type=AuditEventType.UNAUTHORIZED_ACCESS,
@@ -574,6 +580,7 @@ class AuthRegisterAdapter(AuthPublicAdapter):
metadata={"username": register_request.username, "error": str(exc)},
)
db.commit()
context.request.state.tx_committed_by_route = True
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc))
@@ -611,6 +618,7 @@ class AuthChangePasswordAdapter(AuthenticatedApiAdapter):
raise InvalidRequestException("密码长度至少6位")
user.set_password(new_password)
context.db.commit()
context.request.state.tx_committed_by_route = True
logger.info(f"用户修改密码: {user.email}")
return {"message": "密码修改成功"}
@@ -643,6 +651,7 @@ class AuthLogoutAdapter(AuthenticatedApiAdapter):
metadata={"user_id": user.id, "email": user.email},
)
context.db.commit()
context.request.state.tx_committed_by_route = True
logger.info(f"用户登出成功: {user.email}")
src/api/dashboard/routes.py
+1 -1
View File
@@ -303,7 +303,7 @@ class AdminDashboardStatsAdapter(AdminApiAdapter):
yesterday_stats = db.query(StatsDaily).filter(StatsDaily.date == yesterday).first()
if yesterday_stats:
requests_yesterday = yesterday_stats.total_requests
cost_yesterday = yesterday_stats.total_cost
cost_yesterday = float(yesterday_stats.total_cost or 0)
input_tokens_yesterday = yesterday_stats.input_tokens
output_tokens_yesterday = yesterday_stats.output_tokens
cache_creation_yesterday = yesterday_stats.cache_creation_tokens
src/api/user_me/routes.py
+8 -1
View File
@@ -502,6 +502,7 @@ class UpdateProfileAdapter(AuthenticatedApiAdapter):
user.updated_at = datetime.now(timezone.utc)
db.commit()
context.request.state.tx_committed_by_route = True
db.refresh(user)
return {"message": "个人信息更新成功"}
@@ -544,6 +545,7 @@ class ChangePasswordAdapter(AuthenticatedApiAdapter):
user.set_password(request.new_password)
user.updated_at = datetime.now(timezone.utc)
db.commit()
context.request.state.tx_committed_by_route = True
action = "修改" if has_password else "设置"
logger.info(f"用户{action}密码: {user.email}")
return {"message": f"密码{action}成功"}
@@ -735,6 +737,7 @@ class DeleteMyApiKeyAdapter(AuthenticatedApiAdapter):
pre_clean_api_key(context.db, api_key.id)
context.db.delete(api_key)
context.db.commit()
context.request.state.tx_committed_by_route = True
return {"message": "API密钥已删除"}
@@ -756,6 +759,7 @@ class ToggleMyApiKeyAdapter(AuthenticatedApiAdapter):
raise ForbiddenException("该密钥已被管理员锁定，无法修改状态")
api_key.is_active = not api_key.is_active
context.db.commit()
context.request.state.tx_committed_by_route = True
context.db.refresh(api_key)
return {
"id": api_key.id,
@@ -1068,7 +1072,7 @@ class GetUsageAdapter(AuthenticatedApiAdapter):
"input_tokens": r.input_tokens,
"output_tokens": r.output_tokens,
"total_tokens": r.total_tokens,
"cost": r.total_cost_usd,
"cost": float(r.total_cost_usd or 0),
"response_time_ms": r.response_time_ms,
"first_byte_time_ms": r.first_byte_time_ms,
"is_stream": r.is_stream,
@@ -1482,6 +1486,7 @@ class UpdateApiKeyProvidersAdapter(AuthenticatedApiAdapter):
)
api_key.updated_at = datetime.now(timezone.utc)
db.commit()
context.request.state.tx_committed_by_route = True
logger.debug(f"用户 {user.id} 更新API密钥 {self.api_key_id} 的可用提供商")
return {"message": "API密钥可用提供商已更新"}
@@ -1531,6 +1536,7 @@ class UpdateApiKeyCapabilitiesAdapter(AuthenticatedApiAdapter):
api_key.force_capabilities = force_capabilities
api_key.updated_at = datetime.now(timezone.utc)
db.commit()
context.request.state.tx_committed_by_route = True
# 记录审计日志
audit_service.log_event(
@@ -1662,6 +1668,7 @@ class UpdateModelCapabilitySettingsAdapter(AuthenticatedApiAdapter):
user.model_capability_settings = settings
user.updated_at = datetime.now(timezone.utc)
db.commit()
context.request.state.tx_committed_by_route = True
# 清除用户缓存，确保下次读取时获取最新数据
await UserCacheService.invalidate_user_cache(user.id, user.email)
src/models/database.py
+45 -41
View File
@@ -197,7 +197,7 @@ class ApiKey(Base):
# 使用统计
total_requests = Column(Integer, default=0)
total_cost_usd = Column(Float, default=0.0)
total_cost_usd = Column(Numeric(20, 8), default=0.0)
# 钱包体系：余额/额度由 wallets 表统一管理
is_standalone = Column(
@@ -364,29 +364,29 @@ class Usage(Base):
cache_creation_input_tokens_1h = Column(Integer, default=0) # 1h TTL 缓存创建
# 成本计算
input_cost_usd = Column(Float, default=0.0)
output_cost_usd = Column(Float, default=0.0)
cache_cost_usd = Column(Float, default=0.0) # 总缓存成本
cache_creation_cost_usd = Column(Float, default=0.0) # 缓存创建成本
cache_read_cost_usd = Column(Float, default=0.0) # 缓存读取成本
request_cost_usd = Column(Float, default=0.0) # 按次计费成本
total_cost_usd = Column(Float, default=0.0)
input_cost_usd = Column(Numeric(20, 8), default=0.0)
output_cost_usd = Column(Numeric(20, 8), default=0.0)
cache_cost_usd = Column(Numeric(20, 8), default=0.0) # 总缓存成本
cache_creation_cost_usd = Column(Numeric(20, 8), default=0.0) # 缓存创建成本
cache_read_cost_usd = Column(Numeric(20, 8), default=0.0) # 缓存读取成本
request_cost_usd = Column(Numeric(20, 8), default=0.0) # 按次计费成本
total_cost_usd = Column(Numeric(20, 8), default=0.0)
# 真实成本计算（表面成本 × 倍率）
actual_input_cost_usd = Column(Float, default=0.0) # 真实输入成本
actual_output_cost_usd = Column(Float, default=0.0) # 真实输出成本
actual_cache_creation_cost_usd = Column(Float, default=0.0) # 真实缓存创建成本
actual_cache_read_cost_usd = Column(Float, default=0.0) # 真实缓存读取成本
actual_request_cost_usd = Column(Float, default=0.0) # 真实按次计费成本
actual_total_cost_usd = Column(Float, default=0.0) # 真实总成本
rate_multiplier = Column(Float, default=1.0) # 使用的倍率（来自 ProviderAPIKey）
actual_input_cost_usd = Column(Numeric(20, 8), default=0.0) # 真实输入成本
actual_output_cost_usd = Column(Numeric(20, 8), default=0.0) # 真实输出成本
actual_cache_creation_cost_usd = Column(Numeric(20, 8), default=0.0) # 真实缓存创建成本
actual_cache_read_cost_usd = Column(Numeric(20, 8), default=0.0) # 真实缓存读取成本
actual_request_cost_usd = Column(Numeric(20, 8), default=0.0) # 真实按次计费成本
actual_total_cost_usd = Column(Numeric(20, 8), default=0.0) # 真实总成本
rate_multiplier = Column(Numeric(10, 6), default=1.0) # 使用的倍率（来自 ProviderAPIKey）
# 历史价格记录（每1M tokens的美元价格，记录请求时的实际价格）
input_price_per_1m = Column(Float, nullable=True) # 输入单价
output_price_per_1m = Column(Float, nullable=True) # 输出单价
cache_creation_price_per_1m = Column(Float, nullable=True) # 缓存创建单价
cache_read_price_per_1m = Column(Float, nullable=True) # 缓存读取单价
price_per_request = Column(Float, nullable=True) # 按次计费单价（历史记录）
input_price_per_1m = Column(Numeric(20, 8), nullable=True) # 输入单价
output_price_per_1m = Column(Numeric(20, 8), nullable=True) # 输出单价
cache_creation_price_per_1m = Column(Numeric(20, 8), nullable=True) # 缓存创建单价
cache_read_price_per_1m = Column(Numeric(20, 8), nullable=True) # 缓存读取单价
price_per_request = Column(Numeric(20, 8), nullable=True) # 按次计费单价（历史记录）
# 请求详情
request_type = Column(String(50)) # chat, completion, embedding等
@@ -960,8 +960,8 @@ class Provider(ExportMixin, Base):
)
# 月卡配置
monthly_quota_usd = Column(Float, nullable=True) # 月卡总额度
monthly_used_usd = Column(Float, default=0.0) # 本月已用额度
monthly_quota_usd = Column(Numeric(20, 8), nullable=True) # 月卡总额度
monthly_used_usd = Column(Numeric(20, 8), default=0.0) # 本月已用额度
quota_reset_day = Column(Integer, default=30) # 额度重置周期(天数)，例如：7=每周，30=每月
quota_last_reset_at = Column(DateTime(timezone=True), nullable=True) # 上次额度重置时间
quota_expires_at = Column(DateTime(timezone=True), nullable=True) # 月卡过期时间
@@ -1256,7 +1256,9 @@ class GlobalModel(ExportMixin, Base):
display_name = Column(String(100), nullable=False)
# 按次计费配置（每次请求的固定费用，美元）- 可选，与按 token 计费叠加
default_price_per_request = Column(Float, nullable=True, default=None) # 每次请求固定费用
default_price_per_request = Column(
Numeric(20, 8), nullable=True, default=None
) # 每次请求固定费用
# 统一阶梯计费配置（JSON格式）- 必填
# 固定价格也用单阶梯表示: {"tiers": [{"up_to": null, "input_price_per_1m": X, ...}]}
@@ -1364,7 +1366,7 @@ class Model(ExportMixin, Base):
provider_model_mappings = Column(JSON, nullable=True, default=None)
# 按次计费配置（每次请求的固定费用，美元）- 可为空，为空时使用 GlobalModel 的默认值
price_per_request = Column(Float, nullable=True) # 每次请求固定费用
price_per_request = Column(Numeric(20, 8), nullable=True) # 每次请求固定费用
# 阶梯计费配置（JSON格式）- 可为空，为空时使用 GlobalModel 的默认值
tiered_pricing = Column(JSON, nullable=True, default=None)
@@ -1909,6 +1911,8 @@ class ProviderAPIKey(ExportMixin, Base):
nullable=False,
)
__table_args__ = (Index("idx_provider_api_keys_provider_active", "provider_id", "is_active"),)
# 关系
provider = relationship("Provider", back_populates="api_keys")
@@ -2505,8 +2509,8 @@ class StatsHourly(Base):
cache_read_tokens = Column(BigInteger, default=0, nullable=False)
# 成本统计 (USD)
total_cost = Column(Float, default=0.0, nullable=False)
actual_total_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
actual_total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
# 性能统计
avg_response_time_ms = Column(Float, default=0.0, nullable=False)
@@ -2543,7 +2547,7 @@ class StatsHourlyUser(Base):
error_requests = Column(Integer, default=0, nullable=False)
input_tokens = Column(BigInteger, default=0, nullable=False)
output_tokens = Column(BigInteger, default=0, nullable=False)
total_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
created_at = Column(
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False
@@ -2574,7 +2578,7 @@ class StatsHourlyModel(Base):
total_requests = Column(Integer, default=0, nullable=False)
input_tokens = Column(BigInteger, default=0, nullable=False)
output_tokens = Column(BigInteger, default=0, nullable=False)
total_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
avg_response_time_ms = Column(Float, default=0.0, nullable=False)
created_at = Column(
@@ -2606,7 +2610,7 @@ class StatsHourlyProvider(Base):
total_requests = Column(Integer, default=0, nullable=False)
input_tokens = Column(BigInteger, default=0, nullable=False)
output_tokens = Column(BigInteger, default=0, nullable=False)
total_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
created_at = Column(
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False
@@ -2646,12 +2650,12 @@ class StatsDaily(Base):
cache_read_tokens = Column(BigInteger, default=0, nullable=False)
# 成本统计 (USD)
total_cost = Column(Float, default=0.0, nullable=False)
actual_total_cost = Column(Float, default=0.0, nullable=False) # 倍率后成本
input_cost = Column(Float, default=0.0, nullable=False)
output_cost = Column(Float, default=0.0, nullable=False)
cache_creation_cost = Column(Float, default=0.0, nullable=False)
cache_read_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
actual_total_cost = Column(Numeric(20, 8), default=0.0, nullable=False) # 倍率后成本
input_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
output_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
cache_creation_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
cache_read_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
# 性能统计
avg_response_time_ms = Column(Float, default=0.0, nullable=False)
@@ -2706,7 +2710,7 @@ class StatsDailyModel(Base):
cache_read_tokens = Column(BigInteger, default=0, nullable=False)
# 成本统计 (USD)
total_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
# 性能统计
avg_response_time_ms = Column(Float, default=0.0, nullable=False)
@@ -2753,7 +2757,7 @@ class StatsDailyProvider(Base):
cache_read_tokens = Column(BigInteger, default=0, nullable=False)
# 成本统计 (USD)
total_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
# 时间戳
created_at = Column(
@@ -2795,7 +2799,7 @@ class StatsDailyApiKey(Base):
cache_creation_tokens = Column(BigInteger, default=0, nullable=False)
cache_read_tokens = Column(BigInteger, default=0, nullable=False)
total_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
created_at = Column(
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False
@@ -2875,8 +2879,8 @@ class StatsSummary(Base):
all_time_cache_read_tokens = Column(BigInteger, default=0, nullable=False)
# 累计成本统计 (USD)
all_time_cost = Column(Float, default=0.0, nullable=False)
all_time_actual_cost = Column(Float, default=0.0, nullable=False)
all_time_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
all_time_actual_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
# 累计用户/API Key 统计 (快照)
total_users = Column(Integer, default=0, nullable=False)
@@ -2922,7 +2926,7 @@ class StatsUserDaily(Base):
cache_read_tokens = Column(BigInteger, default=0, nullable=False)
# 成本统计 (USD)
total_cost = Column(Float, default=0.0, nullable=False)
total_cost = Column(Numeric(20, 8), default=0.0, nullable=False)
# 时间戳
created_at = Column(
src/services/health/endpoint.py
+12 -2
View File
@@ -16,7 +16,7 @@ from datetime import datetime, timedelta, timezone
from typing import Any
from sqlalchemy import case, func
from sqlalchemy.orm import Session
from sqlalchemy.orm import Session, load_only
from src.core.logger import logger
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint, RequestCandidate
@@ -75,11 +75,21 @@ class EndpointHealthService:
# 收集所有 provider_ids
all_provider_ids = list({ep.provider_id for ep in endpoints})
# 批量查询所有密钥（通过 provider_id 关联）
# 批量查询所有密钥（通过 provider_id 关联，只加载分组和统计所需列）
all_keys = (
(
db.query(ProviderAPIKey)
.filter(ProviderAPIKey.provider_id.in_(all_provider_ids))
.options(
load_only(
ProviderAPIKey.id,
ProviderAPIKey.provider_id,
ProviderAPIKey.is_active,
ProviderAPIKey.api_formats,
ProviderAPIKey.health_by_format,
ProviderAPIKey.circuit_breaker_by_format,
)
)
.all()
)
if all_provider_ids
src/services/health/monitor.py
+13 -7
View File
@@ -810,16 +810,22 @@ class HealthMonitor:
),
).first()
# 统计 Key（需要遍历 JSON 字段计算熔断状态）
keys = db.query(ProviderAPIKey).all()
total_keys = len(keys)
active_keys = sum(1 for k in keys if k.is_active)
# 统计 Key（只加载必要列，避免全字段全表扫描）
key_rows = db.query(
ProviderAPIKey.is_active,
ProviderAPIKey.health_by_format,
ProviderAPIKey.circuit_breaker_by_format,
).all()
total_keys = len(key_rows)
active_keys = 0
unhealthy_keys = 0
circuit_open_keys = 0
for key in keys:
health_by_format = key.health_by_format or {}
circuit_by_format = key.circuit_breaker_by_format or {}
for is_active, health_by_format, circuit_by_format in key_rows:
if is_active:
active_keys += 1
health_by_format = health_by_format or {}
circuit_by_format = circuit_by_format or {}
# 检查是否有任何格式健康度低于 0.5
for fmt, health_data in health_by_format.items():
src/services/model/global_model.py
+8 -12
View File
@@ -8,6 +8,8 @@ from __future__ import annotations
from typing import cast
from sqlalchemy import delete as sa_delete
from sqlalchemy import func
from sqlalchemy.orm import Session, joinedload, load_only
from src.core.exceptions import InvalidRequestException, NotFoundException
@@ -199,21 +201,15 @@ class GlobalModelService:
"""
global_model = GlobalModelService.get_global_model(db, global_model_id)
# 查找所有关联的 Model（使用 global_model.id，预加载 provider 关联）
associated_models = (
db.query(Model)
.options(joinedload(Model.provider))
.filter(Model.global_model_id == global_model.id)
.all()
# 批量删除所有关联的 Provider 模型实现
assoc_count = (
db.query(func.count(Model.id)).filter(Model.global_model_id == global_model.id).scalar()
)
# 级联删除所有关联的 Provider 模型实现
if associated_models:
if assoc_count:
logger.info(
f"删除 GlobalModel {global_model.name} 的 {len(associated_models)} 个关联 Provider 模型"
f"删除 GlobalModel {global_model.name} 的 {assoc_count} 个关联 Provider 模型"
)
for model in associated_models:
db.delete(model)
db.execute(sa_delete(Model).where(Model.global_model_id == global_model.id))
# 删除 GlobalModel
db.delete(global_model)
src/services/provider_keys/key_command_service.py
+4 -4
View File
@@ -11,6 +11,7 @@ from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any
from sqlalchemy import delete as sa_delete
from sqlalchemy.orm import Session
from src.core.crypto import crypto_service
@@ -514,13 +515,12 @@ async def batch_delete_endpoint_keys_response(db: Session, key_ids: list[str]) -
# 收集受影响的 provider_id
affected_provider_ids = {key.provider_id for key in keys if key.provider_id}
# 批量删除，一次提交
# 批量 SQL DELETE，一次提交
success_count = 0
try:
for key in keys:
db.delete(key)
db.execute(sa_delete(ProviderAPIKey).where(ProviderAPIKey.id.in_(list(found_ids))))
db.commit()
success_count = len(keys)
success_count = len(found_ids)
except Exception as exc:
db.rollback()
logger.error("批量删除 Key 提交失败: {}", exc)
src/services/provider_keys/pool_quota_probe_scheduler.py
+19 -3
View File
@@ -254,6 +254,9 @@ class PoolQuotaProbeScheduler:
db = create_session()
try:
providers = db.query(Provider).filter(Provider.is_active == True).all() # noqa: E712
# 先筛选出符合条件的 provider，收集其 ID 和配置
eligible_providers: list[tuple[str, str, int]] = [] # (id, type, interval_seconds)
for provider in providers:
provider_id = str(getattr(provider, "id", "") or "")
provider_type = normalize_provider_type(getattr(provider, "provider_type", ""))
@@ -267,16 +270,29 @@ class PoolQuotaProbeScheduler:
interval_minutes = _normalize_probe_interval_minutes(
pool_cfg.probing_interval_minutes
)
interval_seconds = interval_minutes * 60
eligible_providers.append((provider_id, provider_type, interval_minutes * 60))
keys = (
# 批量查询所有符合条件的 provider 的活跃 keys，避免 N+1
eligible_ids = [p[0] for p in eligible_providers]
all_keys: list[ProviderAPIKey] = []
if eligible_ids:
all_keys = (
db.query(ProviderAPIKey)
.filter(
ProviderAPIKey.provider_id == provider_id,
ProviderAPIKey.provider_id.in_(eligible_ids),
ProviderAPIKey.is_active == True, # noqa: E712
)
.all()
)
# 按 provider_id 分组
keys_by_provider: dict[str, list[ProviderAPIKey]] = {}
for key in all_keys:
pid = str(key.provider_id)
keys_by_provider.setdefault(pid, []).append(key)
for provider_id, provider_type, interval_seconds in eligible_providers:
keys = keys_by_provider.get(provider_id, [])
if not keys:
continue
src/services/scheduling/candidate_builder.py
+13 -1
View File
@@ -97,7 +97,19 @@ class CandidateBuilder:
db.query(Provider)
.options(
# 预加载 Provider 级别的 api_keys
selectinload(Provider.api_keys),
# defer 排除仅后台管理/模型获取用的冷字段，热路径字段全部加载
selectinload(Provider.api_keys).defer(
ProviderAPIKey.note,
ProviderAPIKey.last_error_msg,
ProviderAPIKey.auto_fetch_models,
ProviderAPIKey.locked_models,
ProviderAPIKey.model_include_patterns,
ProviderAPIKey.model_exclude_patterns,
ProviderAPIKey.last_models_fetch_at,
ProviderAPIKey.last_models_fetch_error,
ProviderAPIKey.max_probe_interval_minutes,
ProviderAPIKey.expires_at,
),
# 预加载 endpoints（用于按 api_format 选择请求配置）
selectinload(Provider.endpoints),
# 同时加载 models 和 global_model 关系
src/services/system/maintenance_scheduler.py
+30 -27
View File
@@ -18,10 +18,10 @@
from __future__ import annotations
import asyncio
from datetime import datetime, timedelta, timezone
from datetime import date, datetime, timedelta, timezone
from typing import Any
from sqlalchemy import delete, text
from sqlalchemy import delete, literal_column, text
from src.core.logger import logger
from src.database import create_session
@@ -393,40 +393,43 @@ class MaintenanceScheduler:
check_start_date, datetime.min.time(), tzinfo=timezone.utc
)
# 获取 StatsDaily 和 StatsDailyModel 中已有数据的日期集合
existing_daily_dates = set()
existing_model_dates = set()
existing_provider_dates = set()
# 单次查询获取三张统计表中已有数据的日期（UNION ALL 合并）
existing_daily_dates: set[date] = set()
existing_model_dates: set[date] = set()
existing_provider_dates: set[date] = set()
daily_stats = (
db.query(StatsDaily.date).filter(StatsDaily.date >= check_start_dt).all()
)
for (stat_date,) in daily_stats:
if stat_date.tzinfo is None:
stat_date = stat_date.replace(tzinfo=timezone.utc)
existing_daily_dates.add(stat_date.date())
model_stats = (
db.query(StatsDailyModel.date)
q_daily = db.query(
StatsDaily.date.label("dt"),
literal_column("'daily'").label("src"),
).filter(StatsDaily.date >= check_start_dt)
q_model = (
db.query(
StatsDailyModel.date.label("dt"),
literal_column("'model'").label("src"),
)
.filter(StatsDailyModel.date >= check_start_dt)
.distinct()
.all()
)
for (stat_date,) in model_stats:
if stat_date.tzinfo is None:
stat_date = stat_date.replace(tzinfo=timezone.utc)
existing_model_dates.add(stat_date.date())
provider_stats = (
db.query(StatsDailyProvider.date)
q_provider = (
db.query(
StatsDailyProvider.date.label("dt"),
literal_column("'provider'").label("src"),
)
.filter(StatsDailyProvider.date >= check_start_dt)
.distinct()
.all()
)
for (stat_date,) in provider_stats:
combined = q_daily.union_all(q_model).union_all(q_provider).all()
for stat_date, src in combined:
if stat_date.tzinfo is None:
stat_date = stat_date.replace(tzinfo=timezone.utc)
existing_provider_dates.add(stat_date.date())
d = stat_date.date()
if src == "daily":
existing_daily_dates.add(d)
elif src == "model":
existing_model_dates.add(d)
else:
existing_provider_dates.add(d)
# 找出需要回填的日期
all_dates = set()
src/services/system/stats_aggregator.py
+4 -4
View File
@@ -1339,10 +1339,10 @@ def query_stats_hybrid(
result.output_tokens += stats.output_tokens
result.cache_creation_tokens += stats.cache_creation_tokens
result.cache_read_tokens += stats.cache_read_tokens
result.cache_creation_cost += stats.cache_creation_cost
result.cache_read_cost += stats.cache_read_cost
result.total_cost += stats.total_cost
result.actual_total_cost += stats.actual_total_cost
result.cache_creation_cost += float(stats.cache_creation_cost or 0)
result.cache_read_cost += float(stats.cache_read_cost or 0)
result.total_cost += float(stats.total_cost or 0)
result.actual_total_cost += float(stats.actual_total_cost or 0)
result.total_response_time_ms += (stats.avg_response_time_ms or 0.0) * stats.total_requests
# Realtime per day
src/services/system/sync_stats.py
+3 -2
View File
@@ -37,7 +37,8 @@ class SyncStatsService:
try:
# 获取要同步的API密钥（使用分页避免大数据量问题）
if api_key_id:
api_keys = db.query(ApiKey).filter(ApiKey.id == api_key_id).all()
single_key = db.query(ApiKey).filter(ApiKey.id == api_key_id).first()
api_keys = [single_key] if single_key else []
else:
# 分页处理，避免一次加载所有数据
offset = 0
@@ -106,7 +107,7 @@ class SyncStatsService:
api_key.total_requests = actual_requests
needs_update = True
if abs(api_key.total_cost_usd - actual_cost) > 0.0001:
if abs(float(api_key.total_cost_usd or 0) - actual_cost) > 0.0001:
logger.info(
f"API密钥 {api_key.id} 费用不一致: {api_key.total_cost_usd} -> {actual_cost}"
)
src/services/usage/stream.py
+1 -1
View File
@@ -800,7 +800,7 @@ class StreamUsageTracker:
if usage_record:
try:
# 在 usage_record 仍在会话中时，立即获取所需属性
total_cost = usage_record.total_cost_usd or 0.0
total_cost = float(usage_record.total_cost_usd or 0)
except Exception as e:
logger.warning(f"Failed to access total_cost_usd from usage_record: {e}")
total_cost = 0.0
src/services/user/service.py
+3 -2
View File
@@ -9,7 +9,7 @@ from datetime import datetime, timezone
from typing import Any
from sqlalchemy import and_, func, or_
from sqlalchemy.orm import Session
from sqlalchemy.orm import Session, contains_eager
from src.core.logger import logger
from src.core.validators import EmailValidator, PasswordValidator, UsernameValidator
@@ -481,10 +481,11 @@ class UserService:
if not all_active_provider_ids:
return []
# 查询所有活跃的 Model（关联 GlobalModel）
# 查询所有活跃的 Model（关联 GlobalModel，contains_eager 避免循环中懒加载）
all_models = (
db.query(Model)
.join(GlobalModel, Model.global_model_id == GlobalModel.id)
.options(contains_eager(Model.global_model))
.filter(
and_(
Model.provider_id.in_(all_active_provider_ids),
src/services/wallet/service.py
+26 -20
View File
@@ -338,20 +338,32 @@ class WalletService:
return cls.get_spendable_balance_value(wallet)
@classmethod
def _resolve_wallet_for_usage(cls, db: Session, usage: Usage) -> Wallet | None:
def _resolve_wallet_for_usage(
cls, db: Session, usage: Usage, *, for_update: bool = False
) -> Wallet | None:
"""解析 Usage 对应的钱包。for_update=True 时直接返回加锁的钱包，避免二次查询。"""
wallet_id = None
if usage.wallet_id:
wallet = db.query(Wallet).filter(Wallet.id == usage.wallet_id).first()
if wallet:
return wallet
api_key = None
if usage.api_key_id:
api_key = db.query(ApiKey).filter(ApiKey.id == usage.api_key_id).first()
if api_key and api_key.is_standalone:
return cls.get_or_create_wallet(db, api_key=api_key)
if usage.user_id:
user = db.query(User).filter(User.id == usage.user_id).first()
return cls.get_or_create_wallet(db, user=user, api_key=api_key)
return None
wallet_id = usage.wallet_id
else:
api_key = None
if usage.api_key_id:
api_key = db.query(ApiKey).filter(ApiKey.id == usage.api_key_id).first()
if api_key and api_key.is_standalone:
wallet = cls.get_or_create_wallet(db, api_key=api_key)
wallet_id = wallet.id if wallet else None
if wallet_id is None and usage.user_id:
user = db.query(User).filter(User.id == usage.user_id).first()
wallet = cls.get_or_create_wallet(db, user=user, api_key=api_key)
wallet_id = wallet.id if wallet else None
if wallet_id is None:
return None
query = db.query(Wallet).filter(Wallet.id == wallet_id)
if for_update:
query = query.with_for_update()
return query.first()
@classmethod
def apply_usage_charge(
@@ -365,13 +377,7 @@ class WalletService:
if amount <= Decimal("0"):
return None, None
wallet = cls._resolve_wallet_for_usage(db, usage)
if wallet is None:
return None, None
locked_wallet = (
db.query(Wallet).filter(Wallet.id == wallet.id).with_for_update().one_or_none()
)
locked_wallet = cls._resolve_wallet_for_usage(db, usage, for_update=True)
if locked_wallet is None:
return None, None
tests/api/test_admin_api_key_scope_routes.py
+14 -5
View File
@@ -11,13 +11,13 @@ from fastapi.testclient import TestClient
from src.api.admin.api_keys.routes import (
AdminGetFullKeyAdapter,
AdminToggleApiKeyAdapter,
router as admin_api_keys_router,
)
from src.api.admin.api_keys.routes import router as admin_api_keys_router
from src.api.admin.users.routes import (
AdminGetUserKeyFullKeyAdapter,
AdminToggleUserKeyLockAdapter,
router as admin_users_router,
)
from src.api.admin.users.routes import router as admin_users_router
from src.core.exceptions import InvalidRequestException, NotFoundException
from src.database import get_db
@@ -25,6 +25,7 @@ from src.database import get_db
def _build_context(db: MagicMock) -> SimpleNamespace:
return SimpleNamespace(
db=db,
request=SimpleNamespace(state=SimpleNamespace()),
add_audit_metadata=lambda **_: None,
)
@@ -38,10 +39,13 @@ def _build_admin_users_app(db: MagicMock, monkeypatch: pytest.MonkeyPatch) -> Te
app.include_router(admin_users_router)
app.dependency_overrides[get_db] = lambda: db
async def _fake_pipeline_run(*, adapter: object, http_request: object, db: MagicMock, mode: object) -> object:
async def _fake_pipeline_run(
*, adapter: object, http_request: object, db: MagicMock, mode: object
) -> object:
_ = http_request, mode
context = SimpleNamespace(
db=db,
request=SimpleNamespace(state=SimpleNamespace()),
user=SimpleNamespace(id="admin-1"),
ensure_json_body=lambda: {},
add_audit_metadata=lambda **_: None,
@@ -57,10 +61,13 @@ def _build_admin_api_keys_app(db: MagicMock, monkeypatch: pytest.MonkeyPatch) ->
app.include_router(admin_api_keys_router)
app.dependency_overrides[get_db] = lambda: db
async def _fake_pipeline_run(*, adapter: object, http_request: object, db: MagicMock, mode: object) -> object:
async def _fake_pipeline_run(
*, adapter: object, http_request: object, db: MagicMock, mode: object
) -> object:
_ = http_request, mode
context = SimpleNamespace(
db=db,
request=SimpleNamespace(state=SimpleNamespace()),
user=SimpleNamespace(id="admin-1"),
ensure_json_body=lambda: {},
add_audit_metadata=lambda **_: None,
@@ -243,7 +250,9 @@ def test_standalone_list_route_does_not_expose_is_locked(monkeypatch: pytest.Mon
monkeypatch.setattr(
"src.api.admin.api_keys.routes.WalletService.get_wallet",
lambda _db, user_id=None, api_key_id=None, user=None, api_key=None: SimpleNamespace(id="w-1"),
lambda _db, user_id=None, api_key_id=None, user=None, api_key=None: SimpleNamespace(
id="w-1"
),
)
client = _build_admin_api_keys_app(db, monkeypatch)