From 81576fa0f5919f63f75af1dd7ae3f0b36df5f7e3 Mon Sep 17 00:00:00 2001
From: Niklas <niklas@scaleoutsystems.com>
Date: Mon, 24 Jun 2024 12:36:04 +0200
Subject: [PATCH] Feature/SK-916 | Get active model id (#645)

---
 fedn/network/api/v1/model_routes.py           | 35 +++++++++++++++++++
 .../storage/statestore/stores/model_store.py  | 11 ++++++
 2 files changed, 46 insertions(+)

diff --git a/fedn/network/api/v1/model_routes.py b/fedn/network/api/v1/model_routes.py
index aaea9d733..f227443e0 100644
--- a/fedn/network/api/v1/model_routes.py
+++ b/fedn/network/api/v1/model_routes.py
@@ -598,6 +598,7 @@ def get_ancestors(id: str):
 
 
 @bp.route("/<string:id>/download", methods=["GET"])
+@jwt_auth_required(role="admin")
 def download(id: str):
     """Download
     Downloads the model file of the provided id.
@@ -646,6 +647,7 @@ def download(id: str):
 
 
 @bp.route("/<string:id>/parameters", methods=["GET"])
+@jwt_auth_required(role="admin")
 def get_parameters(id: str):
     """Download
     Downloads parameters of the model of the provided id.
@@ -707,3 +709,36 @@ def get_parameters(id: str):
         return jsonify({"message": f"Entity with id: {id} not found"}), 404
     except Exception:
         return jsonify({"message": "An unexpected error occurred"}), 500
+
+
+@bp.route("/active", methods=["GET"])
+@jwt_auth_required(role="admin")
+def get_active_model():
+    """Get active model
+    Retrieves the active model (id).
+    ---
+    tags:
+        - Models
+    responses:
+        200:
+            description: The active model id
+            schema:
+                type: string
+        500:
+            description: An error occurred
+            schema:
+                type: object
+                properties:
+                    message:
+                        type: string
+    """
+    try:
+        active_model = model_store.get_active()
+
+        response = active_model
+
+        return jsonify(response), 200
+    except EntityNotFound:
+        return jsonify({"message": "No active model found"}), 404
+    except Exception:
+        return jsonify({"message": "An unexpected error occurred"}), 500
diff --git a/fedn/network/storage/statestore/stores/model_store.py b/fedn/network/storage/statestore/stores/model_store.py
index 775eccd83..3048f2a26 100644
--- a/fedn/network/storage/statestore/stores/model_store.py
+++ b/fedn/network/storage/statestore/stores/model_store.py
@@ -209,3 +209,14 @@ def count(self, **kwargs) -> int:
         """
         kwargs["key"] = "models"
         return super().count(**kwargs)
+
+    def get_active(self) -> str:
+        """Get the active model
+        return: The active model id (str)
+        """
+        active_model = self.database[self.collection].find_one({"key": "current_model"})
+
+        if active_model is None:
+            raise EntityNotFound("Active model not found")
+
+        return active_model["model"]