1use crate::{
8 db::{
9 DbSession, PersistedRow, WriteBatchResponse,
10 data::{FieldSlot, StructuralPatch},
11 executor::{EntityAuthority, 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>(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>(patch: &StructuralPatch) -> bool
93where
94 E: PersistedRow + EntityValue,
95{
96 patch.entries().iter().any(|entry| {
97 E::MODEL
98 .fields()
99 .get(entry.slot().index())
100 .is_some_and(|field| {
101 field.insert_generation().is_some() && field.name() != E::MODEL.primary_key.name()
102 })
103 })
104}
105
106impl<C: CanisterKind> DbSession<C> {
107 pub fn insert<E>(&self, entity: E) -> Result<E, InternalError>
109 where
110 E: PersistedRow<Canister = C> + EntityValue,
111 {
112 self.execute_save_entity(|save| save.insert(entity))
113 }
114
115 pub fn create<I>(&self, input: I) -> Result<I::Entity, InternalError>
117 where
118 I: EntityCreateInput,
119 I::Entity: PersistedRow<Canister = C> + EntityValue,
120 {
121 self.execute_save_entity(|save| save.create(input))
122 }
123
124 pub fn insert_many_atomic<E>(
130 &self,
131 entities: impl IntoIterator<Item = E>,
132 ) -> Result<WriteBatchResponse<E>, InternalError>
133 where
134 E: PersistedRow<Canister = C> + EntityValue,
135 {
136 self.execute_save_batch(|save| save.insert_many_atomic(entities))
137 }
138
139 pub fn insert_many_non_atomic<E>(
143 &self,
144 entities: impl IntoIterator<Item = E>,
145 ) -> Result<WriteBatchResponse<E>, InternalError>
146 where
147 E: PersistedRow<Canister = C> + EntityValue,
148 {
149 self.execute_save_batch(|save| save.insert_many_non_atomic(entities))
150 }
151
152 pub fn replace<E>(&self, entity: E) -> Result<E, InternalError>
154 where
155 E: PersistedRow<Canister = C> + EntityValue,
156 {
157 self.execute_save_entity(|save| save.replace(entity))
158 }
159
160 pub fn mutate_structural<E>(
166 &self,
167 key: E::Key,
168 patch: StructuralPatch,
169 mode: MutationMode,
170 ) -> Result<E, InternalError>
171 where
172 E: PersistedRow<Canister = C> + EntityValue,
173 {
174 let (accepted_schema, _) =
175 self.ensure_accepted_schema_snapshot_and_authority(EntityAuthority::for_type::<E>())?;
176 let descriptor =
177 AcceptedRowLayoutRuntimeDescriptor::from_accepted_schema(&accepted_schema)?;
178 validate_structural_patch_absence_policy::<E>(&descriptor, &patch, mode)?;
179
180 self.execute_save_entity(|save| save.apply_structural_mutation(mode, key, patch))
181 }
182
183 pub fn structural_patch<E, I, S>(&self, fields: I) -> Result<StructuralPatch, InternalError>
190 where
191 E: PersistedRow<Canister = C> + EntityValue,
192 I: IntoIterator<Item = (S, Value)>,
193 S: AsRef<str>,
194 {
195 let (accepted_schema, _) =
196 self.ensure_accepted_schema_snapshot_and_authority(EntityAuthority::for_type::<E>())?;
197 let descriptor =
198 AcceptedRowLayoutRuntimeDescriptor::from_accepted_schema(&accepted_schema)?;
199 let mut patch = StructuralPatch::new();
200
201 for (field_name, value) in fields {
205 let field_name = field_name.as_ref();
206 patch = append_accepted_structural_patch_field(
207 E::PATH,
208 &descriptor,
209 patch,
210 field_name,
211 value,
212 )?;
213 }
214
215 Ok(patch)
216 }
217
218 #[cfg(test)]
224 pub(in crate::db) fn replace_structural<E>(
225 &self,
226 key: E::Key,
227 patch: StructuralPatch,
228 ) -> Result<E, InternalError>
229 where
230 E: PersistedRow<Canister = C> + EntityValue,
231 {
232 self.mutate_structural(key, patch, MutationMode::Replace)
233 }
234
235 pub fn replace_many_atomic<E>(
241 &self,
242 entities: impl IntoIterator<Item = E>,
243 ) -> Result<WriteBatchResponse<E>, InternalError>
244 where
245 E: PersistedRow<Canister = C> + EntityValue,
246 {
247 self.execute_save_batch(|save| save.replace_many_atomic(entities))
248 }
249
250 pub fn replace_many_non_atomic<E>(
254 &self,
255 entities: impl IntoIterator<Item = E>,
256 ) -> Result<WriteBatchResponse<E>, InternalError>
257 where
258 E: PersistedRow<Canister = C> + EntityValue,
259 {
260 self.execute_save_batch(|save| save.replace_many_non_atomic(entities))
261 }
262
263 pub fn update<E>(&self, entity: E) -> Result<E, InternalError>
265 where
266 E: PersistedRow<Canister = C> + EntityValue,
267 {
268 self.execute_save_entity(|save| save.update(entity))
269 }
270
271 #[cfg(test)]
277 pub(in crate::db) fn insert_structural<E>(
278 &self,
279 key: E::Key,
280 patch: StructuralPatch,
281 ) -> Result<E, InternalError>
282 where
283 E: PersistedRow<Canister = C> + EntityValue,
284 {
285 self.mutate_structural(key, patch, MutationMode::Insert)
286 }
287
288 #[cfg(test)]
294 pub(in crate::db) fn update_structural<E>(
295 &self,
296 key: E::Key,
297 patch: StructuralPatch,
298 ) -> Result<E, InternalError>
299 where
300 E: PersistedRow<Canister = C> + EntityValue,
301 {
302 self.mutate_structural(key, patch, MutationMode::Update)
303 }
304
305 pub fn update_many_atomic<E>(
311 &self,
312 entities: impl IntoIterator<Item = E>,
313 ) -> Result<WriteBatchResponse<E>, InternalError>
314 where
315 E: PersistedRow<Canister = C> + EntityValue,
316 {
317 self.execute_save_batch(|save| save.update_many_atomic(entities))
318 }
319
320 pub fn update_many_non_atomic<E>(
324 &self,
325 entities: impl IntoIterator<Item = E>,
326 ) -> Result<WriteBatchResponse<E>, InternalError>
327 where
328 E: PersistedRow<Canister = C> + EntityValue,
329 {
330 self.execute_save_batch(|save| save.update_many_non_atomic(entities))
331 }
332}