1use crate::{
8 db::{
9 DbSession, PersistedRow, WriteBatchResponse,
10 data::{FieldSlot, StructuralPatch},
11 executor::MutationMode,
12 schema::{AcceptedFieldAbsencePolicy, AcceptedRowLayoutRuntimeDescriptor},
13 },
14 error::InternalError,
15 traits::{CanisterKind, EntityCreateInput, EntityValue},
16 value::Value,
17};
18
19fn append_accepted_structural_patch_field(
24 entity_path: &'static str,
25 descriptor: &AcceptedRowLayoutRuntimeDescriptor<'_>,
26 patch: StructuralPatch,
27 field_name: &str,
28 value: Value,
29) -> Result<StructuralPatch, InternalError> {
30 let slot = descriptor
31 .field_slot_index_by_name(field_name)
32 .ok_or_else(|| InternalError::mutation_structural_field_unknown(entity_path, field_name))?;
33
34 Ok(patch.set(FieldSlot::from_validated_index(slot), value))
35}
36
37fn validate_structural_patch_schema_policy<E>(
43 descriptor: &AcceptedRowLayoutRuntimeDescriptor<'_>,
44 patch: &StructuralPatch,
45 mode: MutationMode,
46) -> Result<(), InternalError>
47where
48 E: PersistedRow + EntityValue,
49{
50 reject_explicit_generated_fields_from_accepted_patch::<E>(descriptor, patch)?;
51
52 if matches!(mode, MutationMode::Update) {
53 return Ok(());
54 }
55
56 let mut provided_slots = vec![false; descriptor.required_slot_count()];
57 for entry in patch.entries() {
58 let slot = entry.slot().index();
59 if slot < provided_slots.len() {
60 provided_slots[slot] = true;
61 }
62 }
63
64 for field in descriptor.fields() {
68 let slot = usize::from(field.slot().get());
69 if provided_slots.get(slot).copied().unwrap_or(false) {
70 continue;
71 }
72
73 if matches!(field.absence_policy(), AcceptedFieldAbsencePolicy::Required) {
74 return Err(
75 InternalError::mutation_structural_patch_required_field_missing(
76 E::PATH,
77 field.name(),
78 ),
79 );
80 }
81 }
82
83 Ok(())
84}
85
86fn reject_explicit_generated_fields_from_accepted_patch<E>(
92 descriptor: &AcceptedRowLayoutRuntimeDescriptor<'_>,
93 patch: &StructuralPatch,
94) -> Result<(), InternalError>
95where
96 E: PersistedRow + EntityValue,
97{
98 for entry in patch.entries() {
99 let slot = entry.slot().index();
100 let Some(accepted_field) = descriptor.field_for_slot_index(slot) else {
101 continue;
102 };
103 let write_policy = accepted_field.write_policy();
104
105 if write_policy.insert_generation().is_some()
106 && accepted_field.name() != descriptor.primary_key_name()
107 {
108 return Err(InternalError::mutation_generated_field_explicit(
109 E::PATH,
110 accepted_field.name(),
111 ));
112 }
113 }
114
115 Ok(())
116}
117
118impl<C: CanisterKind> DbSession<C> {
119 pub fn insert<E>(&self, entity: E) -> Result<E, InternalError>
121 where
122 E: PersistedRow<Canister = C> + EntityValue,
123 {
124 self.execute_save_entity(|save| save.insert(entity))
125 }
126
127 pub fn create<I>(&self, input: I) -> Result<I::Entity, InternalError>
129 where
130 I: EntityCreateInput,
131 I::Entity: PersistedRow<Canister = C> + EntityValue,
132 {
133 self.execute_save_entity(|save| save.create(input))
134 }
135
136 pub fn insert_many_atomic<E>(
142 &self,
143 entities: impl IntoIterator<Item = E>,
144 ) -> Result<WriteBatchResponse<E>, InternalError>
145 where
146 E: PersistedRow<Canister = C> + EntityValue,
147 {
148 self.execute_save_batch(|save| save.insert_many_atomic(entities))
149 }
150
151 pub fn insert_many_non_atomic<E>(
155 &self,
156 entities: impl IntoIterator<Item = E>,
157 ) -> Result<WriteBatchResponse<E>, InternalError>
158 where
159 E: PersistedRow<Canister = C> + EntityValue,
160 {
161 self.execute_save_batch(|save| save.insert_many_non_atomic(entities))
162 }
163
164 pub fn replace<E>(&self, entity: E) -> Result<E, InternalError>
166 where
167 E: PersistedRow<Canister = C> + EntityValue,
168 {
169 self.execute_save_entity(|save| save.replace(entity))
170 }
171
172 pub fn mutate_structural<E>(
178 &self,
179 key: E::Key,
180 patch: StructuralPatch,
181 mode: MutationMode,
182 ) -> Result<E, InternalError>
183 where
184 E: PersistedRow<Canister = C> + EntityValue,
185 {
186 let accepted_schema = self.ensure_accepted_schema_snapshot::<E>()?;
187 let descriptor =
188 AcceptedRowLayoutRuntimeDescriptor::from_accepted_schema(&accepted_schema)?;
189 descriptor.generated_compatible_row_shape_for_model(E::MODEL)?;
190 validate_structural_patch_schema_policy::<E>(&descriptor, &patch, mode)?;
191
192 let row_decode_contract = descriptor.row_decode_contract();
193 let mutation_row_decode_contract = row_decode_contract.clone();
194
195 self.execute_save_with_checked_accepted_row_contract(
196 row_decode_contract,
197 |save| {
198 save.apply_structural_mutation(mode, key, patch, Some(mutation_row_decode_contract))
199 },
200 std::convert::identity,
201 )
202 }
203
204 pub fn structural_patch<E, I, S>(&self, fields: I) -> Result<StructuralPatch, InternalError>
211 where
212 E: PersistedRow<Canister = C> + EntityValue,
213 I: IntoIterator<Item = (S, Value)>,
214 S: AsRef<str>,
215 {
216 let accepted_schema = self.ensure_accepted_schema_snapshot::<E>()?;
217 let descriptor =
218 AcceptedRowLayoutRuntimeDescriptor::from_accepted_schema(&accepted_schema)?;
219 descriptor.generated_compatible_row_shape_for_model(E::MODEL)?;
220 let mut patch = StructuralPatch::new();
221
222 for (field_name, value) in fields {
226 let field_name = field_name.as_ref();
227 patch = append_accepted_structural_patch_field(
228 E::PATH,
229 &descriptor,
230 patch,
231 field_name,
232 value,
233 )?;
234 }
235
236 Ok(patch)
237 }
238
239 #[cfg(test)]
245 pub(in crate::db) fn replace_structural<E>(
246 &self,
247 key: E::Key,
248 patch: StructuralPatch,
249 ) -> Result<E, InternalError>
250 where
251 E: PersistedRow<Canister = C> + EntityValue,
252 {
253 self.mutate_structural(key, patch, MutationMode::Replace)
254 }
255
256 pub fn replace_many_atomic<E>(
262 &self,
263 entities: impl IntoIterator<Item = E>,
264 ) -> Result<WriteBatchResponse<E>, InternalError>
265 where
266 E: PersistedRow<Canister = C> + EntityValue,
267 {
268 self.execute_save_batch(|save| save.replace_many_atomic(entities))
269 }
270
271 pub fn replace_many_non_atomic<E>(
275 &self,
276 entities: impl IntoIterator<Item = E>,
277 ) -> Result<WriteBatchResponse<E>, InternalError>
278 where
279 E: PersistedRow<Canister = C> + EntityValue,
280 {
281 self.execute_save_batch(|save| save.replace_many_non_atomic(entities))
282 }
283
284 pub fn update<E>(&self, entity: E) -> Result<E, InternalError>
286 where
287 E: PersistedRow<Canister = C> + EntityValue,
288 {
289 self.execute_save_entity(|save| save.update(entity))
290 }
291
292 #[cfg(test)]
298 pub(in crate::db) fn insert_structural<E>(
299 &self,
300 key: E::Key,
301 patch: StructuralPatch,
302 ) -> Result<E, InternalError>
303 where
304 E: PersistedRow<Canister = C> + EntityValue,
305 {
306 self.mutate_structural(key, patch, MutationMode::Insert)
307 }
308
309 #[cfg(test)]
315 pub(in crate::db) fn update_structural<E>(
316 &self,
317 key: E::Key,
318 patch: StructuralPatch,
319 ) -> Result<E, InternalError>
320 where
321 E: PersistedRow<Canister = C> + EntityValue,
322 {
323 self.mutate_structural(key, patch, MutationMode::Update)
324 }
325
326 pub fn update_many_atomic<E>(
332 &self,
333 entities: impl IntoIterator<Item = E>,
334 ) -> Result<WriteBatchResponse<E>, InternalError>
335 where
336 E: PersistedRow<Canister = C> + EntityValue,
337 {
338 self.execute_save_batch(|save| save.update_many_atomic(entities))
339 }
340
341 pub fn update_many_non_atomic<E>(
345 &self,
346 entities: impl IntoIterator<Item = E>,
347 ) -> Result<WriteBatchResponse<E>, InternalError>
348 where
349 E: PersistedRow<Canister = C> + EntityValue,
350 {
351 self.execute_save_batch(|save| save.update_many_non_atomic(entities))
352 }
353}