rama-core 0.3.0

use core::pin::Pin;
use core::task::{Context, Poll};

use pin_project_lite::pin_project;
use rama_error::{BoxError, ErrorContext as _};
use serde::Deserialize;

use crate::futures::{Stream, ready};
use crate::stream::json::config::ParseConfig;
use crate::stream::json::engine::NdjsonEngine;

pin_project! {
    /// Wraps a [Stream] of [Result]s of data blocks, i.e. types that reference as byte array, and offers
    /// a [Stream] implementation over parsed NDJSON-records according to [Deserialize], forwarding
    /// potential errors returned by the wrapped iterator.
    pub struct JsonReadStream<T, S> {
        engine: NdjsonEngine<T>,
        #[pin]
        bytes_stream: S
    }
}

impl<T, S> JsonReadStream<T, S> {
    /// Creates a new fallible NDJSON-stream wrapping the given `bytes_stream` with default
    /// [ParseConfig].
    pub fn new(bytes_stream: S) -> Self {
        Self {
            engine: NdjsonEngine::new(),
            bytes_stream,
        }
    }

    /// Creates a new fallible NDJSON-stream wrapping the given `bytes_stream` with the given
    /// [ParseConfig] to control its behavior. See [ParseConfig] for more details.
    pub fn new_with_config(bytes_stream: S, config: ParseConfig) -> Self {
        Self {
            engine: NdjsonEngine::with_config(config),
            bytes_stream,
        }
    }

    pub fn into_inner(self) -> S {
        self.bytes_stream
    }
}

