1use super::accepted_save_contract_for_descriptor;
8use crate::{
9 db::{
10 DbSession, PersistedRow, WriteBatchResponse,
11 data::{FieldSlot, StructuralPatch},
12 executor::MutationMode,
13 schema::{AcceptedFieldAbsencePolicy, AcceptedRowLayoutRuntimeContract},
14 },
15 error::InternalError,
16 traits::{CanisterKind, EntityCreateInput, EntityValue},
17 value::Value,
18};
19
20fn append_accepted_structural_patch_field(
25 entity_path: &'static str,
26 descriptor: &AcceptedRowLayoutRuntimeContract<'_>,
27 patch: StructuralPatch,
28 field_name: &str,
29 value: Value,
30) -> Result<StructuralPatch, InternalError> {
31 let slot = descriptor
32 .field_slot_index_by_name(field_name)
33 .ok_or_else(|| InternalError::mutation_structural_field_unknown(entity_path, field_name))?;
34
35 Ok(patch.set(FieldSlot::from_validated_index(slot), value))
36}
37
38fn validate_structural_patch_schema_policy<E>(
44 descriptor: &AcceptedRowLayoutRuntimeContract<'_>,
45 patch: &StructuralPatch,
46 mode: MutationMode,
47) -> Result<(), InternalError>
48where
49 E: PersistedRow + EntityValue,
50{
51 reject_explicit_generated_fields_from_accepted_patch::<E>(descriptor, patch)?;
52
53 if matches!(mode, MutationMode::Update) {
54 return Ok(());
55 }
56
57 let mut provided_slots = vec![false; descriptor.required_slot_count()];
58 for entry in patch.entries() {
59 let slot = entry.slot().index();
60 if slot < provided_slots.len() {
61 provided_slots[slot] = true;
62 }
63 }
64
65 for field in descriptor.fields() {
69 let slot = usize::from(field.slot().get());
70 if provided_slots.get(slot).copied().unwrap_or(false) {
71 continue;
72 }
73
74 if matches!(field.absence_policy(), AcceptedFieldAbsencePolicy::Required) {
75 return Err(
76 InternalError::mutation_structural_patch_required_field_missing(
77 E::PATH,
78 field.name(),
79 ),
80 );
81 }
82 }
83
84 Ok(())
85}
86
87fn reject_explicit_generated_fields_from_accepted_patch<E>(
93 descriptor: &AcceptedRowLayoutRuntimeContract<'_>,
94 patch: &StructuralPatch,
95) -> Result<(), InternalError>
96where
97 E: PersistedRow + EntityValue,
98{
99 for entry in patch.entries() {
100 let slot = entry.slot().index();
101 let Some(accepted_field) = descriptor.field_for_slot_index(slot) else {
102 continue;
103 };
104 let write_policy = accepted_field.write_policy();
105
106 if write_policy.insert_generation().is_some()
107 && !descriptor.is_primary_key_field_name(accepted_field.name())
108 {
109 return Err(InternalError::mutation_generated_field_explicit(
110 E::PATH,
111 accepted_field.name(),
112 ));
113 }
114 }
115
116 Ok(())
117}
118
119impl<C: CanisterKind> DbSession<C> {
120 pub fn insert<E>(&self, entity: E) -> Result<E, InternalError>
122 where
123 E: PersistedRow<Canister = C> + EntityValue,
124 {
125 self.execute_save_entity(|save| save.insert(entity))
126 }
127
128 pub fn create<I>(&self, input: I) -> Result<I::Entity, InternalError>
130 where
131 I: EntityCreateInput,
132 I::Entity: PersistedRow<Canister = C> + EntityValue,
133 {
134 self.execute_save_entity(|save| save.create(input))
135 }
136
137 pub fn insert_many_atomic<E>(
145 &self,
146 entities: impl IntoIterator<Item = E>,
147 ) -> Result<WriteBatchResponse<E>, InternalError>
148 where
149 E: PersistedRow<Canister = C> + EntityValue,
150 {
151 self.execute_save_batch(|save| save.insert_many_atomic(entities))
152 }
153
154 pub fn insert_many_non_atomic<E>(
161 &self,
162 entities: impl IntoIterator<Item = E>,
163 ) -> Result<WriteBatchResponse<E>, InternalError>
164 where
165 E: PersistedRow<Canister = C> + EntityValue,
166 {
167 self.execute_save_batch(|save| save.insert_many_non_atomic(entities))
168 }
169
170 pub fn replace<E>(&self, entity: E) -> Result<E, InternalError>
172 where
173 E: PersistedRow<Canister = C> + EntityValue,
174 {
175 self.execute_save_entity(|save| save.replace(entity))
176 }
177
178 pub fn mutate_structural<E>(
184 &self,
185 key: E::Key,
186 patch: StructuralPatch,
187 mode: MutationMode,
188 ) -> Result<E, InternalError>
189 where
190 E: PersistedRow<Canister = C> + EntityValue,
191 {
192 let accepted_schema = self.ensure_accepted_schema_snapshot::<E>()?;
193 let (descriptor, _) = AcceptedRowLayoutRuntimeContract::from_generated_compatible_schema(
194 &accepted_schema,
195 E::MODEL,
196 )?;
197 validate_structural_patch_schema_policy::<E>(&descriptor, &patch, mode)?;
198 let (
199 row_decode_contract,
200 mutation_row_decode_contract,
201 accepted_schema_info,
202 accepted_schema_fingerprint,
203 ) = accepted_save_contract_for_descriptor::<E>(&accepted_schema, &descriptor)?;
204
205 self.execute_save_with_checked_accepted_row_contract(
206 row_decode_contract,
207 accepted_schema_info,
208 accepted_schema_fingerprint,
209 |save| save.apply_structural_mutation(mode, key, patch, mutation_row_decode_contract),
210 std::convert::identity,
211 )
212 }
213
214 pub fn structural_patch<E, I, S>(&self, fields: I) -> Result<StructuralPatch, InternalError>
221 where
222 E: PersistedRow<Canister = C> + EntityValue,
223 I: IntoIterator<Item = (S, Value)>,
224 S: AsRef<str>,
225 {
226 let accepted_schema = self.ensure_accepted_schema_snapshot::<E>()?;
227 let (descriptor, _) = AcceptedRowLayoutRuntimeContract::from_generated_compatible_schema(
228 &accepted_schema,
229 E::MODEL,
230 )?;
231 let mut patch = StructuralPatch::new();
232
233 for (field_name, value) in fields {
237 let field_name = field_name.as_ref();
238 patch = append_accepted_structural_patch_field(
239 E::PATH,
240 &descriptor,
241 patch,
242 field_name,
243 value,
244 )?;
245 }
246
247 Ok(patch)
248 }
249
250 #[cfg(test)]
256 pub(in crate::db) fn replace_structural<E>(
257 &self,
258 key: E::Key,
259 patch: StructuralPatch,
260 ) -> Result<E, InternalError>
261 where
262 E: PersistedRow<Canister = C> + EntityValue,
263 {
264 self.mutate_structural(key, patch, MutationMode::Replace)
265 }
266
267 pub fn replace_many_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_atomic(entities))
282 }
283
284 pub fn replace_many_non_atomic<E>(
291 &self,
292 entities: impl IntoIterator<Item = E>,
293 ) -> Result<WriteBatchResponse<E>, InternalError>
294 where
295 E: PersistedRow<Canister = C> + EntityValue,
296 {
297 self.execute_save_batch(|save| save.replace_many_non_atomic(entities))
298 }
299
300 pub fn update<E>(&self, entity: E) -> Result<E, InternalError>
302 where
303 E: PersistedRow<Canister = C> + EntityValue,
304 {
305 self.execute_save_entity(|save| save.update(entity))
306 }
307
308 #[cfg(test)]
314 pub(in crate::db) fn insert_structural<E>(
315 &self,
316 key: E::Key,
317 patch: StructuralPatch,
318 ) -> Result<E, InternalError>
319 where
320 E: PersistedRow<Canister = C> + EntityValue,
321 {
322 self.mutate_structural(key, patch, MutationMode::Insert)
323 }
324
325 #[cfg(test)]
331 pub(in crate::db) fn update_structural<E>(
332 &self,
333 key: E::Key,
334 patch: StructuralPatch,
335 ) -> Result<E, InternalError>
336 where
337 E: PersistedRow<Canister = C> + EntityValue,
338 {
339 self.mutate_structural(key, patch, MutationMode::Update)
340 }
341
342 pub fn update_many_atomic<E>(
350 &self,
351 entities: impl IntoIterator<Item = E>,
352 ) -> Result<WriteBatchResponse<E>, InternalError>
353 where
354 E: PersistedRow<Canister = C> + EntityValue,
355 {
356 self.execute_save_batch(|save| save.update_many_atomic(entities))
357 }
358
359 pub fn update_many_non_atomic<E>(
366 &self,
367 entities: impl IntoIterator<Item = E>,
368 ) -> Result<WriteBatchResponse<E>, InternalError>
369 where
370 E: PersistedRow<Canister = C> + EntityValue,
371 {
372 self.execute_save_batch(|save| save.update_many_non_atomic(entities))
373 }
374}