use crate::{
db::{
DbSession, DynamicQuery, OutputRow, PreparedOutputRows, PrimaryKeyValue, TypedAdapterError,
TypedEntityBinding, TypedOperationError,
session::{OutputRowProjection, live_page::prepare_live_page_step},
},
traits::CanisterKind,
};
#[doc(hidden)]
pub struct PreparedLivePageOutput {
pub rows: PreparedOutputRows,
pub continuation: Option<String>,
pub work: crate::db::ScalarPageWork,
}
#[doc(hidden)]
pub struct PreparedExactKeyOutput {
pub distinct_rows: Vec<Option<OutputRow>>,
pub positions: Vec<u32>,
}
#[doc(hidden)]
pub struct PreparedLivePageCursor<'session, C>
where
C: CanisterKind,
{
session: &'session DbSession<C>,
binding: TypedEntityBinding,
request: DynamicQuery,
continuation: Option<String>,
exhausted: bool,
}
impl<'session, C> PreparedLivePageCursor<'session, C>
where
C: CanisterKind,
{
const fn new(
session: &'session DbSession<C>,
binding: TypedEntityBinding,
request: DynamicQuery,
) -> Self {
Self {
session,
binding,
request,
continuation: None,
exhausted: false,
}
}
#[must_use]
pub const fn binding(&self) -> &TypedEntityBinding {
&self.binding
}
pub fn execute_page(
&self,
continuation: Option<&str>,
) -> Result<PreparedLivePageOutput, TypedOperationError> {
let page = self
.session
.inner
.execute_public_live_page_for_typed_binding(
self.binding.inner(),
&self.request,
continuation,
)
.map_err(|error| TypedOperationError::Database(crate::Error::from(error)))?
.ok_or_else(stale_binding_error)?;
self.prepare_page(page, continuation)
}
pub fn execute_trusted_page(
&self,
continuation: Option<&str>,
) -> Result<PreparedLivePageOutput, TypedOperationError> {
let page = self
.session
.inner
.execute_trusted_live_page(&self.request, continuation)
.map_err(|error| TypedOperationError::Database(crate::Error::from(error)))?;
self.prepare_page(page, continuation)
}
#[inline(never)]
pub fn next_page(&mut self) -> Result<Option<PreparedOutputRows>, TypedOperationError> {
if self.exhausted {
return Ok(None);
}
let prepared = self.execute_page(self.continuation.as_deref())?;
Ok(Some(self.adopt_page(prepared)))
}
#[inline(never)]
pub fn next_trusted_page(&mut self) -> Result<Option<PreparedOutputRows>, TypedOperationError> {
if self.exhausted {
return Ok(None);
}
let prepared = self.execute_trusted_page(self.continuation.as_deref())?;
Ok(Some(self.adopt_page(prepared)))
}
fn adopt_page(&mut self, prepared: PreparedLivePageOutput) -> PreparedOutputRows {
self.continuation = prepared.continuation;
self.exhausted = self.continuation.is_none();
prepared.rows
}
fn prepare_page(
&self,
page: crate::db::LiveQueryPageOutput,
continuation: Option<&str>,
) -> Result<PreparedLivePageOutput, TypedOperationError> {
let crate::db::LiveQueryPageOutput {
entity,
columns,
rows,
row_count: _,
continuation,
work,
} = prepare_live_page_step(page, continuation)?.into_page();
let rows = self
.session
.prepare_typed_output_rows(&self.binding, entity, columns, rows)?;
Ok(PreparedLivePageOutput {
rows,
continuation,
work,
})
}
}
impl<C: CanisterKind> DbSession<C> {
#[doc(hidden)]
pub fn execute_public_prepared_exact_key_batch(
&self,
binding: &TypedEntityBinding,
keys: &[PrimaryKeyValue],
) -> Result<PreparedExactKeyOutput, TypedOperationError> {
let output = self
.inner
.execute_public_exact_key_batch_for_typed_binding(binding.inner(), keys)
.map_err(|error| TypedOperationError::Database(crate::Error::from(error)))?
.ok_or_else(stale_binding_error)?;
let icydb_core::db::ExactKeyBatchProjectionOutput {
entity,
columns,
distinct_rows,
positions,
} = output;
let projection = OutputRowProjection::new(binding, entity, columns.as_slice())
.map_err(TypedOperationError::Adapter)?;
let distinct_rows = distinct_rows
.into_iter()
.map(|values| {
values
.map(|values| projection.project(values))
.transpose()
.map_err(TypedOperationError::Adapter)
})
.collect::<Result<Vec<_>, _>>()?;
Ok(PreparedExactKeyOutput {
distinct_rows,
positions,
})
}
#[doc(hidden)]
#[must_use]
pub const fn prepare_live_page_cursor(
&self,
binding: TypedEntityBinding,
request: DynamicQuery,
) -> PreparedLivePageCursor<'_, C> {
PreparedLivePageCursor::new(self, binding, request)
}
}
const fn stale_binding_error() -> TypedOperationError {
TypedOperationError::Adapter(TypedAdapterError::StaleBinding)
}