diff --git a/api/chalicelib/core/sessions.py b/api/chalicelib/core/sessions.py index 6c846eb42..054e7fcb2 100644 --- a/api/chalicelib/core/sessions.py +++ b/api/chalicelib/core/sessions.py @@ -1237,3 +1237,15 @@ def count_all(): with pg_client.PostgresClient(unlimited_query=True) as cur: row = cur.execute(query="SELECT COUNT(session_id) AS count FROM public.sessions") return row.get("count", 0) + + +def session_exists(project_id, session_id): + with pg_client.PostgresClient() as cur: + query = cur.mogrify("""SELECT 1 + FROM public.sessions + WHERE session_id=%(session_id)s + AND project_id=%(project_id)s""", + {"project_id": project_id, "session_id": session_id}) + cur.execute(query) + row = cur.fetchone(query) + return row is not None diff --git a/api/routers/core.py b/api/routers/core.py index b3252e34a..92161f960 100644 --- a/api/routers/core.py +++ b/api/routers/core.py @@ -913,6 +913,8 @@ def get_live_session(projectId: int, sessionId: str, background_tasks: Backgroun @app.get('/{projectId}/assist/sessions/{sessionId}/replay', tags=["assist"]) def get_live_session_replay_file(projectId: int, sessionId: str, context: schemas.CurrentContext = Depends(OR_context)): + if isinstance(sessionId, str) or not sessions.session_exists(project_id=projectId, session_id=sessionId): + return {"errors": ["Replay file not found"]} path = assist.get_raw_mob_by_id(project_id=projectId, session_id=sessionId) if path is None: return {"errors": ["Replay file not found"]} diff --git a/ee/api/chalicelib/core/sessions_exp.py b/ee/api/chalicelib/core/sessions_exp.py index add1a790d..81953bcc5 100644 --- a/ee/api/chalicelib/core/sessions_exp.py +++ b/ee/api/chalicelib/core/sessions_exp.py @@ -1542,3 +1542,14 @@ def count_all(): with pg_client.PostgresClient(unlimited_query=True) as cur: row = cur.execute(query="SELECT COUNT(session_id) AS count FROM public.sessions") return row.get("count", 0) + + +def session_exists(project_id, session_id): + with ch_client.ClickHouseClient() as cur: + query = cur.format("""SELECT 1 + FROM public.sessions + WHERE session_id=%(session_id)s + AND project_id=%(project_id)s""", + {"project_id": project_id, "session_id": session_id}) + row = cur.execute(query) + return row is not None