Adding APIs to Create States and Enqueue States

api-server/app/main.py

@@ -21,13 +21,11 @@
 # injecting databases
 from .user.models.user_database_model import User
 from .project.models.project_database_model import Project
-from .satellite.models.satellite_database_model import Satellite
 
 # injecting routers
 from .user.routes import router as user_router
 from .auth.routes import router as auth_router
 from .project.routes import router as project_router
-from .satellite.routes import router as satellite_router
 
 load_dotenv()
 
@@ -46,7 +44,7 @@ async def lifespan(app: FastAPI):
     # os.getenv can return None, which is not valid for dictionary-style access
     db_name = os.getenv("MONGO_DATABASE_NAME", "exosphere-api-server")
     db = client[db_name]
-    await init_beanie(db, document_models=[User, Project, Satellite])
+    await init_beanie(db, document_models=[User, Project])
     logger.info("beanie dbs initialized")
 
     # main logic of the server
@@ -87,4 +85,3 @@ def health() -> dict:
 app.include_router(user_router)
 app.include_router(auth_router)
 app.include_router(project_router)
-app.include_router(satellite_router)

state-manager/app/controller/create_states.py

@@ -0,0 +1,42 @@
+from app.singletons.logs_manager import LogsManager
+from app.models.create_models import CreateRequestModel, CreateResponseModel, ResponseStateModel
+from app.models.state_status_enum import StateStatusEnum
+from app.models.db.state import State
+
+from beanie.operators import In
+from bson import ObjectId
+
+logger = LogsManager().get_logger()
+
+async def create_states(namespace_name: str, body: CreateRequestModel, x_exosphere_request_id: str) -> CreateResponseModel:
+    try:
+        states = []
+        logger.info(f"Creating states for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+
+        for state in body.states:
+            states.append(
+                State(
+                    node_name=state.node_name,
+                    namespace_name=namespace_name,
+                    status=StateStatusEnum.CREATED,
+                    inputs=state.inputs,
+                    outputs={}
+                )
+            )
+
+        inserted_states = await State.insert_many(states)
+
+        logger.info(f"Created states: {inserted_states.inserted_ids}", x_exosphere_request_id=x_exosphere_request_id)
+
+        newStates = await State.find(
+            In(State.id, [ObjectId(id) for id in inserted_states.inserted_ids])
+        ).to_list()
+
+        return CreateResponseModel(
+            status=StateStatusEnum.CREATED,
+            states=[ResponseStateModel(state_id=str(state.id), node_name=state.node_name, inputs=state.inputs, created_at=state.created_at) for state in newStates]
+        )
+
+    except Exception as e:
+        logger.error(f"Error creating states for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+        raise e

state-manager/app/controller/enqueue_states.py

@@ -0,0 +1,53 @@
+from beanie.operators import In
+
+from ..models.enqueue_request import EnqueueRequestModel
+from ..models.enqueue_response import EnqueueResponseModel, StateModel
+from ..models.db.state import State
+from ..models.state_status_enum import StateStatusEnum
+
+from app.singletons.logs_manager import LogsManager
+
+logger = LogsManager().get_logger()
+
+
+async def enqueue_states(namespace_name: str, body: EnqueueRequestModel, x_exosphere_request_id: str) -> EnqueueResponseModel:
+
+    try:
+        logger.info(f"Enqueuing states for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
+
+        states = await State.find(
+            State.namespace_name == namespace_name,
+            State.status == StateStatusEnum.CREATED,
+            In(State.node_name, body.nodes)
+        ).limit(
+            body.batch_size
+        ).to_list()
+
+        if states:
+            await State.find(
+                In(State.id, [state.id for state in states])
+            ).set(
+                {
+                    "status": StateStatusEnum.QUEUED,
+                }
+            ) # type: ignore
+
+        response = EnqueueResponseModel(
+            count=len(states),
+            namespace=namespace_name,
+            status=StateStatusEnum.QUEUED,
+            states=[
+                StateModel(
+                    state_id=str(state.id),
+                    node_name=state.node_name,
+                    inputs=state.inputs,
+                    created_at=state.created_at
+                )
+                for state in states
+            ]
+        )
+        return response
+
+    except Exception as e:
+        logger.error(f"Error enqueuing states for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id, error=e)
+        raise e

state-manager/app/main.py

@@ -16,6 +16,13 @@
     UnhandledExceptionsMiddleware,
 )
 from .middlewares.request_id_middleware import RequestIdMiddleware
+
+# injecting models
+from .models.db.state import State
+from .models.db.namespace import Namespace
+
+# injecting routes
+from .routes import router
 
 load_dotenv()
 
@@ -27,8 +34,8 @@ async def lifespan(app: FastAPI):
 
     # initializing beanie
     client = AsyncMongoClient(os.getenv("MONGO_URI"))
-    db = client[os.getenv("MONGO_DATABASE_NAME")]
-    await init_beanie(db, document_models=[])
+    db = client[os.getenv("MONGO_DATABASE_NAME", "exosphere-state-manager")]
+    await init_beanie(db, document_models=[State, Namespace])
     logger.info("beanie dbs initialized")
 
     # initialize secret
@@ -66,4 +73,6 @@ async def lifespan(app: FastAPI):
 
 @app.get("/health")
 def health() -> dict:
-    return {"message": "OK"}
+    return {"message": "OK"}
+
+app.include_router(router)