Repository: servo/servo
File: components/storage/webstorage/engines/sqlite.rs (Rust)

Code before:
impl SqliteEngine {
    pub fn new(db_dir: &Option<PathBuf>, _pool: Arc<ThreadPool>) -> rusqlite::Result<Self> {
        let connection = match db_dir {
            Some(path) => {
                let path = path.join("webstorage.sqlite");
                Self::init_db(Some(path))?
            },
            None => Self::init_db(None)?,
        };
        // Initialize the database with necessary pragmas
        for pragma in DB_PRAGMAS.iter() {
            let _ = connection.execute(pragma, []);
        }
        Ok(SqliteEngine { connection })
    }

    pub fn init_db(path_opt: Option<PathBuf>) -> rusqlite::Result<Connection> {
        let connection = if let Some(path) = path_opt {
            if let Some(parent) = path.parent() {
                let _ = std::fs::create_dir_all(parent);
            }
            Connection::open(path)?
        } else {
            Connection::open_in_memory()?
        };
        for pragma in DB_INIT_PRAGMAS.iter() {
            let _ = connection.execute(pragma, []);
        }
        connection.execute("CREATE TABLE IF NOT EXISTS data (id INTEGER PRIMARY KEY AUTOINCREMENT, key TEXT, value TEXT);", [])?;
        Ok(connection)
    }
}

impl WebStorageEngine for SqliteEngine {
    type Error = rusqlite::Error;

    fn load(&self) -> Result<OriginEntry, Self::Error> {
        let mut stmt = self.connection.prepare("SELECT key, value FROM data;")?;
        let rows = stmt.query_map([], |row| {
            let key: String = row.get(0)?;
            let value: String = row.get(1)?;
            Ok((key, value))
        })?;

        let mut map = OriginEntry::default();
        for row in rows {
            let (key, value) = row?;
            map.insert(key, value);
        }
        Ok(map)
    }

Review comment (suggestion):
It seems that all `DB_PRAGMAS` and `DB_INIT_PRAGMAS` currently apply without errors to an in-memory database. However, some of these pragmas don’t really make sense in that context, so it might be worth having a separate, explicit list for in-memory databases.

That said, I think this patch is good as is. I would just consider adding a TODO about potentially splitting the pragma lists in the future.

I also assume this was manually tested (ideally, we would have at least a basic unit test for this).


Code after:
impl SqliteEngine {
    pub fn new(db_dir: &Option<PathBuf>, _pool: Arc<ThreadPool>) -> rusqlite::Result<Self> {
        let connection = match db_dir {
            Some(path) => {
                let path = path.join("webstorage.sqlite");
                Self::init_db(Some(&path))?
            },
            None => Self::init_db(None)?,
        };
        Ok(SqliteEngine { connection })
    }

    pub fn init_db(db_path: Option<&PathBuf>) -> rusqlite::Result<Connection> {
        let connection = if let Some(path) = db_path {
            if let Some(parent) = path.parent() {
                let _ = std::fs::create_dir_all(parent);
            }
            let conn = Connection::open(path)?;
            for pragma in DB_INIT_PRAGMAS.iter() {
                let _ = conn.execute(pragma, []);
            }
            for pragma in DB_PRAGMAS.iter() {
                let _ = conn.execute(pragma, []);
            }
            conn
        } else {
            // TODO We probably don't need an in memory implementation at all.
            // WebStorageEnvironment already keeps all key value pairs in memory via its data field.
            // A future refactoring could avoid creating a WebStorageEngine entirely when config_dir is None.
            let conn = Connection::open_in_memory()?;
            for pragma in DB_IN_MEMORY_INIT_PRAGMAS.iter() {
                let _ = conn.execute(pragma, []);
            }
            for pragma in DB_IN_MEMORY_PRAGMAS.iter() {
                let _ = conn.execute(pragma, []);
            }
            conn
        };
        connection.execute("CREATE TABLE IF NOT EXISTS data (id INTEGER PRIMARY KEY AUTOINCREMENT, key TEXT, value TEXT);", [])?;
        Ok(connection)
    }
}

impl WebStorageEngine for SqliteEngine {
    type Error = rusqlite::Error;

    fn load(&self) -> Result<OriginEntry, Self::Error> {
        let mut stmt = self.connection.prepare("SELECT key, value FROM data;")?;
        let rows = stmt.query_map([], |row| {
            let key: String = row.get(0)?;
            let value: String = row.get(1)?;