shaguabijia-app-server/scripts/load_meituan_coupon_tsv.py

"""把 meituan_coupon 的线上采样 TSV 灌进本地 SQLite。

用途:本地开发/调试时,把线上 `meituan_coupon` 表的采样数据(tests/meituan_coupon_data.tsv)
灌进 dev 库(默认 `./data/app.db`),免得每次都实时打美团接口。

TSV 说明:
  - 制表符分隔,每行一条记录,列顺序与线上 PostgreSQL 物理列一致
    (image_size / image_type 是后加的迁移,排在最后两列 —— 与本地 SQLite 一致)。
  - 空字段 = NULL(文件里没有 `\\N` 标记)。
  - 个别记录的文本/JSON 字段内含换行,会把一条逻辑行拆成多物理行 —— 按“累计到 26 列”重组。
  - 文件尾部可能有一条被导出截断的残行(列数不足 / raw JSON 不完整),直接跳过。

datetime 三列(first_seen/last_seen/updated_at)去掉尾部时区偏移(`+08`),
存成 SQLAlchemy 在 SQLite 上用的朴素格式 `YYYY-MM-DD HH:MM:SS.ffffff`,保证 ORM 能读回。

用法:
    python scripts/load_meituan_coupon_tsv.py                     # 默认 TSV + .env 里的库
    python scripts/load_meituan_coupon_tsv.py path/to.tsv         # 指定 TSV
    DATABASE_URL=sqlite:///./data/app.db python scripts/load_meituan_coupon_tsv.py
"""
from __future__ import annotations

import json
import os
import re
import sqlite3
import sys
from pathlib import Path

_PROJECT_ROOT = Path(__file__).resolve().parent.parent

# 列顺序 = TSV 字段顺序 = 本地 SQLite 物理列顺序
COLS = [
    "id", "source", "platform", "biz_line", "city_id", "product_view_sign",
    "sku_view_id", "name", "brand_name", "sell_price_cents", "original_price_cents",
    "head_url", "sale_volume", "sale_volume_num", "commission_percent",
    "commission_amount_cents", "poi_name", "available_poi_num", "delivery_distance_m",
    "dedup_key", "raw", "first_seen", "last_seen", "updated_at", "image_size", "image_type",
]
NCOL = len(COLS)

# 按列做类型转换(空串 -> None)。未列出的列 = 原样字符串(source/city_id/... 等 NOT NULL 文本)。
_INT_COLS = {0, 2, 3, 9, 10, 13, 15, 17, 24}           # id, platform, biz_line, prices, ...
_FLOAT_COLS = {14, 18}                                  # commission_percent, delivery_distance_m
_NULLABLE_STR_COLS = {6, 7, 8, 11, 12, 16, 25}          # sku_view_id, name, brand_name, ...
_DT_COLS = {21, 22, 23}                                 # first_seen, last_seen, updated_at
_RAW_COL = 20

_TZ_SUFFIX = re.compile(r"[+-]\d{2}(:?\d{2})?$")        # 尾部时区偏移 +08 / +08:00 / +0800


def _resolve_sqlite_path() -> Path:
    """从 DATABASE_URL(env 或 .env)解析出 SQLite 文件路径。只支持 sqlite://。"""
    url = os.environ.get("DATABASE_URL", "")
    if not url:
        env = _PROJECT_ROOT / ".env"
        if env.exists():
            for line in env.read_text(encoding="utf-8").splitlines():
                if line.strip().startswith("DATABASE_URL="):
                    url = line.split("=", 1)[1].strip()
                    break
    if not url:
        url = "sqlite:///./data/app.db"
    if not url.startswith("sqlite:"):
        sys.exit(f"仅支持 sqlite:// 库,当前 DATABASE_URL={url!r}")
    rest = url.split("sqlite:///", 1)[1] if "sqlite:///" in url else url.split("sqlite://", 1)[1]
    p = Path(rest)
    if not p.is_absolute():
        p = (_PROJECT_ROOT / rest).resolve()
    return p


