Procházet zdrojové kódy

Merge branch 'v2' of clovis/bdlg2023-sms-back into master

clovis před 1 měsícem
rodič
revize
087d6acabc

+ 4 - 6
app/api/SSEBroadcasterRoute.py

@@ -1,10 +1,9 @@
-import asyncio
 import json
 
 from fastapi import Request, Response
 from fastapi.routing import APIRoute
 
-from app.core.events import broadcast
+from app.core.events import redis_client
 
 
 class SSEBroadcasterRoute(APIRoute):
@@ -12,13 +11,12 @@ class SSEBroadcasterRoute(APIRoute):
         original_route_handler = super().get_route_handler()
 
         async def custom_route_handler(request: Request) -> Response:
-            # 1. Execute the route and let FastAPI serialize the response
+            # Execute the route and let FastAPI serialize the response
             response = await original_route_handler(request)
 
             openapi_extra = self.openapi_extra or {}
             event_name = openapi_extra.get("sse_event")
-
-            # 2. Check if an event is declared and the request was successful
+            # Check if an event is declared and the request was successful
             if event_name and response.status_code in (200, 201):
                 project_id = request.path_params.get("project_id")
                 client_id = request.headers.get("X-Client-ID", "unknown")
@@ -36,7 +34,7 @@ class SSEBroadcasterRoute(APIRoute):
                 )
 
                 channel = f"project_{project_id}"
-                asyncio.create_task(broadcast.publish(channel, payload_str))
+                await redis_client.publish(channel, payload_str)
 
             return response
 

+ 67 - 1
app/api/deps.py

@@ -3,10 +3,12 @@ from collections.abc import Generator
 from uuid import UUID
 
 import jwt
-from fastapi import Depends, HTTPException, Path, status
+from fastapi import Depends, HTTPException, Path, Query, status
 from fastapi.security import OAuth2PasswordBearer
+from fastapi.security.utils import get_authorization_scheme_param
 from sqlalchemy import exists, select
 from sqlalchemy.orm import Session
+from starlette.requests import Request
 
 from app.api.utils import get_project_organization_id
 from app.core import config, security
