Add executed and errored state APIs with corresponding models and controllers (#107)

Author: NiveditJain

state-manager/app/controller/errored_state.py

@@ -0,0 +1,34 @@
+from app.models.errored_models import ErroredRequestModel, ErroredResponseModel
+from bson import ObjectId
+from fastapi import HTTPException, status
+
+from app.models.db.state import State
+from app.models.state_status_enum import StateStatusEnum
+from app.singletons.logs_manager import LogsManager
+
+logger = LogsManager().get_logger()
+
+async def errored_state(namespace_name: str, state_id: ObjectId, body: ErroredRequestModel, x_exosphere_request_id: str) -> ErroredResponseModel:
+
+    try:
+        logger.info(f"Errored state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+
+        state = await State.find_one(State.id == state_id)
+        if not state:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="State not found")
+        
+        if state.status != StateStatusEnum.QUEUED and state.status != StateStatusEnum.EXECUTED:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is not queued or executed")
+        
+        if state.status == StateStatusEnum.EXECUTED:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is already executed")
+        
+        await State.find_one(State.id == state_id).set(
+            {"status": StateStatusEnum.ERRORED, "error": body.error}
+        )
+
+        return ErroredResponseModel(status=StateStatusEnum.ERRORED)
+
+    except Exception as e:
+        logger.error(f"Error errored state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id, error=e)
+        raise e

state-manager/app/controller/executed_state.py

@@ -0,0 +1,31 @@
+from app.models.executed_models import ExecutedRequestModel, ExecutedResponseModel
+from bson import ObjectId
+from fastapi import HTTPException, status
+
+from app.models.db.state import State
+from app.models.state_status_enum import StateStatusEnum
+from app.singletons.logs_manager import LogsManager
+
+logger = LogsManager().get_logger()
+
+async def executed_state(namespace_name: str, state_id: ObjectId, body: ExecutedRequestModel, x_exosphere_request_id: str) -> ExecutedResponseModel:
+
+    try:
+        logger.info(f"Executed state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+
+        state = await State.find_one(State.id == state_id)
+        if not state:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="State not found")
+
+        if state.status != StateStatusEnum.QUEUED:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is not queued")
+        
+        await State.find_one(State.id == state_id).set(
+            {"status": StateStatusEnum.EXECUTED, "outputs": body.outputs}
+        )
+
+        return ExecutedResponseModel(status=StateStatusEnum.EXECUTED)
+
+    except Exception as e:
+        logger.error(f"Error executing state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id, error=e)
+        raise e

state-manager/app/models/db/namespace.py

@@ -5,4 +5,4 @@
 
 class Namespace(BaseDatabaseModel):
 
-    name: Indexed(str, unique=True) = Field(..., description="Name of the namespace")
+    name: Indexed(str, unique=True) = Field(..., description="Name of the namespace") # type: ignore

state-manager/app/models/db/state.py

@@ -1,7 +1,7 @@
 from .base import BaseDatabaseModel
 from ..state_status_enum import StateStatusEnum
 from pydantic import Field
-from typing import Any
+from typing import Any, Optional
 
 
 class State(BaseDatabaseModel):
@@ -10,4 +10,5 @@ class State(BaseDatabaseModel):
     namespace_name: str = Field(..., description="Name of the namespace of the state")
     status: StateStatusEnum = Field(..., description="Status of the state")
     inputs: dict[str, Any] = Field(..., description="Inputs of the state")
-    outputs: dict[str, Any] = Field(..., description="Outputs of the state")
+    outputs: dict[str, Any] = Field(..., description="Outputs of the state")
+    error: Optional[str] = Field(None, description="Error message")

state-manager/app/models/errored_models.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel, Field
+from .state_status_enum import StateStatusEnum
+
+
+class ErroredRequestModel(BaseModel):
+    error: str = Field(..., description="Error message")
+
+
+class ErroredResponseModel(BaseModel):
+    status: StateStatusEnum = Field(..., description="Status of the state")

state-manager/app/models/executed_models.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel, Field
+from typing import Any
+from .state_status_enum import StateStatusEnum
+
+class ExecutedRequestModel(BaseModel):
+    outputs: dict[str, Any] = Field(..., description="Outputs of the state")
+
+
+class ExecutedResponseModel(BaseModel):
+    status: StateStatusEnum = Field(..., description="Status of the state")

state-manager/app/routes.py

@@ -1,5 +1,6 @@
 from fastapi import APIRouter, status, Request, Depends, HTTPException
 from uuid import uuid4
+from bson import ObjectId
 
 from app.utils.check_secret import check_api_key
 from app.singletons.logs_manager import LogsManager
@@ -11,6 +12,13 @@
 from .models.create_models import CreateRequestModel, CreateResponseModel
 from .controller.create_states import create_states
 
+from .models.executed_models import ExecutedRequestModel, ExecutedResponseModel
+from .controller.executed_state import executed_state
+
+from .models.errored_models import ErroredRequestModel, ErroredResponseModel
+from .controller.errored_state import errored_state
+
+
 
 logger = LogsManager().get_logger()
 
@@ -52,4 +60,42 @@ async def create_state(namespace_name: str, body: CreateRequestModel, request: R
         logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
 
-    return await create_states(namespace_name, body, x_exosphere_request_id)
+    return await create_states(namespace_name, body, x_exosphere_request_id)
+
+
+@router.post(
+    "/{state_id}/executed",
+    response_model=ExecutedResponseModel,
+    status_code=status.HTTP_200_OK,
+    response_description="State executed successfully"
+)
+async def executed_state_route(namespace_name: str, state_id: str, body: ExecutedRequestModel, request: Request, api_key: str = Depends(check_api_key)):
+
+    x_exosphere_request_id = getattr(request.state, "x_exosphere_request_id", str(uuid4()))
+
+    if api_key:
+        logger.info(f"API key is valid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+    else:
+        logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
+
+    return await executed_state(namespace_name, ObjectId(state_id), body, x_exosphere_request_id)
+
+
+@router.post(
+    "/{state_id}/errored",
+    response_model=ErroredResponseModel,
+    status_code=status.HTTP_200_OK,
+    response_description="State errored successfully"
+)
+async def errored_state_route(namespace_name: str, state_id: str, body: ErroredRequestModel, request: Request, api_key: str = Depends(check_api_key)):
+
+    x_exosphere_request_id = getattr(request.state, "x_exosphere_request_id", str(uuid4()))
+
+    if api_key:
+        logger.info(f"API key is valid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+    else:
+        logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
+
+    return await errored_state(namespace_name, ObjectId(state_id), body, x_exosphere_request_id)