diff --git a/query-engine/query-engine/src/capture_tracer.rs b/query-engine/query-engine/src/capture_tracer.rs index d91b6c61a1f..964c4599e9d 100644 --- 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>>>, + logs: Arc>>>, } 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 index 4e2a5e35534..24ec898e2a0 100644 --- 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,22 +388,22 @@ fn err_to_http_resp(err: query_core::CoreError) -> Response { } 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) -> Self { + fn new_from_req(id: TraceId, req: &Request) -> 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 { @@ -410,7 +411,7 @@ impl LogCapture { } fn get_capture_from_header(req: &Request) -> 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) -> (Option, Span, LogCapture, Option) { 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) -> (Option, 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") {