fix: Export all visible default edges (#12)

Author: aquamatthias

cloud2sql/collect_plugins.py

@@ -7,14 +7,14 @@
 from queue import Queue
 from threading import Event
 from time import sleep
-from typing import Dict, Optional, List, Any
+from typing import Dict, Optional, List, Any, Tuple
 
 import pkg_resources
 import yaml
 from resotoclient import Kind, Model
 from resotolib.args import Namespace
 from resotolib.baseplugin import BaseCollectorPlugin
-from resotolib.baseresources import BaseResource
+from resotolib.baseresources import BaseResource, EdgeType
 from resotolib.config import Config
 from resotolib.core.actions import CoreFeedback
 from resotolib.core.model_export import node_to_dict
@@ -50,7 +50,6 @@ def collectors(raw_config: Json, feedback: CoreFeedback) -> Dict[str, BaseCollec
 
 
 def configure(path_to_config: Optional[str]) -> Json:
-    # Config.init_default_config()
     if path_to_config:
         with open(path_to_config) as f:
             return yaml.safe_load(f)  # type: ignore
@@ -61,6 +60,30 @@ def collect(collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedba
     # collect cloud data
     feedback.progress_done(collector.cloud, 0, 1)
     collector.collect()
+
+    # group all nodes by kind
+    nodes_by_kind: Dict[str, List[Json]] = defaultdict(list)
+    node: BaseResource
+    for node in collector.graph.nodes:
+        node._graph = collector.graph
+        # create an exported node with the same scheme as resotocore
+        exported = node_to_dict(node)
+        exported["type"] = "node"
+        exported["ancestors"] = {
+            "cloud": {"reported": {"id": node.cloud().name}},
+            "account": {"reported": {"id": node.account().name}},
+            "region": {"reported": {"id": node.region().name}},
+            "zone": {"reported": {"id": node.zone().name}},
+        }
+        nodes_by_kind[node.kind].append(exported)
+
+    # group all edges by kind of from/to
+    edges_by_kind: Dict[Tuple[str, str], List[Json]] = defaultdict(list)
+    for from_node, to_node, key in collector.graph.edges:
+        if key.edge_type == EdgeType.default:
+            edge_node = {"from": from_node.chksum, "to": to_node.chksum, "type": "edge"}
+            edges_by_kind[(from_node.kind, to_node.kind)].append(edge_node)
+
     # read the kinds created from this collector
     kinds = [from_json(m, Kind) for m in collector.graph.export_model(walk_subclasses=False)]
     updater = sql_updater(Model({k.fqn: k for k in kinds}), engine)
@@ -73,25 +96,9 @@ def collect(collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedba
     with engine.connect() as conn:
         with conn.begin():
             # create the ddl metadata from the kinds
-            updater.create_schema(conn, args)
+            updater.create_schema(conn, args, list(edges_by_kind.keys()))
             feedback.progress_done(schema, 1, 1, context=[collector.cloud])
 
-            # group all nodes by kind
-            nodes_by_kind = defaultdict(list)
-            node: BaseResource
-            for node in collector.graph.nodes:
-                node._graph = collector.graph
-                # create an exported node with the same scheme as resotocore
-                exported = node_to_dict(node)
-                exported["type"] = "node"
-                exported["ancestors"] = {
-                    "cloud": {"reported": {"id": node.cloud().name}},
-                    "account": {"reported": {"id": node.account().name}},
-                    "region": {"reported": {"id": node.region().name}},
-                    "zone": {"reported": {"id": node.zone().name}},
-                }
-                nodes_by_kind[node.kind].append(exported)
-
             # insert batches of nodes by kind
             for kind, nodes in nodes_by_kind.items():
                 log.info(f"Inserting {len(nodes)} nodes of kind {kind}")
@@ -100,12 +107,6 @@ def collect(collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedba
                 ne_count += len(nodes)
                 feedback.progress_done(syncdb, ne_count, node_edge_count, context=[collector.cloud])
 
-            # group all nodes by kind of from/to
-            edges_by_kind = defaultdict(list)
-            for from_node, to_node, _ in collector.graph.edges:
-                edge_node = {"from": from_node.chksum, "to": to_node.chksum, "type": "edge"}
-                edges_by_kind[(from_node.kind, to_node.kind)].append(edge_node)
-
             # insert batches of edges by from/to kind
             for from_to, nodes in edges_by_kind.items():
                 log.info(f"Inserting {len(nodes)} edges from {from_to[0]} to {from_to[1]}")

cloud2sql/sql.py

@@ -70,7 +70,7 @@ def insert_edges(self, from_to: Tuple[str, str], nodes: List[Json]) -> Iterator[
         pass
 
     @abstractmethod
-    def create_schema(self, connection: Connection, args: Namespace) -> MetaData:
+    def create_schema(self, connection: Connection, args: Namespace, edges: List[Tuple[str, str]]) -> MetaData:
         pass
 
     @staticmethod
@@ -141,7 +141,7 @@ def base_props_not_visited(kd: Kind) -> Tuple[List[Property], List[str]]:
         prs, scs = base_props_not_visited(kind)
         return prs + carz, scs
 
-    def create_schema(self, connection: Connection, args: Namespace) -> MetaData:
+    def create_schema(self, connection: Connection, args: Namespace, edges: List[Tuple[str, str]]) -> MetaData:
         log.info(f"Create schema for {len(self.table_kinds)} kinds and their relationships")
 
         def table_schema(kind: Kind) -> None:
@@ -182,6 +182,9 @@ def link_table_schema_from_successors(kind: Kind) -> None:
         # step 2: create link tables for all kinds
         for kind in self.table_kinds:
             link_table_schema_from_successors(kind)
+        # step 3: create link tables for all seen edges
+        for from_kind, to_kind in edges:
+            link_table_schema(from_kind, to_kind)
 
         # drop tables if requested
         self.metadata.drop_all(connection)

tests/collect_plugins_test.py

@@ -17,15 +17,23 @@ def test_collect() -> None:
         # get all tables
         metadata = MetaData()
         metadata.reflect(bind=engine)
-        assert set(metadata.tables.keys()) == {
-            "example_account",
-            "example_custom_resource",
-            "example_instance",
-            "example_network",
-            "example_region",
-            "example_volume",
+        expected_counts = {
+            "example_account": 1,
+            "example_custom_resource": 1,
+            "example_instance": 2,
+            "example_network": 2,
+            "example_region": 2,
+            "example_volume": 2,
+            "link_example_account_example_region": 2,
+            "link_example_instance_example_volume": 2,
+            "link_example_network_example_instance": 2,
+            "link_example_region_example_custom_resource": 1,
+            "link_example_region_example_instance": 2,
+            "link_example_region_example_network": 2,
+            "link_example_region_example_volume": 2,
         }
-        # check that there are entries in the tables
+
+        assert set(metadata.tables.keys()) == expected_counts.keys()  # check that there are entries in the tables
         with Session(engine) as session:
             for table in metadata.tables.values():
-                assert session.query(table).count() > 0
+                assert session.query(table).count() == expected_counts[table.name]

tests/conftest.py

@@ -73,7 +73,7 @@ def engine() -> Engine:
 def engine_with_schema(updater: SqlDefaultUpdater, args: Namespace) -> Engine:
     engine = create_engine("sqlite:///:memory:")
     with engine.connect() as connection:
-        updater.create_schema(connection, args)
+        updater.create_schema(connection, args, [])
     return engine
 
 

tests/sql_test.py

@@ -12,7 +12,7 @@
 def test_create_schema(model: Model, engine: Engine) -> None:
     updater: SqlDefaultUpdater = SqlDefaultUpdater(model)
     with engine.connect() as connection:
-        updater.create_schema(connection, Namespace())
+        updater.create_schema(connection, Namespace(), [])
 
     info = MetaData()
     info.reflect(bind=engine)