initial public release

src/api-service/__app__/onefuzzlib/telemetry.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import logging
+import os
+from typing import Any, Dict, Optional, Union
+
+from onefuzztypes.enums import TelemetryData, TelemetryEvent
+from opencensus.ext.azure.log_exporter import AzureLogHandler
+
+LOCAL_CLIENT: Optional[logging.Logger] = None
+CENTRAL_CLIENT: Optional[logging.Logger] = None
+
+
+def _get_client(environ_key: str) -> Optional[logging.Logger]:
+    key = os.environ.get(environ_key)
+    if key is None:
+        return None
+    client = logging.getLogger("onefuzz")
+    client.addHandler(AzureLogHandler(connection_string="InstrumentationKey=%s" % key))
+    return client
+
+
+def _central_client() -> Optional[logging.Logger]:
+    global CENTRAL_CLIENT
+    if not CENTRAL_CLIENT:
+        CENTRAL_CLIENT = _get_client("ONEFUZZ_TELEMETRY")
+    return CENTRAL_CLIENT
+
+
+def _local_client() -> Union[None, Any, logging.Logger]:
+    global LOCAL_CLIENT
+    if not LOCAL_CLIENT:
+        LOCAL_CLIENT = _get_client("APPINSIGHTS_INSTRUMENTATIONKEY")
+    return LOCAL_CLIENT
+
+
+# NOTE: All telemetry that is *NOT* using the ORM telemetry_include should
+# go through this method
+#
+# This provides a point of inspection to know if it's data that is safe to
+# log to the central OneFuzz telemetry point
+def track_event(
+    event: TelemetryEvent, data: Dict[TelemetryData, Union[str, int]]
+) -> None:
+    central = _central_client()
+    local = _local_client()
+
+    if local:
+        serialized = {k.name: v for (k, v) in data.items()}
+        local.info(event.name, extra={"custom_dimensions": serialized})
+
+    if event in TelemetryEvent.can_share() and central:
+        serialized = {
+            k.name: v for (k, v) in data.items() if k in TelemetryData.can_share()
+        }
+        central.info(event.name, extra={"custom_dimensions": serialized})
+
+
+# NOTE: This should *only* be used for logging Telemetry data that uses
+# the ORM telemetry_include method to limit data for telemetry.
+def track_event_filtered(event: TelemetryEvent, data: Any) -> None:
+    central = _central_client()
+    local = _local_client()
+
+    if local:
+        local.info(event.name, extra={"custom_dimensions": data})
+
+    if central and event in TelemetryEvent.can_share():
+        central.info(event.name, extra={"custom_dimensions": data})