From 46fc710bbfd7ef553ed8ccba517c789d0f4f0466 Mon Sep 17 00:00:00 2001
From: Gavin <gaprice@lbl.gov>
Date: Mon, 16 Dec 2024 14:04:02 -0800
Subject: [PATCH] Add callbacks for download and job complete

Not sure if job complete is going to be needed, will delete if not
---
 cdmtaskservice/app.py                |  1 +
 cdmtaskservice/callback_url_paths.py | 17 +++++++++++-
 cdmtaskservice/routes.py             | 39 +++++++++++++++++++++++++++-
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/cdmtaskservice/app.py b/cdmtaskservice/app.py
index 744c31b..693a4b2 100644
--- a/cdmtaskservice/app.py
+++ b/cdmtaskservice/app.py
@@ -68,6 +68,7 @@ def create_app():
     app.include_router(routes.ROUTER_GENERAL)
     app.include_router(routes.ROUTER_JOBS)
     app.include_router(routes.ROUTER_ADMIN)
+    app.include_router(routes.ROUTER_CALLBACKS)
 
     async def build_app_wrapper():
         await app_state.build_app(app, cfg)
diff --git a/cdmtaskservice/callback_url_paths.py b/cdmtaskservice/callback_url_paths.py
index 4cc6209..156dfb3 100644
--- a/cdmtaskservice/callback_url_paths.py
+++ b/cdmtaskservice/callback_url_paths.py
@@ -5,6 +5,7 @@
 
 _CALLBACK = "callback"
 _DOWNLOAD_COMPLETE = "download"
+_JOB_COMPLETE = "job"
 
 
 def get_download_complete_callback(root_url: str = None, job_id: str = None):
@@ -14,7 +15,21 @@ def get_download_complete_callback(root_url: str = None, job_id: str = None):
     root_url - prepend the path with the given root url.
     job_id - suffix the path with a job ID.
     """
+    return _get_callback(_DOWNLOAD_COMPLETE, root_url, job_id)
+
+
+def get_job_complete_callback(root_url: str = None, job_id: str = None):
+    """
+    Get a url or path for a service callback to communicate that a remote job is complete.
+    
+    root_url - prepend the path with the given root url.
+    job_id - suffix the path with a job ID.
+    """
+    return _get_callback(_JOB_COMPLETE, root_url, job_id)
+
+
+def _get_callback(subpath: str, root_url: str = None, job_id: str = None):
     cb = [root_url] if root_url else []
-    cb += [_CALLBACK, _DOWNLOAD_COMPLETE]
+    cb += [_CALLBACK, subpath]
     cb += [job_id] if job_id else []
     return "/".join(cb)
diff --git a/cdmtaskservice/routes.py b/cdmtaskservice/routes.py
index eba429a..2b8c143 100644
--- a/cdmtaskservice/routes.py
+++ b/cdmtaskservice/routes.py
@@ -3,6 +3,7 @@
 """
 
 import datetime
+import logging
 from fastapi import (
     APIRouter,
     Depends,
@@ -12,9 +13,14 @@
 )
 from pydantic import BaseModel, Field
 from typing import Annotated
+
 from cdmtaskservice import app_state
 from cdmtaskservice import kb_auth
 from cdmtaskservice import models
+from cdmtaskservice.callback_url_paths import (
+    get_download_complete_callback,
+    get_job_complete_callback,
+)
 from cdmtaskservice.exceptions import UnauthorizedError
 from cdmtaskservice.git_commit import GIT_COMMIT
 from cdmtaskservice.http_bearer import KBaseHTTPBearer
@@ -25,8 +31,9 @@
 
 # may need to split these into different files if this file gets too big
 ROUTER_GENERAL = APIRouter(tags=["General"])
-ROUTER_ADMIN = APIRouter(tags=["Admin"], prefix="/admin")
 ROUTER_JOBS = APIRouter(tags=["Jobs"], prefix="/jobs")
+ROUTER_ADMIN = APIRouter(tags=["Admin"], prefix="/admin")
+ROUTER_CALLBACKS = APIRouter(tags=["Callbacks"])
 
 _AUTH = KBaseHTTPBearer()
 
@@ -223,5 +230,35 @@ async def get_nersc_client_info(
     )
 
 
+@ROUTER_CALLBACKS.get(
+    f"/{get_download_complete_callback()}/{{job_id}}",
+    summary="Report data download complete",
+    description="Report that data download for a job is complete. This method is not expected "
+        + "to be called by users."
+)
+async def download_complete(
+    r: Request,
+    job_id: _ANN_JOB_ID
+):
+    logging.getLogger(__name__).info(f"Download reported as complete for job {job_id}")
+    # TODO NOW implement
+    raise NotImplementedError()
+
+
+@ROUTER_CALLBACKS.get(
+    f"/{get_job_complete_callback()}/{{job_id}}",
+    summary="Report job complete",
+    description="Report a remote job is complete. This method is not expected "
+        + "to be called by users."
+)
+async def job_complete(
+    r: Request,
+    job_id: _ANN_JOB_ID
+):
+    logging.getLogger(__name__).info(f"Remote job reported as complete for job {job_id}")
+    # TODO JOBS implement when job is complete
+    raise NotImplementedError()
+
+
 class ClientLifeTimeError(Exception):
     """ An error thrown when a client's lifetime is less than required. """