From a3387cfd5fb3c2dc2bf9a42bbcc836a0b15ccc9b Mon Sep 17 00:00:00 2001
From: Glauber Costa <glauber@turso.tech>
Date: Wed, 29 Jan 2025 11:00:10 -0500
Subject: [PATCH] implement the pragma page_count

To do that, we also have to implement the vdbe opcode Pagecount.
---
 COMPAT.md                                     |  4 ++--
 core/translate/mod.rs                         | 11 +++++++++++
 core/vdbe/explain.rs                          |  9 +++++++++
 core/vdbe/insn.rs                             |  5 +++++
 core/vdbe/mod.rs                              | 14 ++++++++++++++
 testing/pragma.test                           |  9 +++++++++
 vendored/sqlite3-parser/src/parser/ast/mod.rs |  2 ++
 7 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/COMPAT.md b/COMPAT.md
index 2b8194695..e48472f6b 100644
--- a/COMPAT.md
+++ b/COMPAT.md
@@ -135,7 +135,7 @@ The current status of Limbo is:
 | PRAGMA mmap_size                 | No         |                                              |
 | PRAGMA module_list               | No         |                                              |
 | PRAGMA optimize                  | No         |                                              |
-| PRAGMA page_count                | No         |                                              |
+| PRAGMA page_count                | Yes        |                                              |
 | PRAGMA page_size                 | No         |                                              |
 | PRAGMA parser_trace              | No         |                                              |
 | PRAGMA pragma_list               | Yes        |                                              |
@@ -504,7 +504,7 @@ Modifiers:
 | OpenWriteAsync | Yes    |         |
 | OpenWriteAwait | Yes    |         |
 | Or             | Yes    |         |
-| Pagecount      | No     |         |
+| Pagecount      | Partial| no temp databases |
 | Param          | No     |         |
 | ParseSchema    | No     |         |
 | Permutation    | No     |         |
diff --git a/core/translate/mod.rs b/core/translate/mod.rs
index ed7dd44cf..26bcf49de 100644
--- a/core/translate/mod.rs
+++ b/core/translate/mod.rs
@@ -640,6 +640,10 @@ fn update_pragma(
             query_pragma(PragmaName::WalCheckpoint, schema, None, header, program)?;
             Ok(())
         }
+        PragmaName::PageCount => {
+            query_pragma(PragmaName::PageCount, schema, None, header, program)?;
+            Ok(())
+        }
         PragmaName::TableInfo => {
             // because we need control over the write parameter for the transaction,
             // this should be unreachable. We have to force-call query_pragma before
@@ -681,6 +685,13 @@ fn query_pragma(
             });
             program.emit_result_row(register, 3);
         }
+        PragmaName::PageCount => {
+            program.emit_insn(Insn::PageCount {
+                db: 0,
+                dest: register,
+            });
+            program.emit_result_row(register, 1);
+        }
         PragmaName::TableInfo => {
             let table = match value {
                 Some(ast::Expr::Name(name)) => {
diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs
index 2d6afea0f..309745b89 100644
--- a/core/vdbe/explain.rs
+++ b/core/vdbe/explain.rs
@@ -1163,6 +1163,15 @@ pub fn insn_to_str(
                 0,
                 String::new(),
             ),
+            Insn::PageCount { db, dest } => (
+                "Pagecount",
+                *db as i32,
+                *dest as i32,
+                0,
+                OwnedValue::build_text(Rc::new("".to_string())),
+                0,
+                "".to_string(),
+            ),
         };
     format!(
         "{:<4}  {:<17}  {:<4}  {:<4}  {:<4}  {:<13}  {:<2}  {}",
diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs
index 97888cc4f..0d5c84a7d 100644
--- a/core/vdbe/insn.rs
+++ b/core/vdbe/insn.rs
@@ -605,6 +605,11 @@ pub enum Insn {
         dest: usize,
     },
     Noop,
+    /// Write the current number of pages in database P1 to memory cell P2.
+    PageCount {
+        db: usize,
+        dest: usize,
+    },
 }
 
 fn cast_text_to_numerical(value: &str) -> OwnedValue {
diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs
index 6d00692ec..5636ab0d3 100644
--- a/core/vdbe/mod.rs
+++ b/core/vdbe/mod.rs
@@ -2406,6 +2406,20 @@ impl Program {
                         state.pc += 1;
                     }
                 }
+                Insn::PageCount { db, dest } => {
+                    if *db > 0 {
+                        // TODO: implement temp databases
+                        todo!("temp databases not implemented yet");
+                    }
+                    // SQLite returns "0" on an empty database, and 2 on the first insertion,
+                    // so we'll mimick that behavior.
+                    let mut pages = pager.db_header.borrow().database_size.into();
+                    if pages == 1 {
+                        pages = 0;
+                    }
+                    state.registers[*dest] = OwnedValue::Integer(pages);
+                    state.pc += 1;
+                }
                 Insn::ParseSchema {
                     db: _,
                     where_clause,
diff --git a/testing/pragma.test b/testing/pragma.test
index 17c33c1b0..436ad48e1 100755
--- a/testing/pragma.test
+++ b/testing/pragma.test
@@ -32,3 +32,12 @@ do_execsql_test pragma-table-info-call-syntax {
 do_execsql_test pragma-table-info-invalid-table {
   PRAGMA table_info=pekka
 } {}
+
+do_execsql_test_on_specific_db ":memory:" pragma-page-count-empty {
+  PRAGMA page_count
+} {0}
+
+do_execsql_test_on_specific_db ":memory:" pragma-page-count-table {
+  CREATE TABLE foo(bar);
+  PRAGMA page_count
+} {2}
diff --git a/vendored/sqlite3-parser/src/parser/ast/mod.rs b/vendored/sqlite3-parser/src/parser/ast/mod.rs
index 26b25c353..f149322b3 100644
--- a/vendored/sqlite3-parser/src/parser/ast/mod.rs
+++ b/vendored/sqlite3-parser/src/parser/ast/mod.rs
@@ -1594,6 +1594,8 @@ pub enum PragmaName {
     CacheSize,
     /// `journal_mode` pragma
     JournalMode,
+    /// Return the total number of pages in the database file.
+    PageCount,
     /// returns information about the columns of a table
     TableInfo,
     /// trigger a checkpoint to run on database(s) if WAL is enabled