Allow spans to be deleted (#185)

* Allow spans to be deleted Allow deleting a specific span or all spans of a trace * Implement delete fn's on D1Store * Remove irrelevant assignments
fiberplane · Aug 27, 2024 · 39b1c10 · 39b1c10
1 parent 0660dc9
commit 39b1c10
Show file tree

Hide file tree

Showing 10 changed files with 354 additions and 67 deletions.
diff --git a/fpx-lib/src/api.rs b/fpx-lib/src/api.rs
@@ -54,19 +54,20 @@ impl Builder {
 
         let router = axum::Router::new()
             .route("/v1/traces", post(handlers::otel::trace_collector_handler))
+            .route("/api/traces", get(handlers::traces::traces_list_handler))
             .route(
-                "/api/traces/:trace_id/spans/:span_id",
-                get(handlers::spans::span_get_handler),
+                "/api/traces/:trace_id",
+                get(handlers::traces::traces_get_handler)
+                    .delete(handlers::traces::traces_delete_handler),
             )
             .route(
                 "/api/traces/:trace_id/spans",
                 get(handlers::spans::span_list_handler),
             )
             .route(
-                "/api/traces/:trace_id",
-                get(handlers::traces::traces_get_handler),
+                "/api/traces/:trace_id/spans/:span_id",
+                get(handlers::spans::span_get_handler).delete(handlers::spans::span_delete_handler),
             )
-            .route("/api/traces", get(handlers::traces::traces_list_handler))
             .route(
                 "/ts-compat/v1/traces/:trace_id/spans",
                 get(handlers::spans::ts_compat_span_list_handler),

diff --git a/fpx-lib/src/api/handlers/spans.rs b/fpx-lib/src/api/handlers/spans.rs
@@ -101,3 +101,40 @@ impl From<DbError> for ApiServerError<SpanListError> {
         ApiServerError::CommonError(CommonError::InternalServerError)
     }
 }
+
+#[tracing::instrument(skip_all)]
+pub async fn span_delete_handler(
+    State(store): State<BoxedStore>,
+    Path((trace_id, span_id)): Path<(String, String)>,
+) -> Result<StatusCode, ApiServerError<SpanDeleteError>> {
+    let tx = store.start_readonly_transaction().await?;
+
+    hex::decode(&trace_id)
+        .map_err(|_| ApiServerError::ServiceError(SpanDeleteError::InvalidTraceId))?;
+    hex::decode(&span_id)
+        .map_err(|_| ApiServerError::ServiceError(SpanDeleteError::InvalidSpanId))?;
+
+    store.span_delete(&tx, &trace_id, &span_id).await?;
+
+    Ok(StatusCode::NO_CONTENT)
+}
+
+#[derive(Debug, Serialize, Deserialize, Error, ApiError)]
+#[serde(tag = "error", content = "details", rename_all = "camelCase")]
+#[non_exhaustive]
+pub enum SpanDeleteError {
+    #[api_error(status_code = StatusCode::BAD_REQUEST)]
+    #[error("Trace ID is invalid")]
+    InvalidTraceId,
+
+    #[api_error(status_code = StatusCode::BAD_REQUEST)]
+    #[error("Trace ID is invalid")]
+    InvalidSpanId,
+}
+
+impl From<DbError> for ApiServerError<SpanDeleteError> {
+    fn from(err: DbError) -> Self {
+        error!(?err, "Failed to list spans from database");
+        ApiServerError::CommonError(CommonError::InternalServerError)
+    }
+}
diff --git a/fpx-lib/src/api/handlers/traces.rs b/fpx-lib/src/api/handlers/traces.rs
@@ -100,3 +100,35 @@ impl From<DbError> for ApiServerError<TraceGetError> {
         ApiServerError::CommonError(CommonError::InternalServerError)
     }
 }
+
+#[tracing::instrument(skip_all)]
+pub async fn traces_delete_handler(
+    State(store): State<BoxedStore>,
+    Path(trace_id): Path<String>,
+) -> Result<StatusCode, ApiServerError<TraceDeleteError>> {
+    let tx = store.start_readonly_transaction().await?;
+
+    hex::decode(&trace_id)
+        .map_err(|_| ApiServerError::ServiceError(TraceDeleteError::InvalidTraceId))?;
+
+    // Retrieve all the spans that are associated with the trace
+    store.span_delete_by_trace(&tx, &trace_id).await?;
+
+    Ok(StatusCode::NO_CONTENT)
+}
+
+#[derive(Debug, Serialize, Deserialize, Error, ApiError)]
+#[serde(tag = "error", content = "details", rename_all = "camelCase")]
+#[non_exhaustive]
+pub enum TraceDeleteError {
+    #[api_error(status_code = StatusCode::BAD_REQUEST)]
+    #[error("Trace ID is invalid")]
+    InvalidTraceId,
+}
+
+impl From<DbError> for ApiServerError<TraceDeleteError> {
+    fn from(err: DbError) -> Self {
+        error!(?err, "Failed to list trace from database");
+        ApiServerError::CommonError(CommonError::InternalServerError)
+    }
+}
diff --git a/fpx-lib/src/data.rs b/fpx-lib/src/data.rs
@@ -69,4 +69,15 @@ pub trait Store: Send + Sync {
         tx: &Transaction,
         // Future improvement could hold sort fields, limits, etc
     ) -> Result<Vec<models::Trace>>;
+
+    /// Delete all spans with a specific trace_id.
+    async fn span_delete_by_trace(&self, tx: &Transaction, trace_id: &str) -> Result<Option<u64>>;
+
+    /// Delete a single span.
+    async fn span_delete(
+        &self,
+        tx: &Transaction,
+        trace_id: &str,
+        span_id: &str,
+    ) -> Result<Option<u64>>;
 }
