1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

use opentelemetry::trace::{SpanContext, TraceContextExt, TraceFlags};
use quaint::ast::{Delete, Insert, Select, Update};
use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;

pub fn trace_parent_to_string(context: &SpanContext) -> String {
    let trace_id = context.trace_id();
    let span_id = context.span_id();

    // see https://www.w3.org/TR/trace-context/#traceparent-header-field-values
    format!("traceparent='00-{trace_id:032x}-{span_id:032x}-01'")
}

pub trait SqlTraceComment: Sized {
    fn append_trace(self, span: &Span) -> Self;
    fn add_trace_id(self, trace_id: Option<&str>) -> Self;
}

macro_rules! sql_trace {
    ($what:ty) => {
        impl SqlTraceComment for $what {
            fn append_trace(self, span: &Span) -> Self {
                let otel_ctx = span.context();
                let span_ref = otel_ctx.span();
                let otel_ctx = span_ref.span_context();

                if otel_ctx.trace_flags() == TraceFlags::SAMPLED {
                    self.comment(trace_parent_to_string(otel_ctx))
                } else {
                    self
                }
            }
            // Temporary method to pass the traceid in an operation
            fn add_trace_id(self, trace_id: Option<&str>) -> Self {
                if let Some(traceparent) = trace_id {
                    if should_sample(&traceparent) {
                        self.comment(format!("traceparent='{}'", traceparent))
                    } else {
                        self
                    }
                } else {
                    self
                }
            }
        }
    };
}

fn should_sample(traceparent: &str) -> bool {
    traceparent.split('-').count() == 4 && traceparent.ends_with("-01")
}

sql_trace!(Insert<'_>);

sql_trace!(Update<'_>);

sql_trace!(Delete<'_>);

sql_trace!(Select<'_>);