uhg_custom_appollo_roouter/plugins/telemetry/config_new/
mod.rs

1use events::EventOn;
2use opentelemetry::KeyValue;
3use opentelemetry::baggage::BaggageExt;
4use opentelemetry::trace::TraceContextExt;
5use opentelemetry::trace::TraceId;
6use opentelemetry_api::Value;
7use paste::paste;
8use tower::BoxError;
9use tracing::Span;
10
11use super::otel::OpenTelemetrySpanExt;
12use super::otlp::TelemetryDataKind;
13use crate::Context;
14use crate::plugins::telemetry::config::AttributeValue;
15use crate::plugins::telemetry::config_new::attributes::DefaultAttributeRequirementLevel;
16
17/// These modules contain a new config structure for telemetry that will progressively move to
18pub(crate) mod attributes;
19pub(crate) mod conditions;
20
21pub(crate) mod cache;
22mod conditional;
23pub(crate) mod cost;
24pub(crate) mod events;
25mod experimental_when_header;
26pub(crate) mod extendable;
27pub(crate) mod graphql;
28pub(crate) mod instruments;
29pub(crate) mod logging;
30pub(crate) mod selectors;
31pub(crate) mod spans;
32
33pub(crate) trait Selectors<Request, Response, EventResponse> {
34    fn on_request(&self, request: &Request) -> Vec<KeyValue>;
35    fn on_response(&self, response: &Response) -> Vec<KeyValue>;
36    fn on_response_event(&self, _response: &EventResponse, _ctx: &Context) -> Vec<KeyValue> {
37        Vec::with_capacity(0)
38    }
39    fn on_error(&self, error: &BoxError, ctx: &Context) -> Vec<KeyValue>;
40    fn on_response_field(
41        &self,
42        _attrs: &mut Vec<KeyValue>,
43        _ty: &apollo_compiler::executable::NamedType,
44        _field: &apollo_compiler::executable::Field,
45        _value: &serde_json_bytes::Value,
46        _ctx: &Context,
47    ) {
48    }
49}
50
51#[allow(dead_code)]
52#[derive(Clone, Copy, Debug, PartialEq)]
53pub(crate) enum Stage {
54    Request,
55    Response,
56    ResponseEvent,
57    ResponseField,
58    Error,
59    Drop,
60}
61
62impl std::fmt::Display for Stage {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        match self {
65            Stage::Request => write!(f, "request"),
66            Stage::Response => write!(f, "response"),
67            Stage::ResponseEvent => write!(f, "response_event"),
68            Stage::ResponseField => write!(f, "response_field"),
69            Stage::Error => write!(f, "error"),
70            Stage::Drop => write!(f, "drop"),
71        }
72    }
73}
74
75impl From<EventOn> for Stage {
76    fn from(value: EventOn) -> Self {
77        match value {
78            EventOn::Request => Self::Request,
79            EventOn::Response => Self::Response,
80            EventOn::EventResponse => Self::ResponseEvent,
81            EventOn::Error => Self::Error,
82        }
83    }
84}
85
86pub(crate) trait Selector: std::fmt::Debug {
87    type Request;
88    type Response;
89    type EventResponse;
90
91    fn on_request(&self, request: &Self::Request) -> Option<opentelemetry::Value>;
92    fn on_response(&self, response: &Self::Response) -> Option<opentelemetry::Value>;
93    fn on_response_event(
94        &self,
95        _response: &Self::EventResponse,
96        _ctx: &Context,
97    ) -> Option<opentelemetry::Value> {
98        None
99    }
100    fn on_error(&self, error: &BoxError, ctx: &Context) -> Option<opentelemetry::Value>;
101    fn on_response_field(
102        &self,
103        _ty: &apollo_compiler::executable::NamedType,
104        _field: &apollo_compiler::executable::Field,
105        _value: &serde_json_bytes::Value,
106        _ctx: &Context,
107    ) -> Option<opentelemetry::Value> {
108        None
109    }
110
111    fn on_drop(&self) -> Option<Value> {
112        None
113    }
114
115    fn is_active(&self, stage: Stage) -> bool;
116}
117
118pub(crate) trait DefaultForLevel {
119    /// Don't call this directly, use `defaults_for_levels` instead.
120    fn defaults_for_level(
121        &mut self,
122        requirement_level: DefaultAttributeRequirementLevel,
123        kind: TelemetryDataKind,
124    );
125    fn defaults_for_levels(
126        &mut self,
127        requirement_level: DefaultAttributeRequirementLevel,
128        kind: TelemetryDataKind,
129    ) {
130        match requirement_level {
131            DefaultAttributeRequirementLevel::None => {}
132            DefaultAttributeRequirementLevel::Required => {
133                self.defaults_for_level(DefaultAttributeRequirementLevel::Required, kind)
134            }
135            DefaultAttributeRequirementLevel::Recommended => {
136                self.defaults_for_level(DefaultAttributeRequirementLevel::Required, kind);
137                self.defaults_for_level(DefaultAttributeRequirementLevel::Recommended, kind);
138            }
139        }
140    }
141}
142
143pub(crate) trait DatadogId {
144    fn to_datadog(&self) -> String;
145}
146impl DatadogId for TraceId {
147    fn to_datadog(&self) -> String {
148        let bytes = &self.to_bytes()[std::mem::size_of::<u64>()..std::mem::size_of::<u128>()];
149        u64::from_be_bytes(bytes.try_into().unwrap()).to_string()
150    }
151}
152
153pub(crate) fn trace_id() -> Option<TraceId> {
154    let context = Span::current().context();
155    let span = context.span();
156    let span_context = span.span_context();
157    if span_context.is_valid() {
158        Some(span_context.trace_id())
159    } else {
160        crate::tracer::TraceId::current().map(|trace_id| TraceId::from(trace_id.to_u128()))
161    }
162}
163
164pub(crate) fn get_baggage(key: &str) -> Option<opentelemetry::Value> {
165    let context = Span::current().context();
166    let baggage = context.baggage();
167    baggage.get(key.to_string()).cloned()
168}
169
170pub(crate) trait ToOtelValue {
171    fn maybe_to_otel_value(&self) -> Option<opentelemetry::Value>;
172}
173impl ToOtelValue for &Option<AttributeValue> {
174    fn maybe_to_otel_value(&self) -> Option<opentelemetry::Value> {
175        self.as_ref().map(|v| v.clone().into())
176    }
177}
178
179macro_rules! impl_to_otel_value {
180    ($type:ty) => {
181        paste! {
182            impl ToOtelValue for $type {
183                fn maybe_to_otel_value(&self) -> Option<opentelemetry::Value> {
184                    match self {
185                        $type::Bool(value) => Some((*value).into()),
186                        $type::Number(value) if value.is_f64() => {
187                            value.as_f64().map(opentelemetry::Value::from)
188                        }
189                        $type::Number(value) if value.is_i64() => {
190                            value.as_i64().map(opentelemetry::Value::from)
191                        }
192                        $type::String(value) => Some(value.as_str().to_string().into()),
193                        $type::Array(value) => {
194                            // Arrays must be uniform in value
195                            if value.iter().all(|v| v.is_i64()) {
196                                Some(opentelemetry::Value::Array(opentelemetry::Array::I64(
197                                    value.iter().filter_map(|v| v.as_i64()).collect(),
198                                )))
199                            } else if value.iter().all(|v| v.is_f64()) {
200                                Some(opentelemetry::Value::Array(opentelemetry::Array::F64(
201                                    value.iter().filter_map(|v| v.as_f64()).collect(),
202                                )))
203                            } else if value.iter().all(|v| v.is_boolean()) {
204                                Some(opentelemetry::Value::Array(opentelemetry::Array::Bool(
205                                    value.iter().filter_map(|v| v.as_bool()).collect(),
206                                )))
207                            } else if value.iter().all(|v| v.is_object()) {
208                                Some(opentelemetry::Value::Array(opentelemetry::Array::String(
209                                    value.iter().map(|v| v.to_string().into()).collect(),
210                                )))
211                            } else if value.iter().all(|v| v.is_string()) {
212                                Some(opentelemetry::Value::Array(opentelemetry::Array::String(
213                                    value
214                                        .iter()
215                                        .filter_map(|v| v.as_str())
216                                        .map(|v| v.to_string().into())
217                                        .collect(),
218                                )))
219                            } else {
220                                Some(serde_json::to_string(value).ok()?.into())
221                            }
222                        }
223                        $type::Object(value) => Some(serde_json::to_string(value).ok()?.into()),
224                        _ => None
225                    }
226                }
227            }
228        }
229    };
230}
231impl_to_otel_value!(serde_json_bytes::Value);
232impl_to_otel_value!(serde_json::Value);
233
234impl From<opentelemetry::Value> for AttributeValue {
235    fn from(value: opentelemetry::Value) -> Self {
236        match value {
237            opentelemetry::Value::Bool(v) => AttributeValue::Bool(v),
238            opentelemetry::Value::I64(v) => AttributeValue::I64(v),
239            opentelemetry::Value::F64(v) => AttributeValue::F64(v),
240            opentelemetry::Value::String(v) => AttributeValue::String(v.into()),
241            opentelemetry::Value::Array(v) => AttributeValue::Array(v.into()),
242        }
243    }
244}
245
246#[cfg(test)]
247mod test {
248    use std::sync::OnceLock;
249
250    use apollo_compiler::Node;
251    use apollo_compiler::ast::FieldDefinition;
252    use apollo_compiler::ast::NamedType;
253    use apollo_compiler::executable::Field;
254    use apollo_compiler::name;
255    use opentelemetry::Context;
256    use opentelemetry::StringValue;
257    use opentelemetry::trace::SpanContext;
258    use opentelemetry::trace::SpanId;
259    use opentelemetry::trace::TraceContextExt;
260    use opentelemetry::trace::TraceFlags;
261    use opentelemetry::trace::TraceId;
262    use opentelemetry::trace::TraceState;
263    use serde_json::json;
264    use tracing::span;
265    use tracing_subscriber::layer::SubscriberExt;
266
267    use crate::plugins::telemetry::config_new::DatadogId;
268    use crate::plugins::telemetry::config_new::ToOtelValue;
269    use crate::plugins::telemetry::config_new::trace_id;
270    use crate::plugins::telemetry::otel;
271
272    pub(crate) fn field() -> &'static Field {
273        static FIELD: OnceLock<Field> = OnceLock::new();
274        FIELD.get_or_init(|| {
275            Field::new(
276                name!("field_name"),
277                Node::new(FieldDefinition {
278                    description: None,
279                    name: name!("field_name"),
280                    arguments: vec![],
281                    ty: apollo_compiler::ty!(field_type),
282                    directives: Default::default(),
283                }),
284            )
285        })
286    }
287    pub(crate) fn ty() -> NamedType {
288        name!("type_name")
289    }
290
291    #[test]
292    fn dd_convert() {
293        let trace_id = TraceId::from_hex("234e10d9e749a0a19e94ac0e4a896aee").unwrap();
294        let dd_id = trace_id.to_datadog();
295        assert_eq!(dd_id, "11426947331925830382");
296    }
297
298    #[test]
299    fn test_trace_id() {
300        // Create a span with a trace ID
301        let subscriber = tracing_subscriber::registry().with(otel::layer());
302        tracing::subscriber::with_default(subscriber, || {
303            let span_context = SpanContext::new(
304                TraceId::from(42),
305                SpanId::from_u64(42),
306                TraceFlags::default(),
307                false,
308                TraceState::default(),
309            );
310            let _context = Context::current()
311                .with_remote_span_context(span_context)
312                .attach();
313            let span = span!(tracing::Level::INFO, "test");
314            let _guard = span.enter();
315            let trace_id = trace_id();
316            assert_eq!(trace_id, Some(TraceId::from_u128(42)));
317        });
318    }
319
320    #[test]
321    fn test_baggage() {
322        // Create a span with a trace ID
323        let subscriber = tracing_subscriber::registry().with(otel::layer());
324        tracing::subscriber::with_default(subscriber, || {
325            let span_context = SpanContext::new(
326                TraceId::from_u128(42),
327                SpanId::from_u64(42),
328                TraceFlags::default(),
329                false,
330                TraceState::default(),
331            );
332            let _context = Context::current()
333                .with_remote_span_context(span_context)
334                .attach();
335            let span = span!(tracing::Level::INFO, "test");
336            let _guard = span.enter();
337            let trace_id = trace_id();
338            assert_eq!(trace_id, Some(TraceId::from_u128(42)));
339        });
340    }
341
342    #[test]
343    fn maybe_to_otel_value() {
344        assert_eq!(json!("string").maybe_to_otel_value(), Some("string".into()));
345        assert_eq!(json!(1).maybe_to_otel_value(), Some(1.into()));
346        assert_eq!(json!(1.0).maybe_to_otel_value(), Some(1.0.into()));
347        assert_eq!(json!(true).maybe_to_otel_value(), Some(true.into()));
348
349        assert_eq!(
350            json!(["string1", "string2"]).maybe_to_otel_value(),
351            Some(opentelemetry::Value::Array(
352                vec![
353                    StringValue::from("string1".to_string()),
354                    StringValue::from("string2".to_string())
355                ]
356                .into()
357            ))
358        );
359        assert_eq!(
360            json!([1, 2]).maybe_to_otel_value(),
361            Some(opentelemetry::Value::Array(vec![1i64, 2i64].into()))
362        );
363        assert_eq!(
364            json!([1.0, 2.0]).maybe_to_otel_value(),
365            Some(opentelemetry::Value::Array(vec![1.0, 2.0].into()))
366        );
367        assert_eq!(
368            json!([true, false]).maybe_to_otel_value(),
369            Some(opentelemetry::Value::Array(vec![true, false].into()))
370        );
371
372        // Arrays must be uniform
373        assert_eq!(
374            json!(["1", 1]).maybe_to_otel_value(),
375            Some(r#"["1",1]"#.to_string().into())
376        );
377        assert_eq!(
378            json!([1.0, 1]).maybe_to_otel_value(),
379            Some(r#"[1.0,1]"#.to_string().into())
380        );
381    }
382}
uhg_custom_appollo_roouter/plugins/telemetry/config_new/mod.rs

uhg_custom_appollo_roouter/plugins/telemetry/config_new/
mod.rs
