feat: Collect from Resoto directly (#47)

Author: aquamatthias

cloud2sql/collect_plugins.py

@@ -33,6 +33,7 @@
 from cloud2sql.arrow.config import ArrowOutputConfig
 from cloud2sql.show_progress import CollectInfo
 from cloud2sql.sql import SqlUpdater, sql_updater
+from cloud2sql.remote_graph import RemoteGraphCollector
 
 try:
     from cloud2sql.arrow.model import ArrowModel
@@ -65,10 +66,18 @@ def collectors(raw_config: Json, feedback: CoreFeedback) -> Dict[str, BaseCollec
             log.info(f"Found collector {plugin_class.cloud} ({plugin_class.__name__})")
             plugin_class.add_config(config)
             plugin = plugin_class()
-            if hasattr(plugin, "core_feedback"):
-                setattr(plugin, "core_feedback", feedback.with_context(plugin.cloud))
             result[plugin_class.cloud] = plugin
 
+    # lookup local plugins
+    if RemoteGraphCollector.cloud in raw_config:
+        log.info(f"Found collector {RemoteGraphCollector.cloud} ({RemoteGraphCollector.__name__})")
+        result[RemoteGraphCollector.cloud] = RemoteGraphCollector()
+        RemoteGraphCollector.add_config(config)
+
+    for plugin in result.values():
+        if hasattr(plugin, "core_feedback"):
+            setattr(plugin, "core_feedback", feedback.with_context(plugin.cloud))
+
     Config.init_default_config()
     Config.running_config.data = {**Config.running_config.data, **Config.read_config(raw_config)}
     return result

cloud2sql/remote_graph.py

@@ -0,0 +1,91 @@
+from typing import Dict, Iterator
+from typing import Optional, ClassVar
+
+from attr import define, field
+from resotoclient import ResotoClient, JsObject
+from resotolib.baseplugin import BaseCollectorPlugin
+from resotolib.baseresources import (
+    BaseResource,
+    Cloud,
+    EdgeType,
+    UnknownZone,
+    UnknownRegion,
+    UnknownAccount,
+)
+from resotolib.config import Config
+from resotolib.core.actions import CoreFeedback
+from resotolib.core.model_export import node_from_dict
+from resotolib.graph import Graph
+from resotolib.json import value_in_path
+from resotolib.logger import log
+from resotolib.types import Json
+
+
+@define
+class RemoteGraphConfig:
+    kind: ClassVar[str] = "remote_graph"
+    resoto_url: str = field(default="https://localhost:8900", metadata={"description": "URL of the resoto server"})
+    psk: Optional[str] = field(default=None, metadata={"description": "Pre-shared key for the resoto server"})
+    graph: str = field(default="resoto", metadata={"description": "Name of the graph to use"})
+    search: Optional[str] = field(
+        default=None, metadata={"description": "Search string to filter resources. None to get all resources."}
+    )
+
+
+carz = {"cloud": Cloud, "account": UnknownAccount, "region": UnknownRegion, "zone": UnknownZone}
+
+
+class RemoteGraphCollector(BaseCollectorPlugin):
+    cloud = "remote_graph"
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.core_feedback: Optional[CoreFeedback] = None
+
+    @staticmethod
+    def add_config(cfg: Config) -> None:
+        cfg.add_config(RemoteGraphConfig)
+
+    def collect(self) -> None:
+        try:
+            self.graph = self._collect_remote_graph()
+        except Exception as ex:
+            if self.core_feedback:
+                self.core_feedback.error(f"Unhandled exception in Remote Plugin: {ex}", log)
+            else:
+                log.error(f"No CoreFeedback available! Unhandled exception in RemoteGraph Plugin: {ex}")
+            raise
+
+    def _collect_remote_graph(self) -> Graph:
+        config: RemoteGraphConfig = Config.remote_graph
+        client = ResotoClient(config.resoto_url, psk=config.psk)
+        search = config.search or "is(graph_root) -[2:]->"
+        return self._collect_from_graph_iterator(client.search_graph(search, graph=config.graph))
+
+    def _collect_from_graph_iterator(self, graph_iterator: Iterator[JsObject]) -> Graph:
+        assert self.core_feedback, "No CoreFeedback available!"
+        graph = Graph()
+        lookup: Dict[str, BaseResource] = {}
+        self.core_feedback.progress_done("Remote Graph", 0, 1)
+
+        def set_carz(jsc: Json, rs: BaseResource) -> None:
+            for ancestor, clazz in carz.items():
+                path = ["ancestors", ancestor, "reported"]
+                if (cv := value_in_path(jsc, path)) and (cid := cv.get("id")) and (cname := cv.get("name")):
+                    resource = lookup.get(cid, clazz(id=cid, name=cname))  # type: ignore
+                    lookup[cid] = resource
+                    setattr(rs, f"_{ancestor}", resource)  # xxx is defined by _xxx property
+
+        for js in graph_iterator:
+            if js.get("type") == "node" and isinstance(js, dict):
+                node = node_from_dict(js)
+                set_carz(js, node)
+                lookup[js["id"]] = node
+                graph.add_node(node)
+            elif js.get("type") == "edge":
+                if (node_from := lookup.get(js["from"])) and (node_to := lookup.get(js["to"])):
+                    graph.add_edge(node_from, node_to, edge_type=EdgeType.default)
+            else:
+                raise ValueError(f"Unknown type: {js.get('type')}")
+        self.core_feedback.progress_done("Remote Graph", 1, 1)
+        return graph

tests/conftest.py

@@ -1,10 +1,13 @@
 from queue import Queue
 from typing import List, Iterator
 
+from resoto_plugin_example_collector import ExampleAccount, ExampleRegion, ExampleInstance, ExampleVolume
 from resotoclient.models import Model, Kind, Property
 from pytest import fixture
 from resotolib.args import Namespace
+from resotolib.baseresources import GraphRoot, Cloud
 from resotolib.core.actions import CoreFeedback
+from resotolib.graph import Graph
 from resotolib.types import Json
 from sqlalchemy.engine import create_engine, Engine
 
@@ -60,6 +63,27 @@ def model() -> Model:
     return Model({k.fqn: k for k in kinds})
 
 
+@fixture
+def example_collector_graph() -> Graph:
+    root = GraphRoot(id="root")
+    graph = Graph(root=root)
+    cloud = Cloud(id="example")
+    account = ExampleAccount(id="example-account")
+    region = ExampleRegion(id="example-region")
+    i1 = ExampleInstance(id="example-instance-1")
+    iv1 = ExampleVolume(id="example-volume-1")
+    i2 = ExampleInstance(id="example-instance-2")
+    iv2 = ExampleVolume(id="example-volume-1")
+    graph.add_resource(root, cloud)
+    graph.add_resource(cloud, account)
+    graph.add_resource(account, region)
+    graph.add_resource(region, i1)
+    graph.add_resource(i1, iv1)
+    graph.add_resource(region, i2)
+    graph.add_resource(i2, iv2)
+    return graph
+
+
 @fixture()
 def args() -> Namespace:
     return Namespace()

tests/remote_graph_test.py

@@ -0,0 +1,25 @@
+from typing import Iterator
+
+from resotoclient import JsObject
+from resotolib.core.actions import CoreFeedback
+from resotolib.core.model_export import node_to_dict
+from resotolib.graph import Graph
+
+from cloud2sql.remote_graph import RemoteGraphCollector
+
+
+def test_remote_graph_collector(example_collector_graph: Graph, core_feedback: CoreFeedback) -> None:
+    def graph_iterator() -> Iterator[JsObject]:
+        for node in example_collector_graph.nodes:
+            fn = node_to_dict(node)
+            fn["type"] = "node"
+            yield fn
+
+        for from_node, to_node in example_collector_graph.edges():
+            yield {"type": "edge", "from": from_node.chksum, "to": to_node.chksum, "edge_type": "default"}
+
+    collector = RemoteGraphCollector()
+    collector.core_feedback = core_feedback
+    graph_again = collector._collect_from_graph_iterator(graph_iterator())
+    assert len(graph_again.nodes) == len(example_collector_graph.nodes)
+    assert len(graph_again.edges) == len(example_collector_graph.edges)