浅念
5a2ff58af3
The scraper was writing a normalized type label ("text"/"unknown"/...) into
Message.type, but the column was Integer, so the storage layer coerced it to
NULL and the raw code was only recoverable from raw_json. Align with design
§4.3, which specifies two fields: `type` (normalized enum) and `type_raw`
(original code).
- Model (Message): type -> String(32) normalized label; add type_raw Integer
for the raw upstream type code.
- Scraper: _normalize_message sets both `type` (label) and `type_raw`
(coerced int from messageType/type); _is_divider now checks raw -2.
- Repos: type_raw joins _INT_FIELDS; Message.type stored as-is (string).
- Migration 96975050fa79: add type_raw + alter type Integer->String(32),
wrapped in batch_alter_table so it runs on SQLite and MySQL.
- env.py: enable render_as_batch on SQLite (dev/test) for portable batch
migrations; MySQL still uses direct ALTER.
- Tests updated end-to-end (scraper, repos, models).
Co-Authored-By: Claude <noreply@anthropic.com>
341 lines
13 KiB
Python
341 lines
13 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 Customer, Message, Staff, WechatAccount
|
|
|
|
__all__ = [
|
|
"upsert_customer",
|
|
"upsert_customers",
|
|
"upsert_messages",
|
|
"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)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WechatAccount + embedded Staff (best-effort)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def upsert_wechat_account(session: AsyncSession, rec: dict[str, Any]) -> None:
|
|
"""Best-effort idempotent upsert for :class:`WechatAccount` (PK = ``wx_id``).
|
|
|
|
The companion :func:`upsert_staff_from_account` covers the embedded staff
|
|
row (``user_id`` + ``user_name``); the two are split so callers can pick
|
|
either or both. A full staff scraper (departments, phones, etc.) is a
|
|
LATER milestone.
|
|
"""
|
|
await _upsert_one(session, WechatAccount, 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`` when
|
|
present) onto :class:`Staff` (PK = ``user_id``). Silently no-ops when the
|
|
record has no ``user_id`` — there is nothing to key the row on. A
|
|
dedicated staff scraper that fills out phone / dept is a LATER milestone.
|
|
"""
|
|
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,
|
|
}
|
|
if "is_on_job" in rec:
|
|
staff_rec["is_on_job"] = rec["is_on_job"]
|
|
await _upsert_one(session, Staff, staff_rec)
|