dbdirectory.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995
  1. from __future__ import annotations
  2. import asyncio
  3. import math
  4. import re
  5. from datetime import UTC, datetime, timedelta
  6. from typing import Any
  7. from uuid import uuid4
  8. from pymongo import ASCENDING, DESCENDING, GEOSPHERE, ReturnDocument
  9. from pymongo.errors import DuplicateKeyError, OperationFailure
  10. from wbb import BOT_PROFILE_ID, control_db
  11. profilesdb = control_db.directory_profiles
  12. locationsdb = control_db.directory_locations
  13. membershipsdb = control_db.directory_memberships
  14. eventsdb = control_db.directory_events
  15. settingsdb = control_db.directory_settings
  16. APPLICATION_NONE = "none"
  17. APPLICATION_PENDING = "pending"
  18. APPLICATION_APPROVED = "approved"
  19. APPLICATION_REJECTED = "rejected"
  20. APPLICATION_REVOKED = "revoked"
  21. APPLICATION_STATUSES = {
  22. APPLICATION_NONE,
  23. APPLICATION_PENDING,
  24. APPLICATION_APPROVED,
  25. APPLICATION_REJECTED,
  26. APPLICATION_REVOKED,
  27. }
  28. PRESENCE_ONLINE = "online"
  29. PRESENCE_OFFLINE = "offline"
  30. DEFAULT_DIRECTORY_SETTINGS: dict[str, Any] = {
  31. "online_duration_hours": 24,
  32. "nearby_radius_km": 20,
  33. "nearby_radius_options_km": [5, 10, 20, 50],
  34. }
  35. _index_lock = asyncio.Lock()
  36. _indexes_ready = False
  37. class DirectoryDataError(ValueError):
  38. def __init__(self, code: str, message: str):
  39. super().__init__(message)
  40. self.code = code
  41. def utc_now() -> datetime:
  42. return datetime.now(UTC)
  43. def aware_utc(value: datetime | None) -> datetime | None:
  44. if value is None:
  45. return None
  46. if value.tzinfo is None:
  47. return value.replace(tzinfo=UTC)
  48. return value.astimezone(UTC)
  49. def display_name(
  50. first_name: str | None,
  51. last_name: str | None,
  52. username: str | None = None,
  53. user_id: int | None = None,
  54. ) -> str:
  55. name = " ".join(
  56. value.strip()
  57. for value in (first_name, last_name)
  58. if value and value.strip()
  59. )
  60. if name:
  61. return name
  62. if username:
  63. return f"@{username}"
  64. return f"用户 {user_id}" if user_id is not None else "未知用户"
  65. async def ensure_directory_indexes() -> None:
  66. global _indexes_ready
  67. if _indexes_ready:
  68. return
  69. async with _index_lock:
  70. if _indexes_ready:
  71. return
  72. await profilesdb.create_index([("user_id", ASCENDING)], unique=True)
  73. await profilesdb.create_index(
  74. [
  75. ("application_status", ASCENDING),
  76. ("listed", ASCENDING),
  77. ("updated_at", DESCENDING),
  78. ]
  79. )
  80. await profilesdb.create_index(
  81. [("presence_status", ASCENDING), ("online_until", ASCENDING)]
  82. )
  83. await locationsdb.create_index([("user_id", ASCENDING)], unique=True)
  84. await locationsdb.create_index([("point", GEOSPHERE)])
  85. await locationsdb.create_index([("updated_at", DESCENDING)])
  86. await membershipsdb.create_index(
  87. [
  88. ("bot_id", ASCENDING),
  89. ("chat_id", ASCENDING),
  90. ("user_id", ASCENDING),
  91. ],
  92. unique=True,
  93. )
  94. await membershipsdb.create_index(
  95. [("user_id", ASCENDING), ("active", ASCENDING), ("verified_at", DESCENDING)]
  96. )
  97. await eventsdb.create_index([("event_id", ASCENDING)], unique=True)
  98. await eventsdb.create_index(
  99. [("idempotency_key", ASCENDING)],
  100. unique=True,
  101. sparse=True,
  102. )
  103. await eventsdb.create_index([("created_at", DESCENDING)])
  104. await eventsdb.create_index([("user_id", ASCENDING), ("created_at", DESCENDING)])
  105. await settingsdb.create_index([("settings_id", ASCENDING)], unique=True)
  106. _indexes_ready = True
  107. async def get_directory_settings() -> dict[str, Any]:
  108. await ensure_directory_indexes()
  109. stored = await settingsdb.find_one({"settings_id": "global"}) or {}
  110. return {
  111. **DEFAULT_DIRECTORY_SETTINGS,
  112. **{
  113. key: stored[key]
  114. for key in DEFAULT_DIRECTORY_SETTINGS
  115. if key in stored
  116. },
  117. }
  118. async def set_directory_settings(values: dict[str, Any]) -> dict[str, Any]:
  119. await ensure_directory_indexes()
  120. current = await get_directory_settings()
  121. target = {**current, **values}
  122. try:
  123. radius_options = sorted(
  124. {
  125. int(value)
  126. for value in target.get("nearby_radius_options_km", [])
  127. if int(value) > 0
  128. }
  129. )
  130. default_radius = int(target.get("nearby_radius_km") or 0)
  131. except (TypeError, ValueError) as exc:
  132. raise DirectoryDataError(
  133. "invalid_radius_options", "附近范围必须是正整数。"
  134. ) from exc
  135. if not radius_options:
  136. raise DirectoryDataError("invalid_radius_options", "附近范围选项不能为空。")
  137. if default_radius not in radius_options:
  138. raise DirectoryDataError(
  139. "invalid_default_radius", "默认附近范围必须属于可选范围。"
  140. )
  141. target.update(
  142. {
  143. "online_duration_hours": 24,
  144. "nearby_radius_km": default_radius,
  145. "nearby_radius_options_km": radius_options,
  146. "updated_at": utc_now(),
  147. }
  148. )
  149. await settingsdb.update_one(
  150. {"settings_id": "global"},
  151. {"$set": target, "$setOnInsert": {"created_at": utc_now()}},
  152. upsert=True,
  153. )
  154. return await get_directory_settings()
  155. async def record_directory_event(
  156. *,
  157. event_type: str,
  158. user_id: int,
  159. actor_id: int | str | None,
  160. actor_name: str = "",
  161. source: str,
  162. reason: str = "",
  163. metadata: dict[str, Any] | None = None,
  164. idempotency_key: str | None = None,
  165. bot_id: str | None = None,
  166. chat_id: int | None = None,
  167. ) -> dict[str, Any] | None:
  168. await ensure_directory_indexes()
  169. document = {
  170. "event_id": uuid4().hex,
  171. "event_type": event_type,
  172. "user_id": int(user_id),
  173. "actor_id": actor_id,
  174. "actor_name": actor_name,
  175. "source": source,
  176. "reason": reason,
  177. "metadata": metadata or {},
  178. "bot_id": bot_id or str(BOT_PROFILE_ID),
  179. "chat_id": int(chat_id) if chat_id is not None else None,
  180. "created_at": utc_now(),
  181. }
  182. if idempotency_key:
  183. document["idempotency_key"] = idempotency_key
  184. try:
  185. await eventsdb.insert_one(document)
  186. except DuplicateKeyError:
  187. return None
  188. return document
  189. async def upsert_directory_identity(
  190. *,
  191. user_id: int,
  192. username: str | None,
  193. first_name: str | None,
  194. last_name: str | None,
  195. source: str = "telegram",
  196. ) -> dict[str, Any]:
  197. await ensure_directory_indexes()
  198. normalized_username = str(username or "").lstrip("@").strip() or None
  199. now = utc_now()
  200. previous = await profilesdb.find_one({"user_id": int(user_id)})
  201. identity = {
  202. "username": normalized_username,
  203. "first_name": str(first_name or "").strip() or None,
  204. "last_name": str(last_name or "").strip() or None,
  205. "display_name": display_name(
  206. first_name,
  207. last_name,
  208. normalized_username,
  209. int(user_id),
  210. ),
  211. }
  212. updates: dict[str, Any] = {**identity, "updated_at": now}
  213. username_removed = bool(previous and previous.get("username") and not normalized_username)
  214. if username_removed and previous.get("application_status") == APPLICATION_APPROVED:
  215. updates.update(
  216. {
  217. "listed": False,
  218. "presence_status": PRESENCE_OFFLINE,
  219. "online_until": None,
  220. }
  221. )
  222. await profilesdb.update_one(
  223. {"user_id": int(user_id)},
  224. {
  225. "$set": updates,
  226. "$setOnInsert": {
  227. "application_status": APPLICATION_NONE,
  228. "listed": False,
  229. "presence_status": PRESENCE_OFFLINE,
  230. "created_at": now,
  231. },
  232. },
  233. upsert=True,
  234. )
  235. if username_removed:
  236. await record_directory_event(
  237. event_type="teacher_username_removed",
  238. user_id=int(user_id),
  239. actor_id=int(user_id),
  240. actor_name=identity["display_name"],
  241. source=source,
  242. reason="技师移除了 Telegram 用户名,系统自动下榜并下线。",
  243. )
  244. return await get_directory_profile(int(user_id)) or {}
  245. async def get_directory_profile(user_id: int) -> dict[str, Any] | None:
  246. await ensure_directory_indexes()
  247. return await profilesdb.find_one({"user_id": int(user_id)})
  248. async def get_directory_location(user_id: int) -> dict[str, Any] | None:
  249. await ensure_directory_indexes()
  250. return await locationsdb.find_one({"user_id": int(user_id)})
  251. def validate_coordinates(longitude: float, latitude: float) -> tuple[float, float]:
  252. try:
  253. lon = float(longitude)
  254. lat = float(latitude)
  255. except (TypeError, ValueError) as exc:
  256. raise DirectoryDataError("invalid_coordinates", "经纬度必须是数字。") from exc
  257. if not math.isfinite(lon) or not -180 <= lon <= 180:
  258. raise DirectoryDataError("invalid_longitude", "经度必须在 -180 到 180 之间。")
  259. if not math.isfinite(lat) or not -90 <= lat <= 90:
  260. raise DirectoryDataError("invalid_latitude", "纬度必须在 -90 到 90 之间。")
  261. return lon, lat
  262. async def save_directory_location(
  263. *,
  264. user_id: int,
  265. longitude: float,
  266. latitude: float,
  267. source: str,
  268. actor_id: int | str,
  269. actor_name: str,
  270. bot_id: str | None = None,
  271. ) -> dict[str, Any]:
  272. await ensure_directory_indexes()
  273. lon, lat = validate_coordinates(longitude, latitude)
  274. now = utc_now()
  275. await locationsdb.update_one(
  276. {"user_id": int(user_id)},
  277. {
  278. "$set": {
  279. "point": {"type": "Point", "coordinates": [lon, lat]},
  280. "longitude": lon,
  281. "latitude": lat,
  282. "source": source,
  283. "source_bot_id": bot_id or str(BOT_PROFILE_ID),
  284. "updated_at": now,
  285. },
  286. "$unset": {
  287. "coordinate_key": "",
  288. "geocode_consent": "",
  289. "geocode_status": "",
  290. "region": "",
  291. "geocoded_at": "",
  292. "geocode_error": "",
  293. },
  294. "$setOnInsert": {"created_at": now},
  295. },
  296. upsert=True,
  297. )
  298. await profilesdb.update_one(
  299. {"user_id": int(user_id)},
  300. {
  301. "$set": {"location_updated_at": now, "updated_at": now},
  302. "$setOnInsert": {
  303. "application_status": APPLICATION_NONE,
  304. "listed": False,
  305. "presence_status": PRESENCE_OFFLINE,
  306. "created_at": now,
  307. },
  308. },
  309. upsert=True,
  310. )
  311. await record_directory_event(
  312. event_type="location_updated",
  313. user_id=int(user_id),
  314. actor_id=actor_id,
  315. actor_name=actor_name,
  316. source=source,
  317. reason="更新位置",
  318. metadata={"location_source": source},
  319. bot_id=bot_id,
  320. )
  321. return await get_directory_location(int(user_id)) or {}
  322. async def clear_directory_location(
  323. *,
  324. user_id: int,
  325. actor_id: int | str,
  326. actor_name: str,
  327. source: str,
  328. reason: str,
  329. ) -> bool:
  330. await ensure_directory_indexes()
  331. result = await locationsdb.delete_one({"user_id": int(user_id)})
  332. await profilesdb.update_one(
  333. {"user_id": int(user_id)},
  334. {
  335. "$set": {
  336. "location_updated_at": None,
  337. "updated_at": utc_now(),
  338. }
  339. },
  340. )
  341. if result.deleted_count:
  342. await record_directory_event(
  343. event_type="location_cleared",
  344. user_id=int(user_id),
  345. actor_id=actor_id,
  346. actor_name=actor_name,
  347. source=source,
  348. reason=reason,
  349. )
  350. return bool(result.deleted_count)
  351. async def submit_teacher_application(
  352. *,
  353. user_id: int,
  354. source: str,
  355. bot_id: str,
  356. chat_id: int,
  357. chat_title: str = "",
  358. ) -> tuple[dict[str, Any], bool]:
  359. await ensure_directory_indexes()
  360. profile = await get_directory_profile(int(user_id))
  361. if not profile or not profile.get("username"):
  362. raise DirectoryDataError(
  363. "username_required", "申请技师前必须先设置 Telegram 用户名。"
  364. )
  365. if profile.get("application_status") == APPLICATION_APPROVED:
  366. raise DirectoryDataError("already_teacher", "你已经是技师,无需重复申请。")
  367. if profile.get("application_status") == APPLICATION_REVOKED:
  368. raise DirectoryDataError(
  369. "teacher_revoked", "技师资格已被撤销,只能由管理员重新批准。"
  370. )
  371. if profile.get("application_status") == APPLICATION_PENDING:
  372. return profile, False
  373. now = utc_now()
  374. await profilesdb.update_one(
  375. {"user_id": int(user_id)},
  376. {
  377. "$set": {
  378. "application_status": APPLICATION_PENDING,
  379. "application_reason": "",
  380. "application_source_bot_id": str(bot_id),
  381. "application_source_chat_id": int(chat_id),
  382. "application_source_chat_title": str(chat_title or "").strip(),
  383. "applied_at": now,
  384. "decided_at": None,
  385. "decided_by": None,
  386. "updated_at": now,
  387. }
  388. },
  389. )
  390. await record_directory_event(
  391. event_type="teacher_application_submitted",
  392. user_id=int(user_id),
  393. actor_id=int(user_id),
  394. actor_name=str(profile.get("display_name") or user_id),
  395. source=source,
  396. reason="申请成为技师",
  397. bot_id=bot_id,
  398. chat_id=int(chat_id),
  399. )
  400. return await get_directory_profile(int(user_id)) or {}, True
  401. async def decide_teacher_application(
  402. *,
  403. user_id: int,
  404. action: str,
  405. actor_id: int | str,
  406. actor_name: str,
  407. source: str,
  408. reason: str = "",
  409. authorization_chat_id: int | None = None,
  410. ) -> dict[str, Any]:
  411. await ensure_directory_indexes()
  412. profile = await get_directory_profile(int(user_id))
  413. if not profile:
  414. raise DirectoryDataError("profile_not_found", "未找到该成员资料。")
  415. action_status = {
  416. "approve": APPLICATION_APPROVED,
  417. "reject": APPLICATION_REJECTED,
  418. "revoke": APPLICATION_REVOKED,
  419. }
  420. if action not in action_status:
  421. raise DirectoryDataError("invalid_decision", "不支持的审批操作。")
  422. if action in {"reject", "revoke"} and not reason.strip():
  423. raise DirectoryDataError("reason_required", "拒绝或撤销时必须填写原因。")
  424. if action == "approve":
  425. if not profile.get("username"):
  426. raise DirectoryDataError(
  427. "username_required", "该成员没有 Telegram 用户名,无法批准。"
  428. )
  429. now = utc_now()
  430. updates: dict[str, Any] = {
  431. "application_status": action_status[action],
  432. "application_reason": reason.strip(),
  433. "decided_at": now,
  434. "decided_by": actor_id,
  435. "decided_by_name": actor_name,
  436. "authorization_chat_id": authorization_chat_id,
  437. "updated_at": now,
  438. }
  439. if action != "approve":
  440. updates.update(
  441. {
  442. "listed": False,
  443. "presence_status": PRESENCE_OFFLINE,
  444. "online_until": None,
  445. }
  446. )
  447. await profilesdb.update_one({"user_id": int(user_id)}, {"$set": updates})
  448. await record_directory_event(
  449. event_type=f"teacher_application_{action}",
  450. user_id=int(user_id),
  451. actor_id=actor_id,
  452. actor_name=actor_name,
  453. source=source,
  454. reason=reason,
  455. metadata={"authorization_chat_id": authorization_chat_id},
  456. chat_id=authorization_chat_id,
  457. )
  458. return await get_directory_profile(int(user_id)) or {}
  459. async def set_teacher_state(
  460. *,
  461. user_id: int,
  462. action: str,
  463. actor_id: int | str,
  464. actor_name: str,
  465. source: str,
  466. reason: str = "",
  467. idempotency_key: str | None = None,
  468. chat_id: int | None = None,
  469. ) -> tuple[dict[str, Any], bool]:
  470. await ensure_directory_indexes()
  471. profile = await get_directory_profile(int(user_id))
  472. if not profile or profile.get("application_status") != APPLICATION_APPROVED:
  473. raise DirectoryDataError("teacher_required", "该成员尚未取得技师资格。")
  474. if action in {"list", "online"}:
  475. if not profile.get("username"):
  476. raise DirectoryDataError("username_required", "请先设置 Telegram 用户名。")
  477. if action == "online" and not profile.get("listed"):
  478. raise DirectoryDataError("listing_required", "请先上榜,再切换为上线状态。")
  479. if action not in {"list", "unlist", "online", "offline"}:
  480. raise DirectoryDataError("invalid_teacher_action", "不支持的技师状态操作。")
  481. event = await record_directory_event(
  482. event_type=f"teacher_{action}",
  483. user_id=int(user_id),
  484. actor_id=actor_id,
  485. actor_name=actor_name,
  486. source=source,
  487. reason=reason or f"技师{action}",
  488. idempotency_key=idempotency_key,
  489. chat_id=chat_id,
  490. )
  491. applied = event is not None
  492. now = utc_now()
  493. updates: dict[str, Any] = {"updated_at": now}
  494. if action == "list":
  495. updates["listed"] = True
  496. updates["listed_at"] = now
  497. elif action == "unlist":
  498. updates.update(
  499. {
  500. "listed": False,
  501. "presence_status": PRESENCE_OFFLINE,
  502. "online_until": None,
  503. }
  504. )
  505. elif action == "online":
  506. updates.update(
  507. {
  508. "presence_status": PRESENCE_ONLINE,
  509. "online_until": now + timedelta(hours=24),
  510. "last_online_at": now,
  511. }
  512. )
  513. else:
  514. updates.update(
  515. {"presence_status": PRESENCE_OFFLINE, "online_until": None}
  516. )
  517. await profilesdb.update_one({"user_id": int(user_id)}, {"$set": updates})
  518. return await get_directory_profile(int(user_id)) or {}, applied
  519. async def expire_teacher_presence() -> int:
  520. await ensure_directory_indexes()
  521. now = utc_now()
  522. expired = 0
  523. while True:
  524. item = await profilesdb.find_one_and_update(
  525. {
  526. "presence_status": PRESENCE_ONLINE,
  527. "online_until": {"$lte": now},
  528. },
  529. {
  530. "$set": {
  531. "presence_status": PRESENCE_OFFLINE,
  532. "online_until": None,
  533. "updated_at": now,
  534. }
  535. },
  536. sort=[("online_until", ASCENDING)],
  537. return_document=ReturnDocument.BEFORE,
  538. )
  539. if not item:
  540. break
  541. expired += 1
  542. previous_until = aware_utc(item.get("online_until"))
  543. await record_directory_event(
  544. event_type="teacher_auto_offline",
  545. user_id=int(item["user_id"]),
  546. actor_id="system",
  547. actor_name="系统",
  548. source="scheduler",
  549. reason="上线状态已超过 24 小时。",
  550. idempotency_key=(
  551. f"teacher-auto-offline:{item['user_id']}:"
  552. f"{previous_until.isoformat() if previous_until else 'unknown'}"
  553. ),
  554. )
  555. return expired
  556. def teacher_is_online(profile: dict[str, Any], *, now: datetime | None = None) -> bool:
  557. current = now or utc_now()
  558. until = aware_utc(profile.get("online_until"))
  559. return bool(
  560. profile.get("presence_status") == PRESENCE_ONLINE
  561. and until
  562. and until > current
  563. )
  564. def _haversine_meters(
  565. longitude_a: float,
  566. latitude_a: float,
  567. longitude_b: float,
  568. latitude_b: float,
  569. ) -> float:
  570. radius = 6_371_008.8
  571. lon_a, lat_a, lon_b, lat_b = map(
  572. math.radians,
  573. (longitude_a, latitude_a, longitude_b, latitude_b),
  574. )
  575. delta_lon = lon_b - lon_a
  576. delta_lat = lat_b - lat_a
  577. value = (
  578. math.sin(delta_lat / 2) ** 2
  579. + math.cos(lat_a) * math.cos(lat_b) * math.sin(delta_lon / 2) ** 2
  580. )
  581. return 2 * radius * math.asin(math.sqrt(value))
  582. async def _list_teachers_fallback(
  583. *,
  584. longitude: float,
  585. latitude: float,
  586. max_distance_meters: float | None,
  587. query: str,
  588. page: int,
  589. page_size: int,
  590. ) -> tuple[list[dict[str, Any]], int]:
  591. profile_filter: dict[str, Any] = {
  592. "application_status": APPLICATION_APPROVED,
  593. "listed": True,
  594. "username": {"$nin": [None, ""]},
  595. }
  596. if query:
  597. pattern = re.compile(re.escape(query.lstrip("@")), re.IGNORECASE)
  598. profile_filter["$or"] = [
  599. {"username": pattern},
  600. {"display_name": pattern},
  601. {"user_id": int(query)} if query.isdigit() else {"user_id": -1},
  602. ]
  603. profiles = {
  604. int(item["user_id"]): item
  605. async for item in profilesdb.find(profile_filter)
  606. }
  607. values: list[dict[str, Any]] = []
  608. if profiles:
  609. async for location in locationsdb.find(
  610. {"user_id": {"$in": list(profiles)}}
  611. ):
  612. point = location.get("point", {}).get("coordinates") or []
  613. if len(point) != 2:
  614. continue
  615. distance = _haversine_meters(
  616. longitude,
  617. latitude,
  618. float(point[0]),
  619. float(point[1]),
  620. )
  621. if max_distance_meters is not None and distance > max_distance_meters:
  622. continue
  623. values.append(
  624. {
  625. **profiles[int(location["user_id"])],
  626. "distance_meters": distance,
  627. "location_updated_at": location.get("updated_at"),
  628. }
  629. )
  630. values.sort(key=lambda item: (item["distance_meters"], int(item["user_id"])))
  631. total = len(values)
  632. start = (page - 1) * page_size
  633. return values[start : start + page_size], total
  634. async def list_directory_teachers(
  635. *,
  636. longitude: float,
  637. latitude: float,
  638. max_distance_meters: float | None = None,
  639. query: str = "",
  640. page: int = 1,
  641. page_size: int = 10,
  642. ) -> tuple[list[dict[str, Any]], int]:
  643. await ensure_directory_indexes()
  644. lon, lat = validate_coordinates(longitude, latitude)
  645. profile_match: dict[str, Any] = {
  646. "profile.application_status": APPLICATION_APPROVED,
  647. "profile.listed": True,
  648. "profile.username": {"$nin": [None, ""]},
  649. }
  650. if query:
  651. pattern = re.compile(re.escape(query.lstrip("@")), re.IGNORECASE)
  652. profile_match["$or"] = [
  653. {"profile.username": pattern},
  654. {"profile.display_name": pattern},
  655. {"profile.user_id": int(query)} if query.isdigit() else {"profile.user_id": -1},
  656. ]
  657. geo_near: dict[str, Any] = {
  658. "near": {"type": "Point", "coordinates": [lon, lat]},
  659. "distanceField": "distance_meters",
  660. "spherical": True,
  661. "key": "point",
  662. }
  663. if max_distance_meters is not None:
  664. geo_near["maxDistance"] = float(max_distance_meters)
  665. pipeline = [
  666. {"$geoNear": geo_near},
  667. {
  668. "$lookup": {
  669. "from": profilesdb.name,
  670. "localField": "user_id",
  671. "foreignField": "user_id",
  672. "as": "profile",
  673. }
  674. },
  675. {"$unwind": "$profile"},
  676. {"$match": profile_match},
  677. {"$sort": {"distance_meters": 1, "user_id": 1}},
  678. {
  679. "$facet": {
  680. "items": [
  681. {"$skip": (max(1, page) - 1) * max(1, page_size)},
  682. {"$limit": max(1, page_size)},
  683. ],
  684. "count": [{"$count": "total"}],
  685. }
  686. },
  687. ]
  688. try:
  689. results = await locationsdb.aggregate(pipeline).to_list(length=1)
  690. result = results[0] if results else {"items": [], "count": []}
  691. items = [
  692. {
  693. **item["profile"],
  694. "distance_meters": item["distance_meters"],
  695. "location_updated_at": item.get("updated_at"),
  696. }
  697. for item in result.get("items", [])
  698. ]
  699. count = result.get("count") or []
  700. return items, int(count[0]["total"]) if count else 0
  701. except (NotImplementedError, OperationFailure):
  702. return await _list_teachers_fallback(
  703. longitude=lon,
  704. latitude=lat,
  705. max_distance_meters=max_distance_meters,
  706. query=query,
  707. page=page,
  708. page_size=page_size,
  709. )
  710. async def list_teacher_applications(
  711. *,
  712. status: str = "",
  713. query: str = "",
  714. source_bot_id: str = "",
  715. source_chat_ids: list[int] | None = None,
  716. page: int = 1,
  717. page_size: int = 20,
  718. ) -> tuple[list[dict[str, Any]], int]:
  719. await ensure_directory_indexes()
  720. filters: dict[str, Any] = {}
  721. if status:
  722. if status not in APPLICATION_STATUSES:
  723. raise DirectoryDataError("invalid_application_status", "申请状态无效。")
  724. filters["application_status"] = status
  725. else:
  726. filters["application_status"] = {"$ne": APPLICATION_NONE}
  727. if source_bot_id:
  728. filters["application_source_bot_id"] = str(source_bot_id)
  729. if source_chat_ids is not None:
  730. filters["application_source_chat_id"] = {
  731. "$in": [int(value) for value in source_chat_ids]
  732. }
  733. if query:
  734. pattern = re.compile(re.escape(query.lstrip("@")), re.IGNORECASE)
  735. filters["$or"] = [
  736. {"username": pattern},
  737. {"display_name": pattern},
  738. {"user_id": int(query)} if query.isdigit() else {"user_id": -1},
  739. ]
  740. total = await profilesdb.count_documents(filters)
  741. items = await (
  742. profilesdb.find(filters)
  743. .sort([("applied_at", DESCENDING), ("updated_at", DESCENDING)])
  744. .skip((page - 1) * page_size)
  745. .limit(page_size)
  746. .to_list(length=page_size)
  747. )
  748. location_map = {
  749. int(item["user_id"]): item
  750. async for item in locationsdb.find(
  751. {"user_id": {"$in": [int(item["user_id"]) for item in items]}}
  752. )
  753. }
  754. for item in items:
  755. location = location_map.get(int(item["user_id"]))
  756. item["location"] = (
  757. {
  758. "source": location.get("source"),
  759. "updated_at": location.get("updated_at"),
  760. }
  761. if location
  762. else None
  763. )
  764. item["online"] = teacher_is_online(item)
  765. return items, total
  766. async def list_directory_profiles(
  767. *,
  768. query: str = "",
  769. application_status: str = "",
  770. listed: bool | None = None,
  771. source_bot_id: str = "",
  772. source_chat_ids: list[int] | None = None,
  773. page: int = 1,
  774. page_size: int = 20,
  775. ) -> tuple[list[dict[str, Any]], int]:
  776. filters: dict[str, Any] = {}
  777. if application_status:
  778. filters["application_status"] = application_status
  779. if listed is not None:
  780. filters["listed"] = listed
  781. if source_bot_id:
  782. filters["application_source_bot_id"] = str(source_bot_id)
  783. if source_chat_ids is not None:
  784. filters["application_source_chat_id"] = {
  785. "$in": [int(value) for value in source_chat_ids]
  786. }
  787. if query:
  788. pattern = re.compile(re.escape(query.lstrip("@")), re.IGNORECASE)
  789. filters["$or"] = [
  790. {"username": pattern},
  791. {"display_name": pattern},
  792. {"user_id": int(query)} if query.isdigit() else {"user_id": -1},
  793. ]
  794. total = await profilesdb.count_documents(filters)
  795. items = await (
  796. profilesdb.find(filters)
  797. .sort([("updated_at", DESCENDING)])
  798. .skip((page - 1) * page_size)
  799. .limit(page_size)
  800. .to_list(length=page_size)
  801. )
  802. locations = {
  803. int(item["user_id"]): item
  804. async for item in locationsdb.find(
  805. {"user_id": {"$in": [int(item["user_id"]) for item in items]}}
  806. )
  807. }
  808. for item in items:
  809. location = locations.get(int(item["user_id"]))
  810. item["location"] = (
  811. {
  812. "source": location.get("source"),
  813. "updated_at": location.get("updated_at"),
  814. }
  815. if location
  816. else None
  817. )
  818. item["online"] = teacher_is_online(item)
  819. return items, total
  820. async def list_directory_locations(
  821. *,
  822. query: str = "",
  823. page: int = 1,
  824. page_size: int = 20,
  825. ) -> tuple[list[dict[str, Any]], int]:
  826. await ensure_directory_indexes()
  827. filters: dict[str, Any] = {}
  828. if query:
  829. pattern = re.compile(re.escape(query.lstrip("@")), re.IGNORECASE)
  830. profile_filters = {
  831. "$or": [
  832. {"username": pattern},
  833. {"display_name": pattern},
  834. {"user_id": int(query)} if query.isdigit() else {"user_id": -1},
  835. ]
  836. }
  837. user_ids = [
  838. int(item["user_id"])
  839. async for item in profilesdb.find(profile_filters, {"user_id": 1})
  840. ]
  841. filters["user_id"] = {"$in": user_ids}
  842. total = await locationsdb.count_documents(filters)
  843. items = await (
  844. locationsdb.find(filters)
  845. .sort("updated_at", DESCENDING)
  846. .skip((page - 1) * page_size)
  847. .limit(page_size)
  848. .to_list(length=page_size)
  849. )
  850. profiles = {
  851. int(item["user_id"]): item
  852. async for item in profilesdb.find(
  853. {"user_id": {"$in": [int(item["user_id"]) for item in items]}}
  854. )
  855. }
  856. for item in items:
  857. item["profile"] = profiles.get(int(item["user_id"]), {})
  858. return items, total
  859. async def list_directory_events(
  860. *,
  861. query: str = "",
  862. event_type: str = "",
  863. page: int = 1,
  864. page_size: int = 20,
  865. ) -> tuple[list[dict[str, Any]], int]:
  866. filters: dict[str, Any] = {}
  867. if event_type:
  868. filters["event_type"] = event_type
  869. if query:
  870. pattern = re.compile(re.escape(query), re.IGNORECASE)
  871. filters["$or"] = [
  872. {"actor_name": pattern},
  873. {"reason": pattern},
  874. {"user_id": int(query)} if query.isdigit() else {"user_id": -1},
  875. ]
  876. total = await eventsdb.count_documents(filters)
  877. items = await (
  878. eventsdb.find(filters)
  879. .sort("created_at", DESCENDING)
  880. .skip((page - 1) * page_size)
  881. .limit(page_size)
  882. .to_list(length=page_size)
  883. )
  884. return items, total
  885. async def upsert_directory_membership(
  886. *,
  887. bot_id: str,
  888. chat_id: int,
  889. user_id: int,
  890. status: str,
  891. active: bool,
  892. chat_title: str = "",
  893. username: str | None = None,
  894. display_name_value: str = "",
  895. verified: bool = False,
  896. ) -> dict[str, Any]:
  897. await ensure_directory_indexes()
  898. now = utc_now()
  899. values: dict[str, Any] = {
  900. "status": status,
  901. "active": bool(active),
  902. "username": username,
  903. "display_name": display_name_value,
  904. "observed_at": now,
  905. "updated_at": now,
  906. }
  907. if chat_title:
  908. values["chat_title"] = chat_title
  909. if verified:
  910. values["verified_at"] = now
  911. await membershipsdb.update_one(
  912. {
  913. "bot_id": str(bot_id),
  914. "chat_id": int(chat_id),
  915. "user_id": int(user_id),
  916. },
  917. {"$set": values, "$setOnInsert": {"created_at": now}},
  918. upsert=True,
  919. )
  920. return await membershipsdb.find_one(
  921. {
  922. "bot_id": str(bot_id),
  923. "chat_id": int(chat_id),
  924. "user_id": int(user_id),
  925. }
  926. ) or {}
  927. async def list_membership_candidates(user_id: int) -> list[dict[str, Any]]:
  928. await ensure_directory_indexes()
  929. return await membershipsdb.find(
  930. {"user_id": int(user_id), "active": True}
  931. ).sort("verified_at", DESCENDING).to_list(length=200)