71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
|
"""Seed chat messages for customers already in the list (so timelines aren't empty).
|
||
|
|
|
||
|
|
Iterates every Customer row in smoke.db, scrapes that conversation's recent
|
||
|
|
messages (WeChat msg/chats or WhatsApp chatRecords), and upserts them. Bounded
|
||
|
|
per-customer so it stays quick.
|
||
|
|
"""
|
||
|
|
import asyncio
|
||
|
|
import sqlite3
|
||
|
|
import time
|
||
|
|
|
||
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||
|
|
|
||
|
|
from app.db import models # noqa: F401
|
||
|
|
from app.scraper.wechat import scrape_wechat_messages
|
||
|
|
from app.scraper.whatsapp import scrape_whatsapp_messages
|
||
|
|
from app.storage.repos import upsert_messages
|
||
|
|
from scripts.smoke_wechat import _push_settings_into_environ, build_client_from_env
|
||
|
|
|
||
|
|
WINDOW_MS = 90 * 86_400_000 # last 90 days
|
||
|
|
MAX_PER_CUSTOMER = 60
|
||
|
|
|
||
|
|
|
||
|
|
async def main() -> None:
|
||
|
|
_push_settings_into_environ()
|
||
|
|
client = build_client_from_env()
|
||
|
|
|
||
|
|
con = sqlite3.connect("smoke.db")
|
||
|
|
rows = con.execute(
|
||
|
|
"SELECT channel, owner_account_id, friend_id FROM customer "
|
||
|
|
"WHERE owner_account_id IS NOT NULL AND friend_id IS NOT NULL"
|
||
|
|
).fetchall()
|
||
|
|
con.close()
|
||
|
|
|
||
|
|
engine = create_async_engine("sqlite+aiosqlite:///./smoke.db")
|
||
|
|
sm = async_sessionmaker(engine, expire_on_commit=False)
|
||
|
|
now = int(time.time() * 1000)
|
||
|
|
total = 0
|
||
|
|
customers_with_msgs = 0
|
||
|
|
|
||
|
|
for channel, owner, friend in rows:
|
||
|
|
recs: list[dict] = []
|
||
|
|
try:
|
||
|
|
if channel == "wechat":
|
||
|
|
agen = scrape_wechat_messages(client, owner, friend, now - WINDOW_MS, now)
|
||
|
|
else:
|
||
|
|
agen = scrape_whatsapp_messages(client, owner, friend, now - WINDOW_MS, now)
|
||
|
|
async for m in agen:
|
||
|
|
recs.append(m)
|
||
|
|
if len(recs) >= MAX_PER_CUSTOMER:
|
||
|
|
break
|
||
|
|
except Exception as e: # noqa: BLE001
|
||
|
|
print(f" [{channel}] {friend}: {type(e).__name__}: {e!r}")
|
||
|
|
continue
|
||
|
|
if recs:
|
||
|
|
async with sm() as session:
|
||
|
|
n = await upsert_messages(session, recs)
|
||
|
|
await session.commit()
|
||
|
|
total += n
|
||
|
|
customers_with_msgs += 1
|
||
|
|
print(f" [{channel}] {friend}: {n} messages")
|
||
|
|
else:
|
||
|
|
print(f" [{channel}] {friend}: (no messages in window)")
|
||
|
|
|
||
|
|
print(f"\nDONE: {total} messages across {customers_with_msgs} customers")
|
||
|
|
await client.aclose()
|
||
|
|
await engine.dispose()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
asyncio.run(main())
|