Coverage for services/data-ingestion-web/src/main.py: 69%

1from datetime import datetime 

2from fastapi import FastAPI 

3from fastapi.middleware.cors import CORSMiddleware 

4from contextlib import asynccontextmanager 

5import logging 

6 

7from .config import settings 

8from .models.health import HealthResponse 

9from .db import get_pool, close_pool 

10from .routers import sessions 

11 

12logging.basicConfig(level=logging.INFO) 

13logger = logging.getLogger(__name__) 

14 

15SERVICE_NAME = "data-ingestion-web" 

16SERVICE_VERSION = "0.1.0" 

17SERVICE_PORT = 8004 

18 

19 

20@asynccontextmanager 

21async def lifespan(app: FastAPI): 

22 """Handle application lifecycle""" 

23 # Startup 

24 logger.info(f"Starting {SERVICE_NAME} v{SERVICE_VERSION}") 

25 try: 

26 await get_pool() 

27 logger.info("Database connection pool initialized") 

28 except Exception as e: 

29 logger.error(f"Failed to initialize database pool: {e}") 

30 

31 yield 

32 

33 # Shutdown 

34 logger.info(f"Shutting down {SERVICE_NAME}") 

35 await close_pool() 

36 logger.info("Database connection pool closed") 

37 

38 

39app = FastAPI( 

40 title=f"Heimdall SDR - {SERVICE_NAME}", 

41 version=SERVICE_VERSION, 

42 lifespan=lifespan, 

43) 

44 

45app.add_middleware( 

46 CORSMiddleware, 

47 allow_origins=["*"], 

48 allow_credentials=True, 

49 allow_methods=["*"], 

50 allow_headers=["*"], 

51) 

52 

53# Include routers 

54app.include_router(sessions.router) 

55 

56# Register routers 

57app.include_router(sessions.router) 

58 

59 

60@app.get("/") 

61async def root(): 

62 return { 

63 "service": SERVICE_NAME, 

64 "status": "running", 

65 "timestamp": datetime.utcnow().isoformat(), 

66 "version": SERVICE_VERSION, 

67 } 

68 

69 

70@app.get("/health") 

71async def health_check(): 

72 return HealthResponse( 

73 status="healthy", 

74 service=SERVICE_NAME, 

75 version=SERVICE_VERSION, 

76 timestamp=datetime.utcnow() 

77 ) 

78 

79 

80@app.get("/ready") 

81async def readiness_check(): 

82 # Check database connection 

83 try: 

84 pool = await get_pool() 

85 async with pool.acquire() as conn: 

86 await conn.fetchval("SELECT 1") 

87 return {"ready": True, "database": "connected"} 

88 except Exception as e: 

89 logger.error(f"Readiness check failed: {e}") 

90 return {"ready": False, "database": "disconnected", "error": str(e)} 

91 

92 

93if __name__ == "__main__": 

94 import uvicorn 

95 uvicorn.run(app, host="0.0.0.0", port=SERVICE_PORT)