1use 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
36pub const GLOBAL_CONFIG_META_KEY: &str = "global_config";
39pub const PROVIDER_CONFIG_META_KEY_PREFIX: &str = "provider_config_";
43
44pub const DEFAULT_SUGGESTION_SCORE: f64 = 0.2;
46
47#[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
60pub(crate) struct SuggestDb {
63 pub conn: Mutex<Connection>,
64
65 pub interrupt_handle: Arc<SqlInterruptHandle>,
71}
72
73impl SuggestDb {
74 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 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 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 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 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 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
163pub(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 pub fn delete_record_data(&mut self, record_id: &SuggestRecordId) -> Result<()> {
191 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 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 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 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 _ => unreachable!(),
493 };
494
495 let prefix = if fts_query.is_prefix_query {
496 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 false
504 };
505
506 Ok(FtsMatchInfo {
507 prefix,
508 stemming: fts_query.match_required_stemming(&fts_content),
509 })
510 }
511
512 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 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 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 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 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 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 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 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 false
834 };
835
836 Ok(FtsMatchInfo {
837 prefix,
838 stemming: fts_query.match_required_stemming(title),
839 })
840 }
841
842 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 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 pub fn insert_amp_suggestions(
968 &mut self,
969 record_id: &SuggestRecordId,
970 suggestions: &[DownloadedAmpSuggestion],
971 enable_fts: bool,
972 ) -> Result<()> {
973 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 pub fn insert_wikipedia_suggestions(
1016 &mut self,
1017 record_id: &SuggestRecordId,
1018 suggestions: &[DownloadedWikipediaSuggestion],
1019 ) -> Result<()> {
1020 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 keyword_insert.execute(suggestion_id, keyword.keyword, None, keyword.rank)?;
1038 }
1039 }
1040 Ok(())
1041 }
1042
1043 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 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 pub fn insert_dynamic_suggestions(
1100 &mut self,
1101 record_id: &SuggestRecordId,
1102 record: &DownloadedDynamicRecord,
1103 suggestions: &[DownloadedDynamicSuggestion],
1104 ) -> Result<()> {
1105 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 "",
1121 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 for (rank, keyword) in suggestion.keywords().enumerate() {
1132 keyword_insert.execute(suggestion_id, &keyword, None, rank)?;
1133 }
1134 }
1135 Ok(())
1136 }
1137
1138 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 pub fn drop_suggestions(&mut self, record_id: &SuggestRecordId) -> Result<()> {
1194 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 self.weather_cache.take();
1272 self.geoname_cache.take();
1273
1274 Ok(())
1275 }
1276
1277 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 pub fn clear(&mut self) -> Result<()> {
1288 Ok(clear_database(self.conn)?)
1289 }
1290
1291 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 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 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 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 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 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 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, )?;
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
1392struct 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
1439pub(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 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 &.advertiser,
1511 amp.block_id,
1512 &.iab_category,
1513 &.impression_url,
1514 &.click_url,
1515 &.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#[derive(Debug, Default, Eq, PartialEq)]
1777pub(crate) struct KeywordsMetrics {
1778 pub(crate) max_len: usize,
1780 pub(crate) max_word_count: usize,
1782}
1783
1784pub(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 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 self.metrics.max_word_count = std::cmp::max(
1820 self.metrics.max_word_count,
1821 transformed_kw.split_whitespace().count(),
1822 );
1823 }
1824
1825 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 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 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 (
1929 "Qu\u{00e9}bec",
1930 KeywordsMetrics {
1931 max_len: 7,
1932 max_word_count: 3,
1933 },
1934 ),
1935 (
1941 "Que\u{0301}bec",
1942 KeywordsMetrics {
1943 max_len: 8,
1944 max_word_count: 3,
1945 },
1946 ),
1947 (
1950 "Carmel-by-the-Sea",
1951 KeywordsMetrics {
1952 max_len: 17,
1953 max_word_count: 4,
1954 },
1955 ),
1956 ];
1957
1958 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 let store = TestStore::new(MockRemoteSettingsClient::default());
1967 store.write(|dao| {
1968 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 let read_metrics_1 = dao.get_keywords_metrics(record_type)?;
1987 assert_eq!(read_metrics_1, tests.last().unwrap().1);
1988
1989 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 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 #[test]
2017 fn keywords_i18n_delete_record() -> anyhow::Result<()> {
2018 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 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 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 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 #[test]
2102 fn keywords_i18n_update_record() -> anyhow::Result<()> {
2103 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 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 let kws_1_new = [
2149 "bbb", "mmm", ];
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 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}