Merge 'implement the MaxPgCount opcode' from Glauber Costa

It is used by the pragma max_page_count, which is also implemented.

Reviewed-by: Preston Thorpe <preston@turso.tech>

Closes #2472

COMPAT.md

@@ -142,7 +142,7 @@ Turso aims to be fully compatible with SQLite, with opt-in features not supporte
 | PRAGMA legacy_alter_table        | No         |                                              |
 | PRAGMA legacy_file_format         | Yes        |                                              |
 | PRAGMA locking_mode              | No         |                                              |
-| PRAGMA max_page_count            | No         |                                              |
+| PRAGMA max_page_count            | Yes        |                                              |
 | PRAGMA mmap_size                 | No         |                                              |
 | PRAGMA module_list               | No         |                                              |
 | PRAGMA optimize                  | No         |                                              |
@@ -479,14 +479,14 @@ Modifiers:
 | IntegrityCk    | Yes    |         |
 | IsNull         | Yes    |         |
 | IsUnique       | No     |         |
-| JournalMode    | No     |         |
+| JournalMode    | Yes    |         |
 | Jump           | Yes    |         |
 | Last           | Yes    |         |
 | Le             | Yes    |         |
 | LoadAnalysis   | No     |         |
 | Lt             | Yes    |         |
 | MakeRecord     | Yes    |         |
-| MaxPgcnt       | No     |         |
+| MaxPgcnt       | Yes    |         |
 | MemMax         | No     |         |
 | Move           | Yes    |         |
 | Multiply       | Yes    |         |

core/pragma.rs

@@ -66,6 +66,13 @@ pub fn pragma_for(pragma: &PragmaName) -> Pragma {
             PragmaFlags::Result0 | PragmaFlags::SchemaReq | PragmaFlags::NoColumns1,
             &["page_size"],
         ),
+        MaxPageCount => Pragma::new(
+            PragmaFlags::NeedSchema
+                | PragmaFlags::Result0
+                | PragmaFlags::SchemaReq
+                | PragmaFlags::NoColumns1,
+            &["max_page_count"],
+        ),
         SchemaVersion => Pragma::new(
             PragmaFlags::NoColumns1 | PragmaFlags::Result0,
             &["schema_version"],

core/storage/pager.rs

@@ -24,6 +24,9 @@ use super::page_cache::{CacheError, CacheResizeResult, DumbLruPageCache, PageCac
 use super::sqlite3_ondisk::begin_write_btree_page;
 use super::wal::CheckpointMode;
 
+/// SQLite's default maximum page count
+const DEFAULT_MAX_PAGE_COUNT: u32 = 0xfffffffe;
+
 #[cfg(not(feature = "omit_autovacuum"))]
 use ptrmap::*;
 
@@ -450,6 +453,8 @@ pub struct Pager {
     pub(crate) page_size: Cell<Option<u32>>,
     reserved_space: OnceCell<u8>,
     free_page_state: RefCell<FreePageState>,
+    /// Maximum number of pages allowed in the database. Default is 1073741823 (SQLite default).
+    max_page_count: Cell<u32>,
     #[cfg(not(feature = "omit_autovacuum"))]
     /// State machine for [Pager::ptrmap_get]
     ptrmap_get_state: RefCell<PtrMapGetState>,
@@ -562,6 +567,7 @@ impl Pager {
             }),
             free_page_state: RefCell::new(FreePageState::Start),
             allocate_page_state: RefCell::new(AllocatePageState::Start),
+            max_page_count: Cell::new(DEFAULT_MAX_PAGE_COUNT),
             #[cfg(not(feature = "omit_autovacuum"))]
             ptrmap_get_state: RefCell::new(PtrMapGetState::Start),
             #[cfg(not(feature = "omit_autovacuum"))]
@@ -572,6 +578,26 @@ impl Pager {
         })
     }
 
+    /// Get the maximum page count for this database
+    pub fn get_max_page_count(&self) -> u32 {
+        self.max_page_count.get()
+    }
+
+    /// Set the maximum page count for this database
+    /// Returns the new maximum page count (may be clamped to current database size)
+    pub fn set_max_page_count(&self, new_max: u32) -> crate::Result<IOResult<u32>> {
+        // Get current database size
+        let current_page_count = match self.with_header(|header| header.database_size.get())? {
+            IOResult::Done(size) => size,
+            IOResult::IO => return Ok(IOResult::IO),
+        };
+
+        // Clamp new_max to be at least the current database size
+        let clamped_max = std::cmp::max(new_max, current_page_count);
+        self.max_page_count.set(clamped_max);
+        Ok(IOResult::Done(clamped_max))
+    }
+
     pub fn set_wal(&mut self, wal: Rc<RefCell<dyn Wal>>) {
         self.wal = Some(wal);
     }
