examples_suggest_cli/
main.rs

Help
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
5use std::{collections::HashMap, sync::Arc};
6
7use anyhow::Result;
8use clap::{Parser, Subcommand, ValueEnum};
9
10use remote_settings::{RemoteSettingsConfig2, RemoteSettingsServer, RemoteSettingsService};
11use suggest::{
12    AmpMatchingStrategy, SuggestIngestionConstraints, SuggestStore, SuggestStoreBuilder,
13    SuggestionProvider, SuggestionProviderConstraints, SuggestionQuery,
14};
15
16static DB_FILENAME: &str = "suggest.db";
17
18const DEFAULT_LOG_FILTER: &str = "suggest::store=info";
19const DEFAULT_LOG_FILTER_VERBOSE: &str = "suggest::store=trace";
20
21#[derive(Debug, Parser)]
22#[command(about, long_about = None)]
23struct Cli {
24    #[arg(short = 's')]
25    remote_settings_server: Option<RemoteSettingsServerArg>,
26    #[arg(short = 'b')]
27    remote_settings_bucket: Option<String>,
28    #[arg(long, short, action)]
29    verbose: bool,
30    // Custom { url: String },
31    #[command(subcommand)]
32    command: Commands,
33}
34
35#[derive(Clone, Debug, ValueEnum)]
36enum RemoteSettingsServerArg {
37    Prod,
38    Stage,
39    Dev,
40}
41
42#[derive(Debug, Subcommand)]
43enum Commands {
44    /// Ingest data
45    Ingest {
46        #[clap(long, short, action)]
47        reingest: bool,
48        #[clap(long, short)]
49        providers: Vec<SuggestionProviderArg>,
50    },
51    /// Query against ingested data
52    Query {
53        #[arg(long, action)]
54        fts_match_info: bool,
55        #[clap(long, short)]
56        provider: Option<SuggestionProviderArg>,
57        /// Input to search
58        input: String,
59        #[clap(long, short)]
60        amp_matching_strategy: Option<AmpMatchingStrategyArg>,
61    },
62}
63
64#[derive(Clone, Debug, ValueEnum)]
65enum AmpMatchingStrategyArg {
66    /// Use keyword matching, without keyword expansion
67    NoKeyword = 1, // Use `1` as the starting discriminant, since the JS code assumes this.
68    /// Use FTS matching
69    Fts,
70    /// Use FTS matching against the title
71    FtsTitle,
72}
73
74impl From<AmpMatchingStrategyArg> for AmpMatchingStrategy {
75    fn from(val: AmpMatchingStrategyArg) -> Self {
76        match val {
77            AmpMatchingStrategyArg::NoKeyword => AmpMatchingStrategy::NoKeywordExpansion,
78            AmpMatchingStrategyArg::Fts => AmpMatchingStrategy::FtsAgainstFullKeywords,
79            AmpMatchingStrategyArg::FtsTitle => AmpMatchingStrategy::FtsAgainstTitle,
80        }
81    }
82}
83
84#[derive(Clone, Debug, ValueEnum)]
85enum SuggestionProviderArg {
86    Amp,
87    Wikipedia,
88    Amo,
89    Yelp,
90    Mdn,
91    Weather,
92    Fakespot,
93}
94
95impl From<SuggestionProviderArg> for SuggestionProvider {
96    fn from(value: SuggestionProviderArg) -> Self {
97        match value {
98            SuggestionProviderArg::Amp => Self::Amp,
99            SuggestionProviderArg::Wikipedia => Self::Wikipedia,
100            SuggestionProviderArg::Amo => Self::Amo,
101            SuggestionProviderArg::Yelp => Self::Yelp,
102            SuggestionProviderArg::Mdn => Self::Mdn,
103            SuggestionProviderArg::Weather => Self::Weather,
104            SuggestionProviderArg::Fakespot => Self::Fakespot,
105        }
106    }
107}
108
109fn main() -> Result<()> {
110    let cli = Cli::parse();
111    env_logger::init_from_env(env_logger::Env::default().filter_or(
112        "RUST_LOG",
113        if cli.verbose {
114            DEFAULT_LOG_FILTER_VERBOSE
115        } else {
116            DEFAULT_LOG_FILTER
117        },
118    ));
119    nss::ensure_initialized();
120    viaduct_reqwest::use_reqwest_backend();
121    let store = build_store(&cli)?;
122    match cli.command {
123        Commands::Ingest {
124            reingest,
125            providers,
126        } => ingest(&store, reingest, providers, cli.verbose),
127        Commands::Query {
128            provider,
129            input,
130            fts_match_info,
131            amp_matching_strategy,
132        } => query(
133            &store,
134            provider,
135            input,
136            fts_match_info,
137            amp_matching_strategy,
138            cli.verbose,
139        ),
140    };
141    Ok(())
142}
143
144fn build_store(cli: &Cli) -> Result<Arc<SuggestStore>> {
145    Ok(Arc::new(SuggestStoreBuilder::default())
146        .data_path(cli_support::cli_data_path(DB_FILENAME))
147        .remote_settings_service(build_remote_settings_service(cli))
148        .build()?)
149}
150
151fn build_remote_settings_service(cli: &Cli) -> Arc<RemoteSettingsService> {
152    let config = RemoteSettingsConfig2 {
153        server: cli.remote_settings_server.as_ref().map(|s| match s {
154            RemoteSettingsServerArg::Dev => RemoteSettingsServer::Dev,
155            RemoteSettingsServerArg::Stage => RemoteSettingsServer::Stage,
156            RemoteSettingsServerArg::Prod => RemoteSettingsServer::Prod,
157        }),
158        bucket_name: cli.remote_settings_bucket.clone(),
159        app_context: None,
160    };
161    let storage_dir = cli_support::cli_data_subdir("remote-settings-data");
162    Arc::new(RemoteSettingsService::new(storage_dir, config))
163}
164
165fn ingest(
166    store: &SuggestStore,
167    reingest: bool,
168    providers: Vec<SuggestionProviderArg>,
169    verbose: bool,
170) {
171    if reingest {
172        print_header("Reingesting data");
173        store.force_reingest();
174    } else {
175        print_header("Ingesting data");
176    }
177    let constraints = if providers.is_empty() {
178        SuggestIngestionConstraints::all_providers()
179    } else {
180        SuggestIngestionConstraints {
181            providers: Some(providers.into_iter().map(Into::into).collect()),
182            ..SuggestIngestionConstraints::default()
183        }
184    };
185
186    let metrics = store
187        .ingest(constraints)
188        .unwrap_or_else(|e| panic!("Error in ingest: {e}"));
189    if verbose && !metrics.ingestion_times.is_empty() {
190        print_header("Ingestion times");
191        let mut ingestion_times = metrics.ingestion_times;
192        let download_times: HashMap<String, u64> = metrics
193            .download_times
194            .into_iter()
195            .map(|s| (s.label, s.value))
196            .collect();
197
198        ingestion_times.sort_by_key(|s| s.value);
199        ingestion_times.reverse();
200        for sample in ingestion_times {
201            let label = &sample.label;
202            let ingestion_time = sample.value / 1000;
203            let download_time = download_times.get(label).unwrap_or(&0) / 1000;
204
205            println!(
206                "{label:30} Download: {download_time:>5}ms    Ingestion: {ingestion_time:>5}ms"
207            );
208        }
209    }
210    print_header("Done");
211}
212
213fn query(
214    store: &SuggestStore,
215    provider: Option<SuggestionProviderArg>,
216    input: String,
217    fts_match_info: bool,
218    amp_matching_strategy: Option<AmpMatchingStrategyArg>,
219    verbose: bool,
220) {
221    let query = SuggestionQuery {
222        providers: match provider {
223            Some(provider) => vec![provider.into()],
224            None => SuggestionProvider::all().to_vec(),
225        },
226        keyword: input,
227        provider_constraints: Some(SuggestionProviderConstraints {
228            amp_alternative_matching: amp_matching_strategy.map(Into::into),
229            ..SuggestionProviderConstraints::default()
230        }),
231        ..SuggestionQuery::default()
232    };
233    let mut results = store
234        .query_with_metrics(query)
235        .unwrap_or_else(|e| panic!("Error querying store: {e}"));
236    if results.suggestions.is_empty() {
237        print_header("No Results");
238    } else {
239        print_header("Results");
240        let count = results.suggestions.len();
241        for suggestion in results.suggestions {
242            let title = suggestion.title();
243            let url = suggestion.url().unwrap_or("[no-url]");
244            let icon = if suggestion.icon_data().is_some() {
245                "with icon"
246            } else {
247                "no icon"
248            };
249            println!("* {title} ({url}) ({icon})");
250            if fts_match_info {
251                if let Some(match_info) = suggestion.fts_match_info() {
252                    println!("   {match_info:?}")
253                } else {
254                    println!("   <no match info>");
255                }
256                println!("+ {} other sugestions", count - 1);
257                break;
258            }
259        }
260    }
261    if verbose {
262        print_header("Query times");
263        results.query_times.sort_by_key(|s| s.value);
264        results.query_times.reverse();
265        for s in results.query_times {
266            println!("{:33} Time: {:>5}us", s.label, s.value);
267        }
268    }
269}
270
271fn print_header(msg: impl Into<String>) {
272    let mut msg = msg.into();
273    if msg.len() % 2 == 1 {
274        msg.push(' ');
275    }
276    let width = (70 - msg.len() - 2) / 2;
277    println!();
278    println!("{} {msg} {}", "=".repeat(width), "=".repeat(width));
279}
examples_suggest_cli/main.rs

examples_suggest_cli/
main.rs
