use crate::{
db::{
DbSession, DynamicQuery, ExhaustiveReadError, GroupedQueryOutput, TypedBindingError,
TypedEntityAdapter, TypedEntityBinding, TypedRowError,
},
traits::{CanisterKind, EntityKey},
types::Id,
};
use candid::CandidType;
use icydb_core::db::{AggregateExpr, FilterExpr, OrderTerm};
use serde::Deserialize;
use std::{error::Error as StdError, fmt, marker::PhantomData};
#[derive(Debug)]
pub enum TypedQueryError {
Database(crate::Error),
Row(TypedRowError),
}
#[derive(CandidType, Clone, Debug, Deserialize, Eq, PartialEq)]
pub struct LivePage<Row> {
pub rows: Vec<Row>,
pub continuation: Option<String>,
pub work: crate::db::ScalarPageWork,
}
#[derive(CandidType, Clone, Debug, Deserialize, Eq, PartialEq)]
pub struct ExhaustivePage<Row> {
pub rows: Vec<Row>,
pub continuation: Option<String>,
pub work: crate::db::ScalarPageWork,
pub proof: crate::db::ReadSetRevisionProof,
}
#[derive(Debug)]
pub enum TypedExhaustiveQueryError {
Exhaustive(ExhaustiveReadError),
Row(TypedRowError),
}
impl fmt::Display for TypedExhaustiveQueryError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Exhaustive(error) => error.fmt(formatter),
Self::Row(error) => error.fmt(formatter),
}
}
}
impl StdError for TypedExhaustiveQueryError {}
impl fmt::Display for TypedQueryError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Database(error) => error.fmt(formatter),
Self::Row(error) => error.fmt(formatter),
}
}
}
impl StdError for TypedQueryError {}
fn typed_query_error_from_binding(error: TypedBindingError) -> TypedQueryError {
match error {
TypedBindingError::Adapter(error) => TypedQueryError::Row(TypedRowError::Adapter(error)),
TypedBindingError::Database(error) => TypedQueryError::Database(error),
}
}
pub struct Query<'session, C, E>
where
C: CanisterKind,
E: TypedEntityAdapter,
{
session: &'session DbSession<C>,
binding: TypedEntityBinding,
request: DynamicQuery,
entity: PhantomData<fn() -> E>,
}
impl<'session, C, E> Query<'session, C, E>
where
C: CanisterKind,
E: TypedEntityAdapter,
{
pub(crate) fn new(session: &'session DbSession<C>) -> Result<Self, TypedBindingError> {
let binding = E::typed_binding(session)?;
let request = DynamicQuery::new(binding.entity());
Ok(Self {
session,
binding,
request,
entity: PhantomData,
})
}
#[must_use]
pub fn filter(mut self, filter: impl Into<FilterExpr>) -> Self {
self.request = self.request.filter(filter);
self
}
#[must_use]
pub fn order_by(mut self, order: OrderTerm) -> Self {
self.request = self.request.order_by(order);
self
}
#[must_use]
pub fn select<I, S>(mut self, fields: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.request = self.request.select(fields);
self
}
#[must_use]
pub fn limit(mut self, limit: u32) -> Self {
self.request = self.request.limit(limit);
self
}
#[must_use]
pub fn group_by(mut self, field: impl Into<String>) -> Self {
self.request = self.request.group_by(field);
self
}
#[must_use]
pub fn aggregate(mut self, aggregate: AggregateExpr) -> Self {
self.request = self.request.aggregate(aggregate);
self
}
#[must_use]
pub fn grouped_limits(mut self, max_groups: u32, max_group_bytes: u32) -> Self {
self.request = self.request.grouped_limits(max_groups, max_group_bytes);
self
}
#[must_use]
pub fn cursor(mut self, cursor: impl Into<String>) -> Self {
self.request = self.request.cursor(cursor);
self
}
pub fn execute_exact_count(self) -> Result<u64, TypedQueryError> {
self.session
.execute_public_typed_exact_count(&self.binding, &self.request)
.map_err(TypedQueryError::Database)?
.ok_or({
TypedQueryError::Row(TypedRowError::Adapter(
crate::db::TypedAdapterError::StaleBinding,
))
})
}
pub fn execute_live_page(
self,
continuation: Option<&str>,
) -> Result<LivePage<E::Row>, TypedQueryError> {
let result = self
.session
.execute_public_typed_live_page(&self.binding, &self.request, continuation)
.map_err(TypedQueryError::Database)?
.ok_or({
TypedQueryError::Row(TypedRowError::Adapter(
crate::db::TypedAdapterError::StaleBinding,
))
})?;
let mut rows = Vec::with_capacity(result.rows.len());
for row_index in 0..result.rows.len() {
let row = self
.session
.typed_live_page_row(&self.binding, &result, row_index)
.map_err(TypedQueryError::Row)?;
rows.push(
E::decode_row(&self.binding, row)
.map_err(|error| TypedQueryError::Row(TypedRowError::Adapter(error)))?,
);
}
Ok(LivePage {
rows,
continuation: result.continuation,
work: result.work,
})
}
pub fn execute_exhaustive_page(
self,
continuation: Option<&str>,
proof: Option<&crate::db::ReadSetRevisionProof>,
) -> Result<ExhaustivePage<E::Row>, TypedExhaustiveQueryError> {
let result = self
.session
.execute_public_typed_exhaustive_page(&self.binding, &self.request, continuation, proof)
.map_err(TypedExhaustiveQueryError::Exhaustive)?
.ok_or({
TypedExhaustiveQueryError::Row(TypedRowError::Adapter(
crate::db::TypedAdapterError::StaleBinding,
))
})?;
let mut rows = Vec::with_capacity(result.rows.len());
for row_index in 0..result.rows.len() {
let row = self
.session
.typed_exhaustive_page_row(&self.binding, &result, row_index)
.map_err(TypedExhaustiveQueryError::Row)?;
rows.push(
E::decode_row(&self.binding, row).map_err(|error| {
TypedExhaustiveQueryError::Row(TypedRowError::Adapter(error))
})?,
);
}
Ok(ExhaustivePage {
rows,
continuation: result.continuation,
work: result.work,
proof: result.proof,
})
}
pub fn execute_grouped(self) -> Result<GroupedQueryOutput, TypedQueryError> {
self.session
.execute_public_typed_dynamic_grouped_query(&self.binding, &self.request)
.map_err(TypedQueryError::Database)?
.ok_or({
TypedQueryError::Row(TypedRowError::Adapter(
crate::db::TypedAdapterError::StaleBinding,
))
})
}
}
impl<C: CanisterKind> DbSession<C> {
pub fn query<E>(&self) -> Result<Query<'_, C, E>, TypedBindingError>
where
E: TypedEntityAdapter,
{
Query::new(self)
}
pub fn get<E>(&self, id: Id<E>) -> Result<Option<E::Row>, TypedQueryError>
where
E: EntityKey + TypedEntityAdapter,
E::Row: Clone,
{
let mut rows = self.get_many::<E>(&[id])?;
rows.pop().ok_or({
TypedQueryError::Row(TypedRowError::Adapter(
crate::db::TypedAdapterError::RowShapeMismatch,
))
})
}
pub fn get_many<E>(&self, ids: &[Id<E>]) -> Result<Vec<Option<E::Row>>, TypedQueryError>
where
E: EntityKey + TypedEntityAdapter,
E::Row: Clone,
{
let binding = E::typed_binding(self).map_err(typed_query_error_from_binding)?;
let result = self
.execute_public_typed_exact_key_batch(&binding, ids)
.map_err(TypedQueryError::Database)?
.ok_or({
TypedQueryError::Row(TypedRowError::Adapter(
crate::db::TypedAdapterError::StaleBinding,
))
})?;
let mut distinct_rows = Vec::with_capacity(result.distinct_rows.len());
for values in result.distinct_rows {
let row = values
.map(|values| {
let row = Self::typed_exact_key_row(
&binding,
result.entity.as_str(),
result.columns.as_slice(),
values,
)
.map_err(TypedQueryError::Row)?;
E::decode_row(&binding, row)
.map_err(|error| TypedQueryError::Row(TypedRowError::Adapter(error)))
})
.transpose()?;
distinct_rows.push(row);
}
result
.positions
.into_iter()
.map(|position| {
let index = usize::try_from(position).map_err(|_| {
TypedQueryError::Row(TypedRowError::Adapter(
crate::db::TypedAdapterError::RowShapeMismatch,
))
})?;
distinct_rows.get(index).cloned().ok_or({
TypedQueryError::Row(TypedRowError::Adapter(
crate::db::TypedAdapterError::RowShapeMismatch,
))
})
})
.collect()
}
}
pub const MAX_TYPED_EXACT_KEY_BATCH_ITEMS: usize = icydb_core::db::MAX_TYPED_EXACT_KEY_BATCH_ITEMS;
pub const MAX_TYPED_EXACT_KEY_BATCH_INPUT_BYTES: usize =
icydb_core::db::MAX_TYPED_EXACT_KEY_BATCH_INPUT_BYTES;
pub const MAX_TYPED_EXACT_KEY_BATCH_STORED_BYTES: usize =
icydb_core::db::MAX_TYPED_EXACT_KEY_BATCH_STORED_BYTES;
pub const MAX_TYPED_EXACT_KEY_BATCH_RESULT_BYTES: usize =
icydb_core::db::MAX_TYPED_EXACT_KEY_BATCH_RESULT_BYTES;