diff --git a/yunke-server/app/storage/__init__.py b/yunke-server/app/storage/__init__.py new file mode 100644 index 0000000..af20576 --- /dev/null +++ b/yunke-server/app/storage/__init__.py @@ -0,0 +1,34 @@ +"""Storage layer for the yunke data-mirror system. + +Bridges the scrapers (``app.scraper.*``, which yield normalized dicts) to the +database (``app.db.models.*``) and the local media store. + +Public surface (M1 / Task 1.7): + +* :mod:`app.storage.repos` — idempotent Core upsert helpers that take an + ``AsyncSession`` plus scraper-shaped dicts and write to the target table. + The upsert is dialect-aware (SQLite for tests, MySQL for prod) so the same + call path works in both environments. +* :mod:`app.storage.object_store` — a local-filesystem object store with + sha256-content-addressed dedup. The interface (``put_bytes`` / ``put_stream`` + / ``put_by_hash`` / ``exists`` / ``get_path``) is the seam a future MinIO/S3 + backend will swap in behind. +""" + +from app.storage.object_store import ObjectStore +from app.storage.repos import ( + upsert_customer, + upsert_customers, + upsert_messages, + upsert_staff_from_account, + upsert_wechat_account, +) + +__all__ = [ + "ObjectStore", + "upsert_customer", + "upsert_customers", + "upsert_messages", + "upsert_staff_from_account", + "upsert_wechat_account", +] diff --git a/yunke-server/app/storage/object_store.py b/yunke-server/app/storage/object_store.py new file mode 100644 index 0000000..09b7e6c --- /dev/null +++ b/yunke-server/app/storage/object_store.py @@ -0,0 +1,155 @@ +"""Local-filesystem object store with sha256 content-addressed dedup. + +This is the M1 storage backend for media blobs (image / voice / video / file +downloads referenced by :class:`app.db.models.Media`). The interface is the +seam a future MinIO/S3 backend will swap in behind: + +* :meth:`ObjectStore.put_bytes` — write bytes under an opaque key. +* :meth:`ObjectStore.put_stream` — write an async-iterable stream under a key. +* :meth:`ObjectStore.put_by_hash` — content-addressed put with sha256 dedup: + computes the hash, and only writes when no file with that hash exists. The + path layout is ``base_dir//`` (a simple + fan-out directory so no single directory holds millions of entries). +* :meth:`ObjectStore.exists` / :meth:`get_path` — predicate / pure-path helpers. + +All disk I/O is dispatched via :func:`asyncio.to_thread`, so the methods are +genuinely async and safe to call from a FastAPI handler / scraper loop without +blocking the event loop on a slow disk. + +The default base directory is :attr:`Settings.media_dir`. +""" + +from __future__ import annotations + +import asyncio +import hashlib +from collections.abc import AsyncIterator +from pathlib import Path + +from app.config import get_settings + +__all__ = ["ObjectStore"] + + +class ObjectStore: + """Async local-filesystem object store with sha256-based dedup.""" + + def __init__(self, base_dir: Path | None = None) -> None: + """Create a store rooted at ``base_dir``. + + Defaults to :attr:`Settings.media_dir` when omitted so callers can + ``ObjectStore()`` and land in the configured media location. + """ + self.base_dir: Path = Path(base_dir) if base_dir is not None else get_settings().media_dir + + # ------------------------------------------------------------------ + # Path helpers (pure, no I/O) + # ------------------------------------------------------------------ + + def get_path(self, key: str) -> Path: + """Return ``base_dir / `` (a pure-path computation; no I/O). + + ``key`` may contain ``/`` separators (e.g. ``"foo/bar.txt"``); they + become real subdirectories on disk. The key is treated as opaque — no + hash is computed. + """ + return self.base_dir / key + + def _hash_path(self, sha256_hex: str, ext: str) -> Path: + """Return ``base_dir / / ``.""" + if ext and not ext.startswith("."): + ext = "." + ext + return self.base_dir / sha256_hex[:2] / f"{sha256_hex}{ext}" + + # ------------------------------------------------------------------ + # Writes + # ------------------------------------------------------------------ + + async def put_bytes(self, key: str, data: bytes) -> Path: + """Write ``data`` under ``key`` (overwrites if present). + + Returns the resolved path. Always writes; for content-addressed dedup + use :meth:`put_by_hash`. + """ + + def _write() -> Path: + path = self.get_path(key) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return path + + return await asyncio.to_thread(_write) + + async def put_stream( + self, key: str, stream: AsyncIterator[bytes] + ) -> Path: + """Concatenate an async-byte-stream under ``key`` (overwrites). + + Chunks are consumed in the event-loop thread; the file open + each + write is dispatched to a worker thread to avoid blocking on disk. + """ + + path = self.get_path(key) + + def _mkdir() -> None: + path.parent.mkdir(parents=True, exist_ok=True) + + await asyncio.to_thread(_mkdir) + + def _write_chunk(buf: list[bytes], f) -> None: + for chunk in buf: + f.write(chunk) + buf.clear() + + # Stream chunks in the event loop but push each disk write to a + # thread so we never block on slow storage mid-iteration. + def _open_and_prep(): + return open(path, "wb") + + f = await asyncio.to_thread(_open_and_prep) + try: + buffer: list[bytes] = [] + async for chunk in stream: + buffer.append(chunk) + # Flush once we have a meaningful amount buffered. + if sum(len(c) for c in buffer) >= 65536: + await asyncio.to_thread(_write_chunk, buffer, f) + if buffer: + await asyncio.to_thread(_write_chunk, buffer, f) + finally: + await asyncio.to_thread(f.close) + + return path + + async def put_by_hash( + self, data: bytes, *, ext: str = "" + ) -> tuple[str, Path]: + """Content-addressed put with sha256 dedup. + + Computes ``sha256(data)``, derives ``base_dir/<2 chars>/``, + and writes the blob **only if no file with that path already exists** + (deduplication across messages that reference the same blob via md5). + + Returns ``(sha256_hex, path)``. On a hit the existing file is left + untouched (no rewrite), so dedup is observable. + """ + sha256_hex = hashlib.sha256(data).hexdigest() + path = self._hash_path(sha256_hex, ext) + + def _write_if_missing() -> Path: + if path.exists(): + return path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return path + + await asyncio.to_thread(_write_if_missing) + return sha256_hex, path + + # ------------------------------------------------------------------ + # Reads / predicates + # ------------------------------------------------------------------ + + async def exists(self, key: str) -> bool: + """True iff ``key`` exists on disk under ``base_dir``.""" + return await asyncio.to_thread(self.get_path(key).exists) diff --git a/yunke-server/app/storage/repos.py b/yunke-server/app/storage/repos.py new file mode 100644 index 0000000..f430b4d --- /dev/null +++ b/yunke-server/app/storage/repos.py @@ -0,0 +1,340 @@ +"""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`` (Message): coerced to ``int``; a non-numeric string (e.g. a + normalized type label like ``"text"``) collapses to ``None`` — the raw code + is still recoverable from ``raw_json``. +* 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__`` 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", + "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 (Message.type) treats that as "no value", and the raw code is still + recoverable from ``raw_json``. + """ + 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.`` 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`` + -> int). 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__`` 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) diff --git a/yunke-server/tests/test_object_store.py b/yunke-server/tests/test_object_store.py new file mode 100644 index 0000000..2bb1104 --- /dev/null +++ b/yunke-server/tests/test_object_store.py @@ -0,0 +1,156 @@ +"""Tests for app.storage.object_store — local-filesystem object store. + +Exercises the async hash-dedup store against a real tmp_path on disk. All disk +I/O is dispatched via ``asyncio.to_thread`` so the methods are genuinely async. +""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +from app.storage.object_store import ObjectStore + +# --------------------------------------------------------------------------- +# put_bytes / get_path / exists +# --------------------------------------------------------------------------- + + +async def test_put_bytes_writes_file_and_exists_reports_true(tmp_path: Path): + """put_bytes writes the blob; exists() then reports True; path matches key.""" + store = ObjectStore(tmp_path) + key = "foo/bar.txt" + data = b"hello" + + assert await store.exists(key) is False + + path = await store.put_bytes(key, data) + assert path == tmp_path / key + assert path.exists() + assert path.read_bytes() == data + + assert await store.exists(key) is True + + +async def test_get_path_returns_base_plus_key(tmp_path: Path): + """get_path is a pure path computation (no I/O). + + ``/`` in the key becomes real path segments (subdirectories). + """ + store = ObjectStore(tmp_path) + assert store.get_path("a/b/c.bin") == tmp_path / "a" / "b" / "c.bin" + assert store.get_path("flat.txt") == tmp_path / "flat.txt" + + +async def test_put_bytes_overwrites_on_change(tmp_path: Path): + """A second put_bytes with new content overwrites the file.""" + store = ObjectStore(tmp_path) + await store.put_bytes("k", b"v1") + await store.put_bytes("k", b"v2") + assert (tmp_path / "k").read_bytes() == b"v2" + + +# --------------------------------------------------------------------------- +# put_by_hash (dedup) +# --------------------------------------------------------------------------- + + +async def test_put_by_hash_returns_sha256_and_path_with_layout(tmp_path: Path): + """Returns (sha256_hex, base/<2 chars>/); file content correct.""" + store = ObjectStore(tmp_path) + data = b"abc" + expected_hash = hashlib.sha256(data).hexdigest() + + h, path = await store.put_by_hash(data, ext=".txt") + + assert h == expected_hash + assert path == tmp_path / expected_hash[:2] / (expected_hash + ".txt") + assert path.exists() + assert path.read_bytes() == data + + +async def test_put_by_hash_dedups_same_bytes_returns_same_path(tmp_path: Path): + """put_by_hash with identical bytes twice -> same (hash, path).""" + store = ObjectStore(tmp_path) + data = b"identical" + + h1, p1 = await store.put_by_hash(data) + h2, p2 = await store.put_by_hash(data) + + assert h1 == h2 + assert p1 == p2 + + +async def test_put_by_hash_does_not_rewrite_existing_file(tmp_path: Path): + """Dedup: a second call for an existing hash must NOT touch the file. + + Proves the implementation checks ``exists`` before writing. We corrupt the + file between calls; a dedup-aware store must leave the corruption in place + (no rewrite), since the blob is already present. + """ + store = ObjectStore(tmp_path) + data = b"original" + + _h1, p1 = await store.put_by_hash(data, ext=".bin") + assert p1.read_bytes() == data + + # Corrupt the on-disk file. + p1.write_bytes(b"CORRUPTED") + + _h2, p2 = await store.put_by_hash(data, ext=".bin") + assert p2 == p1 + # Dedup short-circuit: file content was NOT overwritten. + assert p2.read_bytes() == b"CORRUPTED" + + +async def test_put_by_hash_distinct_bytes_get_distinct_paths(tmp_path: Path): + """Different bytes -> different hash -> different path.""" + store = ObjectStore(tmp_path) + h1, p1 = await store.put_by_hash(b"aaa") + h2, p2 = await store.put_by_hash(b"bbb") + assert h1 != h2 + assert p1 != p2 + + +async def test_put_by_hash_default_ext_empty(tmp_path: Path): + """Default ext='' -> no suffix appended to the file.""" + store = ObjectStore(tmp_path) + data = b"x" + expected_hash = hashlib.sha256(data).hexdigest() + h, p = await store.put_by_hash(data) + assert h == expected_hash + assert p == tmp_path / expected_hash[:2] / expected_hash + + +# --------------------------------------------------------------------------- +# put_stream +# --------------------------------------------------------------------------- + + +async def test_put_stream_writes_concatenated_chunks(tmp_path: Path): + """put_stream consumes an async generator and writes concatenated bytes.""" + + async def gen(): + yield b"chunk1-" + yield b"chunk2-" + yield b"chunk3" + + store = ObjectStore(tmp_path) + path = await store.put_stream("streamed/key", gen()) + assert path == tmp_path / "streamed" / "key" + assert path.read_bytes() == b"chunk1-chunk2-chunk3" + + +# --------------------------------------------------------------------------- +# Default base dir (Settings.media_dir) +# --------------------------------------------------------------------------- + + +def test_object_store_default_base_uses_settings_media_dir(tmp_path: Path, monkeypatch): + """Omitting base_dir falls back to Settings.media_dir.""" + monkeypatch.setenv("MEDIA_DIR", str(tmp_path)) + from app.config import get_settings + + get_settings.cache_clear() + store = ObjectStore() + assert store.base_dir == get_settings().media_dir diff --git a/yunke-server/tests/test_repos.py b/yunke-server/tests/test_repos.py new file mode 100644 index 0000000..cac2721 --- /dev/null +++ b/yunke-server/tests/test_repos.py @@ -0,0 +1,397 @@ +"""Tests for app.storage.repos — idempotent upserts + type coercion. + +These tests build an isolated in-memory SQLite async engine (same pattern as +``tests/test_models.py``) and exercise the public upsert helpers in +``app.storage.repos``. They never touch the production MySQL engine. + +Coverage: + +* Idempotency: same PK twice -> one row, with non-PK fields updated. +* Type coercion: bool fields (``is_send``/``is_del``/``is_block``/``is_on_job``) + accept int / "1"-style str / bool; ``gender`` str -> int; ``None`` preserved. +* Batch dedup: 3 messages with one duplicate ``msg_id`` -> 2 rows in the table. +* Return values: ``upsert_customers`` / ``upsert_messages`` report counts. +* Embedded-staff best-effort upsert from a wechat-account record. +* Direct unit tests on the module-level ``_coerce_bool`` / ``_coerce_int`` + helpers (kept module-level precisely so they can be exercised here). +""" + +from __future__ import annotations + +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app.db.base import Base +from app.db.models import Customer, Message, Staff, WechatAccount +from app.storage.repos import ( + _coerce_bool, + _coerce_int, + upsert_customer, + upsert_customers, + upsert_messages, + upsert_staff_from_account, + upsert_wechat_account, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +async def _make_engine(): + """Create an in-memory async SQLite engine with all tables created.""" + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + return engine + + +# --------------------------------------------------------------------------- +# Coercion helpers (module-level, exercised directly for unit coverage) +# --------------------------------------------------------------------------- + + +def test_coerce_bool_handles_documented_cases(): + """bool fields accept None / bool / int / '1'-style str.""" + assert _coerce_bool(None) is None + assert _coerce_bool(True) is True + assert _coerce_bool(False) is False + assert _coerce_bool(1) is True + assert _coerce_bool(0) is False + assert _coerce_bool("1") is True + assert _coerce_bool("0") is False + + +def test_coerce_int_handles_documented_cases(): + """gender / type-style fields: str -> int; None preserved; 0 preserved.""" + assert _coerce_int(None) is None + assert _coerce_int(2) == 2 + assert _coerce_int("2") == 2 + assert _coerce_int(0) == 0 + assert _coerce_int("0") == 0 + + +def test_coerce_int_returns_none_for_non_numeric_string(): + """A non-numeric string (e.g. normalized 'type' label) -> None. + + The column is INTEGER; a non-numeric value cannot be stored. The raw value + remains recoverable from ``raw_json``. + """ + assert _coerce_int("text") is None + assert _coerce_int("unknown") is None + + +# --------------------------------------------------------------------------- +# Customer upsert +# --------------------------------------------------------------------------- + + +async def test_upsert_customer_is_idempotent_and_updates_version(): + """Upsert same friend_id twice with different version -> 1 row, version updated.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + rec1 = { + "friend_id": "wxid_abc", + "channel": "wechat", + "version": "v1", + "friend_nickname": "Alice", + } + rec2 = { + "friend_id": "wxid_abc", + "channel": "wechat", + "version": "v2", + "friend_nickname": "Alice", + } + + await upsert_customer(session, rec1) + await session.commit() + await upsert_customer(session, rec2) + await session.commit() + + # ``populate_existing=True`` bypasses the identity-map cache so the + # verify query reflects the second upsert's writes (with + # ``expire_on_commit=False`` the loaded instance is otherwise stale). + rows = ( + await session.execute( + sa.select(Customer).execution_options(populate_existing=True) + ) + ).scalars().all() + assert len(rows) == 1 + assert rows[0].friend_id == "wxid_abc" + assert rows[0].version == "v2" + assert rows[0].friend_nickname == "Alice" + finally: + await engine.dispose() + + +async def test_upsert_customer_coerces_gender_and_bool(): + """gender '2' -> int 2; is_block '1' -> True.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + rec = { + "friend_id": "wxid_g1", + "channel": "wechat", + "gender": "2", + "is_block": "1", + } + await upsert_customer(session, rec) + await session.commit() + + fetched = ( + await session.execute( + sa.select(Customer).where(Customer.friend_id == "wxid_g1") + ) + ).scalar_one() + assert fetched.gender == 2 + assert fetched.is_block is True + finally: + await engine.dispose() + + +async def test_upsert_customer_preserves_none_for_bool_and_gender(): + """Missing / None bool & gender stay None (no spurious coercion).""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + rec = { + "friend_id": "wxid_n", + "channel": "wechat", + "gender": None, + "is_block": None, + } + await upsert_customer(session, rec) + await session.commit() + + fetched = ( + await session.execute( + sa.select(Customer).where(Customer.friend_id == "wxid_n") + ) + ).scalar_one() + assert fetched.gender is None + assert fetched.is_block is None + finally: + await engine.dispose() + + +async def test_upsert_customer_keeps_raw_json(): + """raw_json is preserved verbatim (the whole input record).""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + rec = { + "friend_id": "wxid_r", + "channel": "wechat", + "friendWxNickname": "upstream-name", # a non-column key + "nested": {"k": "v"}, + } + await upsert_customer(session, rec) + await session.commit() + + fetched = ( + await session.execute( + sa.select(Customer).where(Customer.friend_id == "wxid_r") + ) + ).scalar_one() + assert fetched.raw_json == rec + finally: + await engine.dispose() + + +async def test_upsert_customers_batch_returns_count(): + """Batch upsert returns the number of records processed.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + recs = [ + {"friend_id": "wxid_b1", "channel": "wechat", "version": "v1"}, + {"friend_id": "wxid_b2", "channel": "wechat", "version": "v1"}, + ] + count = await upsert_customers(session, recs) + await session.commit() + + assert count == 2 + rows = (await session.execute(sa.select(Customer))).scalars().all() + assert {r.friend_id for r in rows} == {"wxid_b1", "wxid_b2"} + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# Message upsert +# --------------------------------------------------------------------------- + + +async def test_upsert_messages_coerces_types(): + """is_send '1' -> True; is_del 0 -> False; same msg_id twice -> 1 row.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + rec = { + "msg_id": "m_coerce", + "channel": "wechat", + "is_send": "1", + "is_del": 0, + "type": "1", + "timestamp": 1_700_000_000_000, + } + await upsert_messages(session, [rec]) + await session.commit() + + fetched = ( + await session.execute( + sa.select(Message).where(Message.msg_id == "m_coerce") + ) + ).scalar_one() + assert fetched.is_send is True + assert fetched.is_del is False + assert fetched.type == 1 + assert fetched.timestamp == 1_700_000_000_000 + + # Upsert same msg_id again with different text -> row updated, not + # duplicated. + rec2 = dict(rec, text="updated") + await upsert_messages(session, [rec2]) + await session.commit() + + # ``populate_existing=True`` bypasses the identity-map cache so we + # observe the second upsert's text value. + rows = ( + await session.execute( + sa.select(Message) + .where(Message.msg_id == "m_coerce") + .execution_options(populate_existing=True) + ) + ).scalars().all() + assert len(rows) == 1 + assert rows[0].text == "updated" + finally: + await engine.dispose() + + +async def test_upsert_messages_dedupes_duplicate_msg_id(): + """3 messages with 1 duplicate msg_id -> 2 rows; dup row updated.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + recs = [ + {"msg_id": "d1", "channel": "wechat", "text": "a"}, + {"msg_id": "d2", "channel": "wechat", "text": "b"}, + {"msg_id": "d1", "channel": "wechat", "text": "a-updated"}, + ] + count = await upsert_messages(session, recs) + await session.commit() + + rows = ( + await session.execute( + sa.select(Message).execution_options(populate_existing=True) + ) + ).scalars().all() + assert len(rows) == 2 + assert {r.msg_id for r in rows} == {"d1", "d2"} + # The duplicate's later text wins. + d1 = next(r for r in rows if r.msg_id == "d1") + assert d1.text == "a-updated" + # Return value reflects records processed (input count). + assert count == 3 + finally: + await engine.dispose() + + +async def test_upsert_messages_empty_list_returns_zero(): + """An empty batch is a no-op and reports 0.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + count = await upsert_messages(session, []) + assert count == 0 + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# WechatAccount + embedded Staff upsert (best-effort) +# --------------------------------------------------------------------------- + + +async def test_upsert_wechat_account_coerces_is_on_job(): + """WechatAccount upsert stores wx_id PK and coerces is_on_job.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + rec = { + "wx_id": "wxid_staff", + "nick_name": "Alice", + "user_id": "u1", + "user_name": "Alice Wang", + "is_on_job": 1, + "contacts_num": 100, + } + await upsert_wechat_account(session, rec) + await session.commit() + + fetched = ( + await session.execute( + sa.select(WechatAccount).where( + WechatAccount.wx_id == "wxid_staff" + ) + ) + ).scalar_one() + assert fetched.nick_name == "Alice" + assert fetched.is_on_job is True + assert fetched.contacts_num == 100 + finally: + await engine.dispose() + + +async def test_upsert_staff_from_account_extracts_user_id_and_name(): + """Embedded staff (user_id + user_name) upsert from an account record.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + rec = { + "wx_id": "wxid_staff", + "user_id": "u1", + "user_name": "Alice Wang", + "is_on_job": 1, + } + await upsert_staff_from_account(session, rec) + await session.commit() + + fetched = ( + await session.execute( + sa.select(Staff).where(Staff.user_id == "u1") + ) + ).scalar_one() + assert fetched.name == "Alice Wang" + assert fetched.is_on_job is True + finally: + await engine.dispose() + + +async def test_upsert_staff_from_account_skips_when_no_user_id(): + """No user_id -> nothing inserted (best-effort, no raise).""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + rec = {"wx_id": "wxid_no_user", "nick_name": "anon"} + # Must not raise. + await upsert_staff_from_account(session, rec) + await session.commit() + + rows = (await session.execute(sa.select(Staff))).scalars().all() + assert rows == [] + finally: + await engine.dispose()