Tolan/dailyCount/run_user_daily_activity.py

#!/usr/bin/env python3
"""
用户日活跃聚合脚本：从 user_events、messages、users 汇总到 user_daily_activity。
配置：dailyCount/setting.conf（MySQL、interval_minutes、active_events）。
用法：
  python run_user_daily_activity.py              # 仅跑当天(UTC)
  python run_user_daily_activity.py --days 7     # 补跑最近7天
  python run_user_daily_activity.py --daemon      # 按 interval_minutes 循环执行
"""
import argparse
import configparser
import os
import sys
from datetime import date, datetime, timedelta, timezone

try:
    import pymysql
except ImportError:
    print("请安装 pymysql: pip install pymysql", file=sys.stderr)
    sys.exit(1)

# 脚本所在目录，用于定位 setting.conf
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DEFAULT_CONF = os.path.join(SCRIPT_DIR, "setting.conf")


def load_config(conf_path: str) -> configparser.ConfigParser:
    if not os.path.isfile(conf_path):
        print(f"配置文件不存在: {conf_path}", file=sys.stderr)
        sys.exit(1)
    # 关闭默认的 % 插值功能，避免像 "page.%__enter" 这样的值触发插值错误
    cfg = configparser.ConfigParser(interpolation=None)
    cfg.read(conf_path, encoding="utf-8")
    return cfg


def get_mysql_config(cfg: configparser.ConfigParser) -> dict:
    return {
        "host": cfg.get("mysql", "host", fallback="127.0.0.1"),
        "port": cfg.getint("mysql", "port", fallback=3306),
        "user": cfg.get("mysql", "user", fallback="root"),
        "password": cfg.get("mysql", "password", fallback=""),
        "database": cfg.get("mysql", "database", fallback="Lyrria"),
        "charset": "utf8mb4",
    }


def parse_active_events(active_events_str: str) -> tuple[str, list]:
    """
    解析 active_events 配置。
    返回 (sql_condition_type, args):
      - ("like", [pattern]) 表示 event_name LIKE %s
      - ("in", [name1, name2, ...]) 表示 event_name IN (%s, %s, ...)
    """
    raw = (active_events_str or "").strip()
    if not raw:
        return "in", []  # 无配置则无匹配
    parts = [p.strip() for p in raw.split(",") if p.strip()]
    if not parts:
        return "in", []
    # 若任一含 % 则整体按 LIKE 处理（通常只配一个 pattern，如 page.%__enter）
    if any("%" in p for p in parts):
        return "like", [parts[0]]
    return "in", parts


def build_active_events_where(cursor, active_events_str: str) -> tuple[str, list]:
    """返回 (WHERE 片段, 参数列表)，用于 user_events 的活跃事件过滤。"""
    cond_type, args = parse_active_events(active_events_str)
    if cond_type == "like":
        return " AND event_name LIKE %s", args
    if cond_type == "in" and args:
        placeholders = ", ".join(["%s"] * len(args))
        return f" AND event_name IN ({placeholders})", args
    return " AND 1=0", []  # 无匹配条件