impl<T, S, B, E> Stream for JsonReadStream<T, S>
where
    for<'deserialize> T: Deserialize<'deserialize>,
    E: Into<BoxError>,
    S: Stream<Item = Result<B, E>>,
    B: AsRef<[u8]>,
{
    type Item = Result<T, BoxError>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        loop {
            if let Some(result) = this.engine.pop() {
                return Poll::Ready(Some(result.context("json-deserialize next value")));
            }

            let bytes = ready!(this.bytes_stream.as_mut().poll_next(cx));

            match bytes {
                Some(Ok(bytes)) => this.engine.input(bytes),
                Some(Err(err)) => {
                    let err = err.into();
                    return Poll::Ready(Some(Err(err)));
                }
                None => {
                    this.engine.finalize();
                    return Poll::Ready(
                        this.engine
                            .pop()
                            .map(|res| res.context("json-deserialize last value")),
                    );
                }
            }
        }
    }
}

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

    use core::convert::Infallible;
    use core::pin::pin;

    use crate::futures::StreamExt;
    use crate::futures::stream;
    use rama_error::BoxErrorExt;

    use tokio_test::assert_pending;
    use tokio_test::task;

    use crate::stream::json::EmptyLineHandling;

    #[derive(Debug, Deserialize, Eq, PartialEq)]
    struct TestStruct {
        key: u64,
        value: u64,
    }

    struct SingleThenPanicIter {
        data: Option<String>,
    }

    impl Iterator for SingleThenPanicIter {
        type Item = Result<String, BoxError>;

        fn next(&mut self) -> Option<Self::Item> {
            Some(Ok(self.data.take().expect("iterator queried twice")))
        }
    }

    #[test]
    fn pending_stream_results_in_pending_item() {
        let mut ndjson_stream: JsonReadStream<(), _> =
            JsonReadStream::new(stream::pending::<Result<&str, BoxError>>());

        let mut next = task::spawn(ndjson_stream.next());

        assert_pending!(next.poll());
    }

    #[test]
    fn empty_stream_results_in_empty_results() {
        let collected = tokio_test::block_on(
            JsonReadStream::<_, _>::new(stream::empty::<Result<&[u8], BoxError>>())
                .collect::<Vec<Result<(), BoxError>>>(),
        );
        assert!(collected.is_empty());
    }

    #[test]
    fn singleton_iter_with_single_json_line() {
        let stream = stream::once(async { Ok::<_, Infallible>("{\"key\":1,\"value\":2}\n") });

        let collected = tokio_test::block_on(
            JsonReadStream::<_, _>::new(stream).collect::<Vec<Result<TestStruct, BoxError>>>(),
        );

        let mut result = collected.into_iter();
        assert_eq!(
            result.next().unwrap().unwrap(),
            TestStruct { key: 1, value: 2 }
        );
        assert!(result.next().is_none());
    }

    #[test]
    fn multiple_iter_items_compose_single_json_line() {
        let stream = stream::iter(vec![
            Ok::<_, Infallible>("{\"key\""),
            Ok::<_, Infallible>(":12,"),
            Ok::<_, Infallible>("\"value\""),
            Ok::<_, Infallible>(":34}\n"),
        ]);

        let collected = tokio_test::block_on(
            JsonReadStream::<_, _>::new(stream).collect::<Vec<Result<TestStruct, BoxError>>>(),
        );

        let mut result = collected.into_iter();
        assert_eq!(
            result.next().unwrap().unwrap(),
            TestStruct { key: 12, value: 34 }
        );
        assert!(result.next().is_none());
    }

    #[tokio::test]
    async fn wrapped_stream_not_queried_while_sufficient_data_remains() {
        let iter = SingleThenPanicIter {
            data: Some("{\"key\":0,\"value\":0}\n{\"key\":0,\"value\":0}\n".to_owned()),
        };
        let mut ndjson_stream = JsonReadStream::<TestStruct, _>::new(stream::iter(iter));

        assert!(ndjson_stream.next().await.is_some());
        assert!(ndjson_stream.next().await.is_some());
    }

    #[tokio::test]
    async fn stream_with_parse_always_config_respects_config() {
        let stream = stream::once(async { Ok::<_, Infallible>("{\"key\":1,\"value\":2}\n\n") });
        let config =
            ParseConfig::default().with_empty_line_handling(EmptyLineHandling::ParseAlways);
        let mut ndjson_stream = pin!(JsonReadStream::<TestStruct, _>::new_with_config(
            stream, config
        ));

        ndjson_stream.next().await.unwrap().unwrap();
        ndjson_stream.next().await.unwrap().unwrap_err();
    }

    #[tokio::test]
    async fn stream_with_ignore_empty_config_respects_config() {
        let stream = stream::once(async { Ok::<_, Infallible>("{\"key\":1,\"value\":2}\n\n") });
        let config =
            ParseConfig::default().with_empty_line_handling(EmptyLineHandling::IgnoreEmpty);
        let mut ndjson_stream = pin!(JsonReadStream::<TestStruct, _>::new_with_config(
            stream, config
        ));

        ndjson_stream.next().await.unwrap().unwrap();
        assert!(ndjson_stream.next().await.is_none());
    }

    #[tokio::test]
    async fn stream_with_parse_rest_handles_valid_finalization() {
        let stream = stream::once(async { Ok::<_, Infallible>("{\"key\":1,\"value\":2}") });
        let config = ParseConfig::default().with_parse_rest(true);
        let mut ndjson_stream = pin!(JsonReadStream::<TestStruct, _>::new_with_config(
            stream, config
        ));

        assert_eq!(
            ndjson_stream.next().await.unwrap().unwrap(),
            TestStruct { key: 1, value: 2 }
        );
        assert!(ndjson_stream.next().await.is_none());
    }

    #[tokio::test]
    async fn stream_with_parse_rest_handles_invalid_finalization() {
        let stream = stream::once(async { Ok::<_, Infallible>("{\"key\":1,") });
        let config = ParseConfig::default().with_parse_rest(true);
        let mut ndjson_stream = pin!(JsonReadStream::<TestStruct, _>::new_with_config(
            stream, config
        ));

        ndjson_stream.next().await.unwrap().unwrap_err();
        assert!(ndjson_stream.next().await.is_none());
    }

    #[tokio::test]
    async fn stream_without_parse_rest_does_not_handle_finalization() {
        let stream = stream::once(async { Ok::<_, Infallible>("some text") });
        let config = ParseConfig::default().with_parse_rest(false);
        let mut ndjson_stream = pin!(JsonReadStream::<TestStruct, _>::new_with_config(
            stream, config
        ));

        assert!(ndjson_stream.next().await.is_none());
    }

    #[test]
    fn fallible_stream_operates_correctly_with_interspersed_errors() {
        let data_vec = vec![
            Err(BoxError::from_static_str("test message 1")),
            Ok("invalid json\n{\"key\":11,\"val"),
            Ok("ue\":22}\n{\"key\":33,\"value\":44}\ninvalid json\n"),
            Err(BoxError::from_static_str("test message 2")),
            Ok("{\"key\":55,\"value\":66}\n"),
        ];
        let data_stream = stream::iter(data_vec);
        let fallible_ndjson_stream = JsonReadStream::<TestStruct, _>::new(data_stream);

        let mut iter = tokio_test::block_on(fallible_ndjson_stream.collect::<Vec<_>>()).into_iter();

        iter.next().unwrap().unwrap_err();
        iter.next().unwrap().unwrap_err();
        assert_eq!(
            TestStruct { key: 11, value: 22 },
            iter.next().unwrap().unwrap()
        );
        assert_eq!(
            TestStruct { key: 33, value: 44 },
            iter.next().unwrap().unwrap()
        );
        iter.next().unwrap().unwrap_err();
        iter.next().unwrap().unwrap_err();
        assert_eq!(
            TestStruct { key: 55, value: 66 },
            iter.next().unwrap().unwrap()
        );
        assert!(iter.next().is_none());
    }
}