diff options
Diffstat (limited to 'third_party/rust/suggest/src/rs.rs')
| -rw-r--r-- | third_party/rust/suggest/src/rs.rs | 223 |
1 files changed, 219 insertions, 4 deletions
diff --git a/third_party/rust/suggest/src/rs.rs b/third_party/rust/suggest/src/rs.rs index 198a8c43f6..4a733ece9d 100644 --- a/third_party/rust/suggest/src/rs.rs +++ b/third_party/rust/suggest/src/rs.rs @@ -31,7 +31,7 @@ //! the new suggestion in their results, and return `Suggestion::T` variants //! as needed. -use std::borrow::Cow; +use std::{borrow::Cow, fmt}; use remote_settings::{GetItemsOptions, RemoteSettingsResponse}; use serde::{Deserialize, Deserializer}; @@ -47,6 +47,20 @@ pub(crate) const REMOTE_SETTINGS_COLLECTION: &str = "quicksuggest"; /// `mozilla-services/quicksuggest-rs` repo. pub(crate) const SUGGESTIONS_PER_ATTACHMENT: u64 = 200; +/// A list of default record types to download if nothing is specified. +/// This currently defaults to all of the record types. +pub(crate) const DEFAULT_RECORDS_TYPES: [SuggestRecordType; 9] = [ + SuggestRecordType::Icon, + SuggestRecordType::AmpWikipedia, + SuggestRecordType::Amo, + SuggestRecordType::Pocket, + SuggestRecordType::Yelp, + SuggestRecordType::Mdn, + SuggestRecordType::Weather, + SuggestRecordType::GlobalConfig, + SuggestRecordType::AmpMobile, +]; + /// A trait for a client that downloads suggestions from Remote Settings. /// /// This trait lets tests use a mock client. @@ -102,6 +116,61 @@ pub(crate) enum SuggestRecord { AmpMobile, } +/// Enum for the different record types that can be consumed. +/// Extracting this from the serialization enum so that we can +/// extend it to get type metadata. +#[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord)] +pub enum SuggestRecordType { + Icon, + AmpWikipedia, + Amo, + Pocket, + Yelp, + Mdn, + Weather, + GlobalConfig, + AmpMobile, +} + +impl From<SuggestRecord> for SuggestRecordType { + fn from(suggest_record: SuggestRecord) -> Self { + match suggest_record { + SuggestRecord::Amo => Self::Amo, + SuggestRecord::AmpWikipedia => Self::AmpWikipedia, + SuggestRecord::Icon => Self::Icon, + SuggestRecord::Mdn => Self::Mdn, + SuggestRecord::Pocket => Self::Pocket, + SuggestRecord::Weather(_) => Self::Weather, + SuggestRecord::Yelp => Self::Yelp, + SuggestRecord::GlobalConfig(_) => Self::GlobalConfig, + SuggestRecord::AmpMobile => Self::AmpMobile, + } + } +} + +impl fmt::Display for SuggestRecordType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Icon => write!(f, "icon"), + Self::AmpWikipedia => write!(f, "data"), + Self::Amo => write!(f, "amo-suggestions"), + Self::Pocket => write!(f, "pocket-suggestions"), + Self::Yelp => write!(f, "yelp-suggestions"), + Self::Mdn => write!(f, "mdn-suggestions"), + Self::Weather => write!(f, "weather"), + Self::GlobalConfig => write!(f, "configuration"), + Self::AmpMobile => write!(f, "amp-mobile-suggestions"), + } + } +} + +impl SuggestRecordType { + /// Return the meta key for the last ingested record. + pub fn last_ingest_meta_key(&self) -> String { + format!("last_quicksuggest_ingest_{}", self) + } +} + /// Represents either a single value, or a list of values. This is used to /// deserialize downloaded attachments. #[derive(Clone, Debug, Deserialize)] @@ -156,16 +225,18 @@ where } /// Fields that are common to all downloaded suggestions. -#[derive(Clone, Debug, Deserialize)] +#[derive(Clone, Debug, Default, Deserialize)] pub(crate) struct DownloadedSuggestionCommonDetails { pub keywords: Vec<String>, pub title: String, pub url: String, pub score: Option<f64>, + #[serde(default)] + pub full_keywords: Vec<(String, usize)>, } /// An AMP suggestion to ingest from an AMP-Wikipedia attachment. -#[derive(Clone, Debug, Deserialize)] +#[derive(Clone, Debug, Default, Deserialize)] pub(crate) struct DownloadedAmpSuggestion { #[serde(flatten)] pub common_details: DownloadedSuggestionCommonDetails, @@ -180,7 +251,7 @@ pub(crate) struct DownloadedAmpSuggestion { } /// A Wikipedia suggestion to ingest from an AMP-Wikipedia attachment. -#[derive(Clone, Debug, Deserialize)] +#[derive(Clone, Debug, Default, Deserialize)] pub(crate) struct DownloadedWikipediaSuggestion { #[serde(flatten)] pub common_details: DownloadedSuggestionCommonDetails, @@ -214,6 +285,34 @@ impl DownloadedAmpWikipediaSuggestion { } } +impl DownloadedSuggestionCommonDetails { + /// Iterate over all keywords for this suggestion + pub fn keywords(&self) -> impl Iterator<Item = AmpKeyword<'_>> { + let full_keywords = self + .full_keywords + .iter() + .flat_map(|(full_keyword, repeat_for)| { + std::iter::repeat(Some(full_keyword.as_str())).take(*repeat_for) + }) + .chain(std::iter::repeat(None)); // In case of insufficient full keywords, just fill in with infinite `None`s + // + self.keywords.iter().zip(full_keywords).enumerate().map( + move |(i, (keyword, full_keyword))| AmpKeyword { + rank: i, + keyword, + full_keyword, + }, + ) + } +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct AmpKeyword<'a> { + pub rank: usize, + pub keyword: &'a str, + pub full_keyword: Option<&'a str>, +} + impl<'de> Deserialize<'de> for DownloadedAmpWikipediaSuggestion { fn deserialize<D>( deserializer: D, @@ -344,3 +443,119 @@ where { String::deserialize(deserializer).map(|s| s.parse().ok()) } + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_full_keywords() { + let suggestion = DownloadedAmpWikipediaSuggestion::Amp(DownloadedAmpSuggestion { + common_details: DownloadedSuggestionCommonDetails { + keywords: vec![ + String::from("f"), + String::from("fo"), + String::from("foo"), + String::from("foo b"), + String::from("foo ba"), + String::from("foo bar"), + ], + full_keywords: vec![(String::from("foo"), 3), (String::from("foo bar"), 3)], + ..DownloadedSuggestionCommonDetails::default() + }, + ..DownloadedAmpSuggestion::default() + }); + + assert_eq!( + Vec::from_iter(suggestion.common_details().keywords()), + vec![ + AmpKeyword { + rank: 0, + keyword: "f", + full_keyword: Some("foo"), + }, + AmpKeyword { + rank: 1, + keyword: "fo", + full_keyword: Some("foo"), + }, + AmpKeyword { + rank: 2, + keyword: "foo", + full_keyword: Some("foo"), + }, + AmpKeyword { + rank: 3, + keyword: "foo b", + full_keyword: Some("foo bar"), + }, + AmpKeyword { + rank: 4, + keyword: "foo ba", + full_keyword: Some("foo bar"), + }, + AmpKeyword { + rank: 5, + keyword: "foo bar", + full_keyword: Some("foo bar"), + }, + ], + ); + } + + #[test] + fn test_missing_full_keywords() { + let suggestion = DownloadedAmpWikipediaSuggestion::Amp(DownloadedAmpSuggestion { + common_details: DownloadedSuggestionCommonDetails { + keywords: vec![ + String::from("f"), + String::from("fo"), + String::from("foo"), + String::from("foo b"), + String::from("foo ba"), + String::from("foo bar"), + ], + // Only the first 3 keywords have full keywords associated with them + full_keywords: vec![(String::from("foo"), 3)], + ..DownloadedSuggestionCommonDetails::default() + }, + ..DownloadedAmpSuggestion::default() + }); + + assert_eq!( + Vec::from_iter(suggestion.common_details().keywords()), + vec![ + AmpKeyword { + rank: 0, + keyword: "f", + full_keyword: Some("foo"), + }, + AmpKeyword { + rank: 1, + keyword: "fo", + full_keyword: Some("foo"), + }, + AmpKeyword { + rank: 2, + keyword: "foo", + full_keyword: Some("foo"), + }, + AmpKeyword { + rank: 3, + keyword: "foo b", + full_keyword: None, + }, + AmpKeyword { + rank: 4, + keyword: "foo ba", + full_keyword: None, + }, + AmpKeyword { + rank: 5, + keyword: "foo bar", + full_keyword: None, + }, + ], + ); + } +} |