google_cloud_lro/internal/
ext.rs

1// Copyright 2026 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::either::Either;
16use super::tracing::Tracing;
17use crate::Poller;
18
19/// Details for tracing a poller.
20#[derive(Clone, Debug, Default)]
21#[non_exhaustive]
22pub struct TracingDetails {
23    pub method_name: &'static str,
24}
25
26/// Options for creating a new poller.
27#[derive(Default)]
28#[non_exhaustive]
29pub struct PollerOptions {
30    pub tracing: Option<TracingDetails>,
31}
32
33pub trait PollerExt<ResponseType, MetadataType> {
34    fn with_options(self, options: PollerOptions) -> impl Poller<ResponseType, MetadataType>;
35}
36
37impl<ResponseType, MetadataType, T> PollerExt<ResponseType, MetadataType> for T
38where
39    T: Poller<ResponseType, MetadataType>,
40    ResponseType: Send,
41    MetadataType: Send,
42{
43    fn with_options(self, options: PollerOptions) -> impl Poller<ResponseType, MetadataType> {
44        if let Some(t) = options.tracing {
45            let method_name = if t.method_name.is_empty() {
46                "google_longrunning::Operations/Wait"
47            } else {
48                t.method_name
49            };
50            let span = tracing::info_span!(
51                "LRO Wait",
52                "otel.name" = method_name,
53                "gcp.rpc.method" = method_name,
54                "gcp.resource.destination.id" = tracing::field::Empty,
55                "otel.status_code" = tracing::field::Empty,
56                "otel.status_description" = tracing::field::Empty,
57                "error.type" = tracing::field::Empty
58            );
59            let traced = Tracing::new(self, span);
60            return Either::Right(traced);
61        }
62        Either::Left(self)
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69    use crate::{PollingResult, sealed};
70    use google_cloud_gax::polling_state::PollingState;
71    use google_cloud_wkt::{Duration, Timestamp};
72    use mockall::mock;
73
74    type ResponseType = Duration;
75    type MetadataType = Timestamp;
76
77    mock! {
78        PollerA {}
79        impl sealed::Poller for PollerA {
80            async fn backoff(&mut self, state: &PollingState);
81        }
82        impl Poller<ResponseType, MetadataType> for PollerA {
83            async fn poll(&mut self) -> Option<PollingResult<ResponseType, MetadataType>>;
84            async fn until_done(self) -> google_cloud_gax::Result<ResponseType>;
85            #[cfg(feature = "unstable-stream")]
86            fn into_stream(
87                self,
88            ) -> impl futures::Stream<Item = PollingResult<ResponseType, MetadataType>> + Unpin;
89        }
90    }
91
92    #[test]
93    fn test_poller_initialization_with_tracing() {
94        let mock = MockPollerA::new();
95        let _poller = mock.with_options(PollerOptions {
96            tracing: Some(TracingDetails {
97                method_name: "test_method",
98            }),
99        });
100    }
101}
google_cloud_lro/internal/ext.rs

google_cloud_lro/internal/
ext.rs