def _reconstruct_rows(text: str) -> tuple[list[list[str]], int]:
    """把文件文本重组成一条条 26 列的逻辑行。返回 (rows, skipped)。

    单个字段内含换行 -> 一条逻辑行被拆成多物理行:累计字段,拆点用 \\n 重新拼回,
    直到凑满 26 列。列数溢出(内嵌 TAB / 错位)或文件尾残行 -> 跳过并计数。
    """
    lines = text.split("\n")
    while lines and lines[-1] == "":
        lines.pop()
    rows: list[list[str]] = []
    skipped = 0
    buf: list[str] = []
    for raw_line in lines:
        parts = raw_line.split("\t")
        if not buf:
            buf = parts
        else:
            buf[-1] += "\n" + parts[0]       # 拼回被换行拆开的字段
            buf.extend(parts[1:])
        if len(buf) == NCOL:
            rows.append(buf)
            buf = []
        elif len(buf) > NCOL:                # 溢出:数据异常,丢弃这段重新开始
            print(f"  [skip] 列数溢出({len(buf)}>{NCOL}),field0={buf[0][:20]!r}")
            skipped += 1
            buf = []
    if buf:                                  # 文件尾被截断的残行
        print(f"  [skip] 尾部残行不足 {NCOL} 列(实 {len(buf)} 列),field0={buf[0][:20]!r}")
        skipped += 1
    return rows, skipped


def _convert(row: list[str]) -> tuple | None:
    """按列类型转换一行;非法(必填 int 为空 / raw 非 JSON)返回 None。"""
    out: list = []
    for i, v in enumerate(row):
        if i in _DT_COLS:
            out.append(_TZ_SUFFIX.sub("", v))
            continue
        if i == _RAW_COL:
            try:
                json.loads(v)
            except Exception as e:
                print(f"  [skip] id={row[0]} raw 非法 JSON: {e}")
                return None
            out.append(v)
            continue
        if i in _INT_COLS:
            out.append(int(v) if v != "" else None)
        elif i in _FLOAT_COLS:
            out.append(float(v) if v != "" else None)
        elif i in _NULLABLE_STR_COLS:
            out.append(v if v != "" else None)
        else:                                # 必填文本列,原样
            out.append(v)
    return tuple(out)


def main() -> None:
    tsv = Path(sys.argv[1]) if len(sys.argv) > 1 else _PROJECT_ROOT / "tests" / "meituan_coupon_data.tsv"
    if not tsv.is_absolute():
        tsv = (_PROJECT_ROOT / tsv).resolve()
    db = _resolve_sqlite_path()
    print(f"TSV: {tsv}")
    print(f"DB : {db}")
    if not tsv.exists():
        sys.exit(f"TSV 不存在: {tsv}")
    if not db.exists():
        sys.exit(f"SQLite 库不存在: {db}(先跑 alembic upgrade head 建表)")

    rows, skipped = _reconstruct_rows(tsv.read_text(encoding="utf-8"))
    print(f"重组逻辑行: {len(rows)}  跳过(残/异常): {skipped}")

    records = []
    bad = 0
    for r in rows:
        rec = _convert(r)
        if rec is None:
            bad += 1
            continue
        records.append(rec)
    print(f"可入库: {len(records)}  转换失败: {bad}")

    placeholders = ",".join(["?"] * NCOL)
    sql = f"INSERT OR REPLACE INTO meituan_coupon ({','.join(COLS)}) VALUES ({placeholders})"
    con = sqlite3.connect(str(db))
    try:
        before = con.execute("SELECT count(*) FROM meituan_coupon").fetchone()[0]
        con.executemany(sql, records)
        con.commit()
        after = con.execute("SELECT count(*) FROM meituan_coupon").fetchone()[0]
    finally:
        con.close()
    print(f"入库前 {before} 行 -> 入库后 {after} 行(本次 {len(records)} 条)")


if __name__ == "__main__":
    main()