54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
|
"""Seed WeChat + WhatsApp accounts (and their embedded staff) into smoke.db."""
|
||
|
|
import asyncio
|
||
|
|
|
||
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||
|
|
|
||
|
|
from app.db import models # noqa: F401
|
||
|
|
from app.scraper.wechat import scrape_departments, scrape_wechat_accounts
|
||
|
|
from app.scraper.whatsapp import scrape_whatsapp_accounts
|
||
|
|
from app.storage.repos import upsert_wechat_account, upsert_whatsapp_account
|
||
|
|
from scripts.smoke_wechat import _push_settings_into_environ, build_client_from_env
|
||
|
|
|
||
|
|
|
||
|
|
async def _root_department_id(client) -> str | None:
|
||
|
|
"""First root department (parent_id falsy) from the live dept tree.
|
||
|
|
|
||
|
|
``wxUser/list`` REQUIRES the company-root ``departmentId`` filter; we pull
|
||
|
|
it from ``scrape_departments`` rather than hardcoding the live id.
|
||
|
|
"""
|
||
|
|
async for d in scrape_departments(client):
|
||
|
|
if not d.get("parent_id"):
|
||
|
|
return d.get("dept_id")
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
async def main() -> None:
|
||
|
|
_push_settings_into_environ()
|
||
|
|
client = build_client_from_env()
|
||
|
|
engine = create_async_engine("sqlite+aiosqlite:///./smoke.db")
|
||
|
|
sm = async_sessionmaker(engine, expire_on_commit=False)
|
||
|
|
|
||
|
|
root_dept = await _root_department_id(client)
|
||
|
|
if not root_dept:
|
||
|
|
print("WARNING: no root department found; skipping wxUser/list scrape")
|
||
|
|
|
||
|
|
wx_n = wa_n = 0
|
||
|
|
async with sm() as session:
|
||
|
|
if root_dept:
|
||
|
|
async for rec in scrape_wechat_accounts(
|
||
|
|
client, department_id=root_dept, page_size=50
|
||
|
|
):
|
||
|
|
await upsert_wechat_account(session, rec)
|
||
|
|
wx_n += 1
|
||
|
|
async for rec in scrape_whatsapp_accounts(client, page_size=50):
|
||
|
|
await upsert_whatsapp_account(session, rec)
|
||
|
|
wa_n += 1
|
||
|
|
await session.commit()
|
||
|
|
print(f"upserted {wx_n} wechat accounts, {wa_n} whatsapp accounts (+ their staff)")
|
||
|
|
await client.aclose()
|
||
|
|
await engine.dispose()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
asyncio.run(main())
|