alopex_sql/storage/
bridge.rs

1use std::collections::{BTreeMap, BTreeSet, HashMap};
2use std::sync::Arc;
3
4use alopex_core::kv::{
5    KVStore, KVTransaction, RangeChangePayload, RangeChangeRecord, ReadAtPoint, ReadAtResult,
6    stage_range_change,
7};
8use alopex_core::types::{Key, TxnMode, Value};
9use alopex_core::vector::hnsw::{HnswIndex, HnswTransactionState};
10use alopex_core::{CanonicalRowKey, Error as CoreError, Result as CoreResult};
11use sha2::{Digest, Sha256};
12
13use crate::catalog::CatalogOverlay;
14use crate::catalog::TableMetadata;
15
16use super::error::Result;
17use super::{IndexStorage, TableStorage};
18
19const SQL_DATA_START: &[u8] = &[0x01];
20const SQL_DATA_END: &[u8] = &[0x03];
21const JOURNAL_EPOCH_PREFIX: &[u8] = b"\x00alopex/range-change/epoch/";
22
23/// Scope used to turn local SQL row and secondary-index changes into one
24/// ordered range-change record.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct RangeChangeJournalScope {
27    range_id: String,
28    generation: u64,
29    index_tables: BTreeMap<u32, u32>,
30}
31
32impl RangeChangeJournalScope {
33    /// Creates a journal scope with the committed range identity and the
34    /// index-to-primary-table mapping used to reconstruct canonical row keys.
35    pub fn new(
36        range_id: impl Into<String>,
37        generation: u64,
38        index_tables: BTreeMap<u32, u32>,
39    ) -> Self {
40        Self {
41            range_id: range_id.into(),
42            generation,
43            index_tables,
44        }
45    }
46
47    /// Returns the local-only scope used by embedded single-node SQL commits.
48    pub fn local(index_tables: BTreeMap<u32, u32>) -> Self {
49        Self::new("local-default", 1, index_tables)
50    }
51}
52
53/// Captures the pre-write SQL keyspace for one local transaction. Calling
54/// [`Self::stage`] before the transaction commits records the exact row/index
55/// delta in that same durable boundary.
56#[derive(Debug, Clone)]
57pub struct LocalRangeChangeJournal {
58    scope: RangeChangeJournalScope,
59    before: BTreeMap<Key, Value>,
60}
61
62impl LocalRangeChangeJournal {
63    /// Captures the row and B-tree secondary-index state visible at transaction
64    /// start. Catalog keys, sequence keys, and journal metadata are excluded.
65    pub fn capture<'txn, T: KVTransaction<'txn>>(
66        transaction: &mut T,
67        scope: RangeChangeJournalScope,
68    ) -> CoreResult<Self> {
69        Ok(Self {
70            scope,
71            before: sql_data_snapshot(transaction)?,
72        })
73    }
74
75    /// Stages one immutable range-change record when SQL row/index state
76    /// changed. No epoch/outbox entry is written when the final state matches
77    /// the captured state, which makes an already-applied replay a no-op.
78    pub fn stage<'txn, T: KVTransaction<'txn>>(
79        self,
80        transaction: &mut T,
81    ) -> CoreResult<Option<RangeChangeRecord>> {
82        let after = sql_data_snapshot(transaction)?;
83        let payload = self.payload(&after)?;
84        if payload.is_empty() {
85            return Ok(None);
86        }
87
88        let epoch_key = journal_epoch_key(&self.scope.range_id);
89        let predecessor_epoch = transaction
90            .get(&epoch_key)?
91            .map(|value| decode_epoch(&value))
92            .transpose()?;
93        let epoch = predecessor_epoch
94            .unwrap_or(0)
95            .checked_add(1)
96            .ok_or_else(|| CoreError::InvalidFormat("range-change epoch overflow".to_string()))?;
97        let record = RangeChangeRecord {
98            range_id: self.scope.range_id.clone(),
99            generation: self.scope.generation,
100            epoch,
101            predecessor_epoch,
102            replay_id: replay_id(&payload)?,
103            payload,
104        };
105
106        stage_range_change(transaction, &record)?;
107        transaction.put(epoch_key, epoch.to_be_bytes().to_vec())?;
108        Ok(Some(record))
109    }
110
111    fn payload(&self, after: &BTreeMap<Key, Value>) -> CoreResult<Vec<RangeChangePayload>> {
112        let keys: BTreeSet<&Key> = self.before.keys().chain(after.keys()).collect();
113        let mut payload = Vec::new();
114        for key in keys {
115            let before = self.before.get(key);
116            let after = after.get(key);
117            if before == after {
118                continue;
119            }
120            match key.first().copied() {
121                Some(0x01) => {
122                    CanonicalRowKey::decode(key).map_err(|error| {
123                        CoreError::InvalidFormat(format!("invalid SQL row key in journal: {error}"))
124                    })?;
125                    match after {
126                        Some(encoded_row) => payload.push(RangeChangePayload::UpsertRow {
127                            row_key: key.clone(),
128                            encoded_row: encoded_row.clone(),
129                        }),
130                        None => payload.push(RangeChangePayload::DeleteRow {
131                            row_key: key.clone(),
132                            tombstone: before.cloned().unwrap_or_default(),
133                        }),
134                    }
135                }
136                Some(0x02) => {
137                    let Some((index_id, row_key)) = self.index_reference(key)? else {
138                        continue;
139                    };
140                    match after {
141                        Some(_) => payload.push(RangeChangePayload::UpsertIndex {
142                            index_id,
143                            index_key: key.clone(),
144                            row_key,
145                        }),
146                        None => payload.push(RangeChangePayload::DeleteIndex {
147                            index_id,
148                            index_key: key.clone(),
149                            row_key,
150                        }),
151                    }
152                }
153                _ => unreachable!("SQL data snapshot includes only row/index keyspaces"),
154            }
155        }
156        Ok(payload)
157    }
158
159    fn index_reference(&self, key: &Key) -> CoreResult<Option<(u32, Key)>> {
160        if key.len() < 13 {
161            return Err(CoreError::InvalidFormat(
162                "invalid SQL secondary-index key in journal".to_string(),
163            ));
164        }
165        let index_id = u32::from_be_bytes(key[1..5].try_into().expect("four bytes"));
166        let Some(table_id) = self.scope.index_tables.get(&index_id).copied() else {
167            // A newly-created index belongs to schema/catalog work. Its
168            // definition is not yet a committed data-plane contract, so do
169            // not emit an unverifiable index payload from this local bridge.
170            return Ok(None);
171        };
172        let row_id = u64::from_be_bytes(key[key.len() - 8..].try_into().expect("eight bytes"));
173        Ok(Some((
174            index_id,
175            CanonicalRowKey::new(table_id, row_id).encode(),
176        )))
177    }
178}
179
180fn sql_data_snapshot<'txn, T: KVTransaction<'txn>>(
181    transaction: &mut T,
182) -> CoreResult<BTreeMap<Key, Value>> {
183    Ok(transaction
184        .scan_range(SQL_DATA_START, SQL_DATA_END)?
185        .collect::<BTreeMap<_, _>>())
186}
187
188fn journal_epoch_key(range_id: &str) -> Key {
189    let mut key = JOURNAL_EPOCH_PREFIX.to_vec();
190    key.extend_from_slice(range_id.as_bytes());
191    key
192}
193
194fn decode_epoch(value: &Value) -> CoreResult<u64> {
195    let bytes: [u8; 8] = value
196        .as_slice()
197        .try_into()
198        .map_err(|_| CoreError::InvalidFormat("invalid range-change epoch state".to_string()))?;
199    Ok(u64::from_be_bytes(bytes))
200}
201
202fn replay_id(payload: &[RangeChangePayload]) -> CoreResult<String> {
203    let encoded = bincode::serialize(payload).map_err(|error| {
204        CoreError::InvalidFormat(format!(
205            "cannot encode range-change replay identity: {error}"
206        ))
207    })?;
208    Ok(format!("{:x}", Sha256::digest(encoded)))
209}
210
211pub trait SqlTxn<'txn, S: KVStore + 'txn> {
212    fn mode(&self) -> TxnMode;
213
214    fn ensure_write_txn(&self) -> CoreResult<()>;
215
216    fn inner_mut(&mut self) -> &mut S::Transaction<'txn>;
217
218    fn hnsw_entry(&mut self, name: &str) -> CoreResult<&HnswIndex>;
219
220    fn hnsw_entry_mut(&mut self, name: &str) -> CoreResult<&mut HnswTxnEntry>;
221
222    fn flush_hnsw(&mut self) -> Result<()>;
223
224    fn abandon_hnsw(&mut self) -> Result<()>;
225
226    fn delete_prefix(&mut self, prefix: &[u8]) -> Result<()>;
227
228    fn table_storage<'a>(
229        &'a mut self,
230        table_meta: &TableMetadata,
231    ) -> TableStorage<'a, 'txn, S::Transaction<'txn>> {
232        TableStorage::new(self.inner_mut(), table_meta)
233    }
234
235    fn index_storage<'a>(
236        &'a mut self,
237        index_id: u32,
238        unique: bool,
239        column_indices: Vec<usize>,
240    ) -> IndexStorage<'a, 'txn, S::Transaction<'txn>> {
241        IndexStorage::new(self.inner_mut(), index_id, unique, column_indices)
242    }
243
244    fn with_table<R, F>(&mut self, table_meta: &TableMetadata, f: F) -> Result<R>
245    where
246        F: FnOnce(&mut TableStorage<'_, 'txn, S::Transaction<'txn>>) -> Result<R>,
247    {
248        let mut storage = self.table_storage(table_meta);
249        f(&mut storage)
250    }
251
252    fn with_index<R, F>(
253        &mut self,
254        index_id: u32,
255        unique: bool,
256        column_indices: Vec<usize>,
257        f: F,
258    ) -> Result<R>
259    where
260        F: FnOnce(&mut IndexStorage<'_, 'txn, S::Transaction<'txn>>) -> Result<R>,
261    {
262        let mut storage = self.index_storage(index_id, unique, column_indices);
263        f(&mut storage)
264    }
265}
266
267/// TxnBridge provides SQL-friendly transaction handles on top of a KVStore.
268pub struct TxnBridge<S: KVStore> {
269    store: Arc<S>,
270}
271
272impl<S: KVStore> TxnBridge<S> {
273    /// Create a new bridge for the given store.
274    pub fn new(store: Arc<S>) -> Self {
275        Self { store }
276    }
277
278    /// Returns the underlying store's runtime statistics.
279    pub fn runtime_stats(&self) -> Option<alopex_core::kv::RuntimeStats> {
280        self.store.runtime_stats()
281    }
282
283    /// Applies a memory-limit setting to the underlying store.
284    pub fn set_memory_limit_bytes(&self, limit: Option<usize>) -> CoreResult<()> {
285        self.store.set_memory_limit_bytes(limit)
286    }
287
288    /// Applies a cache-capacity setting to the underlying store.
289    pub fn set_cache_capacity_bytes(&self, capacity: usize) -> CoreResult<()> {
290        self.store.set_cache_capacity_bytes(capacity)
291    }
292
293    /// Clears the underlying cache.
294    pub fn clear_cache(&self) -> CoreResult<usize> {
295        self.store.clear_cache()
296    }
297
298    /// Begin a read-only transaction.
299    pub fn begin_read(&self) -> Result<SqlTransaction<'_, S>> {
300        let txn = self.store.begin(TxnMode::ReadOnly)?;
301        Ok(SqlTransaction {
302            inner: txn,
303            mode: TxnMode::ReadOnly,
304            read_at: None,
305            hnsw_indices: HashMap::new(),
306        })
307    }
308
309    /// Begin a read-write transaction.
310    pub fn begin_write(&self) -> Result<SqlTransaction<'_, S>> {
311        let txn = self.store.begin(TxnMode::ReadWrite)?;
312        Ok(SqlTransaction {
313            inner: txn,
314            mode: TxnMode::ReadWrite,
315            read_at: None,
316            hnsw_indices: HashMap::new(),
317        })
318    }
319
320    /// Opens a read-only SQL transaction bound to a cluster-issued read point.
321    ///
322    /// No normal local transaction is substituted when the backend cannot
323    /// prove the requested retained snapshot.
324    pub fn begin_read_at(&self, point: ReadAtPoint) -> ReadAtResult<SqlTransaction<'_, S>> {
325        let txn = self.store.begin_read_at(&point)?;
326        Ok(Self::from_read_at(txn, point))
327    }
328
329    /// Wraps an already opened read-at transaction and retains its full fence.
330    pub fn from_read_at<'a>(txn: S::Transaction<'a>, point: ReadAtPoint) -> SqlTransaction<'a, S> {
331        SqlTransaction {
332            inner: txn,
333            mode: TxnMode::ReadOnly,
334            read_at: Some(point),
335            hnsw_indices: HashMap::new(),
336        }
337    }
338
339    pub fn wrap_external<'a, 'b, 'c>(
340        txn: &'a mut S::Transaction<'b>,
341        mode: TxnMode,
342        overlay: &'c mut CatalogOverlay,
343    ) -> BorrowedSqlTransaction<'a, 'b, 'c, S> {
344        BorrowedSqlTransaction {
345            inner: txn,
346            mode,
347            overlay,
348            hnsw_indices: HashMap::new(),
349        }
350    }
351
352    /// Execute a read-only transaction with automatic commit.
353    ///
354    /// The closure receives a `SqlTransaction` that provides access to table and index storage.
355    /// The transaction is automatically committed on success or rolled back on error.
356    pub fn with_read_txn<R, F>(&self, f: F) -> Result<R>
357    where
358        F: FnOnce(&mut SqlTransaction<'_, S>) -> Result<R>,
359    {
360        let mut txn = self.begin_read()?;
361        let result = f(&mut txn)?;
362        txn.commit()?;
363        Ok(result)
364    }
365
366    /// Execute a read-write transaction with automatic commit.
367    ///
368    /// The closure receives a `SqlTransaction` that provides access to table and index storage.
369    /// The transaction is automatically committed on success or rolled back on error.
370    pub fn with_write_txn<R, F>(&self, f: F) -> Result<R>
371    where
372        F: FnOnce(&mut SqlTransaction<'_, S>) -> Result<R>,
373    {
374        let mut txn = self.begin_write()?;
375        let result = f(&mut txn)?;
376        txn.commit()?;
377        Ok(result)
378    }
379
380    /// Execute a read-write transaction with explicit commit control.
381    ///
382    /// Returns `(result, should_commit)` from the closure. If `should_commit` is true,
383    /// the transaction is committed; otherwise it is rolled back.
384    pub fn with_write_txn_explicit<R, F>(&self, f: F) -> Result<R>
385    where
386        F: FnOnce(&mut SqlTransaction<'_, S>) -> Result<(R, bool)>,
387    {
388        let mut txn = self.begin_write()?;
389        let (result, should_commit) = f(&mut txn)?;
390        if should_commit {
391            txn.commit()?;
392        } else {
393            txn.rollback()?;
394        }
395        Ok(result)
396    }
397}
398
399/// SQL transaction wrapping a KV transaction.
400///
401/// Provides SQL-friendly access to table and index storage,
402/// with explicit commit/rollback control.
403pub struct SqlTransaction<'a, S: KVStore + 'a> {
404    inner: S::Transaction<'a>,
405    mode: TxnMode,
406    read_at: Option<ReadAtPoint>,
407    hnsw_indices: HashMap<String, HnswTxnEntry>,
408}
409
410pub struct HnswTxnEntry {
411    pub index: HnswIndex,
412    pub state: HnswTransactionState,
413    pub dirty: bool,
414}
415
416impl<'a, S: KVStore + 'a> SqlTransaction<'a, S> {
417    /// Returns the transaction mode.
418    pub fn mode(&self) -> TxnMode {
419        self.mode
420    }
421
422    /// Returns the fenced read point held by this transaction, if it was
423    /// opened by [`TxnBridge::begin_read_at`].
424    pub fn read_at_point(&self) -> Option<ReadAtPoint> {
425        self.read_at
426    }
427
428    /// Get a table storage handle.
429    ///
430    /// Returns a TableStorage instance that borrows this transaction.
431    /// The returned storage is valid for the duration of the borrow.
432    ///
433    /// The `table_id` is obtained from `table_meta.table_id`.
434    pub fn table_storage<'b>(
435        &'b mut self,
436        table_meta: &TableMetadata,
437    ) -> TableStorage<'b, 'a, S::Transaction<'a>> {
438        TableStorage::new(&mut self.inner, table_meta)
439    }
440
441    /// Get an index storage handle.
442    ///
443    /// Returns an IndexStorage instance that borrows this transaction.
444    /// The returned storage is valid for the duration of the borrow.
445    pub fn index_storage<'b>(
446        &'b mut self,
447        index_id: u32,
448        unique: bool,
449        column_indices: Vec<usize>,
450    ) -> IndexStorage<'b, 'a, S::Transaction<'a>> {
451        IndexStorage::new(&mut self.inner, index_id, unique, column_indices)
452    }
453
454    /// HNSW インデックスのキャッシュエントリを取得する（存在しなければロードする）。
455    #[allow(dead_code)]
456    pub(crate) fn hnsw_entry(&mut self, name: &str) -> CoreResult<&HnswIndex> {
457        if !self.hnsw_indices.contains_key(name) {
458            let index = HnswIndex::load(name, &mut self.inner)?;
459            self.hnsw_indices.insert(
460                name.to_string(),
461                HnswTxnEntry {
462                    index,
463                    state: HnswTransactionState::default(),
464                    dirty: false,
465                },
466            );
467        }
468        Ok(&self.hnsw_indices.get(name).expect("inserted above").index)
469    }
470
471    /// HNSW インデックスのキャッシュエントリを可変で取得する（存在しなければロードする）。
472    pub(crate) fn hnsw_entry_mut(&mut self, name: &str) -> CoreResult<&mut HnswTxnEntry> {
473        if !self.hnsw_indices.contains_key(name) {
474            let index = HnswIndex::load(name, &mut self.inner)?;
475            self.hnsw_indices.insert(
476                name.to_string(),
477                HnswTxnEntry {
478                    index,
479                    state: HnswTransactionState::default(),
480                    dirty: false,
481                },
482            );
483        }
484        Ok(self.hnsw_indices.get_mut(name).expect("inserted above"))
485    }
486
487    pub(crate) fn ensure_write_txn(&self) -> CoreResult<()> {
488        if self.mode != TxnMode::ReadWrite {
489            return Err(CoreError::TxnConflict);
490        }
491        Ok(())
492    }
493
494    /// 内部の KV トランザクションへの可変参照を返す（HNSW ブリッジ向け）。
495    pub(crate) fn inner_mut(&mut self) -> &mut S::Transaction<'a> {
496        &mut self.inner
497    }
498
499    /// Delete all keys with the given prefix.
500    pub fn delete_prefix(&mut self, prefix: &[u8]) -> Result<()> {
501        // Process in small batches to avoid unbounded memory use.
502        const BATCH: usize = 512;
503        loop {
504            let mut keys = Vec::with_capacity(BATCH);
505            {
506                let iter = self.inner.scan_prefix(prefix)?;
507                for (key, _) in iter.take(BATCH) {
508                    keys.push(key);
509                }
510            }
511
512            if keys.is_empty() {
513                break;
514            }
515
516            for key in keys {
517                self.inner.delete(key)?;
518            }
519        }
520
521        Ok(())
522    }
523
524    /// Execute operations on a table within a closure.
525    ///
526    /// This is a convenience method that creates a TableStorage,
527    /// executes the closure, and returns the result.
528    ///
529    /// The `table_id` is obtained from `table_meta.table_id`.
530    pub fn with_table<R, F>(&mut self, table_meta: &TableMetadata, f: F) -> Result<R>
531    where
532        F: FnOnce(&mut TableStorage<'_, 'a, S::Transaction<'a>>) -> Result<R>,
533    {
534        let mut storage = self.table_storage(table_meta);
535        f(&mut storage)
536    }
537
538    /// Execute operations on an index within a closure.
539    ///
540    /// This is a convenience method that creates an IndexStorage,
541    /// executes the closure, and returns the result.
542    pub fn with_index<R, F>(
543        &mut self,
544        index_id: u32,
545        unique: bool,
546        column_indices: Vec<usize>,
547        f: F,
548    ) -> Result<R>
549    where
550        F: FnOnce(&mut IndexStorage<'_, 'a, S::Transaction<'a>>) -> Result<R>,
551    {
552        let mut storage = self.index_storage(index_id, unique, column_indices);
553        f(&mut storage)
554    }
555
556    /// Commit the transaction.
557    ///
558    /// Consumes the transaction. On success, all writes become visible.
559    pub fn commit(mut self) -> Result<()> {
560        self.commit_hnsw()?;
561        self.inner.commit_self()?;
562        Ok(())
563    }
564
565    /// Rollback the transaction.
566    ///
567    /// Consumes the transaction. All pending writes are discarded.
568    pub fn rollback(mut self) -> Result<()> {
569        self.rollback_hnsw()?;
570        self.inner.rollback_self()?;
571        Ok(())
572    }
573
574    fn commit_hnsw(&mut self) -> Result<()> {
575        for entry in self.hnsw_indices.values_mut() {
576            if entry.dirty {
577                entry
578                    .index
579                    .commit_staged(&mut self.inner, &mut entry.state)?;
580            }
581        }
582        self.hnsw_indices.clear();
583        Ok(())
584    }
585
586    fn rollback_hnsw(&mut self) -> Result<()> {
587        for entry in self.hnsw_indices.values_mut() {
588            if entry.dirty {
589                entry.index.rollback(&mut entry.state)?;
590            }
591        }
592        self.hnsw_indices.clear();
593        Ok(())
594    }
595}
596
597impl<'a, S: KVStore + 'a> SqlTxn<'a, S> for SqlTransaction<'a, S> {
598    fn mode(&self) -> TxnMode {
599        self.mode()
600    }
601
602    fn ensure_write_txn(&self) -> CoreResult<()> {
603        self.ensure_write_txn()
604    }
605
606    fn inner_mut(&mut self) -> &mut S::Transaction<'a> {
607        self.inner_mut()
608    }
609
610    fn hnsw_entry(&mut self, name: &str) -> CoreResult<&HnswIndex> {
611        self.hnsw_entry(name)
612    }
613
614    fn hnsw_entry_mut(&mut self, name: &str) -> CoreResult<&mut HnswTxnEntry> {
615        self.hnsw_entry_mut(name)
616    }
617
618    fn flush_hnsw(&mut self) -> Result<()> {
619        self.commit_hnsw()
620    }
621
622    fn abandon_hnsw(&mut self) -> Result<()> {
623        self.rollback_hnsw()
624    }
625
626    fn delete_prefix(&mut self, prefix: &[u8]) -> Result<()> {
627        self.delete_prefix(prefix)
628    }
629}
630
631pub struct BorrowedSqlTransaction<'a, 'b, 'c, S: KVStore + 'b> {
632    inner: &'a mut S::Transaction<'b>,
633    mode: TxnMode,
634    overlay: &'c mut CatalogOverlay,
635    hnsw_indices: HashMap<String, HnswTxnEntry>,
636}
637
638impl<'a, 'b, 'c, S: KVStore + 'b> BorrowedSqlTransaction<'a, 'b, 'c, S> {
639    pub fn mode(&self) -> TxnMode {
640        self.mode
641    }
642
643    pub fn split_parts(&mut self) -> (BorrowedSqlTxn<'_, 'b, S>, &mut CatalogOverlay) {
644        (
645            BorrowedSqlTxn {
646                inner: self.inner,
647                mode: self.mode,
648                hnsw_indices: &mut self.hnsw_indices,
649            },
650            self.overlay,
651        )
652    }
653}
654
655impl<'a, 'b, 'c, S: KVStore + 'b> Drop for BorrowedSqlTransaction<'a, 'b, 'c, S> {
656    fn drop(&mut self) {
657        for entry in self.hnsw_indices.values_mut() {
658            if entry.dirty {
659                let _ = entry.index.rollback(&mut entry.state);
660                entry.dirty = false;
661            }
662        }
663        self.hnsw_indices.clear();
664    }
665}
666
667pub struct BorrowedSqlTxn<'a, 'b, S: KVStore + 'b> {
668    inner: &'a mut S::Transaction<'b>,
669    mode: TxnMode,
670    hnsw_indices: &'a mut HashMap<String, HnswTxnEntry>,
671}
672
673impl<'a, 'b, S: KVStore + 'b> SqlTxn<'b, S> for BorrowedSqlTxn<'a, 'b, S> {
674    fn mode(&self) -> TxnMode {
675        self.mode
676    }
677
678    fn ensure_write_txn(&self) -> CoreResult<()> {
679        if self.mode != TxnMode::ReadWrite {
680            return Err(CoreError::TxnReadOnly);
681        }
682        Ok(())
683    }
684
685    fn inner_mut(&mut self) -> &mut S::Transaction<'b> {
686        self.inner
687    }
688
689    fn hnsw_entry(&mut self, name: &str) -> CoreResult<&HnswIndex> {
690        if !self.hnsw_indices.contains_key(name) {
691            let index = HnswIndex::load(name, self.inner)?;
692            self.hnsw_indices.insert(
693                name.to_string(),
694                HnswTxnEntry {
695                    index,
696                    state: HnswTransactionState::default(),
697                    dirty: false,
698                },
699            );
700        }
701        Ok(&self.hnsw_indices.get(name).expect("inserted above").index)
702    }
703
704    fn hnsw_entry_mut(&mut self, name: &str) -> CoreResult<&mut HnswTxnEntry> {
705        if !self.hnsw_indices.contains_key(name) {
706            let index = HnswIndex::load(name, self.inner)?;
707            self.hnsw_indices.insert(
708                name.to_string(),
709                HnswTxnEntry {
710                    index,
711                    state: HnswTransactionState::default(),
712                    dirty: false,
713                },
714            );
715        }
716        Ok(self.hnsw_indices.get_mut(name).expect("inserted above"))
717    }
718
719    fn flush_hnsw(&mut self) -> Result<()> {
720        for entry in self.hnsw_indices.values_mut() {
721            if entry.dirty {
722                entry.index.commit_staged(self.inner, &mut entry.state)?;
723                entry.dirty = false;
724            }
725        }
726        self.hnsw_indices.clear();
727        Ok(())
728    }
729
730    fn abandon_hnsw(&mut self) -> Result<()> {
731        for entry in self.hnsw_indices.values_mut() {
732            if entry.dirty {
733                entry.index.rollback(&mut entry.state)?;
734                entry.dirty = false;
735            }
736        }
737        self.hnsw_indices.clear();
738        Ok(())
739    }
740
741    fn delete_prefix(&mut self, prefix: &[u8]) -> Result<()> {
742        // Process in small batches to avoid unbounded memory use.
743        const BATCH: usize = 512;
744        loop {
745            let mut keys = Vec::with_capacity(BATCH);
746            {
747                let iter = self.inner.scan_prefix(prefix)?;
748                for (key, _) in iter.take(BATCH) {
749                    keys.push(key);
750                }
751            }
752
753            if keys.is_empty() {
754                break;
755            }
756
757            for key in keys {
758                self.inner.delete(key)?;
759            }
760        }
761
762        Ok(())
763    }
764}
765
766/// Backwards-compatible alias for SqlTransaction.
767pub type TxnContext<'a, S> = SqlTransaction<'a, S>;
768
769#[cfg(test)]
770mod tests {
771    use super::super::SqlValue;
772    use super::*;
773    use crate::catalog::ColumnMetadata;
774    use crate::planner::types::ResolvedType;
775    use crate::storage::KeyEncoder;
776    use alopex_core::kv::memory::MemoryKV;
777    use alopex_core::kv::{KVStore, decode_range_change, journal_key};
778    use alopex_core::types::TxnMode;
779    use std::sync::Arc;
780
781    fn sample_table_meta() -> TableMetadata {
782        TableMetadata::new(
783            "users",
784            vec![
785                ColumnMetadata::new("id", ResolvedType::Integer)
786                    .with_primary_key(true)
787                    .with_not_null(true),
788                ColumnMetadata::new("name", ResolvedType::Text).with_not_null(true),
789            ],
790        )
791        .with_table_id(1)
792    }
793
794    #[test]
795    fn read_txn_mode_is_readonly() {
796        let store = Arc::new(MemoryKV::new());
797        let bridge = TxnBridge::new(store);
798
799        bridge
800            .with_read_txn(|ctx| {
801                assert_eq!(ctx.mode(), TxnMode::ReadOnly);
802                Ok(())
803            })
804            .unwrap();
805    }
806
807    #[test]
808    fn read_at_never_substitutes_a_normal_local_read_transaction() {
809        let bridge = TxnBridge::new(Arc::new(MemoryKV::new()));
810        let point = ReadAtPoint::new(7, 11, 13, 17);
811
812        assert!(matches!(
813            bridge.begin_read_at(point),
814            Err(alopex_core::ReadAtError::Unavailable { .. })
815        ));
816    }
817
818    #[test]
819    fn write_txn_mode_is_readwrite() {
820        let store = Arc::new(MemoryKV::new());
821        let bridge = TxnBridge::new(store);
822
823        bridge
824            .with_write_txn(|ctx| {
825                assert_eq!(ctx.mode(), TxnMode::ReadWrite);
826                Ok(())
827            })
828            .unwrap();
829    }
830
831    #[test]
832    fn commit_persists_changes_and_read_sees_them() {
833        let store = Arc::new(MemoryKV::new());
834        let bridge = TxnBridge::new(store.clone());
835        let meta = sample_table_meta();
836
837        // Write transaction
838        bridge
839            .with_write_txn(|ctx| {
840                ctx.with_table(&meta, |table| {
841                    table.insert(1, &[SqlValue::Integer(1), SqlValue::Text("alice".into())])
842                })
843            })
844            .unwrap();
845
846        // Read transaction
847        let row = bridge
848            .with_read_txn(|ctx| ctx.with_table(&meta, |table| table.get(1)))
849            .unwrap()
850            .unwrap();
851
852        assert_eq!(row[1], SqlValue::Text("alice".into()));
853    }
854
855    #[test]
856    fn rollback_discards_uncommitted_writes() {
857        let store = Arc::new(MemoryKV::new());
858        let bridge = TxnBridge::new(store.clone());
859        let meta = sample_table_meta();
860
861        // Write transaction with explicit rollback
862        bridge
863            .with_write_txn_explicit(|ctx| {
864                ctx.with_table(&meta, |table| {
865                    table.insert(1, &[SqlValue::Integer(1), SqlValue::Text("bob".into())])
866                })?;
867                Ok(((), false)) // false = rollback
868            })
869            .unwrap();
870
871        // Read transaction should not see the row
872        let row = bridge
873            .with_read_txn(|ctx| ctx.with_table(&meta, |table| table.get(1)))
874            .unwrap();
875
876        assert!(row.is_none());
877    }
878
879    #[test]
880    fn conflicting_commits_trigger_transaction_conflict() {
881        let store = Arc::new(MemoryKV::new());
882        let bridge = TxnBridge::new(store);
883        let meta = sample_table_meta();
884
885        // Start first write transaction
886        let mut txn1 = bridge.begin_write().unwrap();
887        {
888            let mut table = txn1.table_storage(&meta);
889            table
890                .insert(1, &[SqlValue::Integer(1), SqlValue::Text("alice".into())])
891                .unwrap();
892        }
893
894        // Start second write transaction
895        let mut txn2 = bridge.begin_write().unwrap();
896        {
897            let mut table = txn2.table_storage(&meta);
898            table
899                .insert(1, &[SqlValue::Integer(1), SqlValue::Text("bob".into())])
900                .unwrap();
901        }
902
903        // Commit first, second should conflict
904        txn1.commit().unwrap();
905        let err = txn2.commit().unwrap_err();
906        assert!(matches!(
907            err,
908            super::super::StorageError::TransactionConflict
909        ));
910    }
911
912    #[test]
913    fn scan_rows_in_transaction() {
914        let store = Arc::new(MemoryKV::new());
915        let bridge = TxnBridge::new(store.clone());
916        let meta = sample_table_meta();
917
918        // Insert multiple rows
919        bridge
920            .with_write_txn(|ctx| {
921                ctx.with_table(&meta, |table| {
922                    for i in 1..=3 {
923                        table.insert(
924                            i,
925                            &[
926                                SqlValue::Integer(i as i32),
927                                SqlValue::Text(format!("user{i}")),
928                            ],
929                        )?;
930                    }
931                    Ok(())
932                })
933            })
934            .unwrap();
935
936        // Scan all rows
937        let rows: Vec<u64> = bridge
938            .with_read_txn(|ctx| {
939                ctx.with_table(&meta, |table| {
940                    let iter = table.scan()?;
941                    let ids: Vec<u64> = iter.filter_map(|r| r.ok().map(|(id, _)| id)).collect();
942                    Ok(ids)
943                })
944            })
945            .unwrap();
946
947        assert_eq!(rows, vec![1, 2, 3]);
948    }
949
950    #[test]
951    fn local_journal_stages_row_and_index_payloads_in_one_commit() {
952        let store = MemoryKV::new();
953        let mut transaction = store.begin(TxnMode::ReadWrite).unwrap();
954        let journal = LocalRangeChangeJournal::capture(
955            &mut transaction,
956            RangeChangeJournalScope::local([(7, 1)].into()),
957        )
958        .unwrap();
959        let row_key = KeyEncoder::row_key(1, 9);
960        let index_key = KeyEncoder::index_key(7, &SqlValue::Text("alice".into()), 9).unwrap();
961        transaction
962            .put(row_key.clone(), b"encoded-row".to_vec())
963            .unwrap();
964        transaction.put(index_key.clone(), Vec::new()).unwrap();
965
966        let record = journal.stage(&mut transaction).unwrap().unwrap();
967        let outbox_key = journal_key(&record);
968        transaction.commit_self().unwrap();
969
970        let mut reader = store.begin(TxnMode::ReadOnly).unwrap();
971        assert_eq!(reader.get(&row_key).unwrap(), Some(b"encoded-row".to_vec()));
972        assert_eq!(
973            decode_range_change(&reader.get(&outbox_key).unwrap().unwrap()).unwrap(),
974            record
975        );
976        assert_eq!(
977            record.payload,
978            vec![
979                RangeChangePayload::UpsertRow {
980                    row_key: row_key.clone(),
981                    encoded_row: b"encoded-row".to_vec(),
982                },
983                RangeChangePayload::UpsertIndex {
984                    index_id: 7,
985                    index_key,
986                    row_key,
987                },
988            ]
989        );
990    }
991
992    #[test]
993    fn failed_journalled_commit_publishes_neither_data_nor_outbox() {
994        let store = MemoryKV::new_with_limit(Some(0));
995        let mut transaction = store.begin(TxnMode::ReadWrite).unwrap();
996        let journal = LocalRangeChangeJournal::capture(
997            &mut transaction,
998            RangeChangeJournalScope::local(Default::default()),
999        )
1000        .unwrap();
1001        let row_key = KeyEncoder::row_key(1, 1);
1002        transaction.put(row_key.clone(), b"row".to_vec()).unwrap();
1003        let record = journal.stage(&mut transaction).unwrap().unwrap();
1004        let outbox_key = journal_key(&record);
1005        assert!(transaction.commit_self().is_err());
1006
1007        let mut reader = store.begin(TxnMode::ReadOnly).unwrap();
1008        assert!(reader.get(&row_key).unwrap().is_none());
1009        assert!(reader.get(&outbox_key).unwrap().is_none());
1010    }
1011
1012    #[test]
1013    fn local_journal_records_tombstones_and_a_contiguous_epoch() {
1014        let store = MemoryKV::new();
1015        let row_key = KeyEncoder::row_key(1, 9);
1016        let index_key = KeyEncoder::index_key(7, &SqlValue::Text("alice".into()), 9).unwrap();
1017        let scope = RangeChangeJournalScope::local([(7, 1)].into());
1018
1019        let first = {
1020            let mut transaction = store.begin(TxnMode::ReadWrite).unwrap();
1021            let journal =
1022                LocalRangeChangeJournal::capture(&mut transaction, scope.clone()).unwrap();
1023            transaction
1024                .put(row_key.clone(), b"encoded-row".to_vec())
1025                .unwrap();
1026            transaction.put(index_key.clone(), Vec::new()).unwrap();
1027            let record = journal.stage(&mut transaction).unwrap().unwrap();
1028            transaction.commit_self().unwrap();
1029            record
1030        };
1031        let second = {
1032            let mut transaction = store.begin(TxnMode::ReadWrite).unwrap();
1033            let journal = LocalRangeChangeJournal::capture(&mut transaction, scope).unwrap();
1034            transaction.delete(row_key.clone()).unwrap();
1035            transaction.delete(index_key.clone()).unwrap();
1036            let record = journal.stage(&mut transaction).unwrap().unwrap();
1037            transaction.commit_self().unwrap();
1038            record
1039        };
1040
1041        assert_eq!(first.epoch, 1);
1042        assert_eq!(first.predecessor_epoch, None);
1043        assert_eq!(second.epoch, 2);
1044        assert_eq!(second.predecessor_epoch, Some(1));
1045        assert_eq!(
1046            second.payload,
1047            vec![
1048                RangeChangePayload::DeleteRow {
1049                    row_key: row_key.clone(),
1050                    tombstone: b"encoded-row".to_vec(),
1051                },
1052                RangeChangePayload::DeleteIndex {
1053                    index_id: 7,
1054                    index_key,
1055                    row_key,
1056                },
1057            ]
1058        );
1059    }
1060
1061    #[test]
1062    fn replay_without_a_new_sql_delta_emits_no_second_record() {
1063        let store = MemoryKV::new();
1064        let row_key = KeyEncoder::row_key(1, 1);
1065        {
1066            let mut transaction = store.begin(TxnMode::ReadWrite).unwrap();
1067            let journal = LocalRangeChangeJournal::capture(
1068                &mut transaction,
1069                RangeChangeJournalScope::local(Default::default()),
1070            )
1071            .unwrap();
1072            transaction.put(row_key.clone(), b"row".to_vec()).unwrap();
1073            assert!(journal.stage(&mut transaction).unwrap().is_some());
1074            transaction.commit_self().unwrap();
1075        }
1076        let mut replay = store.begin(TxnMode::ReadWrite).unwrap();
1077        let journal = LocalRangeChangeJournal::capture(
1078            &mut replay,
1079            RangeChangeJournalScope::local(Default::default()),
1080        )
1081        .unwrap();
1082        assert!(journal.stage(&mut replay).unwrap().is_none());
1083        replay.commit_self().unwrap();
1084    }
1085}
alopex_sql/storage/bridge.rs

alopex_sql/storage/
bridge.rs
