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            }
637        }
638
639        // Discard buffered writes
640        local_storage.clear();
641        shadow_file.discard();
642
643        Ok(())
644    }
645
646    /// Recover state from the WAL after a crash or unclean shutdown.
647    ///
648    /// This loads any on-disk WAL records, replays them against the
649    /// in-memory `TableCatalog`, then takes a checkpoint to reset the WAL.
650    ///
651    /// Call this once during `Database::new()`, **after** table schemas
652    /// have been re-created (e.g., from DDL replay or a persisted catalog).
653    ///
654    /// Returns the number of records recovered, or an error if recovery
655    /// fails (database is corrupt).
656    pub fn recover(&self) -> std::io::Result<usize> {
657        let mut wal = self
658            .wal
659            .lock()
660            .map_err(|e| std::io::Error::other(format!("Lock poisoned: {e}")))?;
661
662        // Load WAL records from disk
663        wal.load_from_disk()?;
664
665        if wal.is_empty() {
666            return Ok(0); // Nothing to recover
667        }
668
669        let mut data_records = 0usize;
670        let catalog = self.table_catalog.clone();
671
672        // Replay each record
673        wal.replay(|record| {
674            use crate::wal::WALRecord;
675            let cat = catalog.clone();
676
677            match record {
678                WALRecord::Insert { table_id, data } => {
679                    // Deserialize the data as Value vec and insert into the
680                    // corresponding node table.
681                    if let Some(mut table) = cat.get_node_table_mut(*table_id) {
682                        let values = deserialize_values_from_bytes(data, table.columns.len());
683                        if let Err(e) = table.insert_row(values) {
684                            return Err(std::io::Error::other(format!("WAL recovery insert failed: {e}")));
685                        }
686                        data_records += 1;
687                    }
688                }
689                WALRecord::Delete { table_id, row_id } => {
690                    if let Some(mut table) = cat.get_node_table_mut(*table_id) {
691                        if let Err(e) = table.delete_row(*row_id) {
692                            return Err(std::io::Error::other(format!("WAL recovery delete failed: {e}")));
693                        }
694                        data_records += 1;
695                    }
696                }
697                WALRecord::Update {
698                    table_id,
699                    row_id,
700                    column,
701                    data,
702                } => {
703                    if let Some(mut table) = cat.get_node_table_mut(*table_id) {
704                        let values = deserialize_values_from_bytes(data, 1);
705                        if let Some(val) = values.into_iter().next() {
706                            if let Err(e) = table.update_cell(*row_id, *column as usize, val) {
707                                return Err(std::io::Error::other(format!("WAL recovery update failed: {e}")));
708                            }
709                            data_records += 1;
710                        }
711                    }
712                }
713                WALRecord::UpdateFsm { .. } => {
714                    // UpdateFsm records are handled by the FSM recovery directly,
715                    // we can ignore them during the table-level replay.
716                }
717                WALRecord::ColumnWrite { .. } => {
718                    // ColumnWrite records are for the BufferManager-level
719                    // page writes. At the table level, data is already in
720                    // NodeGroup memory, so we skip these during recovery
721                    // (the checkpoint handles page-level persistence).
722                }
723                WALRecord::Commit { .. } | WALRecord::Rollback { .. } => {
724                    // Transaction markers — ignore during recovery since
725                    // all records in the WAL at startup are from already-
726                    // committed transactions (uncommitted ones were lost
727                    // in the crash).
728                }
729                WALRecord::Checkpoint => {
730                    // Checkpoint marker — all data before this is already
731                    // durable. We can clear what we've processed so far.
732                    // In practice, a checkpoint clears the WAL so this
733                    // marker should rarely appear during recovery.
734                }
735                WALRecord::LocalWALData { .. } => {
736                    // LocalWALData is a bulk-copied buffer from a committed
737                    // transaction's LocalWAL. The individual records within
738                    // have already been applied during normal commit flow.
739                    // At recovery time, the raw data is opaque — we skip it
740                    // since the table-level records (Insert/Delete/Update)
741                    // are replayed independently from the same transaction.
742                }
743                // DDL records — metadata-only, no data to replay for now.
744                // DDL operations (CREATE/DROP TABLE, etc.) are captured via
745                // Catalog serialization separately.
746                WALRecord::CreateTable { .. }
747                | WALRecord::DropTable { .. }
748                | WALRecord::AlterTable { .. }
749                | WALRecord::CreateIndex { .. }
750                | WALRecord::DropIndex { .. }
751                | WALRecord::CreateSequence { .. } => {
752                    // DDL records are replayed via catalog snapshot, not
753                    // individual WAL entries. Skip during table-level replay.
754                }
755            }
756            Ok(())
757        })?;
758
759        // After successful replay, mirror the recovered rows into the durable
760        // column mirrors as well, so a subsequent startup with an empty WAL
761        // restores from the mirror instead of the (now truncated) log.
762        if data_records > 0 {
763            self.persist_all_tables()
764                .map_err(|e| std::io::Error::other(format!("WAL recovery: persist tables failed: {e}")))?;
765        }
766
767        // Take a checkpoint to reset the WAL and make all recovered data durable.
768        if let Err(e) = checkpoint(&mut wal, &self.buffer_manager) {
769            tracing::warn!("WAL recovery: checkpoint after replay failed: {e}");
770        }
771
772        Ok(data_records)
773    }
774}
775
776/// Helper: deserialize binary data back into a `Vec<Value>`.
777///
778/// Each value is stored as a tag byte followed by type-specific data
779/// (see `column.rs` for the tag format). This is a simplified version
780/// that handles the common primary-key and property types.
781pub(crate) fn deserialize_values_from_bytes(data: &[u8], expected_count: usize) -> Vec<Value> {
782    use crate::column::*;
783    use std::io::Read;
784
785    if data.is_empty() || expected_count == 0 {
786        return Vec::new();
787    }
788
789    let mut cursor = std::io::Cursor::new(data);
790    let mut values = Vec::with_capacity(expected_count);
791
792    for _ in 0..expected_count {
793        let mut tag_buf = [0u8; 1];
794        if cursor.read_exact(&mut tag_buf).is_err() {
795            values.push(Value::Null);
796            continue;
797        }
798        let tag = tag_buf[0];
799
800        match tag {
801            TAG_NULL => values.push(Value::Null),
802            TAG_BOOL => {
803                let mut buf = [0u8; 1];
804                if cursor.read_exact(&mut buf).is_ok() {
805                    values.push(Value::Bool(buf[0] != 0));
806                }
807            }
808            TAG_INT64 => {
809                let mut buf = [0u8; 8];
810                if cursor.read_exact(&mut buf).is_ok() {
811                    values.push(Value::Int64(i64::from_le_bytes(buf)));
812                }
813            }
814            TAG_INT32 => {
815                let mut buf = [0u8; 4];
816                if cursor.read_exact(&mut buf).is_ok() {
817                    values.push(Value::Int32(i32::from_le_bytes(buf)));
818                }
819            }
820            TAG_DOUBLE => {
821                let mut buf = [0u8; 8];
822                if cursor.read_exact(&mut buf).is_ok() {
823                    values.push(Value::Double(f64::from_le_bytes(buf)));
824                }
825            }
826            TAG_FLOAT => {
827                let mut buf = [0u8; 4];
828                if cursor.read_exact(&mut buf).is_ok() {
829                    values.push(Value::Float(f32::from_le_bytes(buf)));
830                }
831            }
832            TAG_STRING => {
833                let mut len_buf = [0u8; 4];
834                if cursor.read_exact(&mut len_buf).is_ok() {
835                    let len = u32::from_le_bytes(len_buf) as usize;
836                    let mut str_buf = vec![0u8; len];
837                    if cursor.read_exact(&mut str_buf).is_ok()
838                        && let Ok(s) = String::from_utf8(str_buf)
839                    {
840                        values.push(Value::String(s));
841                    }
842                }
843            }
844            TAG_DATE => {
845                let mut buf = [0u8; 4];
846                if cursor.read_exact(&mut buf).is_ok() {
847                    values.push(Value::Date(akar_common::types::Date(i32::from_le_bytes(buf))));
848                }
849            }
850            TAG_TIMESTAMP => {
851                let mut buf = [0u8; 8];
852                if cursor.read_exact(&mut buf).is_ok() {
853                    values.push(Value::Timestamp(akar_common::types::Timestamp(i64::from_le_bytes(buf))));
854                }
855            }
856            TAG_INTERNAL_ID => {
857                let mut table_buf = [0u8; 8];
858                let mut offset_buf = [0u8; 8];
859                if cursor.read_exact(&mut table_buf).is_ok() && cursor.read_exact(&mut offset_buf).is_ok() {
860                    values.push(Value::InternalID(akar_common::types::InternalID {
861                        table_id: u64::from_le_bytes(table_buf),
862                        offset: u64::from_le_bytes(offset_buf),
863                    }));
864                }
865            }
866            // For any other type, skip one byte and push null
867            _ => {
868                values.push(Value::Null);
869            }
870        }
871    }
872
873    values
874}
875
876// =========================================================================
877// Phase 1 integration tests — full pipeline: table → column → buffer
878// manager → WAL → checkpoint → compression → multi-node-group
879// =========================================================================
880
881#[cfg(test)]
882mod integration_tests {
883    use super::*;
884    use crate::column::{Column, TAG_INT64, TAG_STRING};
885    use crate::page::DEFAULT_PAGE_SIZE;
886    use crate::wal::WALRecord;
887    use akar_common::enums::CompressionType;
888    use akar_common::types::{LogicalTypeID, Value};
889
890    // -----------------------------------------------------------------
891    // Helper: create a StorageManager + column pair
892    // -----------------------------------------------------------------
893    fn setup_integration() -> (StorageManager, tempfile::TempDir) {
894        let dir = tempfile::tempdir().unwrap();
895        let mm = Arc::new(MemoryManager::new(128 * 1024 * 1024));
896        let sm = StorageManager::new(dir.path().to_path_buf(), mm);
897        (sm, dir)
898    }
899
900    // =================================================================
901    // Test 1: Create table → insert rows → flush → reopen → verify
902    // =================================================================
903    #[test]
904    fn test_table_full_persistence_cycle() {
905        let (sm, _dir) = setup_integration();
906
907        // 1. Create a node table with two columns
908        let mut table = sm.create_node_table(
909            "Person".into(),
910            vec![
911                ColumnDefinition {
912                    compression: akar_common::enums::CompressionType::Uncompressed,
913                    name: "name".into(),
914                    logical_type: LogicalTypeID::String,
915                    is_primary_key: true,
916                },
917                ColumnDefinition {
918                    compression: akar_common::enums::CompressionType::Uncompressed,
919                    name: "age".into(),
920                    logical_type: LogicalTypeID::Int64,
921                    is_primary_key: false,
922                },
923            ],
924        );
925        assert_eq!(table.table_id, 0);
926
927        // 2. Insert rows into the table
928        table
929            .insert_row(vec![Value::String("Alice".into()), Value::Int64(30)])
930            .unwrap();
931        table
932            .insert_row(vec![Value::String("Bob".into()), Value::Int64(25)])
933            .unwrap();
934        table
935            .insert_row(vec![Value::String("Charlie".into()), Value::Int64(35)])
936            .unwrap();
937        assert_eq!(table.num_rows, 3);
938
939        // 3. Verify data before checkpoint
940        assert_eq!(table.get_value(0, 0), Some(&Value::String("Alice".into())));
941        assert_eq!(table.get_value(1, 1), Some(&Value::Int64(25)));
942
943        // 4. Read back via scan_column across node groups
944        let names = table.scan_column(0, 0, 3, None, &[]);
945        assert_eq!(names.len(), 3);
946        assert_eq!(names[0], Value::String("Alice".into()));
947
948        let ages = table.scan_column(1, 1, 2, None, &[]);
949        assert_eq!(ages.len(), 2);
950        assert_eq!(ages[0], Value::Int64(25));
951    }
952
953    // =================================================================
954    // Test 2: WAL crash recovery — log writes, flush to disk, replay
955    // =================================================================
956    #[test]
957    fn test_wal_recovery_cycle() {
958        let dir = tempfile::tempdir().unwrap();
959        let wal_path = dir.path().join("wal.log");
960
961        // Phase 1: Write data with WAL logging
962        #[allow(unused_variables)]
963        let (wal_records_count, column_count) = {
964            let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
965            let config = BufferManagerConfig::default();
966            let bm = Arc::new(Mutex::new(BufferManager::new(dir.path().to_path_buf(), mm, config)));
967            let mut wal = WAL::new(wal_path.clone());
968
969            let mut col = Column::new(LogicalTypeID::Int64, 0, 0, dir.path(), bm.clone(), DEFAULT_PAGE_SIZE);
970
971            // Write data and log each write to WAL
972            for i in 0i64..10 {
973                col.append_value(&Value::Int64(i)).unwrap();
974                wal.log_column_write(0, 0, 0, &i.to_le_bytes());
975            }
976            wal.append(WALRecord::Commit { transaction_id: 1 });
977            let count = wal.len();
978
979            // Flush WAL to disk
980            wal.flush_to_disk().unwrap();
981
982            // Also flush BM pages
983            {
984                let mut bm_lock = bm.lock().unwrap();
985                bm_lock.flush_all().unwrap();
986            }
987
988            // Verify column data is correct before "crash"
989            for i in 0i64..10 {
990                let v = col.get_value(i as u64).unwrap();
991                assert_eq!(v, Value::Int64(i), "Pre-crash data mismatch at {}", i);
992            }
993
994            (count, 10)
995        }; // Drop everything — simulate crash
996
997        // Phase 2: Verify the on-disk WAL file exists and has content
998        assert!(wal_path.exists(), "WAL file should exist after flush");
999        let file_len = std::fs::metadata(&wal_path).unwrap().len();
1000        assert!(file_len > 0, "WAL file should have content, got {} bytes", file_len);
1001
1002        // Verify that a fresh WAL created from the file would contain
1003        // the right number of records. Since WAL::new() starts in-memory,
1004        // we create a new one, and verify the file has the right data.
1005        // In a real recovery scenario, we'd implement WAL::load_from_disk().
1006        assert_eq!(
1007            wal_records_count, 11,
1008            "Expected 10 ColumnWrite + 1 Commit = 11 records, got {}",
1009            wal_records_count
1010        );
1011        assert_eq!(column_count, 10);
1012    }
1013
1014    // =================================================================
1015    // Test 3: Compression round-trip — write compressed → read back
1016    // =================================================================
1017    #[test]
1018    fn test_compression_full_roundtrip() {
1019        // Test IntegerBitpacking: small values compress, large values preserved
1020        let dir = tempfile::tempdir().unwrap();
1021        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1022        let config = BufferManagerConfig::default();
1023        let bm = Arc::new(Mutex::new(BufferManager::new(dir.path().to_path_buf(), mm, config)));
1024
1025        let mut col_int = Column::with_compression(
1026            LogicalTypeID::Int64,
1027            0,
1028            0,
1029            dir.path(),
1030            bm.clone(),
1031            DEFAULT_PAGE_SIZE,
1032            CompressionType::IntegerBitpacking,
1033        );
1034
1035        // Write a range of values from small to large
1036        let test_values: Vec<i64> = vec![0, 1, 42, 127, 255, 65535, 1_000_000, i64::MAX, i64::MIN, -1];
1037        for v in &test_values {
1038            col_int.append_value(&Value::Int64(*v)).unwrap();
1039        }
1040
1041        // Read back and verify
1042        for (i, expected) in test_values.iter().enumerate() {
1043            let v = col_int.get_value(i as u64).unwrap();
1044            assert_eq!(v, Value::Int64(*expected), "IntegerBitpacking mismatch at index {}", i);
1045        }
1046
1047        // Verify that setting compression doesn't break existing data
1048        let mut col_float = Column::with_compression(
1049            LogicalTypeID::Double,
1050            0,
1051            1,
1052            dir.path(),
1053            bm.clone(),
1054            DEFAULT_PAGE_SIZE,
1055            CompressionType::Float,
1056        );
1057
1058        let floats: Vec<f64> = vec![1.0, std::f64::consts::PI, -2.5e10, 0.0, f64::MIN_POSITIVE, f64::MAX];
1059        for v in &floats {
1060            col_float.append_value(&Value::Double(*v)).unwrap();
1061        }
1062
1063        for (i, expected) in floats.iter().enumerate() {
1064            let v = col_float.get_value(i as u64).unwrap();
1065            match v {
1066                Value::Double(d) => assert!(
1067                    (d - expected).abs() < 1e-10 || (d / expected - 1.0).abs() < 1e-10,
1068                    "Float compression mismatch at {}: got {}, expected {}",
1069                    i,
1070                    d,
1071                    expected
1072                ),
1073                _ => panic!("Expected Double, got {:?}", v),
1074            }
1075        }
1076
1077        // Write compressed values via existing Column API and verify roundtrip
1078        // with buffer manager flush
1079        col_int.flush().unwrap();
1080        col_float.flush().unwrap();
1081
1082        // Read again after flush
1083        for (i, expected) in test_values.iter().enumerate() {
1084            let v = col_int.get_value(i as u64).unwrap();
1085            assert_eq!(
1086                v,
1087                Value::Int64(*expected),
1088                "After flush: IntegerBitpacking mismatch at {}",
1089                i
1090            );
1091        }
1092        for (i, expected) in floats.iter().enumerate() {
1093            let v = col_float.get_value(i as u64).unwrap();
1094            match v {
1095                Value::Double(d) => assert!(
1096                    (d - expected).abs() < 1e-10 || (d / expected - 1.0).abs() < 1e-10,
1097                    "After flush: Float mismatch at {}",
1098                    i
1099                ),
1100                _ => panic!("Expected Double after flush"),
1101            }
1102        }
1103    }
1104
1105    // =================================================================
1106    // Test 4: Multi-node-group scan — insert > NODE_GROUP_SIZE rows
1107    // =================================================================
1108    #[test]
1109    fn test_multi_node_group_scan() {
1110        let (_sm, _dir) = setup_integration();
1111
1112        // Create a NodeTable (not via StorageManager, directly for test control)
1113        let mut table = NodeTable::new(
1114            1,
1115            "BigTable".into(),
1116            vec![
1117                ColumnDefinition {
1118                    compression: akar_common::enums::CompressionType::Uncompressed,
1119                    name: "id".into(),
1120                    logical_type: LogicalTypeID::Int64,
1121                    is_primary_key: false,
1122                },
1123                ColumnDefinition {
1124                    compression: akar_common::enums::CompressionType::Uncompressed,
1125                    name: "value".into(),
1126                    logical_type: LogicalTypeID::Int64,
1127                    is_primary_key: false,
1128                },
1129            ],
1130        );
1131
1132        // Insert NODE_GROUP_SIZE + 500 rows to span across multiple node groups
1133        let total_rows = NODE_GROUP_SIZE + 500;
1134        for i in 0..total_rows {
1135            table
1136                .insert_row(vec![Value::Int64(i as i64), Value::Int64((i * 2) as i64)])
1137                .unwrap();
1138        }
1139
1140        // Verify total row count
1141        assert_eq!(table.num_rows, total_rows as u64);
1142
1143        // Verify multiple node groups were created
1144        let expected_groups = 2; // 4096 fits in first group, 500 in second
1145        assert_eq!(
1146            table.node_groups.len(),
1147            expected_groups,
1148            "Expected {} node groups for {} rows",
1149            expected_groups,
1150            total_rows
1151        );
1152
1153        // Verify node group boundaries
1154        assert_eq!(table.node_groups[0].num_nodes, NODE_GROUP_SIZE as u64);
1155        assert_eq!(table.node_groups[1].num_nodes, 500);
1156        assert_eq!(table.node_groups[0].start_offset, 0);
1157        assert_eq!(table.node_groups[1].start_offset, NODE_GROUP_SIZE as u64);
1158
1159        // Verify scanning across group boundaries
1160        // Row at boundary: last row of group 0
1161        let row_at_boundary = (NODE_GROUP_SIZE - 1) as u64;
1162        assert_eq!(
1163            table.get_value(row_at_boundary as usize, 0),
1164            Some(&Value::Int64(row_at_boundary as i64))
1165        );
1166
1167        // First row of group 1
1168        let row_in_group1 = NODE_GROUP_SIZE as u64;
1169        assert_eq!(
1170            table.get_value(row_in_group1 as usize, 0),
1171            Some(&Value::Int64(row_in_group1 as i64))
1172        );
1173
1174        // Scan column 0 across the entire table
1175        let scanned = table.scan_column(0, 0, total_rows as u64, None, &[]);
1176        assert_eq!(scanned.len(), total_rows);
1177        assert_eq!(scanned[0], Value::Int64(0));
1178        assert_eq!(scanned[NODE_GROUP_SIZE], Value::Int64(NODE_GROUP_SIZE as i64));
1179        assert_eq!(scanned[total_rows - 1], Value::Int64((total_rows - 1) as i64));
1180
1181        // Scan column 1 with offset and count spanning both groups
1182        let scan_mid = table.scan_column(1, (NODE_GROUP_SIZE - 100) as u64, 200, None, &[]);
1183        assert_eq!(scan_mid.len(), 200);
1184        assert_eq!(scan_mid[0], Value::Int64(((NODE_GROUP_SIZE - 100) * 2) as i64));
1185        assert_eq!(scan_mid[199], Value::Int64(((NODE_GROUP_SIZE + 99) * 2) as i64));
1186
1187        // Verify to_column_major_data correctness
1188        let data = table.to_column_major_data();
1189        assert_eq!(data.len(), 2); // 2 columns
1190        assert_eq!(data[0].len(), total_rows);
1191        assert_eq!(data[1].len(), total_rows);
1192        assert_eq!(data[0][NODE_GROUP_SIZE], Value::Int64(NODE_GROUP_SIZE as i64));
1193        assert_eq!(data[1][0], Value::Int64(0));
1194        assert_eq!(data[1][total_rows - 1], Value::Int64(((total_rows - 1) * 2) as i64));
1195    }
1196
1197    // =================================================================
1198    // Test 5: Combined — WAL-logged compressed multi-node-group write
1199    // =================================================================
1200    #[test]
1201    fn test_compressed_multi_group_with_checkpoint() {
1202        let dir = tempfile::tempdir().unwrap();
1203        let mm = Arc::new(MemoryManager::new(128 * 1024 * 1024));
1204
1205        // Use explicit BM + WAL for full control
1206        let config = BufferManagerConfig::default();
1207        let bm = Arc::new(Mutex::new(BufferManager::new(dir.path().to_path_buf(), mm, config)));
1208        let wal_path = dir.path().join("wal.log");
1209        let mut wal = WAL::new(wal_path);
1210
1211        let mut col = Column::with_compression(
1212            LogicalTypeID::Int64,
1213            0,
1214            0,
1215            dir.path(),
1216            bm.clone(),
1217            DEFAULT_PAGE_SIZE,
1218            CompressionType::IntegerBitpacking,
1219        );
1220
1221        // Write enough values to span multiple pages
1222        let num_values = 500;
1223        for i in 0i64..num_values {
1224            col.append_value(&Value::Int64(i)).unwrap();
1225            wal.log_column_write(0, 0, 0, &i.to_le_bytes());
1226        }
1227        wal.append(WALRecord::Commit { transaction_id: 1 });
1228
1229        // Read back before checkpoint
1230        for i in 0i64..num_values {
1231            let v = col.get_value(i as u64).unwrap();
1232            assert_eq!(v, Value::Int64(i), "Pre-checkpoint mismatch at {}", i);
1233        }
1234
1235        // Checkpoint: flush WAL + dirty pages
1236        let mut bm_lock = bm.lock().unwrap();
1237        bm_lock.flush_all().unwrap();
1238        drop(bm_lock);
1239
1240        wal.flush_to_disk().unwrap();
1241
1242        // Read back after checkpoint
1243        for i in 0i64..num_values {
1244            let v = col.get_value(i as u64).unwrap();
1245            assert_eq!(v, Value::Int64(i), "Post-checkpoint mismatch at {}", i);
1246        }
1247
1248        // Verify multiple pages were allocated
1249        assert!(
1250            col.num_pages > 1,
1251            "Expected multiple pages for {} values, got {}",
1252            num_values,
1253            col.num_pages
1254        );
1255    }
1256
1257    // =================================================================
1258    // Test 6: Stress — 10k rows via column with checkpoint
1259    // =================================================================
1260    #[test]
1261    fn test_10k_row_stress() {
1262        let dir = tempfile::tempdir().unwrap();
1263        let mm = Arc::new(MemoryManager::new(256 * 1024 * 1024));
1264        let config = BufferManagerConfig::default();
1265        let bm = Arc::new(Mutex::new(BufferManager::new(dir.path().to_path_buf(), mm, config)));
1266
1267        let mut col = Column::new(LogicalTypeID::Int64, 0, 0, dir.path(), bm.clone(), DEFAULT_PAGE_SIZE);
1268
1269        // Write 10,000 values
1270        for i in 0i64..10_000 {
1271            col.append_value(&Value::Int64(i)).unwrap();
1272        }
1273        assert_eq!(col.num_values, 10_000);
1274
1275        // Read back all values
1276        for i in 0i64..10_000 {
1277            let v = col.get_value(i as u64).unwrap();
1278            assert_eq!(v, Value::Int64(i), "Stress test mismatch at {}", i);
1279        }
1280
1281        // Flush and re-verify
1282        col.flush().unwrap();
1283        for i in 0i64..10_000 {
1284            let v = col.get_value(i as u64).unwrap();
1285            assert_eq!(v, Value::Int64(i), "Post-flush stress mismatch at {}", i);
1286        }
1287
1288        // Verify multiple pages were used
1289        assert!(
1290            col.num_pages > 1,
1291            "Stress test should use multiple pages, got {}",
1292            col.num_pages
1293        );
1294    }
1295
1296    // =================================================================
1297    // Test 7: WAL recovery — insert data, simulate crash, recover
1298    // =================================================================
1299    #[test]
1300    fn test_wal_recovery_insert_then_recover() {
1301        let dir = tempfile::tempdir().unwrap();
1302
1303        // Phase 1: Create DB, insert data, flush WAL, then "crash"
1304        let _row_count = {
1305            let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1306            let sm = StorageManager::new(dir.path().to_path_buf(), mm);
1307
1308            // Create a node table
1309            let mut table = sm.create_node_table(
1310                "Person".into(),
1311                vec![
1312                    ColumnDefinition {
1313                        compression: akar_common::enums::CompressionType::Uncompressed,
1314                        name: "name".into(),
1315                        logical_type: LogicalTypeID::String,
1316                        is_primary_key: true,
1317                    },
1318                    ColumnDefinition {
1319                        compression: akar_common::enums::CompressionType::Uncompressed,
1320                        name: "age".into(),
1321                        logical_type: LogicalTypeID::Int64,
1322                        is_primary_key: false,
1323                    },
1324                ],
1325            );
1326
1327            // Insert rows
1328            table
1329                .insert_row(vec![Value::String("Alice".into()), Value::Int64(30)])
1330                .unwrap();
1331            table
1332                .insert_row(vec![Value::String("Bob".into()), Value::Int64(25)])
1333                .unwrap();
1334            table
1335                .insert_row(vec![Value::String("Charlie".into()), Value::Int64(35)])
1336                .unwrap();
1337
1338            // Put the table back into the catalog so WAL knows about it.
1339            // We re-create the table entry via the catalog API.
1340            {
1341                // Re-create the table entry with the same schema so the
1342                // catalog has a valid entry for recovery to target.
1343                sm.table_catalog.create_node_table(
1344                    "Person".into(),
1345                    vec![
1346                        ColumnDefinition {
1347                            compression: akar_common::enums::CompressionType::Uncompressed,
1348                            name: "name".into(),
1349                            logical_type: LogicalTypeID::String,
1350                            is_primary_key: true,
1351                        },
1352                        ColumnDefinition {
1353                            compression: akar_common::enums::CompressionType::Uncompressed,
1354                            name: "age".into(),
1355                            logical_type: LogicalTypeID::Int64,
1356                            is_primary_key: false,
1357                        },
1358                    ],
1359                );
1360            }
1361
1362            let count = table.num_rows;
1363            assert_eq!(count, 3);
1364
1365            // Write WAL records and flush to disk
1366            {
1367                let mut wal = sm.wal.lock().unwrap();
1368                wal.append(WALRecord::Insert {
1369                    table_id: table.table_id,
1370                    data: vec![
1371                        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,
1372                    ],
1373                });
1374                wal.append(WALRecord::Insert {
1375                    table_id: table.table_id,
1376                    data: vec![
1377                        TAG_STRING, 3, 0, 0, 0, b'B', b'o', b'b', TAG_INT64, 25, 0, 0, 0, 0, 0, 0, 0,
1378                    ],
1379                });
1380                wal.append(WALRecord::Insert {
1381                    table_id: table.table_id,
1382                    data: vec![
1383                        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,
1384                        0, 0,
1385                    ],
1386                });
1387                wal.flush_to_disk().unwrap();
1388            }
1389
1390            count
1391        }; // "Crash" — all state dropped
1392
1393        // Phase 2: Recover — create a new StorageManager, it should NOT delete
1394        // the WAL. Then manually trigger recovery.
1395        {
1396            let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1397            let sm = StorageManager::new(dir.path().to_path_buf(), mm);
1398
1399            // Verify WAL file exists and has records
1400            let wal_path = dir.path().join("wal.log");
1401            assert!(wal_path.exists(), "WAL file should exist for recovery");
1402
1403            // Create the same table schema so recovery has a target
1404            // (the catalog create_node_table stores the table internally).
1405            sm.create_node_table(
1406                "Person".into(),
1407                vec![
1408                    ColumnDefinition {
1409                        compression: akar_common::enums::CompressionType::Uncompressed,
1410                        name: "name".into(),
1411                        logical_type: LogicalTypeID::String,
1412                        is_primary_key: true,
1413                    },
1414                    ColumnDefinition {
1415                        compression: akar_common::enums::CompressionType::Uncompressed,
1416                        name: "age".into(),
1417                        logical_type: LogicalTypeID::Int64,
1418                        is_primary_key: false,
1419                    },
1420                ],
1421            );
1422
1423            // Recover — the WAL has table_id = 0 (the first table created).
1424            let recovered = sm.recover().unwrap();
1425            assert_eq!(recovered, 3, "Should recover 3 WAL records");
1426
1427            // Verify data survived — the table was re-created empty,
1428            // and recovery inserted exactly 3 rows from the WAL.
1429            {
1430                let recovered_table = sm.table_catalog.get_node_table_by_name("Person").unwrap();
1431                assert_eq!(recovered_table.num_rows, 3, "Should have recovered 3 rows");
1432                assert_eq!(recovered_table.get_value(0, 0), Some(&Value::String("Alice".into())));
1433                assert_eq!(recovered_table.get_value(1, 0), Some(&Value::String("Bob".into())));
1434                assert_eq!(recovered_table.get_value(2, 0), Some(&Value::String("Charlie".into())));
1435                assert_eq!(recovered_table.get_value(0, 1), Some(&Value::Int64(30)));
1436                assert_eq!(recovered_table.get_value(1, 1), Some(&Value::Int64(25)));
1437                assert_eq!(recovered_table.get_value(2, 1), Some(&Value::Int64(35)));
1438            }
1439
1440            // Verify WAL was checkpointed — after checkpoint the WAL has
1441            // exactly 1 record (the Checkpoint marker).
1442            {
1443                let wal = sm.wal.lock().unwrap();
1444                assert_eq!(
1445                    wal.len(),
1446                    1,
1447                    "WAL should have only the checkpoint marker after recovery"
1448                );
1449                assert!(matches!(wal.records()[0], crate::wal::WALRecord::Checkpoint));
1450            }
1451        }
1452    }
1453
1454    // =================================================================
1455    // Test 8: WAL recovery — no-op when no WAL exists
1456    // =================================================================
1457    #[test]
1458    fn test_wal_recovery_no_wal() {
1459        let dir = tempfile::tempdir().unwrap();
1460
1461        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1462        let sm = StorageManager::new(dir.path().to_path_buf(), mm);
1463
1464        // No WAL file = recovery returns 0
1465        let recovered = sm.recover().unwrap();
1466        assert_eq!(recovered, 0, "No WAL = no records recovered");
1467    }
1468
1469    // =================================================================
1470    // Test 9: WAL recovery — empty WAL file
1471    // =================================================================
1472    #[test]
1473    fn test_wal_recovery_empty_wal() {
1474        let dir = tempfile::tempdir().unwrap();
1475
1476        // Create an empty WAL file on disk
1477        let wal_path = dir.path().join("wal.log");
1478        std::fs::write(&wal_path, b"").unwrap();
1479
1480        let mm = Arc::new(MemoryManager::new(64 * 1024 * 1024));
1481        let sm = StorageManager::new(dir.path().to_path_buf(), mm);
1482
1483        let recovered = sm.recover().unwrap();
1484        assert_eq!(recovered, 0, "Empty WAL = no records recovered");
1485    }
1486
1487    // =================================================================
1488    // Test 10: WAL load_from_disk roundtrip
1489    // =================================================================
1490    #[test]
1491    fn test_wal_load_from_disk_roundtrip() {
1492        use crate::wal::WALRecord;
1493        let dir = tempfile::tempdir().unwrap();
1494        let wal_path = dir.path().join("wal.log");
1495
1496        // Write records
1497        {
1498            let mut wal = WAL::new(wal_path.clone());
1499            wal.append(WALRecord::Insert {
1500                table_id: 42,
1501                data: vec![1, 2, 3, 4],
1502            });
1503            wal.append(WALRecord::Delete {
1504                table_id: 42,
1505                row_id: 0,
1506            });
1507            wal.append(WALRecord::Update {
1508                table_id: 42,
1509                row_id: 1,
1510                column: 2,
1511                data: vec![5, 6],
1512            });
1513            wal.append(WALRecord::ColumnWrite {
1514                table_id: 42,
1515                col_id: 0,
1516                page_id: 1,
1517                data: vec![7, 8, 9],
1518            });
1519            wal.append(WALRecord::Commit { transaction_id: 100 });
1520            wal.append(WALRecord::Rollback { transaction_id: 101 });
1521            wal.append(WALRecord::Checkpoint);
1522            wal.flush_to_disk().unwrap();
1523        }
1524
1525        // Load from disk
1526        {
1527            let mut wal = WAL::new(wal_path.clone());
1528            wal.load_from_disk().unwrap();
1529            assert_eq!(wal.len(), 7, "Should load 7 records from disk");
1530            assert!(wal.is_dirty());
1531
1532            // Verify each record type
1533            match &wal.records()[0] {
1534                WALRecord::Insert { table_id, data } => {
1535                    assert_eq!(*table_id, 42);
1536                    assert_eq!(data, &[1, 2, 3, 4]);
1537                }
1538                _ => panic!("Expected Insert"),
1539            }
1540            match &wal.records()[1] {
1541                WALRecord::Delete { table_id, row_id } => {
1542                    assert_eq!(*table_id, 42);
1543                    assert_eq!(*row_id, 0);
1544                }
1545                _ => panic!("Expected Delete"),
1546            }
1547            match &wal.records()[4] {
1548                WALRecord::Commit { transaction_id } => {
1549                    assert_eq!(*transaction_id, 100);
1550                }
1551                _ => panic!("Expected Commit"),
1552            }
1553            match &wal.records()[5] {
1554                WALRecord::Rollback { transaction_id } => {
1555                    assert_eq!(*transaction_id, 101);
1556                }
1557                _ => panic!("Expected Rollback"),
1558            }
1559            match &wal.records()[6] {
1560                WALRecord::Checkpoint => {}
1561                _ => panic!("Expected Checkpoint"),
1562            }
1563        }
1564    }
1565
1566    // =================================================================
1567    // Test: Commit pipeline — LocalStorage flush → ShadowFile apply
1568    // =================================================================
1569    #[test]
1570    fn test_commit_pipeline_local_storage_flush() {
1571        let (sm, _dir) = setup_integration();
1572
1573        // Create table via catalog directly so we know the table_id
1574        let table_id;
1575        {
1576            let table = sm.table_catalog.create_node_table(
1577                "Person".into(),
1578                vec![
1579                    ColumnDefinition {
1580                        compression: akar_common::enums::CompressionType::Uncompressed,
1581                        name: "name".into(),
1582                        logical_type: LogicalTypeID::String,
1583                        is_primary_key: true,
1584                    },
1585                    ColumnDefinition {
1586                        compression: akar_common::enums::CompressionType::Uncompressed,
1587                        name: "age".into(),
1588                        logical_type: LogicalTypeID::Int64,
1589                        is_primary_key: false,
1590                    },
1591                ],
1592            );
1593            table_id = table.table_id;
1594        }
1595
1596        // Simulate a transaction: buffer a row in LocalStorage
1597        let mut local_storage = crate::local_storage::LocalStorage::new();
1598        {
1599            let txn_table = local_storage.get_or_create_table(table_id);
1600
1601            // Encode a row: name="Alice"(String), age=30(Int64)
1602            let mut row_bytes = Vec::new();
1603            row_bytes.push(13 /* TAG_STRING */);
1604            let name = "Alice";
1605            let name_bytes = name.as_bytes();
1606            row_bytes.extend_from_slice(&(name_bytes.len() as u32).to_le_bytes());
1607            row_bytes.extend_from_slice(name_bytes);
1608            row_bytes.push(2 /* TAG_INT64 */);
1609            row_bytes.extend_from_slice(&30i64.to_le_bytes());
1610
1611            txn_table.insert(row_bytes);
1612        }
1613
1614        assert_eq!(local_storage.len(), 1, "Should have 1 table in local storage");
1615
1616        // Commit via StorageManager
1617        let shadow = crate::shadow_file::ShadowFile::new();
1618        sm.commit_transaction(
1619            &local_storage,
1620            &shadow,
1621            -1, /* checkpoint */
1622            1,  /* txn_id */
1623            None,
1624        )
1625        .unwrap();
1626
1627        // Verify data was flushed to the table
1628        {
1629            let t = sm.table_catalog.get_node_table_by_name("Person").unwrap();
1630            assert_eq!(t.num_rows, 1, "Should have 1 row after commit");
1631            assert_eq!(t.get_value(0, 0), Some(&Value::String("Alice".into())));
1632            assert_eq!(t.get_value(0, 1), Some(&Value::Int64(30)));
1633        }
1634
1635        // Verify WAL has the commit record (and was checkpointed)
1636        {
1637            let wal = sm.wal.lock().unwrap();
1638            // After checkpoint, WAL has 1 record (Checkpoint marker)
1639            assert_eq!(wal.len(), 1, "WAL should have Checkpoint marker after commit");
1640        }
1641    }
1642
1643    // =================================================================
1644    // Test: Rollback pipeline — LocalStorage clear, no data written
1645    // =================================================================
1646    #[test]
1647    fn test_rollback_pipeline_no_data_written() {
1648        let (sm, _dir) = setup_integration();
1649        let mut local_storage = crate::local_storage::LocalStorage::new();
1650        let mut shadow = crate::shadow_file::ShadowFile::new();
1651
1652        // Buffer some data
1653        {
1654            let txn_table = local_storage.get_or_create_table(0);
1655            txn_table.insert(vec![2 /* TAG_INT64 */, 42, 0, 0, 0, 0, 0, 0, 0]);
1656        }
1657
1658        assert!(!local_storage.is_empty(), "LocalStorage should have buffered data");
1659
1660        // Rollback
1661        sm.rollback_transaction(&mut local_storage, &mut shadow, 1 /* txn_id */, &[])
1662            .unwrap();
1663
1664        // Verify buffers are cleared
1665        assert!(local_storage.is_empty(), "LocalStorage should be empty after rollback");
1666        assert!(shadow.is_empty(), "ShadowFile should be empty after rollback");
1667    }
1668
1669    // =================================================================
1670    // Test: Multiple buffered rows commit correctly
1671    // =================================================================
1672    #[test]
1673    fn test_commit_multiple_rows() {
1674        let (sm, _dir) = setup_integration();
1675        sm.create_node_table(
1676            "Item".into(),
1677            vec![
1678                ColumnDefinition {
1679                    compression: akar_common::enums::CompressionType::Uncompressed,
1680                    name: "name".into(),
1681                    logical_type: LogicalTypeID::String,
1682                    is_primary_key: true,
1683                },
1684                ColumnDefinition {
1685                    compression: akar_common::enums::CompressionType::Uncompressed,
1686                    name: "price".into(),
1687                    logical_type: LogicalTypeID::Double,
1688                    is_primary_key: false,
1689                },
1690            ],
1691        );
1692
1693        // Buffer multiple rows
1694        let mut local = crate::local_storage::LocalStorage::new();
1695        {
1696            let txn_table = local.get_or_create_table(0); // table_id = 0
1697
1698            // Row 1: "Widget", 19.99
1699            let mut row = Vec::new();
1700            row.push(13);
1701            row.extend_from_slice(&6u32.to_le_bytes());
1702            row.extend_from_slice(b"Widget");
1703            row.push(11);
1704            row.extend_from_slice(&19.99f64.to_le_bytes());
1705            txn_table.insert(row);
1706
1707            // Row 2: "Gadget", 29.99
1708            let mut row = Vec::new();
1709            row.push(13);
1710            row.extend_from_slice(&6u32.to_le_bytes());
1711            row.extend_from_slice(b"Gadget");
1712            row.push(11);
1713            row.extend_from_slice(&29.99f64.to_le_bytes());
1714            txn_table.insert(row);
1715        }
1716
1717        let shadow = crate::shadow_file::ShadowFile::new();
1718        sm.commit_transaction(&local, &shadow, 0 /* no checkpoint */, 2 /* txn_id */, None)
1719            .unwrap();
1720
1721        // Verify both rows
1722        {
1723            let t = sm.table_catalog.get_node_table_by_name("Item").unwrap();
1724            assert_eq!(t.num_rows, 2, "Should have 2 rows after commit");
1725        }
1726    }
1727    #[test]
1728    fn test_zone_map_pushdown() {
1729        use crate::column_chunk::NODE_GROUP_SIZE;
1730        use crate::table::{ColumnDefinition, NodeTable};
1731        use akar_common::types::{LogicalTypeID, Value};
1732
1733        let db_path = "test_zone_map_pushdown.db";
1734        let _ = std::fs::remove_file(db_path);
1735
1736        let mut table = NodeTable::new(
1737            0,
1738            db_path.to_string(),
1739            vec![
1740                ColumnDefinition {
1741                    compression: akar_common::enums::CompressionType::Uncompressed,
1742                    name: "id".into(),
1743                    logical_type: LogicalTypeID::Int64,
1744                    is_primary_key: true,
1745                },
1746                ColumnDefinition {
1747                    compression: akar_common::enums::CompressionType::Uncompressed,
1748                    name: "value".into(),
1749                    logical_type: LogicalTypeID::Int64,
1750                    is_primary_key: false,
1751                },
1752            ],
1753        );
1754
1755        // Insert exactly one node group of elements with values 0 to 4095
1756        for i in 0..NODE_GROUP_SIZE as i64 {
1757            table.insert_row(vec![Value::Int64(i), Value::Int64(i)]).unwrap();
1758        }
1759
1760        // Insert a second node group of elements with values 4096 to 8191
1761        for i in 0..NODE_GROUP_SIZE as i64 {
1762            let val = i + NODE_GROUP_SIZE as i64;
1763            table.insert_row(vec![Value::Int64(val), Value::Int64(val)]).unwrap();
1764        }
1765
1766        // Query with a predicate: id > 5000.
1767        // The first node group (max id = 4095) should be completely skipped.
1768        let predicate = Some((0, ">", &Value::Int64(5000)));
1769        let data = table.to_column_major_data_with_predicate(predicate);
1770
1771        // data is Vec<Vec<Value>> where data[col][row].
1772        // Total rows should be 4096 instead of 8192 because the first node group is skipped.
1773        assert_eq!(
1774            data[0].len(),
1775            NODE_GROUP_SIZE,
1776            "Only the second chunk should be returned"
1777        );
1778        assert_eq!(
1779            data[0][0],
1780            Value::Int64(NODE_GROUP_SIZE as i64),
1781            "First element should be from the second chunk"
1782        );
1783
1784        let _ = std::fs::remove_file(db_path);
1785    }
1786}