diff --git a/fastapi/Dockerfile b/fastapi/Dockerfile index a9bc39a..23b0b24 100644 --- a/fastapi/Dockerfile +++ b/fastapi/Dockerfile @@ -25,9 +25,6 @@ COPY . . # Create a new user and switch to that user RUN useradd -ms /bin/bash pgbouncer && \ chown -R pgbouncer:pgbouncer /etc/pgbouncer - -RUN pip install gevent - CMD python -c "\ import os;\ from urllib.parse import urlparse;\ @@ -48,6 +45,5 @@ print(f'export DB_NAME={result.path[1:]}')" > /tmp/env.sh && \ echo "auth_file = /etc/pgbouncer/userlist.txt" >> /etc/pgbouncer/pgbouncer.ini && \ echo "client_login_timeout = 120" >> /etc/pgbouncer/pgbouncer.ini && \ gosu pgbouncer pgbouncer /etc/pgbouncer/pgbouncer.ini & \ - gunicorn app.main:app --workers 3 --worker-class gevent --bind 0.0.0.0:80 --timeout 120 --keep-alive 120 - + exec uvicorn app.main:app --host 0.0.0.0 --port 80 --workers 2 EXPOSE 80 \ No newline at end of file diff --git a/fastapi/app/main.py b/fastapi/app/main.py index 53991d9..3dcd6cf 100644 --- a/fastapi/app/main.py +++ b/fastapi/app/main.py @@ -18,7 +18,7 @@ from datetime import timedelta, date, datetime -from fastapi import FastAPI, Request, Response, Depends, HTTPException, status, Query +from fastapi import FastAPI, Request, Response, Depends, HTTPException, status, Query, Websocket,WebSocketDisconnect from fastapi import Path as FastAPIPath # from fastapi import FastAPI, Request, Response, Depends, HTTPException, status from fastapi.encoders import jsonable_encoder @@ -398,8 +398,48 @@ async def get_list_of_field_values(agency_id: AgencyIdEnum, field: VehiclePositi raise HTTPException(status_code=404, detail="Data not found") return data - - +@app.websocket("/ws/{agency_id}/vehicle_positions") +async def websocket_endpoint(websocket: WebSocket, agency_id: str, async_db: AsyncSession = Depends(get_async_db)): + await websocket.accept() + try: + while True: + try: + data = await asyncio.wait_for(crud.get_all_data_async(async_db, models.VehiclePositions, agency_id), timeout=120) + if data is not None: + await websocket.send_json(data) + await asyncio.sleep(10) + # Send a ping every 10 seconds + await websocket.send_json({"type": "ping"}) + except asyncio.TimeoutError: + raise HTTPException(status_code=408, detail="Request timed out") + except WebSocketDisconnect: + # Handle the WebSocket disconnect event + print("WebSocket disconnected") + + +@app.websocket("/ws/{agency_id}/vehicle_positions/{field}/{ids}") +async def websocket_vehicle_positions_by_ids(websocket: WebSocket, agency_id: AgencyIdEnum, field: VehiclePositionsFieldsEnum, ids: str, async_db: AsyncSession = Depends(get_async_db)): + await websocket.accept() + model = models.VehiclePositions + ids = ids.split(',') + try: + while True: + data = {} + for id in ids: + try: + result = await asyncio.wait_for(crud.get_data_async(async_db, model, agency_id.value, field.value, id), timeout=120) + if result is not None: + data[id] = result + except asyncio.TimeoutError: + raise HTTPException(status_code=408, detail="Request timed out") + if data: + await websocket.send_json(data) + await asyncio.sleep(5) + # Send a ping every 5 seconds + await websocket.send_json({"type": "ping"}) + except WebSocketDisconnect: + # Handle the WebSocket disconnect event + print("WebSocket disconnected") ##### todo: Needs to be tested