Add CaptureSettings

prisma · Nov 22, 2024 · a367cfb · a367cfb
1 parent 3453dd6
commit a367cfb
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 5 deletions.
diff --git a/libs/telemetry/src/capturing/ng/exporter.rs b/libs/telemetry/src/capturing/ng/exporter.rs
@@ -1,5 +1,6 @@
-use std::{borrow::Cow, collections::HashMap, sync::Arc};
+use std::{borrow::Cow, collections::HashMap, str::FromStr, sync::Arc};
 
+use enumflags2::{bitflags, BitFlags};
 use serde::Serialize;
 use tokio::sync::{mpsc, oneshot, Mutex, RwLock};
 
@@ -65,6 +66,47 @@ pub struct Trace {
     pub events: Vec<ExportedEvent>,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[bitflags]
+#[repr(u8)]
+pub enum CaptureTarget {
+    TraceEvents,
+    DebugEvents,
+    InfoEvents,
+    WarnEvents,
+    ErrorEvents,
+    QueryEvents,
+    Spans,
+}
+
+impl From<LogLevel> for CaptureTarget {
+    fn from(value: LogLevel) -> Self {
+        match value {
+            LogLevel::Trace => Self::TraceEvents,
+            LogLevel::Debug => Self::DebugEvents,
+            LogLevel::Info => Self::InfoEvents,
+            LogLevel::Warn => Self::WarnEvents,
+            LogLevel::Error => Self::ErrorEvents,
+            LogLevel::Query => Self::QueryEvents,
+        }
+    }
+}
+
+impl FromStr for CaptureTarget {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "tracing" => Ok(Self::Spans),
+            _ => Ok(s.parse::<LogLevel>()?.into()),
+        }
+    }
+}
+
+pub struct CaptureSettings {
+    targets: BitFlags<CaptureTarget>,
+}
+
 #[derive(Clone)]
 pub struct Exporter(Arc<Inner>);
 

diff --git a/libs/telemetry/src/models.rs b/libs/telemetry/src/models.rs
@@ -1,13 +1,15 @@
-use opentelemetry::{sdk::export::trace::SpanData, KeyValue, Value};
-use serde::Serialize;
-use serde_json::json;
 use std::{
     borrow::Cow,
     collections::HashMap,
     str::FromStr,
     time::{Duration, SystemTime},
 };
 
+use enumflags2::bitflags;
+use opentelemetry::{sdk::export::trace::SpanData, KeyValue, Value};
+use serde::Serialize;
+use serde_json::json;
+
 const ACCEPT_ATTRIBUTES: &[&str] = &[
     "db.system",
     "db.statement",
@@ -16,7 +18,9 @@ const ACCEPT_ATTRIBUTES: &[&str] = &[
     "itx_id",
 ];
 
-#[derive(Serialize, Debug, Clone, Copy)]
+#[derive(Serialize, Debug, Clone, Copy, PartialEq, Eq)]
+#[bitflags]
+#[repr(u8)]
 pub enum LogLevel {
     Trace,
     Debug,