Revision control
Copy as Markdown
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
*/
use crate::{
interest::InterestVectorKind,
schema::RelevancyConnectionInitializer,
url_hash::{hash_url, UrlHash},
Interest, InterestVector, Result,
};
use interrupt_support::SqlInterruptScope;
use rusqlite::{Connection, OpenFlags};
use sql_support::{ConnExt, LazyDb};
use std::path::Path;
/// A thread-safe wrapper around an SQLite connection to the Relevancy database
pub struct RelevancyDb {
reader: LazyDb<RelevancyConnectionInitializer>,
writer: LazyDb<RelevancyConnectionInitializer>,
}
impl RelevancyDb {
pub fn new(path: impl AsRef<Path>) -> Self {
// Note: use `SQLITE_OPEN_READ_WRITE` for both read and write connections.
// Even if we're opening a read connection, we may need to do a write as part of the
// initialization process.
//
// The read-only nature of the connection is enforced by the fact that [RelevancyDb::read] uses a
// shared ref to the `RelevancyDao`.
let db_open_flags = OpenFlags::SQLITE_OPEN_URI
| OpenFlags::SQLITE_OPEN_NO_MUTEX
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_READ_WRITE;
Self {
reader: LazyDb::new(path.as_ref(), db_open_flags, RelevancyConnectionInitializer),
writer: LazyDb::new(path.as_ref(), db_open_flags, RelevancyConnectionInitializer),
}
}
pub fn close(&self) {
self.reader.close(true);
self.writer.close(true);
}
pub fn interrupt(&self) {
self.reader.interrupt();
self.writer.interrupt();
}
#[cfg(test)]
pub fn new_for_test() -> Self {
use std::sync::atomic::{AtomicU32, Ordering};
static COUNTER: AtomicU32 = AtomicU32::new(0);
let count = COUNTER.fetch_add(1, Ordering::Relaxed);
Self::new(format!("file:test{count}.sqlite?mode=memory&cache=shared"))
}
/// Accesses the Suggest database in a transaction for reading.
pub fn read<T>(&self, op: impl FnOnce(&RelevancyDao) -> Result<T>) -> Result<T> {
let (mut conn, scope) = self.reader.lock()?;
let tx = conn.transaction()?;
let dao = RelevancyDao::new(&tx, scope);
op(&dao)
}
/// Accesses the Suggest database in a transaction for reading and writing.
pub fn read_write<T>(&self, op: impl FnOnce(&mut RelevancyDao) -> Result<T>) -> Result<T> {
let (mut conn, scope) = self.writer.lock()?;
let tx = conn.transaction()?;
let mut dao = RelevancyDao::new(&tx, scope);
let result = op(&mut dao)?;
tx.commit()?;
Ok(result)
}
}
/// A data access object (DAO) that wraps a connection to the Relevancy database
///
/// Methods that only read from the database take an immutable reference to
/// `self` (`&self`), and methods that write to the database take a mutable
/// reference (`&mut self`).
pub struct RelevancyDao<'a> {
pub conn: &'a Connection,
pub scope: SqlInterruptScope,
}
impl<'a> RelevancyDao<'a> {
fn new(conn: &'a Connection, scope: SqlInterruptScope) -> Self {
Self { conn, scope }
}
/// Return Err(Interrupted) if we were interrupted
pub fn err_if_interrupted(&self) -> Result<()> {
Ok(self.scope.err_if_interrupted()?)
}
/// Associate a URL with an interest
pub fn add_url_interest(&mut self, url_hash: UrlHash, interest: Interest) -> Result<()> {
let sql = "
INSERT OR REPLACE INTO url_interest(url_hash, interest_code)
VALUES (?, ?)
";
self.conn.execute(sql, (url_hash, interest as u32))?;
Ok(())
}
/// Get an interest vector for a URL
pub fn get_url_interest_vector(&self, url: &str) -> Result<InterestVector> {
let hash = match hash_url(url) {
Some(u) => u,
None => return Ok(InterestVector::default()),
};
let mut stmt = self.conn.prepare_cached(
"
SELECT interest_code
FROM url_interest
WHERE url_hash=?
",
)?;
let interests = stmt.query_and_then((hash,), |row| -> Result<Interest> {
row.get::<_, u32>(0)?.try_into()
})?;
let mut interest_vec = InterestVector::default();
for interest in interests {
interest_vec[interest?] += 1
}
Ok(interest_vec)
}
/// Do we need to load the interest data?
pub fn need_to_load_url_interests(&self) -> Result<bool> {
// TODO: we probably will need a better check than this.
Ok(self
.conn
.query_one("SELECT NOT EXISTS (SELECT 1 FROM url_interest)")?)
}
/// Update the frecency user interest vector based on a new measurement.
///
/// Right now this completely replaces the interest vector with the new data. At some point,
/// we may switch to incrementally updating it instead.
pub fn update_frecency_user_interest_vector(&self, interests: &InterestVector) -> Result<()> {
let mut stmt = self.conn.prepare(
"
INSERT OR REPLACE INTO user_interest(kind, interest_code, count)
VALUES (?, ?, ?)
",
)?;
for (interest, count) in interests.as_vec() {
stmt.execute((InterestVectorKind::Frecency, interest, count))?;
}
Ok(())
}
pub fn get_frecency_user_interest_vector(&self) -> Result<InterestVector> {
let mut stmt = self
.conn
.prepare("SELECT interest_code, count FROM user_interest WHERE kind = ?")?;
let mut interest_vec = InterestVector::default();
let rows = stmt.query_and_then((InterestVectorKind::Frecency,), |row| {
crate::Result::Ok((
Interest::try_from(row.get::<_, u32>(0)?)?,
row.get::<_, u32>(1)?,
))
})?;
for row in rows {
let (interest_code, count) = row?;
interest_vec.set(interest_code, count);
}
Ok(interest_vec)
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_store_frecency_user_interest_vector() {
let db = RelevancyDb::new_for_test();
// Initially the interest vector should be blank
assert_eq!(
db.read_write(|dao| dao.get_frecency_user_interest_vector())
.unwrap(),
InterestVector::default()
);
let interest_vec = InterestVector {
animals: 2,
autos: 1,
news: 5,
..InterestVector::default()
};
db.read_write(|dao| dao.update_frecency_user_interest_vector(&interest_vec))
.unwrap();
assert_eq!(
db.read_write(|dao| dao.get_frecency_user_interest_vector())
.unwrap(),
interest_vec,
);
}
#[test]
fn test_update_frecency_user_interest_vector() {
let db = RelevancyDb::new_for_test();
let interest_vec1 = InterestVector {
animals: 2,
autos: 1,
news: 5,
..InterestVector::default()
};
let interest_vec2 = InterestVector {
animals: 1,
career: 3,
..InterestVector::default()
};
// Update the first interest vec, then the second one
db.read_write(|dao| dao.update_frecency_user_interest_vector(&interest_vec1))
.unwrap();
db.read_write(|dao| dao.update_frecency_user_interest_vector(&interest_vec2))
.unwrap();
// The current behavior is the second one should replace the first
assert_eq!(
db.read_write(|dao| dao.get_frecency_user_interest_vector())
.unwrap(),
interest_vec2,
);
}
}