Add parquet support (#10)

Author: meln1k

README.md

@@ -76,6 +76,12 @@ destinations:
             key: value
 ```
 
+#### Apache Parquet
+
+```
+parquet:///path/to/parquet/directory
+```
+
 #### My database is not listed here
 
 cloud2sql uses SQLAlchemy to connect to the database. If your database is not listed here, you can check if it is supported in [SQLAlchemy Dialects](https://docs.sqlalchemy.org/en/20/dialects/index.html).

cloud2sql/__main__.py

@@ -1,6 +1,5 @@
 from logging import getLogger
 from typing import Optional
-
 from resotolib.args import Namespace, ArgumentParser
 from resotolib.logger import setup_logger
 from sqlalchemy import create_engine
@@ -44,7 +43,7 @@ def parse_args() -> Namespace:
     return args  # type: ignore
 
 
-def collect(engine: Engine, args: Namespace, sender: AnalyticsEventSender) -> None:
+def collect(engine: Optional[Engine], args: Namespace, sender: AnalyticsEventSender) -> None:
     try:
         collect_from_plugins(engine, args, sender)
     except Exception as e:
@@ -59,7 +58,8 @@ def main() -> None:
         setup_logger("resoto.cloud2sql", level=args.log_level, force=True)
         sender = NoEventSender() if args.analytics_opt_out else PosthogEventSender()
         config = configure(args.config)
-        engine = create_engine(db_string_from_config(config))
+        is_parquet = next(iter(config["destinations"].keys()), None) == "parquet"
+        engine = None if is_parquet else create_engine(db_string_from_config(config))
         collect(engine, args, sender)
     except Exception as e:
         if args.debug:  # raise exception and show complete tracelog

cloud2sql/collect_plugins.py

@@ -2,12 +2,14 @@
 import multiprocessing
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, Future
+import concurrent.futures
 from contextlib import suppress
 from logging import getLogger
 from queue import Queue
 from threading import Event
 from time import sleep
-from typing import Dict, Optional, List, Any, Tuple
+from typing import Dict, Optional, List, Any, Tuple, Set
+from pathlib import Path
 
 import pkg_resources
 import yaml
@@ -25,9 +27,12 @@
 from rich.live import Live
 from sqlalchemy.engine import Engine
 
+
 from cloud2sql.analytics import AnalyticsEventSender
 from cloud2sql.show_progress import CollectInfo
 from cloud2sql.sql import SqlUpdater, sql_updater
+from cloud2sql.parquet import ParquetModel, ParquetWriter
+
 
 log = getLogger("resoto.cloud2sql")
 
@@ -65,6 +70,73 @@ def configure(path_to_config: Optional[str]) -> Json:
 
 
 def collect(
+    collector: BaseCollectorPlugin, engine: Optional[Engine], feedback: CoreFeedback, args: Namespace, config: Json
+) -> Tuple[str, int, int]:
+    if engine:
+        return collect_sql(collector, engine, feedback, args)
+    else:
+        return collect_parquet(collector, feedback, config)
+
+
+def prepare_node(node: BaseResource, collector: BaseCollectorPlugin) -> Json:
+    node._graph = collector.graph
+    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}},
+    }
+    return exported
+
+
+def collect_parquet(collector: BaseCollectorPlugin, feedback: CoreFeedback, config: Json) -> Tuple[str, int, int]:
+    # collect cloud data
+    feedback.progress_done(collector.cloud, 0, 1)
+    collector.collect()
+    # read the kinds created from this collector
+    kinds = [from_json(m, Kind) for m in collector.graph.export_model(walk_subclasses=False)]
+    model = ParquetModel(Model({k.fqn: k for k in kinds}))
+    node_edge_count = len(collector.graph.nodes) + len(collector.graph.edges)
+    ne_current = 0
+    progress_update = node_edge_count // 100
+    feedback.progress_done("sync_db", 0, node_edge_count, context=[collector.cloud])
+
+    # group all edges by kind of from/to
+    edges_by_kind: Set[Tuple[str, str]] = set()
+    for from_node, to_node, key in collector.graph.edges:
+        if key.edge_type == EdgeType.default:
+            edges_by_kind.add((from_node.kind, to_node.kind))
+    # create the ddl metadata from the kinds
+    model.create_schema(list(edges_by_kind))
+    # ingest the data
+    parquet_conf = config.get("destinations", {}).get("parquet")
+    assert parquet_conf
+    parquet_path = Path(parquet_conf["path"])
+    parquet_batch_size = int(parquet_conf["batch_size"])
+    writer = ParquetWriter(model, parquet_path, parquet_batch_size)
+    node: BaseResource
+    for node in sorted(collector.graph.nodes, key=lambda n: n.kind):
+        exported = prepare_node(node, collector)
+        writer.insert_node(exported)
+        ne_current += 1
+        if ne_current % progress_update == 0:
+            feedback.progress_done("sync_db", ne_current, node_edge_count, context=[collector.cloud])
+    for from_node, to_node, key in collector.graph.edges:
+        if key.edge_type == EdgeType.default:
+            writer.insert_node({"from": from_node.chksum, "to": to_node.chksum, "type": "edge"})
+            ne_current += 1
+            if ne_current % progress_update == 0:
+                feedback.progress_done("sync_db", ne_current, node_edge_count, context=[collector.cloud])
+
+    writer.close()
+
+    feedback.progress_done(collector.cloud, 1, 1)
+    return collector.cloud, len(collector.graph.nodes), len(collector.graph.edges)
+
+
+def collect_sql(
     collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedback, args: Namespace
 ) -> Tuple[str, int, int]:
     # collect cloud data
@@ -75,16 +147,8 @@ def collect(
     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}},
-        }
+        exported = prepare_node(node, collector)
         nodes_by_kind[node.kind].append(exported)
 
     # group all edges by kind of from/to
@@ -138,33 +202,35 @@ def show_messages(core_messages: Queue[Json], end: Event) -> None:
         rich_print(message)
 
 
-def collect_from_plugins(engine: Engine, args: Namespace, sender: AnalyticsEventSender) -> None:
+def collect_from_plugins(engine: Optional[Engine], args: Namespace, sender: AnalyticsEventSender) -> None:
     # the multiprocessing manager is used to share data between processes
     mp_manager = multiprocessing.Manager()
     core_messages: Queue[Json] = mp_manager.Queue()
     feedback = CoreFeedback("cloud2sql", "collect", "collect", core_messages)
     raw_config = configure(args.config)
     sources = raw_config["sources"]
     all_collectors = collectors(sources, feedback)
-    analytics = {"total": len(all_collectors), "engine": engine.dialect.name} | {name: 1 for name in all_collectors}
+    engine_name = engine.dialect.name if engine else "parquet"
+    analytics = {"total": len(all_collectors), "engine": engine_name} | {name: 1 for name in all_collectors}
     end = Event()
     with ThreadPoolExecutor(max_workers=4) as executor:
         try:
             if args.show == "progress":
                 executor.submit(show_messages, core_messages, end)
             futures: List[Future[Any]] = []
             for collector in all_collectors.values():
-                futures.append(executor.submit(collect, collector, engine, feedback, args))
+                futures.append(executor.submit(collect, collector, engine, feedback, args, raw_config))
             for future in concurrent.futures.as_completed(futures):
                 name, nodes, edges = future.result()
                 analytics[f"{name}_nodes"] = nodes
                 analytics[f"{name}_edges"] = edges
             sender.capture("collect", **analytics)
             # when all collectors are done, we can swap all temp tables
-            swap_tables = "Make latest snapshot available"
-            feedback.progress_done(swap_tables, 0, 1)
-            SqlUpdater.swap_temp_tables(engine)
-            feedback.progress_done(swap_tables, 1, 1)
+            if engine:
+                swap_tables = "Make latest snapshot available"
+                feedback.progress_done(swap_tables, 0, 1)
+                SqlUpdater.swap_temp_tables(engine)
+                feedback.progress_done(swap_tables, 1, 1)
         except Exception as e:
             # set end and wait for live to finish, otherwise the cursor is not reset
             end.set()

cloud2sql/parquet.py

@@ -0,0 +1,159 @@
+from resotoclient.models import Kind, Model, JsObject
+from typing import Dict, List, Any, NamedTuple, Optional, Tuple
+import pyarrow as pa
+from cloud2sql.schema_utils import (
+    base_kinds,
+    get_table_name,
+    get_link_table_name,
+    kind_properties,
+    insert_node,
+)
+import pyarrow.parquet as pq
+from pathlib import Path
+from dataclasses import dataclass
+
+
+class ParquetModel:
+    def __init__(self, model: Model):
+        self.model = model
+        self.table_kinds = [
+            kind
+            for kind in model.kinds.values()
+            if kind.aggregate_root and kind.runtime_kind is None and kind.fqn not in base_kinds
+        ]
+        self.schemas: Dict[str, pa.Schema] = {}
+
+    def _parquet_type(self, kind: str) -> pa.lib.DataType:
+        if kind.startswith("dict") or "[]" in kind:
+            return pa.string()  # dicts and lists are converted to json strings
+        elif kind == "int32":
+            return pa.int32()
+        elif kind == "int64":
+            return pa.int64()
+        elif kind == "float":
+            pa.float32()
+        elif kind == "double":
+            return pa.float64()
+        elif kind in {"string", "datetime", "date", "duration"}:
+            return pa.string()
+        elif kind == "boolean":
+            return pa.bool_()
+        else:
+            return pa.string()
+
+    def create_schema(self, edges: List[Tuple[str, str]]) -> None:
+        def table_schema(kind: Kind) -> None:
+            table_name = get_table_name(kind.fqn, with_tmp_prefix=False)
+            if table_name not in self.schemas:
+                properties, _ = kind_properties(kind, self.model)
+                schema = pa.schema(
+                    [
+                        pa.field("_id", pa.string()),
+                        *[pa.field(p.name, self._parquet_type(p.kind)) for p in properties],
+                    ]
+                )
+                self.schemas[table_name] = schema
+
+        def link_table_schema(from_kind: str, to_kind: str) -> None:
+            from_table = get_table_name(from_kind, with_tmp_prefix=False)
+            to_table = get_table_name(to_kind, with_tmp_prefix=False)
+            link_table = get_link_table_name(from_kind, to_kind, with_tmp_prefix=False)
+            if link_table not in self.schemas and from_table in self.schemas and to_table in self.schemas:
+                schema = pa.schema(
+                    [
+                        pa.field("from_id", pa.string()),
+                        pa.field("to_id", pa.string()),
+                    ]
+                )
+                self.schemas[link_table] = schema
+
+        def link_table_schema_from_successors(kind: Kind) -> None:
+            _, successors = kind_properties(kind, self.model)
+            # create link table for all linked entities
+            for successor in successors:
+                link_table_schema(kind.fqn, successor)
+
+        # step 1: create tables for all kinds
+        for kind in self.table_kinds:
+            table_schema(kind)
+        # 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)
+
+        return None
+
+
+class WriteResult(NamedTuple):
+    table_name: str
+
+
+@dataclass
+class ParquetBatch:
+    rows: List[Dict[str, Any]]
+    schema: pa.Schema
+    writer: pq.ParquetWriter
+
+
+class ParquetWriter:
+    def __init__(
+        self,
+        model: ParquetModel,
+        result_directory: Path,
+        rows_per_batch: int,
+    ):
+        self.model = model
+        self.kind_by_id: Dict[str, str] = {}
+        self.batches: Dict[str, ParquetBatch] = {}
+        self.rows_per_batch = rows_per_batch
+        self.result_directory = result_directory
+
+    def insert_value(self, table_name: str, values: Any) -> Optional[WriteResult]:
+        if self.model.schemas.get(table_name):
+
+            def ensure_path(path: Path) -> Path:
+                path.mkdir(parents=True, exist_ok=True)
+                return path
+
+            batch = self.batches.get(
+                table_name,
+                ParquetBatch(
+                    [],
+                    self.model.schemas[table_name],
+                    pq.ParquetWriter(
+                        Path(ensure_path(self.result_directory), f"{table_name}.parquet"),
+                        self.model.schemas[table_name],
+                    ),
+                ),
+            )
+
+            batch.rows.append(values)
+            self.batches[table_name] = batch
+            return WriteResult(table_name)
+        return None
+
+    def write_batch_bundle(self, batch: ParquetBatch) -> None:
+        rows = batch.rows
+        batch.rows = []
+        pa_table = pa.Table.from_pylist(rows, batch.schema)
+        batch.writer.write_table(pa_table)
+
+    def insert_node(self, node: JsObject) -> None:
+        result = insert_node(
+            node,
+            self.kind_by_id,
+            self.insert_value,
+            with_tmp_prefix=False,
+            flatten=True,
+        )
+        should_write_batch = result and len(self.batches[result.table_name].rows) > self.rows_per_batch
+        if result and should_write_batch:
+            batch = self.batches[result.table_name]
+            self.write_batch_bundle(batch)
+
+    def close(self) -> None:
+        for batch in self.batches.values():
+            self.write_batch_bundle(batch)
+            batch.writer.close()