telemetry_rust/middleware/aws/
instrumentation.rs

1use aws_types::request_id::RequestId;
2use std::{error::Error, future::Future};
3
4use super::{AwsSpan, AwsSpanBuilder};
5use crate::future::{InstrumentedFuture, InstrumentedFutureContext};
6
7impl<T, E> InstrumentedFutureContext<Result<T, E>> for AwsSpan
8where
9    T: RequestId,
10    E: RequestId + Error,
11{
12    fn on_result(self, result: &Result<T, E>) {
13        self.end(result);
14    }
15}
16
17/// Trait for instrumenting AWS futures with automatic span management.
18///
19/// This trait provides a convenient way to wrap AWS SDK futures with OpenTelemetry
20/// instrumentation, automatically handling span creation, error recording, and cleanup.
21///
22/// # Example
23///
24/// ```rust
25/// use aws_sdk_dynamodb::{Client as DynamoClient, types::AttributeValue};
26/// use telemetry_rust::middleware::aws::{AwsInstrument, DynamodbSpanBuilder};
27///
28/// async fn query_table() {
29///     let config = aws_config::load_from_env().await;
30///     let dynamo_client = DynamoClient::new(&config);
31///     let res = dynamo_client
32///         .query()
33///         .table_name("table_name")
34///         .index_name("my_index")
35///         .key_condition_expression("PK = :pk")
36///         .expression_attribute_values(":pk", AttributeValue::S("Test".to_string()))
37///         .send()
38///         .instrument(DynamodbSpanBuilder::get_item("table_name"))
39///         .await;
40///     println!("DynamoDB response: {res:#?}");
41/// }
42/// ```
43pub trait AwsInstrument<T, E, F>
44where
45    T: RequestId,
46    E: RequestId + Error,
47    F: Future<Output = Result<T, E>>,
48{
49    /// Instruments the future with an AWS span.
50    ///
51    /// Creates an instrumented future that will automatically start a span when polled
52    /// and properly handle the result when the future completes.
53    ///
54    /// # Arguments
55    ///
56    /// * `span` - A span builder or span to use for instrumentation
57    ///
58    /// # Returns
59    ///
60    /// An instrumented future that will record AWS operation details
61    fn instrument<'a>(
62        self,
63        span: impl Into<AwsSpanBuilder<'a>>,
64    ) -> InstrumentedFuture<F, AwsSpan>;
65}
66
67impl<T, E, F> AwsInstrument<T, E, F> for F
68where
69    T: RequestId,
70    E: RequestId + Error,
71    F: Future<Output = Result<T, E>>,
72{
73    fn instrument<'a>(
74        self,
75        span: impl Into<AwsSpanBuilder<'a>>,
76    ) -> InstrumentedFuture<F, AwsSpan> {
77        let span = span.into().start();
78        InstrumentedFuture::new(self, span)
79    }
80}
telemetry_rust/middleware/aws/instrumentation.rs

telemetry_rust/middleware/aws/
instrumentation.rs
