Give orchest api access to the analytics module

lib/python/orchest-internals/_orchest/internals/analytics.py

@@ -30,6 +30,7 @@ class Event(Enum):
     CONFIRM_SHOW = "confirm show"
     CRONJOB_PAUSE = "cron-job pause"
     CRONJOB_RESUME = "cron-job resume"
+    DEBUG_PING = "debug ping"
     ENVIRONMENT_BUILD_CANCEL = "environment-build cancel"
     ENVIRONMENT_BUILD_START = "environment-build start"
     HEARTBEAT_TRIGGER = "heartbeat trigger"
@@ -147,6 +148,7 @@ def send_event(
     telemetry_uuid = app.config.get("TELEMETRY_UUID")
     if telemetry_uuid is None:
         app.logger.error("No telemetry uuid found, won't send telemetry event.")
+        return False
 
     if not _posthog_initialized:
         _initialize_posthog()

services/orchest-api/app/app/__init__.py

@@ -31,6 +31,7 @@
 from app.apis.namespace_runs import AbortPipelineRun
 from app.apis.namespace_sessions import StopInteractiveSession
 from app.connections import db
+from app.core.notifications import analytics as api_analytics
 from app.core.scheduler import add_recurring_jobs_to_scheduler
 from app.models import (
     EnvironmentImageBuild,
@@ -105,6 +106,10 @@ def create_app(
             settings.save()
             app.config.update(settings.as_dict())
 
+    # Keep analytics subscribed to all events of interest.
+    with app.app_context():
+        api_analytics.upsert_analytics_subscriptions()
+
     # Create a background scheduler (in a daemon thread) for every
     # gunicorn worker. The individual schedulers do not cause duplicate
     # execution because all jobs of the all the schedulers read state

services/orchest-api/app/app/core/events.py

@@ -4,9 +4,12 @@
 accordingly based on any subscribers subscribed to the event type that
 happened.
 """
-from app import models, utils
+from app import models
+from app import utils
+from app import utils as app_utils
 from app.connections import db
 from app.core import notifications
+from app.core.notifications import analytics as api_analytics
 
 _logger = utils.get_logger()
 
@@ -29,15 +32,26 @@ def _register_event(ev: models.Event) -> None:
         ev.type, project_uuid=project_uuid, job_uuid=job_uuid
     )
     for sub in subscribers:
+        if isinstance(sub, models.AnalyticsSubscriber):
+            if app_utils.OrchestSettings()["TELEMETRY_DISABLED"]:
+                _logger.info(
+                    "Telemetry is disabled, skipping event delivery to analytics."
+                )
+                continue
+            notification_payload = api_analytics.generate_payload_for_analytics(ev)
+        else:
+            notification_payload = ev.to_notification_payload()
+
         _logger.info(
             f"Scheduling delivery for event {ev.uuid}, event type: {ev.type} for "
-            f"deliveree {sub.uuid}."
+            f"deliveree {sub.uuid} of type {sub.type}."
         )
+
         delivery = models.Delivery(
             event=ev.uuid,
             deliveree=sub.uuid,
             status="SCHEDULED",
-            notification_payload=ev.to_notification_payload(),
+            notification_payload=notification_payload,
         )
         db.session.add(delivery)
 

services/orchest-api/app/app/core/notifications/__init__.py

@@ -14,7 +14,7 @@
 from app import models
 from app import utils as app_utils
 from app.connections import db
-from app.core.notifications import webhooks
+from app.core.notifications import analytics, webhooks
 
 logger = app_utils.get_logger()
 
@@ -105,16 +105,17 @@ def get_subscribers_subscribed_to_event(
 def process_notifications_deliveries_task() -> None:
     logger.info("Processing notifications deliveries")
 
+    delivery_filter = [
+        models.Delivery.status.in_(["SCHEDULED", "RESCHEDULED"]),
+        models.Delivery.scheduled_at <= datetime.datetime.now(datetime.timezone.utc),
+    ]
+
     # Note that we don't select for update here, said locking will
     # happen on a single delivery basis.
     webhook_deliveries = (
         (db.session.query(models.Delivery.uuid))
         .join(models.Webhook, models.Webhook.uuid == models.Delivery.deliveree)
-        .filter(
-            models.Delivery.status.in_(["SCHEDULED", "RESCHEDULED"]),
-            models.Delivery.scheduled_at
-            <= datetime.datetime.now(datetime.timezone.utc),
-        )
+        .filter(*delivery_filter)
         .order_by(models.Delivery.scheduled_at)
         .all()
     )
@@ -127,3 +128,23 @@ def process_notifications_deliveries_task() -> None:
         # Don't let failures affect other deliveries.
         except Exception as e:
             logger.error(e)
+
+    analytics_deliveries = (
+        (db.session.query(models.Delivery.uuid))
+        .join(
+            models.AnalyticsSubscriber,
+            models.AnalyticsSubscriber.uuid == models.Delivery.deliveree,
+        )
+        .filter(*delivery_filter)
+        .order_by(models.Delivery.scheduled_at)
+        .all()
+    )
+
+    logger.info(f"Found {len(analytics_deliveries)} analytics deliveries to deliver.")
+
+    for delivery in analytics_deliveries:
+        try:
+            analytics.deliver(delivery.uuid)
+        # Don't let failures affect other deliveries.
+        except Exception as e:
+            logger.error(e)

services/orchest-api/app/app/core/notifications/analytics.py

@@ -0,0 +1,204 @@
+"""Implements the logic that sends events to analytics.
+
+This is a special case when it comes to notifications modules.
+Essentially, there can only be one subscriber of the type
+AnalyticsSubscriber, and the actual data like URL and secrets are stored
+in the code base since currently the analytics module is shared across
+multiple services.
+
+"""
+from flask import current_app
+
+from _orchest.internals import analytics
+from app import models
+from app import utils as app_utils
+from app.connections import db
+from app.core.notifications import utils as notification_utils
+
+logger = app_utils.get_logger()
+
+
+def upsert_analytics_subscriptions() -> None:
+    """Makes sure the analytics backend is subscribed to all events.
+
+    Commits to the db.
+    """
+
+    # Use an hardcoded uuid to make sure only 1 such subscriber exists,
+    # to avoid race conditions.
+    uuid = "c9075806-be93-4b87-9d77-9a376d347b3a"
+    subscriber = models.AnalyticsSubscriber.query.with_for_update().first()
+    if subscriber is None:
+        subscriber = models.AnalyticsSubscriber(uuid=uuid)
+        db.session.add(subscriber)
+        db.session.flush()
+
+    existing_subs = {
+        sub.event_type
+        for sub in models.Subscription.query.filter(
+            models.Subscription.subscriber_uuid == uuid
+        ).all()
+    }
+
+    event_types = [event_type.name for event_type in models.EventType.query.all()]
+    new_subscriptions = []
+    for event_type in event_types:
+        if event_type not in existing_subs:
+            new_subscriptions.append({"event_type": event_type})
+    subs = notification_utils.subscription_specs_to_subscriptions(
+        uuid, new_subscriptions
+    )
+    db.session.bulk_save_objects(subs)
+    db.session.commit()
+
+
+def send_test_ping_delivery() -> bool:
+    """Used internally for testing, send a ping event.
+
+    Returns:
+        True if the delivery was made, False otherwise.
+    """
+    # TODO: this is needed because a depency of this PR actually lives
+    # in the controller branch, remove this.
+    current_app.config["TELEMETRY_DISABLED"] = False
+    current_app.config["TELEMETRY_UUID"] = "e2e2abb5-3c54-4209-a245-e9fa37810824"
+    return analytics.send_event(current_app, analytics.Event.DEBUG_PING, {})
+
+
+# The ones mapped to DEBUG_PING are events which are not currently
+# defined in analytics, to be later expanded in a future commit.
+_event_type_to_analytics_event_enum = {
+    "project:one-off-job:created": analytics.Event.JOB_CREATE,
+    "project:one-off-job:started": analytics.Event.DEBUG_PING,
+    "project:one-off-job:deleted": analytics.Event.JOB_DELETE,
+    "project:one-off-job:cancelled": analytics.Event.JOB_CANCEL,
+    "project:one-off-job:failed": analytics.Event.DEBUG_PING,
+    "project:one-off-job:succeeded": analytics.Event.DEBUG_PING,
+    "project:one-off-job:pipeline-run:created": analytics.Event.DEBUG_PING,
+    "project:one-off-job:pipeline-run:started": analytics.Event.DEBUG_PING,
+    "project:one-off-job:pipeline-run:cancelled": analytics.Event.JOB_PIPELINE_RUN_CANCEL,  # noqa
+    "project:one-off-job:pipeline-run:failed": analytics.Event.DEBUG_PING,
+    "project:one-off-job:pipeline-run:deleted": analytics.Event.JOB_PIPELINE_RUN_DELETE,
+    "project:one-off-job:pipeline-run:succeeded": analytics.Event.DEBUG_PING,
+    "project:cron-job:created": analytics.Event.JOB_CREATE,
+    "project:cron-job:started": analytics.Event.DEBUG_PING,
+    "project:cron-job:deleted": analytics.Event.JOB_DELETE,
+    "project:cron-job:cancelled": analytics.Event.JOB_CANCEL,
+    "project:cron-job:failed": analytics.Event.DEBUG_PING,
+    "project:cron-job:paused": analytics.Event.CRONJOB_PAUSE,
+    "project:cron-job:unpaused": analytics.Event.CRONJOB_RESUME,
+    "project:cron-job:run:started": analytics.Event.DEBUG_PING,
+    "project:cron-job:run:succeeded": analytics.Event.DEBUG_PING,
+    "project:cron-job:run:failed": analytics.Event.DEBUG_PING,
+    "project:cron-job:run:pipeline-run:created": analytics.Event.DEBUG_PING,
+    "project:cron-job:run:pipeline-run:started": analytics.Event.DEBUG_PING,
+    "project:cron-job:run:pipeline-run:cancelled": analytics.Event.JOB_PIPELINE_RUN_CANCEL,  # noqa
+    "project:cron-job:run:pipeline-run:failed": analytics.Event.DEBUG_PING,
+    "project:cron-job:run:pipeline-run:deleted": analytics.Event.JOB_PIPELINE_RUN_DELETE,  # noqa
+    "project:cron-job:run:pipeline-run:succeeded": analytics.Event.DEBUG_PING,
+}
+
+
+def generate_payload_for_analytics(event: models.Event) -> dict:
+    """Creates an analytics module compatible payload.
+
+    Acts as a compatibility layer between orchest-api events and what
+    the shared analytics module expects, and also provides some "old"
+    fields to the analytics BE, to avoid breaking changes.
+    """
+
+    analytics_payload = event.to_notification_payload()
+    event_type = analytics_payload["type"]
+
+    if event_type.startswith("project:cron-job:") or event_type.startswith(
+        "project:one-off-job:"
+    ):
+        analytics_payload["job_uuid"] = analytics_payload["job"]["uuid"]
+
+    if event_type.startswith("project:cron-job:run:pipeline-run:"):
+        analytics_payload["run_uuid"] = analytics_payload["job"]["run"]["pipeline_run"][
+            "uuid"
+        ]
+    elif event_type.startswith("project:one-off-job:pipeline-run:"):
+        analytics_payload["run_uuid"] = analytics_payload["job"]["pipeline_run"]["uuid"]
+
+    if event_type in ["project:cron-job:created", "project:one-off-job:created"]:
+        analytics_payload["snapshot_size"] = 10
+        job: models.Job = models.Job.query.filter(
+            models.Job.project_uuid == analytics_payload["project"]["uuid"],
+            models.Job.uuid == analytics_payload["job"]["uuid"],
+        ).one()
+        analytics_payload["job_definition"] = {
+            "parameters": job.parameters,
+            "project_uuid": job.project_uuid,
+            "pipeline_uuid": job.pipeline_uuid,
+            "draft": True,
+            "uuid": job.uuid,
+            "pipeline_run_spec": {"run_type": "full", "uuids": []},
+        }
+
+    return analytics_payload
+
+
+def deliver(delivery_uuid: str) -> None:
+    """Delivers an analytics delivery. Will commit to the database.
+
+    If the delivery fails it's rescheduled in the future with a capped
+    exponential backoff. This function will commit to the database to
+    ensure that each delivery get its own transaction.
+
+    Args:
+        delivery_uuid: Delivery to be delivered, the associated
+            deliveree must be an AnalyticsSubscriber.
+
+    Raises:
+        ValueError: If the associated deliveree isn't an
+            AnalyticsSubscriber.
+
+    """
+    logger.info(f"Delivering {delivery_uuid}.")
+
+    delivery = (
+        models.Delivery.query.with_for_update(skip_locked=True)
+        .filter(
+            models.Delivery.uuid == delivery_uuid,
+            models.Delivery.status.in_(["SCHEDULED", "RESCHEDULED"]),
+        )
+        .first()
+    )
+    if delivery is None:
+        logger.info(f"No need to deliver {delivery_uuid}.")
+        return
+
+    deliveree = models.AnalyticsSubscriber.query.filter(
+        models.AnalyticsSubscriber.uuid == delivery.deliveree
+    ).first()
+    if deliveree is None:
+        raise ValueError("Deliveree of delivery isn't of type AnalyticsSubscriber.")
+
+    try:
+        # TODO: this is needed because a depency of this PR actually
+        # lives in the controller branch, remove this.
+        current_app.config["TELEMETRY_DISABLED"] = False
+        current_app.config["TELEMETRY_UUID"] = "e2e2abb5-3c54-4209-a245-e9fa37810824"
+
+        payload = delivery.notification_payload
+        analytics_event_type = _event_type_to_analytics_event_enum.get(
+            payload["type"], analytics.Event.DEBUG_PING
+        )
+        analytics.send_event(current_app, analytics_event_type, payload)
+        db.session.delete(delivery)
+    except Exception as e:
+        logger.error(e)
+        delivery.reschedule()
+
+        # Analytics is subscribed to all events, avoid having an unbound
+        # number of said deliveries in case things go wrong with the
+        # analytics back-end or if we move away from that.
+        if delivery.n_delivery_attempts > 5:
+            db.session.delete(delivery)
+            logger.info(f"Deleting {delivery_uuid}, couldn't deliver it.")
+        else:
+            logger.info(f"Rescheduling {delivery_uuid} at {delivery.scheduled_at}.")
+
+    db.session.commit()

services/orchest-api/app/app/models.py

@@ -1512,16 +1512,16 @@ class ContentType(enum.Enum):
 
     __tablename__ = None
 
-    url = db.Column(db.String(), nullable=False)
+    url = db.Column(db.String(), nullable=True)
 
-    name = db.Column(db.String(100), nullable=False)
+    name = db.Column(db.String(100), nullable=True)
 
-    verify_ssl = db.Column(db.Boolean(), nullable=False)
+    verify_ssl = db.Column(db.Boolean(), nullable=True)
 
     # Used to calculate the HMAC digest of the payload and sign it.
-    secret = deferred(db.Column(db.String(), nullable=False))
+    secret = deferred(db.Column(db.String(), nullable=True))
 
-    content_type = db.Column(db.String(50), nullable=False)
+    content_type = db.Column(db.String(50), nullable=True)
 
     def is_slack_webhook(self) -> bool:
         return self.url.startswith("https://hooks.slack.com/")
@@ -1537,6 +1537,15 @@ def is_teams_webhook(self) -> bool:
     }
 
 
+class AnalyticsSubscriber(Subscriber):
+
+    __tablename__ = None
+
+    __mapper_args__ = {
+        "polymorphic_identity": "analytics",
+    }
+
+
 class Subscription(BaseModel):
     __tablename__ = "subscriptions"