use crate::{
db::{
DbSession, PersistedRow, WriteBatchResponse,
data::{FieldSlot, StructuralPatch},
executor::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>(descriptor, 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>(
descriptor: &AcceptedRowLayoutRuntimeDescriptor<'_>,
patch: &StructuralPatch,
) -> bool
where
E: PersistedRow + EntityValue,
{
patch.entries().iter().any(|entry| {
let slot = entry.slot().index();
let Some(accepted_field) = descriptor.field_for_slot_index(slot) else {
return false;
};
let Some(generated_field) = E::MODEL.fields().get(slot) else {
return false;
};
generated_field.insert_generation().is_some()
&& accepted_field.name() != descriptor.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::<E>()?;
let descriptor =
AcceptedRowLayoutRuntimeDescriptor::from_accepted_schema(&accepted_schema)?;
descriptor.generated_compatible_row_shape_for_model(E::MODEL)?;
validate_structural_patch_absence_policy::<E>(&descriptor, &patch, mode)?;
self.execute_save_with_checked_accepted_schema(
|save| save.apply_structural_mutation(mode, key, patch),
std::convert::identity,
)
}
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::<E>()?;
let descriptor =
AcceptedRowLayoutRuntimeDescriptor::from_accepted_schema(&accepted_schema)?;
descriptor.generated_compatible_row_shape_for_model(E::MODEL)?;
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))
}
}