fixes from review

prisma · Aug 15, 2022 · e66186a · e66186a
1 parent a667bc6
commit e66186a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 14 deletions.
diff --git a/query-engine/query-engine/src/capture_tracer.rs b/query-engine/query-engine/src/capture_tracer.rs
@@ -1,12 +1,12 @@
 use async_trait::async_trait;
 use opentelemetry::{
     global,
-    sdk::{self},
     sdk::{
+        self,
         export::trace::{ExportResult, SpanData, SpanExporter},
         propagation::TraceContextPropagator,
     },
-    trace::TracerProvider,
+    trace::{TraceId, TracerProvider},
 };
 use query_core::spans_to_json;
 use std::fmt::Debug;
@@ -58,7 +58,7 @@ impl PipelineBuilder {
 /// A [`CaptureExporter`] that sends spans to stdout.
 #[derive(Debug, Clone)]
 pub struct CaptureExporter {
-    logs: Arc<Mutex<HashMap<String, Vec<SpanData>>>>,
+    logs: Arc<Mutex<HashMap<TraceId, Vec<SpanData>>>>,
 }
 
 impl CaptureExporter {
@@ -68,12 +68,12 @@ impl CaptureExporter {
         }
     }
 
-    pub async fn capture(&self, trace_id: String) {
+    pub async fn capture(&self, trace_id: TraceId) {
         let mut logs = self.logs.lock().await;
         logs.insert(trace_id, Vec::new());
     }
 
-    pub async fn get(&self, trace_id: String) -> String {
+    pub async fn get(&self, trace_id: TraceId) -> String {
         let mut logs = self.logs.lock().await;
         if let Some(spans) = logs.remove(&trace_id) {
             spans_to_json(&spans)
@@ -96,7 +96,7 @@ impl SpanExporter for CaptureExporter {
 
         let mut logs = self.logs.lock().await;
         for span in batch {
-            let trace_id = span.span_context.trace_id().to_string();
+            let trace_id = span.span_context.trace_id();
 
             if let Some(spans) = logs.get_mut(&trace_id) {
                 spans.push(span)

diff --git a/query-engine/query-engine/src/server/mod.rs b/query-engine/query-engine/src/server/mod.rs
@@ -1,6 +1,7 @@
 use crate::{opt::PrismaOpt, state::State, PrismaResult};
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{header::CONTENT_TYPE, Body, HeaderMap, Method, Request, Response, Server, StatusCode};
+use opentelemetry::trace::TraceId;
 use opentelemetry::{global, propagation::Extractor, Context};
 use query_core::{get_trace_id_from_context, MetricFormat};
 use query_core::{schema::QuerySchemaRenderer, TxId};
@@ -387,30 +388,30 @@ fn err_to_http_resp(err: query_core::CoreError) -> Response<Body> {
 }
 
 struct LogCapture {
-    id: String,
+    id: TraceId,
     capture: bool,
 }
 
 impl LogCapture {
-    fn new(id: String, capture: bool) -> Self {
+    fn new(id: TraceId, capture: bool) -> Self {
         Self { id, capture }
     }
 
-    fn new_from_req(id: String, req: &Request<Body>) -> Self {
+    fn new_from_req(id: TraceId, req: &Request<Body>) -> Self {
         let should_capture = Self::get_capture_from_header(req);
         Self::new(id, should_capture)
     }
 
-    fn id(&self) -> String {
-        self.id.to_string()
+    fn id(&self) -> TraceId {
+        self.id
     }
 
     fn should_capture(&self) -> bool {
         self.capture
     }
 
     fn get_capture_from_header(req: &Request<Body>) -> bool {
-        match req.headers().get("CAPTURE_LOGS") {
+        match req.headers().get("PRISMA-CAPTURE-LOGS") {
             Some(header) => {
                 if let Ok(capture_logs) = header.to_str() {
                     capture_logs == "true"
@@ -423,6 +424,15 @@ impl LogCapture {
     }
 }
 
+impl Default for LogCapture {
+    fn default() -> Self {
+        Self {
+            id: TraceId::from_hex("0").unwrap(),
+            capture: false,
+        }
+    }
+}
+
 fn process_gql_req_headers(req: &Request<Body>) -> (Option<TxId>, Span, LogCapture, Option<String>) {
     let tx_id = get_transaction_id_from_header(req);
     let (span, log_capture) = if tx_id.is_none() {
@@ -432,11 +442,11 @@ fn process_gql_req_headers(req: &Request<Body>) -> (Option<TxId>, Span, LogCaptu
         let span = info_span!("prisma:engine", user_facing = true);
         span.set_parent(cx);
 
-        let log_capture = LogCapture::new_from_req(trace_id.to_string(), req);
+        let log_capture = LogCapture::new_from_req(trace_id, req);
 
         (span, log_capture)
     } else {
-        (Span::none(), LogCapture::new(Default::default(), false))
+        (Span::none(), LogCapture::default())
     };
 
     let traceparent = match req.headers().get("traceparent") {