Skip to content

Commit

Permalink
feedback changes
Browse files Browse the repository at this point in the history
  • Loading branch information
@tiftran
tiftran committed Oct 3, 2023
1 parent 029f80a commit 528868e
Show file tree
Hide file tree
Showing 5 changed files with 270 additions and 145 deletions.
155 changes: 97 additions & 58 deletions components/suggest/src/db.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,18 @@
* file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/

use std::{borrow::Cow, path::Path, sync::Arc};
use std::{path::Path, sync::Arc};

use interrupt_support::{SqlInterruptHandle, SqlInterruptScope};
use parking_lot::Mutex;
use rusqlite::{
named_params,
types::{FromSql, ToSql},
Connection, OpenFlags,
Connection, OpenFlags, Row,
};
use sql_support::{open_database::open_database_with_flags, ConnExt};

use crate::pocket_confidence::PocketKeywordConfidence;
use crate::rs::{DownloadedAmoSuggestion, DownloadedPocketSuggestion};
use crate::{
keyword::full_keyword,
Expand Down Expand Up @@ -124,31 +125,28 @@ impl<'a> SuggestDao<'a> {

/// Fetches suggestions that match the given keyword from the database.
pub fn fetch_by_keyword(&self, keyword: &str) -> Result<Vec<Suggestion>> {
let mut pocket_keyword = Cow::from(keyword);
if !keyword.contains(' ') {
pocket_keyword.to_mut().push(' ');
}
self.conn.query_rows_and_then_cached(
"SELECT s.id, k.rank, s.title, s.url, s.provider, false as is_top_pick
let (keyword_prefix, keyword_suffix) = match keyword.find(' ') {
Some(index) => {
let (prefix, suffix) = keyword.split_at(index);
(prefix.to_string(), suffix[1..].to_string())
}
None => (keyword.to_string(), String::new()),
};
Ok(self.conn.query_rows_and_then_cached(
"SELECT s.id, k.rank, s.title, s.url, s.provider, -1 as confidence_type
FROM suggestions s
JOIN keywords k ON k.suggestion_id = s.id
WHERE k.keyword = :keyword
UNION
SELECT s.id, k.rank, s.title, s.url, s.provider, true as is_top_pick
FROM suggestions s
JOIN pocket_high_confidence_keywords k ON k.suggestion_id = s.id
WHERE k.keyword = :keyword
UNION
SELECT s.id, k.rank, s.title, s.url, s.provider, false as is_top_pick
SELECT s.id, -1 as rank, s.title, s.url, s.provider, k.confidence_type
FROM suggestions s
JOIN pocket_low_confidence_keywords k ON k.suggestion_id = s.id
WHERE k.keyword BETWEEN :pocket_keyword AND :pocket_keyword || x'ffff'
LIMIT 1",
JOIN pocket_keywords k ON k.suggestion_id = s.id
WHERE k.keyword_prefix = :keyword_prefix",
named_params! {
":keyword": keyword,
":pocket_keyword": pocket_keyword
":keyword_prefix": keyword_prefix,
},
|row| -> Result<Suggestion>{
|row: &Row<'_> | -> Result<Option<Suggestion>>{
let suggestion_id: i64 = row.get("id")?;
let title = row.get("title")?;
let raw_url = row.get::<_, String>("url")?;
Expand Down Expand Up @@ -179,7 +177,7 @@ impl<'a> SuggestDao<'a> {
let cooked_url = cook_raw_suggestion_url(&raw_url);
let raw_click_url = row.get::<_, String>("click_url")?;
let cooked_click_url = cook_raw_suggestion_url(&raw_click_url);
Ok(Suggestion::Amp {
Ok(Some(Suggestion::Amp {
block_id: row.get("block_id")?,
advertiser: row.get("advertiser")?,
iab_category: row.get("iab_category")?,
Expand All @@ -191,7 +189,7 @@ impl<'a> SuggestDao<'a> {
impression_url: row.get("impression_url")?,
click_url: cooked_click_url,
raw_click_url,
})
}))
}
)
},
Expand All @@ -206,12 +204,12 @@ impl<'a> SuggestDao<'a> {
},
true,
)?;
Ok(Suggestion::Wikipedia {
Ok(Some(Suggestion::Wikipedia {
title,
url: raw_url,
full_keyword: full_keyword(keyword, &keywords),
icon,
})
}))
}
SuggestionProvider::Amo => {
self.conn.query_row_and_then(
Expand All @@ -222,7 +220,7 @@ impl<'a> SuggestDao<'a> {
":suggestion_id": suggestion_id
},
|row| {
Ok(Suggestion::Amo{
Ok(Some(Suggestion::Amo{
title,
url: raw_url,
icon_url: row.get("icon_url")?,
Expand All @@ -231,31 +229,52 @@ impl<'a> SuggestDao<'a> {
number_of_ratings: row.get("number_of_ratings")?,
guid: row.get("guid")?,
score: row.get("score")?,
})
}))
})
},
SuggestionProvider::Pocket => {
let is_top_pick = row.get("is_top_pick")?;
let confidence_type = row.get("confidence_type")?;
self.conn.query_row_and_then(
"SELECT p.score
FROM pocket_custom_details p
WHERE p.suggestion_id = :suggestion_id",
"SELECT keyword_suffix FROM pocket_keywords WHERE suggestion_id = :suggestion_id AND keyword_prefix = :keyword_prefix",
named_params! {
":suggestion_id": suggestion_id
":suggestion_id": suggestion_id,
":keyword_prefix": keyword_prefix
},
|row| {
Ok(Suggestion::Pocket {
title,
url: raw_url,
score: row.get("score")?,
is_top_pick
})
let suffixes_match = match confidence_type {
PocketKeywordConfidence::Low => row.get::<_, String>("keyword_suffix")?.starts_with(&keyword_suffix),
PocketKeywordConfidence::High => row.get::<_, String>("keyword_suffix")? == keyword_suffix,
};
if suffixes_match{
self.conn.query_row_and_then(
"SELECT p.score
FROM pocket_custom_details p
WHERE p.suggestion_id = :suggestion_id",
named_params! {
":suggestion_id": suggestion_id
},
|row| {
Ok(Some(Suggestion::Pocket {
title,
url: raw_url,
score: row.get("score")?,
is_top_pick: match confidence_type{
PocketKeywordConfidence::Low => false,
PocketKeywordConfidence::High => true,
}
}))
}
)
}
else{
Ok(None)
}
}
)
}
}
},
)
}
)?.into_iter().flatten().collect())
}