@@ -59,6 +61,70 @@ async def get_current_user(
     return user
 
 
+async def get_token_flexible(
+    request: Request,
+    token: str | None = Query(default=None),
+) -> str:
+    """Same as reusable_oauth2, but also accepts ?token=... in the query
+    string — needed because native EventSource cannot set custom headers."""
+    auth_header = request.headers.get("Authorization")
+    if auth_header:
+        scheme, param = get_authorization_scheme_param(auth_header)
+        if scheme.lower() == "bearer" and param:
+            return param
+    if token:
+        return token
+    raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Not authenticated")
+
+
+async def get_current_user_flexible(
+    session: Session = Depends(get_session),
+    token: str = Depends(get_token_flexible),
+) -> User:
+    # identical body to get_current_user, just sourcing `token` differently
+    try:
+        payload = jwt.decode(token, config.settings.SECRET_KEY, algorithms=[security.JWT_ALGORITHM])
+    except jwt.DecodeError:
+        raise HTTPException(status.HTTP_403_FORBIDDEN, "Could not validate credentials.")
+
+    token_data = security.JWTTokenPayload(**payload)
+    if token_data.refresh:
+        raise HTTPException(
+            status.HTTP_403_FORBIDDEN, "Could not validate credentials, cannot use refresh token"
+        )
+
+    now = int(time.time())
+    if now < token_data.issued_at or now > token_data.expires_at:
+        raise HTTPException(
+            status.HTTP_403_FORBIDDEN,
+            "Could not validate credentials, token expired or not yet valid",
+        )
+
+    result = session.execute(select(User).where(User.id == token_data.sub))
+    user = result.scalars().first()
+    if not user:
+        raise HTTPException(status_code=404, detail="User not found.")
+    return user
+
+
+def require_org_role_sse(*allowed_roles: OrgRole):
+    def dependency(
+        project_id: UUID = Path(...),
+        session: Session = Depends(get_session),
+        current_user: User = Depends(get_current_user_flexible),
+    ) -> User:
+        if current_user.global_role == GlobalRole.SUPER_ADMIN:
+            return current_user
+        if not _has_org_role(session, current_user.id, project_id, *allowed_roles):
+            get_project_organization_id(session, project_id)
+            raise HTTPException(
+                status.HTTP_403_FORBIDDEN, "Insufficient permissions for this organization"
+            )
+        return current_user
+
+    return dependency
+
+
 def require_super_admin(current_user: User = Depends(get_current_user)) -> User:
     if current_user.global_role != GlobalRole.SUPER_ADMIN:
         raise HTTPException(status.HTTP_403_FORBIDDEN, "Requires super_admin")

+ 21 - 14
app/api/endpoints/project_stream.py

@@ -4,7 +4,7 @@ from fastapi import APIRouter, Depends, Request
 from fastapi.responses import StreamingResponse
 
 from app.api import deps
-from app.core.events import broadcast
+from app.core.events import redis_client
 from app.models import OrgRole, User
 
 router = APIRouter()
@@ -15,8 +15,11 @@ async def project_stream(
     project_id: str,
     client_id: str,
     request: Request,
+    token: str,
     current_user: User = Depends(
-        deps.require_org_role(OrgRole.ORG_ADMIN, OrgRole.RESPO_BENEVOLE, OrgRole.RESPO_COMMISSION)
+        deps.require_org_role_sse(
+            OrgRole.ORG_ADMIN, OrgRole.RESPO_BENEVOLE, OrgRole.RESPO_COMMISSION
+        )
     ),
 ):
     """
@@ -25,25 +28,29 @@ async def project_stream(
     **Required Query Parameters:**
     - `client_id`: A unique UUID generated by the frontend to prevent echoing.
     """
+    print(f"create event generator for {client_id}")
 
     async def event_generator():
         channel = f"project_{project_id}"
-
-        # Subscribe to the Redis channel specific to this project
-        async with broadcast.subscribe(channel=channel) as subscriber:
-            async for event in subscriber:
-                # If the user closes the tab or navigates away, cleanly close the connection
+        pubsub = redis_client.pubsub()
+        await pubsub.subscribe(channel)
+        print(f"subscribed to {channel}")
+        try:
+            async for event in pubsub.listen():
+                if event["type"] != "message":
+                    continue
                 if await request.is_disconnected():
+                    print("disconnected")
                     break
-                # Parse the raw string payload from Redis
-                message = json.loads(event.message)
-
-                # Echo prevention: Do not send the event back to the client that caused it
+                message = json.loads(event["data"])
                 if message.get("client_id") == client_id:
                     continue
+                yield f"data: {json.dumps({'event': message['event'], 'item': message['item']})}\n\n"
+        finally:
+            await pubsub.unsubscribe(channel)
+            await pubsub.aclose()
+            print(f"unsubscribed from {channel}")
 
-                # Yield the properly formatted Server-Sent Event (SSE)
-                yield f"event: {message['event']}\ndata: {json.dumps(message['item'])}\n\n"
-
+    print(f"return event generator for {client_id}")
     # Return as text/event-stream so the browser knows it's an ongoing connection
     return StreamingResponse(event_generator(), media_type="text/event-stream")

+ 2 - 3
app/core/events.py

@@ -1,6 +1,5 @@
-from broadcaster import Broadcast
+import redis.asyncio as redis
 
 from app.core.config import settings
 
-# Instantiate the broadcaster using the URL from your environment variables
-broadcast = Broadcast(settings.REDIS_URL)
+redis_client = redis.from_url(settings.REDIS_URL, decode_responses=True)

+ 3 - 4
app/main.py

@@ -10,14 +10,13 @@ from fastapi.openapi.utils import get_openapi
 
 from app.api.api import api_router
 from app.core import config
-from app.core.events import broadcast
+from app.core.events import redis_client
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    await broadcast.connect()
     yield
-    await broadcast.disconnect()
+    await redis_client.aclose()
 
 
 app = FastAPI(
@@ -125,5 +124,5 @@ def custom_openapi():
     return app.openapi_schema
 
 
-# 5. Override the default OpenAPI method
+# Override the default OpenAPI method
 app.openapi = custom_openapi