Add PRAGMA key to set the encryption key

If set, set the key for the connection

core/lib.rs

@@ -76,6 +76,8 @@ use std::{
 };
 #[cfg(feature = "fs")]
 use storage::database::DatabaseFile;
+#[cfg(feature = "encryption")]
+pub use storage::encryption::EncryptionKey;
 use storage::page_cache::DumbLruPageCache;
 use storage::pager::{AtomicDbState, DbState};
 use storage::sqlite3_ondisk::PageSize;
@@ -425,6 +427,8 @@ impl Database {
             view_transaction_states: RefCell::new(HashMap::new()),
             metrics: RefCell::new(ConnectionMetrics::new()),
             is_nested_stmt: Cell::new(false),
+            #[cfg(feature = "encryption")]
+            encryption_key: RefCell::new(None),
         });
         let builtin_syms = self.builtin_syms.borrow();
         // add built-in extensions symbols to the connection to prevent having to load each time
@@ -852,6 +856,8 @@ pub struct Connection {
     /// Whether the connection is executing a statement initiated by another statement.
     /// Generally this is only true for ParseSchema.
     is_nested_stmt: Cell<bool>,
+    #[cfg(feature = "encryption")]
+    encryption_key: RefCell<Option<EncryptionKey>>,
 }
 
 impl Connection {
@@ -1925,6 +1931,14 @@ impl Connection {
     pub fn get_syms_vtab_mods(&self) -> std::collections::HashSet<String> {
         self.syms.borrow().vtab_modules.keys().cloned().collect()
     }
+
+    #[cfg(feature = "encryption")]
+    pub fn set_encryption_key(&self, key: Option<EncryptionKey>) {
+        tracing::trace!("setting encryption key for connection");
+        *self.encryption_key.borrow_mut() = key.clone();
+        let pager = self.pager.borrow();
+        pager.set_encryption_key(key);
+    }
 }
 
 pub struct Statement {

core/pragma.rs

@@ -107,6 +107,10 @@ pub fn pragma_for(pragma: &PragmaName) -> Pragma {
             &["query_only"],
         ),
         FreelistCount => Pragma::new(PragmaFlags::Result0, &["freelist_count"]),
+        EncryptionKey => Pragma::new(
+            PragmaFlags::Result0 | PragmaFlags::SchemaReq | PragmaFlags::NoColumns1,
+            &["key"],
+        ),
     }
 }
 

core/translate/pragma.rs

@@ -7,11 +7,16 @@ use std::sync::Arc;
 use turso_sqlite3_parser::ast::{self, ColumnDefinition, Expr, Literal, Name};
 use turso_sqlite3_parser::ast::{PragmaName, QualifiedName};
 
+use super::integrity_check::translate_integrity_check;
 use crate::pragma::pragma_for;
 use crate::schema::Schema;
+#[cfg(feature = "encryption")]
+use crate::storage::encryption::EncryptionKey;
 use crate::storage::pager::AutoVacuumMode;
+use crate::storage::pager::Pager;
 use crate::storage::sqlite3_ondisk::CacheSize;
 use crate::storage::wal::CheckpointMode;
+use crate::translate::emitter::TransactionMode;
 use crate::translate::schema::translate_create_table;
 use crate::util::{normalize_ident, parse_signed_number, parse_string, IOExt as _};
 use crate::vdbe::builder::{ProgramBuilder, ProgramBuilderOpts};
@@ -20,10 +25,6 @@ use crate::{bail_parse_error, CaptureDataChangesMode, LimboError, SymbolTable, V
 use std::str::FromStr;
 use strum::IntoEnumIterator;
 
-use super::integrity_check::translate_integrity_check;
-use crate::storage::pager::Pager;
-use crate::translate::emitter::TransactionMode;
-
 fn list_pragmas(program: &mut ProgramBuilder) {
     for x in PragmaName::iter() {
         let register = program.emit_string8_new_reg(x.to_string());
@@ -309,6 +310,15 @@ fn update_pragma(
             connection,
             program,
         ),
+        PragmaName::EncryptionKey => {
+            #[cfg(feature = "encryption")]
+            {
+                let value = parse_string(&value)?;
+                let key = EncryptionKey::from_string(&value);
+                connection.set_encryption_key(Some(key));
+            }
+            Ok((program, TransactionMode::None))
+        }
     }
 }
 
@@ -566,6 +576,23 @@ fn query_pragma(
             program.add_pragma_result_column(pragma.to_string());
             Ok((program, TransactionMode::None))
         }
+        PragmaName::EncryptionKey => {
+            #[cfg(feature = "encryption")]
+            {
+                let msg = {
+                    if connection.encryption_key.borrow().is_some() {
+                        "encryption key is set for this session"
+                    } else {
+                        "encryption key is not set for this session"
+                    }
+                };
+                let register = program.alloc_register();
+                program.emit_string8(msg.to_string(), register);
+                program.emit_result_row(register, 1);
+                program.add_pragma_result_column(pragma.to_string());
+            }
+            Ok((program, TransactionMode::None))
+        }
     }
 }
 

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

@@ -1829,6 +1829,11 @@ pub enum PragmaName {
     IntegrityCheck,
     /// `journal_mode` pragma
     JournalMode,
+    /// encryption key for encrypted databases. This is just called `key` because most
+    /// extensions use this name instead of `encryption_key`.
+    #[strum(serialize = "key")]
+    #[cfg_attr(feature = "serde", serde(rename = "key"))]
+    EncryptionKey,
     /// Noop as per SQLite docs
     LegacyFileFormat,
     /// Set or get the maximum number of pages in the database file.