pgtask-otel 0.1.3

OpenTelemetry conventions and propagation for pgtask
Documentation
use opentelemetry::{
    Context, global,
    propagation::{Extractor, Injector, TextMapCompositePropagator},
};
use opentelemetry_sdk::propagation::{BaggagePropagator, TraceContextPropagator};
use serde_json::{Map, Value};
use tracing::Span;
use tracing_opentelemetry::{OpenTelemetrySpanExt, SetParentError};

pub fn configure_propagation() {
    global::set_text_map_propagator(TextMapCompositePropagator::new(vec![
        Box::new(TraceContextPropagator::new()),
        Box::new(BaggagePropagator::new()),
    ]));
}

pub fn inject_context(headers: &Map<String, Value>, context: &Context) -> Map<String, Value> {
    let mut headers = headers.clone();
    global::get_text_map_propagator(|propagator| propagator.inject_context(context, &mut HeaderInjector(&mut headers)));
    headers
}

pub fn inject_span_context(headers: &Map<String, Value>, span: &Span) -> Map<String, Value> {
    inject_context(headers, &span.context())
}

pub fn set_parent_from_headers(span: &Span, headers: &Map<String, Value>) -> Result<(), SetParentError> {
    let context = global::get_text_map_propagator(|propagator| propagator.extract(&HeaderExtractor(headers)));
    span.set_parent(context)
}

struct HeaderInjector<'a>(&'a mut Map<String, Value>);

impl Injector for HeaderInjector<'_> {
    fn set(&mut self, key: &str, value: String) {
        self.0.insert(key.to_owned(), Value::String(value));
    }
}

struct HeaderExtractor<'a>(&'a Map<String, Value>);

impl Extractor for HeaderExtractor<'_> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).and_then(Value::as_str)
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(String::as_str).collect()
    }
}