/// Inserts all suggestions from a downloaded AMO attachment into
Expand Down Expand Up @@ -502,41 +521,61 @@ impl<'a> SuggestDao<'a> {
":score": suggestion.score,
},
)?;
for (index, keyword) in suggestion.low_confidence_keywords.iter().enumerate() {
for (_, keyword) in suggestion.low_confidence_keywords.iter().enumerate() {
let (keyword_prefix, keyword_suffix) = match keyword.find(' ') {
Some(index) => {
let (prefix, suffix) = keyword.split_at(index);
(prefix.to_string(), suffix[1..].to_string())
}
None => (keyword.clone(), String::new()),
};
self.conn.execute(
"INSERT INTO pocket_low_confidence_keywords(
keyword,
suggestion_id,
rank
"INSERT INTO pocket_keywords(
keyword_prefix,
keyword_suffix,
confidence_type,
suggestion_id
)
VALUES(
:keyword,
:suggestion_id,
:rank
:keyword_prefix,
:keyword_suffix,
:confidence_type,
:suggestion_id
)",
named_params! {
":keyword": keyword,
":rank": index,
":keyword_prefix": keyword_prefix,
":keyword_suffix": keyword_suffix,
":confidence_type": 0,
":suggestion_id": suggestion_id,
},
)?;
}

