use crate::{
db::{
DbSession, PersistedRow, WriteBatchResponse,
data::{FieldSlot, StructuralPatch},
executor::{EntityAuthority, MutationMode},
schema::{AcceptedFieldAbsencePolicy, AcceptedRowLayoutRuntimeDescriptor},
},
error::InternalError,
traits::{CanisterKind, EntityCreateInput, EntityValue},
value::Value,
};
fn append_accepted_structural_patch_field(
entity_path: &'static str,
descriptor: &AcceptedRowLayoutRuntimeDescriptor<'_>,
patch: StructuralPatch,
field_name: &str,
value: Value,
) -> Result<StructuralPatch, InternalError> {
let slot = descriptor
.field_slot_index_by_name(field_name)
.ok_or_else(|| InternalError::mutation_structural_field_unknown(entity_path, field_name))?;
Ok(patch.set(FieldSlot::from_validated_index(slot), value))
}
fn validate_structural_patch_absence_policy<E>(
descriptor: &AcceptedRowLayoutRuntimeDescriptor<'_>,
patch: &StructuralPatch,
mode: MutationMode,
) -> Result<(), InternalError>
where
E: PersistedRow + EntityValue,
{
if matches!(mode, MutationMode::Update) {
return Ok(());
}
if patch_explicitly_writes_generated_field::<E>(patch) {
return Ok(());
}
let mut provided_slots = vec![false; descriptor.required_slot_count()];
for entry in patch.entries() {
let slot = entry.slot().index();
if slot < provided_slots.len() {
provided_slots[slot] = true;
}
}
for field in descriptor.fields() {
let slot = usize::from(field.slot().get());
if provided_slots.get(slot).copied().unwrap_or(false) {
continue;
}
if matches!(field.absence_policy(), AcceptedFieldAbsencePolicy::Required) {
return Err(
InternalError::mutation_structural_patch_required_field_missing(
E::PATH,
field.name(),
),
);
}
}
Ok(())
}
fn patch_explicitly_writes_generated_field<E>(patch: &StructuralPatch) -> bool
where
E: PersistedRow + EntityValue,
{
patch.entries().iter().any(|entry| {
E::MODEL
.fields()
.get(entry.slot().index())
.is_some_and(|field| {
field.insert_generation().is_some() && field.name() != E::MODEL.primary_key.name()
})
})
}
impl<C: CanisterKind> DbSession<C> {
pub fn insert<E>(&self, entity: E) -> Result<E, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_entity(|save| save.insert(entity))
}
pub fn create<I>(&self, input: I) -> Result<I::Entity, InternalError>
where
I: EntityCreateInput,
I::Entity: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_entity(|save| save.create(input))
}
pub fn insert_many_atomic<E>(
&self,
entities: impl IntoIterator<Item = E>,
) -> Result<WriteBatchResponse<E>, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_batch(|save| save.insert_many_atomic(entities))
}
pub fn insert_many_non_atomic<E>(
&self,
entities: impl IntoIterator<Item = E>,
) -> Result<WriteBatchResponse<E>, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_batch(|save| save.insert_many_non_atomic(entities))
}
pub fn replace<E>(&self, entity: E) -> Result<E, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_entity(|save| save.replace(entity))
}
pub fn mutate_structural<E>(
&self,
key: E::Key,
patch: StructuralPatch,
mode: MutationMode,
) -> Result<E, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
let (accepted_schema, _) =
self.ensure_accepted_schema_snapshot_and_authority(EntityAuthority::for_type::<E>())?;
let descriptor =
AcceptedRowLayoutRuntimeDescriptor::from_accepted_schema(&accepted_schema)?;
validate_structural_patch_absence_policy::<E>(&descriptor, &patch, mode)?;
self.execute_save_entity(|save| save.apply_structural_mutation(mode, key, patch))
}
pub fn structural_patch<E, I, S>(&self, fields: I) -> Result<StructuralPatch, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
I: IntoIterator<Item = (S, Value)>,
S: AsRef<str>,
{
let (accepted_schema, _) =
self.ensure_accepted_schema_snapshot_and_authority(EntityAuthority::for_type::<E>())?;
let descriptor =
AcceptedRowLayoutRuntimeDescriptor::from_accepted_schema(&accepted_schema)?;
let mut patch = StructuralPatch::new();
for (field_name, value) in fields {
let field_name = field_name.as_ref();
patch = append_accepted_structural_patch_field(
E::PATH,
&descriptor,
patch,
field_name,
value,
)?;
}
Ok(patch)
}
#[cfg(test)]
pub(in crate::db) fn replace_structural<E>(
&self,
key: E::Key,
patch: StructuralPatch,
) -> Result<E, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.mutate_structural(key, patch, MutationMode::Replace)
}
pub fn replace_many_atomic<E>(
&self,
entities: impl IntoIterator<Item = E>,
) -> Result<WriteBatchResponse<E>, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_batch(|save| save.replace_many_atomic(entities))
}
pub fn replace_many_non_atomic<E>(
&self,
entities: impl IntoIterator<Item = E>,
) -> Result<WriteBatchResponse<E>, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_batch(|save| save.replace_many_non_atomic(entities))
}
pub fn update<E>(&self, entity: E) -> Result<E, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_entity(|save| save.update(entity))
}
#[cfg(test)]
pub(in crate::db) fn insert_structural<E>(
&self,
key: E::Key,
patch: StructuralPatch,
) -> Result<E, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.mutate_structural(key, patch, MutationMode::Insert)
}
#[cfg(test)]
pub(in crate::db) fn update_structural<E>(
&self,
key: E::Key,
patch: StructuralPatch,
) -> Result<E, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.mutate_structural(key, patch, MutationMode::Update)
}
pub fn update_many_atomic<E>(
&self,
entities: impl IntoIterator<Item = E>,
) -> Result<WriteBatchResponse<E>, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_batch(|save| save.update_many_atomic(entities))
}
pub fn update_many_non_atomic<E>(
&self,
entities: impl IntoIterator<Item = E>,
) -> Result<WriteBatchResponse<E>, InternalError>
where
E: PersistedRow<Canister = C> + EntityValue,
{
self.execute_save_batch(|save| save.update_many_non_atomic(entities))
}
}