Skip to content

Commit 26e73f4

Browse files
axisrowclaude
andauthored
refactor(db): dedupe Z.AI migrations, account row mapping, emoji query (#1135) (#1284)
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: axisrow <axisrow@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent d1796e4 commit 26e73f4

3 files changed

Lines changed: 79 additions & 112 deletions

File tree

src/database/migrations.py

Lines changed: 37 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from __future__ import annotations
22

33
import logging
4-
from collections.abc import Mapping, Sequence
4+
from collections.abc import Callable, Mapping, Sequence
55
from pathlib import Path
66
from typing import cast
77

@@ -710,14 +710,21 @@ async def _migrate_vec_to_portable(db: aiosqlite.Connection) -> None:
710710
_ = db
711711

712712

713-
async def _migrate_zai_legacy_base_url(db: aiosqlite.Connection) -> None:
714-
"""Rewrite legacy Z.AI Anthropic-compatible base_url to the OpenAI-compatible default."""
715-
import json
713+
async def _rewrite_zai_base_urls(
714+
db: aiosqlite.Connection,
715+
*,
716+
should_rewrite: Callable[[str], bool],
717+
new_base_url: str,
718+
log_message: str,
719+
) -> None:
720+
"""Rewrite ``base_url`` of stored Z.AI provider configs matching a predicate.
716721
717-
from src.agent.provider_registry import (
718-
ZAI_GENERAL_BASE_URL,
719-
is_zai_legacy_anthropic_base_url,
720-
)
722+
Shared body of the two Z.AI base_url migrations: load the
723+
``agent_deepagents_providers_v1`` JSON, rewrite ``plain_fields.base_url``
724+
(clearing ``last_validation_error``) on every zai item whose current raw
725+
value satisfies ``should_rewrite``, and save only when something changed.
726+
"""
727+
import json
721728

722729
cur = await db.execute(
723730
"SELECT value FROM settings WHERE key = 'agent_deepagents_providers_v1' LIMIT 1"
@@ -740,8 +747,8 @@ async def _migrate_zai_legacy_base_url(db: aiosqlite.Connection) -> None:
740747
if not isinstance(plain, dict):
741748
continue
742749
current = str(plain.get("base_url", "") or "")
743-
if is_zai_legacy_anthropic_base_url(current):
744-
plain["base_url"] = ZAI_GENERAL_BASE_URL
750+
if should_rewrite(current):
751+
plain["base_url"] = new_base_url
745752
item["last_validation_error"] = ""
746753
changed = True
747754

@@ -752,49 +759,34 @@ async def _migrate_zai_legacy_base_url(db: aiosqlite.Connection) -> None:
752759
"UPDATE settings SET value = ? WHERE key = 'agent_deepagents_providers_v1'",
753760
(json.dumps(data, ensure_ascii=False),),
754761
)
755-
logger.info("Migrated legacy Z.AI Anthropic-compatible base_url to %s", ZAI_GENERAL_BASE_URL)
762+
logger.info(log_message, new_base_url)
756763

757764

758-
async def _migrate_zai_empty_base_url_to_coding(db: aiosqlite.Connection) -> None:
759-
"""Backfill empty Z.AI base_url values to the Coding Plan endpoint."""
760-
import json
761-
762-
from src.agent.provider_registry import ZAI_CODING_BASE_URL
765+
async def _migrate_zai_legacy_base_url(db: aiosqlite.Connection) -> None:
766+
"""Rewrite legacy Z.AI Anthropic-compatible base_url to the OpenAI-compatible default."""
767+
from src.agent.provider_registry import (
768+
ZAI_GENERAL_BASE_URL,
769+
is_zai_legacy_anthropic_base_url,
770+
)
763771

764-
cur = await db.execute(
765-
"SELECT value FROM settings WHERE key = 'agent_deepagents_providers_v1' LIMIT 1"
772+
await _rewrite_zai_base_urls(
773+
db,
774+
should_rewrite=is_zai_legacy_anthropic_base_url,
775+
new_base_url=ZAI_GENERAL_BASE_URL,
776+
log_message="Migrated legacy Z.AI Anthropic-compatible base_url to %s",
766777
)
767-
row = await cur.fetchone()
768-
if not row or not row["value"]:
769-
return
770-
try:
771-
data = json.loads(row["value"])
772-
except (json.JSONDecodeError, TypeError):
773-
return
774-
if not isinstance(data, list):
775-
return
776778

777-
changed = False
778-
for item in data:
779-
if not isinstance(item, dict) or item.get("provider") != "zai":
780-
continue
781-
plain = item.get("plain_fields")
782-
if not isinstance(plain, dict):
783-
continue
784-
current = (str(plain.get("base_url", "") or "")).strip().rstrip("/")
785-
if current == "":
786-
plain["base_url"] = ZAI_CODING_BASE_URL
787-
item["last_validation_error"] = ""
788-
changed = True
789779

790-
if not changed:
791-
return
780+
async def _migrate_zai_empty_base_url_to_coding(db: aiosqlite.Connection) -> None:
781+
"""Backfill empty Z.AI base_url values to the Coding Plan endpoint."""
782+
from src.agent.provider_registry import ZAI_CODING_BASE_URL
792783

793-
await db.execute(
794-
"UPDATE settings SET value = ? WHERE key = 'agent_deepagents_providers_v1'",
795-
(json.dumps(data, ensure_ascii=False),),
784+
await _rewrite_zai_base_urls(
785+
db,
786+
should_rewrite=lambda current: current.strip().rstrip("/") == "",
787+
new_base_url=ZAI_CODING_BASE_URL,
788+
log_message="Migrated empty Z.AI base_url to %s",
796789
)
797-
logger.info("Migrated empty Z.AI base_url to %s", ZAI_CODING_BASE_URL)
798790

799791

800792
async def _migrate_tool_permission_key(

src/database/repositories/accounts.py

Lines changed: 24 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,25 @@ async def get_account_summaries(self, active_only: bool = False) -> list[Account
342342
for row in rows
343343
]
344344

345+
def _account_from_row(self, row, session_string: str) -> Account:
346+
"""Build an [`Account`][src.models.Account] from a row + pre-resolved session.
347+
348+
Shared by the listing paths that decrypt every row's session up front
349+
(`get_accounts`, `get_live_usable_accounts`); the caller decides whether
350+
a decrypt failure aborts the whole list (`get_accounts`) or skips just
351+
that account (`get_live_usable_accounts`).
352+
"""
353+
return Account(
354+
id=row["id"],
355+
phone=row["phone"],
356+
session_string=session_string,
357+
is_primary=bool(row["is_primary"]),
358+
is_active=bool(row["is_active"]),
359+
is_premium=bool(row["is_premium"]) if row["is_premium"] is not None else False,
360+
flood_wait_until=parse_datetime(row["flood_wait_until"]),
361+
created_at=parse_datetime(row["created_at"]),
362+
)
363+
345364
async def get_accounts(self, active_only: bool = False) -> list[Account]:
346365
"""Полные [`Account`][src.models.Account] с расшифрованными сессиями для живого использования.
347366
@@ -362,19 +381,7 @@ async def get_accounts(self, active_only: bool = False) -> list[Account]:
362381
for row in rows:
363382
raw_session = str(row["session_string"] or "")
364383
session_string = self._decrypt_session_for_live_use(raw_session, str(row["phone"]))
365-
366-
accounts.append(
367-
Account(
368-
id=row["id"],
369-
phone=row["phone"],
370-
session_string=session_string,
371-
is_primary=bool(row["is_primary"]),
372-
is_active=bool(row["is_active"]),
373-
is_premium=bool(row["is_premium"]) if row["is_premium"] is not None else False,
374-
flood_wait_until=parse_datetime(row["flood_wait_until"]),
375-
created_at=parse_datetime(row["created_at"]),
376-
)
377-
)
384+
accounts.append(self._account_from_row(row, session_string))
378385

379386
return accounts
380387

@@ -389,21 +396,10 @@ async def get_decrypted_session(
389396
``phone`` must be given. A decrypt failure on the *target* still raises
390397
:class:`AccountSessionDecryptError` (the caller wants to know).
391398
"""
392-
if (account_id is None) == (phone is None):
393-
raise ValueError("provide exactly one of account_id / phone")
394-
if account_id is not None:
395-
cur = await self._db.execute(
396-
"SELECT phone, session_string FROM accounts WHERE id = ?", (account_id,)
397-
)
398-
else:
399-
cur = await self._db.execute(
400-
"SELECT phone, session_string FROM accounts WHERE phone = ?", (phone,)
401-
)
402-
row = await cur.fetchone()
403-
if row is None:
399+
exported = await self.get_session_export(account_id=account_id, phone=phone)
400+
if exported is None:
404401
return None
405-
raw_session = str(row["session_string"] or "")
406-
return self._decrypt_session_for_live_use(raw_session, str(row["phone"]))
402+
return exported[1]
407403

408404
async def get_session_export(
409405
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
457453
)
458454
continue
459455

460-
accounts.append(
461-
Account(
462-
id=row["id"],
463-
phone=row["phone"],
464-
session_string=session_string,
465-
is_primary=bool(row["is_primary"]),
466-
is_active=bool(row["is_active"]),
467-
is_premium=bool(row["is_premium"]) if row["is_premium"] is not None else False,
468-
flood_wait_until=parse_datetime(row["flood_wait_until"]),
469-
created_at=parse_datetime(row["created_at"]),
470-
)
471-
)
456+
accounts.append(self._account_from_row(row, session_string))
472457

473458
return accounts
474459

src/database/repositories/messages.py

Lines changed: 18 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1362,35 +1362,25 @@ async def get_trending_emojis(self, limit: int = 10, days: int | None = None) ->
13621362
Returns:
13631363
List of ``{"emoji": str, "count": int}`` dicts ordered by count desc.
13641364
"""
1365+
params: list = []
1366+
where_extra = ""
13651367
if days is not None:
1366-
cur = await self._db.execute(
1367-
"""
1368-
SELECT mr.emoji, SUM(mr.count) AS total
1369-
FROM message_reactions mr
1370-
JOIN messages m ON mr.channel_id = m.channel_id AND mr.message_id = m.message_id
1371-
LEFT JOIN channels c ON m.channel_id = c.channel_id
1372-
WHERE (c.is_filtered IS NULL OR c.is_filtered = 0)
1373-
AND m.collected_at >= datetime('now', ?)
1374-
GROUP BY mr.emoji
1375-
ORDER BY total DESC
1376-
LIMIT ?
1377-
""",
1378-
(f"-{days} days", limit),
1379-
)
1380-
else:
1381-
cur = await self._db.execute(
1382-
"""
1383-
SELECT mr.emoji, SUM(mr.count) AS total
1384-
FROM message_reactions mr
1385-
JOIN messages m ON mr.channel_id = m.channel_id AND mr.message_id = m.message_id
1386-
LEFT JOIN channels c ON m.channel_id = c.channel_id
1387-
WHERE (c.is_filtered IS NULL OR c.is_filtered = 0)
1388-
GROUP BY mr.emoji
1389-
ORDER BY total DESC
1390-
LIMIT ?
1391-
""",
1392-
(limit,),
1393-
)
1368+
where_extra = " AND m.collected_at >= datetime('now', ?)"
1369+
params.append(f"-{days} days")
1370+
params.append(limit)
1371+
cur = await self._db.execute(
1372+
f"""
1373+
SELECT mr.emoji, SUM(mr.count) AS total
1374+
FROM message_reactions mr
1375+
JOIN messages m ON mr.channel_id = m.channel_id AND mr.message_id = m.message_id
1376+
LEFT JOIN channels c ON m.channel_id = c.channel_id
1377+
WHERE (c.is_filtered IS NULL OR c.is_filtered = 0){where_extra}
1378+
GROUP BY mr.emoji
1379+
ORDER BY total DESC
1380+
LIMIT ?
1381+
""",
1382+
tuple(params),
1383+
)
13941384
rows = await cur.fetchall()
13951385
return [{"emoji": r["emoji"], "count": r["total"]} for r in rows]
13961386

0 commit comments

Comments
 (0)