mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
perf: 依赖数据库 CASCADE/SET NULL 替代手动清理关联表，缩短删除事务
- 批量删除移除 cleanup_key_references 手动清理，改为依赖 FK CASCADE/SET NULL - video_tasks.key_id FK 增加 ondelete="SET NULL"，附带幂等迁移脚本 - _sync_delete 增加 statement_timeout 和任务级超时保护 - 批量导入在每次 await 前释放闲置 DB 连接，按批次提交写入避免长事务 - 前端轮询改为先查后等，首次查询不再多等一个间隔
This commit is contained in:
fawney19
210
tests/api/test_provider_oauth_batch_import.py
Normal file
210
tests/api/test_provider_oauth_batch_import.py
Normal file
@@ -0,0 +1,210 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from itertools import count
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from src.api.admin import provider_oauth as oauthmod
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_standard_batch_import_releases_db_connection_before_network_await(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
release_calls: list[str] = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_require_oauth_template",
|
||||
lambda _provider_type: SimpleNamespace(
|
||||
oauth=SimpleNamespace(
|
||||
token_url="https://example.com/oauth/token",
|
||||
client_id="client-id",
|
||||
client_secret=None,
|
||||
scopes=[],
|
||||
)
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_parse_standard_oauth_import_entries",
|
||||
lambda _raw: [{"refresh_token": "r" * 120}],
|
||||
)
|
||||
monkeypatch.setattr(oauthmod, "_get_provider_api_formats", lambda _provider: [])
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_release_batch_import_db_connection_before_await",
|
||||
lambda _db: release_calls.append("release"),
|
||||
)
|
||||
|
||||
async def _fake_post_oauth_token(**_kwargs: object) -> httpx.Response:
|
||||
raise RuntimeError("upstream unavailable")
|
||||
|
||||
monkeypatch.setattr(oauthmod, "post_oauth_token", _fake_post_oauth_token)
|
||||
|
||||
db = MagicMock()
|
||||
|
||||
result = await oauthmod._batch_import_standard_oauth_internal(
|
||||
provider_id="provider-1",
|
||||
provider_type="codex",
|
||||
provider=SimpleNamespace(endpoints=[]), # type: ignore[arg-type]
|
||||
raw_credentials="ignored",
|
||||
db=db,
|
||||
concurrency=1,
|
||||
)
|
||||
|
||||
assert result.total == 1
|
||||
assert result.success == 0
|
||||
assert result.failed == 1
|
||||
assert release_calls
|
||||
db.commit.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_standard_batch_import_commits_successes_in_chunks(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
key_ids = count(1)
|
||||
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_PROVIDER_OAUTH_BATCH_IMPORT_COMMIT_BATCH_SIZE",
|
||||
2,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_require_oauth_template",
|
||||
lambda _provider_type: SimpleNamespace(
|
||||
oauth=SimpleNamespace(
|
||||
token_url="https://example.com/oauth/token",
|
||||
client_id="client-id",
|
||||
client_secret=None,
|
||||
scopes=[],
|
||||
)
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_parse_standard_oauth_import_entries",
|
||||
lambda _raw: [{"refresh_token": f"r-{idx}" + ("x" * 120)} for idx in range(3)],
|
||||
)
|
||||
monkeypatch.setattr(oauthmod, "_get_provider_api_formats", lambda _provider: ["responses"])
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_release_batch_import_db_connection_before_await",
|
||||
lambda _db: None,
|
||||
)
|
||||
|
||||
async def _fake_post_oauth_token(**_kwargs: object) -> httpx.Response:
|
||||
idx = next(key_ids)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": f"access-{idx}",
|
||||
"refresh_token": f"refresh-{idx}",
|
||||
"expires_in": 3600,
|
||||
},
|
||||
request=httpx.Request("POST", "https://example.com/oauth/token"),
|
||||
)
|
||||
|
||||
async def _fake_enrich_auth_config(**kwargs: object) -> dict[str, object]:
|
||||
auth_config = dict(kwargs["auth_config"]) # type: ignore[call-overload]
|
||||
auth_config["email"] = f"user-{next(key_ids)}@example.com"
|
||||
return auth_config
|
||||
|
||||
created_ids = count(1)
|
||||
monkeypatch.setattr(oauthmod, "post_oauth_token", _fake_post_oauth_token)
|
||||
monkeypatch.setattr(oauthmod, "enrich_auth_config", _fake_enrich_auth_config)
|
||||
monkeypatch.setattr(oauthmod, "_check_duplicate_oauth_account", lambda *_args, **_kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_create_oauth_key",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(id=f"key-{next(created_ids)}"),
|
||||
)
|
||||
|
||||
db = MagicMock()
|
||||
|
||||
result = await oauthmod._batch_import_standard_oauth_internal(
|
||||
provider_id="provider-1",
|
||||
provider_type="example",
|
||||
provider=SimpleNamespace(endpoints=[]), # type: ignore[arg-type]
|
||||
raw_credentials="ignored",
|
||||
db=db,
|
||||
concurrency=1,
|
||||
)
|
||||
|
||||
assert result.total == 3
|
||||
assert result.success == 3
|
||||
assert result.failed == 0
|
||||
assert db.commit.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kiro_batch_import_releases_db_connection_before_refresh(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
release_calls: list[str] = []
|
||||
|
||||
class FakeKiroAuthConfig:
|
||||
def __init__(self, data: dict[str, object]) -> None:
|
||||
self._data = dict(data)
|
||||
self.provider_type = str(data.get("provider_type") or "")
|
||||
self.email = data.get("email") if isinstance(data.get("email"), str) else None
|
||||
self.auth_method = (
|
||||
data.get("auth_method") if isinstance(data.get("auth_method"), str) else "social"
|
||||
)
|
||||
self.refresh_token = str(data.get("refresh_token") or "")
|
||||
|
||||
@staticmethod
|
||||
def validate_required_fields(_cred: dict[str, object]) -> tuple[bool, str | None]:
|
||||
return True, None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, object]) -> "FakeKiroAuthConfig":
|
||||
return cls(data)
|
||||
|
||||
def to_dict(self) -> dict[str, object]:
|
||||
return dict(self._data)
|
||||
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_parse_kiro_import_input",
|
||||
lambda _raw: [{"refresh_token": "r" * 120, "auth_method": "social"}],
|
||||
)
|
||||
monkeypatch.setattr(oauthmod, "_get_provider_api_formats", lambda _provider: [])
|
||||
monkeypatch.setattr(
|
||||
oauthmod,
|
||||
"_release_batch_import_db_connection_before_await",
|
||||
lambda _db: release_calls.append("release"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.provider.adapters.kiro.models.credentials.KiroAuthConfig",
|
||||
FakeKiroAuthConfig,
|
||||
)
|
||||
|
||||
async def _fake_refresh_access_token(*_args: object, **_kwargs: object) -> tuple[str, object]:
|
||||
raise RuntimeError("refresh token reused")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.services.provider.adapters.kiro.token_manager.refresh_access_token",
|
||||
_fake_refresh_access_token,
|
||||
)
|
||||
|
||||
db = MagicMock()
|
||||
|
||||
result = await oauthmod._batch_import_kiro_internal(
|
||||
provider_id="provider-1",
|
||||
provider=SimpleNamespace(endpoints=[]), # type: ignore[arg-type]
|
||||
raw_credentials="ignored",
|
||||
db=db,
|
||||
concurrency=1,
|
||||
)
|
||||
|
||||
assert result.total == 1
|
||||
assert result.success == 0
|
||||
assert result.failed == 1
|
||||
assert release_calls
|
||||
db.commit.assert_not_called()
|
||||
@@ -4,6 +4,7 @@ import asyncio
|
||||
from collections.abc import Coroutine, Generator
|
||||
from concurrent.futures import Future
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -96,3 +97,65 @@ async def test_run_batch_delete_waits_for_progress_updates_before_completion(
|
||||
"message": "1 keys deleted",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def test_sync_delete_reports_progress_after_each_batch(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
class _FakeColumn:
|
||||
def __eq__(self, other: object) -> tuple[str, object]: # type: ignore[override]
|
||||
return ("eq", other)
|
||||
|
||||
def in_(self, values: list[str]) -> tuple[str, tuple[str, ...]]:
|
||||
return ("in", tuple(values))
|
||||
|
||||
class _FakeProviderAPIKey:
|
||||
provider_id = _FakeColumn()
|
||||
id = _FakeColumn()
|
||||
|
||||
class _FakeDeleteStatement:
|
||||
def where(self, *_conditions: object) -> "_FakeDeleteStatement":
|
||||
return self
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self) -> None:
|
||||
self.rowcounts = [2, 1]
|
||||
self.commits = 0
|
||||
self.closed = False
|
||||
|
||||
def execute(self, _statement: object) -> SimpleNamespace:
|
||||
# SET LOCAL statement_timeout 不消耗 rowcount
|
||||
if hasattr(_statement, "text"):
|
||||
return SimpleNamespace(rowcount=0)
|
||||
return SimpleNamespace(rowcount=self.rowcounts.pop(0))
|
||||
|
||||
def commit(self) -> None:
|
||||
self.commits += 1
|
||||
|
||||
def rollback(self) -> None:
|
||||
raise AssertionError("rollback should not be called")
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
session = _FakeSession()
|
||||
progress_updates: list[int] = []
|
||||
|
||||
monkeypatch.setattr(taskmod, "_CLEANUP_BATCH_SIZE", 2)
|
||||
monkeypatch.setattr("src.database.create_session", lambda: session)
|
||||
monkeypatch.setattr(
|
||||
"src.models.database.ProviderAPIKey",
|
||||
_FakeProviderAPIKey,
|
||||
)
|
||||
monkeypatch.setattr(taskmod, "sa_delete", lambda _model: _FakeDeleteStatement())
|
||||
|
||||
affected = taskmod._sync_delete(
|
||||
"provider-1",
|
||||
["key-1", "key-2", "key-3"],
|
||||
progress_updates.append,
|
||||
)
|
||||
|
||||
assert affected == 3
|
||||
assert progress_updates == [2, 3]
|
||||
assert session.commits == 2
|
||||
assert session.closed is True
|
||||
|
||||
Reference in New Issue
Block a user