record tags in a batch (#977)

* record tags in a batch

* structured tag instead of tuple

---------

Co-authored-by: Din <[email protected]>

Author: skull8888888

app-server/src/ch/tags.rs

@@ -2,9 +2,10 @@ use anyhow::Result;
 use chrono::Utc;
 use clickhouse::Row;
 use serde::{Deserialize, Serialize};
+use tracing::instrument;
 use uuid::Uuid;
 
-use crate::db::tags::TagSource;
+use crate::db::tags::{SpanTag, TagSource};
 
 use super::utils::chrono_to_nanoseconds;
 
@@ -79,3 +80,44 @@ pub async fn insert_tag(
         }
     }
 }
+
+#[instrument(skip(client, tags))]
+pub async fn insert_tags_batch(client: clickhouse::Client, tags: &[SpanTag]) -> Result<()> {
+    if tags.is_empty() {
+        return Ok(());
+    }
+
+    let ch_insert = client.insert("tags");
+    match ch_insert {
+        Ok(mut ch_insert) => {
+            for span_tag in tags {
+                let id = Uuid::new_v4();
+                let tag = CHTag::new(
+                    span_tag.project_id,
+                    id,
+                    span_tag.name.clone(),
+                    span_tag.source.clone(),
+                    span_tag.span_id,
+                );
+                ch_insert.write(&tag).await?;
+            }
+
+            let ch_insert_end_res = ch_insert.end().await;
+            match ch_insert_end_res {
+                Ok(_) => Ok(()),
+                Err(e) => {
+                    return Err(anyhow::anyhow!(
+                        "Clickhouse batch tag insertion failed: {:?}",
+                        e
+                    ));
+                }
+            }
+        }
+        Err(e) => {
+            return Err(anyhow::anyhow!(
+                "Failed to insert tags batch into Clickhouse: {:?}",
+                e
+            ));
+        }
+    }
+}

app-server/src/db/tags.rs

@@ -1,9 +1,30 @@
 use serde::{Deserialize, Serialize};
+use uuid::Uuid;
 
-#[derive(sqlx::Type, Serialize, Deserialize, Clone, PartialEq)]
+#[derive(Debug, sqlx::Type, Serialize, Deserialize, Clone, PartialEq)]
 #[sqlx(type_name = "tag_source")]
 pub enum TagSource {
     MANUAL,
     AUTO,
     CODE,
 }
+
+/// Structured representation of a span tag for batch operations
+#[derive(Debug, Clone)]
+pub struct SpanTag {
+    pub project_id: Uuid,
+    pub name: String,
+    pub source: TagSource,
+    pub span_id: Uuid,
+}
+
+impl SpanTag {
+    pub fn new(project_id: Uuid, name: String, source: TagSource, span_id: Uuid) -> Self {
+        Self {
+            project_id,
+            name,
+            source,
+            span_id,
+        }
+    }
+}

app-server/src/traces/consumer.rs

@@ -24,7 +24,13 @@ use crate::{
         spans::CHSpan,
         traces::{CHTrace, TraceAggregation, upsert_traces_batch},
     },
-    db::{DB, events::Event, spans::Span, trace::upsert_trace_statistics_batch},
+    db::{
+        DB,
+        events::Event,
+        spans::Span,
+        tags::{SpanTag, TagSource},
+        trace::upsert_trace_statistics_batch,
+    },
     evaluators::{get_evaluators_by_path, push_to_evaluators_queue},
     features::{Feature, is_feature_enabled},
     mq::{
@@ -38,7 +44,7 @@ use crate::{
         events::record_events,
         limits::update_workspace_limit_exceeded_by_project_id,
         provider::convert_span_to_provider_format,
-        utils::{get_llm_usage_for_span, prepare_span_for_recording, record_tags},
+        utils::{get_llm_usage_for_span, prepare_span_for_recording},
     },
 };
 
@@ -427,19 +433,22 @@ async fn process_batch(
         }
     }
 
-    for span in &stripped_spans {
-        if let Err(e) = record_tags(
-            clickhouse.clone(),
-            &span.tags,
-            &span.span_id,
-            &span.project_id,
-        )
-        .await
-        {
+    // Collect all tags from all spans for batch insertion
+    let tags_batch: Vec<SpanTag> = stripped_spans
+        .iter()
+        .flat_map(|span| {
+            span.tags.iter().map(move |tag| {
+                SpanTag::new(span.project_id, tag.clone(), TagSource::CODE, span.span_id)
+            })
+        })
+        .collect();
+
+    // Record all tags in a single batch
+    if !tags_batch.is_empty() {
+        if let Err(e) = crate::ch::tags::insert_tags_batch(clickhouse.clone(), &tags_batch).await {
             log::error!(
-                "Failed to record tags to DB. span_id [{}], project_id [{}]: {:?}",
-                span.span_id,
-                span.project_id,
+                "Failed to record tags to DB for batch of {} tags: {:?}",
+                tags_batch.len(),
                 e
             );
         }

app-server/src/traces/utils.rs

@@ -13,7 +13,7 @@ use crate::{
 
 use crate::{
     cache::Cache,
-    db::{DB, spans::Span, tags::TagSource},
+    db::{DB, spans::Span},
     language_model::costs::estimate_cost_by_provider_name,
 };
 
@@ -97,31 +97,6 @@ pub async fn get_llm_usage_for_span(
     }
 }
 
-#[instrument(skip(clickhouse, tags, span_id, project_id))]
-pub async fn record_tags(
-    clickhouse: clickhouse::Client,
-    tags: &[String],
-    span_id: &Uuid,
-    project_id: &Uuid,
-) -> anyhow::Result<()> {
-    if tags.is_empty() {
-        return Ok(());
-    }
-
-    for tag_name in tags {
-        crate::ch::tags::insert_tag(
-            clickhouse.clone(),
-            *project_id,
-            tag_name.clone(),
-            TagSource::CODE,
-            *span_id,
-        )
-        .await?;
-    }
-
-    Ok(())
-}
-
 pub fn skip_span_name(name: &str) -> bool {
     SKIP_SPAN_NAME_REGEX.is_match(name)
 }