Add SmartUp account pool parameter configuration

This commit is contained in:
SmartUp Developer
2026-07-04 11:59:21 +08:00
parent 0cce5d53fe
commit 5a20e792fe
5 changed files with 229 additions and 21 deletions
+41 -11
View File
@@ -1151,7 +1151,7 @@ def set_website_accounts_concurrency(
db: Session = Depends(get_db),
_=Depends(get_current_user),
):
"""批量设置当前网站已导入账号的并发数"""
"""批量设置当前网站已导入账号参数。"""
website = db.query(Website).filter(Website.id == wid).first()
if not website:
raise HTTPException(404, "website not found")
@@ -1204,6 +1204,16 @@ def set_website_accounts_concurrency(
bulk_candidates = []
skipped_items = []
retry_status_codes = sorted({
int(code)
for code in body.pool_mode_retry_status_codes
if isinstance(code, int) and 100 <= code <= 599
})
pool_credentials = {
"pool_mode": body.pool_mode,
"pool_mode_retry_count": body.pool_mode_retry_count,
"pool_mode_retry_status_codes": retry_status_codes if body.pool_mode else [],
}
for aid, cand in candidates.items():
remote_acc = remote_map.get(aid)
@@ -1220,13 +1230,19 @@ def set_website_accounts_concurrency(
acc_name = remote_acc.get("name")
old_con = remote_acc.get("concurrency")
remote_credentials = remote_acc.get("credentials")
if not isinstance(remote_credentials, dict):
remote_credentials = {}
old_pool_mode = remote_credentials.get("pool_mode")
try:
int(aid)
bulk_candidates.append({
"account_id": aid,
"account_name": acc_name,
"old_concurrency": old_con
"old_concurrency": old_con,
"old_pool_mode": old_pool_mode,
"credentials": remote_credentials,
})
except ValueError:
skipped_items.append(SetConcurrencyItem(
@@ -1234,6 +1250,8 @@ def set_website_accounts_concurrency(
account_name=acc_name,
old_concurrency=old_con,
target_concurrency=body.concurrency,
old_pool_mode=old_pool_mode if isinstance(old_pool_mode, bool) else None,
target_pool_mode=body.pool_mode if body.configure_pool_mode else None,
status="skipped",
message="账号 ID 无法转换为数字,跳过批量更新"
))
@@ -1253,12 +1271,13 @@ def set_website_accounts_concurrency(
bulk_success = False
bulk_error_msg = ""
try:
c.bulk_update_accounts(bulk_ids, {"concurrency": body.concurrency})
bulk_success = True
except Exception as e:
bulk_error_msg = str(e)
logger.warning("bulk_update_accounts failed, falling back to individual updates: %s", e)
if not body.configure_pool_mode:
try:
c.bulk_update_accounts(bulk_ids, {"concurrency": body.concurrency})
bulk_success = True
except Exception as e:
bulk_error_msg = str(e)
logger.warning("bulk_update_accounts failed, falling back to individual updates: %s", e)
if bulk_success:
for item in bulk_candidates:
@@ -1274,14 +1293,22 @@ def set_website_accounts_concurrency(
for item in bulk_candidates:
aid = item["account_id"]
try:
c.update_account(aid, {"concurrency": body.concurrency})
update_payload = {"concurrency": body.concurrency}
if body.configure_pool_mode:
update_payload["credentials"] = {
**item["credentials"],
**pool_credentials,
}
c.update_account(aid, update_payload)
items.append(SetConcurrencyItem(
account_id=aid,
account_name=item["account_name"],
old_concurrency=item["old_concurrency"],
target_concurrency=body.concurrency,
old_pool_mode=item["old_pool_mode"] if isinstance(item["old_pool_mode"], bool) else None,
target_pool_mode=body.pool_mode if body.configure_pool_mode else None,
status="success",
message="成功 (逐个更新)"
message="成功 (逐个更新)" if not body.configure_pool_mode else "成功"
))
except Exception as e:
items.append(SetConcurrencyItem(
@@ -1289,6 +1316,8 @@ def set_website_accounts_concurrency(
account_name=item["account_name"],
old_concurrency=item["old_concurrency"],
target_concurrency=body.concurrency,
old_pool_mode=item["old_pool_mode"] if isinstance(item["old_pool_mode"], bool) else None,
target_pool_mode=body.pool_mode if body.configure_pool_mode else None,
status="failed",
message=f"更新失败: {e}"
))
@@ -1306,7 +1335,8 @@ def set_website_accounts_concurrency(
if skip_count:
msg_parts.append(f"跳过 {skip_count}")
message = "设置账号并发执行完毕:" + "".join(msg_parts)
action_name = "设置账号参数" if body.configure_pool_mode else "设置账号并发"
message = f"{action_name}执行完毕:" + "".join(msg_parts)
return SetConcurrencyResponse(
success=success,
+6
View File
@@ -272,6 +272,10 @@ class CleanupInvalidAccountsExecuteResponse(BaseModel):
class SetConcurrencyRequest(BaseModel):
concurrency: int = Field(default=100, ge=1, le=1000)
configure_pool_mode: bool = False
pool_mode: bool = True
pool_mode_retry_count: int = Field(default=1, ge=0, le=10)
pool_mode_retry_status_codes: list[int] = Field(default_factory=lambda: [429, 502, 503, 529])
class SetConcurrencyItem(BaseModel):
@@ -279,6 +283,8 @@ class SetConcurrencyItem(BaseModel):
account_name: Optional[str] = None
old_concurrency: Optional[int] = None
target_concurrency: int
old_pool_mode: Optional[bool] = None
target_pool_mode: Optional[bool] = None
status: str # "success" | "skipped" | "failed"
message: str
+88
View File
@@ -233,6 +233,94 @@ def test_set_website_accounts_concurrency_fallback(db_session, monkeypatch):
assert closed_count == 1
def test_set_website_accounts_concurrency_pool_mode_merges_credentials(db_session, monkeypatch):
w = Website(
id=1,
name="Sub2Api site",
site_type="sub2api",
base_url="http://sub2api",
api_prefix="api/v1",
auth_type="bearer",
auth_config_json='{"token": "tok1"}'
)
db_session.add(w)
up = Upstream(id=10, name="Upstream A", base_url="http://upstream-a")
db_session.add(up)
db_session.commit()
k1 = UpstreamGeneratedKey(
id=101,
upstream_id=up.id,
group_id="g1",
key_name="key-101",
key_value="val-101",
status="active",
imported_website_id=w.id,
imported_account_id="1001",
)
db_session.add(k1)
db_session.commit()
bulk_calls = []
update_calls = []
class MockClient:
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
return False
def list_accounts(self):
return [{
"id": 1001,
"name": "acc-1001",
"concurrency": 10,
"credentials": {
"base_url": "https://relay.example.com",
"api_key": "sk-old",
"pool_mode": False,
},
}]
def extract_id(self, val):
return str(val.get("id"))
def bulk_update_accounts(self, account_ids, body):
bulk_calls.append((account_ids, body))
def update_account(self, account_id, body):
update_calls.append((account_id, body))
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClient())
req = SetConcurrencyRequest(
concurrency=80,
configure_pool_mode=True,
pool_mode=True,
pool_mode_retry_count=1,
pool_mode_retry_status_codes=[529, 429, 999, 502, 429],
)
res = set_website_accounts_concurrency(wid=w.id, body=req, db=db_session)
assert res.success is True
assert "设置账号参数执行完毕" in res.message
assert bulk_calls == []
assert update_calls == [("1001", {
"concurrency": 80,
"credentials": {
"base_url": "https://relay.example.com",
"api_key": "sk-old",
"pool_mode": True,
"pool_mode_retry_count": 1,
"pool_mode_retry_status_codes": [429, 502, 529],
},
})]
assert res.items[0].old_pool_mode is False
assert res.items[0].target_pool_mode is True
def test_set_website_accounts_concurrency_list_accounts_none(db_session, monkeypatch):
w = Website(
id=1,