402 lines
14 KiB
Python
402 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
为指定企业补「订单运营 + 分销推广」种子数据，量级与 wechat_users 池大致匹配（取总用户约 1/12 订单条数作中间值）。
|
||
|
||
- orders：enterpriseId 对齐；orderNo 前缀 SEED5_ 可幂等跳过
|
||
- 为已支付/已完成订单写入一条 test_results（带 orderId），供分销「产品佣金分布」按测评类型归类
|
||
- distribution_agents + distribution_bindings（enterprise 维度）+ commission_records
|
||
|
||
读 api/.env 的 DATABASE_*。
|
||
|
||
用法:
|
||
python3 scripts/seed_ent5_orders_and_distribution_mysql.py --enterprise-id 5
|
||
python3 scripts/seed_ent5_orders_and_distribution_mysql.py --enterprise-id 5 --dry-run
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hashlib
|
||
import random
|
||
import re
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
|
||
try:
|
||
import pymysql
|
||
except ImportError:
|
||
sys.exit("需要: pip install pymysql")
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
ENV_PATH = ROOT / "api" / ".env"
|
||
ORDER_PREFIX = "SEED5_"
|
||
|
||
|
||
def load_env(path: Path) -> dict[str, str]:
|
||
cfg: dict[str, str] = {}
|
||
if not path.is_file():
|
||
return cfg
|
||
for line in path.read_text(encoding="utf-8").splitlines():
|
||
line = line.strip()
|
||
if not line or line.startswith("#"):
|
||
continue
|
||
m = re.match(r"^([A-Z0-9_]+)\s*=\s*(.*)$", line)
|
||
if m:
|
||
cfg[m.group(1)] = m.group(2).strip()
|
||
return cfg
|
||
|
||
|
||
def pick_type(i: int) -> str:
|
||
r = (i * 7919) % 100
|
||
if r < 55:
|
||
return "face"
|
||
if r < 72:
|
||
return "mbti"
|
||
if r < 86:
|
||
return "disc"
|
||
return "pdp"
|
||
|
||
|
||
def product_title(pt: str) -> str:
|
||
return {
|
||
"face": "AI人脸性格分析完整报告",
|
||
"mbti": "MBTI 职业性格测评",
|
||
"disc": "DISC 行为风格测评",
|
||
"pdp": "PDP 天赋特质测评",
|
||
}.get(pt, "测评服务")
|
||
|
||
|
||
def mini_result_json(pt: str) -> str:
|
||
import json
|
||
|
||
if pt == "mbti":
|
||
return json.dumps({"mbtiType": "INTJ", "type": "INTJ"}, ensure_ascii=False)
|
||
if pt == "disc":
|
||
return json.dumps({"description": {"type": "D型"}}, ensure_ascii=False)
|
||
if pt == "pdp":
|
||
return json.dumps({"description": {"type": "老虎型"}}, ensure_ascii=False)
|
||
return json.dumps(
|
||
{"overview": "订单关联占位", "mbti": {"type": "INTJ"}},
|
||
ensure_ascii=False,
|
||
)
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--enterprise-id", type=int, default=5)
|
||
ap.add_argument("--dry-run", action="store_true")
|
||
ap.add_argument("--inviters", type=int, default=80)
|
||
ap.add_argument("--bindings", type=int, default=520)
|
||
args = ap.parse_args()
|
||
eid = args.enterprise_id
|
||
|
||
env = load_env(ENV_PATH)
|
||
host = env.get("DATABASE_HOSTNAME", "127.0.0.1")
|
||
port = int(env.get("DATABASE_HOSTPORT", "3306"))
|
||
user = env.get("DATABASE_USERNAME", "root")
|
||
password = env.get("DATABASE_PASSWORD", "")
|
||
database = env.get("DATABASE_DATABASE", "mbti")
|
||
prefix = env.get("DATABASE_PREFIX", "mbti_")
|
||
orders_t = f"{prefix}orders"
|
||
tr_t = f"{prefix}test_results"
|
||
wu_t = f"{prefix}wechat_users"
|
||
ag_t = f"{prefix}distribution_agents"
|
||
bd_t = f"{prefix}distribution_bindings"
|
||
cr_t = f"{prefix}commission_records"
|
||
|
||
conn = pymysql.connect(
|
||
host=host,
|
||
port=port,
|
||
user=user,
|
||
password=password,
|
||
database=database,
|
||
charset="utf8mb4",
|
||
autocommit=False,
|
||
)
|
||
rnd = random.Random(20260330 + eid)
|
||
now = int(time.time())
|
||
|
||
try:
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
f"SELECT COUNT(*) FROM `{wu_t}` WHERE `enterpriseId`=%s",
|
||
(eid,),
|
||
)
|
||
n_users = int(cur.fetchone()[0])
|
||
cur.execute(
|
||
f"SELECT `id`,`nickname` FROM `{wu_t}` WHERE `enterpriseId`=%s ORDER BY `id` ASC",
|
||
(eid,),
|
||
)
|
||
rows = cur.fetchall()
|
||
user_ids = [int(r[0]) for r in rows]
|
||
nick_map = {int(r[0]): (r[1] or "") for r in rows}
|
||
|
||
if not user_ids:
|
||
print({"ok": False, "error": "no users"})
|
||
return
|
||
|
||
target_orders = max(400, min(1400, n_users // 12))
|
||
existing_seed = 0
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
f"SELECT COUNT(*) FROM `{orders_t}` WHERE `orderNo` LIKE %s",
|
||
(ORDER_PREFIX + "%",),
|
||
)
|
||
existing_seed = int(cur.fetchone()[0])
|
||
|
||
if existing_seed >= target_orders * 0.9 and not args.dry_run:
|
||
print(
|
||
{
|
||
"ok": True,
|
||
"skipped": True,
|
||
"reason": "SEED5 orders already near target",
|
||
"existing_seed_orders": existing_seed,
|
||
"target_orders": target_orders,
|
||
}
|
||
)
|
||
return
|
||
|
||
need_orders = max(0, target_orders - existing_seed)
|
||
if args.dry_run:
|
||
print(
|
||
{
|
||
"ok": True,
|
||
"dry_run": True,
|
||
"n_users": n_users,
|
||
"target_orders": target_orders,
|
||
"existing_seed": existing_seed,
|
||
"would_insert_orders": need_orders,
|
||
"inviters": min(args.inviters, len(user_ids) // 4),
|
||
"bindings_cap": min(args.bindings, len(user_ids) // 2),
|
||
}
|
||
)
|
||
return
|
||
|
||
inviter_n = min(args.inviters, max(20, len(user_ids) // 6))
|
||
inviters = user_ids[:inviter_n]
|
||
pool_bind = user_ids[inviter_n + 50 : inviter_n + 50 + args.bindings]
|
||
if len(pool_bind) < 100:
|
||
pool_bind = user_ids[inviter_n + 10 :]
|
||
|
||
# 1) distribution_agents
|
||
inviter_to_agent: dict[int, int] = {}
|
||
with conn.cursor() as cur:
|
||
for uid in inviters:
|
||
cur.execute(f"SELECT id FROM `{ag_t}` WHERE userId=%s LIMIT 1", (uid,))
|
||
ex = cur.fetchone()
|
||
if ex:
|
||
inviter_to_agent[uid] = int(ex[0])
|
||
continue
|
||
name = (nick_map.get(uid) or f"用户{uid}")[:100]
|
||
cur.execute(
|
||
f"""
|
||
INSERT INTO `{ag_t}`
|
||
(`userId`,`agentName`,`contactPhone`,`contactEmail`,`totalOrders`,`totalCommission`,
|
||
`availableCommission`,`status`,`createdAt`,`updatedAt`)
|
||
VALUES (%s,%s,NULL,NULL,0,0,0,1,%s,%s)
|
||
""",
|
||
(uid, name, now, now),
|
||
)
|
||
inviter_to_agent[uid] = int(cur.lastrowid)
|
||
conn.commit()
|
||
|
||
# 2) orders + linked test_results (paid/completed only)
|
||
order_rows: list[tuple] = []
|
||
expire_at = now + 86400 * 400
|
||
ins_order = f"""
|
||
INSERT INTO `{orders_t}`
|
||
(`orderNo`,`userId`,`enterpriseId`,`productType`,`productTitle`,`amount`,`status`,
|
||
`payMethod`,`payTime`,`createdAt`,`updatedAt`)
|
||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
|
||
"""
|
||
ins_tr = f"""
|
||
INSERT INTO `{tr_t}`
|
||
(`userId`,`testType`,`resultData`,`requiresPayment`,`isPaid`,`paidAmount`,`paidAt`,
|
||
`createdAt`,`updatedAt`,`enterpriseId`,`testScope`,`orderId`)
|
||
VALUES (%s,%s,%s,0,1,%s,%s,%s,%s,%s,'enterprise',%s)
|
||
"""
|
||
|
||
for k in range(need_orders):
|
||
uid = user_ids[(k * 9973) % len(user_ids)]
|
||
pt = pick_type(k)
|
||
amt = rnd.choice([100, 150, 200, 299, 399, 101, 201])
|
||
st = rnd.choices(
|
||
["paid", "completed", "pending"],
|
||
weights=[42, 28, 30],
|
||
k=1,
|
||
)[0]
|
||
ts = now - rnd.randint(0, 86400 * 20)
|
||
pay_t = ts if st in ("paid", "completed") else None
|
||
h = hashlib.md5(f"{uid}{k}{ts}".encode()).hexdigest()[:10]
|
||
ono = f"{ORDER_PREFIX}{ts}{uid}{k}{h}"[:48]
|
||
order_rows.append(
|
||
(ono, uid, eid, pt, product_title(pt), amt, st, "wechat", pay_t, ts, ts)
|
||
)
|
||
|
||
oid_map: list[int] = []
|
||
with conn.cursor() as cur:
|
||
for row in order_rows:
|
||
cur.execute(ins_order, row)
|
||
oid_map.append(int(cur.lastrowid))
|
||
oid = oid_map[-1]
|
||
st = row[6]
|
||
if st not in ("paid", "completed"):
|
||
continue
|
||
uid = row[1]
|
||
pt = row[3]
|
||
amt = row[5]
|
||
ts = row[9]
|
||
pay_t = row[8] or ts
|
||
cur.execute(
|
||
ins_tr,
|
||
(
|
||
uid,
|
||
pt,
|
||
mini_result_json(pt),
|
||
amt,
|
||
pay_t,
|
||
ts,
|
||
ts,
|
||
eid,
|
||
oid,
|
||
),
|
||
)
|
||
conn.commit()
|
||
|
||
# 3) bindings
|
||
bindings_done = 0
|
||
binding_ids: list[tuple[int, int, int]] = [] # id, inviter, invitee
|
||
with conn.cursor() as cur:
|
||
for i, invitee in enumerate(pool_bind[: args.bindings]):
|
||
inv = inviters[i % len(inviters)]
|
||
if invitee == inv:
|
||
continue
|
||
cur.execute(
|
||
f"""
|
||
SELECT id FROM `{bd_t}`
|
||
WHERE inviteeId=%s AND scope='enterprise' AND enterpriseId=%s
|
||
""",
|
||
(invitee, eid),
|
||
)
|
||
if cur.fetchone():
|
||
continue
|
||
tsb = now - rnd.randint(86400, 86400 * 60)
|
||
cur.execute(
|
||
f"""
|
||
INSERT INTO `{bd_t}`
|
||
(`inviterId`,`inviteeId`,`scope`,`enterpriseId`,`expireAt`,`status`,
|
||
`prevInviterId`,`overriddenAt`,`createdAt`,`updatedAt`)
|
||
VALUES (%s,%s,'enterprise',%s,%s,'active',NULL,NULL,%s,%s)
|
||
""",
|
||
(inv, invitee, eid, expire_at, tsb, tsb),
|
||
)
|
||
bid = int(cur.lastrowid)
|
||
binding_ids.append((bid, inv, invitee))
|
||
bindings_done += 1
|
||
conn.commit()
|
||
|
||
# 4) commission_records（与订单/绑定挂钩）
|
||
ins_cr = f"""
|
||
INSERT INTO `{cr_t}`
|
||
(`agentId`,`orderId`,`commissionRate`,`commissionAmount`,`status`,`paidAt`,
|
||
`createdAt`,`updatedAt`,`scope`,`enterpriseId`,`inviterId`,`inviteeId`,
|
||
`bindingId`,`orderAmount`,`commissionFen`,`frozenAt`,`unfrozenAt`,
|
||
`testResultId`,`commissionSource`)
|
||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,'enterprise',%s,%s,%s,%s,%s,%s,NULL,NULL,NULL,%s)
|
||
"""
|
||
cr_n = 0
|
||
with conn.cursor() as cur:
|
||
for j, oid in enumerate(oid_map[: min(len(oid_map), 450)]):
|
||
row = order_rows[j]
|
||
if row[6] not in ("paid", "completed"):
|
||
continue
|
||
uid_buyer = row[1]
|
||
inv = inviters[j % len(inviters)]
|
||
if uid_buyer == inv:
|
||
inv = inviters[(j + 1) % len(inviters)]
|
||
agent_id = inviter_to_agent.get(inv)
|
||
if not agent_id:
|
||
continue
|
||
bid_row = next((b for b in binding_ids if b[2] == uid_buyer), None)
|
||
bid = bid_row[0] if bid_row else None
|
||
oamt = row[5]
|
||
cfen = max(10, min(500, int(oamt * rnd.uniform(0.08, 0.22))))
|
||
st_f = rnd.choices(["paid", "frozen"], weights=[82, 18], k=1)[0]
|
||
paid_at = now - rnd.randint(0, 86400 * 14) if st_f == "paid" else None
|
||
tsr = paid_at or (now - rnd.randint(0, 86400 * 10))
|
||
camt = round(cfen / 100, 2)
|
||
cur.execute(
|
||
ins_cr,
|
||
(
|
||
agent_id,
|
||
oid,
|
||
10.0,
|
||
camt,
|
||
st_f,
|
||
paid_at,
|
||
tsr,
|
||
tsr,
|
||
eid,
|
||
inv,
|
||
uid_buyer,
|
||
bid,
|
||
oamt,
|
||
cfen,
|
||
"payment",
|
||
),
|
||
)
|
||
cr_n += 1
|
||
|
||
extra = min(120, len(binding_ids))
|
||
for j in range(extra):
|
||
bid, inv, invitee = binding_ids[j]
|
||
agent_id = inviter_to_agent.get(inv)
|
||
if not agent_id:
|
||
continue
|
||
cfen = rnd.randint(20, 180)
|
||
st_f = "paid"
|
||
tsr = now - rnd.randint(0, 86400 * 7)
|
||
cur.execute(
|
||
ins_cr,
|
||
(
|
||
agent_id,
|
||
None,
|
||
0.0,
|
||
round(cfen / 100, 2),
|
||
st_f,
|
||
tsr,
|
||
tsr,
|
||
tsr,
|
||
eid,
|
||
inv,
|
||
invitee,
|
||
bid,
|
||
0,
|
||
cfen,
|
||
"seed_invite",
|
||
),
|
||
)
|
||
cr_n += 1
|
||
|
||
conn.commit()
|
||
|
||
print(
|
||
{
|
||
"ok": True,
|
||
"enterpriseId": eid,
|
||
"n_users": n_users,
|
||
"inserted_orders": len(oid_map),
|
||
"bindings_inserted": bindings_done,
|
||
"commission_rows": cr_n,
|
||
"note": "已支付订单各带 1 条 test_results(orderId)，概览「已完成测试」会增加；分销饼图按订单关联测评类型归类。",
|
||
}
|
||
)
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|