use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use futures::stream::{Stream, StreamExt};
use crate::error::DbError;
use crate::pool::QueryStreamItem;
use crate::value::Value;
pub type RowResult = HashMap<String, Value>;
pub trait AsyncRowStream: Send {
fn next_row<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Option<Result<RowResult, DbError>>> + Send + 'a>>;
fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
}
pub struct CursorRowStream<S: Stream<Item = QueryStreamItem> + Send + Unpin> {
inner: S,
closed: bool,
}
impl<S: Stream<Item = QueryStreamItem> + Send + Unpin> CursorRowStream<S> {
#[must_use]
pub fn new(stream: S) -> Self {
Self {
inner: stream,
closed: false,
}
}
}
impl<S: Stream<Item = QueryStreamItem> + Send + Unpin> AsyncRowStream for CursorRowStream<S> {
fn next_row<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Option<Result<RowResult, DbError>>> + Send + 'a>> {
if self.closed {
return Box::pin(async { None });
}
Box::pin(async move { self.inner.next().await })
}
fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
self.closed = true;
Ok(())
})
}
}
pub struct BoxedCursorRowStream<'a> {
inner: Pin<Box<dyn Stream<Item = QueryStreamItem> + Send + 'a>>,
closed: bool,
}
impl<'a> BoxedCursorRowStream<'a> {
#[must_use]
pub fn new(stream: Pin<Box<dyn Stream<Item = QueryStreamItem> + Send + 'a>>) -> Self {
Self {
inner: stream,
closed: false,
}
}
}
impl<'a> AsyncRowStream for BoxedCursorRowStream<'a> {
fn next_row<'b>(
&'b mut self,
) -> Pin<Box<dyn Future<Output = Option<Result<RowResult, DbError>>> + Send + 'b>> {
if self.closed {
return Box::pin(async { None });
}
Box::pin(async move { self.inner.next().await })
}
fn close<'b>(&'b mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'b>> {
Box::pin(async move {
self.closed = true;
Ok(())
})
}
}
impl<'a> AsyncRowStream for Box<dyn AsyncRowStream + 'a> {
fn next_row<'b>(
&'b mut self,
) -> Pin<Box<dyn Future<Output = Option<Result<RowResult, DbError>>> + Send + 'b>> {
(**self).next_row()
}
fn close<'b>(&'b mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'b>> {
(**self).close()
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::stream;
fn make_row(id: i64) -> RowResult {
let mut row = HashMap::new();
row.insert("id".to_string(), Value::I64(id));
row
}
#[tokio::test]
async fn test_next_row_yields_rows() {
let rows: Vec<QueryStreamItem> = vec![Ok(make_row(1)), Ok(make_row(2)), Ok(make_row(3))];
let stream = stream::iter(rows);
let mut s = CursorRowStream::new(stream);
let row = s.next_row().await.unwrap().unwrap();
assert_eq!(row.get("id"), Some(&Value::I64(1)));
let row = s.next_row().await.unwrap().unwrap();
assert_eq!(row.get("id"), Some(&Value::I64(2)));
let row = s.next_row().await.unwrap().unwrap();
assert_eq!(row.get("id"), Some(&Value::I64(3)));
}
#[tokio::test]
async fn test_stream_end_none() {
let rows: Vec<QueryStreamItem> = vec![Ok(make_row(1))];
let stream = stream::iter(rows);
let mut s = CursorRowStream::new(stream);
s.next_row().await;
assert!(s.next_row().await.is_none());
}
#[tokio::test]
async fn test_error_propagation() {
let rows: Vec<QueryStreamItem> = vec![Err(DbError::QueryError("test error".into()))];
let stream = stream::iter(rows);
let mut s = CursorRowStream::new(stream);
let result = s.next_row().await;
assert!(matches!(result, Some(Err(_))));
}
#[tokio::test]
async fn test_close_then_none() {
let rows: Vec<QueryStreamItem> = vec![Ok(make_row(1)), Ok(make_row(2))];
let stream = stream::iter(rows);
let mut s = CursorRowStream::new(stream);
s.close().await.unwrap();
assert!(s.next_row().await.is_none());
}
#[tokio::test]
async fn test_cursor_row_stream_degradation() {
let rows: Vec<QueryStreamItem> = vec![Ok(make_row(1))];
let stream = stream::iter(rows);
let mut s: BoxedCursorRowStream<'_> = BoxedCursorRowStream::new(Box::pin(stream));
let row = s.next_row().await.unwrap().unwrap();
assert_eq!(row.get("id"), Some(&Value::I64(1)));
assert!(s.next_row().await.is_none());
}
#[tokio::test]
async fn test_drop_releases() {
let rows: Vec<QueryStreamItem> = vec![Ok(make_row(1))];
let stream = stream::iter(rows);
let s = CursorRowStream::new(stream);
drop(s);
}
}