Added a test function for SSE
This commit is contained in:
@ -1,25 +1,27 @@
|
||||
import json
|
||||
import asyncio
|
||||
from fastapi import APIRouter, Depends
|
||||
from sqlalchemy.orm import Session
|
||||
from starlette.responses import StreamingResponse
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
from app.models import JobStatus, Jobs as JobModel
|
||||
from app.dependencies import get_db
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
async def job_event_generator(db):
|
||||
async def job_event_generator(db: Session):
|
||||
while True:
|
||||
# Fetch jobs with status TODO
|
||||
jobs = db.query(JobModel).filter(JobModel.status == JobStatus.TODO).all()
|
||||
|
||||
if jobs:
|
||||
yield json.dumps([job.parameters for job in jobs]) + "\n\n"
|
||||
await asyncio.sleep(5)
|
||||
# It's recommended to explicitly communicate IDs clearly
|
||||
job_payload = [{"id": job.id, "parameters": job.parameters} for job in jobs]
|
||||
yield f"data: {json.dumps(job_payload)}\n\n"
|
||||
|
||||
await asyncio.sleep(5) # A reasonable heartbeat/refresh
|
||||
|
||||
|
||||
@router.get("/jobs/stream")
|
||||
async def stream_jobs(db: Session = Depends(get_db)):
|
||||
# test
|
||||
return StreamingResponse(job_event_generator(db), media_type="text/event-stream")
|
||||
|
Reference in New Issue
Block a user