only write tags to clickhouse (#903)

* only write tags to clickhouse

* fix build

Author: dinmukhamedm

app-server/src/api/v1/evaluators.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use actix_web::{
     HttpResponse, post,
     web::{Data, Json},
@@ -12,9 +14,10 @@ use crate::{
         DB,
         evaluators::{EvaluatorScoreSource, insert_evaluator_score},
         project_api_keys::ProjectApiKey,
-        spans::{get_root_span_id, is_span_in_project},
     },
+    query_engine::QueryEngine,
     routes::types::ResponseResult,
+    sql::{self, ClickhouseReadonlyClient},
 };
 
 #[derive(Deserialize)]
@@ -55,9 +58,14 @@ pub async fn create_evaluator_score(
     req: Json<CreateEvaluatorScoreRequest>,
     db: Data<DB>,
     clickhouse: Data<clickhouse::Client>,
+    clickhouse_ro: Data<Option<Arc<ClickhouseReadonlyClient>>>,
+    query_engine: Data<Arc<QueryEngine>>,
     project_api_key: ProjectApiKey,
 ) -> ResponseResult {
     let req = req.into_inner();
+    let clickhouse_ro = clickhouse_ro.as_ref().clone().unwrap();
+    let query_engine = query_engine.as_ref().clone();
+    let clickhouse = clickhouse.as_ref().clone();
 
     // Extract common fields from both variants
     let (name, metadata, score, source) = match &req {
@@ -85,11 +93,21 @@ pub async fn create_evaluator_score(
 
     let span_id = match &req {
         CreateEvaluatorScoreRequest::WithTraceId(req) => {
-            get_root_span_id(&db.pool, &req.trace_id, &project_api_key.project_id).await?
+            sql::queries::get_top_span_id(
+                clickhouse_ro,
+                query_engine,
+                req.trace_id,
+                project_api_key.project_id,
+            )
+            .await?
         }
         CreateEvaluatorScoreRequest::WithSpanId(req) => {
-            let exists =
-                is_span_in_project(&db.pool, &req.span_id, &project_api_key.project_id).await?;
+            let exists = crate::ch::spans::is_span_in_project(
+                clickhouse.clone(),
+                req.span_id,
+                project_api_key.project_id,
+            )
+            .await?;
             if !exists {
                 return Ok(HttpResponse::NotFound().body("No matching spans found"));
             }
@@ -118,7 +136,7 @@ pub async fn create_evaluator_score(
     .await?;
 
     let _ = insert_evaluator_score_ch(
-        clickhouse.into_inner().as_ref().clone(),
+        clickhouse.clone(),
         score_id,
         project_id,
         name,

app-server/src/api/v1/tag.rs

@@ -1,8 +1,12 @@
+use std::sync::Arc;
+
 use crate::{
     ch::spans::append_tags_to_span,
-    db::{self, DB, project_api_keys::ProjectApiKey, tags::TagSource},
+    db::{DB, project_api_keys::ProjectApiKey, tags::TagSource},
+    query_engine::QueryEngine,
     routes::types::ResponseResult,
-    tags::insert_or_update_tag,
+    sql::{self, ClickhouseReadonlyClient},
+    tags::create_tag,
 };
 use actix_web::{
     HttpResponse, post,
@@ -39,6 +43,8 @@ pub async fn tag_trace(
     req: Json<TagRequest>,
     db: web::Data<DB>,
     clickhouse: web::Data<clickhouse::Client>,
+    clickhouse_ro: web::Data<Option<Arc<ClickhouseReadonlyClient>>>,
+    query_engine: web::Data<Arc<QueryEngine>>,
     project_api_key: ProjectApiKey,
 ) -> ResponseResult {
     let req = req.into_inner();
@@ -49,15 +55,27 @@ pub async fn tag_trace(
     if names.is_empty() {
         return Ok(HttpResponse::BadRequest().body("No names provided"));
     }
+    let clickhouse_ro = clickhouse_ro.as_ref().clone().unwrap();
+    let query_engine = query_engine.as_ref().clone();
+    let clickhouse = clickhouse.as_ref().clone();
+
     let span_id = match &req {
         TagRequest::WithTraceId(req) => {
-            db::spans::get_root_span_id(&db.pool, &req.trace_id, &project_api_key.project_id)
-                .await?
+            sql::queries::get_top_span_id(
+                clickhouse_ro,
+                query_engine,
+                req.trace_id,
+                project_api_key.project_id,
+            )
+            .await?
         }
         TagRequest::WithSpanId(req) => {
-            let exists =
-                db::spans::is_span_in_project(&db.pool, &req.span_id, &project_api_key.project_id)
-                    .await?;
+            let exists = crate::ch::spans::is_span_in_project(
+                clickhouse.clone(),
+                req.span_id,
+                project_api_key.project_id,
+            )
+            .await?;
             if !exists {
                 return Ok(HttpResponse::NotFound().body("No matching spans found"));
             }
@@ -69,26 +87,21 @@ pub async fn tag_trace(
         return Ok(HttpResponse::NotFound().body("No matching spans found"));
     };
 
-    let clickhouse = clickhouse.as_ref().clone();
-
     let futures = names
         .iter()
         .map(|name| {
-            insert_or_update_tag(
+            create_tag(
                 &db.pool,
                 clickhouse.clone(),
                 project_api_key.project_id,
-                Uuid::new_v4(),
                 span_id,
-                None,
-                None,
                 name.clone(),
                 TagSource::CODE,
             )
         })
         .collect::<Vec<_>>();
 
-    let tags = futures_util::future::try_join_all(futures).await?;
+    let tag_ids = futures_util::future::try_join_all(futures).await?;
 
     append_tags_to_span(
         clickhouse.clone(),
@@ -98,14 +111,12 @@ pub async fn tag_trace(
     )
     .await?;
 
-    let response = tags
+    let response = tag_ids
         .iter()
-        .map(|tag| {
+        .map(|id| {
             serde_json::json!({
-                "id": tag.id,
-                "spanId": tag.span_id,
-                "createdAt": tag.created_at,
-                "updatedAt": tag.updated_at,
+                "id": id,
+                "spanId": span_id,
             })
         })
         .collect::<Vec<_>>();

app-server/src/ch/spans.rs

@@ -218,3 +218,18 @@ pub async fn append_tags_to_span(
 
     Ok(())
 }
+
+pub async fn is_span_in_project(
+    clickhouse: clickhouse::Client,
+    span_id: Uuid,
+    project_id: Uuid,
+) -> Result<bool> {
+    let result = clickhouse
+        .query("SELECT count(*) FROM spans WHERE span_id = ? AND project_id = ?")
+        .bind(span_id)
+        .bind(project_id)
+        .fetch_one::<u64>()
+        .await?;
+
+    Ok(result > 0)
+}

app-server/src/ch/tags.rs

@@ -36,17 +36,11 @@ pub struct CHTag {
 }
 
 impl CHTag {
-    pub fn new(
-        project_id: Uuid,
-        class_id: Uuid,
-        id: Uuid,
-        name: String,
-        source: TagSource,
-        span_id: Uuid,
-    ) -> Self {
+    pub fn new(project_id: Uuid, id: Uuid, name: String, source: TagSource, span_id: Uuid) -> Self {
         Self {
             project_id,
-            class_id,
+            // TODO: Remove this once we drop the class_id column
+            class_id: Uuid::nil(),
             created_at: chrono_to_nanoseconds(Utc::now()),
             id,
             name,
@@ -59,13 +53,12 @@ impl CHTag {
 pub async fn insert_tag(
     client: clickhouse::Client,
     project_id: Uuid,
-    class_id: Uuid,
     id: Uuid,
     name: String,
     source: TagSource,
     span_id: Uuid,
 ) -> Result<()> {
-    let tag = CHTag::new(project_id, class_id, id, name, source, span_id);
+    let tag = CHTag::new(project_id, id, name, source, span_id);
     let ch_insert = client.insert("tags");
     match ch_insert {
         Ok(mut ch_insert) => {

app-server/src/db/spans.rs

@@ -4,7 +4,6 @@ use anyhow::Result;
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-use sqlx::PgPool;
 use uuid::Uuid;
 
 use crate::traces::spans::SpanAttributes;
@@ -63,39 +62,6 @@ pub struct Span {
     pub output_url: Option<String>,
 }
 
-pub async fn get_root_span_id(
-    pool: &PgPool,
-    trace_id: &Uuid,
-    project_id: &Uuid,
-) -> Result<Option<Uuid>> {
-    let span_id = sqlx::query_scalar::<_, Uuid>(
-        "SELECT span_id FROM spans
-        WHERE trace_id = $1
-        AND project_id = $2
-        AND parent_span_id IS NULL
-        ORDER BY start_time ASC
-        LIMIT 1",
-    )
-    .bind(trace_id)
-    .bind(project_id)
-    .fetch_optional(pool)
-    .await?;
-
-    Ok(span_id)
-}
-
-pub async fn is_span_in_project(pool: &PgPool, span_id: &Uuid, project_id: &Uuid) -> Result<bool> {
-    let exists = sqlx::query_scalar::<_, bool>(
-        "SELECT EXISTS(SELECT 1 FROM spans WHERE span_id = $1 AND project_id = $2)",
-    )
-    .bind(span_id)
-    .bind(project_id)
-    .fetch_one(pool)
-    .await?;
-
-    Ok(exists)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;