feat: Add analytics tracking (#11)

Author: aquamatthias

cloud2sql/__main__.py

@@ -1,9 +1,12 @@
 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
 from sqlalchemy.engine import Engine
+
+from cloud2sql.analytics import PosthogEventSender, NoEventSender, AnalyticsEventSender
 from cloud2sql.collect_plugins import collect_from_plugins
 
 # Will fail in case snowflake is not installed - which is fine.
@@ -16,7 +19,11 @@
 
 
 def parse_args() -> Namespace:
-    parser = ArgumentParser(epilog="Collect data from cloud providers and store it in a database")
+    parser = ArgumentParser(
+        description="Collect data from cloud providers and store it in a database",
+        epilog="Synchronizes cloud data to a database",
+        env_args_prefix="CLOUD2SQL_",
+    )
     parser.add_argument("--debug", action="store_true", help="Enable debug logging")
     parser.add_argument("--config", help="Path to config file", required=True)
     parser.add_argument(
@@ -30,30 +37,41 @@ def parse_args() -> Namespace:
         help="The database url. See https://docs.sqlalchemy.org/en/20/core/engines.html.",
         required=True,
     )
+    parser.add_argument(
+        "--analytics-opt-out",
+        default=False,
+        action="store_true",
+        help="Do not send anonimized analytics data",
+    )
     args = parser.parse_args()
     args.log_level = "CRITICAL" if args.show != "log" else "DEBUG" if args.debug else "INFO"
     return args  # type: ignore
 
 
-def collect(engine: Engine, args: Namespace) -> None:
+def collect(engine: Engine, args: Namespace, sender: AnalyticsEventSender) -> None:
     try:
-        collect_from_plugins(engine, args)
+        collect_from_plugins(engine, args, sender)
     except Exception as e:
         log.error("Error during collection", e)
         print(f"Error syncing data to database: {e}")
 
 
 def main() -> None:
     args = parse_args()
+    sender: Optional[AnalyticsEventSender] = None
     try:
-        setup_logger("cloud2sql", level=args.log_level, force=True)
+        setup_logger("resoto.cloud2sql", level=args.log_level, force=True)
+        sender = NoEventSender() if args.analytics_opt_out else PosthogEventSender()
         engine = create_engine(args.db)
-        collect(engine, args)
+        collect(engine, args, sender)
     except Exception as e:
         if args.debug:  # raise exception and show complete tracelog
             raise e
         else:
             print(f"Error syncing data to database: {e}")
+    finally:
+        if sender:
+            sender.flush()
 
 
 if __name__ == "__main__":

cloud2sql/analytics.py

@@ -0,0 +1,62 @@
+import logging
+import uuid
+from abc import abstractmethod, ABC
+from datetime import datetime
+from typing import Union
+
+from requests import get
+from posthog import Client
+
+log = logging.getLogger("resoto.cloud2sql")
+
+
+class AnalyticsEventSender(ABC):
+    """
+    Abstract event sender used to send analytics events.
+    """
+
+    @abstractmethod
+    def capture(self, kind: str, **context: Union[str, int, float]) -> None:
+        pass
+
+    def flush(self) -> None:
+        pass
+
+
+class NoEventSender(AnalyticsEventSender):
+    """
+    Use this sender to not emit any events other than writing it to the log file.
+    """
+
+    def capture(self, kind: str, **context: Union[str, int, float]) -> None:
+        log.debug(f"Capture event {kind}: {context}")
+
+
+class PosthogEventSender(AnalyticsEventSender):
+    """
+    This sender is used to collect analytics events in posthog.
+    """
+
+    def __init__(self) -> None:
+        self.client = Client(host="https://analytics.some.engineering", api_key="")
+        self.uid = uuid.uuid4()
+
+    def capture(self, kind: str, **context: Union[str, int, float]) -> None:
+        if self.client.api_key == "":
+            api_key = get("https://cdn.some.engineering/posthog/public_api_key").text.strip()
+            self.client.api_key = api_key
+            for consumer in self.client.consumers:
+                consumer.api_key = api_key
+        self.client.capture(
+            distinct_id=self.uid,
+            event="cloud2sql." + kind,
+            properties={
+                **context,
+                "source": "cloud2sql",
+                "run_id": self.uid,
+            },
+            timestamp=datetime.now(),
+        )
+
+    def flush(self) -> None:
+        self.client.flush()

cloud2sql/collect_plugins.py

@@ -25,6 +25,7 @@
 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
 
@@ -56,7 +57,9 @@ def configure(path_to_config: Optional[str]) -> Json:
     return {}
 
 
-def collect(collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedback, args: Namespace) -> None:
+def collect(
+    collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedback, args: Namespace
+) -> Tuple[str, int, int]:
     # collect cloud data
     feedback.progress_done(collector.cloud, 0, 1)
     collector.collect()
@@ -115,26 +118,27 @@ def collect(collector: BaseCollectorPlugin, engine: Engine, feedback: CoreFeedba
                 ne_count += len(nodes)
                 feedback.progress_done(syncdb, ne_count, node_edge_count, context=[collector.cloud])
     feedback.progress_done(collector.cloud, 1, 1)
+    return collector.cloud, len(collector.graph.nodes), len(collector.graph.edges)
 
 
 def show_messages(core_messages: Queue[Json], end: Event) -> None:
     info = CollectInfo()
     while not end.is_set():
-        with Live(info.render(), auto_refresh=False, transient=True) as live:
+        with Live(info.render(), transient=True):
             with suppress(Exception):
                 info.handle_message(core_messages.get(timeout=1))
-                live.update(info.render())
     for message in info.rendered_messages():
         rich_print(message)
 
 
-def collect_from_plugins(engine: Engine, args: Namespace) -> None:
+def collect_from_plugins(engine: 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)
     all_collectors = collectors(raw_config, feedback)
+    analytics = {"total": len(all_collectors), "engine": engine.dialect.name} | {name: 1 for name in all_collectors}
     end = Event()
     with ThreadPoolExecutor(max_workers=4) as executor:
         try:
@@ -144,7 +148,10 @@ def collect_from_plugins(engine: Engine, args: Namespace) -> None:
             for collector in all_collectors.values():
                 futures.append(executor.submit(collect, collector, engine, feedback, args))
             for future in concurrent.futures.as_completed(futures):
-                future.result()
+                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)
@@ -153,7 +160,9 @@ def collect_from_plugins(engine: Engine, args: Namespace) -> None:
         except Exception as e:
             # set end and wait for live to finish, otherwise the cursor is not reset
             end.set()
+            sender.capture("error", error=str(e))
             sleep(1)
+            sender.flush()
             log.error("An error occurred", exc_info=e)
             print(f"Encountered Error. Giving up. {e}")
             emergency_shutdown()

requirements.txt

@@ -5,6 +5,8 @@ SQLAlchemy==1.4.45
 PyYAML>=6.0
 rich>=12.6.0
 resotoclient>=1.2.1
+posthog>=2.2.0
+requests>=2.28.1
 
 # bundle popular db drivers
 psycopg2-binary>=2.9.5 # postgres

tests/collect_plugins_test.py

@@ -4,6 +4,7 @@
 from sqlalchemy import MetaData, create_engine
 from sqlalchemy.orm import Session
 
+from cloud2sql.analytics import NoEventSender
 from cloud2sql.collect_plugins import collect_from_plugins
 
 
@@ -13,7 +14,7 @@ def test_collect() -> None:
         cfg = f"{tmp}/config.yml"
         with open(cfg, "w") as f:
             f.write("example: {}\n")
-        collect_from_plugins(engine, Namespace(config=cfg, show="none"))
+        collect_from_plugins(engine, Namespace(config=cfg, show="none"), NoEventSender())
         # get all tables
         metadata = MetaData()
         metadata.reflect(bind=engine)