From 084749b6f906ed9332ba451e11b667a94af46ed2 Mon Sep 17 00:00:00 2001
From: Taha Yassine Kraiem <tahayk2@gmail.com>
Date: Tue, 28 Jan 2025 16:21:55 +0100
Subject: [PATCH] refactor(chalice): return file_key with first mob

---
 api/chalicelib/core/modules/__init__.py       |  6 ++++-
 .../core/sessions/sessions_replay.py          | 19 ++--------------
 ee/api/chalicelib/core/modules/__init__.py    | 22 ++++++++++++++++++-
 3 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/api/chalicelib/core/modules/__init__.py b/api/chalicelib/core/modules/__init__.py
index 9b76dd704..37928a80a 100644
--- a/api/chalicelib/core/modules/__init__.py
+++ b/api/chalicelib/core/modules/__init__.py
@@ -1,2 +1,6 @@
 TENANT_CONDITION = "TRUE"
-MOB_KEY=""
\ No newline at end of file
+MOB_KEY = ""
+
+
+def get_file_key(project_id, session_id):
+    pass
diff --git a/api/chalicelib/core/sessions/sessions_replay.py b/api/chalicelib/core/sessions/sessions_replay.py
index a796dbb09..7c1f7aff4 100644
--- a/api/chalicelib/core/sessions/sessions_replay.py
+++ b/api/chalicelib/core/sessions/sessions_replay.py
@@ -4,7 +4,7 @@ from chalicelib.core import events, metadata, events_mobile, \
 from chalicelib.core.sessions import sessions_mobs, sessions_devtool
 from chalicelib.utils import errors_helper
 from chalicelib.utils import pg_client, helper
-from chalicelib.core.modules import MOB_KEY
+from chalicelib.core.modules import MOB_KEY, get_file_key
 
 
 def __is_mobile_session(platform):
@@ -21,23 +21,8 @@ def __group_metadata(session, project_metadata):
 
 
 def get_pre_replay(project_id, session_id):
-    with pg_client.PostgresClient() as cur:
-        query = cur.mogrify(
-            f"""\
-            SELECT encode(file_key,'hex') AS file_key
-            FROM public.sessions
-            WHERE project_id = %(project_id)s
-                AND session_id = %(session_id)s;""",
-            {"project_id": project_id, "session_id": session_id}
-        )
-        cur.execute(query=query)
-
-        data = cur.fetchone()
-        file_key = None
-        if data is not None:
-            file_key = data['file_key']
     return {
-        'fileKey': file_key,
+        **get_file_key(project_id=project_id, session_id=session_id),
         'domURL': [sessions_mobs.get_first_url(project_id=project_id, session_id=session_id, check_existence=False)]}
 
 
diff --git a/ee/api/chalicelib/core/modules/__init__.py b/ee/api/chalicelib/core/modules/__init__.py
index d99c6d8b6..77e73fd1f 100644
--- a/ee/api/chalicelib/core/modules/__init__.py
+++ b/ee/api/chalicelib/core/modules/__init__.py
@@ -1,2 +1,22 @@
 TENANT_CONDITION = "tenant_id = %(tenant_id)s"
-MOB_KEY="encode(file_key,'hex') AS file_key,"
\ No newline at end of file
+MOB_KEY = "encode(file_key,'hex') AS file_key,"
+
+
+def get_file_key(project_id, session_id):
+    from chalicelib.utils import pg_client
+    with pg_client.PostgresClient() as cur:
+        query = cur.mogrify(
+            f"""\
+            SELECT encode(file_key,'hex') AS file_key
+            FROM public.sessions
+            WHERE project_id = %(project_id)s
+                AND session_id = %(session_id)s;""",
+            {"project_id": project_id, "session_id": session_id}
+        )
+        cur.execute(query=query)
+
+        data = cur.fetchone()
+        file_key = None
+        if data is not None:
+            file_key = data['file_key']
+    return {"fileKey": file_key}