1use reifydb_catalog::catalog::Catalog;
5use reifydb_codec::row::{
6 bytes::{EncodedBytes, RowBuilder},
7 ringbuffer::EncodedRingBufferRow,
8 shape::{RowFamily, RowShape},
9};
10use reifydb_core::{
11 common::{ChangeVersion, CommitVersion},
12 delta::SetPre,
13 interface::{
14 catalog::{
15 object::ObjectId,
16 ringbuffer::{RingBuffer, RingBufferMetadata},
17 },
18 change::{Change, ChangeOrigin, Diff},
19 },
20 key::{
21 any::TaggedKey,
22 row::{PartitionedRowKey, RowKey},
23 },
24 partition::{PartitionError, partition_col_indices, partitions_of_rows},
25 row::row_shape_from_columns,
26 value::batch::from_encoded_bytes,
27};
28use reifydb_transaction::{
29 interceptor::{WithInterceptors, ringbuffer_row::RingBufferRowInterceptor},
30 transaction::{Transaction, admin::AdminTransaction, command::CommandTransaction, read::Read, write::Write},
31};
32use reifydb_value::{
33 util::cowvec::CowVec,
34 value::{Value, datetime::DateTime, partition::Partition, row_number::RowNumber},
35};
36use smallvec::smallvec;
37
38use crate::Result;
39
40pub(crate) fn ringbuffer_key(
41 ringbuffer: &RingBuffer,
42 partition: Option<Partition>,
43 row_number: RowNumber,
44) -> TaggedKey {
45 match partition {
46 None => RowKey::new(ringbuffer.id, row_number).into(),
47 Some(partition) => PartitionedRowKey::new(ringbuffer.id, partition, row_number).into(),
48 }
49}
50
51fn moved_partitions(
52 ringbuffer: &RingBuffer,
53 shape: &RowShape,
54 partitions: &[Partition],
55 rows: &[EncodedBytes],
56) -> Result<bool> {
57 let indices = partition_col_indices(&ringbuffer.columns, &ringbuffer.partition_by);
58 Ok(!indices.is_empty() && partitions_of_rows(shape, &indices, rows)? != partitions)
59}
60
61fn ringbuffer_change(ringbuffer: &RingBuffer, diff: Diff) -> Change {
62 Change {
63 origin: ChangeOrigin::Object(ObjectId::ringbuffer(ringbuffer.id)),
64 version: ChangeVersion::from(CommitVersion(0)),
65 diffs: smallvec![diff],
66 changed_at: DateTime::default(),
67 }
68}
69
70fn build_ringbuffer_insert_change(
71 ringbuffer: &RingBuffer,
72 shape: &RowShape,
73 ids: &[RowNumber],
74 rows: &[EncodedBytes],
75) -> Result<Change> {
76 Ok(ringbuffer_change(ringbuffer, Diff::insert(from_encoded_bytes(shape, ids, rows)?)))
77}
78
79fn build_ringbuffer_update_change(
80 ringbuffer: &RingBuffer,
81 shape: &RowShape,
82 ids: &[RowNumber],
83 pres: &[EncodedBytes],
84 posts: &[EncodedBytes],
85) -> Result<Change> {
86 Ok(ringbuffer_change(
87 ringbuffer,
88 Diff::update(from_encoded_bytes(shape, ids, pres)?, from_encoded_bytes(shape, ids, posts)?),
89 ))
90}
91
92fn build_ringbuffer_remove_change(
93 ringbuffer: &RingBuffer,
94 shape: &RowShape,
95 ids: &[RowNumber],
96 rows: &[EncodedBytes],
97) -> Result<Change> {
98 Ok(ringbuffer_change(ringbuffer, Diff::remove(from_encoded_bytes(shape, ids, rows)?)))
99}
100
101pub fn apply_ringbuffer_partition_metadata_after_delete(
102 catalog: &Catalog,
103 txn: &mut Transaction<'_>,
104 ringbuffer: &RingBuffer,
105 partition_key: &[Value],
106 mut partition: RingBufferMetadata,
107 deleted: u64,
108 min_remaining_row: Option<u64>,
109) -> Result<()> {
110 if deleted == 0 {
111 return Ok(());
112 }
113 let remaining_count = partition.count.saturating_sub(deleted);
114 if remaining_count == 0 {
115 catalog.remove_partition_metadata(txn, ringbuffer, partition_key)
116 } else {
117 partition.count = remaining_count;
118 partition.head = min_remaining_row.unwrap();
119 catalog.save_partition_metadata(txn, ringbuffer, partition_key, &partition)
120 }
121}
122
123pub trait RingBufferOperations {
124 fn insert_ringbuffer(
125 &mut self,
126 ringbuffer: &RingBuffer,
127 shape: &RowShape,
128 partitions: &[Partition],
129 ids: &[RowNumber],
130 rows: &[EncodedBytes],
131 ) -> Result<Vec<EncodedBytes>>;
132
133 fn update_ringbuffer(
134 &mut self,
135 ringbuffer: &RingBuffer,
136 partitions: &[Partition],
137 ids: &[RowNumber],
138 rows: &[EncodedBytes],
139 pres: &[EncodedBytes],
140 ) -> Result<Vec<EncodedBytes>>;
141
142 fn remove_from_ringbuffer(
143 &mut self,
144 ringbuffer: &RingBuffer,
145 partitions: &[Partition],
146 ids: &[RowNumber],
147 displayed: Vec<Option<EncodedBytes>>,
148 ) -> Result<Vec<EncodedBytes>>;
149}
150
151fn insert_rows<T: WithInterceptors + Write + Read>(
152 txn: &mut T,
153 ringbuffer: &RingBuffer,
154 shape: &RowShape,
155 partitions: &[Partition],
156 ids: &[RowNumber],
157 rows: &[EncodedBytes],
158) -> Result<Vec<EncodedBytes>> {
159 assert_eq!(ids.len(), rows.len(), "ids/rows length mismatch");
160 if ids.is_empty() {
161 return Ok(Vec::new());
162 }
163
164 let keys: Vec<TaggedKey> = ids
165 .iter()
166 .enumerate()
167 .map(|(idx, &row_number)| ringbuffer_key(ringbuffer, partitions.get(idx).copied(), row_number))
168 .collect();
169
170 let written: Vec<EncodedBytes> = if txn.ringbuffer_row_pre_insert_interceptors().is_empty() {
171 rows.to_vec()
172 } else {
173 let mut builders: Vec<_> =
174 rows.iter().map(|bytes| EncodedRingBufferRow::from(bytes.clone()).thaw()).collect();
175 RingBufferRowInterceptor::pre_insert(txn, ringbuffer, &mut builders)?;
176 let written: Vec<EncodedBytes> = builders.into_iter().map(|builder| builder.freeze_bytes()).collect();
177 if moved_partitions(ringbuffer, shape, partitions, &written)? {
178 return Err(PartitionError::PartitionColumnChangedByHook {
179 object: ObjectId::ringbuffer(ringbuffer.id),
180 }
181 .into());
182 }
183 written
184 };
185
186 txn.set_with_pre(
187 keys.into_iter()
188 .zip(written.iter().cloned())
189 .map(|(key, bytes)| (key, bytes, SetPre::Absent))
190 .collect(),
191 )?;
192
193 if !txn.ringbuffer_row_post_insert_interceptors().is_empty() {
194 RingBufferRowInterceptor::post_insert(txn, ringbuffer, ids, &written)?;
195 }
196
197 let object = ObjectId::ringbuffer(ringbuffer.id);
198 txn.touch_flow_object(object);
199 if txn.wants_flow_diff(object) {
200 txn.track_flow_change(build_ringbuffer_insert_change(ringbuffer, shape, ids, &written)?);
201 }
202
203 Ok(written)
204}
205
206fn update_rows<T: WithInterceptors + Write + Read>(
207 txn: &mut T,
208 ringbuffer: &RingBuffer,
209 partitions: &[Partition],
210 ids: &[RowNumber],
211 rows: &[EncodedBytes],
212 pres: &[EncodedBytes],
213) -> Result<Vec<EncodedBytes>> {
214 assert_eq!(ids.len(), rows.len(), "ids/rows length mismatch");
215 assert_eq!(ids.len(), pres.len(), "ids/pres length mismatch");
216 let shape = row_shape_from_columns(RowFamily::RingBuffer, &ringbuffer.columns);
217
218 let keys: Vec<TaggedKey> = ids
219 .iter()
220 .enumerate()
221 .map(|(idx, &id)| ringbuffer_key(ringbuffer, partitions.get(idx).copied(), id))
222 .collect();
223
224 let posts: Vec<EncodedBytes> = if txn.ringbuffer_row_pre_update_interceptors().is_empty() {
225 rows.to_vec()
226 } else {
227 let mut builders: Vec<_> =
228 rows.iter().map(|bytes| EncodedRingBufferRow::from(bytes.clone()).thaw()).collect();
229 RingBufferRowInterceptor::pre_update(txn, ringbuffer, ids, &mut builders)?;
230 let posts: Vec<EncodedBytes> = builders.into_iter().map(|builder| builder.freeze_bytes()).collect();
231 if moved_partitions(ringbuffer, &shape, partitions, &posts)? {
232 return Err(PartitionError::ImmutablePartitionColumn {
233 object: ObjectId::ringbuffer(ringbuffer.id),
234 }
235 .into());
236 }
237 posts
238 };
239
240 txn.set(keys.into_iter().zip(posts.iter().cloned()).collect())?;
241
242 if !txn.ringbuffer_row_post_update_interceptors().is_empty() {
243 RingBufferRowInterceptor::post_update(txn, ringbuffer, ids, &posts, pres)?;
244 }
245
246 let object = ObjectId::ringbuffer(ringbuffer.id);
247 txn.touch_flow_object(object);
248 if txn.wants_flow_diff(object) {
249 txn.track_flow_change(build_ringbuffer_update_change(ringbuffer, &shape, ids, pres, &posts)?);
250 }
251
252 Ok(posts)
253}
254
255fn remove_rows<T: WithInterceptors + Write + Read>(
256 txn: &mut T,
257 ringbuffer: &RingBuffer,
258 partitions: &[Partition],
259 ids: &[RowNumber],
260 displayed: Vec<Option<EncodedBytes>>,
261) -> Result<Vec<EncodedBytes>> {
262 assert_eq!(ids.len(), displayed.len(), "ids/displayed length mismatch");
263 let mut displayed_rows = Vec::with_capacity(ids.len());
264 let mut removed_ids = Vec::new();
265 let mut matched_keys = Vec::new();
266 let mut matched_displayed = Vec::new();
267
268 for (idx, (&id, displayed)) in ids.iter().zip(displayed).enumerate() {
269 let Some(displayed) = displayed else {
270 displayed_rows.push(EncodedBytes(CowVec::new(vec![])));
271 continue;
272 };
273
274 removed_ids.push(id);
275 matched_keys.push(ringbuffer_key(ringbuffer, partitions.get(idx).copied(), id));
276 matched_displayed.push(displayed.clone());
277 displayed_rows.push(displayed);
278 }
279 if removed_ids.is_empty() {
280 return Ok(displayed_rows);
281 }
282
283 let committed: Vec<Option<EncodedBytes>> =
284 txn.get_committed(&matched_keys)?.into_iter().map(|row| row.map(|v| v.bytes)).collect();
285 let removed: Vec<EncodedBytes> = committed
286 .iter()
287 .zip(&matched_displayed)
288 .map(|(committed, displayed)| committed.clone().unwrap_or_else(|| displayed.clone()))
289 .collect();
290
291 if !txn.ringbuffer_row_pre_delete_interceptors().is_empty() {
292 RingBufferRowInterceptor::pre_delete(txn, ringbuffer, &removed_ids)?;
293 }
294
295 txn.remove(matched_keys.into_iter().zip(removed.iter().cloned().map(Some)).collect())?;
296
297 if !txn.ringbuffer_row_post_delete_interceptors().is_empty() {
298 RingBufferRowInterceptor::post_delete(txn, ringbuffer, &removed_ids, &removed)?;
299 }
300
301 let object = ObjectId::ringbuffer(ringbuffer.id);
302 txn.touch_flow_object(object);
303 if txn.wants_flow_diff(object) {
304 let shape = row_shape_from_columns(RowFamily::RingBuffer, &ringbuffer.columns);
305 txn.track_flow_change(build_ringbuffer_remove_change(ringbuffer, &shape, &removed_ids, &removed)?);
306 }
307
308 Ok(displayed_rows)
309}
310
311impl RingBufferOperations for CommandTransaction {
312 fn insert_ringbuffer(
313 &mut self,
314 ringbuffer: &RingBuffer,
315 shape: &RowShape,
316 partitions: &[Partition],
317 ids: &[RowNumber],
318 rows: &[EncodedBytes],
319 ) -> Result<Vec<EncodedBytes>> {
320 insert_rows(self, ringbuffer, shape, partitions, ids, rows)
321 }
322
323 fn update_ringbuffer(
324 &mut self,
325 ringbuffer: &RingBuffer,
326 partitions: &[Partition],
327 ids: &[RowNumber],
328 rows: &[EncodedBytes],
329 pres: &[EncodedBytes],
330 ) -> Result<Vec<EncodedBytes>> {
331 update_rows(self, ringbuffer, partitions, ids, rows, pres)
332 }
333
334 fn remove_from_ringbuffer(
335 &mut self,
336 ringbuffer: &RingBuffer,
337 partitions: &[Partition],
338 ids: &[RowNumber],
339 displayed: Vec<Option<EncodedBytes>>,
340 ) -> Result<Vec<EncodedBytes>> {
341 remove_rows(self, ringbuffer, partitions, ids, displayed)
342 }
343}
344
345impl RingBufferOperations for AdminTransaction {
346 fn insert_ringbuffer(
347 &mut self,
348 ringbuffer: &RingBuffer,
349 shape: &RowShape,
350 partitions: &[Partition],
351 ids: &[RowNumber],
352 rows: &[EncodedBytes],
353 ) -> Result<Vec<EncodedBytes>> {
354 insert_rows(self, ringbuffer, shape, partitions, ids, rows)
355 }
356
357 fn update_ringbuffer(
358 &mut self,
359 ringbuffer: &RingBuffer,
360 partitions: &[Partition],
361 ids: &[RowNumber],
362 rows: &[EncodedBytes],
363 pres: &[EncodedBytes],
364 ) -> Result<Vec<EncodedBytes>> {
365 update_rows(self, ringbuffer, partitions, ids, rows, pres)
366 }
367
368 fn remove_from_ringbuffer(
369 &mut self,
370 ringbuffer: &RingBuffer,
371 partitions: &[Partition],
372 ids: &[RowNumber],
373 displayed: Vec<Option<EncodedBytes>>,
374 ) -> Result<Vec<EncodedBytes>> {
375 remove_rows(self, ringbuffer, partitions, ids, displayed)
376 }
377}
378
379impl RingBufferOperations for Transaction<'_> {
380 fn insert_ringbuffer(
381 &mut self,
382 ringbuffer: &RingBuffer,
383 shape: &RowShape,
384 partitions: &[Partition],
385 ids: &[RowNumber],
386 rows: &[EncodedBytes],
387 ) -> Result<Vec<EncodedBytes>> {
388 match self {
389 Transaction::Command(txn) => txn.insert_ringbuffer(ringbuffer, shape, partitions, ids, rows),
390 Transaction::Admin(txn) => txn.insert_ringbuffer(ringbuffer, shape, partitions, ids, rows),
391 Transaction::Test(t) => t.inner.insert_ringbuffer(ringbuffer, shape, partitions, ids, rows),
392 Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
393 }
394 }
395
396 fn update_ringbuffer(
397 &mut self,
398 ringbuffer: &RingBuffer,
399 partitions: &[Partition],
400 ids: &[RowNumber],
401 rows: &[EncodedBytes],
402 pres: &[EncodedBytes],
403 ) -> Result<Vec<EncodedBytes>> {
404 match self {
405 Transaction::Command(txn) => txn.update_ringbuffer(ringbuffer, partitions, ids, rows, pres),
406 Transaction::Admin(txn) => txn.update_ringbuffer(ringbuffer, partitions, ids, rows, pres),
407 Transaction::Test(t) => t.inner.update_ringbuffer(ringbuffer, partitions, ids, rows, pres),
408 Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
409 }
410 }
411
412 fn remove_from_ringbuffer(
413 &mut self,
414 ringbuffer: &RingBuffer,
415 partitions: &[Partition],
416 ids: &[RowNumber],
417 displayed: Vec<Option<EncodedBytes>>,
418 ) -> Result<Vec<EncodedBytes>> {
419 match self {
420 Transaction::Command(txn) => txn.remove_from_ringbuffer(ringbuffer, partitions, ids, displayed),
421 Transaction::Admin(txn) => txn.remove_from_ringbuffer(ringbuffer, partitions, ids, displayed),
422 Transaction::Test(t) => t.inner.remove_from_ringbuffer(ringbuffer, partitions, ids, displayed),
423 Transaction::Query(_) => panic!("Write operations not supported on Query transaction"),
424 }
425 }
426}