create event definitions for new events from traces (#965)

Author: dinmukhamedm

app-server/src/cache/keys.rs

@@ -7,5 +7,5 @@ pub const PROJECT_CACHE_KEY: &str = "project";
 pub const WORKSPACE_LIMITS_CACHE_KEY: &str = "workspace_limits";
 pub const PROJECT_EVALUATORS_BY_PATH_CACHE_KEY: &str = "project_evaluators_by_path";
 pub const SUMMARY_TRIGGER_SPANS_CACHE_KEY: &str = "summary_trigger_spans";
-
+pub const PROJECT_EVENT_NAMES_CACHE_KEY: &str = "project_event_names";
 pub const WORKSPACE_BYTES_USAGE_CACHE_KEY: &str = "workspace_bytes_usage";

app-server/src/db/event_definitions.rs

@@ -0,0 +1,32 @@
+use anyhow::Result;
+use sqlx::PgPool;
+use uuid::Uuid;
+
+pub async fn get_event_definition_names(pool: &PgPool, project_id: &Uuid) -> Result<Vec<String>> {
+    let event_definitions =
+        sqlx::query_scalar::<_, String>("SELECT name FROM event_definitions WHERE project_id = $1")
+            .bind(project_id)
+            .fetch_all(pool)
+            .await?;
+
+    Ok(event_definitions)
+}
+
+pub async fn insert_event_definition_names(
+    pool: &PgPool,
+    project_id: &Uuid,
+    names: &Vec<String>,
+) -> Result<()> {
+    sqlx::query(
+        "INSERT INTO event_definitions (name, project_id)
+        SELECT UNNEST($1::text[]) AS name, $2 AS project_id
+        ON CONFLICT (name, project_id) DO NOTHING
+        ",
+    )
+    .bind(names)
+    .bind(project_id)
+    .execute(pool)
+    .await?;
+
+    Ok(())
+}

app-server/src/db/mod.rs

@@ -5,6 +5,7 @@ pub mod agent_messages;
 pub mod datasets;
 pub mod evaluations;
 pub mod evaluators;
+pub mod event_definitions;
 pub mod events;
 pub mod prices;
 pub mod project_api_keys;

app-server/src/traces/consumer.rs

@@ -361,7 +361,15 @@ async fn process_batch(
         log::error!("Failed to ack MQ delivery (batch): {:?}", e);
     });
 
-    let total_events_ingested_bytes = match record_events(clickhouse.clone(), &all_events).await {
+    let total_events_ingested_bytes = match record_events(
+        cache.clone(),
+        db.clone(),
+        project_id,
+        clickhouse.clone(),
+        &all_events,
+    )
+    .await
+    {
         Ok(bytes) => bytes,
         Err(e) => {
             log::error!("Failed to record events: {:?}", e);

app-server/src/traces/events.rs

@@ -1,18 +1,71 @@
+use std::collections::HashSet;
+use std::sync::Arc;
+
 use anyhow::Result;
+use uuid::Uuid;
 
 use crate::{
+    cache::{Cache, CacheTrait, keys::PROJECT_EVENT_NAMES_CACHE_KEY},
     ch::{self, events::CHEvent},
-    db::events::Event,
+    db::{self, DB, events::Event},
 };
 
-///
 pub async fn record_events(
+    cache: Arc<Cache>,
+    db: Arc<DB>,
+    project_id: Uuid,
     clickhouse: clickhouse::Client,
     event_payloads: &Vec<Event>,
 ) -> Result<usize> {
     let ch_events = event_payloads
         .iter()
         .map(|e| CHEvent::from_db_event(e))
         .collect::<Vec<CHEvent>>();
+    let event_names = ch_events
+        .iter()
+        .map(|e| e.name.clone())
+        .collect::<Vec<String>>();
+    tokio::spawn(async move {
+        let _ = insert_event_definition_names(db.clone(), cache.clone(), &project_id, event_names)
+            .await
+            .map_err(|e| log::error!("Failed to insert event definition names: {:?}", e));
+    });
     ch::events::insert_events(clickhouse, ch_events).await
 }
+
+async fn insert_event_definition_names(
+    db: Arc<DB>,
+    cache: Arc<Cache>,
+    project_id: &Uuid,
+    names: Vec<String>,
+) -> Result<()> {
+    let cache_key = format!("{PROJECT_EVENT_NAMES_CACHE_KEY}:{}", project_id);
+    let cached_names = cache.get::<Vec<String>>(&cache_key).await;
+    match cached_names {
+        Ok(Some(cached_names)) => {
+            let cached_names_set = HashSet::<String>::from_iter(cached_names);
+            let names_set = HashSet::from_iter(names);
+            if cached_names_set.is_superset(&names_set) {
+                return Ok(());
+            }
+            let new_names = cached_names_set
+                .union(&names_set)
+                .cloned()
+                .collect::<Vec<String>>();
+            // Found in cache, but this event name is new, insert it into the database and cache
+            db::event_definitions::insert_event_definition_names(&db.pool, project_id, &new_names)
+                .await?;
+            cache.insert(&cache_key, new_names).await?;
+            Ok(())
+        }
+        Err(_) | Ok(None) => {
+            // Not found in cache, insert it into the database, update the cache
+            db::event_definitions::insert_event_definition_names(&db.pool, project_id, &names)
+                .await?;
+            let new_names =
+                db::event_definitions::get_event_definition_names(&db.pool, project_id).await?;
+            cache.insert(&cache_key, new_names).await?;
+            Ok(())
+        }
+    }
+}