main.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  1. """FastAPI 应用入口。"""
  2. from __future__ import annotations
  3. import asyncio
  4. import logging
  5. from contextlib import asynccontextmanager
  6. from typing import AsyncIterator
  7. from fastapi import FastAPI
  8. from app.config import settings
  9. from app.docs_auth import setup_protected_docs
  10. from app.exception_handlers import register_exception_handlers
  11. from app.openapi import setup_custom_openapi
  12. from app.response import ApiResponse
  13. from app.routers import search
  14. from app.schemas import CacheRefreshResult, CacheStatus
  15. from app.services.cache import (
  16. WORKER_ID,
  17. get_snapshot,
  18. leader_heartbeat,
  19. refresh_cache,
  20. scheduler_loop,
  21. try_become_leader,
  22. wait_cache_ready,
  23. )
  24. from app.services.redis_client import (
  25. KEY_LOCK_LEADER,
  26. close_redis,
  27. release_lock,
  28. )
  29. logging.basicConfig(
  30. level=logging.INFO,
  31. format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
  32. )
  33. logger = logging.getLogger(__name__)
  34. @asynccontextmanager
  35. async def lifespan(app: FastAPI) -> AsyncIterator[None]:
  36. """启动:选举 leader → 仅 leader 刷新缓存与跑定时任务;关闭:释放资源。"""
  37. is_leader = await try_become_leader()
  38. tasks: list[asyncio.Task[None]] = []
  39. if is_leader:
  40. logger.info("[%s] 当选 leader,开始首次刷新", WORKER_ID)
  41. await refresh_cache()
  42. tasks.append(
  43. asyncio.create_task(scheduler_loop(), name="faq-scheduler")
  44. )
  45. tasks.append(
  46. asyncio.create_task(leader_heartbeat(), name="leader-heartbeat")
  47. )
  48. logger.info("[%s] 定时任务与 leader 心跳已启动", WORKER_ID)
  49. else:
  50. logger.info("[%s] 不是 leader,等待缓存就绪…", WORKER_ID)
  51. ready = await wait_cache_ready()
  52. if ready:
  53. logger.info("[%s] 共享缓存已就绪", WORKER_ID)
  54. else:
  55. logger.warning(
  56. "[%s] 等待缓存就绪超时,搜索将以空 sec_ids 降级运行",
  57. WORKER_ID,
  58. )
  59. try:
  60. yield
  61. finally:
  62. logger.info("[%s] 关闭:取消后台任务…", WORKER_ID)
  63. for task in tasks:
  64. task.cancel()
  65. for task in tasks:
  66. try:
  67. await task
  68. except asyncio.CancelledError:
  69. pass
  70. if is_leader:
  71. # 主动释放 leader 锁,让其他 worker 可以更快接管
  72. await release_lock(KEY_LOCK_LEADER, WORKER_ID)
  73. await close_redis()
  74. # 仅开发环境暴露 API 文档;生产环境一律关闭,避免泄露接口结构
  75. _docs_enabled = settings.app_env.lower() == "dev"
  76. app = FastAPI(
  77. title="Zendesk FAQ Keyword Search",
  78. description="封装 Zendesk Help Center 搜索;FAQ",
  79. version="1.0.0",
  80. lifespan=lifespan,
  81. # 关闭默认 docs,由 setup_protected_docs 注册受 HTTP Basic 保护的等价路由
  82. docs_url=None,
  83. redoc_url=None,
  84. openapi_url=None,
  85. )
  86. # 注册全局异常处理器(统一错误响应格式 {code, data, msg})
  87. register_exception_handlers(app)
  88. # 覆写 OpenAPI 文档:让所有错误响应(4xx/5xx)使用统一信封格式
  89. setup_custom_openapi(app)
  90. # 注册受 HTTP Basic 认证保护的 /faq/docs、/faq/redoc、/faq/openapi.json
  91. setup_protected_docs(app)
  92. # 路由
  93. app.include_router(search.router)
  94. @app.get(
  95. "/faq/health",
  96. response_model=ApiResponse[dict[str, str]],
  97. tags=["meta"],
  98. )
  99. async def health() -> ApiResponse[dict[str, str]]:
  100. """健康检查。"""
  101. return ApiResponse.ok(data={"status": "ok"})
  102. @app.get(
  103. "/faq/health/cache",
  104. response_model=ApiResponse[CacheStatus],
  105. tags=["meta"],
  106. )
  107. async def cache_status() -> ApiResponse[CacheStatus]:
  108. """缓存状态查看。"""
  109. snap = await get_snapshot()
  110. return ApiResponse.ok(
  111. data=CacheStatus(
  112. sec_ids_count=len(snap["sec_ids"]),
  113. last_updated_at=snap["last_updated_at"],
  114. next_refresh_at=snap["next_refresh_at"],
  115. )
  116. )
  117. @app.post(
  118. "/faq/admin/cache/refresh",
  119. response_model=ApiResponse[CacheRefreshResult],
  120. tags=["meta"],
  121. )
  122. async def refresh_cache_now() -> ApiResponse[CacheRefreshResult]:
  123. """手动触发刷新(运维 / 测试用途)。
  124. 使用 force=True 跳过去重锁,确保管理员请求立即生效。
  125. """
  126. await refresh_cache(force=True)
  127. snap = await get_snapshot()
  128. return ApiResponse.ok(
  129. data=CacheRefreshResult(
  130. sec_ids_count=len(snap["sec_ids"]),
  131. last_updated_at=snap["last_updated_at"],
  132. ),
  133. msg="缓存刷新成功",
  134. )