def run_once(conn, activity_date: date, active_events_where: str, active_events_args: list) -> int:
    """对给定 UTC 日期执行一次聚合，返回写入/更新的行数。"""
    date_str = activity_date.isoformat()
    cursor = conn.cursor()

    # 统一使用 UTC 日期：先设置会话时区
    cursor.execute("SET SESSION time_zone = '+00:00'")

    # 1) 当天有事件或消息的所有 user_id
    cursor.execute(
        """
        SELECT user_id FROM user_events WHERE DATE(event_time) = %s
        UNION
        SELECT user_id FROM messages WHERE DATE(created_at) = %s AND role = 'user'
        """,
        (date_str, date_str),
    )
    user_ids = [row[0] for row in cursor.fetchall()]
    if not user_ids:
        return 0

    # 2) 活跃事件统计：active_flag、session_count（仅统计活跃事件且 session_id 非空）
    active_where = " DATE(event_time) = %s " + active_events_where
    active_params = [date_str] + active_events_args
    cursor.execute(
        """
        SELECT user_id,
               LEAST(1, COUNT(*)) AS active_flag,
               COUNT(DISTINCT session_id) AS session_count
        FROM user_events
        WHERE """ + active_where + """
        GROUP BY user_id
        """,
        active_params,
    )
    active_map = {row[0]: {"active_flag": int(row[1]), "session_count": int(row[2])} for row in cursor.fetchall()}

    # 3) 当天用户消息数 (role='user')
    cursor.execute(
        """
        SELECT user_id, COUNT(*) AS message_count
        FROM messages
        WHERE DATE(created_at) = %s AND role = 'user'
        GROUP BY user_id
        """,
        (date_str,),
    )
    message_map = {row[0]: int(row[1]) for row in cursor.fetchall()}

    # 4) 用户快照：first_seen_date, user_segment, org_id
    placeholders = ", ".join(["%s"] * len(user_ids))
    cursor.execute(
        f"""
        SELECT user_id,
               DATE(created_at) AS first_seen_date,
               user_segment,
               org_id
        FROM users
        WHERE user_id IN ({placeholders})
        """,
        user_ids,
    )
    user_snapshot = {row[0]: {"first_seen_date": row[1], "user_segment": row[2], "org_id": row[3]} for row in cursor.fetchall()}

    # 5) 合并并写入
    rows = []
    for uid in user_ids:
        active = active_map.get(uid, {"active_flag": 0, "session_count": 0})
        msg_count = message_map.get(uid, 0)
        snap = user_snapshot.get(uid, {})
        first_seen = snap.get("first_seen_date")
        first_seen_str = first_seen.isoformat() if first_seen else None
        seg = snap.get("user_segment") or None
        org = snap.get("org_id") or None
        rows.append(
            (
                date_str,
                uid,
                active["active_flag"],
                active["session_count"],
                msg_count,
                first_seen_str,
                seg,
                org,
            )
        )

    insert_sql = """
    INSERT INTO user_daily_activity
    (activity_date, user_id, active_flag, session_count, message_count,
     first_seen_date, user_segment_snapshot, org_id_snapshot)
    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
    ON DUPLICATE KEY UPDATE
        active_flag = VALUES(active_flag),
        session_count = VALUES(session_count),
        message_count = VALUES(message_count),
        user_segment_snapshot = VALUES(user_segment_snapshot),
        org_id_snapshot = VALUES(org_id_snapshot),
        first_seen_date = IFNULL(first_seen_date, VALUES(first_seen_date)),
        updated_at = CURRENT_TIMESTAMP
    """
    cursor.executemany(insert_sql, rows)
    conn.commit()
    count = cursor.rowcount if cursor.rowcount else len(rows)
    cursor.close()
    return count


def main() -> None:
    parser = argparse.ArgumentParser(description="用户日活跃聚合 -> user_daily_activity")
    parser.add_argument("--config", default=DEFAULT_CONF, help="setting.conf 路径")
    parser.add_argument("--days", type=int, default=0, help="补跑最近 N 天（含今天），0 表示仅今天")
    parser.add_argument(
        "--once",
        action="store_true",
        help="仅执行一次（默认按 interval_minutes 守护循环执行）",
    )
    args = parser.parse_args()

    cfg = load_config(args.config)
    mysql_kw = get_mysql_config(cfg)
    active_events_str = cfg.get("user_daily_activity", "active_events", fallback="page.%__enter")
    interval_minutes = cfg.getint("schedule", "interval_minutes", fallback=60)

    active_where, active_args = build_active_events_where(None, active_events_str)

    def do_run() -> None:
        conn = pymysql.connect(**mysql_kw)
        try:
            today_utc = datetime.now(timezone.utc).date()
            if args.days > 0:
                dates = [today_utc - timedelta(days=i) for i in range(args.days)]
            else:
                dates = [today_utc]
            total = 0
            for d in dates:
                n = run_once(conn, d, active_where, active_args)
                total += n
                print(f"[user_daily_activity] {d.isoformat()} 写入/更新 {n} 条")
            if total > 0:
                print(f"[user_daily_activity] 合计 {total} 条")
        finally:
            conn.close()

    if args.once:
        # 仅执行一次（可结合 --days 使用）
        do_run()
    else:
        # 默认守护模式：根据 interval_minutes 周期性执行
        import time
        print(f"[user_daily_activity] 守护模式，每 {interval_minutes} 分钟执行一次")
        while True:
            do_run()
            time.sleep(interval_minutes * 60)


if __name__ == "__main__":
    main()