Files
yunke/yunke-server/app/storage/repos.py
T
浅念 2a74d3888d feat: wechat group+member and talker scrapers
Add scrape_wechat_groups (POST /wechatweb/chatroom/getChatroomList) and
scrape_wechat_talkers (POST /wechatweb/wxFriends/talkerList), both reusing
the page-number paginator. Groups parse the semicolon-joined memberIds into
a member_ids_list so upsert_chat_group can mirror GroupMember rows
(composite PK channel/room_id/member_id) idempotently alongside the
ChatGroup row. Talkers are a sync-planning input (member_count>0 marks a
group conversation); no dedicated table.

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-01 13:54:16 +08:00

446 lines
17 KiB
Python

"""Idempotent upsert repositories bridging scraper dicts to ORM tables.
Every public helper takes an ``AsyncSession`` plus one-or-more scraper-shaped
dicts (the snake_case records that ``app.scraper.*`` yields) and writes them
to the corresponding table via **dialect-aware Core upsert**, so the same call
path is correct on:
* **in-memory SQLite** — used by the test suite (``sqlite+aiosqlite``); uses
``sqlalchemy.dialects.sqlite.insert(...).on_conflict_do_update(...)``, and
* **MySQL** — production (``mysql+asyncmy``); uses
``sqlalchemy.dialects.mysql.insert(...).on_duplicate_key_update(...)``.
The dialect is detected at write time via ``session.bind.name``, so callers
do not need to know which backend they target.
Type coercion
-------------
Upstream CRM values are not always typed the way the columns want them. The
following coercions are applied when building column values from the input
dict (see :func:`_coerce_bool`, :func:`_coerce_int`):
* Booleans (``is_block`` / ``is_send`` / ``is_del`` / ``is_on_job``):
``1`` / ``"1"`` / ``True`` -> ``True``; ``0`` / ``"0"`` / ``False`` ->
``False``; ``None`` preserved.
* ``gender``: ``"0"`` / ``"1"`` / ``"2"`` (str) -> ``int`` when not ``None``.
* ``type_raw`` (Message): the raw upstream type code, coerced to
``int | None``. ``Message.type`` is a normalized label string and is
stored as-is (no int coercion).
* Timestamps (``last_chat_time`` / ``timestamp`` / ``last_update`` / ...):
passed through as-is (ms-epoch ``int`` or ``None``).
* ``raw_json``: the original record dict, preserved verbatim. When the input
already carries a ``raw_json`` key (scraper output does), that value wins;
otherwise the input dict itself is stored.
Sharding note
-------------
``app.db.models.message_table_name`` exists for the future per-month sharded
``message_<channel>_<YYYYMM>`` layout. For M1, the :class:`Message` model maps
to a single physical ``message`` table and :func:`upsert_messages` writes to
that one table. Multi-table routing is a LATER task (YAGNI for M1).
"""
from __future__ import annotations
from typing import Any
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase
from app.db.models import (
ChatGroup,
Customer,
Department,
GroupMember,
Message,
Staff,
WechatAccount,
)
__all__ = [
"upsert_chat_group",
"upsert_customer",
"upsert_customers",
"upsert_department",
"upsert_messages",
"upsert_staff",
"upsert_staff_from_account",
"upsert_wechat_account",
]
# ---------------------------------------------------------------------------
# Type coercion helpers (module-level so they are unit-testable)
# ---------------------------------------------------------------------------
# Columns whose source values are int-as-bool and must land as actual bool.
_BOOL_FIELDS: frozenset[str] = frozenset(
{"is_block", "is_send", "is_del", "is_on_job", "is_upload"}
)
# Columns whose source values may arrive as numeric strings.
_INT_FIELDS: frozenset[str] = frozenset(
{
"gender",
"type_raw",
"pic_type",
"friend_state",
"delete_flag",
"direction",
"incoming_call",
"call_status",
"qc_status",
}
)
def _coerce_bool(value: Any) -> bool | None:
"""Coerce int/str/bool CRM bool fields to a real bool.
``None`` is preserved. ``1``/``"1"``/``True`` -> ``True``;
``0``/``"0"``/``False`` -> ``False``. Common textual forms
(``"true"``/``"false"``/``"yes"``/``"no"``) are also accepted; any other
non-empty string falls back to truthiness so a stray value never raises.
"""
if value is None:
return None
if isinstance(value, bool):
return value
if isinstance(value, int):
return value != 0
if isinstance(value, str):
s = value.strip().lower()
if s in ("1", "true", "yes", "y", "t"):
return True
if s in ("0", "false", "no", "n", "f", ""):
return False
# Last-resort: try numeric compare, else truthiness of the raw string.
try:
return int(s) != 0
except ValueError:
return bool(s)
# Floats / other: Python truthiness.
return bool(value)
def _coerce_int(value: Any) -> int | None:
"""Coerce a possibly-string numeric value to ``int``; ``None`` preserved.
``"2"`` -> ``2``; ``2`` -> ``2``; ``0``/``"0"`` -> ``0``; ``None`` ->
``None``. A non-numeric string (e.g. ``"text"``) returns ``None`` — the
caller (e.g. Message.type_raw) treats that as "no value".
"""
if value is None:
return None
if isinstance(value, bool):
return int(value)
if isinstance(value, int):
return value
if isinstance(value, float):
return int(value)
if isinstance(value, str):
s = value.strip()
if not s:
return None
try:
return int(s)
except ValueError:
return None
return None
# ---------------------------------------------------------------------------
# Dialect-aware Core upsert
# ---------------------------------------------------------------------------
def _pk_col_names(model: type[DeclarativeBase]) -> list[str]:
"""Return the names of the model's primary-key columns."""
return [c.name for c in model.__table__.primary_key.columns]
def _build_column_values(
model: type[DeclarativeBase], rec: dict[str, Any]
) -> dict[str, Any]:
"""Pick rec keys that match model columns, applying type coercion.
* Only keys that correspond to real columns on ``model`` are kept (extra
upstream keys are dropped, since they would cause Core insert() to raise
on unknown columns).
* Boolean / int-as-str fields are coerced via :func:`_coerce_bool` /
:func:`_coerce_int`.
* ``raw_json`` defaults to the entire input dict when not explicitly
provided, so the verbatim upstream record is always recoverable.
"""
col_names = {c.name for c in model.__table__.columns}
values: dict[str, Any] = {}
for col_name in col_names:
if col_name in rec:
v = rec[col_name]
if col_name in _BOOL_FIELDS:
v = _coerce_bool(v)
elif col_name in _INT_FIELDS:
v = _coerce_int(v)
values[col_name] = v
# Preserve the verbatim source record. If the scraper already set raw_json
# (it does), that wins; otherwise the input dict itself is the source of
# truth.
if "raw_json" in col_names and "raw_json" not in values:
values["raw_json"] = rec
return values
def _build_upsert_stmt(
model: type[DeclarativeBase], values: dict[str, Any], dialect_name: str
) -> sa.Insert:
"""Build a dialect-specific idempotent Core upsert for one row.
On conflict the statement updates **every non-PK column** from the
incoming row, so re-running a scrape yields one final row per PK with the
freshest values (counts, version, names, raw_json, updated_at, ...).
Parameters
----------
model:
The ORM model class to upsert into.
values:
Column-name -> value dict (output of :func:`_build_column_values`).
dialect_name:
One of ``"sqlite"`` / ``"mysql"`` (the two backends this project
targets). Any other dialect raises ``ValueError``.
"""
pk_cols = _pk_col_names(model)
# Non-PK columns to refresh on conflict (volatile fields).
update_col_names = [c.name for c in model.__table__.columns if c.name not in pk_cols]
# NB: access ``excluded`` / ``inserted`` on the BARE insert (before
# ``.values(...)``) and reuse that accessor when building the SET dict.
# If the accessor is taken on the post-``.values()`` statement, the
# ``excluded.<col>`` references do not bind to the supplied values at
# compile time and the ON CONFLICT DO UPDATE silently fails to refresh
# the row (a SQLAlchemy quirk; verified on 2.0.x).
if dialect_name == "sqlite":
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
bare = sqlite_insert(model)
excluded = bare.excluded
set_ = {name: excluded[name] for name in update_col_names}
return bare.values(**values).on_conflict_do_update(
index_elements=pk_cols, set_=set_
)
if dialect_name == "mysql":
from sqlalchemy.dialects.mysql import insert as mysql_insert
bare = mysql_insert(model)
inserted = bare.inserted
update_dict = {name: inserted[name] for name in update_col_names}
return bare.values(**values).on_duplicate_key_update(**update_dict)
raise ValueError(
f"Unsupported dialect for idempotent upsert: {dialect_name!r}. "
"Only 'sqlite' (tests) and 'mysql' (prod) are supported."
)
async def _dialect_name(session: AsyncSession) -> str:
"""Return the dialect name of the session's bind (e.g. 'sqlite'/'mysql').
Uses ``session.bind.name`` (the AsyncEngine); falls back to
``session.get_bind().dialect.name`` if ``bind`` is not set on the session.
"""
bind = session.bind
if bind is not None and hasattr(bind, "name"):
return bind.name # type: ignore[no-any-return]
return session.get_bind().dialect.name
async def _upsert_one(session: AsyncSession, model: type[DeclarativeBase], rec: dict[str, Any]) -> None:
"""Upsert a single record into ``model``."""
values = _build_column_values(model, rec)
# If the record carries no values for any column (e.g. empty dict), there
# is nothing to write — skip the round-trip. PK-only rows still go through
# because their PK is in `values`.
if not values:
return
dialect = await _dialect_name(session)
stmt = _build_upsert_stmt(model, values, dialect)
await session.execute(stmt)
# ---------------------------------------------------------------------------
# Public upsert helpers
# ---------------------------------------------------------------------------
async def upsert_customer(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Idempotently upsert one :class:`Customer` (PK = ``friend_id``).
Same ``friend_id`` twice -> one row, with non-PK fields refreshed from the
later record.
"""
await _upsert_one(session, Customer, rec)
async def upsert_customers(session: AsyncSession, recs: list[dict[str, Any]]) -> int:
"""Batch upsert for :class:`Customer`. Returns the number of records processed.
Idempotent: duplicate ``friend_id`` values within the batch collapse to a
single row (the last write wins) at the DB layer.
"""
for rec in recs:
await _upsert_one(session, Customer, rec)
return len(recs)
async def upsert_messages(session: AsyncSession, recs: list[dict[str, Any]]) -> int:
"""Batch idempotent upsert for :class:`Message` (PK = ``msg_id``).
Applies the documented coercions (``is_send``/``is_del`` -> bool,
``type_raw`` -> int). ``Message.type`` is the normalized label string and
is stored as-is. Returns the number of records processed (input length);
the DB collapses duplicate ``msg_id`` values to one row.
.. note::
For M1 the ``Message`` model maps to a single physical ``message``
table. Real per-month sharding across
``message_<channel>_<YYYYMM>`` tables (see
:func:`app.db.models.message_table_name`) is a LATER task; do not
implement multi-table routing here.
"""
for rec in recs:
await _upsert_one(session, Message, rec)
return len(recs)
# ---------------------------------------------------------------------------
# Department upsert (Task 2.0)
# ---------------------------------------------------------------------------
async def upsert_department(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Idempotently upsert one :class:`Department` (PK = ``dept_id``).
Same ``dept_id`` twice -> one row, with non-PK fields (``parent_id`` /
``name`` / ``path`` / ``raw_json``) refreshed from the later record.
"""
await _upsert_one(session, Department, rec)
# ---------------------------------------------------------------------------
# WechatAccount + embedded Staff (best-effort)
# ---------------------------------------------------------------------------
async def upsert_staff(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Idempotently upsert one :class:`Staff` row from a staff-shaped dict.
PK = ``user_id``. Same ``user_id`` twice -> one row, with non-PK fields
refreshed. Applies the documented coercions (``is_on_job`` -> bool).
This is the primitive staff upsert; the account-scoped helper
:func:`upsert_staff_from_account` (which extracts the staff fields from a
wechat-account record) delegates here, and :func:`upsert_wechat_account`
calls that helper so an account upsert automatically refreshes its
embedded staff row too.
"""
await _upsert_one(session, Staff, rec)
async def upsert_wechat_account(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Idempotently upsert a :class:`WechatAccount` AND its embedded Staff row.
PK = ``wx_id``. The companion :func:`upsert_staff_from_account` covers the
embedded staff row (``user_id`` + ``user_name``); this helper calls it
automatically so an account scrape that lands accounts also lands the
owning staff — callers do not need to upsert staff separately.
The staff upsert is best-effort: an account record with no ``user_id``
(e.g. an orphaned / unassigned wechat number) still writes the account
row; the staff upsert silently no-ops.
"""
await _upsert_one(session, WechatAccount, rec)
await upsert_staff_from_account(session, rec)
async def upsert_staff_from_account(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Best-effort upsert of the Staff row embedded in a wechat-account record.
Maps the account's ``user_id`` + ``user_name`` (plus ``is_on_job`` /
``phone`` / ``sale_phone`` when present) onto :class:`Staff`
(PK = ``user_id``) and delegates to :func:`upsert_staff`. Silently no-ops
when the record has no ``user_id`` — there is nothing to key the row on.
Per Task 2.1, dept linkage for staff is inferred ELSEWHERE (from
customers' ``deptRoute``); staff with no customers have no dept. This
helper therefore does not populate ``dept_id`` — that is a later concern.
"""
user_id = rec.get("user_id")
if not user_id:
return
staff_rec: dict[str, Any] = {
"user_id": user_id,
"name": rec.get("user_name"),
"raw_json": rec,
}
# Forward any Staff-relevant fields the account record happens to carry.
# ``phone`` and ``sale_phone`` are usually absent on the wxUser payload
# (they come from the WhatsApp account payload); carry them through when
# present so a future shared caller does not silently drop them.
for key in ("is_on_job", "phone", "sale_phone"):
if key in rec:
staff_rec[key] = rec[key]
await upsert_staff(session, staff_rec)
# ---------------------------------------------------------------------------
# ChatGroup + GroupMember upsert (M2)
# ---------------------------------------------------------------------------
async def upsert_chat_group(session: AsyncSession, rec: dict[str, Any]) -> None:
"""Idempotently upsert one :class:`ChatGroup` AND its :class:`GroupMember` rows.
PK = ``room_id``. The group row is upserted first; then, for every id in
the record's ``member_ids_list`` (the scraper-parsed, empties-dropped
member wxid list), a :class:`GroupMember` row is upserted against the
composite PK ``(channel, room_id, member_id)``. Re-scraping a group
therefore refreshes the group row and re-asserts its members without
duplication (idempotent).
The ``member_ids_list`` key is the parsed list shape produced by
:func:`app.scraper.wechat._normalize_chat_group`; it is NOT a
``ChatGroup`` column, so it is dropped by :func:`_build_column_values`
when writing the group row (the raw semicolon-joined string is preserved
on ``ChatGroup.member_ids``).
Per-member ``nickname`` / ``avatar`` / ``display_name`` are left ``None``
here — Phase-0 does not yet resolve them (they need a separate per-member
API; see data dictionary §9.6). The repo writes a minimal member row keyed
by the composite PK so membership is captured now and enriched later.
"""
await _upsert_one(session, ChatGroup, rec)
channel = rec.get("channel") or "wechat"
room_id = rec.get("room_id")
# No room_id -> nothing to key member rows on (defensive; the scraper
# always sets room_id). Skip the member loop rather than raising so a
# malformed record still lands its (PK-less) group upsert attempt above.
if not room_id:
return
member_ids: list[Any] = list(rec.get("member_ids_list") or [])
for member_id in member_ids:
# Defensive: never insert an empty/falsy member id (would collide on
# an empty-PK row). The scraper already drops empties, but a hand-
# built record could slip one through.
if not member_id:
continue
member_rec: dict[str, Any] = {
"channel": channel,
"room_id": room_id,
"member_id": member_id,
}
await _upsert_one(session, GroupMember, member_rec)