use async_trait::async_trait;
use serde_json::Value;
use crate::cosmos::CosmosError;
#[async_trait]
pub(crate) trait QueryStreamInner: Send {
async fn next_page(&mut self) -> Result<Option<Vec<Value>>, CosmosError>;
fn continuation_token(&self) -> Option<&str>;
}
pub struct QueryStream {
pub(crate) inner: Box<dyn QueryStreamInner>,
}
impl QueryStream {
pub(crate) fn new(inner: Box<dyn QueryStreamInner>) -> Self {
Self { inner }
}
pub async fn next_page(&mut self) -> Result<Option<Vec<Value>>, CosmosError> {
self.inner.next_page().await
}
pub fn continuation_token(&self) -> Option<&str> {
self.inner.continuation_token()
}
}
pub(crate) struct VecQueryStream {
results: Option<Vec<Value>>,
}
impl VecQueryStream {
pub(crate) fn new(results: Vec<Value>) -> Self {
Self {
results: Some(results),
}
}
}
#[async_trait]
impl QueryStreamInner for VecQueryStream {
async fn next_page(&mut self) -> Result<Option<Vec<Value>>, CosmosError> {
Ok(self.results.take())
}
fn continuation_token(&self) -> Option<&str> {
None
}
}