use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
use orion_error::conversion::ToStructError;
use orion_error::{OrionError, StructError, UnifiedReason};
use std::error::Error as StdError;
#[derive(Debug, Clone, PartialEq, OrionError)]
pub enum SourceReason {
#[orion_error(message = "end of stream", identity = "sys.wf_connector.eof")]
EOF,
#[orion_error(message = "no data available", identity = "sys.wf_connector.not_data")]
NotData,
#[orion_error(message = "I/O error", identity = "sys.wf_connector.io")]
Io,
#[orion_error(message = "connection error", identity = "sys.wf_connector.connect")]
Connect,
#[orion_error(message = "decode error", identity = "sys.wf_connector.decode")]
Decode,
#[orion_error(
message = "connector not found",
identity = "sys.wf_connector.not_found"
)]
NotFound,
#[orion_error(transparent)]
General(UnifiedReason),
}
impl SourceReason {
pub fn err_detail<S: Into<String>>(self, detail: S) -> SourceError {
self.to_err().with_detail(detail.into())
}
pub fn err_source<E>(self, source: E) -> SourceError
where
E: StdError + Send + Sync + 'static,
{
self.to_err().with_source(source)
}
}
pub type SourceError = StructError<SourceReason>;
pub type SourceResult<T> = Result<T, SourceError>;
#[async_trait]
pub trait BatchSource: Send {
async fn start(&mut self) -> SourceResult<()> {
Ok(())
}
async fn receive_batch(&mut self) -> SourceResult<Vec<RecordBatch>>;
async fn close(&mut self) -> SourceResult<()> {
Ok(())
}
fn identifier(&self) -> &str;
}