Merge pull request #396 from LACMTA:2023-api-optimization

Reimplement websockets

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

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