From 815966cb2500231c8e114f2cd92aad420883aeeb Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Wed, 1 May 2024 04:39:45 +0530
Subject: [PATCH] Unify, modularize DB adapters to get automation metadata by
 user further

---
 src/khoj/database/adapters/__init__.py | 43 ++++++++++++++------------
 src/khoj/routers/api.py                | 20 +++---------
 2 files changed, 28 insertions(+), 35 deletions(-)

diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py
index ee72a76a..048df839 100644
--- a/src/khoj/database/adapters/__init__.py
+++ b/src/khoj/database/adapters/__init__.py
@@ -929,21 +929,31 @@ class AutomationAdapters:
             if automation.id.startswith(f"automation_{user.uuid}_"):
                 yield automation
 
+    @staticmethod
+    def get_automation_metadata(user: KhojUser, automation: Job):
+        # Perform validation checks
+        # Check if user is allowed to delete this automation id
+        if not automation.id.startswith(f"automation_{user.uuid}_"):
+            raise ValueError("Invalid automation id")
+
+        automation_metadata = json.loads(automation.name)
+        crontime = automation_metadata["crontime"]
+        timezone = automation.next_run_time.strftime("%Z")
+        schedule = f"{cron_descriptor.get_description(crontime)} {timezone}"
+        return {
+            "id": automation.id,
+            "subject": automation_metadata["subject"],
+            "query_to_run": re.sub(r"^/automated_task\s*", "", automation_metadata["query_to_run"]),
+            "scheduling_request": automation_metadata["scheduling_request"],
+            "schedule": schedule,
+            "crontime": crontime,
+            "next": automation.next_run_time.strftime("%Y-%m-%d %I:%M %p %Z"),
+        }
+
     @staticmethod
     def get_automations_metadata(user: KhojUser):
         for automation in AutomationAdapters.get_automations(user):
-            automation_metadata = json.loads(automation.name)
-            crontime = automation_metadata["crontime"]
-            timezone = automation.next_run_time.strftime("%Z")
-            schedule = f"{cron_descriptor.get_description(crontime)} {timezone}"
-            yield {
-                "id": automation.id,
-                "subject": automation_metadata["subject"],
-                "query_to_run": re.sub(r"^/automated_task\s*", "", automation_metadata["query_to_run"]),
-                "scheduling_request": automation_metadata["scheduling_request"],
-                "schedule": schedule,
-                "next": automation.next_run_time.strftime("%Y-%m-%d %I:%M %p %Z"),
-            }
+            yield AutomationAdapters.get_automation_metadata(user, automation)
 
     @staticmethod
     def get_automation(user: KhojUser, automation_id: str) -> Job:
@@ -964,12 +974,7 @@ class AutomationAdapters:
         automation: Job = AutomationAdapters.get_automation(user, automation_id)
 
         # Collate info about user automation to be deleted
-        automation_metadata = json.loads(automation.name)
-        automation_info = {
-            "id": automation.id,
-            "name": automation_metadata["query_to_run"],
-            "next": automation.next_run_time.strftime("%Y-%m-%d %I:%M %p %Z"),
-        }
+        automation_metadata = AutomationAdapters.get_automation_metadata(user, automation)
 
         automation.remove()
-        return automation_info
+        return automation_metadata
diff --git a/src/khoj/routers/api.py b/src/khoj/routers/api.py
index d3a0ef1d..b7c5d5ab 100644
--- a/src/khoj/routers/api.py
+++ b/src/khoj/routers/api.py
@@ -470,16 +470,7 @@ async def post_automation(
         )
 
     # Collate info about the created user automation
-    schedule = f'{cron_descriptor.get_description(crontime)} {automation.next_run_time.strftime("%Z")}'
-    automation_info = {
-        "id": automation.id,
-        "subject": subject,
-        "query_to_run": query_to_run,
-        "scheduling_request": query_to_run,
-        "schedule": schedule,
-        "crontime": crontime,
-        "next": automation.next_run_time.strftime("%Y-%m-%d %I:%M %p %Z"),
-    }
+    automation_info = AutomationAdapters.get_automation_metadata(user, automation)
 
     # Return information about the created automation as a JSON response
     return Response(content=json.dumps(automation_info), media_type="application/json", status_code=200)
@@ -538,12 +529,9 @@ def edit_job(
     if automation.trigger != trigger:
         automation.reschedule(trigger=trigger)
 
-    # Collate info about the modified user automation
-    automation_info = {
-        "id": automation.id,
-        "name": automation.name,
-        "next": automation.next_run_time.strftime("%Y-%m-%d %H:%MS"),
-    }
+    # Collate info about the updated user automation
+    automation = AutomationAdapters.get_automation(user, automation.id)
+    automation_info = AutomationAdapters.get_automation_metadata(user, automation)
 
     # Return modified automation information as a JSON response
     return Response(content=json.dumps(automation_info), media_type="application/json", status_code=200)