1use reifydb_codec::row::{
5 bytes::{EncodedBytes, RowBuilder},
6 shape::{RowFamily, RowShape},
7 table::EncodedTableRowBuilder,
8};
9use reifydb_core::{
10 common::CommitVersion,
11 interface::{
12 catalog::{object::ObjectId, table::Table},
13 change::{Change, ChangeOrigin, Diff},
14 },
15 partition::PartitionError,
16 row::row_shape_from_columns,
17 value::column::columns::Columns,
18};
19use reifydb_transaction::{
20 change::{RowChange, TableRowInsertion},
21 interceptor::table_row::TableRowInterceptor,
22 transaction::{Transaction, admin::AdminTransaction, command::CommandTransaction},
23};
24use reifydb_value::value::{datetime::DateTime, partition::Partition, row_number::RowNumber};
25use smallvec::smallvec;
26
27use crate::{
28 Result,
29 partition::{row_key_from_partition, table_partition_of_row, table_row_key},
30};
31
32fn build_table_insert_change(
33 table: &Table,
34 shape: &RowShape,
35 ids: &[RowNumber],
36 bytes_slice: &[EncodedBytes],
37) -> Change {
38 Change {
39 origin: ChangeOrigin::Object(ObjectId::Table(table.id)),
40 version: CommitVersion(0),
41 diffs: smallvec![Diff::insert(Columns::from_encoded_bytes(shape, ids, bytes_slice))],
42 changed_at: DateTime::default(),
43 }
44}
45
46fn build_table_update_change(
47 table: &Table,
48 shape: &RowShape,
49 ids: &[RowNumber],
50 pres: &[EncodedBytes],
51 posts: &[EncodedBytes],
52) -> Change {
53 Change {
54 origin: ChangeOrigin::Object(ObjectId::Table(table.id)),
55 version: CommitVersion(0),
56 diffs: smallvec![Diff::update(
57 Columns::from_encoded_bytes(shape, ids, pres),
58 Columns::from_encoded_bytes(shape, ids, posts),
59 )],
60 changed_at: DateTime::default(),
61 }
62}
63
64fn build_table_remove_change(
65 table: &Table,
66 shape: &RowShape,
67 ids: &[RowNumber],
68 bytes_slice: &[EncodedBytes],
69) -> Change {
70 Change {
71 origin: ChangeOrigin::Object(ObjectId::Table(table.id)),
72 version: CommitVersion(0),
73 diffs: smallvec![Diff::remove(Columns::from_encoded_bytes(shape, ids, bytes_slice))],
74 changed_at: DateTime::default(),
75 }
76}
77
78pub trait TableOperations {
79 fn insert_table(
80 &mut self,
81 table: &Table,
82 shape: &RowShape,
83 ids: &[RowNumber],
84 rows: &mut [EncodedTableRowBuilder],
85 ) -> Result<()>;
86
87 fn update_table(
88 &mut self,
89 table: &Table,
90 ids: &[RowNumber],
91 partitions: &[Partition],
92 rows: &mut [EncodedTableRowBuilder],
93 ) -> Result<Vec<(RowNumber, EncodedBytes)>>;
94
95 fn remove_from_table(
96 &mut self,
97 table: &Table,
98 ids: &[RowNumber],
99 partitions: &[Partition],
100 ) -> Result<Vec<(RowNumber, EncodedBytes)>>;
101}
102
103impl TableOperations for CommandTransaction {
104 fn insert_table(
105 &mut self,
106 table: &Table,
107 shape: &RowShape,
108 ids: &[RowNumber],
109 rows: &mut [EncodedTableRowBuilder],
110 ) -> Result<()> {
111 assert_eq!(ids.len(), rows.len(), "ids/rows length mismatch");
112 if ids.is_empty() {
113 return Ok(());
114 }
115
116 TableRowInterceptor::pre_insert(self, table, ids, rows)?;
117
118 let frozen: Vec<EncodedBytes> = rows.iter().map(|row| row.clone().freeze_bytes()).collect();
119
120 for (row, &row_number) in frozen.iter().zip(ids.iter()) {
121 self.set(&table_row_key(table, shape, row, row_number), row.clone())?;
122 }
123
124 TableRowInterceptor::post_insert(self, table, ids, &frozen)?;
125
126 let row_changes: Vec<RowChange> = ids
127 .iter()
128 .zip(frozen.iter())
129 .map(|(&row_number, row)| {
130 RowChange::TableInsert(TableRowInsertion {
131 table_id: table.id,
132 row_number,
133 encoded: row.clone(),
134 })
135 })
136 .collect();
137 self.track_row_change(&row_changes);
138
139 self.track_flow_change(build_table_insert_change(table, shape, ids, &frozen));
140
141 Ok(())
142 }
143
144 fn update_table(
145 &mut self,
146 table: &Table,
147 ids: &[RowNumber],
148 partitions: &[Partition],
149 rows: &mut [EncodedTableRowBuilder],
150 ) -> Result<Vec<(RowNumber, EncodedBytes)>> {
151 assert_eq!(ids.len(), rows.len(), "ids/rows length mismatch");
152 if ids.is_empty() {
153 return Ok(Vec::new());
154 }
155
156 let shape = row_shape_from_columns(RowFamily::Table, &table.columns);
157
158 TableRowInterceptor::pre_update(self, table, ids, rows)?;
159
160 if !table.partition_by.is_empty() {
161 for (idx, row) in rows.iter().enumerate() {
162 if let Some(&expected) = partitions.get(idx)
163 && table_partition_of_row(table, &shape, row) != expected
164 {
165 return Err(PartitionError::ImmutablePartitionColumn {
166 object: ObjectId::Table(table.id),
167 }
168 .into());
169 }
170 }
171 }
172
173 let mut matched_indices: Vec<usize> = Vec::with_capacity(ids.len());
174 let mut pres: Vec<EncodedBytes> = Vec::with_capacity(ids.len());
175 for (idx, &row_number) in ids.iter().enumerate() {
176 let key = row_key_from_partition(table.id, partitions.get(idx).copied(), row_number);
177 let pre = match self.get(&key)? {
178 Some(v) => v.bytes,
179 None => continue,
180 };
181 if self.get_committed(&key)?.is_some() {
182 self.mark_preexisting(&key)?;
183 }
184 self.set(&key, rows[idx].clone().freeze())?;
185 matched_indices.push(idx);
186 pres.push(pre);
187 }
188
189 if matched_indices.is_empty() {
190 return Ok(Vec::new());
191 }
192
193 let matched_ids: Vec<RowNumber> = matched_indices.iter().map(|&i| ids[i]).collect();
194 let matched_posts: Vec<EncodedBytes> =
195 matched_indices.iter().map(|&i| rows[i].clone().freeze_bytes()).collect();
196
197 TableRowInterceptor::post_update(self, table, &matched_ids, &matched_posts, &pres)?;
198
199 self.track_flow_change(build_table_update_change(table, &shape, &matched_ids, &pres, &matched_posts));
200
201 Ok(matched_ids.into_iter().zip(matched_posts).collect())
202 }
203
204 fn remove_from_table(
205 &mut self,
206 table: &Table,
207 ids: &[RowNumber],
208 partitions: &[Partition],
209 ) -> Result<Vec<(RowNumber, EncodedBytes)>> {
210 if ids.is_empty() {
211 return Ok(Vec::new());
212 }
213
214 let mut matched_ids: Vec<RowNumber> = Vec::with_capacity(ids.len());
215 let mut matched_partitions: Vec<Option<Partition>> = Vec::with_capacity(ids.len());
216 let mut displayed_bytes_vec: Vec<EncodedBytes> = Vec::with_capacity(ids.len());
217 let mut pre_for_cdc_bytes_vec: Vec<EncodedBytes> = Vec::with_capacity(ids.len());
218 for (idx, &row_number) in ids.iter().enumerate() {
219 let partition = partitions.get(idx).copied();
220 let key = row_key_from_partition(table.id, partition, row_number);
221 let displayed = match self.get(&key)? {
222 Some(v) => v.bytes,
223 None => continue,
224 };
225 let committed = self.get_committed(&key)?.map(|v| v.bytes);
226 let pre_for_cdc = committed.clone().unwrap_or_else(|| displayed.clone());
227 matched_ids.push(row_number);
228 matched_partitions.push(partition);
229 displayed_bytes_vec.push(displayed);
230 pre_for_cdc_bytes_vec.push(pre_for_cdc);
231 }
232
233 if matched_ids.is_empty() {
234 return Ok(Vec::new());
235 }
236
237 TableRowInterceptor::pre_delete(self, table, &matched_ids)?;
238
239 for (i, &row_number) in matched_ids.iter().enumerate() {
240 let key = row_key_from_partition(table.id, matched_partitions[i], row_number);
241 if self.get_committed(&key)?.is_some() {
242 self.mark_preexisting(&key)?;
243 }
244 self.remove_with_pre(&key, pre_for_cdc_bytes_vec[i].clone())?;
245 }
246
247 TableRowInterceptor::post_delete(self, table, &matched_ids, &pre_for_cdc_bytes_vec)?;
248
249 let shape = row_shape_from_columns(RowFamily::Table, &table.columns);
250 self.track_flow_change(build_table_remove_change(table, &shape, &matched_ids, &pre_for_cdc_bytes_vec));
251
252 Ok(matched_ids.into_iter().zip(displayed_bytes_vec).collect())
253 }
254}
255
256impl TableOperations for AdminTransaction {
257 fn insert_table(
258 &mut self,
259 table: &Table,
260 shape: &RowShape,
261 ids: &[RowNumber],
262 rows: &mut [EncodedTableRowBuilder],
263 ) -> Result<()> {
264 assert_eq!(ids.len(), rows.len(), "ids/rows length mismatch");
265 if ids.is_empty() {
266 return Ok(());
267 }
268
269 TableRowInterceptor::pre_insert(self, table, ids, rows)?;
270
271 let frozen: Vec<EncodedBytes> = rows.iter().map(|row| row.clone().freeze_bytes()).collect();
272
273 for (row, &row_number) in frozen.iter().zip(ids.iter()) {
274 self.set(&table_row_key(table, shape, row, row_number), row.clone())?;
275 }
276
277 TableRowInterceptor::post_insert(self, table, ids, &frozen)?;
278
279 let row_changes: Vec<RowChange> = ids
280 .iter()
281 .zip(frozen.iter())
282 .map(|(&row_number, row)| {
283 RowChange::TableInsert(TableRowInsertion {
284 table_id: table.id,
285 row_number,
286 encoded: row.clone(),
287 })
288 })
289 .collect();
290 self.track_row_change(&row_changes);
291
292 self.track_flow_change(build_table_insert_change(table, shape, ids, &frozen));
293
294 Ok(())
295 }
296
297 fn update_table(
298 &mut self,
299 table: &Table,
300 ids: &[RowNumber],
301 partitions: &[Partition],
302 rows: &mut [EncodedTableRowBuilder],
303 ) -> Result<Vec<(RowNumber, EncodedBytes)>> {
304 assert_eq!(ids.len(), rows.len(), "ids/rows length mismatch");
305 if ids.is_empty() {
306 return Ok(Vec::new());
307 }
308
309 let shape = row_shape_from_columns(RowFamily::Table, &table.columns);
310
311 TableRowInterceptor::pre_update(self, table, ids, rows)?;
312
313 if !table.partition_by.is_empty() {
314 for (idx, row) in rows.iter().enumerate() {
315 if let Some(&expected) = partitions.get(idx)
316 && table_partition_of_row(table, &shape, row) != expected
317 {
318 return Err(PartitionError::ImmutablePartitionColumn {
319 object: ObjectId::Table(table.id),
320 }
321 .into());
322 }
323 }
324 }
325
326 let mut matched_indices: Vec<usize> = Vec::with_capacity(ids.len());
327 let mut pres: Vec<EncodedBytes> = Vec::with_capacity(ids.len());
328 for (idx, &row_number) in ids.iter().enumerate() {
329 let key = row_key_from_partition(table.id, partitions.get(idx).copied(), row_number);
330 let pre = match self.get(&key)? {
331 Some(v) => v.bytes,
332 None => continue,
333 };
334 if self.get_committed(&key)?.is_some() {
335 self.mark_preexisting(&key)?;
336 }
337 self.set(&key, rows[idx].clone().freeze())?;
338 matched_indices.push(idx);
339 pres.push(pre);
340 }
341
342 if matched_indices.is_empty() {
343 return Ok(Vec::new());
344 }
345
346 let matched_ids: Vec<RowNumber> = matched_indices.iter().map(|&i| ids[i]).collect();
347 let matched_posts: Vec<EncodedBytes> =
348 matched_indices.iter().map(|&i| rows[i].clone().freeze_bytes()).collect();
349
350 TableRowInterceptor::post_update(self, table, &matched_ids, &matched_posts, &pres)?;
351
352 self.track_flow_change(build_table_update_change(table, &shape, &matched_ids, &pres, &matched_posts));
353
354 Ok(matched_ids.into_iter().zip(matched_posts).collect())
355 }
356
357 fn remove_from_table(
358 &mut self,
359 table: &Table,
360 ids: &[RowNumber],
361 partitions: &[Partition],
362 ) -> Result<Vec<(RowNumber, EncodedBytes)>> {
363 if ids.is_empty() {
364 return Ok(Vec::new());
365 }
366
367 let mut matched_ids: Vec<RowNumber> = Vec::with_capacity(ids.len());
368 let mut matched_partitions: Vec<Option<Partition>> = Vec::with_capacity(ids.len());
369 let mut displayed_bytes_vec: Vec<EncodedBytes> = Vec::with_capacity(ids.len());
370 let mut pre_for_cdc_bytes_vec: Vec<EncodedBytes> = Vec::with_capacity(ids.len());
371 for (idx, &row_number) in ids.iter().enumerate() {
372 let partition = partitions.get(idx).copied();
373 let key = row_key_from_partition(table.id, partition, row_number);
374 let displayed = match self.get(&key)? {
375 Some(v) => v.bytes,
376 None => continue,
377 };
378 let committed = self.get_committed(&key)?.map(|v| v.bytes);
379 let pre_for_cdc = committed.clone().unwrap_or_else(|| displayed.clone());
380 matched_ids.push(row_number);
381 matched_partitions.push(partition);
382 displayed_bytes_vec.push(displayed);
383 pre_for_cdc_bytes_vec.push(pre_for_cdc);
384 }
385
386 if matched_ids.is_empty() {
387 return Ok(Vec::new());
388 }
389
390 TableRowInterceptor::pre_delete(self, table, &matched_ids)?;
391
392 for (i, &row_number) in matched_ids.iter().enumerate() {
393 let key = row_key_from_partition(table.id, matched_partitions[i], row_number);
394 if self.get_committed(&key)?.is_some() {
395 self.mark_preexisting(&key)?;
396 }
397 self.remove_with_pre(&key, pre_for_cdc_bytes_vec[i].clone())?;
398 }
399
400 TableRowInterceptor::post_delete(self, table, &matched_ids, &pre_for_cdc_bytes_vec)?;
401
402 let shape = row_shape_from_columns(RowFamily::Table, &table.columns);
403 self.track_flow_change(build_table_remove_change(table, &shape, &matched_ids, &pre_for_cdc_bytes_vec));
404
405 Ok(matched_ids.into_iter().zip(displayed_bytes_vec).collect())
406 }
407}
408
409impl TableOperations for Transaction<'_> {
410 fn insert_table(
411 &mut self,
412 table: &Table,
413 shape: &RowShape,
414 ids: &[RowNumber],
415 rows: &mut [EncodedTableRowBuilder],
416 ) -> Result<()> {
417 match self {
418 Transaction::Command(txn) => txn.insert_table(table, shape, ids, rows),
419 Transaction::Admin(txn) => txn.insert_table(table, shape, ids, rows),
420 Transaction::Test(t) => t.inner.insert_table(table, shape, ids, rows),
421 Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
422 }
423 }
424
425 fn update_table(
426 &mut self,
427 table: &Table,
428 ids: &[RowNumber],
429 partitions: &[Partition],
430 rows: &mut [EncodedTableRowBuilder],
431 ) -> Result<Vec<(RowNumber, EncodedBytes)>> {
432 match self {
433 Transaction::Command(txn) => txn.update_table(table, ids, partitions, rows),
434 Transaction::Admin(txn) => txn.update_table(table, ids, partitions, rows),
435 Transaction::Test(t) => t.inner.update_table(table, ids, partitions, rows),
436 Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
437 }
438 }
439
440 fn remove_from_table(
441 &mut self,
442 table: &Table,
443 ids: &[RowNumber],
444 partitions: &[Partition],
445 ) -> Result<Vec<(RowNumber, EncodedBytes)>> {
446 match self {
447 Transaction::Command(txn) => txn.remove_from_table(table, ids, partitions),
448 Transaction::Admin(txn) => txn.remove_from_table(table, ids, partitions),
449 Transaction::Test(t) => t.inner.remove_from_table(table, ids, partitions),
450 Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
451 }
452 }
453}