for (index, keyword) in suggestion.high_confidence_keywords.iter().enumerate() {
for (_, keyword) in suggestion.high_confidence_keywords.iter().enumerate() {
let (keyword_prefix, keyword_suffix) = match keyword.find(' ') {
Some(index) => {
let (prefix, suffix) = keyword.split_at(index);
(prefix.to_string(), suffix[1..].to_string())
}
None => (keyword.clone(), String::new()),
};
self.conn.execute(
"INSERT INTO pocket_high_confidence_keywords(
keyword,
suggestion_id,
rank
"INSERT INTO pocket_keywords(
keyword_prefix,
keyword_suffix,
confidence_type,
suggestion_id
)
VALUES(
:keyword,
:suggestion_id,
:rank
:keyword_prefix,
:keyword_suffix,
:confidence_type,
:suggestion_id
)",
named_params! {
":keyword": keyword,
":rank": index,
":keyword_prefix": keyword_prefix,
":keyword_suffix": keyword_suffix,
":confidence_type": 1,
":suggestion_id": suggestion_id,
},
)?;
Expand Down
1 change: 1 addition & 0 deletions components/suggest/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use remote_settings::RemoteSettingsConfig;
mod db;
mod error;
mod keyword;
mod pocket_confidence;
mod provider;
mod rs;
mod schema;
Expand Down
37 changes: 37 additions & 0 deletions components/suggest/src/pocket_confidence.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
use rusqlite::types::{FromSql, FromSqlError, FromSqlResult, ToSqlOutput, ValueRef};
use rusqlite::{Result as RusqliteResult, ToSql};

/// Classification of Pocket confidence keywords.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
#[repr(u8)]
pub enum PocketKeywordConfidence {
Low = 0,
High = 1,
}

impl FromSql for PocketKeywordConfidence {
fn column_result(value: ValueRef<'_>) -> FromSqlResult<Self> {
let v = value.as_i64()?;
u8::try_from(v)
.ok()
.and_then(PocketKeywordConfidence::from_u8)
.ok_or_else(|| FromSqlError::OutOfRange(v))
}
}

impl PocketKeywordConfidence {
#[inline]
pub(crate) fn from_u8(v: u8) -> Option<Self> {
match v {
0 => Some(PocketKeywordConfidence::Low),
1 => Some(PocketKeywordConfidence::High),
_ => None,
}
}
}

impl ToSql for PocketKeywordConfidence {
fn to_sql(&self) -> RusqliteResult<ToSqlOutput<'_>> {
Ok(ToSqlOutput::from(*self as u8))
}
}
19 changes: 6 additions & 13 deletions components/suggest/src/schema.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,19 +21,12 @@ pub const SQL: &str = "
PRIMARY KEY (keyword, suggestion_id)
) WITHOUT ROWID;
CREATE TABLE pocket_low_confidence_keywords(
keyword TEXT NOT NULL,
suggestion_id INTEGER NOT NULL REFERENCES suggestions(id) ON DELETE CASCADE,
rank INTEGER NOT NULL,
PRIMARY KEY (keyword, suggestion_id)
);
CREATE TABLE pocket_high_confidence_keywords(
keyword TEXT NOT NULL,
suggestion_id INTEGER NOT NULL REFERENCES suggestions(id) ON DELETE CASCADE,
rank INTEGER NOT NULL,
PRIMARY KEY (keyword, suggestion_id)
);
CREATE TABLE pocket_keywords(
keyword_prefix TEXT NOT NULL,
keyword_suffix TEXT NOT NULL DEFAULT '',
confidence_type INTEGER NOT NULL,
suggestion_id INTEGER NOT NULL REFERENCES suggestions(id),
PRIMARY KEY (keyword_prefix, keyword_suffix, suggestion_id) );
CREATE UNIQUE INDEX keywords_suggestion_id_rank ON keywords(suggestion_id, rank);
Expand Down
Loading

0 comments on commit 528868e

Please sign in to comment.