From 2fe6159eaa94770b6ca1f491eaa14988e607ba2e Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Fri, 26 Jul 2024 16:59:08 -0700
Subject: [PATCH] add MCDecoratorSyncResponse model, fix typing errors for
 cronjobs?

---
 backend/btrixcloud/operator/cronjobs.py | 36 +++++++++++++------------
 backend/btrixcloud/operator/models.py   | 11 +++++++-
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/backend/btrixcloud/operator/cronjobs.py b/backend/btrixcloud/operator/cronjobs.py
index 773b27546..9b179bf9a 100644
--- a/backend/btrixcloud/operator/cronjobs.py
+++ b/backend/btrixcloud/operator/cronjobs.py
@@ -1,11 +1,11 @@
 """ Operator handler for crawl CronJobs """
 
 from uuid import UUID
-from typing import Optional, Any
+from typing import Optional
 import yaml
 
 from btrixcloud.utils import to_k8s_date, dt_now
-from .models import MCDecoratorSyncData, CJS
+from .models import MCDecoratorSyncData, CJS, MCDecoratorSyncResponse
 from .baseoperator import BaseOperator
 
 from ..models import CrawlConfig
@@ -20,12 +20,14 @@ def init_routes(self, app):
         """init routes for crawl CronJob decorator"""
 
         @app.post("/op/cronjob/sync")
-        async def mc_sync_cronjob_crawls(data: MCDecoratorSyncData):
+        async def mc_sync_cronjob_crawls(
+            data: MCDecoratorSyncData,
+        ) -> MCDecoratorSyncResponse:
             return await self.sync_cronjob_crawl(data)
 
     def get_finished_response(
         self, metadata: dict[str, str], set_status=True, finished: Optional[str] = None
-    ):
+    ) -> MCDecoratorSyncResponse:
         """get final response to indicate cronjob created job is finished"""
 
         if not finished:
@@ -40,14 +42,12 @@ def get_finished_response(
                 "completionTime": finished,
             }
 
-        res = {
-            "attachments": [],
+        return MCDecoratorSyncResponse(
+            attachments=[],
             # set on job to match default behavior when job finishes
-            "annotations": {"finished": finished},
-            "status": status,
-        }
-        print(res)
-        return res
+            annotations={"finished": finished},
+            status=status,
+        )
 
     # pylint: disable=too-many-arguments
     async def make_new_crawljob(
@@ -58,7 +58,7 @@ async def make_new_crawljob(
         crawl_id: str,
         metadata: dict[str, str],
         state: Optional[str],
-    ) -> list[dict[str, Any]]:
+    ) -> MCDecoratorSyncResponse:
         """declare new CrawlJob from cid, based on db data"""
         # cronjob doesn't exist yet
         crawlconfig: CrawlConfig
@@ -127,9 +127,11 @@ async def make_new_crawljob(
             profile_filename=profile_filename or "",
         )
 
-        return list(yaml.safe_load_all(crawljob))
+        return MCDecoratorSyncResponse(attachments=list(yaml.safe_load_all(crawljob)))
 
-    async def sync_cronjob_crawl(self, data: MCDecoratorSyncData):
+    async def sync_cronjob_crawl(
+        self, data: MCDecoratorSyncData
+    ) -> MCDecoratorSyncResponse:
         """create crawljobs from a job object spawned by cronjob"""
 
         metadata = data.object["metadata"]
@@ -163,7 +165,7 @@ async def sync_cronjob_crawl(self, data: MCDecoratorSyncData):
         crawljob_id = f"crawljob-{crawl_id}"
 
         if crawljob_id not in crawljobs:
-            attachments = await self.make_new_crawljob(
+            response = await self.make_new_crawljob(
                 UUID(cid),
                 UUID(oid) if oid else None,
                 UUID(userid) if userid else None,
@@ -182,6 +184,6 @@ async def sync_cronjob_crawl(self, data: MCDecoratorSyncData):
             except KeyError:
                 pass
 
-            attachments = [crawljob]
+            response = MCDecoratorSyncResponse(attachments=[crawljob])
 
-        return {"attachments": attachments}
+        return response
diff --git a/backend/btrixcloud/operator/models.py b/backend/btrixcloud/operator/models.py
index dd0c84655..13a04918d 100644
--- a/backend/btrixcloud/operator/models.py
+++ b/backend/btrixcloud/operator/models.py
@@ -2,7 +2,7 @@
 
 from collections import defaultdict
 from uuid import UUID
-from typing import Optional, DefaultDict, Literal, Annotated
+from typing import Optional, DefaultDict, Literal, Annotated, Any
 from pydantic import BaseModel, Field
 from kubernetes.utils import parse_quantity
 from btrixcloud.models import StorageRef, TYPE_ALL_CRAWL_STATES
@@ -53,6 +53,15 @@ class MCDecoratorSyncData(BaseModel):
     finalizing: bool = False
 
 
+# ============================================================================
+class MCDecoratorSyncResponse(BaseModel):
+    """Response model for decoratorcontroller sync api"""
+
+    attachments: list[dict[str, Any]]
+    status: Optional[dict[str, Any]] = None
+    annotations: Optional[dict[str, str]] = None
+
+
 # ============================================================================
 class CrawlSpec(BaseModel):
     """spec from k8s CrawlJob object"""