Implement TypeScript API compatible AttributeMap (#199)

* Implement TypeScript API compatible AttributeMap * Slightly refactor fn's --------- Co-authored-by: Benno van den Berg <[email protected]>
fiberplane · Aug 28, 2024 · 0d6eddf · 0d6eddf
1 parent cfabf65
commit 0d6eddf
Show file tree

Hide file tree

Showing 5 changed files with 135 additions and 32 deletions.
diff --git a/fpx-lib/src/api/handlers/spans.rs b/fpx-lib/src/api/handlers/spans.rs
@@ -1,5 +1,5 @@
 use crate::api::errors::{ApiServerError, CommonError};
-use crate::api::models::{Span, TypeScriptCompatSpan};
+use crate::api::models::{ts_compat::TypeScriptCompatSpan, Span};
 use crate::data::{BoxedStore, DbError};
 use axum::extract::{Path, State};
 use axum::Json;

diff --git a/fpx-lib/src/api/handlers/traces.rs b/fpx-lib/src/api/handlers/traces.rs
@@ -1,5 +1,5 @@
 use crate::api::errors::{ApiServerError, CommonError};
-use crate::api::models::{TraceSummary, TypeScriptCompatTrace};
+use crate::api::models::{ts_compat::TypeScriptCompatTrace, TraceSummary};
 use crate::data::{BoxedStore, DbError};
 use axum::extract::{Path, State};
 use axum::Json;

diff --git a/fpx-lib/src/api/models.rs b/fpx-lib/src/api/models.rs
@@ -2,6 +2,7 @@ use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 mod otel;
+pub mod ts_compat;
 
 pub use otel::*;
 
@@ -140,32 +141,3 @@ impl From<SpanAdded> for ServerMessage {
         ServerMessageDetails::SpanAdded(Box::new(val)).into()
     }
 }
