main.py 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243
  1. from typing import Any
  2. from fastapi import Depends, FastAPI, Header, HTTPException, Query
  3. from fastapi.middleware.cors import CORSMiddleware
  4. from pydantic import BaseModel, Field
  5. from .auth import create_token, revoke_token, validate_token, verify_credentials
  6. from .config import settings
  7. from .db import ensure_annotation_table
  8. from .services.data_service import data_service
  9. from .services.pretrain_service import pretrain_service
  10. class TimePointInput(BaseModel):
  11. index: int = Field(ge=0)
  12. sampleTime: str
  13. files: dict[str, Any]
  14. class WaveWindowInput(BaseModel):
  15. pointName: str
  16. measurementTypes: list[str]
  17. points: list[TimePointInput]
  18. maxPoints: int = Field(default=200000, ge=256, le=200000)
  19. noSampling: bool = False
  20. class AnnotationInput(BaseModel):
  21. wave_file_id: int = Field(gt=0)
  22. label: str = Field(min_length=1, max_length=8)
  23. period_start: int = Field(gt=0)
  24. period_end: int = Field(gt=0)
  25. sample_index_start: int = Field(ge=0)
  26. sample_index_end: int = Field(ge=0)
  27. class AnnotationQueryInput(BaseModel):
  28. wave_file_ids: list[int] = Field(default=[])
  29. class PretrainRecordInput(BaseModel):
  30. id: int = Field(gt=0)
  31. point_name: str = Field(min_length=1)
  32. measurement_type: str = Field(min_length=1)
  33. rpm: float | None = None
  34. sample_time: str = Field(min_length=1)
  35. file_name: str = Field(default="")
  36. class LoginInput(BaseModel):
  37. username: str = Field(min_length=1)
  38. password: str = Field(min_length=1)
  39. app = FastAPI(
  40. title="压缩机故障预测波形服务",
  41. version="0.1.0",
  42. )
  43. app.add_middleware(
  44. CORSMiddleware,
  45. allow_origins=[settings.cors_origin, "http://127.0.0.1:5173"],
  46. allow_credentials=True,
  47. allow_methods=["*"],
  48. allow_headers=["*"],
  49. )
  50. def require_auth(authorization: str | None = Header(default=None)) -> str:
  51. token = (authorization or "").removeprefix("Bearer ").strip()
  52. if not validate_token(token):
  53. raise HTTPException(status_code=401, detail="未登录或登录已过期")
  54. return token
  55. @app.on_event("startup")
  56. def startup() -> None:
  57. try:
  58. ensure_annotation_table()
  59. except Exception:
  60. # 数据库不可达时跳过建表,演示模式仍可正常启动。
  61. pass
  62. @app.post("/api/login")
  63. def login(payload: LoginInput) -> dict[str, Any]:
  64. if not verify_credentials(payload.username, payload.password):
  65. raise HTTPException(status_code=401, detail="用户名或密码错误")
  66. return {"token": create_token(), "username": payload.username}
  67. @app.post("/api/logout")
  68. def logout(authorization: str | None = Header(default=None)) -> dict[str, Any]:
  69. token = (authorization or "").removeprefix("Bearer ").strip()
  70. revoke_token(token)
  71. return {"ok": True}
  72. @app.get("/api/health")
  73. def health(_auth: str = Depends(require_auth)) -> dict[str, Any]:
  74. return data_service.health()
  75. @app.get("/api/query-options")
  76. def query_options(_auth: str = Depends(require_auth)) -> dict[str, Any]:
  77. try:
  78. return data_service.query_options()
  79. except Exception as error:
  80. raise HTTPException(status_code=503, detail=str(error)) from error
  81. @app.get("/api/time-points")
  82. def time_points(
  83. point_name: str = Query(min_length=1),
  84. measurement_types: list[str] = Query(default=[]),
  85. min_time: str | None = None,
  86. max_time: str | None = None,
  87. _auth: str = Depends(require_auth),
  88. ) -> dict[str, Any]:
  89. try:
  90. return data_service.time_points(point_name, measurement_types, min_time, max_time)
  91. except ValueError as error:
  92. raise HTTPException(status_code=400, detail=str(error)) from error
  93. except Exception as error:
  94. raise HTTPException(status_code=503, detail=str(error)) from error
  95. @app.post("/api/wave-window")
  96. def wave_window(payload: WaveWindowInput, _auth: str = Depends(require_auth)) -> dict[str, Any]:
  97. try:
  98. return data_service.wave_window(
  99. payload.pointName,
  100. payload.measurementTypes,
  101. [point.model_dump() if hasattr(point, "model_dump") else point.dict() for point in payload.points],
  102. payload.maxPoints,
  103. payload.noSampling,
  104. )
  105. except ValueError as error:
  106. raise HTTPException(status_code=400, detail=str(error)) from error
  107. except Exception as error:
  108. raise HTTPException(status_code=503, detail=str(error)) from error
  109. @app.get("/api/wave-files/{wave_file_id}/periods/{period_number}")
  110. def period_detail(
  111. wave_file_id: int,
  112. period_number: int,
  113. _auth: str = Depends(require_auth),
  114. ) -> dict[str, Any]:
  115. try:
  116. return data_service.period_detail(wave_file_id, period_number)
  117. except ValueError as error:
  118. raise HTTPException(status_code=400, detail=str(error)) from error
  119. except Exception as error:
  120. raise HTTPException(status_code=503, detail=str(error)) from error
  121. @app.get("/api/annotation-config")
  122. def annotation_config(_auth: str = Depends(require_auth)) -> dict[str, Any]:
  123. return data_service.annotation_config()
  124. @app.get("/api/annotations")
  125. def list_annotations(
  126. wave_file_ids: list[int] = Query(default=[]),
  127. _auth: str = Depends(require_auth),
  128. ) -> dict[str, Any]:
  129. try:
  130. return data_service.list_annotations(wave_file_ids)
  131. except ValueError as error:
  132. raise HTTPException(status_code=400, detail=str(error)) from error
  133. except Exception as error:
  134. raise HTTPException(status_code=503, detail=str(error)) from error
  135. @app.post("/api/annotations/query")
  136. def query_annotations(
  137. payload: AnnotationQueryInput,
  138. _auth: str = Depends(require_auth),
  139. ) -> dict[str, Any]:
  140. try:
  141. return data_service.list_annotations(payload.wave_file_ids)
  142. except ValueError as error:
  143. raise HTTPException(status_code=400, detail=str(error)) from error
  144. except Exception as error:
  145. raise HTTPException(status_code=503, detail=str(error)) from error
  146. @app.post("/api/annotations")
  147. def create_annotation(payload: AnnotationInput, _auth: str = Depends(require_auth)) -> dict[str, Any]:
  148. try:
  149. return data_service.create_annotation(payload.model_dump())
  150. except ValueError as error:
  151. raise HTTPException(status_code=400, detail=str(error)) from error
  152. except Exception as error:
  153. raise HTTPException(status_code=503, detail=str(error)) from error
  154. @app.delete("/api/annotations/{annotation_id}")
  155. def delete_annotation(annotation_id: int, _auth: str = Depends(require_auth)) -> dict[str, Any]:
  156. try:
  157. return data_service.delete_annotation(annotation_id)
  158. except ValueError as error:
  159. raise HTTPException(status_code=404, detail=str(error)) from error
  160. except Exception as error:
  161. raise HTTPException(status_code=503, detail=str(error)) from error
  162. @app.get("/api/pretrain-records")
  163. def pretrain_records_status(
  164. ids: list[int] = Query(default=[]),
  165. _auth: str = Depends(require_auth),
  166. ) -> dict[str, Any]:
  167. try:
  168. return {"recorded": pretrain_service.recorded_ids(ids)}
  169. except Exception as error:
  170. raise HTTPException(status_code=503, detail=str(error)) from error
  171. @app.post("/api/pretrain-records")
  172. def create_pretrain_record(
  173. payload: PretrainRecordInput,
  174. _auth: str = Depends(require_auth),
  175. ) -> dict[str, Any]:
  176. try:
  177. recorded, already = pretrain_service.record(payload.model_dump())
  178. except ValueError as error:
  179. raise HTTPException(status_code=400, detail=str(error)) from error
  180. except Exception as error:
  181. raise HTTPException(status_code=503, detail=str(error)) from error
  182. return {"recorded": recorded, "already": already}
  183. @app.delete("/api/pretrain-records/{file_id}")
  184. def delete_pretrain_record(file_id: int, _auth: str = Depends(require_auth)) -> dict[str, Any]:
  185. try:
  186. deleted = pretrain_service.remove(file_id)
  187. except ValueError as error:
  188. raise HTTPException(status_code=400, detail=str(error)) from error
  189. except Exception as error:
  190. raise HTTPException(status_code=503, detail=str(error)) from error
  191. if not deleted:
  192. raise HTTPException(status_code=404, detail="CSV 中不存在该记录")
  193. return {"deleted": deleted}