dynamo-runtime 1.4.0

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use super::maybe_error::MaybeError;
use crate::error::DynamoError;
use anyhow::Result;
use serde::{Deserialize, Serialize};

pub trait AnnotationsProvider {
    fn annotations(&self) -> Option<Vec<String>>;
    fn has_annotation(&self, annotation: &str) -> bool {
        self.annotations()
            .map(|annotations| annotations.iter().any(|a| a == annotation))
            .unwrap_or(false)
    }
}

/// Our services have the option of returning an "annotated" stream, which allows use
/// to include additional information with each delta. This is useful for debugging,
/// performance benchmarking, and improved observability.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct Annotated<R> {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<R>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub event: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub comment: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<DynamoError>,
}

impl<R> Annotated<R> {
    fn cloned_error(&self) -> Option<DynamoError> {
        self.is_error().then(|| {
            self.error.clone().unwrap_or_else(|| {
                DynamoError::msg(
                    self.comment
                        .as_ref()
                        .filter(|comments| !comments.is_empty())
                        .map(|comments| comments.join(", "))
                        .unwrap_or_else(|| "unknown error".to_string()),
                )
            })
        })
    }

    /// Create a new annotated stream from the given error string
    pub fn from_error(error: impl Into<String>) -> Self {
        Self {
            data: None,
            id: None,
            event: Some("error".to_string()),
            comment: None,
            error: Some(DynamoError::msg(error)),
        }
    }

    /// Create a new annotated stream from the given data
    pub fn from_data(data: R) -> Self {
        Self {
            data: Some(data),
            id: None,
            event: None,
            comment: None,
            error: None,
        }
    }

    /// Add an annotation to the stream
    ///
    /// Annotations populate the `event` field and the `comment` field
    pub fn from_annotation<S: Serialize>(
        name: impl Into<String>,
        value: &S,
    ) -> Result<Self, serde_json::Error> {
        Ok(Self {
            data: None,
            id: None,
            event: Some(name.into()),
            comment: Some(vec![serde_json::to_string(value)?]),
            error: None,
        })
    }

    /// Convert to a [`Result<Self, String>`]
    /// If [`Self::event`] is "error", return an error message
    pub fn ok(self) -> Result<Self, String> {
        if let Some(error) = self.cloned_error() {
            return Err(error.to_string());
        }
        Ok(self)
    }

    /// Consume the annotation and return its optional payload while preserving
    /// structured errors.
    pub fn into_data(self) -> Result<Option<R>, DynamoError> {
        if let Some(error) = self.cloned_error() {
            return Err(error);
        }
        Ok(self.data)
    }

    pub fn is_ok(&self) -> bool {
        self.event.as_deref() != Some("error")
    }

    pub fn is_event(&self) -> bool {
        self.event.is_some()
    }

    pub fn transfer<U: Serialize>(self, data: Option<U>) -> Annotated<U> {
        Annotated::<U> {
            data,
            id: self.id,
            event: self.event,
            comment: self.comment,
            error: self.error,
        }
    }

    /// Apply a mapping/transformation to the data field
    /// If the mapping fails, the error is returned as an annotated stream
    pub fn map_data<U, F>(self, transform: F) -> Annotated<U>
    where
        F: FnOnce(R) -> Result<U, String>,
    {
        match self.data.map(transform).transpose() {
            Ok(data) => Annotated::<U> {
                data,
                id: self.id,
                event: self.event,
                comment: self.comment,
                error: self.error,
            },
            Err(e) => Annotated::from_error(e),
        }
    }

    pub fn is_error(&self) -> bool {
        self.event.as_deref() == Some("error")
    }

    pub fn into_result(self) -> Result<Option<R>> {
        self.into_data().map_err(anyhow::Error::new)
    }
}

impl<R> MaybeError for Annotated<R>
where
    R: for<'de> Deserialize<'de>,
{
    fn from_err(err: impl std::error::Error + 'static) -> Self {
        Self {
            data: None,
            id: None,
            event: Some("error".to_string()),
            comment: None,
            error: Some(DynamoError::from(
                Box::new(err) as Box<dyn std::error::Error + 'static>
            )),
        }
    }

    fn err(&self) -> Option<DynamoError> {
        self.cloned_error()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_maybe_error() {
        let annotated = Annotated::from_data("Test data".to_string());
        assert!(annotated.err().is_none());
        assert!(annotated.is_ok());

        let annotated = Annotated::<String>::from_error("Test error 2".to_string());
        assert!(annotated.err().is_some());
        assert!(annotated.is_err());

        let dynamo_err = DynamoError::msg("Test error 3");
        let annotated = Annotated::<String>::from_err(dynamo_err);
        assert!(annotated.is_err());
    }

    #[test]
    fn test_from_err() {
        let err = DynamoError::msg("connection lost");
        let annotated = Annotated::<String>::from_err(err);

        assert!(annotated.is_err());
        let err = annotated.err().unwrap();
        assert!(err.to_string().contains("connection lost"));
    }

    #[test]
    fn test_comment_only_error_fallback() {
        for (comments, expected) in [
            (vec![], "unknown error"),
            (
                vec!["first".to_string(), "second".to_string()],
                "first, second",
            ),
        ] {
            let annotated = Annotated::<String> {
                data: None,
                id: None,
                event: Some("error".to_string()),
                comment: Some(comments),
                error: None,
            };

            assert_eq!(annotated.err().unwrap().message(), expected);
        }
    }

    #[test]
    fn test_error_serialization() {
        let err = DynamoError::msg("test error");
        let annotated = Annotated::<String>::from_err(err);

        // Serialize and deserialize
        let json = serde_json::to_string(&annotated).unwrap();
        let deserialized: Annotated<String> = serde_json::from_str(&json).unwrap();

        assert!(deserialized.is_err());
        assert!(
            deserialized
                .err()
                .unwrap()
                .to_string()
                .contains("test error")
        );
    }

    #[test]
    fn test_transfer_preserves_error() {
        let err = DynamoError::msg("request timed out");
        let annotated = Annotated::<String>::from_err(err);

        let transferred: Annotated<i32> = annotated.transfer(None);
        assert!(transferred.err().is_some());
    }

    #[test]
    fn test_ok_method() {
        let err = DynamoError::msg("connection lost");
        let annotated = Annotated::<String>::from_err(err);

        let result = annotated.ok();
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("connection lost"));
    }

    #[test]
    fn test_into_data_preserves_error_type() {
        use crate::error::{BackendError, ErrorType};

        let error = DynamoError::builder()
            .error_type(ErrorType::Backend(BackendError::InvalidArgument))
            .message("invalid request")
            .build();
        let annotated = Annotated::<String>::from_err(error);

        let error = annotated.into_data().unwrap_err();
        assert_eq!(
            error.error_type(),
            ErrorType::Backend(BackendError::InvalidArgument)
        );
        assert_eq!(error.message(), "invalid request");
    }

    #[test]
    fn test_into_result() {
        let err = DynamoError::msg("connection lost");
        let annotated = Annotated::<String>::from_err(err);

        let result = annotated.into_result();
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("connection lost"));
    }
}