-
-#[derive(Serialize)]
-#[serde(rename_all = "camelCase")]
-pub struct TypeScriptCompatSpan {
-    pub span_id: Option<String>,
-    pub trace_id: Option<String>,
-    pub created_at: time::OffsetDateTime,
-    pub updated_at: time::OffsetDateTime,
-    pub parsed_payload: Span,
-}
-
-impl From<crate::data::models::Span> for TypeScriptCompatSpan {
-    fn from(span: crate::data::models::Span) -> Self {
-        Self {
-            span_id: Some(span.span_id.clone()),
-            trace_id: Some(span.trace_id.clone()),
-            created_at: span.end_time.into(),
-            updated_at: span.end_time.into(),
-            parsed_payload: span.clone().into(),
-        }
-    }
-}
-
-#[derive(Serialize)]
-#[serde(rename_all = "camelCase")]
-pub struct TypeScriptCompatTrace {
-    pub trace_id: String,
-    pub spans: Vec<TypeScriptCompatSpan>,
-}
diff --git a/fpx-lib/src/api/models/otel.rs b/fpx-lib/src/api/models/otel.rs
@@ -231,7 +231,7 @@ impl From<StatusCode> for SpanStatusCode {
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, Default, PartialEq)]
-pub struct AttributeMap(BTreeMap<String, Option<AttributeValue>>);
+pub struct AttributeMap(pub BTreeMap<String, Option<AttributeValue>>);
 
 impl From<KeyValueList> for AttributeMap {
     fn from(value: KeyValueList) -> Self {

diff --git a/fpx-lib/src/api/models/ts_compat.rs b/fpx-lib/src/api/models/ts_compat.rs
@@ -0,0 +1,131 @@
+use super::{AttributeMap, AttributeValue, SpanEvent, SpanKind};
+use opentelemetry_proto::tonic::trace::v1::span::Link;
+use opentelemetry_proto::tonic::trace::v1::Status;
+use serde::{Deserialize, Serialize};
+use std::collections::BTreeMap;
+
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TypeScriptCompatSpan {
+    pub span_id: Option<String>,
+    pub trace_id: Option<String>,
+    pub created_at: time::OffsetDateTime,
+    pub updated_at: time::OffsetDateTime,
+    pub parsed_payload: TypeScriptCompatOtelSpan,
+}
+
+impl From<crate::data::models::Span> for TypeScriptCompatSpan {
+    fn from(span: crate::data::models::Span) -> Self {
+        Self {
+            span_id: Some(span.span_id),
+            trace_id: Some(span.trace_id),
+            created_at: span.end_time.into(),
+            updated_at: span.end_time.into(),
+            parsed_payload: span.inner.0.into(),
+        }
+    }
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TypeScriptCompatTrace {
+    pub trace_id: String,
+    pub spans: Vec<TypeScriptCompatSpan>,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
+pub struct TypeScriptCompatOtelSpan {
+    pub trace_id: String,
+    pub span_id: String,
+    pub parent_span_id: Option<String>,
+
+    pub name: String,
+    pub trace_state: String,
+    pub flags: u32,
+    pub kind: SpanKind,
+
+    pub scope_name: Option<String>,
+    pub scope_version: Option<String>,
+
+    #[serde(with = "time::serde::rfc3339")]
+    pub start_time: time::OffsetDateTime,
+
+    #[serde(with = "time::serde::rfc3339")]
+    pub end_time: time::OffsetDateTime,
+
+    pub attributes: TypeScriptCompatAttributeMap,
+    pub scope_attributes: Option<TypeScriptCompatAttributeMap>,
+    pub resource_attributes: Option<TypeScriptCompatAttributeMap>,
+
+    pub status: Option<Status>,
+    pub events: Vec<SpanEvent>,
+    pub links: Vec<Link>,
+}
+
+impl From<crate::api::models::otel::Span> for TypeScriptCompatOtelSpan {
+    fn from(span: crate::api::models::otel::Span) -> Self {
+        Self {
+            trace_id: span.trace_id,
+            span_id: span.span_id,
+            parent_span_id: span.parent_span_id,
+            name: span.name,
+            trace_state: span.trace_state,
+            flags: span.flags,
+            kind: span.kind,
+            scope_name: span.scope_name,
+            scope_version: span.scope_version,
+            start_time: span.start_time,
+            end_time: span.end_time,
+            attributes: span.attributes.into(),
+            scope_attributes: span.scope_attributes.map(Into::into),
+            resource_attributes: span.resource_attributes.map(Into::into),
+            status: span.status,
+            events: span.events,
+            links: span.links,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, Default, PartialEq)]
+pub struct TypeScriptCompatAttributeMap(BTreeMap<String, Option<serde_json::Value>>);
+
+impl From<AttributeMap> for TypeScriptCompatAttributeMap {
+    fn from(attr_map: AttributeMap) -> Self {
+        let result = attr_map
+            .0
+            .into_iter()
+            .map(|(key, value)| (key, value.map(Into::into)))
+            .collect();
+
+        TypeScriptCompatAttributeMap(result)
+    }
+}
+
+impl From<crate::api::models::otel::AttributeMap> for serde_json::Map<String, serde_json::Value> {
+    fn from(attr_map: crate::api::models::otel::AttributeMap) -> Self {
+        attr_map
+            .0
+            .into_iter()
+            .map(|(key, value)| {
+                (
+                    key,
+                    value.map(Into::into).unwrap_or(serde_json::Value::Null),
+                )
+            })
+            .collect()
+    }
+}
+
+impl From<AttributeValue> for serde_json::Value {
+    fn from(value: AttributeValue) -> Self {
+        match value {
+            AttributeValue::String(value) => value.into(),
+            AttributeValue::Bool(value) => value.into(),
+            AttributeValue::Int(value) => value.into(),
+            AttributeValue::Double(value) => value.into(),
+            AttributeValue::Array(values) => values.into(),
+            AttributeValue::KeyValueList(value) => serde_json::Value::Object(value.into()),
+            AttributeValue::Bytes(value) => value.into(),
+        }
+    }
+}