Skip to main content

akar_storage/
lib.rs

1//! Akar storage engine.
2//!
3//! Disk-based columnar storage with buffer management, WAL, compression, and indexing.
4
5pub mod art_index;
6pub mod art_key;
7pub mod art_node;
8pub mod buffer_manager;
9pub mod checkpoint;
10pub mod column;
11pub mod column_chunk;
12pub mod compression;
13pub mod csr;
14pub mod csv_reader;
15pub mod free_space_manager;
16pub mod hyperloglog;
17pub mod ice_format;
18pub mod index;
19pub mod lazy_scanner;
20pub mod local_storage;
21pub mod local_wal;
22pub mod node_group;
23pub mod npy_reader;
24pub mod page;
25pub mod page_manager;
26#[cfg(feature = "parquet")]
27pub mod parquet_reader;
28#[cfg(feature = "parquet")]
29pub mod parquet_writer;
30pub mod persistence;
31pub mod predicate;
32pub mod roaring_bitmap;
33pub mod shadow_file;
34pub mod spiller;
35pub mod stats;
36pub mod string_dictionary;
37pub mod table;
38pub mod undo_buffer;
39pub mod update_info;
40pub mod vector_index;
41pub mod version_info;
42pub mod wal;
43pub mod wal_replayer;
44
45use akar_common::error::StorageError;
46use akar_common::memory::MemoryManager;
47use akar_common::types::Value;
48use akar_vector::hnsw::DistanceMetric;
49use buffer_manager::{BufferManager, BufferManagerConfig};
50use checkpoint::checkpoint;
51use std::path::PathBuf;
52use std::sync::{Arc, Mutex};
53use wal::WAL;
54
55pub use art_index::ArtPrimaryKeyIndex;
56pub use art_key::ArtKey;
57pub use column_chunk::{ColumnChunk, NODE_GROUP_SIZE};
58pub use index::{HashIndex, IndexKey, OnDiskHashIndex};
59pub use local_storage::LocalStorage;
60pub use local_wal::LocalWAL;
61pub use node_group::NodeGroup;
62pub use page_manager::PageManager;
63pub use persistence::TablePersistence;
64pub use shadow_file::ShadowFile;
65pub use spiller::{MultiWayStreamMerge, SpillFile, Spiller};
66pub use string_dictionary::StringDictionary;
67pub use table::{ColumnDefinition, NodeTable, RelTable, TableCatalog};
68pub use undo_buffer::UndoBuffer;
69pub use vector_index::{VectorIndexTable, extract_f64_list_from_value};
70pub use wal_replayer::{ReplayResult, WALReplayer};
71
72/// Convert a catalog column definition into a storage-level column definition.
73///
74/// `CatalogColumn.default_value` is handled at the binder/schema layer and does
75/// not need to be materialized in the storage table.
76impl From<&akar_catalog::CatalogColumn> for ColumnDefinition {
77    fn from(c: &akar_catalog::CatalogColumn) -> Self {
78        ColumnDefinition {
79            name: c.name.clone(),
80            logical_type: c.logical_type,
81            is_primary_key: c.is_primary_key,
82            compression: c.compression,
83        }
84    }
85}
86
87/// The storage manager — root of the storage engine.
88#[allow(dead_code)]
89pub struct StorageManager {
90    db_path: PathBuf,
91    buffer_manager: Arc<Mutex<BufferManager>>,
92    wal: Arc<Mutex<WAL>>,
93    memory_manager: Arc<MemoryManager>,
94    /// Page manager for allocation/deallocation.
95    page_manager: Option<Arc<PageManager>>,
96    /// Lock-free table catalog using DashMap internally.
97    pub(crate) table_catalog: Arc<TableCatalog>,
98    /// Durable column mirrors for in-memory tables (P45.4).
99    table_persistence: TablePersistence,
100}
101
102/// Storage info returned by CALL storage_info().
103#[derive(Debug, Clone)]
104pub struct StorageInfo {
105    pub db_path: String,
106    pub page_size: usize,
107    pub total_pages: u64,
108    pub free_pages: u64,
109}
110
111/// Buffer manager info returned by CALL bm_info().
112#[derive(Debug, Clone)]
113pub struct BufferInfo {
114    pub total_memory: usize,
115    pub used_memory: usize,
116    pub num_pinned: usize,
117}
118
119/// File info returned by CALL file_info() / CALL disk_size_info().
120#[derive(Debug, Clone)]
121pub struct FileInfo {
122    pub total_file_size: u64,
123    pub num_data_pages: u64,
124    pub wal_size: u64,
125}
126
127/// FSM info returned by CALL free_space_info().
128#[derive(Debug, Clone)]
129pub struct FsmInfo {
130    pub total_free_pages: u64,
131    pub num_entries: usize,
132}
133
134impl StorageManager {
135    pub fn new(db_path: PathBuf, memory_manager: Arc<MemoryManager>) -> Self {
136        // Ensure the database directory exists (ignore error for :memory: mode)
137        let _ = std::fs::create_dir_all(&db_path);
138
139        let config = BufferManagerConfig::default();
140        let bm = BufferManager::new(db_path.clone(), memory_manager.clone(), config);
141        let wal_path = if db_path.to_string_lossy() == ":memory:" {
142            // In-memory mode: use a temp dir for the WAL
143            let tmp = std::env::temp_dir().join("akar-wal");
144            let _ = std::fs::create_dir_all(&tmp);
145            tmp.join("wal.log")
146        } else {
147            db_path.join("wal.log")
148        };
149        // Do NOT delete the WAL file — it may contain un-recovered data from
150        // a previous session. Recovery is triggered later by `Database::new()`
151        // via the `recover()` method.
152        let wal = WAL::new(wal_path);
153        let fsm = Arc::new(free_space_manager::FreeSpaceManager::new());
154        let existing_pages = 0u64; // Will be determined by file metadata
155        let pm = PageManager::new(db_path.clone(), page::DEFAULT_PAGE_SIZE, existing_pages, fsm);
156        Self {
157            db_path,
158            buffer_manager: Arc::new(Mutex::new(bm)),
159            wal: Arc::new(Mutex::new(wal)),
160            memory_manager,
161            page_manager: Some(Arc::new(pm)),
162            table_catalog: Arc::new(TableCatalog::new()),
163            table_persistence: TablePersistence::new(),
164        }
165    }
166
167    /// Open (or create) a database at `db_path`, initializing all storage
168    /// subsystems and replaying the WAL if necessary.
169    ///
170    /// This is the primary entry point for storage initialization.
171    /// After opening, call `recover()` to replay any uncommitted WAL records.
172    pub fn open(db_path: PathBuf, memory_manager: Arc<MemoryManager>) -> Self {
173        Self::new(db_path, memory_manager)
174    }
175
176    /// Get a reference to the page manager, if available.
177    pub fn page_manager(&self) -> Option<&Arc<PageManager>> {
178        self.page_manager.as_ref()
179    }
180
181    pub fn buffer_manager(&self) -> &Arc<Mutex<BufferManager>> {
182        &self.buffer_manager
183    }
184
185    pub fn wal(&self) -> &Arc<Mutex<WAL>> {
186        &self.wal
187    }
188
189    pub fn db_path(&self) -> &PathBuf {
190        &self.db_path
191    }
192
193    /// Get a reference to the table catalog for reading/writing table data.
194    pub fn table_catalog(&self) -> Arc<TableCatalog> {
195        self.table_catalog.clone()
196    }
197
198    /// Flush all node + rel tables into their durable column mirrors.
199    ///
200    /// Called after every write (commit or single-writer DML) and at
201    /// checkpoint time so committed rows survive restarts (P45.4).
202    pub fn persist_all_tables(&self) -> Result<(), StorageError> {
203        if self.db_path.to_string_lossy() == ":memory:" {
204            return Ok(()); // In-memory databases have nothing to persist
205        }
206        let page_size = self.buffer_manager.lock().unwrap().page_size();
207        self.table_persistence
208            .persist_all(&self.table_catalog, &self.db_path, &self.buffer_manager, page_size)
209    }
210
211    /// Load all persisted tables from their durable column mirrors.
212    ///
213    /// Called during `Database::new()` AFTER tables are restored from the
214    /// persisted catalog. Returns the number of tables that had persisted data.
215    pub fn load_persisted_tables(&self) -> Result<usize, StorageError> {
216        if self.db_path.to_string_lossy() == ":memory:" {
217            return Ok(0); // In-memory databases have no persisted mirrors
218        }
219        let page_size = self.buffer_manager.lock().unwrap().page_size();
220        self.table_persistence
221            .load_all(&self.table_catalog, &self.db_path, &self.buffer_manager, page_size)
222    }
223
224    /// Delete the durable column mirror for a dropped table.
225    pub fn drop_table_persistence(&self, table_id: u64) {
226        self.table_persistence
227            .remove(table_id, &self.db_path, &self.buffer_manager);
228    }
229
230    /// Log a column write to the WAL before applying it to the BufferManager.
231    pub fn log_column_write(&self, table_id: u64, col_id: u32, page_id: u64, data: &[u8]) {
232        let mut wal = self.wal.lock().unwrap();
233        wal.log_column_write(table_id, col_id, page_id, data);
234    }
235
236    /// Create a node table in the catalog and return its ID.
237    pub fn create_node_table(&self, name: String, columns: Vec<ColumnDefinition>) -> NodeTable {
238        self.table_catalog.create_node_table(name, columns)
239    }
240
241    /// Restore a node table at a specific table ID during recovery from a
242    /// persisted catalog. Optionally recreates an ART primary-key index and
243    /// registers its file with the BufferManager.
244    pub fn restore_node_table(
245        &self,
246        table_id: u64,
247        name: String,
248        columns: Vec<ColumnDefinition>,
249        index_name: Option<&str>,
250    ) -> NodeTable {
251        let table = self
252            .table_catalog
253            .create_node_table_with_id(table_id, name.clone(), columns);
254
255        if let Some(index_name) = index_name {
256            // Register the index file with the BufferManager for persistence
257            let mut bm = self.buffer_manager.lock().unwrap();
258            let full_path = self.db_path.join(format!("{index_name}.art"));
259            if !bm.is_file_registered(index_name) {
260                bm.register_file(index_name, full_path);
261            }
262            drop(bm);
263
264            let _ = self.table_catalog.create_art_index(&name, index_name);
265        }
266        table
267    }
268
269    /// Restore a rel table at a specific table ID during recovery from a
270    /// persisted catalog.
271    pub fn restore_rel_table(
272        &self,
273        table_id: u64,
274        name: String,
275        src_table_id: u64,
276        dst_table_id: u64,
277        columns: Vec<ColumnDefinition>,
278    ) -> RelTable {
279        self.table_catalog
280            .create_rel_table_with_id(table_id, name, src_table_id, dst_table_id, columns)
281    }
282
283    /// Create a vector index in the catalog and register its file with the BufferManager.
284    pub fn create_vector_index(
285        &self,
286        name: String,
287        table_name: String,
288        column_name: String,
289        metric: DistanceMetric,
290        dimensions: u32,
291    ) -> VectorIndexTable {
292        let table = self
293            .table_catalog
294            .create_vector_index(name, table_name, column_name, metric, dimensions);
295
296        // Register the index file with the BufferManager
297        let mut bm = self.buffer_manager.lock().unwrap();
298        table.register_file(&mut bm, &self.db_path);
299
300        table
301    }
302
303    /// Get a vector index by name.
304    pub fn get_vector_index_by_name(&self, name: &str) -> Option<dashmap::mapref::one::Ref<'_, u64, VectorIndexTable>> {
305        self.table_catalog.get_vector_index_by_name(name)
306    }
307
308    /// Get a mutable vector index by name.
309    pub fn get_vector_index_by_name_mut(
310        &self,
311        name: &str,
312    ) -> Option<dashmap::mapref::one::RefMut<'_, u64, VectorIndexTable>> {
313        self.table_catalog.get_vector_index_by_name_mut(name)
314    }
315
316    /// Create an ART (Adaptive Radix Tree) index on a node table.
317    /// Delegates to TableCatalog and registers the index file with BufferManager.
318    pub fn create_art_index(&self, table_name: &str, index_name: &str) -> Result<(), StorageError> {
319        self.table_catalog.create_art_index(table_name, index_name)?;
320
321        // Register the index file with the BufferManager for persistence
322        let mut bm = self
323            .buffer_manager
324            .lock()
325            .map_err(|e| StorageError::BufferManager(format!("Lock poisoned: {e}")))?;
326        let full_path = self.db_path.join(format!("{index_name}.art"));
327        let file_name = index_name.to_string();
328        if !bm.is_file_registered(&file_name) {
329            bm.register_file(&file_name, full_path);
330        }
331        drop(bm);
332
333        Ok(())
334    }
335
336    /// Drop an ART index from a node table.
337    pub fn drop_art_index(&self, table_name: &str, _index_name: &str) -> Result<(), StorageError> {
338        self.table_catalog.drop_art_index(table_name)
339    }
340
341    /// Get the ART index for a node table (cloned copy for read-only access).
342    pub fn get_art_index(&self, table_name: &str) -> Option<crate::ArtPrimaryKeyIndex> {
343        self.table_catalog.get_art_index(table_name)
344    }
345
346    /// Create a rel table in the catalog.
347    pub fn create_rel_table(
348        &self,
349        name: String,
350        src_table_id: u64,
351        dst_table_id: u64,
352        columns: Vec<ColumnDefinition>,
353    ) -> RelTable {
354        self.table_catalog
355            .create_rel_table(name, src_table_id, dst_table_id, columns)
356    }
357
358    /// Get the total size of the WAL in bytes.
359    pub fn wal_size(&self) -> usize {
360        self.wal.lock().unwrap().total_size()
361    }
362
363    /// Perform a checkpoint: flush WAL + dirty pages to disk.
364    pub fn checkpoint(&self) -> std::io::Result<checkpoint::CheckpointResult> {
365        let mut wal = self
366            .wal
367            .lock()
368            .map_err(|e| std::io::Error::other(format!("Lock poisoned: {e}")))?;
369        checkpoint(&mut wal, &self.buffer_manager)
370    }
371
372    /// Conditionally trigger a checkpoint based on the given threshold.
373    ///
374    /// This is called after every DML/DDL operation from `Connection::query()`.
375    ///
376    /// Semantics:
377    /// - `threshold < 0` (e.g., -1): checkpoint after every write (every DML/DDL).
378    /// - `threshold == 0`: never auto-checkpoint (manual only via `CHECKPOINT`).
379    /// - `threshold > 0`: checkpoint when `wal_size() > threshold` (bytes).
380    ///
381    /// Returns `true` if a checkpoint was triggered.
382    pub fn maybe_checkpoint(
383        &self,
384        threshold: i64,
385        drain_fn: Option<&dyn Fn(std::time::Duration) -> bool>,
386    ) -> std::io::Result<bool> {
387        if threshold == 0 {
388            return Ok(false); // Auto-checkpoint disabled
389        }
390
391        let should_checkpoint = if threshold < 0 {
392            // Always checkpoint after every write (default behavior)
393            true
394        } else {
395            self.wal_size() > threshold as usize
396        };
397
398        if should_checkpoint {
399            let _ = self.checkpoint_with_drain(drain_fn)?;
400            Ok(true)
401        } else {
402            Ok(false)
403        }
404    }
405
406    /// Perform a checkpoint with transaction drain.
407    ///
408    /// Two-phase drain:
409    /// 1. Call the `drain_fn` callback to stop new transactions and wait for active ones
410    /// 2. Perform the checkpoint (WAL flush + BM flush)
411    ///
412    /// This is the concurrent-writer-safe checkpoint. Use this instead of
413    /// plain `checkpoint()` when concurrent writes are enabled.
414    ///
415    /// If `drain_fn` is `None`, the drain is skipped (backwards-compatible default).
416    /// If the drain times out, the checkpoint proceeds anyway — this is safe because
417    /// the WAL will capture any in-flight writes.
418    pub fn checkpoint_with_drain(
419        &self,
420        drain_fn: Option<&dyn Fn(std::time::Duration) -> bool>,
421    ) -> std::io::Result<crate::checkpoint::CheckpointResult> {
422        // Phase 1: Stop new transactions and drain active ones
423        if let Some(drain) = drain_fn {
424            let drained = drain(std::time::Duration::from_secs(30));
425            if !drained {
426                tracing::warn!("Checkpoint drain timed out — proceeding with best-effort checkpoint");
427            }
428        }
429
430        // Persist in-memory tables to their durable column mirrors so the
431        // checkpoint's BufferManager flush writes them to disk too (P45.4).
432        // Done before the checkpoint so a crash after WAL truncation still
433        // leaves the mirror consistent with the tables.
434        if let Err(e) = self.persist_all_tables() {
435            tracing::warn!("Persist tables before checkpoint failed: {e}");
436        }
437
438        // Phase 2: Do the actual checkpoint
439        let mut wal = self
440            .wal
441            .lock()
442            .map_err(|e| std::io::Error::other(format!("Lock poisoned: {e}")))?;
443        crate::checkpoint::checkpoint(&mut wal, &self.buffer_manager)
444    }
445
446    /// Get storage-level information for diagnostics.
447    pub fn storage_info(&self) -> StorageInfo {
448        let total_pages = self.page_manager.as_ref().map(|pm| pm.total_pages()).unwrap_or(0);
449        let free_pages = 0u64; // FSM query could be added later
450        StorageInfo {
451            db_path: self.db_path.to_string_lossy().to_string(),
452            page_size: self
453                .page_manager
454                .as_ref()
455                .map(|pm| pm.page_size())
456                .unwrap_or(page::DEFAULT_PAGE_SIZE),
457            total_pages,
458            free_pages,
459        }
460    }
461
462    /// Buffer manager statistics for CALL bm_info().
463    pub fn buffer_info(&self) -> BufferInfo {
464        let bm = self.buffer_manager.lock().unwrap();
465        let stats = bm.stats();
466        let page_size = bm.page_size();
467        BufferInfo {
468            total_memory: stats.num_frames * page_size,
469            used_memory: (stats.num_frames - (stats.num_frames - stats.pinned_frames - stats.dirty_frames)) * page_size,
470            num_pinned: stats.pinned_frames,
471        }
472    }
473
474    /// File-level statistics for CALL file_info() / CALL disk_size_info().
475    pub fn file_info(&self) -> FileInfo {
476        let db_path = &self.db_path;
477        let wal_path = db_path.join("wal.log");
478        let wal_size = std::fs::metadata(&wal_path).map(|m| m.len()).unwrap_or(0);
479        let data_size = std::fs::read_dir(db_path)
480            .map(|entries| {
481                entries
482                    .filter_map(|e| e.ok())
483                    .filter(|e| e.path().extension().map(|x| x == "data").unwrap_or(false))
484                    .map(|e| e.metadata().map(|m| m.len()).unwrap_or(0))
485                    .sum::<u64>()
486            })
487            .unwrap_or(0);
488        let page_size = self
489            .page_manager
490            .as_ref()
491            .map(|pm| pm.page_size())
492            .unwrap_or(page::DEFAULT_PAGE_SIZE) as u64;
493        FileInfo {
494            total_file_size: data_size + wal_size,
495            num_data_pages: data_size / page_size.max(1),
496            wal_size,
497        }
498    }
499
500    /// FSM statistics for CALL free_space_info().
501    pub fn fsm_info(&self) -> FsmInfo {
502        let total_pages = self.page_manager.as_ref().map(|pm| pm.total_pages()).unwrap_or(0);
503        let free_pages = 0u64; // FSM query later
504        FsmInfo {
505            total_free_pages: free_pages,
506            num_entries: total_pages as usize,
507        }
508    }
509
510    /// Commit a write transaction's data to storage.
511    ///
512    /// Orchestrates the full commit pipeline:
513    /// 1. Append `Commit` record to the WAL (write-ahead log)
514    /// 2. Flush node/rel tables to their durable column mirrors (P45.4)
515    /// 3. Flush `LocalStorage` buffered writes to the actual tables
516    /// 4. Apply `ShadowFile` copy-on-write pages to the BufferManager
517    /// 5. Optionally checkpoint if the WAL threshold is met
518    ///
519    /// # Arguments
520    ///
521    /// * `local_storage` — the transaction's write buffer (consumed on success).
522    /// * `shadow_file` — the transaction's COW page buffer.
523    /// * `checkpoint_threshold` — passed to `maybe_checkpoint()`; use -1 for
524    ///   always-checkpoint, 0 for never, N for byte-based threshold.
525    /// * `drain_fn` — optional callback to drain active transactions before checkpoint.
526    ///
527    /// Returns `Ok(())` if the commit pipeline succeeded.
528    pub fn commit_transaction(
529        &self,
530        local_storage: &crate::local_storage::LocalStorage,
531        shadow_file: &crate::shadow_file::ShadowFile,
532        checkpoint_threshold: i64,
533        txn_id: u64,
534        drain_fn: Option<&dyn Fn(std::time::Duration) -> bool>,
535    ) -> Result<(), StorageError> {
536        // Step 1: Write-ahead log the commit
537        {
538            let mut wal = self
539                .wal
540                .lock()
541                .map_err(|e| StorageError::Wal(format!("Lock poisoned: {e}")))?;
542            wal.append(crate::wal::WALRecord::Commit { transaction_id: txn_id });
543            wal.flush_to_disk()
544                .map_err(|e| StorageError::Wal(format!("WAL flush failed during commit: {e}")))?;
545        }
546
547        // Step 2: Flush in-memory tables to their durable column mirrors so
548        // committed rows survive process restarts.
549        self.persist_all_tables()
550            .map_err(|e| StorageError::LocalStorage(format!("table persist failed during commit: {e}")))?;
551
552        // Step 3: Flush local storage buffers to the actual tables
553        // Pass txn_id so inserts/deletes are recorded in VersionInfo for MVCC
554        let _commit_undo_records = local_storage.flush_to_tables(&self.table_catalog, Some(txn_id))?;
555        // Undo records generated during commit for potential rollback-on-failure.
556        // Currently unused since commit is atomic, but stored for future use
557        // (e.g., partial-commit recovery).
558        tracing::debug!(
559            "commit_transaction: generated {} undo records for txn#{}",
560            _commit_undo_records.len(),
561            txn_id
562        );
563
564        // Step 4: Apply shadow pages to the BufferManager
565        shadow_file
566            .apply(&self.buffer_manager)
567            .map_err(|e| StorageError::ShadowFile(format!("ShadowFile apply failed during commit: {e}")))?;
568
569        // Step 5: Auto-checkpoint if needed
570        if let Err(e) = self.maybe_checkpoint(checkpoint_threshold, drain_fn) {
571            tracing::warn!("Checkpoint after commit failed: {e}");
572            // Non-fatal — data is already in tables and WAL
573        }
574
575        Ok(())
576    }
577
578    /// Roll back a write transaction, discarding all pending changes.
579    ///
580    /// Clears the local storage buffer, discards shadow pages, and
581    /// applies undo records to restore pre-write state.
582    /// The caller should also call `TransactionManager::rollback()` to
583    /// update the transaction's status and release locks.
584    ///
585    /// `undo_records` — accumulated undo records from the transaction.
586    ///   Applied in reverse order to restore overwritten data.
587    ///
588    /// Returns `Ok(())` on success.
589    pub fn rollback_transaction(
590        &self,
591        local_storage: &mut crate::local_storage::LocalStorage,
592        shadow_file: &mut crate::shadow_file::ShadowFile,
593        txn_id: u64,
594        undo_records: &[akar_transaction::UndoRecord],
595    ) -> Result<(), StorageError> {
596        // Log the rollback to WAL
597        {
598            let mut wal = self
599                .wal
600                .lock()
601                .map_err(|e| StorageError::Wal(format!("Lock poisoned: {e}")))?;
602            wal.append(crate::wal::WALRecord::Rollback { transaction_id: txn_id });
603            let _ = wal.flush_to_disk();
604        }
605
606        // Apply undo records in reverse order to restore pre-write state
607        for record in undo_records.iter().rev() {
608            if let Some(mut table) = self.table_catalog.get_node_table_mut(record.table_id) {
609                match record.undo_type {
610                    akar_transaction::UndoType::Update => {
611                        let values = deserialize_values_from_bytes(&record.old_data, 1);
612                        if let Some(val) = values.into_iter().next() {
613                            table
614                                .update_cell(record.row_id, record.column as usize, val)
615                                .map_err(|e| {
616                                    StorageError::Undo(format!(
617                                        "Undo failed for table {} row {}: {e}",
618                                        record.table_id, record.row_id
619                                    ))
620                                })?;
621                        }
622                    }
623                    akar_transaction::UndoType::Insert => {
624                        // Rollback an insert: delete the row
625                        let _ = table.delete_row(record.row_id);
626                    }
627                    akar_transaction::UndoType::Delete => {
628                        // Rollback a delete: restore all column values
629                        let num_cols = table.columns.len();
630                        let values = deserialize_values_from_bytes(&record.old_data, num_cols);
631                        for (col_idx, val) in values.into_iter().enumerate() {
632                            let _ = table.update_cell(record.row_id, col_idx, val);
633                        }
634                    }
635                }
636            } else if let Some(mut rel) = self.table_catalog.get_rel_table_mut(record.table_id) {
637                // Rel-table undo (P52.18): inserts delete the edge, updates
638                // restore the property cell, deletes restore src/dst + props.
639                match record.undo_type {
640                    akar_transaction::UndoType::Update => {
641                        let values = deserialize_values_from_bytes(&record.old_data, 1);
642                        if let Some(val) = values.into_iter().next() {
643                            let _ = rel.update_cell(record.row_id as usize, record.column as usize, val);
644                        }
645                    }
646                    akar_transaction::UndoType::Insert => {
647                        // Rollback an edge insert: tombstone the edge.
648                        let _ = rel.delete_edge(record.row_id as usize);
649                    }
650                    akar_transaction::UndoType::Delete => {
651                        // Rollback an edge delete: restore src/dst + properties.
652                        let num_cols = rel.columns.len();
653                        let values = deserialize_values_from_bytes(&record.old_data, num_cols + 2);
654                        let mut iter = values.into_iter();
655                        let src = match iter.next() {
656                            Some(Value::UInt64(v)) => v,
657                            Some(Value::Int64(v)) if v >= 0 => v as u64,
658                            _ => u64::MAX,
659                        };
660                        let dst = match iter.next() {
661                            Some(Value::UInt64(v)) => v,
662                            Some(Value::Int64(v)) if v >= 0 => v as u64,
663                            _ => u64::MAX,
664                        };
665                        let props: Vec<_> = iter.collect();
666                        let _ = rel.restore_deleted_edge(record.row_id as usize, src, dst, props);
667                    }
668                }
669            }
670        }
671
672        // Discard buffered writes
673        local_storage.clear();
674        shadow_file.discard();
675
676        Ok(())
677    }
678
679    /// Recover state from the WAL after a crash or unclean shutdown.
680    ///
681    /// This loads any on-disk WAL records, replays them against the
682    /// in-memory `TableCatalog`, then takes a checkpoint to reset the WAL.
683    ///
684    /// Call this once during `Database::new()`, **after** table schemas
685    /// have been re-created (e.g., from DDL replay or a persisted catalog).
686    ///
687    /// Returns the number of records recovered, or an error if recovery
688    /// fails (database is corrupt).
689    pub fn recover(&self) -> std::io::Result<usize> {
690        let mut wal = self
691            .wal
692            .lock()
693            .map_err(|e| std::io::Error::other(format!("Lock poisoned: {e}")))?;
694
695        // Load WAL records from disk
696        wal.load_from_disk()?;
697
698        if wal.is_empty() {
699            return Ok(0); // Nothing to recover
700        }
701
702        let mut data_records = 0usize;
703        let catalog = self.table_catalog.clone();
704
705        // Replay each record
706        wal.replay(|record| {
707            use crate::wal::WALRecord;
708            let cat = catalog.clone();
709
710            match record {
711                WALRecord::Insert { table_id, data } => {
712                    // Deserialize the data as Value vec and insert into the
713                    // corresponding node table.
714                    if let Some(mut table) = cat.get_node_table_mut(*table_id) {
715                        let values = deserialize_values_from_bytes(data, table.columns.len());
716                        if let Err(e) = table.insert_row(values) {
717                            return Err(std::io::Error::other(format!("WAL recovery insert failed: {e}")));
718                        }
719                        data_records += 1;
720                    }
721                }
722                WALRecord::Delete { table_id, row_id } => {
723                    if let Some(mut table) = cat.get_node_table_mut(*table_id) {
724                        if let Err(e) = table.delete_row(*row_id) {
725                            return Err(std::io::Error::other(format!("WAL recovery delete failed: {e}")));
726                        }
727                        data_records += 1;
728                    }
729                }
730                WALRecord::Update {
731                    table_id,
732                    row_id,
733                    column,
734                    data,
735                } => {
736                    if let Some(mut table) = cat.get_node_table_mut(*table_id) {
737                        let values = deserialize_values_from_bytes(data, 1);
738                        if let Some(val) = values.into_iter().next() {
739                            if let Err(e) = table.update_cell(*row_id, *column as usize, val) {
740                                return Err(std::io::Error::other(format!("WAL recovery update failed: {e}")));
741                            }
742                            data_records += 1;
743                        }
744                    }
745                }
746                WALRecord::UpdateFsm { .. } => {
747                    // UpdateFsm records are handled by the FSM recovery directly,
748                    // we can ignore them during the table-level replay.
749                }
750                WALRecord::ColumnWrite { .. } => {
751                    // ColumnWrite records are for the BufferManager-level
752                    // page writes. At the table level, data is already in
753                    // NodeGroup memory, so we skip these during recovery
754                    // (the checkpoint handles page-level persistence).
755                }
756                WALRecord::Commit { .. } | WALRecord::Rollback { .. } => {
757                    // Transaction markers — ignore during recovery since
758                    // all records in the WAL at startup are from already-
759                    // committed transactions (uncommitted ones were lost
760                    // in the crash).
761                }
762                WALRecord::Checkpoint => {
763                    // Checkpoint marker — all data before this is already
764                    // durable. We can clear what we've processed so far.
765                    // In practice, a checkpoint clears the WAL so this
766                    // marker should rarely appear during recovery.
767                }
768                WALRecord::LocalWALData { .. } => {
769                    // LocalWALData is a bulk-copied buffer from a committed
770                    // transaction's LocalWAL. The individual records within
771                    // have already been applied during normal commit flow.
772                    // At recovery time, the raw data is opaque — we skip it
773                    // since the table-level records (Insert/Delete/Update)
774                    // are replayed independently from the same transaction.
775                }
776                // DDL records — metadata-only, no data to replay for now.
777                // DDL operations (CREATE/DROP TABLE, etc.) are captured via
778                // Catalog serialization separately.
779                WALRecord::CreateTable { .. }
780                | WALRecord::DropTable { .. }
781                | WALRecord::AlterTable { .. }
782                | WALRecord::CreateIndex { .. }
783                | WALRecord::DropIndex { .. }
784                | WALRecord::CreateSequence { .. } => {
785                    // DDL records are replayed via catalog snapshot, not
786                    // individual WAL entries. Skip during table-level replay.
787                }
788            }
789            Ok(())
790        })?;
791
792        // After successful replay, mirror the recovered rows into the durable
793        // column mirrors as well, so a subsequent startup with an empty WAL
794        // restores from the mirror instead of the (now truncated) log.
795        if data_records > 0 {
796            self.persist_all_tables()
797                .map_err(|e| std::io::Error::other(format!("WAL recovery: persist tables failed: {e}")))?;
798        }
799
800        // Take a checkpoint to reset the WAL and make all recovered data durable.
801        if let Err(e) = checkpoint(&mut wal, &self.buffer_manager) {
802            tracing::warn!("WAL recovery: checkpoint after replay failed: {e}");
803        }
804
805        Ok(data_records)
806    }
807}
808
809/// Helper: deserialize binary data back into a `Vec<Value>`.
810///
811/// Each value is stored as a tag byte followed by type-specific data
812/// (see `column.rs` for the tag format). This is a simplified version
813/// that handles the common primary-key and property types.
814pub(crate) fn deserialize_values_from_bytes(data: &[u8], expected_count: usize) -> Vec<Value> {
815    use crate::column::*;
816    use std::io::Read;
817
818    if data.is_empty() || expected_count == 0 {
819        return Vec::new();
820    }
821
822    let mut cursor = std::io::Cursor::new(data);
823    let mut values = Vec::with_capacity(expected_count);
824
825    for _ in 0..expected_count {
826        let mut tag_buf = [0u8; 1];
827        if cursor.read_exact(&mut tag_buf).is_err() {
828            values.push(Value::Null);
829            continue;
830        }
831        let tag = tag_buf[0];
832
833        match tag {
834            TAG_NULL => values.push(Value::Null),
835            TAG_BOOL => {
836                let mut buf = [0u8; 1];
837                if cursor.read_exact(&mut buf).is_ok() {
838                    values.push(Value::Bool(buf[0] != 0));
839                }
840            }
841            TAG_INT64 => {
842                let mut buf = [0u8; 8];
843                if cursor.read_exact(&mut buf).is_ok() {
844                    values.push(Value::Int64(i64::from_le_bytes(buf)));
845                }
846            }
847            TAG_INT32 => {
848                let mut buf = [0u8; 4];
849                if cursor.read_exact(&mut buf).is_ok() {
850                    values.push(Value::Int32(i32::from_le_bytes(buf)));
851                }
852            }
853            TAG_DOUBLE => {
854                let mut buf = [0u8; 8];
855                if cursor.read_exact(&mut buf).is_ok() {
856                    values.push(Value::Double(f64::from_le_bytes(buf)));
857                }
858            }
859            TAG_FLOAT => {
860                let mut buf = [0u8; 4];
861                if cursor.read_exact(&mut buf).is_ok() {
862                    values.push(Value::Float(f32::from_le_bytes(buf)));
863                }
864            }
865            TAG_STRING => {
866                let mut len_buf = [0u8; 4];
867                if cursor.read_exact(&mut len_buf).is_ok() {
868                    let len = u32::from_le_bytes(len_buf) as usize;
869                    let mut str_buf = vec![0u8; len];
870                    if cursor.read_exact(&mut str_buf).is_ok()
871                        && let Ok(s) = String::from_utf8(str_buf)
872                    {
873                        values.push(Value::String(s));
874                    }
875                }
876            }
877            TAG_DATE => {
878                let mut buf = [0u8; 4];
879                if cursor.read_exact(&mut buf).is_ok() {
880                    values.push(Value::Date(akar_common::types::Date(i32::from_le_bytes(buf))));
881                }
882            }
883            TAG_TIMESTAMP => {
884                let mut buf = [0u8; 8];
885                if cursor.read_exact(&mut buf).is_ok() {
886                    values.push(Value::Timestamp(akar_common::types::Timestamp(i64::from_le_bytes(buf))));
887                }
888            }
889            TAG_INTERNAL_ID => {
890                let mut table_buf = [0u8; 8];
891                let mut offset_buf = [0u8; 8];
892                if cursor.read_exact(&mut table_buf).is_ok() && cursor.read_exact(&mut offset_buf).is_ok() {
893                    values.push(Value::InternalID(akar_common::types::InternalID {
894                        table_id: u64::from_le_bytes(table_buf),
895                        offset: u64::from_le_bytes(offset_buf),
896                    }));
897                }
898            }
899            // For any other type, skip one byte and push null
900            _ => {
901                values.push(Value::Null);
902            }
903        }
904    }
905
906    values
907}
908
909// =========================================================================
910// Phase 1 integration tests — full pipeline: table → column → buffer
911// manager → WAL → checkpoint → compression → multi-node-group
912// =========================================================================
913
914#[cfg(test)]
915mod integration_tests {
916    use super::*;
917    use crate::column::{Column, TAG_INT64, TAG_STRING};
918    use crate::page::DEFAULT_PAGE_SIZE;
919    use crate::wal::WALRecord;
920    use akar_common::enums::CompressionType;
921    use akar_common::types::{LogicalTypeID, Value};
922
923    // -----------------------------------------------------------------
924    // Helper: create a StorageManager + column pair
925    // -----------------------------------------------------------------
926    fn setup_integration() -> (StorageManager, tempfile::TempDir) {
927        let dir = tempfile::tempdir().unwrap();
928        let mm = Arc::new(MemoryManager::new(128 * 1024 * 1024));
929        let sm = StorageManager::new(dir.path().to_path_buf(), mm);
930        (sm, dir)
931    }
932
933    // =================================================================
934    // Test 1: Create table → insert rows → flush → reopen → verify
935    // =================================================================
936    #[test]
937    fn test_table_full_persistence_cycle() {
938        let (sm, _dir) = setup_integration();
939
940        // 1. Create a node table with two columns
941        let mut table = sm.create_node_table(
942            "Person".into(),
943            vec![
944                ColumnDefinition {
945                    compression: akar_common::enums::CompressionType::Uncompressed,
946                    name: "name".into(),
947                    logical_type: LogicalTypeID::String,
948                    is_primary_key: true,
949                },
950                ColumnDefinition {
951                    compression: akar_common::enums::CompressionType::Uncompressed,
952                    name: "age".into(),
953                    logical_type: LogicalTypeID::Int64,
954                    is_primary_key: false,
955                },
956            ],
957        );
958        assert_eq!(table.table_id, 0);
959
960        // 2. Insert rows into the table
961        table
962            .insert_row(vec![Value::String("Alice".into()), Value::Int64(30)])
963            .unwrap();
964        table
965            .insert_row(vec![Value::String("Bob".into()), Value::Int64(25)])
966            .unwrap();
967        table
968            .insert_row(vec![Value::String("Charlie".into()), Value::Int64(35)])
969            .unwrap();
970        assert_eq!(table.num_rows, 3);
971
972        // 3. Verify data before checkpoint
973        assert_eq!(table.get_value(0, 0), Some(&Value::String("Alice".into())));
974        assert_eq!(table.get_value(1, 1), Some(&Value::Int64(25)));
975
976        // 4. Read back via scan_column across node groups
977        let names = table.scan_column(0, 0, 3, None, &[]);
978        assert_eq!(names.len(), 3);
979        assert_eq!(names[0], Value::String("Alice".into()));
980
981        let ages = table.scan_column(1, 1, 2, None, &[]);
982        assert_eq!(ages.len(), 2);
983        assert_eq!(ages[0], Value::Int64(25));
984    }
985
986    // =================================================================
987    // Test 2: WAL crash recovery — log writes, flush to disk, replay
988    // =================================================================
989    #[test]
990    fn test_wal_recovery_cycle() {
991        let dir = tempfile::tempdir().unwrap();
992        let wal_path = dir.path().join("wal.log");
993
994        // Phase 1: Write data with WAL logging
995        #[allow(unused_variables)]
996        let (wal_records_count, column_count) = {
997            let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
998            let config = BufferManagerConfig::default();
999            let bm = Arc::new(Mutex::new(BufferManager::new(dir.path().to_path_buf(), mm, config)));
1000            let mut wal = WAL::new(wal_path.clone());
1001
1002            let mut col = Column::new(LogicalTypeID::Int64, 0, 0, dir.path(), bm.clone(), DEFAULT_PAGE_SIZE);
1003
1004            // Write data and log each write to WAL
1005            for i in 0i64..10 {
1006                col.append_value(&Value::Int64(i)).unwrap();
1007                wal.log_column_write(0, 0, 0, &i.to_le_bytes());
1008            }
1009            wal.append(WALRecord::Commit { transaction_id: 1 });
1010            let count = wal.len();
1011
1012            // Flush WAL to disk
1013            wal.flush_to_disk().unwrap();
1014
1015            // Also flush BM pages
1016            {
1017                let mut bm_lock = bm.lock().unwrap();
1018                bm_lock.flush_all().unwrap();
1019            }
1020
1021            // Verify column data is correct before "crash"
1022            for i in 0i64..10 {
1023                let v = col.get_value(i as u64).unwrap();
1024                assert_eq!(v, Value::Int64(i), "Pre-crash data mismatch at {}", i);
1025            }
1026
1027            (count, 10)
1028        }; // Drop everything — simulate crash
1029
1030        // Phase 2: Verify the on-disk WAL file exists and has content
1031        assert!(wal_path.exists(), "WAL file should exist after flush");
1032        let file_len = std::fs::metadata(&wal_path).unwrap().len();
1033        assert!(file_len > 0, "WAL file should have content, got {} bytes", file_len);
1034
1035        // Verify that a fresh WAL created from the file would contain
1036        // the right number of records. Since WAL::new() starts in-memory,
1037        // we create a new one, and verify the file has the right data.
1038        // In a real recovery scenario, we'd implement WAL::load_from_disk().
1039        assert_eq!(
1040            wal_records_count, 11,
1041            "Expected 10 ColumnWrite + 1 Commit = 11 records, got {}",
1042            wal_records_count
1043        );
1044        assert_eq!(column_count, 10);
1045    }
1046
1047    // =================================================================
1048    // Test 3: Compression round-trip — write compressed → read back
1049    // =================================================================
1050    #[test]
1051    fn test_compression_full_roundtrip() {
1052        // Test IntegerBitpacking: small values compress, large values preserved
1053        let dir = tempfile::tempdir().unwrap();
1054        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1055        let config = BufferManagerConfig::default();
1056        let bm = Arc::new(Mutex::new(BufferManager::new(dir.path().to_path_buf(), mm, config)));
1057
1058        let mut col_int = Column::with_compression(
1059            LogicalTypeID::Int64,
1060            0,
1061            0,
1062            dir.path(),
1063            bm.clone(),
1064            DEFAULT_PAGE_SIZE,
1065            CompressionType::IntegerBitpacking,
1066        );
1067
1068        // Write a range of values from small to large
1069        let test_values: Vec<i64> = vec![0, 1, 42, 127, 255, 65535, 1_000_000, i64::MAX, i64::MIN, -1];
1070        for v in &test_values {
1071            col_int.append_value(&Value::Int64(*v)).unwrap();
1072        }
1073
1074        // Read back and verify
1075        for (i, expected) in test_values.iter().enumerate() {
1076            let v = col_int.get_value(i as u64).unwrap();
1077            assert_eq!(v, Value::Int64(*expected), "IntegerBitpacking mismatch at index {}", i);
1078        }
1079
1080        // Verify that setting compression doesn't break existing data
1081        let mut col_float = Column::with_compression(
1082            LogicalTypeID::Double,
1083            0,
1084            1,
1085            dir.path(),
1086            bm.clone(),
1087            DEFAULT_PAGE_SIZE,
1088            CompressionType::Float,
1089        );
1090
1091        let floats: Vec<f64> = vec![1.0, std::f64::consts::PI, -2.5e10, 0.0, f64::MIN_POSITIVE, f64::MAX];
1092        for v in &floats {
1093            col_float.append_value(&Value::Double(*v)).unwrap();
1094        }
1095
1096        for (i, expected) in floats.iter().enumerate() {
1097            let v = col_float.get_value(i as u64).unwrap();
1098            match v {
1099                Value::Double(d) => assert!(
1100                    (d - expected).abs() < 1e-10 || (d / expected - 1.0).abs() < 1e-10,
1101                    "Float compression mismatch at {}: got {}, expected {}",
1102                    i,
1103                    d,
1104                    expected
1105                ),
1106                _ => panic!("Expected Double, got {:?}", v),
1107            }
1108        }
1109
1110        // Write compressed values via existing Column API and verify roundtrip
1111        // with buffer manager flush
1112        col_int.flush().unwrap();
1113        col_float.flush().unwrap();
1114
1115        // Read again after flush
1116        for (i, expected) in test_values.iter().enumerate() {
1117            let v = col_int.get_value(i as u64).unwrap();
1118            assert_eq!(
1119                v,
1120                Value::Int64(*expected),
1121                "After flush: IntegerBitpacking mismatch at {}",
1122                i
1123            );
1124        }
1125        for (i, expected) in floats.iter().enumerate() {
1126            let v = col_float.get_value(i as u64).unwrap();
1127            match v {
1128                Value::Double(d) => assert!(
1129                    (d - expected).abs() < 1e-10 || (d / expected - 1.0).abs() < 1e-10,
1130                    "After flush: Float mismatch at {}",
1131                    i
1132                ),
1133                _ => panic!("Expected Double after flush"),
1134            }
1135        }
1136    }
1137
1138    // =================================================================
1139    // Test 4: Multi-node-group scan — insert > NODE_GROUP_SIZE rows
1140    // =================================================================
1141    #[test]
1142    fn test_multi_node_group_scan() {
1143        let (_sm, _dir) = setup_integration();
1144
1145        // Create a NodeTable (not via StorageManager, directly for test control)
1146        let mut table = NodeTable::new(
1147            1,
1148            "BigTable".into(),
1149            vec![
1150                ColumnDefinition {
1151                    compression: akar_common::enums::CompressionType::Uncompressed,
1152                    name: "id".into(),
1153                    logical_type: LogicalTypeID::Int64,
1154                    is_primary_key: false,
1155                },
1156                ColumnDefinition {
1157                    compression: akar_common::enums::CompressionType::Uncompressed,
1158                    name: "value".into(),
1159                    logical_type: LogicalTypeID::Int64,
1160                    is_primary_key: false,
1161                },
1162            ],
1163        );
1164
1165        // Insert NODE_GROUP_SIZE + 500 rows to span across multiple node groups
1166        let total_rows = NODE_GROUP_SIZE + 500;
1167        for i in 0..total_rows {
1168            table
1169                .insert_row(vec![Value::Int64(i as i64), Value::Int64((i * 2) as i64)])
1170                .unwrap();
1171        }
1172
1173        // Verify total row count
1174        assert_eq!(table.num_rows, total_rows as u64);
1175
1176        // Verify multiple node groups were created
1177        let expected_groups = 2; // 4096 fits in first group, 500 in second
1178        assert_eq!(
1179            table.node_groups.len(),
1180            expected_groups,
1181            "Expected {} node groups for {} rows",
1182            expected_groups,
1183            total_rows
1184        );
1185
1186        // Verify node group boundaries
1187        assert_eq!(table.node_groups[0].num_nodes, NODE_GROUP_SIZE as u64);
1188        assert_eq!(table.node_groups[1].num_nodes, 500);
1189        assert_eq!(table.node_groups[0].start_offset, 0);
1190        assert_eq!(table.node_groups[1].start_offset, NODE_GROUP_SIZE as u64);
1191
1192        // Verify scanning across group boundaries
1193        // Row at boundary: last row of group 0
1194        let row_at_boundary = (NODE_GROUP_SIZE - 1) as u64;
1195        assert_eq!(
1196            table.get_value(row_at_boundary as usize, 0),
1197            Some(&Value::Int64(row_at_boundary as i64))
1198        );
1199
1200        // First row of group 1
1201        let row_in_group1 = NODE_GROUP_SIZE as u64;
1202        assert_eq!(
1203            table.get_value(row_in_group1 as usize, 0),
1204            Some(&Value::Int64(row_in_group1 as i64))
1205        );
1206
1207        // Scan column 0 across the entire table
1208        let scanned = table.scan_column(0, 0, total_rows as u64, None, &[]);
1209        assert_eq!(scanned.len(), total_rows);
1210        assert_eq!(scanned[0], Value::Int64(0));
1211        assert_eq!(scanned[NODE_GROUP_SIZE], Value::Int64(NODE_GROUP_SIZE as i64));
1212        assert_eq!(scanned[total_rows - 1], Value::Int64((total_rows - 1) as i64));
1213
1214        // Scan column 1 with offset and count spanning both groups
1215        let scan_mid = table.scan_column(1, (NODE_GROUP_SIZE - 100) as u64, 200, None, &[]);
1216        assert_eq!(scan_mid.len(), 200);
1217        assert_eq!(scan_mid[0], Value::Int64(((NODE_GROUP_SIZE - 100) * 2) as i64));
1218        assert_eq!(scan_mid[199], Value::Int64(((NODE_GROUP_SIZE + 99) * 2) as i64));
1219
1220        // Verify to_column_major_data correctness
1221        let data = table.to_column_major_data();
1222        assert_eq!(data.len(), 2); // 2 columns
1223        assert_eq!(data[0].len(), total_rows);
1224        assert_eq!(data[1].len(), total_rows);
1225        assert_eq!(data[0][NODE_GROUP_SIZE], Value::Int64(NODE_GROUP_SIZE as i64));
1226        assert_eq!(data[1][0], Value::Int64(0));
1227        assert_eq!(data[1][total_rows - 1], Value::Int64(((total_rows - 1) * 2) as i64));
1228    }
1229
1230    // =================================================================
1231    // Test 5: Combined — WAL-logged compressed multi-node-group write
1232    // =================================================================
1233    #[test]
1234    fn test_compressed_multi_group_with_checkpoint() {
1235        let dir = tempfile::tempdir().unwrap();
1236        let mm = Arc::new(MemoryManager::new(128 * 1024 * 1024));
1237
1238        // Use explicit BM + WAL for full control
1239        let config = BufferManagerConfig::default();
1240        let bm = Arc::new(Mutex::new(BufferManager::new(dir.path().to_path_buf(), mm, config)));
1241        let wal_path = dir.path().join("wal.log");
1242        let mut wal = WAL::new(wal_path);
1243
1244        let mut col = Column::with_compression(
1245            LogicalTypeID::Int64,
1246            0,
1247            0,
1248            dir.path(),
1249            bm.clone(),
1250            DEFAULT_PAGE_SIZE,
1251            CompressionType::IntegerBitpacking,
1252        );
1253
1254        // Write enough values to span multiple pages
1255        let num_values = 500;
1256        for i in 0i64..num_values {
1257            col.append_value(&Value::Int64(i)).unwrap();
1258            wal.log_column_write(0, 0, 0, &i.to_le_bytes());
1259        }
1260        wal.append(WALRecord::Commit { transaction_id: 1 });
1261
1262        // Read back before checkpoint
1263        for i in 0i64..num_values {
1264            let v = col.get_value(i as u64).unwrap();
1265            assert_eq!(v, Value::Int64(i), "Pre-checkpoint mismatch at {}", i);
1266        }
1267
1268        // Checkpoint: flush WAL + dirty pages
1269        let mut bm_lock = bm.lock().unwrap();
1270        bm_lock.flush_all().unwrap();
1271        drop(bm_lock);
1272
1273        wal.flush_to_disk().unwrap();
1274
1275        // Read back after checkpoint
1276        for i in 0i64..num_values {
1277            let v = col.get_value(i as u64).unwrap();
1278            assert_eq!(v, Value::Int64(i), "Post-checkpoint mismatch at {}", i);
1279        }
1280
1281        // Verify multiple pages were allocated
1282        assert!(
1283            col.num_pages > 1,
1284            "Expected multiple pages for {} values, got {}",
1285            num_values,
1286            col.num_pages
1287        );
1288    }
1289
1290    // =================================================================
1291    // Test 6: Stress — 10k rows via column with checkpoint
1292    // =================================================================
1293    #[test]
1294    fn test_10k_row_stress() {
1295        let dir = tempfile::tempdir().unwrap();
1296        let mm = Arc::new(MemoryManager::new(256 * 1024 * 1024));
1297        let config = BufferManagerConfig::default();
1298        let bm = Arc::new(Mutex::new(BufferManager::new(dir.path().to_path_buf(), mm, config)));
1299
1300        let mut col = Column::new(LogicalTypeID::Int64, 0, 0, dir.path(), bm.clone(), DEFAULT_PAGE_SIZE);
1301
1302        // Write 10,000 values
1303        for i in 0i64..10_000 {
1304            col.append_value(&Value::Int64(i)).unwrap();
1305        }
1306        assert_eq!(col.num_values, 10_000);
1307
1308        // Read back all values
1309        for i in 0i64..10_000 {
1310            let v = col.get_value(i as u64).unwrap();
1311            assert_eq!(v, Value::Int64(i), "Stress test mismatch at {}", i);
1312        }
1313
1314        // Flush and re-verify
1315        col.flush().unwrap();
1316        for i in 0i64..10_000 {
1317            let v = col.get_value(i as u64).unwrap();
1318            assert_eq!(v, Value::Int64(i), "Post-flush stress mismatch at {}", i);
1319        }
1320
1321        // Verify multiple pages were used
1322        assert!(
1323            col.num_pages > 1,
1324            "Stress test should use multiple pages, got {}",
1325            col.num_pages
1326        );
1327    }
1328
1329    // =================================================================
1330    // Test 7: WAL recovery — insert data, simulate crash, recover
1331    // =================================================================
1332    #[test]
1333    fn test_wal_recovery_insert_then_recover() {
1334        let dir = tempfile::tempdir().unwrap();
1335
1336        // Phase 1: Create DB, insert data, flush WAL, then "crash"
1337        let _row_count = {
1338            let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1339            let sm = StorageManager::new(dir.path().to_path_buf(), mm);
1340
1341            // Create a node table
1342            let mut table = sm.create_node_table(
1343                "Person".into(),
1344                vec![
1345                    ColumnDefinition {
1346                        compression: akar_common::enums::CompressionType::Uncompressed,
1347                        name: "name".into(),
1348                        logical_type: LogicalTypeID::String,
1349                        is_primary_key: true,
1350                    },
1351                    ColumnDefinition {
1352                        compression: akar_common::enums::CompressionType::Uncompressed,
1353                        name: "age".into(),
1354                        logical_type: LogicalTypeID::Int64,
1355                        is_primary_key: false,
1356                    },
1357                ],
1358            );
1359
1360            // Insert rows
1361            table
1362                .insert_row(vec![Value::String("Alice".into()), Value::Int64(30)])
1363                .unwrap();
1364            table
1365                .insert_row(vec![Value::String("Bob".into()), Value::Int64(25)])
1366                .unwrap();
1367            table
1368                .insert_row(vec![Value::String("Charlie".into()), Value::Int64(35)])
1369                .unwrap();
1370
1371            // Put the table back into the catalog so WAL knows about it.
1372            // We re-create the table entry via the catalog API.
1373            {
1374                // Re-create the table entry with the same schema so the
1375                // catalog has a valid entry for recovery to target.
1376                sm.table_catalog.create_node_table(
1377                    "Person".into(),
1378                    vec![
1379                        ColumnDefinition {
1380                            compression: akar_common::enums::CompressionType::Uncompressed,
1381                            name: "name".into(),
1382                            logical_type: LogicalTypeID::String,
1383                            is_primary_key: true,
1384                        },
1385                        ColumnDefinition {
1386                            compression: akar_common::enums::CompressionType::Uncompressed,
1387                            name: "age".into(),
1388                            logical_type: LogicalTypeID::Int64,
1389                            is_primary_key: false,
1390                        },
1391                    ],
1392                );
1393            }
1394
1395            let count = table.num_rows;
1396            assert_eq!(count, 3);
1397
1398            // Write WAL records and flush to disk
1399            {
1400                let mut wal = sm.wal.lock().unwrap();
1401                wal.append(WALRecord::Insert {
1402                    table_id: table.table_id,
1403                    data: vec![
1404                        TAG_STRING, 5, 0, 0, 0, b'A', b'l', b'i', b'c', b'e', TAG_INT64, 30, 0, 0, 0, 0, 0, 0, 0,
1405                    ],
1406                });
1407                wal.append(WALRecord::Insert {
1408                    table_id: table.table_id,
1409                    data: vec![
1410                        TAG_STRING, 3, 0, 0, 0, b'B', b'o', b'b', TAG_INT64, 25, 0, 0, 0, 0, 0, 0, 0,
1411                    ],
1412                });
1413                wal.append(WALRecord::Insert {
1414                    table_id: table.table_id,
1415                    data: vec![
1416                        TAG_STRING, 7, 0, 0, 0, b'C', b'h', b'a', b'r', b'l', b'i', b'e', TAG_INT64, 35, 0, 0, 0, 0, 0,
1417                        0, 0,
1418                    ],
1419                });
1420                wal.flush_to_disk().unwrap();
1421            }
1422
1423            count
1424        }; // "Crash" — all state dropped
1425
1426        // Phase 2: Recover — create a new StorageManager, it should NOT delete
1427        // the WAL. Then manually trigger recovery.
1428        {
1429            let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1430            let sm = StorageManager::new(dir.path().to_path_buf(), mm);
1431
1432            // Verify WAL file exists and has records
1433            let wal_path = dir.path().join("wal.log");
1434            assert!(wal_path.exists(), "WAL file should exist for recovery");
1435
1436            // Create the same table schema so recovery has a target
1437            // (the catalog create_node_table stores the table internally).
1438            sm.create_node_table(
1439                "Person".into(),
1440                vec![
1441                    ColumnDefinition {
1442                        compression: akar_common::enums::CompressionType::Uncompressed,
1443                        name: "name".into(),
1444                        logical_type: LogicalTypeID::String,
1445                        is_primary_key: true,
1446                    },
1447                    ColumnDefinition {
1448                        compression: akar_common::enums::CompressionType::Uncompressed,
1449                        name: "age".into(),
1450                        logical_type: LogicalTypeID::Int64,
1451                        is_primary_key: false,
1452                    },
1453                ],
1454            );
1455
1456            // Recover — the WAL has table_id = 0 (the first table created).
1457            let recovered = sm.recover().unwrap();
1458            assert_eq!(recovered, 3, "Should recover 3 WAL records");
1459
1460            // Verify data survived — the table was re-created empty,
1461            // and recovery inserted exactly 3 rows from the WAL.
1462            {
1463                let recovered_table = sm.table_catalog.get_node_table_by_name("Person").unwrap();
1464                assert_eq!(recovered_table.num_rows, 3, "Should have recovered 3 rows");
1465                assert_eq!(recovered_table.get_value(0, 0), Some(&Value::String("Alice".into())));
1466                assert_eq!(recovered_table.get_value(1, 0), Some(&Value::String("Bob".into())));
1467                assert_eq!(recovered_table.get_value(2, 0), Some(&Value::String("Charlie".into())));
1468                assert_eq!(recovered_table.get_value(0, 1), Some(&Value::Int64(30)));
1469                assert_eq!(recovered_table.get_value(1, 1), Some(&Value::Int64(25)));
1470                assert_eq!(recovered_table.get_value(2, 1), Some(&Value::Int64(35)));
1471            }
1472
1473            // Verify WAL was checkpointed — after checkpoint the WAL has
1474            // exactly 1 record (the Checkpoint marker).
1475            {
1476                let wal = sm.wal.lock().unwrap();
1477                assert_eq!(
1478                    wal.len(),
1479                    1,
1480                    "WAL should have only the checkpoint marker after recovery"
1481                );
1482                assert!(matches!(wal.records()[0], crate::wal::WALRecord::Checkpoint));
1483            }
1484        }
1485    }
1486
1487    // =================================================================
1488    // Test 8: WAL recovery — no-op when no WAL exists
1489    // =================================================================
1490    #[test]
1491    fn test_wal_recovery_no_wal() {
1492        let dir = tempfile::tempdir().unwrap();
1493
1494        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1495        let sm = StorageManager::new(dir.path().to_path_buf(), mm);
1496
1497        // No WAL file = recovery returns 0
1498        let recovered = sm.recover().unwrap();
1499        assert_eq!(recovered, 0, "No WAL = no records recovered");
1500    }
1501
1502    // =================================================================
1503    // Test 9: WAL recovery — empty WAL file
1504    // =================================================================
1505    #[test]
1506    fn test_wal_recovery_empty_wal() {
1507        let dir = tempfile::tempdir().unwrap();
1508
1509        // Create an empty WAL file on disk
1510        let wal_path = dir.path().join("wal.log");
1511        std::fs::write(&wal_path, b"").unwrap();
1512
1513        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1514        let sm = StorageManager::new(dir.path().to_path_buf(), mm);
1515
1516        let recovered = sm.recover().unwrap();
1517        assert_eq!(recovered, 0, "Empty WAL = no records recovered");
1518    }
1519
1520    // =================================================================
1521    // Test 10: WAL load_from_disk roundtrip
1522    // =================================================================
1523    #[test]
1524    fn test_wal_load_from_disk_roundtrip() {
1525        use crate::wal::WALRecord;
1526        let dir = tempfile::tempdir().unwrap();
1527        let wal_path = dir.path().join("wal.log");
1528
1529        // Write records
1530        {
1531            let mut wal = WAL::new(wal_path.clone());
1532            wal.append(WALRecord::Insert {
1533                table_id: 42,
1534                data: vec![1, 2, 3, 4],
1535            });
1536            wal.append(WALRecord::Delete {
1537                table_id: 42,
1538                row_id: 0,
1539            });
1540            wal.append(WALRecord::Update {
1541                table_id: 42,
1542                row_id: 1,
1543                column: 2,
1544                data: vec![5, 6],
1545            });
1546            wal.append(WALRecord::ColumnWrite {
1547                table_id: 42,
1548                col_id: 0,
1549                page_id: 1,
1550                data: vec![7, 8, 9],
1551            });
1552            wal.append(WALRecord::Commit { transaction_id: 100 });
1553            wal.append(WALRecord::Rollback { transaction_id: 101 });
1554            wal.append(WALRecord::Checkpoint);
1555            wal.flush_to_disk().unwrap();
1556        }
1557
1558        // Load from disk
1559        {
1560            let mut wal = WAL::new(wal_path.clone());
1561            wal.load_from_disk().unwrap();
1562            assert_eq!(wal.len(), 7, "Should load 7 records from disk");
1563            assert!(wal.is_dirty());
1564
1565            // Verify each record type
1566            match &wal.records()[0] {
1567                WALRecord::Insert { table_id, data } => {
1568                    assert_eq!(*table_id, 42);
1569                    assert_eq!(data, &[1, 2, 3, 4]);
1570                }
1571                _ => panic!("Expected Insert"),
1572            }
1573            match &wal.records()[1] {
1574                WALRecord::Delete { table_id, row_id } => {
1575                    assert_eq!(*table_id, 42);
1576                    assert_eq!(*row_id, 0);
1577                }
1578                _ => panic!("Expected Delete"),
1579            }
1580            match &wal.records()[4] {
1581                WALRecord::Commit { transaction_id } => {
1582                    assert_eq!(*transaction_id, 100);
1583                }
1584                _ => panic!("Expected Commit"),
1585            }
1586            match &wal.records()[5] {
1587                WALRecord::Rollback { transaction_id } => {
1588                    assert_eq!(*transaction_id, 101);
1589                }
1590                _ => panic!("Expected Rollback"),
1591            }
1592            match &wal.records()[6] {
1593                WALRecord::Checkpoint => {}
1594                _ => panic!("Expected Checkpoint"),
1595            }
1596        }
1597    }
1598
1599    // =================================================================
1600    // Test: Commit pipeline — LocalStorage flush → ShadowFile apply
1601    // =================================================================
1602    #[test]
1603    fn test_commit_pipeline_local_storage_flush() {
1604        let (sm, _dir) = setup_integration();
1605
1606        // Create table via catalog directly so we know the table_id
1607        let table_id;
1608        {
1609            let table = sm.table_catalog.create_node_table(
1610                "Person".into(),
1611                vec![
1612                    ColumnDefinition {
1613                        compression: akar_common::enums::CompressionType::Uncompressed,
1614                        name: "name".into(),
1615                        logical_type: LogicalTypeID::String,
1616                        is_primary_key: true,
1617                    },
1618                    ColumnDefinition {
1619                        compression: akar_common::enums::CompressionType::Uncompressed,
1620                        name: "age".into(),
1621                        logical_type: LogicalTypeID::Int64,
1622                        is_primary_key: false,
1623                    },
1624                ],
1625            );
1626            table_id = table.table_id;
1627        }
1628
1629        // Simulate a transaction: buffer a row in LocalStorage
1630        let mut local_storage = crate::local_storage::LocalStorage::new();
1631        {
1632            let txn_table = local_storage.get_or_create_table(table_id);
1633
1634            // Encode a row: name="Alice"(String), age=30(Int64)
1635            let mut row_bytes = Vec::new();
1636            row_bytes.push(13 /* TAG_STRING */);
1637            let name = "Alice";
1638            let name_bytes = name.as_bytes();
1639            row_bytes.extend_from_slice(&(name_bytes.len() as u32).to_le_bytes());
1640            row_bytes.extend_from_slice(name_bytes);
1641            row_bytes.push(2 /* TAG_INT64 */);
1642            row_bytes.extend_from_slice(&30i64.to_le_bytes());
1643
1644            txn_table.insert(row_bytes);
1645        }
1646
1647        assert_eq!(local_storage.len(), 1, "Should have 1 table in local storage");
1648
1649        // Commit via StorageManager
1650        let shadow = crate::shadow_file::ShadowFile::new();
1651        sm.commit_transaction(
1652            &local_storage,
1653            &shadow,
1654            -1, /* checkpoint */
1655            1,  /* txn_id */
1656            None,
1657        )
1658        .unwrap();
1659
1660        // Verify data was flushed to the table
1661        {
1662            let t = sm.table_catalog.get_node_table_by_name("Person").unwrap();
1663            assert_eq!(t.num_rows, 1, "Should have 1 row after commit");
1664            assert_eq!(t.get_value(0, 0), Some(&Value::String("Alice".into())));
1665            assert_eq!(t.get_value(0, 1), Some(&Value::Int64(30)));
1666        }
1667
1668        // Verify WAL has the commit record (and was checkpointed)
1669        {
1670            let wal = sm.wal.lock().unwrap();
1671            // After checkpoint, WAL has 1 record (Checkpoint marker)
1672            assert_eq!(wal.len(), 1, "WAL should have Checkpoint marker after commit");
1673        }
1674    }
1675
1676    // =================================================================
1677    // Test: Rollback pipeline — LocalStorage clear, no data written
1678    // =================================================================
1679    #[test]
1680    fn test_rollback_pipeline_no_data_written() {
1681        let (sm, _dir) = setup_integration();
1682        let mut local_storage = crate::local_storage::LocalStorage::new();
1683        let mut shadow = crate::shadow_file::ShadowFile::new();
1684
1685        // Buffer some data
1686        {
1687            let txn_table = local_storage.get_or_create_table(0);
1688            txn_table.insert(vec![2 /* TAG_INT64 */, 42, 0, 0, 0, 0, 0, 0, 0]);
1689        }
1690
1691        assert!(!local_storage.is_empty(), "LocalStorage should have buffered data");
1692
1693        // Rollback
1694        sm.rollback_transaction(&mut local_storage, &mut shadow, 1 /* txn_id */, &[])
1695            .unwrap();
1696
1697        // Verify buffers are cleared
1698        assert!(local_storage.is_empty(), "LocalStorage should be empty after rollback");
1699        assert!(shadow.is_empty(), "ShadowFile should be empty after rollback");
1700    }
1701
1702    // =================================================================
1703    // Test: Multiple buffered rows commit correctly
1704    // =================================================================
1705    #[test]
1706    fn test_commit_multiple_rows() {
1707        let (sm, _dir) = setup_integration();
1708        sm.create_node_table(
1709            "Item".into(),
1710            vec![
1711                ColumnDefinition {
1712                    compression: akar_common::enums::CompressionType::Uncompressed,
1713                    name: "name".into(),
1714                    logical_type: LogicalTypeID::String,
1715                    is_primary_key: true,
1716                },
1717                ColumnDefinition {
1718                    compression: akar_common::enums::CompressionType::Uncompressed,
1719                    name: "price".into(),
1720                    logical_type: LogicalTypeID::Double,
1721                    is_primary_key: false,
1722                },
1723            ],
1724        );
1725
1726        // Buffer multiple rows
1727        let mut local = crate::local_storage::LocalStorage::new();
1728        {
1729            let txn_table = local.get_or_create_table(0); // table_id = 0
1730
1731            // Row 1: "Widget", 19.99
1732            let mut row = Vec::new();
1733            row.push(13);
1734            row.extend_from_slice(&6u32.to_le_bytes());
1735            row.extend_from_slice(b"Widget");
1736            row.push(11);
1737            row.extend_from_slice(&19.99f64.to_le_bytes());
1738            txn_table.insert(row);
1739
1740            // Row 2: "Gadget", 29.99
1741            let mut row = Vec::new();
1742            row.push(13);
1743            row.extend_from_slice(&6u32.to_le_bytes());
1744            row.extend_from_slice(b"Gadget");
1745            row.push(11);
1746            row.extend_from_slice(&29.99f64.to_le_bytes());
1747            txn_table.insert(row);
1748        }
1749
1750        let shadow = crate::shadow_file::ShadowFile::new();
1751        sm.commit_transaction(&local, &shadow, 0 /* no checkpoint */, 2 /* txn_id */, None)
1752            .unwrap();
1753
1754        // Verify both rows
1755        {
1756            let t = sm.table_catalog.get_node_table_by_name("Item").unwrap();
1757            assert_eq!(t.num_rows, 2, "Should have 2 rows after commit");
1758        }
1759    }
1760    #[test]
1761    fn test_zone_map_pushdown() {
1762        use crate::column_chunk::NODE_GROUP_SIZE;
1763        use crate::table::{ColumnDefinition, NodeTable};
1764        use akar_common::types::{LogicalTypeID, Value};
1765
1766        let db_path = "test_zone_map_pushdown.db";
1767        let _ = std::fs::remove_file(db_path);
1768
1769        let mut table = NodeTable::new(
1770            0,
1771            db_path.to_string(),
1772            vec![
1773                ColumnDefinition {
1774                    compression: akar_common::enums::CompressionType::Uncompressed,
1775                    name: "id".into(),
1776                    logical_type: LogicalTypeID::Int64,
1777                    is_primary_key: true,
1778                },
1779                ColumnDefinition {
1780                    compression: akar_common::enums::CompressionType::Uncompressed,
1781                    name: "value".into(),
1782                    logical_type: LogicalTypeID::Int64,
1783                    is_primary_key: false,
1784                },
1785            ],
1786        );
1787
1788        // Insert exactly one node group of elements with values 0 to 4095
1789        for i in 0..NODE_GROUP_SIZE as i64 {
1790            table.insert_row(vec![Value::Int64(i), Value::Int64(i)]).unwrap();
1791        }
1792
1793        // Insert a second node group of elements with values 4096 to 8191
1794        for i in 0..NODE_GROUP_SIZE as i64 {
1795            let val = i + NODE_GROUP_SIZE as i64;
1796            table.insert_row(vec![Value::Int64(val), Value::Int64(val)]).unwrap();
1797        }
1798
1799        // Query with a predicate: id > 5000.
1800        // The first node group (max id = 4095) should be completely skipped.
1801        let predicate = Some((0, ">", &Value::Int64(5000)));
1802        let data = table.to_column_major_data_with_predicate(predicate);
1803
1804        // data is Vec<Vec<Value>> where data[col][row].
1805        // Total rows should be 4096 instead of 8192 because the first node group is skipped.
1806        assert_eq!(
1807            data[0].len(),
1808            NODE_GROUP_SIZE,
1809            "Only the second chunk should be returned"
1810        );
1811        assert_eq!(
1812            data[0][0],
1813            Value::Int64(NODE_GROUP_SIZE as i64),
1814            "First element should be from the second chunk"
1815        );
1816
1817        let _ = std::fs::remove_file(db_path);
1818    }
1819}