Feat: Telemetry - Added "install" event type to track failed installs (#115)

Co-authored-by: bindipankhudi <[email protected]>
Co-authored-by: Aaron ("AJ") Steers <[email protected]>

Author: 3 people

airbyte/_util/telemetry.py

@@ -166,12 +166,17 @@ def _get_analytics_id() -> str | None:
 _ANALYTICS_ID = _get_analytics_id()
 
 
-class SyncState(str, Enum):
+class EventState(str, Enum):
     STARTED = "started"
     FAILED = "failed"
     SUCCEEDED = "succeeded"
 
 
+class EventType(str, Enum):
+    INSTALL = "install"
+    SYNC = "sync"
+
+
 @dataclass
 class CacheTelemetryInfo:
     type: str
@@ -198,6 +203,14 @@ def from_source(cls, source: Source) -> SourceTelemetryInfo:
             version=source.executor.reported_version,
         )
 
+    @classmethod
+    def from_name(cls, name: str) -> SourceTelemetryInfo:
+        return cls(
+            name=name,
+            executor_type="unknown",
+            version="unknown",
+        )
+
 
 def one_way_hash(
     string_to_hash: Any,  # noqa: ANN401  # Allow Any type
@@ -229,9 +242,10 @@ def get_env_flags() -> dict[str, Any]:
 
 
 def send_telemetry(
-    source: Source,
+    source: Source | str,
     cache: CacheBase | None,
-    state: SyncState,
+    state: EventState,
+    event_type: EventType,
     number_of_records: int | None = None,
     exception: Exception | None = None,
 ) -> None:
@@ -241,7 +255,6 @@ def send_telemetry(
 
     payload_props: dict[str, str | int | dict] = {
         "session_id": PYAIRBYTE_SESSION_ID,
-        "source": asdict(SourceTelemetryInfo.from_source(source)),
         "cache": asdict(CacheTelemetryInfo.from_cache(cache)),
         "state": state,
         "version": get_version(),
@@ -250,6 +263,12 @@ def send_telemetry(
         "application_hash": one_way_hash(meta.get_application_name()),
         "flags": get_env_flags(),
     }
+
+    if isinstance(source, str):
+        payload_props["source"] = asdict(SourceTelemetryInfo.from_name(source))
+    else:
+        payload_props["source"] = asdict(SourceTelemetryInfo.from_source(source))
+
     if exception:
         if isinstance(exception, exc.AirbyteError):
             payload_props["exception"] = exception.safe_logging_dict()
@@ -267,7 +286,7 @@ def send_telemetry(
             auth=(PYAIRBYTE_APP_TRACKING_KEY, ""),
             json={
                 "anonymousId": _get_analytics_id(),
-                "event": "sync",
+                "event": event_type,
                 "properties": payload_props,
                 "timestamp": datetime.datetime.utcnow().isoformat(),  # noqa: DTZ003
             },

airbyte/sources/base.py

@@ -29,10 +29,7 @@
 from airbyte import exceptions as exc
 from airbyte._util import protocol_util
 from airbyte._util.name_normalizers import normalize_records
-from airbyte._util.telemetry import (
-    SyncState,
-    send_telemetry,
-)
+from airbyte._util.telemetry import EventState, EventType, send_telemetry
 from airbyte.caches.util import get_default_cache
 from airbyte.datasets._lazy import LazyDataset
 from airbyte.progress import progress
@@ -515,7 +512,8 @@ def _log_sync_start(
         send_telemetry(
             source=self,
             cache=cache,
-            state=SyncState.STARTED,
+            state=EventState.STARTED,
+            event_type=EventType.SYNC,
         )
 
     def _log_sync_success(
@@ -528,8 +526,9 @@ def _log_sync_success(
         send_telemetry(
             source=self,
             cache=cache,
-            state=SyncState.SUCCEEDED,
+            state=EventState.SUCCEEDED,
             number_of_records=self._processed_records,
+            event_type=EventType.SYNC,
         )
 
     def _log_sync_failure(
@@ -541,11 +540,12 @@ def _log_sync_failure(
         """Log the failure of a sync operation."""
         print(f"Failed `{self.name}` read operation at {pendulum.now().format('HH:mm:ss')}.")
         send_telemetry(
-            state=SyncState.FAILED,
+            state=EventState.FAILED,
             source=self,
             cache=cache,
             number_of_records=self._processed_records,
             exception=exception,
+            event_type=EventType.SYNC,
         )
 
     def read(

airbyte/sources/util.py

@@ -10,6 +10,7 @@
 
 from airbyte import exceptions as exc
 from airbyte._executor import PathExecutor, VenvExecutor
+from airbyte._util.telemetry import EventState, EventType, send_telemetry
 from airbyte.sources.base import Source
 from airbyte.sources.registry import ConnectorMetadata, get_connector_metadata
 
@@ -106,28 +107,48 @@ def get_source(
     metadata: ConnectorMetadata | None = None
     try:
         metadata = get_connector_metadata(name)
-    except exc.AirbyteConnectorNotRegisteredError:
+    except exc.AirbyteConnectorNotRegisteredError as ex:
         if not pip_url:
+            _log_install_state(name, state=EventState.FAILED, exception=ex)
             # We don't have a pip url or registry entry, so we can't install the connector
             raise
 
-    executor = VenvExecutor(
-        name=name,
-        metadata=metadata,
-        target_version=version,
-        pip_url=pip_url,
-    )
-    if install_if_missing:
-        executor.ensure_installation()
+    try:
+        executor = VenvExecutor(
+            name=name,
+            metadata=metadata,
+            target_version=version,
+            pip_url=pip_url,
+        )
+        if install_if_missing:
+            executor.ensure_installation()
 
-    return Source(
-        name=name,
-        config=config,
-        streams=streams,
-        executor=executor,
-    )
+        return Source(
+            name=name,
+            config=config,
+            streams=streams,
+            executor=executor,
+        )
+    except Exception as e:
+        _log_install_state(name, state=EventState.FAILED, exception=e)
+        raise
 
 
 __all__ = [
     "get_source",
 ]
+
+
+def _log_install_state(
+    name: str,
+    state: EventState,
+    exception: Exception | None = None,
+) -> None:
+    """Log an install event."""
+    send_telemetry(
+        source=name,
+        cache=None,
+        state=state,
+        event_type=EventType.INSTALL,
+        exception=exception,
+    )

tests/unit_tests/test_anonymous_usage_stats.py

@@ -34,8 +34,9 @@ def test_telemetry_track(monkeypatch):
     telemetry.send_telemetry(
         source=source_test,
         cache=cache,
-        state="started",
+        state=telemetry.EventState.STARTED,
         number_of_records=0,
+        event_type=telemetry.EventType.SYNC,
     )
 
     # Check that one request was made
@@ -78,8 +79,9 @@ def test_do_not_track(monkeypatch, do_not_track):
     telemetry.send_telemetry(
         source=source_test,
         cache=cache,
-        state="started",
+        state=telemetry.EventState.STARTED,
         number_of_records=0,
+        event_type=telemetry.EventType.SYNC,
     )
 
     # Check that zero requests were made, because DO_NOT_TRACK is set