@@ -1993,6 +2019,15 @@ impl Pager {
                 }
                 AllocatePageState::AllocateNewPage { current_db_size } => {
                     let new_db_size = *current_db_size + 1;
+
+                    // Check if allocating a new page would exceed the maximum page count
+                    let max_page_count = self.get_max_page_count();
+                    if new_db_size > max_page_count {
+                        return Err(LimboError::DatabaseFull(
+                            "database or disk is full".to_string(),
+                        ));
+                    }
+
                     // FIXME: should reserve page cache entry before modifying the database
                     let page = allocate_new_page(new_db_size as usize, &self.buffer_pool, 0);
                     {

core/translate/pragma.rs

@@ -120,14 +120,23 @@ fn update_pragma(
             let year = chrono::Local::now().year();
             bail_parse_error!("It's {year}. UTF-8 won.");
         }
-        PragmaName::JournalMode => query_pragma(
-            PragmaName::JournalMode,
-            schema,
-            None,
-            pager,
-            connection,
-            program,
-        ),
+        PragmaName::JournalMode => {
+            // For JournalMode, when setting a value, we use the opcode
+            let mode_str = match value {
+                Expr::Name(name) => name.as_str().to_string(),
+                _ => parse_string(&value)?,
+            };
+
+            let result_reg = program.alloc_register();
+            program.emit_insn(Insn::JournalMode {
+                db: 0,
+                dest: result_reg,
+                new_mode: Some(mode_str),
+            });
+            program.emit_result_row(result_reg, 1);
+            program.add_pragma_result_column("journal_mode".into());
+            Ok((program, TransactionMode::None))
+        }
         PragmaName::LegacyFileFormat => Ok((program, TransactionMode::None)),
         PragmaName::WalCheckpoint => query_pragma(
             PragmaName::WalCheckpoint,
@@ -145,6 +154,24 @@ fn update_pragma(
             connection,
             program,
         ),
+        PragmaName::MaxPageCount => {
+            let data = parse_signed_number(&value)?;
+            let max_page_count_value = match data {
+                Value::Integer(i) => i as usize,
+                Value::Float(f) => f as usize,
+                _ => unreachable!(),
+            };
+
+            let result_reg = program.alloc_register();
+            program.emit_insn(Insn::MaxPgcnt {
+                db: 0,
+                dest: result_reg,
+                new_max: max_page_count_value,
+            });
+            program.emit_result_row(result_reg, 1);
+            program.add_pragma_result_column("max_page_count".into());
+            Ok((program, TransactionMode::Write))
+        }
         PragmaName::UserVersion => {
             let data = parse_signed_number(&value)?;
             let version_value = match data {
@@ -331,7 +358,12 @@ fn query_pragma(
             Ok((program, TransactionMode::None))
         }
         PragmaName::JournalMode => {
-            program.emit_string8("wal".into(), register);
+            // Use the JournalMode opcode to get the current journal mode
+            program.emit_insn(Insn::JournalMode {
+                db: 0,
+                dest: register,
+                new_mode: None,
+            });
             program.emit_result_row(register, 1);
             program.add_pragma_result_column(pragma.to_string());
             Ok((program, TransactionMode::None))
@@ -368,6 +400,16 @@ fn query_pragma(
             program.add_pragma_result_column(pragma.to_string());
             Ok((program, TransactionMode::Read))
         }
+        PragmaName::MaxPageCount => {
+            program.emit_insn(Insn::MaxPgcnt {
+                db: 0,
+                dest: register,
+                new_max: 0, // 0 means just return current max
+            });
+            program.emit_result_row(register, 1);
+            program.add_pragma_result_column(pragma.to_string());
+            Ok((program, TransactionMode::Read))
+        }
         PragmaName::TableInfo => {
             let table = match value {
                 Some(ast::Expr::Name(name)) => {

core/vdbe/execute.rs

@@ -8315,6 +8315,76 @@ fn stringify_register(reg: &mut Register) -> bool {
     }
 }
 
+pub fn op_max_pgcnt(
+    program: &Program,
+    state: &mut ProgramState,
+    insn: &Insn,
+    pager: &Rc<Pager>,
+    mv_store: Option<&Arc<MvStore>>,
+) -> Result<InsnFunctionStepResult> {
+    load_insn!(MaxPgcnt { db, dest, new_max }, insn);
+
+    if *db > 0 {
+        return Err(LimboError::InternalError(
+            "temp/attached databases not implemented yet".to_string(),
+        ));
+    }
+
+    let result_value = if *new_max == 0 {
+        // If new_max is 0, just return current maximum without changing it
+        pager.get_max_page_count()
+    } else {
+        // Set new maximum page count (will be clamped to current database size)
+        match pager.set_max_page_count(*new_max as u32)? {
+            IOResult::Done(new_max_count) => new_max_count,
+            IOResult::IO => return Ok(InsnFunctionStepResult::IO),
+        }
+    };
+
+    state.registers[*dest] = Register::Value(Value::Integer(result_value.into()));
+    state.pc += 1;
+    Ok(InsnFunctionStepResult::Step)
+}
+
+pub fn op_journal_mode(
+    program: &Program,
+    state: &mut ProgramState,
+    insn: &Insn,
+    pager: &Rc<Pager>,
+    mv_store: Option<&Arc<MvStore>>,
+) -> Result<InsnFunctionStepResult> {
+    load_insn!(JournalMode { db, dest, new_mode }, insn);
+    if *db > 0 {
+        return Err(LimboError::InternalError(
+            "temp/attached databases not implemented yet".to_string(),
+        ));
+    }
+
+    // Currently, Turso only supports WAL mode
+    // If a new mode is specified, we validate it but always return "wal"
+    if let Some(mode) = new_mode {
+        let mode_lower = mode.to_lowercase();
+        // Valid journal modes in SQLite are: delete, truncate, persist, memory, wal, off
+        // We accept any valid mode but always use WAL
+        match mode_lower.as_str() {
+            "delete" | "truncate" | "persist" | "memory" | "wal" | "off" => {
+                // Mode is valid, but we stay in WAL mode
+            }
+            _ => {
+                // Invalid journal mode
+                return Err(LimboError::ParseError(format!(
+                    "Unknown journal mode: {mode}"
+                )));
+            }
+        }
+    }
+
+    // Always return "wal" as the current journal mode
+    state.registers[*dest] = Register::Value(Value::build_text("wal"));
+    state.pc += 1;
+    Ok(InsnFunctionStepResult::Step)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

core/vdbe/explain.rs

@@ -1645,6 +1645,25 @@ pub fn insn_to_str(
                 0,
                 format!("drop_column({table}, {column_index})"),
             ),
+            Insn::MaxPgcnt { db, dest, new_max } => (
+                "MaxPgcnt",
+                *db as i32,
+                *dest as i32,
+                *new_max as i32,
+                Value::build_text(""),
+                0,
+                format!("r[{dest}]=max_page_count(db[{db}],{new_max})"),
+            ),
+            Insn::JournalMode { db, dest, new_mode } => (
+                "JournalMode",
+                *db as i32,
+                *dest as i32,
+                0,
+                Value::build_text(new_mode.as_ref().unwrap_or(&String::new())),
+                0,
+                format!("r[{dest}]=journal_mode(db[{db}]{})",
+                    new_mode.as_ref().map_or(String::new(), |m| format!(",'{m}'"))),
+            ),
             Insn::CollSeq { reg, collation } => (
                 "CollSeq",
                 reg.unwrap_or(0) as i32,

core/vdbe/insn.rs

@@ -1034,6 +1034,23 @@ pub enum Insn {
         table: String,
         column_index: usize,
     },
+    /// Try to set the maximum page count for database P1 to the value in P3.
+    /// Do not let the maximum page count fall below the current page count and
+    /// do not change the maximum page count value if P3==0.
+    /// Store the maximum page count after the change in register P2.
+    MaxPgcnt {
+        db: usize,      // P1: database index
+        dest: usize,    // P2: output register
+        new_max: usize, // P3: new maximum page count (0 = just return current)
+    },
+    /// Get or set the journal mode for database P1.
+    /// If P3 is not null, it contains the new journal mode string.
+    /// Store the resulting journal mode in register P2.
+    JournalMode {
+        db: usize,                // P1: database index
+        dest: usize,              // P2: output register for result
+        new_mode: Option<String>, // P3: new journal mode (if setting)
+    },
 }
 
 impl Insn {
@@ -1165,6 +1182,8 @@ impl Insn {
             Insn::IntegrityCk { .. } => execute::op_integrity_check,
             Insn::RenameTable { .. } => execute::op_rename_table,
             Insn::DropColumn { .. } => execute::op_drop_column,
+            Insn::MaxPgcnt { .. } => execute::op_max_pgcnt,
+            Insn::JournalMode { .. } => execute::op_journal_mode,
         }
     }
 }

testing/pragma.test

@@ -296,3 +296,32 @@ products|id
 products|name
 products|price
 }
+
+do_execsql_test_on_specific_db ":memory:" pragma-max-page-count-default {
+  PRAGMA max_page_count
+} {4294967294}
+
+do_execsql_test_on_specific_db ":memory:" pragma-max-page-count-set-large {
+  PRAGMA max_page_count = 1000;
+  PRAGMA max_page_count
+} {1000
+1000}
+
+do_execsql_test_on_specific_db ":memory:" pragma-max-page-count-set-zero-ignored {
+  PRAGMA max_page_count = 0;
+  PRAGMA max_page_count
+} {4294967294
+4294967294
+}
+
+do_execsql_test_on_specific_db ":memory:" pragma-max-page-count-clamping-with-data {
+  CREATE TABLE test (id INTEGER);
+  PRAGMA page_count;
+  PRAGMA max_page_count = 1;
+} {2
+2}
+
+do_execsql_test_in_memory_any_error pragma-max-page-count-enforcement-error {
+  PRAGMA max_page_count = 1;
+  CREATE TABLE test (id INTEGER)
+}

vendored/sqlite3-parser/src/parser/ast/mod.rs

@@ -1810,6 +1810,8 @@ pub enum PragmaName {
     JournalMode,
     /// Noop as per SQLite docs
     LegacyFileFormat,
+    /// Set or get the maximum number of pages in the database file.
+    MaxPageCount,
     /// Return the total number of pages in the database file.
     PageCount,
     /// Return the page size of the database in bytes.