main.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. from __future__ import annotations
  2. from contextlib import asynccontextmanager
  3. from typing import AsyncGenerator
  4. from fastapi import FastAPI
  5. from fastapi.middleware.cors import CORSMiddleware
  6. from app.config import settings
  7. from app.db import close_pool, init_pool, get_pool
  8. from app.middleware.logging import LoggingMiddleware
  9. from app.services.ws_hub import hub as ws_hub # noqa: F401
  10. from app.services.scheduler import start_scheduler, stop_scheduler
  11. @asynccontextmanager
  12. async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
  13. await init_pool()
  14. pool = get_pool()
  15. # 建 users 表并初始化 admin
  16. await pool.execute("""
  17. CREATE TABLE IF NOT EXISTS crawl.users (
  18. id BIGSERIAL PRIMARY KEY,
  19. username VARCHAR(100) NOT NULL UNIQUE,
  20. password_hash TEXT NOT NULL,
  21. created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
  22. )
  23. """)
  24. from app.routers.auth import ensure_admin_user
  25. await ensure_admin_user()
  26. # 清理上次进程意外退出遗留的僵死任务
  27. cleaned = await pool.execute(
  28. """
  29. UPDATE scrape_jobs SET status = 'failed', error = '服务重启,任务中断', updated_at = NOW()
  30. WHERE status IN ('pending', 'running')
  31. """
  32. )
  33. import logging
  34. logging.getLogger(__name__).info(f"[startup] 清理僵死任务: {cleaned}")
  35. await start_scheduler()
  36. yield
  37. await stop_scheduler()
  38. await close_pool()
  39. app = FastAPI(title="Sentinel Lens", lifespan=lifespan)
  40. app.add_middleware(
  41. CORSMiddleware,
  42. allow_origins=settings.allowed_origins,
  43. allow_credentials=True,
  44. allow_methods=["*"],
  45. allow_headers=["*"],
  46. )
  47. app.add_middleware(LoggingMiddleware)
  48. # Router registration
  49. from app.routers import stats # noqa: E402
  50. from app.routers import logs # noqa: E402
  51. from app.routers import scrape # noqa: E402
  52. from app.routers import public # noqa: E402
  53. from app.routers import ws # noqa: E402
  54. from app.routers import models # noqa: E402
  55. from app.routers import schedule # noqa: E402
  56. from app.routers import discounts # noqa: E402
  57. from app.routers import auth # noqa: E402
  58. from app.routers import scrape_stats # noqa: E402
  59. app.include_router(stats.router, prefix="/api")
  60. app.include_router(logs.router, prefix="/api")
  61. app.include_router(scrape.router, prefix="/api")
  62. app.include_router(public.router, prefix="/api/public")
  63. app.include_router(models.router, prefix="/api")
  64. app.include_router(schedule.router, prefix="/api")
  65. app.include_router(discounts.router, prefix="/api")
  66. app.include_router(auth.router, prefix="/api")
  67. app.include_router(scrape_stats.router, prefix="/api")
  68. app.include_router(ws.router)
  69. @app.get("/")
  70. async def health_check() -> dict:
  71. return {"status": "ok", "service": "sentinel-lens"}