feat: Support Snowflake (#9)

Author: aquamatthias

.github/workflows/build_and_publish.yml

@@ -21,7 +21,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v4
         with:
-          python-version: '3.11'
+          python-version: '3.10'
           architecture: 'x64'
 
       - name: Restore dependency cache

cloud2sql/__main__.py

@@ -4,9 +4,14 @@
 from resotolib.logger import setup_logger
 from sqlalchemy import create_engine
 from sqlalchemy.engine import Engine
-
 from cloud2sql.collect_plugins import collect_from_plugins
 
+# Will fail in case snowflake is not installed - which is fine.
+try:
+    from cloud2sql.snowflake import SnowflakeUpdater  # noqa:F401
+except ImportError:
+    pass
+
 log = getLogger("cloud2sql")
 
 
@@ -40,9 +45,15 @@ def collect(engine: Engine, args: Namespace) -> None:
 
 def main() -> None:
     args = parse_args()
-    setup_logger("cloud2sql", level=args.log_level, force=True)
-    engine = create_engine(args.db)
-    collect(engine, args)
+    try:
+        setup_logger("cloud2sql", level=args.log_level, force=True)
+        engine = create_engine(args.db)
+        collect(engine, args)
+    except Exception as e:
+        if args.debug:  # raise exception and show complete tracelog
+            raise e
+        else:
+            print(f"Error syncing data to database: {e}")
 
 
 if __name__ == "__main__":

cloud2sql/collect_plugins.py

@@ -1,5 +1,6 @@
 import concurrent
 import multiprocessing
+from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, Future
 from contextlib import suppress
 from logging import getLogger
@@ -25,9 +26,9 @@
 from sqlalchemy.engine import Engine
 
 from cloud2sql.show_progress import CollectInfo
-from cloud2sql.sql import SqlModel, SqlUpdater
+from cloud2sql.sql import SqlUpdater, sql_updater
 
-log = getLogger("cloud2sql")
+log = getLogger("resoto.cloud2sql")
 
 
 def collectors(raw_config: Json, feedback: CoreFeedback) -> Dict[str, BaseCollectorPlugin]:
@@ -62,24 +63,25 @@ def collect(collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedba
     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 = SqlModel(Model({k.fqn: k for k in kinds}))
+    updater = sql_updater(Model({k.fqn: k for k in kinds}), engine)
     node_edge_count = len(collector.graph.nodes) + len(collector.graph.edges)
-    ne_count = iter(range(0, node_edge_count))
-    progress_update = max(node_edge_count // 100, 50)
+    ne_count = 0
     schema = f"create temp tables {engine.dialect.name}"
     syncdb = f"synchronize {engine.dialect.name}"
     feedback.progress_done(schema, 0, 1, context=[collector.cloud])
     feedback.progress_done(syncdb, 0, node_edge_count, context=[collector.cloud])
     with engine.connect() as conn:
         with conn.begin():
             # create the ddl metadata from the kinds
-            model.create_schema(conn, args)
+            updater.create_schema(conn, args)
             feedback.progress_done(schema, 1, 1, context=[collector.cloud])
-            # ingest the data
-            updater = SqlUpdater(model)
+
+            # 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"] = {
@@ -88,17 +90,29 @@ def collect(collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedba
                     "region": {"reported": {"id": node.region().name}},
                     "zone": {"reported": {"id": node.zone().name}},
                 }
-                stmt = updater.insert_node(exported)
-                if stmt is not None:
-                    conn.execute(stmt)
-                if (nx := next(ne_count)) % progress_update == 0:
-                    feedback.progress_done(syncdb, nx, node_edge_count, context=[collector.cloud])
+                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}")
+                for insert in updater.insert_nodes(kind, nodes):
+                    conn.execute(insert)
+                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:
-                stmt = updater.insert_node({"from": from_node.chksum, "to": to_node.chksum, "type": "edge"})
-                if stmt is not None:
-                    conn.execute(stmt)
-                if (nx := next(ne_count)) % progress_update == 0:
-                    feedback.progress_done(syncdb, nx, node_edge_count, context=[collector.cloud])
+                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]}")
+                for insert in updater.insert_edges(from_to, nodes):
+                    conn.execute(insert)
+                ne_count += len(nodes)
+                feedback.progress_done(syncdb, ne_count, node_edge_count, context=[collector.cloud])
     feedback.progress_done(collector.cloud, 1, 1)
 
 
