dynamo_runtime/protocols/
annotated.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use super::maybe_error::MaybeError;
5use crate::error::DynamoError;
6use anyhow::Result;
7use serde::{Deserialize, Serialize};
8
9pub trait AnnotationsProvider {
10    fn annotations(&self) -> Option<Vec<String>>;
11    fn has_annotation(&self, annotation: &str) -> bool {
12        self.annotations()
13            .map(|annotations| annotations.iter().any(|a| a == annotation))
14            .unwrap_or(false)
15    }
16}
17
18/// Our services have the option of returning an "annotated" stream, which allows use
19/// to include additional information with each delta. This is useful for debugging,
20/// performance benchmarking, and improved observability.
21#[derive(Serialize, Deserialize, Clone, Debug)]
22pub struct Annotated<R> {
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub data: Option<R>,
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub id: Option<String>,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub event: Option<String>,
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub comment: Option<Vec<String>>,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub error: Option<DynamoError>,
33}
34
35impl<R> Annotated<R> {
36    fn cloned_error(&self) -> Option<DynamoError> {
37        self.is_error().then(|| {
38            self.error.clone().unwrap_or_else(|| {
39                DynamoError::msg(
40                    self.comment
41                        .as_ref()
42                        .filter(|comments| !comments.is_empty())
43                        .map(|comments| comments.join(", "))
44                        .unwrap_or_else(|| "unknown error".to_string()),
45                )
46            })
47        })
48    }
49
50    /// Create a new annotated stream from the given error string
51    pub fn from_error(error: impl Into<String>) -> Self {
52        Self {
53            data: None,
54            id: None,
55            event: Some("error".to_string()),
56            comment: None,
57            error: Some(DynamoError::msg(error)),
58        }
59    }
60
61    /// Create a new annotated stream from the given data
62    pub fn from_data(data: R) -> Self {
63        Self {
64            data: Some(data),
65            id: None,
66            event: None,
67            comment: None,
68            error: None,
69        }
70    }
71
72    /// Add an annotation to the stream
73    ///
74    /// Annotations populate the `event` field and the `comment` field
75    pub fn from_annotation<S: Serialize>(
76        name: impl Into<String>,
77        value: &S,
78    ) -> Result<Self, serde_json::Error> {
79        Ok(Self {
80            data: None,
81            id: None,
82            event: Some(name.into()),
83            comment: Some(vec![serde_json::to_string(value)?]),
84            error: None,
85        })
86    }
87
88    /// Convert to a [`Result<Self, String>`]
89    /// If [`Self::event`] is "error", return an error message
90    pub fn ok(self) -> Result<Self, String> {
91        if let Some(error) = self.cloned_error() {
92            return Err(error.to_string());
93        }
94        Ok(self)
95    }
96
97    /// Consume the annotation and return its optional payload while preserving
98    /// structured errors.
99    pub fn into_data(self) -> Result<Option<R>, DynamoError> {
100        if let Some(error) = self.cloned_error() {
101            return Err(error);
102        }
103        Ok(self.data)
104    }
105
106    pub fn is_ok(&self) -> bool {
107        self.event.as_deref() != Some("error")
108    }
109
110    pub fn is_event(&self) -> bool {
111        self.event.is_some()
112    }
113
114    pub fn transfer<U: Serialize>(self, data: Option<U>) -> Annotated<U> {
115        Annotated::<U> {
116            data,
117            id: self.id,
118            event: self.event,
119            comment: self.comment,
120            error: self.error,
121        }
122    }
123
124    /// Apply a mapping/transformation to the data field
125    /// If the mapping fails, the error is returned as an annotated stream
126    pub fn map_data<U, F>(self, transform: F) -> Annotated<U>
127    where
128        F: FnOnce(R) -> Result<U, String>,
129    {
130        match self.data.map(transform).transpose() {
131            Ok(data) => Annotated::<U> {
132                data,
133                id: self.id,
134                event: self.event,
135                comment: self.comment,
136                error: self.error,
137            },
138            Err(e) => Annotated::from_error(e),
139        }
140    }
141
142    pub fn is_error(&self) -> bool {
143        self.event.as_deref() == Some("error")
144    }
145
146    pub fn into_result(self) -> Result<Option<R>> {
147        self.into_data().map_err(anyhow::Error::new)
148    }
149}
150
151impl<R> MaybeError for Annotated<R>
152where
153    R: for<'de> Deserialize<'de>,
154{
155    fn from_err(err: impl std::error::Error + 'static) -> Self {
156        Self {
157            data: None,
158            id: None,
159            event: Some("error".to_string()),
160            comment: None,
161            error: Some(DynamoError::from(
162                Box::new(err) as Box<dyn std::error::Error + 'static>
163            )),
164        }
165    }
166
167    fn err(&self) -> Option<DynamoError> {
168        self.cloned_error()
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    #[test]
177    fn test_maybe_error() {
178        let annotated = Annotated::from_data("Test data".to_string());
179        assert!(annotated.err().is_none());
180        assert!(annotated.is_ok());
181
182        let annotated = Annotated::<String>::from_error("Test error 2".to_string());
183        assert!(annotated.err().is_some());
184        assert!(annotated.is_err());
185
186        let dynamo_err = DynamoError::msg("Test error 3");
187        let annotated = Annotated::<String>::from_err(dynamo_err);
188        assert!(annotated.is_err());
189    }
190
191    #[test]
192    fn test_from_err() {
193        let err = DynamoError::msg("connection lost");
194        let annotated = Annotated::<String>::from_err(err);
195
196        assert!(annotated.is_err());
197        let err = annotated.err().unwrap();
198        assert!(err.to_string().contains("connection lost"));
199    }
200
201    #[test]
202    fn test_comment_only_error_fallback() {
203        for (comments, expected) in [
204            (vec![], "unknown error"),
205            (
206                vec!["first".to_string(), "second".to_string()],
207                "first, second",
208            ),
209        ] {
210            let annotated = Annotated::<String> {
211                data: None,
212                id: None,
213                event: Some("error".to_string()),
214                comment: Some(comments),
215                error: None,
216            };
217
218            assert_eq!(annotated.err().unwrap().message(), expected);
219        }
220    }
221
222    #[test]
223    fn test_error_serialization() {
224        let err = DynamoError::msg("test error");
225        let annotated = Annotated::<String>::from_err(err);
226
227        // Serialize and deserialize
228        let json = serde_json::to_string(&annotated).unwrap();
229        let deserialized: Annotated<String> = serde_json::from_str(&json).unwrap();
230
231        assert!(deserialized.is_err());
232        assert!(
233            deserialized
234                .err()
235                .unwrap()
236                .to_string()
237                .contains("test error")
238        );
239    }
240
241    #[test]
242    fn test_transfer_preserves_error() {
243        let err = DynamoError::msg("request timed out");
244        let annotated = Annotated::<String>::from_err(err);
245
246        let transferred: Annotated<i32> = annotated.transfer(None);
247        assert!(transferred.err().is_some());
248    }
249
250    #[test]
251    fn test_ok_method() {
252        let err = DynamoError::msg("connection lost");
253        let annotated = Annotated::<String>::from_err(err);
254
255        let result = annotated.ok();
256        assert!(result.is_err());
257        assert!(result.unwrap_err().contains("connection lost"));
258    }
259
260    #[test]
261    fn test_into_data_preserves_error_type() {
262        use crate::error::{BackendError, ErrorType};
263
264        let error = DynamoError::builder()
265            .error_type(ErrorType::Backend(BackendError::InvalidArgument))
266            .message("invalid request")
267            .build();
268        let annotated = Annotated::<String>::from_err(error);
269
270        let error = annotated.into_data().unwrap_err();
271        assert_eq!(
272            error.error_type(),
273            ErrorType::Backend(BackendError::InvalidArgument)
274        );
275        assert_eq!(error.message(), "invalid request");
276    }
277
278    #[test]
279    fn test_into_result() {
280        let err = DynamoError::msg("connection lost");
281        let annotated = Annotated::<String>::from_err(err);
282
283        let result = annotated.into_result();
284        assert!(result.is_err());
285        assert!(result.unwrap_err().to_string().contains("connection lost"));
286    }
287}
dynamo_runtime/protocols/annotated.rs

dynamo_runtime/protocols/
annotated.rs
