From bb01b2195cbf03e41e978a7c616e0177984e1eb2 Mon Sep 17 00:00:00 2001 From: axisrow Date: Fri, 10 Jul 2026 11:49:47 +0800 Subject: [PATCH] refactor(db): dedupe Z.AI migrations, account row mapping, emoji query (#1135) migrations.py: the twin Z.AI base_url rewriters (_migrate_zai_legacy_base_url and _migrate_zai_empty_base_url_to_coding) shared an identical load-JSON / walk-zai-items / save-on-change body. Extracted _rewrite_zai_base_urls taking a should_rewrite predicate + target URL; each public migration is now a thin caller. Both run only when something changed (unchanged), and each keeps its own provider_registry import + log message. accounts.py: the Account(...) constructor block was duplicated verbatim in get_accounts and get_live_usable_accounts. Extracted _account_from_row(row, session_string); the two methods differ only in how they handle a decrypt failure (raise vs skip), which stays in the caller. get_decrypted_session reimplemented as a delegate to get_session_export so the SELECT + decrypt + row-binding (#1145 consistency) is no longer forked. messages.py: get_trending_emojis had two near-identical SQL bodies differing only by a collected_at window fragment. Collapsed into one query with an optional WHERE fragment + parameterized params list. All three jscpd clones from the #1135 target list are removed. migrations/accounts/messages unit + repository tests pass (400 + 481). Part of #1135 (axis 5 of #1130). Co-Authored-By: Claude Fable 5 --- src/database/migrations.py | 82 ++++++++++++--------------- src/database/repositories/accounts.py | 63 ++++++++------------ src/database/repositories/messages.py | 46 ++++++--------- 3 files changed, 79 insertions(+), 112 deletions(-) diff --git a/src/database/migrations.py b/src/database/migrations.py index d59c1f42..c6c66a42 100644 --- a/src/database/migrations.py +++ b/src/database/migrations.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from pathlib import Path import aiosqlite @@ -708,14 +708,21 @@ async def _migrate_vec_to_portable(db: aiosqlite.Connection) -> None: _ = db -async def _migrate_zai_legacy_base_url(db: aiosqlite.Connection) -> None: - """Rewrite legacy Z.AI Anthropic-compatible base_url to the OpenAI-compatible default.""" - import json +async def _rewrite_zai_base_urls( + db: aiosqlite.Connection, + *, + should_rewrite: Callable[[str], bool], + new_base_url: str, + log_message: str, +) -> None: + """Rewrite ``base_url`` of stored Z.AI provider configs matching a predicate. - from src.agent.provider_registry import ( - ZAI_GENERAL_BASE_URL, - is_zai_legacy_anthropic_base_url, - ) + Shared body of the two Z.AI base_url migrations: load the + ``agent_deepagents_providers_v1`` JSON, rewrite ``plain_fields.base_url`` + (clearing ``last_validation_error``) on every zai item whose current raw + value satisfies ``should_rewrite``, and save only when something changed. + """ + import json cur = await db.execute( "SELECT value FROM settings WHERE key = 'agent_deepagents_providers_v1' LIMIT 1" @@ -738,8 +745,8 @@ async def _migrate_zai_legacy_base_url(db: aiosqlite.Connection) -> None: if not isinstance(plain, dict): continue current = str(plain.get("base_url", "") or "") - if is_zai_legacy_anthropic_base_url(current): - plain["base_url"] = ZAI_GENERAL_BASE_URL + if should_rewrite(current): + plain["base_url"] = new_base_url item["last_validation_error"] = "" changed = True @@ -750,49 +757,34 @@ async def _migrate_zai_legacy_base_url(db: aiosqlite.Connection) -> None: "UPDATE settings SET value = ? WHERE key = 'agent_deepagents_providers_v1'", (json.dumps(data, ensure_ascii=False),), ) - logger.info("Migrated legacy Z.AI Anthropic-compatible base_url to %s", ZAI_GENERAL_BASE_URL) + logger.info(log_message, new_base_url) -async def _migrate_zai_empty_base_url_to_coding(db: aiosqlite.Connection) -> None: - """Backfill empty Z.AI base_url values to the Coding Plan endpoint.""" - import json - - from src.agent.provider_registry import ZAI_CODING_BASE_URL +async def _migrate_zai_legacy_base_url(db: aiosqlite.Connection) -> None: + """Rewrite legacy Z.AI Anthropic-compatible base_url to the OpenAI-compatible default.""" + from src.agent.provider_registry import ( + ZAI_GENERAL_BASE_URL, + is_zai_legacy_anthropic_base_url, + ) - cur = await db.execute( - "SELECT value FROM settings WHERE key = 'agent_deepagents_providers_v1' LIMIT 1" + await _rewrite_zai_base_urls( + db, + should_rewrite=is_zai_legacy_anthropic_base_url, + new_base_url=ZAI_GENERAL_BASE_URL, + log_message="Migrated legacy Z.AI Anthropic-compatible base_url to %s", ) - row = await cur.fetchone() - if not row or not row["value"]: - return - try: - data = json.loads(row["value"]) - except (json.JSONDecodeError, TypeError): - return - if not isinstance(data, list): - return - changed = False - for item in data: - if not isinstance(item, dict) or item.get("provider") != "zai": - continue - plain = item.get("plain_fields") - if not isinstance(plain, dict): - continue - current = (str(plain.get("base_url", "") or "")).strip().rstrip("/") - if current == "": - plain["base_url"] = ZAI_CODING_BASE_URL - item["last_validation_error"] = "" - changed = True - if not changed: - return +async def _migrate_zai_empty_base_url_to_coding(db: aiosqlite.Connection) -> None: + """Backfill empty Z.AI base_url values to the Coding Plan endpoint.""" + from src.agent.provider_registry import ZAI_CODING_BASE_URL - await db.execute( - "UPDATE settings SET value = ? WHERE key = 'agent_deepagents_providers_v1'", - (json.dumps(data, ensure_ascii=False),), + await _rewrite_zai_base_urls( + db, + should_rewrite=lambda current: current.strip().rstrip("/") == "", + new_base_url=ZAI_CODING_BASE_URL, + log_message="Migrated empty Z.AI base_url to %s", ) - logger.info("Migrated empty Z.AI base_url to %s", ZAI_CODING_BASE_URL) async def _migrate_tool_permission_key( diff --git a/src/database/repositories/accounts.py b/src/database/repositories/accounts.py index 7eba755d..dae6a3a9 100644 --- a/src/database/repositories/accounts.py +++ b/src/database/repositories/accounts.py @@ -342,6 +342,25 @@ async def get_account_summaries(self, active_only: bool = False) -> list[Account for row in rows ] + def _account_from_row(self, row, session_string: str) -> Account: + """Build an [`Account`][src.models.Account] from a row + pre-resolved session. + + Shared by the listing paths that decrypt every row's session up front + (`get_accounts`, `get_live_usable_accounts`); the caller decides whether + a decrypt failure aborts the whole list (`get_accounts`) or skips just + that account (`get_live_usable_accounts`). + """ + return Account( + id=row["id"], + phone=row["phone"], + session_string=session_string, + is_primary=bool(row["is_primary"]), + is_active=bool(row["is_active"]), + is_premium=bool(row["is_premium"]) if row["is_premium"] is not None else False, + flood_wait_until=parse_datetime(row["flood_wait_until"]), + created_at=parse_datetime(row["created_at"]), + ) + async def get_accounts(self, active_only: bool = False) -> list[Account]: """Полные [`Account`][src.models.Account] с расшифрованными сессиями для живого использования. @@ -362,19 +381,7 @@ async def get_accounts(self, active_only: bool = False) -> list[Account]: for row in rows: raw_session = str(row["session_string"] or "") session_string = self._decrypt_session_for_live_use(raw_session, str(row["phone"])) - - accounts.append( - Account( - id=row["id"], - phone=row["phone"], - session_string=session_string, - is_primary=bool(row["is_primary"]), - is_active=bool(row["is_active"]), - is_premium=bool(row["is_premium"]) if row["is_premium"] is not None else False, - flood_wait_until=parse_datetime(row["flood_wait_until"]), - created_at=parse_datetime(row["created_at"]), - ) - ) + accounts.append(self._account_from_row(row, session_string)) return accounts @@ -389,21 +396,10 @@ async def get_decrypted_session( ``phone`` must be given. A decrypt failure on the *target* still raises :class:`AccountSessionDecryptError` (the caller wants to know). """ - if (account_id is None) == (phone is None): - raise ValueError("provide exactly one of account_id / phone") - if account_id is not None: - cur = await self._db.execute( - "SELECT phone, session_string FROM accounts WHERE id = ?", (account_id,) - ) - else: - cur = await self._db.execute( - "SELECT phone, session_string FROM accounts WHERE phone = ?", (phone,) - ) - row = await cur.fetchone() - if row is None: + exported = await self.get_session_export(account_id=account_id, phone=phone) + if exported is None: return None - raw_session = str(row["session_string"] or "") - return self._decrypt_session_for_live_use(raw_session, str(row["phone"])) + return exported[1] async def get_session_export( self, *, account_id: int | None = None, phone: str | None = None @@ -457,18 +453,7 @@ async def get_live_usable_accounts(self, active_only: bool = False) -> list[Acco ) continue - accounts.append( - Account( - id=row["id"], - phone=row["phone"], - session_string=session_string, - is_primary=bool(row["is_primary"]), - is_active=bool(row["is_active"]), - is_premium=bool(row["is_premium"]) if row["is_premium"] is not None else False, - flood_wait_until=parse_datetime(row["flood_wait_until"]), - created_at=parse_datetime(row["created_at"]), - ) - ) + accounts.append(self._account_from_row(row, session_string)) return accounts diff --git a/src/database/repositories/messages.py b/src/database/repositories/messages.py index ad29354c..db1eac2d 100644 --- a/src/database/repositories/messages.py +++ b/src/database/repositories/messages.py @@ -1361,35 +1361,25 @@ async def get_trending_emojis(self, limit: int = 10, days: int | None = None) -> Returns: List of ``{"emoji": str, "count": int}`` dicts ordered by count desc. """ + params: list = [] + where_extra = "" if days is not None: - cur = await self._db.execute( - """ - SELECT mr.emoji, SUM(mr.count) AS total - FROM message_reactions mr - JOIN messages m ON mr.channel_id = m.channel_id AND mr.message_id = m.message_id - LEFT JOIN channels c ON m.channel_id = c.channel_id - WHERE (c.is_filtered IS NULL OR c.is_filtered = 0) - AND m.collected_at >= datetime('now', ?) - GROUP BY mr.emoji - ORDER BY total DESC - LIMIT ? - """, - (f"-{days} days", limit), - ) - else: - cur = await self._db.execute( - """ - SELECT mr.emoji, SUM(mr.count) AS total - FROM message_reactions mr - JOIN messages m ON mr.channel_id = m.channel_id AND mr.message_id = m.message_id - LEFT JOIN channels c ON m.channel_id = c.channel_id - WHERE (c.is_filtered IS NULL OR c.is_filtered = 0) - GROUP BY mr.emoji - ORDER BY total DESC - LIMIT ? - """, - (limit,), - ) + where_extra = " AND m.collected_at >= datetime('now', ?)" + params.append(f"-{days} days") + params.append(limit) + cur = await self._db.execute( + f""" + SELECT mr.emoji, SUM(mr.count) AS total + FROM message_reactions mr + JOIN messages m ON mr.channel_id = m.channel_id AND mr.message_id = m.message_id + LEFT JOIN channels c ON m.channel_id = c.channel_id + WHERE (c.is_filtered IS NULL OR c.is_filtered = 0){where_extra} + GROUP BY mr.emoji + ORDER BY total DESC + LIMIT ? + """, + tuple(params), + ) rows = await cur.fetchall() return [{"emoji": r["emoji"], "count": r["total"]} for r in rows]