From 269b979abd7687d2c118455968414f45b092b344 Mon Sep 17 00:00:00 2001
From: Maxim V4S <maxim.d@v4scale.com>
Date: Mon, 5 Aug 2024 16:17:26 +0300
Subject: [PATCH] feat: qualibration graph representation

---
 qualibrate_runner/api/dependencies.py | 28 +++++++++++++++++
 qualibrate_runner/api/routes.py       | 44 ++++++++++++++++++++++++---
 qualibrate_runner/core/run_job.py     |  8 +++--
 3 files changed, 74 insertions(+), 6 deletions(-)

diff --git a/qualibrate_runner/api/dependencies.py b/qualibrate_runner/api/dependencies.py
index d489da7..f0624d5 100644
--- a/qualibrate_runner/api/dependencies.py
+++ b/qualibrate_runner/api/dependencies.py
@@ -2,6 +2,7 @@
 from typing import Annotated, Mapping, cast
 
 from fastapi import Depends, HTTPException
+from qualibrate.qualibration_graph import QualibrationGraph
 from qualibrate.qualibration_library import QualibrationLibrary
 from qualibrate.qualibration_node import QualibrationNode
 
@@ -33,6 +34,13 @@ def get_nodes(
     return cast(Mapping[str, QualibrationNode], library.get_nodes())
 
 
+@cache
+def get_graphs(
+    library: Annotated[QualibrationLibrary, Depends(get_library)],
+) -> Mapping[str, QualibrationGraph]:
+    return cast(Mapping[str, QualibrationGraph], library.get_graphs())
+
+
 def get_node(
     name: str,
     nodes: Annotated[Mapping[str, QualibrationNode], Depends(get_nodes)],
@@ -41,3 +49,23 @@ def get_node(
     if node is None:
         raise HTTPException(status_code=422, detail=f"Unknown node name {name}")
     return node
+
+
+def get_graph(
+    name: str,
+    graphs: Annotated[Mapping[str, QualibrationGraph], Depends(get_graphs)],
+) -> QualibrationGraph:
+    graph = graphs.get(name)
+    if graph is None:
+        raise HTTPException(
+            status_code=422, detail=f"Unknown graph name {name}"
+        )
+    return graph
+
+
+CACHED_DEPENDENCIES = (get_library, get_nodes, get_graphs)
+
+
+def cache_clear() -> None:
+    for func in CACHED_DEPENDENCIES:
+        func.cache_clear()
diff --git a/qualibrate_runner/api/routes.py b/qualibrate_runner/api/routes.py
index 85d887f..898edd7 100644
--- a/qualibrate_runner/api/routes.py
+++ b/qualibrate_runner/api/routes.py
@@ -1,13 +1,21 @@
-from typing import Annotated, Any, Mapping, Optional, Type, cast
+from typing import Annotated, Any, Mapping, Optional, Sequence, Type, cast
 
 from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, status
 from pydantic import BaseModel
+from qualibrate.qualibration_graph import QualibrationGraph
 from qualibrate.qualibration_node import QualibrationNode
 
 from qualibrate_runner.api.dependencies import (
+    cache_clear,
     get_library,
     get_state,
 )
+from qualibrate_runner.api.dependencies import (
+    get_graph as get_qgraph,
+)
+from qualibrate_runner.api.dependencies import (
+    get_graphs as get_qgraphs,
+)
 from qualibrate_runner.api.dependencies import (
     get_node as get_qnode,
 )
@@ -54,18 +62,32 @@ def submit_run(
 
 @base_router.get("/get_nodes")
 def get_nodes(
-    nodes: Annotated[Mapping[str, Any], Depends(get_qnodes)],
+    nodes: Annotated[Mapping[str, QualibrationNode], Depends(get_qnodes)],
     settings: Annotated[QualibrateRunnerSettings, Depends(get_settings)],
     rescan: bool = False,
 ) -> Mapping[str, Any]:
     if rescan:
-        get_library.cache_clear()
-        get_qnodes.cache_clear()
+        cache_clear()
         library = get_library(settings)
         nodes = get_qnodes(library)
     return {node_name: node.serialize() for node_name, node in nodes.items()}
 
 
+@base_router.get("/get_graphs")
+def get_graphs(
+    graphs: Annotated[Mapping[str, QualibrationNode], Depends(get_qgraphs)],
+    settings: Annotated[QualibrateRunnerSettings, Depends(get_settings)],
+    rescan: bool = False,
+) -> Mapping[str, Any]:
+    if rescan:
+        cache_clear()
+        library = get_library(settings)
+        graphs = get_qnodes(library)
+    return {
+        graph_name: graph.serialize() for graph_name, graph in graphs.items()
+    }
+
+
 @base_router.get("/get_node")
 def get_node(
     node: Annotated[QualibrationNode, Depends(get_qnode)],
@@ -73,6 +95,20 @@ def get_node(
     return cast(Mapping[str, Any], node.serialize())
 
 
+@base_router.get("/get_graph")
+def get_graph(
+    graph: Annotated[QualibrationGraph, Depends(get_qgraph)],
+) -> Mapping[str, Any]:
+    return cast(Mapping[str, Any], graph.serialize())
+
+
+@base_router.get("/get_graph/cytoscape")
+def get_graph_cytoscape(
+    graph: Annotated[QualibrationGraph, Depends(get_qgraph)],
+) -> Sequence[Mapping[str, Any]]:
+    return cast(Sequence[Mapping[str, Any]], graph.cytoscape_representation())
+
+
 @base_router.get("/last_run")
 def get_last_run(
     state: Annotated[State, Depends(get_state)],
diff --git a/qualibrate_runner/core/run_job.py b/qualibrate_runner/core/run_job.py
index c0a3541..1013a6e 100644
--- a/qualibrate_runner/core/run_job.py
+++ b/qualibrate_runner/core/run_job.py
@@ -34,8 +34,12 @@ def run_job(
         idx=-1,
     )
     try:
-        node = QualibrationLibrary.active_library.nodes[node.name]
-        node.run_node(node.parameters_class(**passed_input_parameters))
+        library = QualibrationLibrary.active_library
+
+        node = library.nodes[node.name]
+        library.run_node(
+            node.name, node.parameters_class(**passed_input_parameters)
+        )
     except Exception as ex:
         state.last_run = LastRun(
             name=state.last_run.name,