suggest/
db.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::{cell::OnceCell, path::Path, sync::Arc};
7
8use interrupt_support::{SqlInterruptHandle, SqlInterruptScope};
9use parking_lot::{Mutex, MutexGuard};
10use rusqlite::{
11    named_params,
12    types::{FromSql, ToSql},
13    Connection,
14};
15use sql_support::{open_database, repeat_sql_vars, ConnExt};
16
17use crate::{
18    config::{SuggestGlobalConfig, SuggestProviderConfig},
19    error::RusqliteResultExt,
20    fakespot,
21    geoname::GeonameCache,
22    provider::{AmpMatchingStrategy, SuggestionProvider},
23    query::{full_keywords_to_fts_content, FtsQuery},
24    rs::{
25        DownloadedAmoSuggestion, DownloadedAmpSuggestion, DownloadedDynamicRecord,
26        DownloadedDynamicSuggestion, DownloadedFakespotSuggestion, DownloadedMdnSuggestion,
27        DownloadedWikipediaSuggestion, Record, SuggestRecordId, SuggestRecordType,
28    },
29    schema::{clear_database, SuggestConnectionInitializer},
30    suggestion::{cook_raw_suggestion_url, FtsMatchInfo, Suggestion},
31    util::{full_keyword, i18n_transform, split_keyword},
32    weather::WeatherCache,
33    Result, SuggestionQuery,
34};
35
36/// The metadata key whose value is a JSON string encoding a
37/// `SuggestGlobalConfig`, which contains global Suggest configuration data.
38pub const GLOBAL_CONFIG_META_KEY: &str = "global_config";
39/// Prefix of metadata keys whose values are JSON strings encoding
40/// `SuggestProviderConfig`, which contains per-provider configuration data. The
41/// full key is this prefix plus the `SuggestionProvider` value as a u8.
42pub const PROVIDER_CONFIG_META_KEY_PREFIX: &str = "provider_config_";
43
44// Default value when Suggestion does not have a value for score
45pub const DEFAULT_SUGGESTION_SCORE: f64 = 0.2;
46
47/// The database connection type.
48#[derive(Clone, Copy)]
49pub(crate) enum ConnectionType {
50    ReadOnly,
51    ReadWrite,
52}
53
54#[derive(Default, Clone)]
55pub struct Sqlite3Extension {
56    pub library: String,
57    pub entry_point: Option<String>,
58}
59
60/// A thread-safe wrapper around an SQLite connection to the Suggest database,
61/// and its interrupt handle.
62pub(crate) struct SuggestDb {
63    pub conn: Mutex<Connection>,
64
65    /// An object that's used to interrupt an ongoing database operation.
66    ///
67    /// When this handle is interrupted, the thread that's currently accessing
68    /// the database will be told to stop and release the `conn` lock as soon
69    /// as possible.
70    pub interrupt_handle: Arc<SqlInterruptHandle>,
71}
72
73impl SuggestDb {
74    /// Opens a read-only or read-write connection to a Suggest database at the
75    /// given path.
76    pub fn open(
77        path: impl AsRef<Path>,
78        extensions_to_load: &[Sqlite3Extension],
79        type_: ConnectionType,
80    ) -> Result<Self> {
81        let conn = open_database::open_database_with_flags(
82            path,
83            match type_ {
84                ConnectionType::ReadWrite => open_database::read_write_flags(),
85                ConnectionType::ReadOnly => open_database::read_only_flags(),
86            },
87            &SuggestConnectionInitializer::new(extensions_to_load),
88        )?;
89        Ok(Self::with_connection(conn))
90    }
91
92    fn with_connection(conn: Connection) -> Self {
93        let interrupt_handle = Arc::new(SqlInterruptHandle::new(&conn));
94        Self {
95            conn: Mutex::new(conn),
96            interrupt_handle,
97        }
98    }
99
100    /// Accesses the Suggest database for reading.
101    pub fn read<T>(&self, op: impl FnOnce(&SuggestDao) -> Result<T>) -> Result<T> {
102        let conn = self.conn.lock();
103        let scope = self.interrupt_handle.begin_interrupt_scope()?;
104        let dao = SuggestDao::new(&conn, &scope);
105        op(&dao)
106    }
107
108    /// Accesses the Suggest database in a transaction for reading and writing.
109    pub fn write<T>(&self, op: impl FnOnce(&mut SuggestDao) -> Result<T>) -> Result<T> {
110        let mut conn = self.conn.lock();
111        let scope = self.interrupt_handle.begin_interrupt_scope()?;
112        let tx = conn.transaction()?;
113        let mut dao = SuggestDao::new(&tx, &scope);
114        let result = op(&mut dao)?;
115        tx.commit()?;
116        Ok(result)
117    }
118
119    /// Create a new write scope.
120    ///
121    /// This enables performing multiple `write()` calls with the same shared interrupt scope.
122    /// This is important for things like ingestion, where you want the operation to be interrupted
123    /// if [Self::interrupt_handle::interrupt] is called after the operation starts.  Calling
124    /// [Self::write] multiple times during the operation risks missing a call that happens after
125    /// between those calls.
126    pub fn write_scope(&self) -> Result<WriteScope> {
127        Ok(WriteScope {
128            conn: self.conn.lock(),
129            scope: self.interrupt_handle.begin_interrupt_scope()?,
130        })
131    }
132}
133
134pub(crate) struct WriteScope<'a> {
135    pub conn: MutexGuard<'a, Connection>,
136    pub scope: SqlInterruptScope,
137}
138
139impl WriteScope<'_> {
140    /// Accesses the Suggest database in a transaction for reading and writing.
141    pub fn write<T>(&mut self, op: impl FnOnce(&mut SuggestDao) -> Result<T>) -> Result<T> {
142        let tx = self.conn.transaction()?;
143        let mut dao = SuggestDao::new(&tx, &self.scope);
144        let result = op(&mut dao)?;
145        tx.commit()?;
146        Ok(result)
147    }
148
149    /// Accesses the Suggest database in a transaction for reading only
150    pub fn read<T>(&mut self, op: impl FnOnce(&SuggestDao) -> Result<T>) -> Result<T> {
151        let tx = self.conn.transaction()?;
152        let dao = SuggestDao::new(&tx, &self.scope);
153        let result = op(&dao)?;
154        tx.commit()?;
155        Ok(result)
156    }
157
158    pub fn err_if_interrupted(&self) -> Result<()> {
159        Ok(self.scope.err_if_interrupted()?)
160    }
161}
162
163/// A data access object (DAO) that wraps a connection to the Suggest database
164/// with methods for reading and writing suggestions, icons, and metadata.
165///
166/// Methods that only read from the database take an immutable reference to
167/// `self` (`&self`), and methods that write to the database take a mutable
168/// reference (`&mut self`).
169pub(crate) struct SuggestDao<'a> {
170    pub conn: &'a Connection,
171    pub scope: &'a SqlInterruptScope,
172    pub weather_cache: OnceCell<WeatherCache>,
173    pub geoname_cache: OnceCell<GeonameCache>,
174}
175
176impl<'a> SuggestDao<'a> {
177    fn new(conn: &'a Connection, scope: &'a SqlInterruptScope) -> Self {
178        Self {
179            conn,
180            scope,
181            weather_cache: std::cell::OnceCell::new(),
182            geoname_cache: std::cell::OnceCell::new(),
183        }
184    }
185
186    // =============== High level API ===============
187    //
188    //  These methods combine several low-level calls into one logical operation.
189
190    pub fn delete_record_data(&mut self, record_id: &SuggestRecordId) -> Result<()> {
191        // Drop either the icon or suggestions, records only contain one or the other
192        match record_id.as_icon_id() {
193            Some(icon_id) => self.drop_icon(icon_id)?,
194            None => self.drop_suggestions(record_id)?,
195        };
196        Ok(())
197    }
198
199    // =============== Low level API ===============
200    //
201    //  These methods implement CRUD operations
202
203    pub fn get_ingested_records(&self) -> Result<Vec<IngestedRecord>> {
204        let mut stmt = self
205            .conn
206            .prepare_cached("SELECT id, collection, type, last_modified FROM ingested_records")?;
207        let rows = stmt.query_and_then((), IngestedRecord::from_row)?;
208        rows.collect()
209    }
210
211    pub fn update_ingested_records(
212        &mut self,
213        collection: &str,
214        new_records: &[&Record],
215        updated_records: &[&Record],
216        deleted_records: &[&IngestedRecord],
217    ) -> Result<()> {
218        let mut delete_stmt = self
219            .conn
220            .prepare_cached("DELETE FROM ingested_records WHERE collection = ? AND id = ?")?;
221        for deleted in deleted_records {
222            delete_stmt.execute((collection, deleted.id.as_str()))?;
223        }
224
225        let mut insert_stmt = self.conn.prepare_cached(
226            "INSERT OR REPLACE INTO ingested_records(id, collection, type, last_modified) VALUES(?, ?, ?, ?)",
227        )?;
228        for record in new_records.iter().chain(updated_records) {
229            insert_stmt.execute((
230                record.id.as_str(),
231                collection,
232                record.record_type().as_str(),
233                record.last_modified,
234            ))?;
235        }
236        Ok(())
237    }
238
239    /// Update the DB so that we re-ingest all records on the next ingestion.
240    ///
241    /// We hack this by setting the last_modified time to 1 so that the next time around we always
242    /// re-ingest the record.
243    pub fn force_reingest(&mut self) -> Result<()> {
244        self.conn
245            .prepare_cached("UPDATE ingested_records SET last_modified=1")?
246            .execute(())?;
247        Ok(())
248    }
249
250    pub fn suggestions_table_empty(&self) -> Result<bool> {
251        Ok(self
252            .conn
253            .query_one::<bool>("SELECT NOT EXISTS (SELECT 1 FROM suggestions)")?)
254    }
255
256    /// Fetches Suggestions of type Amp provider that match the given query
257    pub fn fetch_amp_suggestions(&self, query: &SuggestionQuery) -> Result<Vec<Suggestion>> {
258        let strategy = query
259            .provider_constraints
260            .as_ref()
261            .and_then(|c| c.amp_alternative_matching.as_ref());
262        match strategy {
263            None => self.fetch_amp_suggestions_using_keywords(query, true),
264            Some(AmpMatchingStrategy::NoKeywordExpansion) => {
265                self.fetch_amp_suggestions_using_keywords(query, false)
266            }
267            Some(AmpMatchingStrategy::FtsAgainstFullKeywords) => {
268                self.fetch_amp_suggestions_using_fts(query, "full_keywords")
269            }
270            Some(AmpMatchingStrategy::FtsAgainstTitle) => {
271                self.fetch_amp_suggestions_using_fts(query, "title")
272            }
273        }
274    }
275
276    pub fn fetch_amp_suggestions_using_keywords(
277        &self,
278        query: &SuggestionQuery,
279        allow_keyword_expansion: bool,
280    ) -> Result<Vec<Suggestion>> {
281        let keyword_lowercased = &query.keyword.to_lowercase();
282        let where_extra = if allow_keyword_expansion {
283            ""
284        } else {
285            "AND INSTR(CONCAT(fk.full_keyword, ' '), k.keyword) != 0"
286        };
287        let suggestions = self.conn.query_rows_and_then_cached(
288            &format!(
289                r#"
290                SELECT
291                  s.id,
292                  k.rank,
293                  s.title,
294                  s.url,
295                  s.provider,
296                  s.score,
297                  fk.full_keyword
298                FROM
299                  suggestions s
300                JOIN
301                  keywords k
302                  ON k.suggestion_id = s.id
303                LEFT JOIN
304                  full_keywords fk
305                  ON k.full_keyword_id = fk.id
306                WHERE
307                  s.provider = :provider
308                  AND k.keyword = :keyword
309                  {where_extra}
310                  AND NOT EXISTS (
311                    -- For AMP suggestions dismissed with the deprecated URL-based dismissal API,
312                    -- `dismissed_suggestions.url` will be the suggestion URL. With the new
313                    -- `Suggestion`-based API, it will be the full keyword.
314                    SELECT 1 FROM dismissed_suggestions WHERE url IN (fk.full_keyword, s.url)
315                  )
316                "#
317            ),
318            named_params! {
319                ":keyword": keyword_lowercased,
320                ":provider": SuggestionProvider::Amp,
321            },
322            |row| -> Result<Suggestion> {
323                let suggestion_id: i64 = row.get("id")?;
324                let title = row.get("title")?;
325                let raw_url: String = row.get("url")?;
326                let score: f64 = row.get("score")?;
327                let full_keyword_from_db: Option<String> = row.get("full_keyword")?;
328
329                self.conn.query_row_and_then(
330                    r#"
331                    SELECT
332                      amp.advertiser,
333                      amp.block_id,
334                      amp.iab_category,
335                      amp.impression_url,
336                      amp.click_url,
337                      i.data AS icon,
338                      i.mimetype AS icon_mimetype
339                    FROM
340                      amp_custom_details amp
341                    LEFT JOIN
342                      icons i ON amp.icon_id = i.id
343                    WHERE
344                      amp.suggestion_id = :suggestion_id
345                    "#,
346                    named_params! {
347                        ":suggestion_id": suggestion_id
348                    },
349                    |row| {
350                        let cooked_url = cook_raw_suggestion_url(&raw_url);
351                        let raw_click_url = row.get::<_, String>("click_url")?;
352                        let cooked_click_url = cook_raw_suggestion_url(&raw_click_url);
353
354                        Ok(Suggestion::Amp {
355                            block_id: row.get("block_id")?,
356                            advertiser: row.get("advertiser")?,
357                            iab_category: row.get("iab_category")?,
358                            title,
359                            url: cooked_url,
360                            raw_url,
361                            full_keyword: full_keyword_from_db.unwrap_or_default(),
362                            icon: row.get("icon")?,
363                            icon_mimetype: row.get("icon_mimetype")?,
364                            impression_url: row.get("impression_url")?,
365                            click_url: cooked_click_url,
366                            raw_click_url,
367                            score,
368                            fts_match_info: None,
369                        })
370                    },
371                )
372            },
373        )?;
374        Ok(suggestions)
375    }
376
377    pub fn fetch_amp_suggestions_using_fts(
378        &self,
379        query: &SuggestionQuery,
380        fts_column: &str,
381    ) -> Result<Vec<Suggestion>> {
382        let fts_query = query.fts_query();
383        let match_arg = &fts_query.match_arg;
384        let suggestions = self.conn.query_rows_and_then_cached(
385            &format!(
386                r#"
387                SELECT
388                  s.id,
389                  s.title,
390                  s.url,
391                  s.provider,
392                  s.score
393                FROM
394                  suggestions s
395                JOIN
396                  amp_fts fts
397                  ON fts.rowid = s.id
398                WHERE
399                  s.provider = :provider
400                  AND amp_fts match '{fts_column}: {match_arg}'
401                AND NOT EXISTS (SELECT 1 FROM dismissed_suggestions WHERE url=s.url)
402                ORDER BY rank
403                LIMIT 1
404                "#
405            ),
406            named_params! {
407                ":provider": SuggestionProvider::Amp,
408            },
409            |row| -> Result<Suggestion> {
410                let suggestion_id: i64 = row.get("id")?;
411                let title: String = row.get("title")?;
412                let raw_url: String = row.get("url")?;
413                let score: f64 = row.get("score")?;
414
415                self.conn.query_row_and_then(
416                    r#"
417                    SELECT
418                      amp.advertiser,
419                      amp.block_id,
420                      amp.iab_category,
421                      amp.impression_url,
422                      amp.click_url,
423                      i.data AS icon,
424                      i.mimetype AS icon_mimetype
425                    FROM
426                      amp_custom_details amp
427                    LEFT JOIN
428                      icons i ON amp.icon_id = i.id
429                    WHERE
430                      amp.suggestion_id = :suggestion_id
431                    "#,
432                    named_params! {
433                        ":suggestion_id": suggestion_id
434                    },
435                    |row| {
436                        let cooked_url = cook_raw_suggestion_url(&raw_url);
437                        let raw_click_url = row.get::<_, String>("click_url")?;
438                        let cooked_click_url = cook_raw_suggestion_url(&raw_click_url);
439                        let match_info = self.fetch_amp_fts_match_info(
440                            &fts_query,
441                            suggestion_id,
442                            fts_column,
443                            &title,
444                        )?;
445
446                        Ok(Suggestion::Amp {
447                            block_id: row.get("block_id")?,
448                            advertiser: row.get("advertiser")?,
449                            iab_category: row.get("iab_category")?,
450                            title,
451                            url: cooked_url,
452                            raw_url,
453                            full_keyword: query.keyword.clone(),
454                            icon: row.get("icon")?,
455                            icon_mimetype: row.get("icon_mimetype")?,
456                            impression_url: row.get("impression_url")?,
457                            click_url: cooked_click_url,
458                            raw_click_url,
459                            score,
460                            fts_match_info: Some(match_info),
461                        })
462                    },
463                )
464            },
465        )?;
466        Ok(suggestions)
467    }
468
469    fn fetch_amp_fts_match_info(
470        &self,
471        fts_query: &FtsQuery<'_>,
472        suggestion_id: i64,
473        fts_column: &str,
474        title: &str,
475    ) -> Result<FtsMatchInfo> {
476        let fts_content = match fts_column {
477            "title" => title.to_lowercase(),
478            "full_keywords" => {
479                let full_keyword_list: Vec<String> = self.conn.query_rows_and_then(
480                    "
481                    SELECT fk.full_keyword
482                    FROM full_keywords fk
483                    JOIN keywords k on fk.id == k.full_keyword_id
484                    WHERE k.suggestion_id = ?
485                    ",
486                    (suggestion_id,),
487                    |row| row.get(0),
488                )?;
489                full_keywords_to_fts_content(full_keyword_list.iter().map(String::as_str))
490            }
491            // fts_column comes from the code above and we know there's only 2 possibilities
492            _ => unreachable!(),
493        };
494
495        let prefix = if fts_query.is_prefix_query {
496            // If the query was a prefix match query then test if the query without the prefix
497            // match would have also matched.  If not, then this counts as a prefix match.
498            let sql = "SELECT 1 FROM amp_fts WHERE rowid = ? AND amp_fts MATCH ?";
499            let params = (&suggestion_id, &fts_query.match_arg_without_prefix_match);
500            !self.conn.exists(sql, params)?
501        } else {
502            // If not, then it definitely wasn't a prefix match
503            false
504        };
505
506        Ok(FtsMatchInfo {
507            prefix,
508            stemming: fts_query.match_required_stemming(&fts_content),
509        })
510    }
511
512    /// Fetches Suggestions of type Wikipedia provider that match the given query
513    pub fn fetch_wikipedia_suggestions(&self, query: &SuggestionQuery) -> Result<Vec<Suggestion>> {
514        let keyword_lowercased = &query.keyword.to_lowercase();
515        let suggestions = self.conn.query_rows_and_then_cached(
516            r#"
517            SELECT
518              s.id,
519              k.rank,
520              s.title,
521              s.url
522            FROM
523              suggestions s
524            JOIN
525              keywords k
526              ON k.suggestion_id = s.id
527            WHERE
528              s.provider = :provider
529              AND k.keyword = :keyword
530              AND NOT EXISTS (SELECT 1 FROM dismissed_suggestions WHERE url=s.url)
531            "#,
532            named_params! {
533                ":keyword": keyword_lowercased,
534                ":provider": SuggestionProvider::Wikipedia
535            },
536            |row| -> Result<Suggestion> {
537                let suggestion_id: i64 = row.get("id")?;
538                let title = row.get("title")?;
539                let raw_url = row.get::<_, String>("url")?;
540
541                let keywords: Vec<String> = self.conn.query_rows_and_then_cached(
542                    "SELECT keyword FROM keywords
543                     WHERE suggestion_id = :suggestion_id AND rank >= :rank
544                     ORDER BY rank ASC",
545                    named_params! {
546                        ":suggestion_id": suggestion_id,
547                        ":rank": row.get::<_, i64>("rank")?,
548                    },
549                    |row| row.get(0),
550                )?;
551                let (icon, icon_mimetype) = self
552                    .conn
553                    .try_query_row(
554                        "SELECT i.data, i.mimetype
555                     FROM icons i
556                     JOIN wikipedia_custom_details s ON s.icon_id = i.id
557                     WHERE s.suggestion_id = :suggestion_id
558                     LIMIT 1",
559                        named_params! {
560                            ":suggestion_id": suggestion_id
561                        },
562                        |row| -> Result<_> {
563                            Ok((
564                                row.get::<_, Option<Vec<u8>>>(0)?,
565                                row.get::<_, Option<String>>(1)?,
566                            ))
567                        },
568                        true,
569                    )?
570                    .unwrap_or((None, None));
571
572                Ok(Suggestion::Wikipedia {
573                    title,
574                    url: raw_url,
575                    full_keyword: full_keyword(keyword_lowercased, &keywords),
576                    icon,
577                    icon_mimetype,
578                })
579            },
580        )?;
581        Ok(suggestions)
582    }
583
584    /// Query for suggestions using the keyword prefix and provider
585    fn map_prefix_keywords<T>(
586        &self,
587        query: &SuggestionQuery,
588        provider: &SuggestionProvider,
589        mut mapper: impl FnMut(&rusqlite::Row, &str) -> Result<T>,
590    ) -> Result<Vec<T>> {
591        let keyword_lowercased = &query.keyword.to_lowercase();
592        let (keyword_prefix, keyword_suffix) = split_keyword(keyword_lowercased);
593        let suggestions_limit = query.limit.unwrap_or(-1);
594        self.conn.query_rows_and_then_cached(
595            r#"
596                SELECT
597                  s.id,
598                  MAX(k.rank) AS rank,
599                  s.title,
600                  s.url,
601                  s.provider,
602                  s.score,
603                  k.keyword_suffix
604                FROM
605                  suggestions s
606                JOIN
607                  prefix_keywords k
608                  ON k.suggestion_id = s.id
609                WHERE
610                  k.keyword_prefix = :keyword_prefix
611                  AND (k.keyword_suffix BETWEEN :keyword_suffix AND :keyword_suffix || x'FFFF')
612                  AND s.provider = :provider
613                  AND NOT EXISTS (SELECT 1 FROM dismissed_suggestions WHERE url=s.url)
614                GROUP BY
615                  s.id
616                ORDER BY
617                  s.score DESC,
618                  rank DESC
619                LIMIT
620                  :suggestions_limit
621                "#,
622            &[
623                (":keyword_prefix", &keyword_prefix as &dyn ToSql),
624                (":keyword_suffix", &keyword_suffix as &dyn ToSql),
625                (":provider", provider as &dyn ToSql),
626                (":suggestions_limit", &suggestions_limit as &dyn ToSql),
627            ],
628            |row| mapper(row, keyword_suffix),
629        )
630    }
631
632    /// Fetches Suggestions of type Amo provider that match the given query
633    pub fn fetch_amo_suggestions(&self, query: &SuggestionQuery) -> Result<Vec<Suggestion>> {
634        let suggestions = self
635            .map_prefix_keywords(
636                query,
637                &SuggestionProvider::Amo,
638                |row, keyword_suffix| -> Result<Option<Suggestion>> {
639                    let suggestion_id: i64 = row.get("id")?;
640                    let title = row.get("title")?;
641                    let raw_url = row.get::<_, String>("url")?;
642                    let score = row.get::<_, f64>("score")?;
643
644                    let full_suffix = row.get::<_, String>("keyword_suffix")?;
645                    full_suffix
646                        .starts_with(keyword_suffix)
647                        .then(|| {
648                            self.conn.query_row_and_then(
649                                r#"
650                                SELECT
651                                  amo.description,
652                                  amo.guid,
653                                  amo.rating,
654                                  amo.icon_url,
655                                  amo.number_of_ratings
656                                FROM
657                                  amo_custom_details amo
658                                WHERE
659                                  amo.suggestion_id = :suggestion_id
660                                "#,
661                                named_params! {
662                                    ":suggestion_id": suggestion_id
663                                },
664                                |row| {
665                                    Ok(Suggestion::Amo {
666                                        title,
667                                        url: raw_url,
668                                        icon_url: row.get("icon_url")?,
669                                        description: row.get("description")?,
670                                        rating: row.get("rating")?,
671                                        number_of_ratings: row.get("number_of_ratings")?,
672                                        guid: row.get("guid")?,
673                                        score,
674                                    })
675                                },
676                            )
677                        })
678                        .transpose()
679                },
680            )?
681            .into_iter()
682            .flatten()
683            .collect();
684        Ok(suggestions)
685    }
686
687    /// Fetches suggestions for MDN
688    pub fn fetch_mdn_suggestions(&self, query: &SuggestionQuery) -> Result<Vec<Suggestion>> {
689        let suggestions = self
690            .map_prefix_keywords(
691                query,
692                &SuggestionProvider::Mdn,
693                |row, keyword_suffix| -> Result<Option<Suggestion>> {
694                    let suggestion_id: i64 = row.get("id")?;
695                    let title = row.get("title")?;
696                    let raw_url = row.get::<_, String>("url")?;
697                    let score = row.get::<_, f64>("score")?;
698
699                    let full_suffix = row.get::<_, String>("keyword_suffix")?;
700                    full_suffix
701                        .starts_with(keyword_suffix)
702                        .then(|| {
703                            self.conn.query_row_and_then(
704                                r#"
705                                SELECT
706                                    description
707                                FROM
708                                    mdn_custom_details
709                                WHERE
710                                    suggestion_id = :suggestion_id
711                                "#,
712                                named_params! {
713                                    ":suggestion_id": suggestion_id
714                                },
715                                |row| {
716                                    Ok(Suggestion::Mdn {
717                                        title,
718                                        url: raw_url,
719                                        description: row.get("description")?,
720                                        score,
721                                    })
722                                },
723                            )
724                        })
725                        .transpose()
726                },
727            )?
728            .into_iter()
729            .flatten()
730            .collect();
731
732        Ok(suggestions)
733    }
734
735    /// Fetches Fakespot suggestions
736    pub fn fetch_fakespot_suggestions(&self, query: &SuggestionQuery) -> Result<Vec<Suggestion>> {
737        let fts_query = query.fts_query();
738        let sql = r#"
739            SELECT
740                s.id,
741                s.title,
742                s.url,
743                s.score,
744                f.fakespot_grade,
745                f.product_id,
746                f.rating,
747                f.total_reviews,
748                i.data,
749                i.mimetype,
750                f.keywords,
751                f.product_type
752            FROM
753                suggestions s
754            JOIN
755                fakespot_fts fts
756                ON fts.rowid = s.id
757            JOIN
758                fakespot_custom_details f
759                ON f.suggestion_id = s.id
760            LEFT JOIN
761                icons i
762                ON i.id = f.icon_id
763            WHERE
764                fakespot_fts MATCH ?
765            ORDER BY
766                s.score DESC
767            "#
768        .to_string();
769
770        // Store the list of results plus the suggestion id for calculating the FTS match info
771        let mut results =
772            self.conn
773                .query_rows_and_then_cached(&sql, (&fts_query.match_arg,), |row| {
774                    let id: usize = row.get(0)?;
775                    let score = fakespot::FakespotScore::new(
776                        &query.keyword,
777                        row.get(10)?,
778                        row.get(11)?,
779                        row.get(3)?,
780                    )
781                    .as_suggest_score();
782                    Result::Ok((
783                        Suggestion::Fakespot {
784                            title: row.get(1)?,
785                            url: row.get(2)?,
786                            score,
787                            fakespot_grade: row.get(4)?,
788                            product_id: row.get(5)?,
789                            rating: row.get(6)?,
790                            total_reviews: row.get(7)?,
791                            icon: row.get(8)?,
792                            icon_mimetype: row.get(9)?,
793                            match_info: None,
794                        },
795                        id,
796                    ))
797                })?;
798        // Sort the results, then add the FTS match info to the first one
799        // For performance reasons, this is only calculated for the result with the highest score.
800        // We assume that only one that will be shown to the user and therefore the only one we'll
801        // collect metrics for.
802        results.sort();
803        if let Some((suggestion, id)) = results.first_mut() {
804            match suggestion {
805                Suggestion::Fakespot {
806                    match_info, title, ..
807                } => {
808                    *match_info = Some(self.fetch_fakespot_fts_match_info(&fts_query, *id, title)?);
809                }
810                _ => unreachable!(),
811            }
812        }
813        Ok(results
814            .into_iter()
815            .map(|(suggestion, _)| suggestion)
816            .collect())
817    }
818
819    fn fetch_fakespot_fts_match_info(
820        &self,
821        fts_query: &FtsQuery<'_>,
822        suggestion_id: usize,
823        title: &str,
824    ) -> Result<FtsMatchInfo> {
825        let prefix = if fts_query.is_prefix_query {
826            // If the query was a prefix match query then test if the query without the prefix
827            // match would have also matched.  If not, then this counts as a prefix match.
828            let sql = "SELECT 1 FROM fakespot_fts WHERE rowid = ? AND fakespot_fts MATCH ?";
829            let params = (&suggestion_id, &fts_query.match_arg_without_prefix_match);
830            !self.conn.exists(sql, params)?
831        } else {
832            // If not, then it definitely wasn't a prefix match
833            false
834        };
835
836        Ok(FtsMatchInfo {
837            prefix,
838            stemming: fts_query.match_required_stemming(title),
839        })
840    }
841
842    /// Fetches dynamic suggestions
843    pub fn fetch_dynamic_suggestions(&self, query: &SuggestionQuery) -> Result<Vec<Suggestion>> {
844        let Some(suggestion_types) = query
845            .provider_constraints
846            .as_ref()
847            .and_then(|c| c.dynamic_suggestion_types.as_ref())
848        else {
849            return Ok(vec![]);
850        };
851
852        let keyword = query.keyword.to_lowercase();
853        let params = rusqlite::params_from_iter(
854            std::iter::once(&SuggestionProvider::Dynamic as &dyn ToSql)
855                .chain(std::iter::once(&keyword as &dyn ToSql))
856                .chain(suggestion_types.iter().map(|t| t as &dyn ToSql)),
857        );
858        self.conn.query_rows_and_then_cached(
859            &format!(
860                r#"
861                SELECT
862                  s.url,
863                  s.score,
864                  d.suggestion_type,
865                  d.json_data
866                FROM
867                  suggestions s
868                JOIN
869                  dynamic_custom_details d
870                  ON d.suggestion_id = s.id
871                JOIN
872                  keywords k
873                  ON k.suggestion_id = s.id
874                WHERE
875                  s.provider = ?
876                  AND k.keyword = ?
877                  AND d.suggestion_type IN ({})
878                  AND NOT EXISTS (SELECT 1 FROM dismissed_suggestions WHERE url = s.url)
879                ORDER BY
880                  s.score ASC, d.suggestion_type ASC, s.id ASC
881                "#,
882                repeat_sql_vars(suggestion_types.len())
883            ),
884            params,
885            |row| -> Result<Suggestion> {
886                let dismissal_key: String = row.get("url")?;
887                let json_data: Option<String> = row.get("json_data")?;
888                Ok(Suggestion::Dynamic {
889                    suggestion_type: row.get("suggestion_type")?,
890                    data: match json_data {
891                        None => None,
892                        Some(j) => serde_json::from_str(&j)?,
893                    },
894                    score: row.get("score")?,
895                    dismissal_key: (!dismissal_key.is_empty()).then_some(dismissal_key),
896                })
897            },
898        )
899    }
900
901    pub fn are_suggestions_ingested_for_record(&self, record_id: &SuggestRecordId) -> Result<bool> {
902        Ok(self.conn.exists(
903            r#"
904            SELECT
905              id
906            FROM
907              suggestions
908            WHERE
909              record_id = :record_id
910            "#,
911            named_params! {
912                ":record_id": record_id.as_str(),
913            },
914        )?)
915    }
916
917    pub fn is_amp_fts_data_ingested(&self, record_id: &SuggestRecordId) -> Result<bool> {
918        Ok(self.conn.exists(
919            r#"
920            SELECT 1
921            FROM suggestions s
922            JOIN amp_fts fts
923              ON fts.rowid = s.id
924            WHERE s.record_id = :record_id
925            "#,
926            named_params! {
927                ":record_id": record_id.as_str(),
928            },
929        )?)
930    }
931
932    /// Inserts all suggestions from a downloaded AMO attachment into
933    /// the database.
934    pub fn insert_amo_suggestions(
935        &mut self,
936        record_id: &SuggestRecordId,
937        suggestions: &[DownloadedAmoSuggestion],
938    ) -> Result<()> {
939        let mut suggestion_insert = SuggestionInsertStatement::new(self.conn)?;
940        let mut amo_insert = AmoInsertStatement::new(self.conn)?;
941        let mut prefix_keyword_insert = PrefixKeywordInsertStatement::new(self.conn)?;
942        for suggestion in suggestions {
943            self.scope.err_if_interrupted()?;
944            let suggestion_id = suggestion_insert.execute(
945                record_id,
946                &suggestion.title,
947                &suggestion.url,
948                suggestion.score,
949                SuggestionProvider::Amo,
950            )?;
951            amo_insert.execute(suggestion_id, suggestion)?;
952            for (index, keyword) in suggestion.keywords.iter().enumerate() {
953                let (keyword_prefix, keyword_suffix) = split_keyword(keyword);
954                prefix_keyword_insert.execute(
955                    suggestion_id,
956                    None,
957                    keyword_prefix,
958                    keyword_suffix,
959                    index,
960                )?;
961            }
962        }
963        Ok(())
964    }
965
966    /// Inserts suggestions from an AMP attachment into the database.
967    pub fn insert_amp_suggestions(
968        &mut self,
969        record_id: &SuggestRecordId,
970        suggestions: &[DownloadedAmpSuggestion],
971        enable_fts: bool,
972    ) -> Result<()> {
973        // Prepare statements outside of the loop.  This results in a large performance
974        // improvement on a fresh ingest, since there are so many rows.
975        let mut suggestion_insert = SuggestionInsertStatement::new(self.conn)?;
976        let mut amp_insert = AmpInsertStatement::new(self.conn)?;
977        let mut keyword_insert = KeywordInsertStatement::new(self.conn)?;
978        let mut fts_insert = AmpFtsInsertStatement::new(self.conn)?;
979        for suggestion in suggestions {
980            self.scope.err_if_interrupted()?;
981            let suggestion_id = suggestion_insert.execute(
982                record_id,
983                &suggestion.title,
984                &suggestion.url,
985                suggestion.score.unwrap_or(DEFAULT_SUGGESTION_SCORE),
986                SuggestionProvider::Amp,
987            )?;
988            amp_insert.execute(suggestion_id, suggestion)?;
989            if enable_fts {
990                fts_insert.execute(
991                    suggestion_id,
992                    &suggestion.full_keywords_fts_column(),
993                    &suggestion.title,
994                )?;
995            }
996            let mut full_keyword_inserter = FullKeywordInserter::new(self.conn, suggestion_id);
997            for keyword in suggestion.keywords() {
998                let full_keyword_id = if let Some(full_keyword) = keyword.full_keyword {
999                    Some(full_keyword_inserter.maybe_insert(full_keyword)?)
1000                } else {
1001                    None
1002                };
1003                keyword_insert.execute(
1004                    suggestion_id,
1005                    keyword.keyword,
1006                    full_keyword_id,
1007                    keyword.rank,
1008                )?;
1009            }
1010        }
1011        Ok(())
1012    }
1013
1014    /// Inserts suggestions from a Wikipedia attachment into the database.
1015    pub fn insert_wikipedia_suggestions(
1016        &mut self,
1017        record_id: &SuggestRecordId,
1018        suggestions: &[DownloadedWikipediaSuggestion],
1019    ) -> Result<()> {
1020        // Prepare statements outside of the loop.  This results in a large performance
1021        // improvement on a fresh ingest, since there are so many rows.
1022        let mut suggestion_insert = SuggestionInsertStatement::new(self.conn)?;
1023        let mut wiki_insert = WikipediaInsertStatement::new(self.conn)?;
1024        let mut keyword_insert = KeywordInsertStatement::new(self.conn)?;
1025        for suggestion in suggestions {
1026            self.scope.err_if_interrupted()?;
1027            let suggestion_id = suggestion_insert.execute(
1028                record_id,
1029                &suggestion.title,
1030                &suggestion.url,
1031                suggestion.score.unwrap_or(DEFAULT_SUGGESTION_SCORE),
1032                SuggestionProvider::Wikipedia,
1033            )?;
1034            wiki_insert.execute(suggestion_id, suggestion)?;
1035            for keyword in suggestion.keywords() {
1036                // Don't update `full_keywords`, see bug 1876217.
1037                keyword_insert.execute(suggestion_id, keyword.keyword, None, keyword.rank)?;
1038            }
1039        }
1040        Ok(())
1041    }
1042
1043    /// Inserts all suggestions from a downloaded MDN attachment into
1044    /// the database.
1045    pub fn insert_mdn_suggestions(
1046        &mut self,
1047        record_id: &SuggestRecordId,
1048        suggestions: &[DownloadedMdnSuggestion],
1049    ) -> Result<()> {
1050        let mut suggestion_insert = SuggestionInsertStatement::new(self.conn)?;
1051        let mut mdn_insert = MdnInsertStatement::new(self.conn)?;
1052        let mut prefix_keyword_insert = PrefixKeywordInsertStatement::new(self.conn)?;
1053        for suggestion in suggestions {
1054            self.scope.err_if_interrupted()?;
1055            let suggestion_id = suggestion_insert.execute(
1056                record_id,
1057                &suggestion.title,
1058                &suggestion.url,
1059                suggestion.score,
1060                SuggestionProvider::Mdn,
1061            )?;
1062            mdn_insert.execute(suggestion_id, suggestion)?;
1063            for (index, keyword) in suggestion.keywords.iter().enumerate() {
1064                let (keyword_prefix, keyword_suffix) = split_keyword(keyword);
1065                prefix_keyword_insert.execute(
1066                    suggestion_id,
1067                    None,
1068                    keyword_prefix,
1069                    keyword_suffix,
1070                    index,
1071                )?;
1072            }
1073        }
1074        Ok(())
1075    }
1076
1077    /// Inserts all suggestions from a downloaded Fakespot attachment into the database.
1078    pub fn insert_fakespot_suggestions(
1079        &mut self,
1080        record_id: &SuggestRecordId,
1081        suggestions: &[DownloadedFakespotSuggestion],
1082    ) -> Result<()> {
1083        let mut suggestion_insert = SuggestionInsertStatement::new(self.conn)?;
1084        let mut fakespot_insert = FakespotInsertStatement::new(self.conn)?;
1085        for suggestion in suggestions {
1086            let suggestion_id = suggestion_insert.execute(
1087                record_id,
1088                &suggestion.title,
1089                &suggestion.url,
1090                suggestion.score,
1091                SuggestionProvider::Fakespot,
1092            )?;
1093            fakespot_insert.execute(suggestion_id, suggestion)?;
1094        }
1095        Ok(())
1096    }
1097
1098    /// Inserts dynamic suggestion records data into the database.
1099    pub fn insert_dynamic_suggestions(
1100        &mut self,
1101        record_id: &SuggestRecordId,
1102        record: &DownloadedDynamicRecord,
1103        suggestions: &[DownloadedDynamicSuggestion],
1104    ) -> Result<()> {
1105        // `suggestion.keywords()` can yield duplicates for dynamic
1106        // suggestions, so ignore failures on insert in the uniqueness
1107        // constraint on `(suggestion_id, keyword)`.
1108        let mut keyword_insert = KeywordInsertStatement::with_details(
1109            self.conn,
1110            "keywords",
1111            Some(InsertConflictResolution::Ignore),
1112        )?;
1113        let mut suggestion_insert = SuggestionInsertStatement::new(self.conn)?;
1114        let mut dynamic_insert = DynamicInsertStatement::new(self.conn)?;
1115        for suggestion in suggestions {
1116            self.scope.err_if_interrupted()?;
1117            let suggestion_id = suggestion_insert.execute(
1118                record_id,
1119                // title - Not used by dynamic suggestions.
1120                "",
1121                // url - Dynamic suggestions store their dismissal key here
1122                // instead.
1123                suggestion.dismissal_key.as_deref().unwrap_or(""),
1124                record.score.unwrap_or(DEFAULT_SUGGESTION_SCORE),
1125                SuggestionProvider::Dynamic,
1126            )?;
1127            dynamic_insert.execute(suggestion_id, &record.suggestion_type, suggestion)?;
1128
1129            // Dynamic suggestions don't use `rank` but `(suggestion_id, rank)`
1130            // must be unique since there's an index on that tuple.
1131            for (rank, keyword) in suggestion.keywords().enumerate() {
1132                keyword_insert.execute(suggestion_id, &keyword, None, rank)?;
1133            }
1134        }
1135        Ok(())
1136    }
1137
1138    /// Inserts or replaces an icon for a suggestion into the database.
1139    pub fn put_icon(&mut self, icon_id: &str, data: &[u8], mimetype: &str) -> Result<()> {
1140        self.conn.execute(
1141            "INSERT OR REPLACE INTO icons(
1142                 id,
1143                 data,
1144                 mimetype
1145             )
1146             VALUES(
1147                 :id,
1148                 :data,
1149                 :mimetype
1150             )",
1151            named_params! {
1152                ":id": icon_id,
1153                ":data": data,
1154                ":mimetype": mimetype,
1155            },
1156        )?;
1157        Ok(())
1158    }
1159
1160    pub fn insert_dismissal(&self, key: &str) -> Result<()> {
1161        self.conn.execute(
1162            "INSERT OR IGNORE INTO dismissed_suggestions(url)
1163             VALUES(:url)",
1164            named_params! {
1165                ":url": key,
1166            },
1167        )?;
1168        Ok(())
1169    }
1170
1171    pub fn clear_dismissals(&self) -> Result<()> {
1172        self.conn.execute("DELETE FROM dismissed_suggestions", ())?;
1173        Ok(())
1174    }
1175
1176    pub fn has_dismissal(&self, key: &str) -> Result<bool> {
1177        Ok(self.conn.exists(
1178            "SELECT 1 FROM dismissed_suggestions WHERE url = :url",
1179            named_params! {
1180                ":url": key,
1181            },
1182        )?)
1183    }
1184
1185    pub fn any_dismissals(&self) -> Result<bool> {
1186        Ok(self
1187            .conn
1188            .exists("SELECT 1 FROM dismissed_suggestions LIMIT 1", ())?)
1189    }
1190
1191    /// Deletes all suggestions associated with a Remote Settings record from
1192    /// the database.
1193    pub fn drop_suggestions(&mut self, record_id: &SuggestRecordId) -> Result<()> {
1194        // If you update this, you probably need to update
1195        // `schema::clear_database()` too!
1196        //
1197        // Call `err_if_interrupted` before each statement since these have historically taken a
1198        // long time and caused shutdown hangs.
1199
1200        self.scope.err_if_interrupted()?;
1201        self.conn.execute_cached(
1202            "DELETE FROM keywords WHERE suggestion_id IN (SELECT id from suggestions WHERE record_id = :record_id)",
1203            named_params! { ":record_id": record_id.as_str() },
1204        )?;
1205        self.scope.err_if_interrupted()?;
1206        self.conn.execute_cached(
1207            "DELETE FROM keywords_i18n WHERE suggestion_id IN (SELECT id from suggestions WHERE record_id = :record_id)",
1208            named_params! { ":record_id": record_id.as_str() },
1209        )?;
1210        self.scope.err_if_interrupted()?;
1211        self.conn.execute_cached(
1212            "DELETE FROM full_keywords WHERE suggestion_id IN (SELECT id from suggestions WHERE record_id = :record_id)",
1213            named_params! { ":record_id": record_id.as_str() },
1214        )?;
1215        self.scope.err_if_interrupted()?;
1216        self.conn.execute_cached(
1217            "DELETE FROM prefix_keywords WHERE suggestion_id IN (SELECT id from suggestions WHERE record_id = :record_id)",
1218            named_params! { ":record_id": record_id.as_str() },
1219        )?;
1220        self.scope.err_if_interrupted()?;
1221        self.conn.execute_cached(
1222            "DELETE FROM keywords_metrics WHERE record_id = :record_id",
1223            named_params! { ":record_id": record_id.as_str() },
1224        )?;
1225        self.scope.err_if_interrupted()?;
1226        self.conn.execute_cached(
1227            "
1228            DELETE FROM fakespot_fts
1229            WHERE rowid IN (SELECT id from suggestions WHERE record_id = :record_id)
1230            ",
1231            named_params! { ":record_id": record_id.as_str() },
1232        )?;
1233        self.scope.err_if_interrupted()?;
1234        self.conn.execute_cached(
1235            "DELETE FROM suggestions WHERE record_id = :record_id",
1236            named_params! { ":record_id": record_id.as_str() },
1237        )?;
1238        self.scope.err_if_interrupted()?;
1239        self.conn.execute_cached(
1240            "DELETE FROM yelp_subjects WHERE record_id = :record_id",
1241            named_params! { ":record_id": record_id.as_str() },
1242        )?;
1243        self.scope.err_if_interrupted()?;
1244        self.conn.execute_cached(
1245            "DELETE FROM yelp_modifiers WHERE record_id = :record_id",
1246            named_params! { ":record_id": record_id.as_str() },
1247        )?;
1248        self.scope.err_if_interrupted()?;
1249        self.conn.execute_cached(
1250            "DELETE FROM yelp_custom_details WHERE record_id = :record_id",
1251            named_params! { ":record_id": record_id.as_str() },
1252        )?;
1253        self.scope.err_if_interrupted()?;
1254        self.conn.execute_cached(
1255            "DELETE FROM geonames WHERE record_id = :record_id",
1256            named_params! { ":record_id": record_id.as_str() },
1257        )?;
1258        self.scope.err_if_interrupted()?;
1259        self.conn.execute_cached(
1260            "DELETE FROM geonames_alternates WHERE record_id = :record_id",
1261            named_params! { ":record_id": record_id.as_str() },
1262        )?;
1263        self.scope.err_if_interrupted()?;
1264        self.conn.execute_cached(
1265            "DELETE FROM geonames_metrics WHERE record_id = :record_id",
1266            named_params! { ":record_id": record_id.as_str() },
1267        )?;
1268
1269        // Invalidate these caches since we might have deleted a record their
1270        // contents are based on.
1271        self.weather_cache.take();
1272        self.geoname_cache.take();
1273
1274        Ok(())
1275    }
1276
1277    /// Deletes an icon for a suggestion from the database.
1278    pub fn drop_icon(&mut self, icon_id: &str) -> Result<()> {
1279        self.conn.execute_cached(
1280            "DELETE FROM icons WHERE id = :id",
1281            named_params! { ":id": icon_id },
1282        )?;
1283        Ok(())
1284    }
1285
1286    /// Clears the database, removing all suggestions, icons, and metadata.
1287    pub fn clear(&mut self) -> Result<()> {
1288        Ok(clear_database(self.conn)?)
1289    }
1290
1291    /// Returns the value associated with a metadata key.
1292    pub fn get_meta<T: FromSql>(&self, key: &str) -> Result<Option<T>> {
1293        Ok(self.conn.try_query_one(
1294            "SELECT value FROM meta WHERE key = :key",
1295            named_params! { ":key": key },
1296            true,
1297        )?)
1298    }
1299
1300    /// Sets the value for a metadata key.
1301    pub fn put_meta(&mut self, key: &str, value: impl ToSql) -> Result<()> {
1302        self.conn.execute_cached(
1303            "INSERT OR REPLACE INTO meta(key, value) VALUES(:key, :value)",
1304            named_params! { ":key": key, ":value": value },
1305        )?;
1306        Ok(())
1307    }
1308
1309    /// Stores global Suggest configuration data.
1310    pub fn put_global_config(&mut self, config: &SuggestGlobalConfig) -> Result<()> {
1311        self.put_meta(GLOBAL_CONFIG_META_KEY, serde_json::to_string(config)?)
1312    }
1313
1314    /// Gets the stored global Suggest configuration data or a default config if
1315    /// none is stored.
1316    pub fn get_global_config(&self) -> Result<SuggestGlobalConfig> {
1317        self.get_meta::<String>(GLOBAL_CONFIG_META_KEY)?
1318            .map_or_else(
1319                || Ok(SuggestGlobalConfig::default()),
1320                |json| Ok(serde_json::from_str(&json)?),
1321            )
1322    }
1323
1324    /// Stores configuration data for a given provider.
1325    pub fn put_provider_config(
1326        &mut self,
1327        provider: SuggestionProvider,
1328        config: &SuggestProviderConfig,
1329    ) -> Result<()> {
1330        self.put_meta(
1331            &provider_config_meta_key(provider),
1332            serde_json::to_string(config)?,
1333        )
1334    }
1335
1336    /// Gets the stored configuration data for a given provider or None if none
1337    /// is stored.
1338    pub fn get_provider_config(
1339        &self,
1340        provider: SuggestionProvider,
1341    ) -> Result<Option<SuggestProviderConfig>> {
1342        self.get_meta::<String>(&provider_config_meta_key(provider))?
1343            .map_or_else(|| Ok(None), |json| Ok(serde_json::from_str(&json)?))
1344    }
1345
1346    /// Gets keywords metrics for a record type.
1347    pub fn get_keywords_metrics(&self, record_type: SuggestRecordType) -> Result<KeywordsMetrics> {
1348        let data = self.conn.try_query_row(
1349            r#"
1350            SELECT
1351                max(max_len) AS len,
1352                max(max_word_count) AS word_count
1353            FROM
1354                keywords_metrics
1355            WHERE
1356                record_type = :record_type
1357            "#,
1358            named_params! {
1359                ":record_type": record_type,
1360            },
1361            |row| -> Result<(usize, usize)> { Ok((row.get("len")?, row.get("word_count")?)) },
1362            true, // cache
1363        )?;
1364        Ok(data
1365            .map(|(max_len, max_word_count)| KeywordsMetrics {
1366                max_len,
1367                max_word_count,
1368            })
1369            .unwrap_or_default())
1370    }
1371}
1372
1373#[derive(Debug, PartialEq, Eq, Hash)]
1374pub struct IngestedRecord {
1375    pub id: SuggestRecordId,
1376    pub collection: String,
1377    pub record_type: String,
1378    pub last_modified: u64,
1379}
1380
1381impl IngestedRecord {
1382    fn from_row(row: &rusqlite::Row) -> Result<Self> {
1383        Ok(Self {
1384            id: SuggestRecordId::new(row.get("id")?),
1385            collection: row.get("collection")?,
1386            record_type: row.get("type")?,
1387            last_modified: row.get("last_modified")?,
1388        })
1389    }
1390}
1391
1392/// Helper struct to get full_keyword_ids for a suggestion
1393///
1394/// `FullKeywordInserter` handles repeated full keywords efficiently.  The first instance will
1395/// cause a row to be inserted into the database.  Subsequent instances will return the same
1396/// full_keyword_id.
1397struct FullKeywordInserter<'a> {
1398    conn: &'a Connection,
1399    suggestion_id: i64,
1400    last_inserted: Option<(&'a str, i64)>,
1401}
1402
1403impl<'a> FullKeywordInserter<'a> {
1404    fn new(conn: &'a Connection, suggestion_id: i64) -> Self {
1405        Self {
1406            conn,
1407            suggestion_id,
1408            last_inserted: None,
1409        }
1410    }
1411
1412    fn maybe_insert(&mut self, full_keyword: &'a str) -> rusqlite::Result<i64> {
1413        match self.last_inserted {
1414            Some((s, id)) if s == full_keyword => Ok(id),
1415            _ => {
1416                let full_keyword_id = self.conn.query_row_and_then(
1417                    "INSERT INTO full_keywords(
1418                        suggestion_id,
1419                        full_keyword
1420                     )
1421                     VALUES(
1422                        :suggestion_id,
1423                        :keyword
1424                     )
1425                     RETURNING id",
1426                    named_params! {
1427                        ":keyword": full_keyword,
1428                        ":suggestion_id": self.suggestion_id,
1429                    },
1430                    |row| row.get(0),
1431                )?;
1432                self.last_inserted = Some((full_keyword, full_keyword_id));
1433                Ok(full_keyword_id)
1434            }
1435        }
1436    }
1437}
1438
1439// ======================== Statement types ========================
1440//
1441// During ingestion we can insert hundreds of thousands of rows.  These types enable speedups by
1442// allowing us to prepare a statement outside a loop and use it many times inside the loop.
1443//
1444// Each type wraps [Connection::prepare] and [Statement] to provide a simplified interface,
1445// tailored to a specific query.
1446//
1447// This pattern is applicable for whenever we execute the same query repeatedly in a loop.
1448// The impact scales with the number of loop iterations, which is why we currently don't do this
1449// for providers like Mdn and Weather, which have relatively small number of records
1450// compared to Amp/Wikipedia.
1451
1452pub(crate) struct SuggestionInsertStatement<'conn>(rusqlite::Statement<'conn>);
1453
1454impl<'conn> SuggestionInsertStatement<'conn> {
1455    pub(crate) fn new(conn: &'conn Connection) -> Result<Self> {
1456        Ok(Self(conn.prepare(
1457            "INSERT INTO suggestions(
1458                 record_id,
1459                 title,
1460                 url,
1461                 score,
1462                 provider
1463             )
1464             VALUES(?, ?, ?, ?, ?)
1465             RETURNING id",
1466        )?))
1467    }
1468
1469    /// Execute the insert and return the `suggestion_id` for the new row
1470    pub(crate) fn execute(
1471        &mut self,
1472        record_id: &SuggestRecordId,
1473        title: &str,
1474        url: &str,
1475        score: f64,
1476        provider: SuggestionProvider,
1477    ) -> Result<i64> {
1478        self.0
1479            .query_row(
1480                (record_id.as_str(), title, url, score, provider as u8),
1481                |row| row.get(0),
1482            )
1483            .with_context("suggestion insert")
1484    }
1485}
1486
1487struct AmpInsertStatement<'conn>(rusqlite::Statement<'conn>);
1488
1489impl<'conn> AmpInsertStatement<'conn> {
1490    fn new(conn: &'conn Connection) -> Result<Self> {
1491        Ok(Self(conn.prepare(
1492            "INSERT INTO amp_custom_details(
1493                 suggestion_id,
1494                 advertiser,
1495                 block_id,
1496                 iab_category,
1497                 impression_url,
1498                 click_url,
1499                 icon_id
1500             )
1501             VALUES(?, ?, ?, ?, ?, ?, ?)
1502             ",
1503        )?))
1504    }
1505
1506    fn execute(&mut self, suggestion_id: i64, amp: &DownloadedAmpSuggestion) -> Result<()> {
1507        self.0
1508            .execute((
1509                suggestion_id,
1510                &amp.advertiser,
1511                amp.block_id,
1512                &amp.iab_category,
1513                &amp.impression_url,
1514                &amp.click_url,
1515                &amp.icon_id,
1516            ))
1517            .with_context("amp insert")?;
1518        Ok(())
1519    }
1520}
1521
1522struct WikipediaInsertStatement<'conn>(rusqlite::Statement<'conn>);
1523
1524impl<'conn> WikipediaInsertStatement<'conn> {
1525    fn new(conn: &'conn Connection) -> Result<Self> {
1526        Ok(Self(conn.prepare(
1527            "INSERT INTO wikipedia_custom_details(
1528                 suggestion_id,
1529                 icon_id
1530             )
1531             VALUES(?, ?)
1532             ",
1533        )?))
1534    }
1535
1536    fn execute(
1537        &mut self,
1538        suggestion_id: i64,
1539        wikipedia: &DownloadedWikipediaSuggestion,
1540    ) -> Result<()> {
1541        self.0
1542            .execute((suggestion_id, &wikipedia.icon_id))
1543            .with_context("wikipedia insert")?;
1544        Ok(())
1545    }
1546}
1547
1548struct AmoInsertStatement<'conn>(rusqlite::Statement<'conn>);
1549
1550impl<'conn> AmoInsertStatement<'conn> {
1551    fn new(conn: &'conn Connection) -> Result<Self> {
1552        Ok(Self(conn.prepare(
1553            "INSERT INTO amo_custom_details(
1554                 suggestion_id,
1555                 description,
1556                 guid,
1557                 icon_url,
1558                 rating,
1559                 number_of_ratings
1560             )
1561             VALUES(?, ?, ?, ?, ?, ?)
1562             ",
1563        )?))
1564    }
1565
1566    fn execute(&mut self, suggestion_id: i64, amo: &DownloadedAmoSuggestion) -> Result<()> {
1567        self.0
1568            .execute((
1569                suggestion_id,
1570                &amo.description,
1571                &amo.guid,
1572                &amo.icon_url,
1573                &amo.rating,
1574                amo.number_of_ratings,
1575            ))
1576            .with_context("amo insert")?;
1577        Ok(())
1578    }
1579}
1580
1581struct MdnInsertStatement<'conn>(rusqlite::Statement<'conn>);
1582
1583impl<'conn> MdnInsertStatement<'conn> {
1584    fn new(conn: &'conn Connection) -> Result<Self> {
1585        Ok(Self(conn.prepare(
1586            "INSERT INTO mdn_custom_details(
1587                 suggestion_id,
1588                 description
1589             )
1590             VALUES(?, ?)
1591             ",
1592        )?))
1593    }
1594
1595    fn execute(&mut self, suggestion_id: i64, mdn: &DownloadedMdnSuggestion) -> Result<()> {
1596        self.0
1597            .execute((suggestion_id, &mdn.description))
1598            .with_context("mdn insert")?;
1599        Ok(())
1600    }
1601}
1602
1603struct FakespotInsertStatement<'conn>(rusqlite::Statement<'conn>);
1604
1605impl<'conn> FakespotInsertStatement<'conn> {
1606    fn new(conn: &'conn Connection) -> Result<Self> {
1607        Ok(Self(conn.prepare(
1608            "INSERT INTO fakespot_custom_details(
1609                 suggestion_id,
1610                 fakespot_grade,
1611                 product_id,
1612                 keywords,
1613                 product_type,
1614                 rating,
1615                 total_reviews,
1616                 icon_id
1617             )
1618             VALUES(?, ?, ?, ?, ?, ?, ?, ?)
1619             ",
1620        )?))
1621    }
1622
1623    fn execute(
1624        &mut self,
1625        suggestion_id: i64,
1626        fakespot: &DownloadedFakespotSuggestion,
1627    ) -> Result<()> {
1628        let icon_id = fakespot
1629            .product_id
1630            .split_once('-')
1631            .map(|(vendor, _)| format!("fakespot-{vendor}"));
1632        self.0
1633            .execute((
1634                suggestion_id,
1635                &fakespot.fakespot_grade,
1636                &fakespot.product_id,
1637                &fakespot.keywords.to_lowercase(),
1638                &fakespot.product_type.to_lowercase(),
1639                fakespot.rating,
1640                fakespot.total_reviews,
1641                icon_id,
1642            ))
1643            .with_context("fakespot insert")?;
1644        Ok(())
1645    }
1646}
1647
1648struct DynamicInsertStatement<'conn>(rusqlite::Statement<'conn>);
1649
1650impl<'conn> DynamicInsertStatement<'conn> {
1651    fn new(conn: &'conn Connection) -> Result<Self> {
1652        Ok(Self(conn.prepare(
1653            "INSERT INTO dynamic_custom_details(
1654                 suggestion_id,
1655                 suggestion_type,
1656                 json_data
1657             )
1658             VALUES(?, ?, ?)
1659             ",
1660        )?))
1661    }
1662
1663    fn execute(
1664        &mut self,
1665        suggestion_id: i64,
1666        suggestion_type: &str,
1667        suggestion: &DownloadedDynamicSuggestion,
1668    ) -> Result<()> {
1669        self.0
1670            .execute((
1671                suggestion_id,
1672                suggestion_type,
1673                match &suggestion.data {
1674                    None => None,
1675                    Some(d) => Some(serde_json::to_string(&d)?),
1676                },
1677            ))
1678            .with_context("dynamic insert")?;
1679        Ok(())
1680    }
1681}
1682
1683pub(crate) struct KeywordInsertStatement<'conn>(rusqlite::Statement<'conn>);
1684
1685impl<'conn> KeywordInsertStatement<'conn> {
1686    pub(crate) fn new(conn: &'conn Connection) -> Result<Self> {
1687        Self::with_details(conn, "keywords", None)
1688    }
1689
1690    pub(crate) fn with_details(
1691        conn: &'conn Connection,
1692        table: &str,
1693        conflict_resolution: Option<InsertConflictResolution>,
1694    ) -> Result<Self> {
1695        Ok(Self(conn.prepare(&format!(
1696            r#"
1697            INSERT {} INTO {}(
1698                suggestion_id,
1699                keyword,
1700                full_keyword_id,
1701                rank
1702            )
1703            VALUES(?, ?, ?, ?)
1704            "#,
1705            conflict_resolution.as_ref().map(|r| r.as_str()).unwrap_or_default(),
1706            table,
1707        ))?))
1708    }
1709
1710    pub(crate) fn execute(
1711        &mut self,
1712        suggestion_id: i64,
1713        keyword: &str,
1714        full_keyword_id: Option<i64>,
1715        rank: usize,
1716    ) -> Result<()> {
1717        self.0
1718            .execute((suggestion_id, keyword, full_keyword_id, rank))
1719            .with_context("keyword insert")?;
1720        Ok(())
1721    }
1722}
1723
1724pub(crate) enum InsertConflictResolution {
1725    Ignore,
1726}
1727
1728impl InsertConflictResolution {
1729    fn as_str(&self) -> &str {
1730        match self {
1731            InsertConflictResolution::Ignore => "OR IGNORE",
1732        }
1733    }
1734}
1735
1736struct PrefixKeywordInsertStatement<'conn>(rusqlite::Statement<'conn>);
1737
1738impl<'conn> PrefixKeywordInsertStatement<'conn> {
1739    fn new(conn: &'conn Connection) -> Result<Self> {
1740        Ok(Self(conn.prepare(
1741            "INSERT INTO prefix_keywords(
1742                 suggestion_id,
1743                 confidence,
1744                 keyword_prefix,
1745                 keyword_suffix,
1746                 rank
1747             )
1748             VALUES(?, ?, ?, ?, ?)
1749             ",
1750        )?))
1751    }
1752
1753    fn execute(
1754        &mut self,
1755        suggestion_id: i64,
1756        confidence: Option<u8>,
1757        keyword_prefix: &str,
1758        keyword_suffix: &str,
1759        rank: usize,
1760    ) -> Result<()> {
1761        self.0
1762            .execute((
1763                suggestion_id,
1764                confidence.unwrap_or(0),
1765                keyword_prefix,
1766                keyword_suffix,
1767                rank,
1768            ))
1769            .with_context("prefix keyword insert")?;
1770        Ok(())
1771    }
1772}
1773
1774/// Metrics for a set of keywords. Use `KeywordsMetricsUpdater` to write metrics
1775/// to the store and `SuggestDao::get_keywords_metrics` to read them.
1776#[derive(Debug, Default, Eq, PartialEq)]
1777pub(crate) struct KeywordsMetrics {
1778    /// The max byte count (not chars) of all keywords in the set.
1779    pub(crate) max_len: usize,
1780    /// The max word count of all keywords in the set.
1781    pub(crate) max_word_count: usize,
1782}
1783
1784/// This can be used to keep a running tally of metrics as you process keywords
1785/// and then to write the metrics to the store. Make a `KeywordsMetricsUpdater`,
1786/// call `update` on it as you process each keyword, and then call `finish` to
1787/// write the metrics to the store.
1788pub(crate) struct KeywordsMetricsUpdater {
1789    metrics: KeywordsMetrics,
1790}
1791
1792impl KeywordsMetricsUpdater {
1793    pub(crate) fn new() -> Self {
1794        Self {
1795            metrics: KeywordsMetrics::default(),
1796        }
1797    }
1798
1799    /// Updates the metrics with those of the given keyword.
1800    ///
1801    /// The keyword's metrics are calculated as the max of its metrics as-is and
1802    /// its metrics when it's transformed using `i18n_transform`. For example,
1803    /// `"Qu\u{00e9}bec"` ("Québec" with single two-byte 'é' char) has 7 bytes,
1804    /// so its `len` is 7, but the `len` of `i18n_transform("Qu\u{00e9}bec")` is
1805    /// 6, since the 'é' is transformed to an ASCII 'e'. The keyword's `max_len`
1806    /// will therefore be the max of 7 and 6, which is 7.
1807    pub(crate) fn update(&mut self, keyword: &str) {
1808        let transformed_kw = i18n_transform(keyword);
1809        self.metrics.max_len = std::cmp::max(
1810            self.metrics.max_len,
1811            std::cmp::max(transformed_kw.len(), keyword.len()),
1812        );
1813
1814        // Getting the word count is costly, so we get only the transformed word
1815        // count under the assumption that it will always be at least as large
1816        // as the the original word count. That's currently true because the
1817        // only relevant transform that `i18n_transform` performs is to replace
1818        // hyphens with spaces.
1819        self.metrics.max_word_count = std::cmp::max(
1820            self.metrics.max_word_count,
1821            transformed_kw.split_whitespace().count(),
1822        );
1823    }
1824
1825    /// Inserts keywords metrics into the database. This assumes you have a
1826    /// cache object inside the `cache` cell that caches the metrics. It will be
1827    /// cleared since it will be invalidated by the metrics update.
1828    pub(crate) fn finish<T>(
1829        &self,
1830        conn: &Connection,
1831        record_id: &SuggestRecordId,
1832        record_type: SuggestRecordType,
1833        cache: &mut OnceCell<T>,
1834    ) -> Result<()> {
1835        let mut insert_stmt = conn.prepare(
1836            r#"
1837            INSERT OR REPLACE INTO keywords_metrics(
1838                record_id,
1839                record_type,
1840                max_len,
1841                max_word_count
1842            )
1843            VALUES(?, ?, ?, ?)
1844            "#,
1845        )?;
1846        insert_stmt
1847            .execute((
1848                record_id.as_str(),
1849                record_type,
1850                self.metrics.max_len,
1851                self.metrics.max_word_count,
1852            ))
1853            .with_context("keywords metrics insert")?;
1854
1855        // We just made some insertions that might invalidate the data in the
1856        // cache. Clear it so it's repopulated the next time it's accessed.
1857        cache.take();
1858
1859        Ok(())
1860    }
1861}
1862
1863pub(crate) struct AmpFtsInsertStatement<'conn>(rusqlite::Statement<'conn>);
1864
1865impl<'conn> AmpFtsInsertStatement<'conn> {
1866    pub(crate) fn new(conn: &'conn Connection) -> Result<Self> {
1867        Ok(Self(conn.prepare(
1868            "INSERT INTO amp_fts(rowid, full_keywords, title)
1869             VALUES(?, ?, ?)
1870             ",
1871        )?))
1872    }
1873
1874    pub(crate) fn execute(
1875        &mut self,
1876        suggestion_id: i64,
1877        full_keywords: &str,
1878        title: &str,
1879    ) -> Result<()> {
1880        self.0
1881            .execute((suggestion_id, full_keywords, title))
1882            .with_context("amp fts insert")?;
1883        Ok(())
1884    }
1885}
1886
1887fn provider_config_meta_key(provider: SuggestionProvider) -> String {
1888    format!("{}{}", PROVIDER_CONFIG_META_KEY_PREFIX, provider as u8)
1889}
1890
1891#[cfg(test)]
1892mod tests {
1893    use super::*;
1894    use crate::{store::tests::TestStore, testing::*, SuggestIngestionConstraints};
1895
1896    #[test]
1897    fn keywords_metrics_updater() -> anyhow::Result<()> {
1898        // (test keyword, expected `KeywordsMetrics`)
1899        //
1900        // Tests are cumulative!
1901        let tests = [
1902            (
1903                "abc",
1904                KeywordsMetrics {
1905                    max_len: 3,
1906                    max_word_count: 1,
1907                },
1908            ),
1909            (
1910                "a b",
1911                KeywordsMetrics {
1912                    max_len: 3,
1913                    max_word_count: 2,
1914                },
1915            ),
1916            (
1917                "a b c",
1918                KeywordsMetrics {
1919                    max_len: 5,
1920                    max_word_count: 3,
1921                },
1922            ),
1923            // "Québec" with single 'é' char:
1924            // With `i18n_transform`: len = 6
1925            // Without `i18n_transform`: len = 7
1926            // => Length for this string should be max(6, 7) = 7, which is a new
1927            //    overall max
1928            (
1929                "Qu\u{00e9}bec",
1930                KeywordsMetrics {
1931                    max_len: 7,
1932                    max_word_count: 3,
1933                },
1934            ),
1935            // "Québec" with ASCII 'e' followed by combining acute accent:
1936            // With `i18n_transform`, len = 6
1937            // Without `i18n_transform`, len = 8
1938            // => Length for this string should be max(6, 8) = 8, which is a new
1939            //    overall max
1940            (
1941                "Que\u{0301}bec",
1942                KeywordsMetrics {
1943                    max_len: 8,
1944                    max_word_count: 3,
1945                },
1946            ),
1947            // Each '-' should be replaced with a space, so the word count for
1948            // this string should be 4, which is a new overall max
1949            (
1950                "Carmel-by-the-Sea",
1951                KeywordsMetrics {
1952                    max_len: 17,
1953                    max_word_count: 4,
1954                },
1955            ),
1956        ];
1957
1958        // Make an updater and call it with each test in turn.
1959        let mut updater = KeywordsMetricsUpdater::new();
1960        for (test_kw, expected_metrics) in &tests {
1961            updater.update(test_kw);
1962            assert_eq!(&updater.metrics, expected_metrics);
1963        }
1964
1965        // Make a store and finish the updater so it updates the metrics table.
1966        let store = TestStore::new(MockRemoteSettingsClient::default());
1967        store.write(|dao| {
1968            // Make a dummy cache cell. It doesn't matter what it contains, we
1969            // just want to make sure it's empty after finishing the updater.
1970            let mut dummy_cache = OnceCell::new();
1971            dummy_cache.set("test").expect("dummy cache set");
1972            assert_ne!(dummy_cache.get(), None);
1973
1974            let record_type = SuggestRecordType::Wikipedia;
1975            updater.finish(
1976                dao.conn,
1977                &SuggestRecordId::new("test-record-1".to_string()),
1978                record_type,
1979                &mut dummy_cache,
1980            )?;
1981
1982            assert_eq!(dummy_cache.get(), None);
1983
1984            // Read back the metrics and make sure they match the ones in the
1985            // final test.
1986            let read_metrics_1 = dao.get_keywords_metrics(record_type)?;
1987            assert_eq!(read_metrics_1, tests.last().unwrap().1);
1988
1989            // Update the metrics one more time and finish them again.
1990            updater.update("a very long keyword with many words");
1991            let new_expected = KeywordsMetrics {
1992                max_len: 35,
1993                max_word_count: 7,
1994            };
1995            assert_eq!(updater.metrics, new_expected);
1996
1997            updater.finish(
1998                dao.conn,
1999                &SuggestRecordId::new("test-record-2".to_string()),
2000                record_type,
2001                &mut dummy_cache,
2002            )?;
2003
2004            // Read back the new metrics and make sure they match.
2005            let read_metrics_2 = dao.get_keywords_metrics(record_type)?;
2006            assert_eq!(read_metrics_2, new_expected);
2007
2008            Ok(())
2009        })?;
2010
2011        Ok(())
2012    }
2013
2014    // Keywords in `keywords_i18n` should be dropped when their records are
2015    // deleted.
2016    #[test]
2017    fn keywords_i18n_delete_record() -> anyhow::Result<()> {
2018        // Add some records whose keywords are stored in `keywords_i18n`. We'll
2019        // use weather records.
2020        let kws_1 = ["aaa", "bbb", "ccc"];
2021        let kws_2 = ["yyy", "zzz"];
2022        let mut store = TestStore::new(
2023            MockRemoteSettingsClient::default()
2024                .with_record(SuggestionProvider::Weather.record(
2025                    "weather-1",
2026                    json!({
2027                        "score": 0.24,
2028                        "keywords": kws_1,
2029                    }),
2030                ))
2031                .with_record(SuggestionProvider::Weather.record(
2032                    "weather-2",
2033                    json!({
2034                        "score": 0.24,
2035                        "keywords": kws_2,
2036                    }),
2037                )),
2038        );
2039        store.ingest(SuggestIngestionConstraints {
2040            providers: Some(vec![SuggestionProvider::Weather]),
2041            ..SuggestIngestionConstraints::all_providers()
2042        });
2043
2044        // Make sure all keywords are present.
2045        assert_eq!(
2046            store.count_rows("keywords_i18n") as usize,
2047            kws_1.len() + kws_2.len()
2048        );
2049
2050        for q in kws_1.iter().chain(kws_2.iter()) {
2051            assert_eq!(
2052                store.fetch_suggestions(SuggestionQuery::weather(q)),
2053                vec![Suggestion::Weather {
2054                    score: 0.24,
2055                    city: None,
2056                }],
2057                "query: {:?}",
2058                q
2059            );
2060        }
2061
2062        // Delete the first record.
2063        store
2064            .client_mut()
2065            .delete_record(SuggestionProvider::Weather.empty_record("weather-1"));
2066        store.ingest(SuggestIngestionConstraints {
2067            providers: Some(vec![SuggestionProvider::Weather]),
2068            ..SuggestIngestionConstraints::all_providers()
2069        });
2070
2071        // Its keywords should be dropped and the keywords from the second
2072        // record should still be present.
2073        assert_eq!(store.count_rows("keywords_i18n") as usize, kws_2.len());
2074
2075        for q in kws_1 {
2076            assert_eq!(
2077                store.fetch_suggestions(SuggestionQuery::weather(q)),
2078                vec![],
2079                "query: {:?}",
2080                q
2081            );
2082        }
2083        for q in kws_2 {
2084            assert_eq!(
2085                store.fetch_suggestions(SuggestionQuery::weather(q)),
2086                vec![Suggestion::Weather {
2087                    score: 0.24,
2088                    city: None,
2089                }],
2090                "query: {:?}",
2091                q
2092            );
2093        }
2094
2095        Ok(())
2096    }
2097
2098    // Keywords in `keywords_i18n` should be dropped when their records are
2099    // updated and the keywords are no longer present, and new keywords should
2100    // be added.
2101    #[test]
2102    fn keywords_i18n_update_record() -> anyhow::Result<()> {
2103        // Add some records whose keywords are stored in `keywords_i18n`. We'll
2104        // use weather records.
2105        let kws_1 = ["aaa", "bbb", "ccc"];
2106        let kws_2 = ["yyy", "zzz"];
2107        let mut store = TestStore::new(
2108            MockRemoteSettingsClient::default()
2109                .with_record(SuggestionProvider::Weather.record(
2110                    "weather-1",
2111                    json!({
2112                        "score": 0.24,
2113                        "keywords": kws_1,
2114                    }),
2115                ))
2116                .with_record(SuggestionProvider::Weather.record(
2117                    "weather-2",
2118                    json!({
2119                        "score": 0.24,
2120                        "keywords": kws_2,
2121                    }),
2122                )),
2123        );
2124        store.ingest(SuggestIngestionConstraints {
2125            providers: Some(vec![SuggestionProvider::Weather]),
2126            ..SuggestIngestionConstraints::all_providers()
2127        });
2128
2129        // Make sure all keywords are present.
2130        assert_eq!(
2131            store.count_rows("keywords_i18n") as usize,
2132            kws_1.len() + kws_2.len()
2133        );
2134
2135        for q in kws_1.iter().chain(kws_2.iter()) {
2136            assert_eq!(
2137                store.fetch_suggestions(SuggestionQuery::weather(q)),
2138                vec![Suggestion::Weather {
2139                    score: 0.24,
2140                    city: None,
2141                }],
2142                "query: {:?}",
2143                q
2144            );
2145        }
2146
2147        // Update the first record.
2148        let kws_1_new = [
2149            "bbb", // keyword from the old record
2150            "mmm", // new keyword
2151        ];
2152        store
2153            .client_mut()
2154            .update_record(SuggestionProvider::Weather.record(
2155                "weather-1",
2156                json!({
2157                    "score": 0.24,
2158                    "keywords": kws_1_new,
2159                }),
2160            ));
2161        store.ingest(SuggestIngestionConstraints {
2162            providers: Some(vec![SuggestionProvider::Weather]),
2163            ..SuggestIngestionConstraints::all_providers()
2164        });
2165
2166        // Check all keywords.
2167        assert_eq!(
2168            store.count_rows("keywords_i18n") as usize,
2169            kws_1_new.len() + kws_2.len()
2170        );
2171
2172        for q in ["aaa", "ccc"] {
2173            assert_eq!(
2174                store.fetch_suggestions(SuggestionQuery::weather(q)),
2175                vec![],
2176                "query: {:?}",
2177                q
2178            );
2179        }
2180        for q in kws_1_new.iter().chain(kws_2.iter()) {
2181            assert_eq!(
2182                store.fetch_suggestions(SuggestionQuery::weather(q)),
2183                vec![Suggestion::Weather {
2184                    score: 0.24,
2185                    city: None,
2186                }],
2187                "query: {:?}",
2188                q
2189            );
2190        }
2191
2192        Ok(())
2193    }
2194}