main.py 3.2 KB

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