@@ -131,7 +145,10 @@ def collect_from_plugins(engine: Engine, args: Namespace) -> None:
             for future in concurrent.futures.as_completed(futures):
                 future.result()
             # when all collectors are done, we can swap all temp tables
-            SqlModel.swap_temp_tables(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/collect_resoto.py

@@ -0,0 +1,93 @@
+import json
+import logging
+from typing import Any, List, Iterator
+
+from resotoclient import Model
+from resotoclient.models import Property
+from resotolib.types import Json
+from snowflake.sqlalchemy import ARRAY, OBJECT
+from sqlalchemy import Integer, Float, String, Boolean, column
+from sqlalchemy import select
+from sqlalchemy.sql import Values
+from sqlalchemy.sql.dml import ValuesBase
+
+from cloud2sql.sql import SqlDefaultUpdater, DialectUpdater
+
+log = logging.getLogger("resoto.cloud2sql.snowflake")
+
+
+def kind_to_snowflake_type(kind_name: str, model: Model) -> Any:  # Type[TypeEngine[Any]]
+    """
+    Map internal kinds to snowflake types.
+    More or less the default mapping, but with some special cases for OBJECT and ARRAY types.
+    """
+    kind = model.kinds.get(kind_name)
+    if "[]" in kind_name:
+        return ARRAY
+    elif kind_name.startswith("dict"):
+        return OBJECT
+    elif kind_name == "any":
+        return OBJECT
+    elif kind_name in ("int32", "int64"):
+        return Integer
+    elif kind_name in "float":
+        return Float
+    elif kind_name in "double":
+        return Float  # use Double with sqlalchemy 2
+    elif kind_name in ("string", "date", "datetime", "duration"):
+        return String
+    elif kind_name == "boolean":
+        return Boolean
+    elif kind.runtime_kind is not None:  # refined simple type like enum
+        return kind_to_snowflake_type(kind.runtime_kind, model)
+    elif kind.properties:  # complex kind
+        return OBJECT
+    else:
+        raise ValueError(f"Not able to handle kind {kind_name}")
+
+
+class SnowflakeUpdater(SqlDefaultUpdater):
+    """
+    This updater synchronizes resource data to snowflake https://www.snowflake.com
+    Snowflake needs special handling, since it does not support default json or array types.
+    It also does not understand json or array types as bind parameters.
+    This updater handles those shortcomings by using special insert statements.
+    """
+
+    def __init__(self, model: Model, **args: Any) -> None:
+        super().__init__(model, **args)
+        self.column_types_fn = kind_to_snowflake_type
+
+    def insert_nodes(self, kind: str, nodes: List[Json]) -> Iterator[ValuesBase]:
+        kp, _ = self.kind_properties(self.model.kinds[kind])
+        kind_props = [Property("_id", "string")] + kp
+        select_array = []
+        column_definitions = []
+        prop_is_json = {}
+
+        # Inserting structured data into Snowflake requires a bit of work. General scheme:
+        # insert into TBL(col_string, col_json) SELECT column1, parse_json(column2) from values('a', '{"b":1}');
+        # All json and array elements need to be json encoded and parsed on the server side again.
+        for num, prop in enumerate(kind_props):
+            name = f"column{num+1}"
+            select_array.append(prop.name)
+            snowflake_kind = kind_to_snowflake_type(prop.kind, self.model)
+            if snowflake_kind in (ARRAY, OBJECT):
+                column_definitions.append(column(f"parse_json({name})", is_literal=True))
+                prop_is_json[prop.name] = True
+            else:
+                column_definitions.append(column(name))
+
+        def values_tuple(node: Json) -> List[Any]:
+            nj = self.node_to_json(node)
+            # make sure to use the same order as in select_array
+            return [json.dumps(nj.get(p.name)) if prop_is_json.get(p.name) else nj.get(p.name) for p in kind_props]
+
+        if (table := self.metadata.tables.get(self.table_name(kind))) is not None:
+            for batch in (nodes[i : i + self.insert_batch_size] for i in range(0, len(nodes), self.insert_batch_size)):
+                converted = [values_tuple(node) for node in batch]
+                yield table.insert().from_select(select_array, select(Values(*column_definitions).data(converted)))
+
+
+# register this updater for the snowflake dialect, when snowflake is installed
+DialectUpdater["snowflake"] = SnowflakeUpdater