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_absence_policy<E>(
43 descriptor: &AcceptedRowLayoutRuntimeDescriptor<'_>,
44 patch: &StructuralPatch,
45 mode: MutationMode,
46) -> Result<(), InternalError>
47where
48 E: PersistedRow + EntityValue,
49{
50 if matches!(mode, MutationMode::Update) {
51 return Ok(());
52 }
53 if patch_explicitly_writes_generated_field::<E>(descriptor, patch) {
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 patch_explicitly_writes_generated_field<E>(
93 descriptor: &AcceptedRowLayoutRuntimeDescriptor<'_>,
94 patch: &StructuralPatch,
95) -> bool
96where
97 E: PersistedRow + EntityValue,
98{
99 patch.entries().iter().any(|entry| {
100 let slot = entry.slot().index();
101 let Some(accepted_field) = descriptor.field_for_slot_index(slot) else {
102 return false;
103 };
104 let Some(generated_field) = E::MODEL.fields().get(slot) else {
105 return false;
106 };
107
108 generated_field.insert_generation().is_some()
109 && accepted_field.name() != descriptor.primary_key_name()
110 })
111}
112
113impl<C: CanisterKind> DbSession<C> {
114 pub fn insert<E>(&self, entity: E) -> Result<E, InternalError>
116 where
117 E: PersistedRow<Canister = C> + EntityValue,
118 {
119 self.execute_save_entity(|save| save.insert(entity))
120 }
121
122 pub fn create<I>(&self, input: I) -> Result<I::Entity, InternalError>
124 where
125 I: EntityCreateInput,
126 I::Entity: PersistedRow<Canister = C> + EntityValue,
127 {
128 self.execute_save_entity(|save| save.create(input))
129 }
130
131 pub fn insert_many_atomic<E>(
137 &self,
138 entities: impl IntoIterator<Item = E>,
139 ) -> Result<WriteBatchResponse<E>, InternalError>
140 where
141 E: PersistedRow<Canister = C> + EntityValue,
142 {
143 self.execute_save_batch(|save| save.insert_many_atomic(entities))
144 }
145
146 pub fn insert_many_non_atomic<E>(
150 &self,
151 entities: impl IntoIterator<Item = E>,
152 ) -> Result<WriteBatchResponse<E>, InternalError>
153 where
154 E: PersistedRow<Canister = C> + EntityValue,
155 {
156 self.execute_save_batch(|save| save.insert_many_non_atomic(entities))
157 }
158
159 pub fn replace<E>(&self, entity: E) -> Result<E, InternalError>
161 where
162 E: PersistedRow<Canister = C> + EntityValue,
163 {
164 self.execute_save_entity(|save| save.replace(entity))
165 }
166
167 pub fn mutate_structural<E>(
173 &self,
174 key: E::Key,
175 patch: StructuralPatch,
176 mode: MutationMode,
177 ) -> Result<E, InternalError>
178 where
179 E: PersistedRow<Canister = C> + EntityValue,
180 {
181 let accepted_schema = self.ensure_accepted_schema_snapshot::<E>()?;
182 let descriptor =
183 AcceptedRowLayoutRuntimeDescriptor::from_generated_compatible_accepted_schema(
184 &accepted_schema,
185 E::MODEL,
186 )?;
187 validate_structural_patch_absence_policy::<E>(&descriptor, &patch, mode)?;
188
189 self.execute_save_with_checked_accepted_schema(
190 |save| save.apply_structural_mutation(mode, key, patch),
191 std::convert::identity,
192 )
193 }
194
195 pub fn structural_patch<E, I, S>(&self, fields: I) -> Result<StructuralPatch, InternalError>
202 where
203 E: PersistedRow<Canister = C> + EntityValue,
204 I: IntoIterator<Item = (S, Value)>,
205 S: AsRef<str>,
206 {
207 let accepted_schema = self.ensure_accepted_schema_snapshot::<E>()?;
208 let descriptor =
209 AcceptedRowLayoutRuntimeDescriptor::from_generated_compatible_accepted_schema(
210 &accepted_schema,
211 E::MODEL,
212 )?;
213 let mut patch = StructuralPatch::new();
214
215 for (field_name, value) in fields {
219 let field_name = field_name.as_ref();
220 patch = append_accepted_structural_patch_field(
221 E::PATH,
222 &descriptor,
223 patch,
224 field_name,
225 value,
226 )?;
227 }
228
229 Ok(patch)
230 }
231
232 #[cfg(test)]
238 pub(in crate::db) fn replace_structural<E>(
239 &self,
240 key: E::Key,
241 patch: StructuralPatch,
242 ) -> Result<E, InternalError>
243 where
244 E: PersistedRow<Canister = C> + EntityValue,
245 {
246 self.mutate_structural(key, patch, MutationMode::Replace)
247 }
248
249 pub fn replace_many_atomic<E>(
255 &self,
256 entities: impl IntoIterator<Item = E>,
257 ) -> Result<WriteBatchResponse<E>, InternalError>
258 where
259 E: PersistedRow<Canister = C> + EntityValue,
260 {
261 self.execute_save_batch(|save| save.replace_many_atomic(entities))
262 }
263
264 pub fn replace_many_non_atomic<E>(
268 &self,
269 entities: impl IntoIterator<Item = E>,
270 ) -> Result<WriteBatchResponse<E>, InternalError>
271 where
272 E: PersistedRow<Canister = C> + EntityValue,
273 {
274 self.execute_save_batch(|save| save.replace_many_non_atomic(entities))
275 }
276
277 pub fn update<E>(&self, entity: E) -> Result<E, InternalError>
279 where
280 E: PersistedRow<Canister = C> + EntityValue,
281 {
282 self.execute_save_entity(|save| save.update(entity))
283 }
284
285 #[cfg(test)]
291 pub(in crate::db) fn insert_structural<E>(
292 &self,
293 key: E::Key,
294 patch: StructuralPatch,
295 ) -> Result<E, InternalError>
296 where
297 E: PersistedRow<Canister = C> + EntityValue,
298 {
299 self.mutate_structural(key, patch, MutationMode::Insert)
300 }
301
302 #[cfg(test)]
308 pub(in crate::db) fn update_structural<E>(
309 &self,
310 key: E::Key,
311 patch: StructuralPatch,
312 ) -> Result<E, InternalError>
313 where
314 E: PersistedRow<Canister = C> + EntityValue,
315 {
316 self.mutate_structural(key, patch, MutationMode::Update)
317 }
318
319 pub fn update_many_atomic<E>(
325 &self,
326 entities: impl IntoIterator<Item = E>,
327 ) -> Result<WriteBatchResponse<E>, InternalError>
328 where
329 E: PersistedRow<Canister = C> + EntityValue,
330 {
331 self.execute_save_batch(|save| save.update_many_atomic(entities))
332 }
333
334 pub fn update_many_non_atomic<E>(
338 &self,
339 entities: impl IntoIterator<Item = E>,
340 ) -> Result<WriteBatchResponse<E>, InternalError>
341 where
342 E: PersistedRow<Canister = C> + EntityValue,
343 {
344 self.execute_save_batch(|save| save.update_many_non_atomic(entities))
345 }
346}