suggest/
provider.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
4 */
5
6use std::{
7    collections::{HashMap, HashSet},
8    fmt,
9};
10
11use rusqlite::{
12    types::{FromSql, FromSqlError, FromSqlResult, ToSql, ToSqlOutput, ValueRef},
13    Result as RusqliteResult,
14};
15
16use crate::rs::{Collection, SuggestRecordType};
17
18#[cfg(test)]
19use serde_json::Value as JsonValue;
20
21#[cfg(test)]
22use crate::testing::{MockAttachment, MockIcon, MockRecord};
23
24/// Record types from these providers will be ingested when consumers do not
25/// specify providers in `SuggestIngestionConstraints`.
26pub(crate) const DEFAULT_INGEST_PROVIDERS: [SuggestionProvider; 5] = [
27    SuggestionProvider::Amp,
28    SuggestionProvider::Wikipedia,
29    SuggestionProvider::Amo,
30    SuggestionProvider::Yelp,
31    SuggestionProvider::Mdn,
32];
33
34/// A provider is a source of search suggestions.
35/// Please preserve the integer values after removing or adding providers.
36/// Provider configs are associated with integer keys stored in the database.
37#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, uniffi::Enum)]
38#[repr(u8)]
39pub enum SuggestionProvider {
40    Amp = 1,
41    Wikipedia = 2,
42    Amo = 3,
43    Yelp = 5,
44    Mdn = 6,
45    Weather = 7,
46    Fakespot = 8,
47    Dynamic = 9,
48}
49
50impl fmt::Display for SuggestionProvider {
51    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52        match self {
53            Self::Amp => write!(f, "amp"),
54            Self::Wikipedia => write!(f, "wikipedia"),
55            Self::Amo => write!(f, "amo"),
56            Self::Yelp => write!(f, "yelp"),
57            Self::Mdn => write!(f, "mdn"),
58            Self::Weather => write!(f, "weather"),
59            Self::Fakespot => write!(f, "fakespot"),
60            Self::Dynamic => write!(f, "dynamic"),
61        }
62    }
63}
64
65impl FromSql for SuggestionProvider {
66    fn column_result(value: ValueRef<'_>) -> FromSqlResult<Self> {
67        let v = value.as_i64()?;
68        u8::try_from(v)
69            .ok()
70            .and_then(SuggestionProvider::from_u8)
71            .ok_or_else(|| FromSqlError::OutOfRange(v))
72    }
73}
74
75impl SuggestionProvider {
76    pub fn all() -> [Self; 8] {
77        [
78            Self::Amp,
79            Self::Wikipedia,
80            Self::Amo,
81            Self::Yelp,
82            Self::Mdn,
83            Self::Weather,
84            Self::Fakespot,
85            Self::Dynamic,
86        ]
87    }
88
89    #[inline]
90    pub(crate) fn from_u8(v: u8) -> Option<Self> {
91        match v {
92            1 => Some(Self::Amp),
93            2 => Some(Self::Wikipedia),
94            3 => Some(Self::Amo),
95            5 => Some(Self::Yelp),
96            6 => Some(Self::Mdn),
97            7 => Some(Self::Weather),
98            8 => Some(Self::Fakespot),
99            9 => Some(Self::Dynamic),
100            _ => None,
101        }
102    }
103
104    /// The collection that stores the provider's primary record.
105    pub(crate) fn primary_collection(&self) -> Collection {
106        match self {
107            Self::Amp => Collection::Amp,
108            Self::Fakespot => Collection::Fakespot,
109            _ => Collection::Other,
110        }
111    }
112
113    /// The provider's primary record type.
114    pub(crate) fn primary_record_type(&self) -> SuggestRecordType {
115        match self {
116            Self::Amp => SuggestRecordType::Amp,
117            Self::Wikipedia => SuggestRecordType::Wikipedia,
118            Self::Amo => SuggestRecordType::Amo,
119            Self::Yelp => SuggestRecordType::Yelp,
120            Self::Mdn => SuggestRecordType::Mdn,
121            Self::Weather => SuggestRecordType::Weather,
122            Self::Fakespot => SuggestRecordType::Fakespot,
123            Self::Dynamic => SuggestRecordType::Dynamic,
124        }
125    }
126
127    /// Other record types and their collections that the provider depends on.
128    fn secondary_record_types(&self) -> Option<HashMap<Collection, HashSet<SuggestRecordType>>> {
129        match self {
130            Self::Amp => Some(HashMap::from([(
131                Collection::Amp,
132                HashSet::from([SuggestRecordType::Icon]),
133            )])),
134            Self::Wikipedia => Some(HashMap::from([(
135                Collection::Other,
136                HashSet::from([SuggestRecordType::Icon]),
137            )])),
138            Self::Yelp => Some(HashMap::from([(
139                Collection::Other,
140                HashSet::from([
141                    SuggestRecordType::Icon,
142                    SuggestRecordType::Geonames,
143                    SuggestRecordType::GeonamesAlternates,
144                ]),
145            )])),
146            Self::Weather => Some(HashMap::from([(
147                Collection::Other,
148                HashSet::from([
149                    SuggestRecordType::Geonames,
150                    SuggestRecordType::GeonamesAlternates,
151                ]),
152            )])),
153            Self::Fakespot => Some(HashMap::from([(
154                Collection::Fakespot,
155                HashSet::from([SuggestRecordType::Icon]),
156            )])),
157            _ => None,
158        }
159    }
160
161    /// All record types and their collections that the provider depends on,
162    /// including primary and secondary records.
163    pub(crate) fn record_types_by_collection(
164        &self,
165    ) -> HashMap<Collection, HashSet<SuggestRecordType>> {
166        let mut rts = self.secondary_record_types().unwrap_or_default();
167        rts.entry(self.primary_collection())
168            .or_default()
169            .insert(self.primary_record_type());
170        rts
171    }
172}
173
174impl ToSql for SuggestionProvider {
175    fn to_sql(&self) -> RusqliteResult<ToSqlOutput<'_>> {
176        Ok(ToSqlOutput::from(*self as u8))
177    }
178}
179
180#[cfg(test)]
181impl SuggestionProvider {
182    pub fn record(&self, record_id: &str, attachment: JsonValue) -> MockRecord {
183        self.full_record(record_id, None, Some(MockAttachment::Json(attachment)))
184    }
185
186    pub fn empty_record(&self, record_id: &str) -> MockRecord {
187        self.full_record(record_id, None, None)
188    }
189
190    pub fn full_record(
191        &self,
192        record_id: &str,
193        inline_data: Option<JsonValue>,
194        attachment: Option<MockAttachment>,
195    ) -> MockRecord {
196        MockRecord {
197            collection: self.primary_collection(),
198            record_type: self.primary_record_type(),
199            id: record_id.to_string(),
200            inline_data,
201            attachment,
202        }
203    }
204
205    pub fn icon(&self, icon: MockIcon) -> MockRecord {
206        MockRecord {
207            collection: self.primary_collection(),
208            record_type: SuggestRecordType::Icon,
209            id: format!("icon-{}", icon.id),
210            inline_data: None,
211            attachment: Some(MockAttachment::Icon(icon)),
212        }
213    }
214}
215
216/// Some providers manage multiple suggestion subtypes. Queries, ingests, and
217/// other operations on those providers must be constrained to a desired subtype.
218#[derive(Clone, Default, Debug, uniffi::Record)]
219pub struct SuggestionProviderConstraints {
220    /// Which dynamic suggestions should we fetch or ingest? Corresponds to the
221    /// `suggestion_type` value in dynamic suggestions remote settings records.
222    #[uniffi(default = None)]
223    pub dynamic_suggestion_types: Option<Vec<String>>,
224    /// Which strategy should we use for the AMP queries?
225    /// Use None for the default strategy.
226    #[uniffi(default = None)]
227    pub amp_alternative_matching: Option<AmpMatchingStrategy>,
228}
229
230#[derive(Clone, Debug, uniffi::Enum)]
231pub enum AmpMatchingStrategy {
232    /// Disable keywords added via keyword expansion.
233    /// This eliminates keywords that for terms related to the "real" keywords, for example
234    /// misspellings like "underarmor" instead of "under armor"'.
235    NoKeywordExpansion = 1, // The desktop consumer assumes this starts at `1`
236    /// Use FTS matching against the full keywords, joined together.
237    FtsAgainstFullKeywords,
238    /// Use FTS matching against the title field
239    FtsAgainstTitle,
240}
241
242impl AmpMatchingStrategy {
243    pub fn uses_fts(&self) -> bool {
244        matches!(self, Self::FtsAgainstFullKeywords | Self::FtsAgainstTitle)
245    }
246}