Bug 1893086 - Port the potential impressions framework to the suggest…

… component

components/suggest/src/benchmarks/query.rs

@@ -40,7 +40,7 @@ impl BenchmarkWithInput for QueryBenchmark {
         InputType(SuggestionQuery {
             providers: vec![self.provider],
             keyword: self.query.to_string(),
-            limit: None,
+            ..SuggestionQuery::default()
         })
     }
 

components/suggest/src/db.rs

@@ -13,7 +13,7 @@ use rusqlite::{
     types::{FromSql, ToSql},
     Connection, OpenFlags, OptionalExtension,
 };
-use sql_support::{open_database::open_database_with_flags, ConnExt};
+use sql_support::{open_database::open_database_with_flags, repeat_sql_vars, ConnExt};
 
 use crate::{
     config::{SuggestGlobalConfig, SuggestProviderConfig},
@@ -24,8 +24,9 @@ use crate::{
     provider::SuggestionProvider,
     rs::{
         DownloadedAmoSuggestion, DownloadedAmpSuggestion, DownloadedAmpWikipediaSuggestion,
-        DownloadedFakespotSuggestion, DownloadedMdnSuggestion, DownloadedPocketSuggestion,
-        DownloadedWeatherData, DownloadedWikipediaSuggestion, Record, SuggestRecordId,
+        DownloadedExposureSuggestion, DownloadedFakespotSuggestion, DownloadedMdnSuggestion,
+        DownloadedPocketSuggestion, DownloadedWeatherData, DownloadedWikipediaSuggestion, Record,
+        SuggestRecordId,
     },
     schema::{clear_database, SuggestConnectionInitializer},
     suggestion::{cook_raw_suggestion_url, AmpSuggestionType, Suggestion},
@@ -805,6 +806,67 @@ impl<'a> SuggestDao<'a> {
         )
     }
 
+    /// Fetches exposure suggestions
+    pub fn fetch_exposure_suggestions(&self, query: &SuggestionQuery) -> Result<Vec<Suggestion>> {
+        // A single exposure suggestion can be spread across multiple remote
+        // settings records, for example if it has very many keywords. On ingest
+        // we will insert one row in `exposure_custom_details` and one row in
+        // `suggestions` per record, but that's only an implementation detail.
+        // Logically, and for consumers, there's only ever at most one exposure
+        // suggestion with a given exposure suggestion type.
+        //
+        // Why do insertions this way? It's how other suggestions work, and it
+        // lets us perform relational operations on suggestions, records, and
+        // keywords. For example, when a record is deleted we can look up its ID
+        // in `suggestions`, join the keywords table on the suggestion ID, and
+        // delete the keywords that were added by that record.
+
+        let Some(suggestion_types) = query
+            .provider_constraints
+            .as_ref()
+            .and_then(|c| c.exposure_suggestion_types.as_ref())
+        else {
+            return Ok(vec![]);
+        };
+
+        let keyword = query.keyword.to_lowercase();
+        let params = rusqlite::params_from_iter(
+            std::iter::once(&SuggestionProvider::Exposure as &dyn ToSql)
+                .chain(std::iter::once(&keyword as &dyn ToSql))
+                .chain(suggestion_types.iter().map(|t| t as &dyn ToSql)),
+        );
+        self.conn.query_rows_and_then_cached(
+            &format!(
+                r#"
+                    SELECT DISTINCT
+                      d.type
+                    FROM
+                      suggestions s
+                    JOIN
+                      exposure_custom_details d
+                      ON d.suggestion_id = s.id
+                    JOIN
+                      keywords k
+                      ON k.suggestion_id = s.id
+                    WHERE
+                      s.provider = ?
+                      AND k.keyword = ?
+                      AND d.type IN ({})
+                    ORDER BY
+                      d.type
+                    "#,
+                repeat_sql_vars(suggestion_types.len())
+            ),
+            params,
+            |row| -> Result<Suggestion> {
+                Ok(Suggestion::Exposure {
+                    suggestion_type: row.get("type")?,
+                    score: 1.0,
+                })
+            },
+        )
+    }
+
     /// Inserts all suggestions from a downloaded AMO attachment into
     /// the database.
     pub fn insert_amo_suggestions(
@@ -1058,6 +1120,39 @@ impl<'a> SuggestDao<'a> {
         Ok(())
     }
 
+    /// Inserts exposure suggestion records data into the database.
+    pub fn insert_exposure_suggestions(
+        &mut self,
+        record_id: &SuggestRecordId,
+        suggestion_type: &str,
+        suggestions: &[DownloadedExposureSuggestion],
+    ) -> Result<()> {
+        // `suggestion.keywords()` can yield duplicates for exposure
+        // suggestions, so ignore failures on insert in the uniqueness
+        // constraint on `(suggestion_id, keyword)`.
+        let mut keyword_insert = KeywordInsertStatement::new_with_or_ignore(self.conn)?;
+        let mut suggestion_insert = SuggestionInsertStatement::new(self.conn)?;
+        let mut exposure_insert = ExposureInsertStatement::new(self.conn)?;
+        for suggestion in suggestions {
+            self.scope.err_if_interrupted()?;
+            let suggestion_id = suggestion_insert.execute(
+                record_id,
+                "", // title, not used by exposure suggestions
+                "", // url, not used by exposure suggestions
+                DEFAULT_SUGGESTION_SCORE,
+                SuggestionProvider::Exposure,
+            )?;
+            exposure_insert.execute(suggestion_id, suggestion_type)?;
+
+            // Exposure suggestions don't use `rank` but `(suggestion_id, rank)`
+            // must be unique since there's an index on that tuple.
+            for (rank, keyword) in suggestion.keywords().enumerate() {
+                keyword_insert.execute(suggestion_id, &keyword, None, rank)?;
+            }
+        }
+        Ok(())
+    }
+
     /// Inserts or replaces an icon for a suggestion into the database.
     pub fn put_icon(&mut self, icon_id: &str, data: &[u8], mimetype: &str) -> Result<()> {
         self.conn.execute(
@@ -1498,6 +1593,28 @@ impl<'conn> FakespotInsertStatement<'conn> {
     }
 }
 
+struct ExposureInsertStatement<'conn>(rusqlite::Statement<'conn>);
+
+impl<'conn> ExposureInsertStatement<'conn> {
+    fn new(conn: &'conn Connection) -> Result<Self> {
+        Ok(Self(conn.prepare(
+            "INSERT INTO exposure_custom_details(
+                 suggestion_id,
+                 type
+             )
+             VALUES(?, ?)
+             ",
+        )?))
+    }
+
+    fn execute(&mut self, suggestion_id: i64, suggestion_type: &str) -> Result<()> {
+        self.0
+            .execute((suggestion_id, suggestion_type))
+            .with_context("exposure insert")?;
+        Ok(())
+    }
+}
+
 struct KeywordInsertStatement<'conn>(rusqlite::Statement<'conn>);
 
 impl<'conn> KeywordInsertStatement<'conn> {
@@ -1514,6 +1631,19 @@ impl<'conn> KeywordInsertStatement<'conn> {
         )?))
     }
 