diff --git a/fpx-lib/src/data/sql.rs b/fpx-lib/src/data/sql.rs
@@ -86,4 +86,21 @@ impl SqlBuilder {
             ",
         )
     }
+
+    /// Delete all spans with a specific trace_id.
+    ///
+    /// Query parameters:
+    /// - $1: trace_id
+    pub fn span_delete_by_trace(&self) -> String {
+        String::from("DELETE FROM spans WHERE trace_id=$1")
+    }
+
+    /// Delete a specific span.
+    ///
+    /// Query parameters:
+    /// - $1: trace_id
+    /// - $2: span_id
+    pub fn span_delete(&self) -> String {
+        String::from("DELETE FROM spans WHERE trace_id=$1 AND span_id=$2")
+    }
 }
diff --git a/fpx-workers/src/data.rs b/fpx-workers/src/data.rs
@@ -5,7 +5,7 @@ use serde::Deserialize;
 use std::sync::Arc;
 use wasm_bindgen::JsValue;
 use worker::send::SendFuture;
-use worker::D1Database;
+use worker::{D1Database, D1ResultMeta};
 
 pub struct D1Store {
     database: Arc<D1Database>,
@@ -24,17 +24,17 @@ impl D1Store {
     where
         T: for<'a> Deserialize<'a>,
     {
-        let prepared_statement = self.database.prepare(query);
-
-        let result = prepared_statement
+        let prepared_statement = self
+            .database
+            .prepare(query)
             .bind(values)
-            .map_err(|err| DbError::InternalError(err.to_string()))?; // TODO: Correct error
+            .map_err(|err| DbError::InternalError(err.to_string()))?; // TODO: Correct error;
 
-        let result = result
+        let result = prepared_statement
             .first(None)
             .await
-            .map_err(|err| DbError::InternalError(err.to_string()))?
-            .ok_or(DbError::NotFound)?; // TODO: Correct error
+            .map_err(|err| DbError::InternalError(err.to_string()))? // TODO: Correct error;
+            .ok_or(DbError::NotFound)?;
 
         Ok(result)
     }
@@ -43,13 +43,13 @@ impl D1Store {
     where
         T: for<'a> Deserialize<'a>,
     {
-        let prepared_statement = self.database.prepare(query);
-
-        let result = prepared_statement
+        let prepared_statement = self
+            .database
+            .prepare(query)
             .bind(values)
-            .map_err(|err| DbError::InternalError(err.to_string()))?; // TODO: Correct error
+            .map_err(|err| DbError::InternalError(err.to_string()))?; // TODO: Correct error;
 
-        let result = result
+        let result = prepared_statement
             .all()
             .await
             .map_err(|err| DbError::InternalError(err.to_string()))? // TODO: Correct error
@@ -156,4 +156,63 @@ impl Store for D1Store {
         })
         .await
     }
+
+    /// Delete all spans with a specific trace_id.
+    async fn span_delete_by_trace(&self, _tx: &Transaction, trace_id: &str) -> Result<Option<u64>> {
+        SendFuture::new(async {
+            let prepared_statement = self
+                .database
+                .prepare(self.sql_builder.span_delete_by_trace())
+                .bind(&[trace_id.into()])
+                .map_err(|err| DbError::InternalError(err.to_string()))?;
+
+            let results = prepared_statement
+                .run()
+                .await
+                .map_err(|err| DbError::InternalError(err.to_string()))?;
+
+            if let Ok(Some(D1ResultMeta {
+                rows_written: Some(rows_written),
+                ..
+            })) = results.meta()
+            {
+                Ok(Some(rows_written as u64))
+            } else {
+                Ok(None)
+            }
+        })
+        .await
+    }
+
+    /// Delete a single span.
+    async fn span_delete(
+        &self,
+        _tx: &Transaction,
+        trace_id: &str,
+        span_id: &str,
+    ) -> Result<Option<u64>> {
+        SendFuture::new(async {
+            let prepared_statement = self
+                .database
+                .prepare(self.sql_builder.span_delete())
+                .bind(&[trace_id.into(), span_id.into()])
+                .map_err(|err| DbError::InternalError(err.to_string()))?;
+
+            let results = prepared_statement
+                .run()
+                .await
+                .map_err(|err| DbError::InternalError(err.to_string()))?;
+
+            if let Ok(Some(D1ResultMeta {
+                rows_written: Some(rows_written),
+                ..
+            })) = results.meta()
+            {
+                Ok(Some(rows_written as u64))
+            } else {
+                Ok(None)
+            }
+        })
+        .await
+    }
 }