use crate::WyrdClient;
use arrow::record_batch::RecordBatch;
use wyrd_spec::request_id::RequestId;
#[cfg(feature = "internal")]
use wyrd_spec::vala::api::BifrostQueryRequest;
use wyrd_spec::vala::api::{
BifrostTableDescription, BifrostTableEntry, CancelRunningQueryResponse, QueryParam,
RegisterOutcome, RunningQuerySummary,
};
use crate::bifrost::facade::{BifrostOptions, QueryResult};
use crate::bifrost::query::{BifrostClientError, QueryResultStream};
#[cfg(feature = "internal")]
use crate::bifrost::query::{CollectedQueryLimits, CollectedQueryResult};
use crate::bifrost::table::{Correlation, TableConfig};
pub struct Bifrost {
inner: crate::bifrost::Bifrost,
}
impl Bifrost {
pub fn connect(options: BifrostOptions) -> Result<Self, BifrostClientError> {
Ok(Self {
inner: block_on(crate::bifrost::Bifrost::connect(options))?,
})
}
pub fn connect_with(
client: &WyrdClient,
options: BifrostOptions,
) -> Result<Self, BifrostClientError> {
Ok(Self {
inner: block_on(crate::bifrost::Bifrost::connect_with(client, options))?,
})
}
pub fn register(&self) -> Result<RegisterOutcome, BifrostClientError> {
block_on(self.inner.register())
}
pub fn use_table(&self, table: TableConfig) -> Option<TableConfig> {
self.inner.use_table(table)
}
pub fn use_table_by_name(&self, fqn: &str) -> Result<(), BifrostClientError> {
block_on(self.inner.use_table_by_name(fqn))
}
#[must_use]
pub fn table(&self) -> Option<TableConfig> {
self.inner.table()
}
pub fn insert(&self, row: Vec<u8>, correlation: Correlation) -> Result<(), BifrostClientError> {
self.inner.insert(row, correlation)
}
pub fn flush(&self) -> Result<(), BifrostClientError> {
block_on(self.inner.flush())
}
pub fn shutdown(&self) -> Result<(), BifrostClientError> {
block_on(self.inner.shutdown())
}
pub fn sql(
&self,
query: &str,
params: &[QueryParam],
) -> Result<QueryResult, BifrostClientError> {
block_on(self.inner.sql(query, params))
}
pub fn sql_as<T: serde::de::DeserializeOwned>(
&self,
query: &str,
params: &[QueryParam],
) -> Result<Vec<T>, BifrostClientError> {
block_on(self.inner.sql_as(query, params))
}
pub fn trace_timeline(
&self,
trace_id: wyrd_spec::vala::ids::TraceId,
) -> Result<Vec<wyrd_spec::vala::trace::SpanRecord>, BifrostClientError> {
block_on(self.inner.trace_timeline(trace_id))
}
pub fn stream(
&self,
query: &str,
params: &[QueryParam],
deadline: Option<std::time::Duration>,
) -> Result<BlockingQueryStream, BifrostClientError> {
Ok(BlockingQueryStream {
inner: block_on(self.inner.stream(query, params, deadline))?,
})
}
pub fn list_tables(&self) -> Result<Vec<BifrostTableEntry>, BifrostClientError> {
block_on(self.inner.list_tables())
}
pub fn describe_table(
&self,
namespace: &str,
name: &str,
) -> Result<BifrostTableDescription, BifrostClientError> {
block_on(self.inner.describe_table(namespace, name))
}
#[cfg(feature = "internal")]
pub fn query(
&self,
request: &BifrostQueryRequest,
) -> Result<BlockingQueryStream, BifrostClientError> {
Ok(BlockingQueryStream {
inner: block_on(self.inner.query(request))?,
})
}
#[cfg(feature = "internal")]
pub fn collect_bounded(
&self,
request: &BifrostQueryRequest,
limits: CollectedQueryLimits,
) -> Result<CollectedQueryResult, BifrostClientError> {
block_on(self.inner.collect_bounded(request, limits))
}
pub fn running(&self) -> Result<Vec<RunningQuerySummary>, BifrostClientError> {
block_on(self.inner.running())
}
pub fn status(
&self,
request_id: &RequestId,
) -> Result<RunningQuerySummary, BifrostClientError> {
block_on(self.inner.status(request_id))
}
pub fn cancel(
&self,
request_id: &RequestId,
) -> Result<CancelRunningQueryResponse, BifrostClientError> {
block_on(self.inner.cancel(request_id))
}
#[must_use]
#[cfg(feature = "internal")]
pub fn into_async(self) -> crate::bifrost::Bifrost {
self.inner
}
}
pub struct BlockingQueryStream {
inner: QueryResultStream,
}
impl BlockingQueryStream {
#[must_use]
pub fn terminal(&self) -> Option<&wyrd_spec::vala::api::QueryTerminalFrame> {
self.inner.terminal()
}
}
impl Iterator for BlockingQueryStream {
type Item = Result<RecordBatch, BifrostClientError>;
fn next(&mut self) -> Option<Self::Item> {
match block_on(self.inner.next_batch()) {
Ok(Some(batch)) => Some(Ok(batch)),
Ok(None) => None,
Err(error) => Some(Err(error)),
}
}
}
fn block_on<T>(future: impl Future<Output = T>) -> T {
wyrd_runtime::runtime().block_on(future)
}