+    fn new_with_or_ignore(conn: &'conn Connection) -> Result<Self> {
+        Ok(Self(conn.prepare(
+            "INSERT OR IGNORE INTO keywords(
+                 suggestion_id,
+                 keyword,
+                 full_keyword_id,
+                 rank
+             )
+             VALUES(?, ?, ?, ?)
+             ",
+        )?))
+    }
+
     fn execute(
         &mut self,
         suggestion_id: i64,

components/suggest/src/lib.rs

@@ -26,7 +26,7 @@ mod yelp;
 pub use config::{SuggestGlobalConfig, SuggestProviderConfig};
 pub use error::SuggestApiError;
 pub use metrics::{LabeledTimingSample, SuggestIngestionMetrics};
-pub use provider::SuggestionProvider;
+pub use provider::{SuggestionProvider, SuggestionProviderConstraints};
 pub use query::{QueryWithMetricsResult, SuggestionQuery};
 pub use store::{InterruptKind, SuggestIngestionConstraints, SuggestStore, SuggestStoreBuilder};
 pub use suggestion::{raw_suggestion_url_matches, Suggestion};

components/suggest/src/provider.rs

@@ -12,6 +12,17 @@ use rusqlite::{
 
 use crate::rs::SuggestRecordType;
 
+/// Record types from these providers will be ingested when consumers do not
+/// specify providers in `SuggestIngestionConstraints`.
+pub(crate) const DEFAULT_INGEST_PROVIDERS: [SuggestionProvider; 6] = [
+    SuggestionProvider::Amp,
+    SuggestionProvider::Wikipedia,
+    SuggestionProvider::Amo,
+    SuggestionProvider::Yelp,
+    SuggestionProvider::Mdn,
+    SuggestionProvider::AmpMobile,
+];
+
 /// A provider is a source of search suggestions.
 #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
 #[repr(u8)]
@@ -25,6 +36,7 @@ pub enum SuggestionProvider {
     Weather = 7,
     AmpMobile = 8,
     Fakespot = 9,
+    Exposure = 10,
 }
 
 impl fmt::Display for SuggestionProvider {
@@ -39,6 +51,7 @@ impl fmt::Display for SuggestionProvider {
             Self::Weather => write!(f, "weather"),
             Self::AmpMobile => write!(f, "ampmobile"),
             Self::Fakespot => write!(f, "fakespot"),
+            Self::Exposure => write!(f, "exposure"),
         }
     }
 }
@@ -54,7 +67,7 @@ impl FromSql for SuggestionProvider {
 }
 
 impl SuggestionProvider {
-    pub fn all() -> [Self; 9] {
+    pub fn all() -> [Self; 10] {
         [
             Self::Amp,
             Self::Wikipedia,
@@ -65,36 +78,39 @@ impl SuggestionProvider {
             Self::Weather,
             Self::AmpMobile,
             Self::Fakespot,
+            Self::Exposure,
         ]
     }
 
     #[inline]
     pub(crate) fn from_u8(v: u8) -> Option<Self> {
         match v {
-            1 => Some(SuggestionProvider::Amp),
-            2 => Some(SuggestionProvider::Wikipedia),
-            3 => Some(SuggestionProvider::Amo),
-            4 => Some(SuggestionProvider::Pocket),
-            5 => Some(SuggestionProvider::Yelp),
-            6 => Some(SuggestionProvider::Mdn),
-            7 => Some(SuggestionProvider::Weather),
-            8 => Some(SuggestionProvider::AmpMobile),
-            9 => Some(SuggestionProvider::Fakespot),
+            1 => Some(Self::Amp),
+            2 => Some(Self::Wikipedia),
+            3 => Some(Self::Amo),
+            4 => Some(Self::Pocket),
+            5 => Some(Self::Yelp),
+            6 => Some(Self::Mdn),
+            7 => Some(Self::Weather),
+            8 => Some(Self::AmpMobile),
+            9 => Some(Self::Fakespot),
+            10 => Some(Self::Exposure),
             _ => None,
         }
     }
 
     pub(crate) fn record_type(&self) -> SuggestRecordType {
         match self {
-            SuggestionProvider::Amp => SuggestRecordType::AmpWikipedia,
-            SuggestionProvider::Wikipedia => SuggestRecordType::AmpWikipedia,
-            SuggestionProvider::Amo => SuggestRecordType::Amo,
-            SuggestionProvider::Pocket => SuggestRecordType::Pocket,
-            SuggestionProvider::Yelp => SuggestRecordType::Yelp,
-            SuggestionProvider::Mdn => SuggestRecordType::Mdn,
-            SuggestionProvider::Weather => SuggestRecordType::Weather,
-            SuggestionProvider::AmpMobile => SuggestRecordType::AmpMobile,
-            SuggestionProvider::Fakespot => SuggestRecordType::Fakespot,
+            Self::Amp => SuggestRecordType::AmpWikipedia,
+            Self::Wikipedia => SuggestRecordType::AmpWikipedia,
+            Self::Amo => SuggestRecordType::Amo,
+            Self::Pocket => SuggestRecordType::Pocket,
+            Self::Yelp => SuggestRecordType::Yelp,
+            Self::Mdn => SuggestRecordType::Mdn,
+            Self::Weather => SuggestRecordType::Weather,
+            Self::AmpMobile => SuggestRecordType::AmpMobile,
+            Self::Fakespot => SuggestRecordType::Fakespot,
+            Self::Exposure => SuggestRecordType::Exposure,
         }
     }
 }
@@ -104,3 +120,8 @@ impl ToSql for SuggestionProvider {
         Ok(ToSqlOutput::from(*self as u8))
     }
 }
+
+#[derive(Clone, Default, Debug)]
+pub struct SuggestionProviderConstraints {
+    pub exposure_suggestion_types: Option<Vec<String>>,
+}