examples_remote_settings_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 anyhow::Result;
6use clap::{Parser, Subcommand, ValueEnum};
7use std::path::PathBuf;
8
9use dump::client::CollectionDownloader;
10use remote_settings::{RemoteSettingsConfig2, RemoteSettingsServer, RemoteSettingsService};
11
12const DEFAULT_LOG_FILTER: &str = "remote_settings=info";
13const DEFAULT_LOG_FILTER_VERBOSE: &str = "remote_settings=trace";
14
15#[derive(Debug, Parser)]
16#[command(about, long_about = None)]
17struct Cli {
18    #[arg(short = 's')]
19    server: Option<RemoteSettingsServerArg>,
20    #[arg(short = 'b')]
21    bucket: Option<String>,
22    #[arg(short = 'd')]
23    storage_dir: Option<String>,
24    #[arg(long, short, action)]
25    verbose: bool,
26    #[command(subcommand)]
27    command: Commands,
28}
29
30#[derive(Clone, Debug, ValueEnum)]
31enum RemoteSettingsServerArg {
32    Prod,
33    Stage,
34    Dev,
35}
36
37#[derive(Debug, Subcommand)]
38enum Commands {
39    /// Sync collections
40    Sync {
41        #[clap(required = true)]
42        collections: Vec<String>,
43    },
44    /// Query against ingested data
45    Get {
46        collection: String,
47        #[arg(long)]
48        sync_if_empty: bool,
49    },
50    /// Download and combine all remote settings collections
51    DumpSync {
52        /// Root path of the repository
53        #[arg(short, long, default_value = ".")]
54        path: PathBuf,
55
56        /// Dry run - don't write any files
57        #[arg(long, default_value_t = false)]
58        dry_run: bool,
59    },
60    /// Download a single collection to the dumps directory
61    DumpGet {
62        /// Bucket name
63        #[arg(long, required = true)]
64        bucket: String,
65
66        /// Collection name
67        #[arg(long, required = true)]
68        collection_name: String,
69
70        /// Root path of the repository
71        #[arg(short, long, default_value = ".")]
72        path: PathBuf,
73    },
74}
75
76fn main() -> Result<()> {
77    let cli = Cli::parse();
78    env_logger::init_from_env(env_logger::Env::default().filter_or(
79        "RUST_LOG",
80        if cli.verbose {
81            DEFAULT_LOG_FILTER_VERBOSE
82        } else {
83            DEFAULT_LOG_FILTER
84        },
85    ));
86    nss::ensure_initialized();
87    viaduct_reqwest::use_reqwest_backend();
88    let service = build_service(&cli)?;
89    match cli.command {
90        Commands::Sync { collections } => sync(service, collections),
91        Commands::Get {
92            collection,
93            sync_if_empty,
94        } => {
95            get_records(service, collection, sync_if_empty);
96            Ok(())
97        }
98        Commands::DumpSync { path, dry_run } => {
99            let downloader = CollectionDownloader::new(path);
100            let runtime = tokio::runtime::Runtime::new()?;
101            runtime.block_on(downloader.run(dry_run))
102        }
103        Commands::DumpGet {
104            bucket,
105            collection_name,
106            path,
107        } => {
108            let downloader = CollectionDownloader::new(path);
109            let runtime = tokio::runtime::Runtime::new()?;
110            runtime.block_on(downloader.download_single(&bucket, &collection_name))
111        }
112    }
113}
114
115fn build_service(cli: &Cli) -> Result<RemoteSettingsService> {
116    let config = RemoteSettingsConfig2 {
117        server: cli.server.as_ref().map(|s| match s {
118            RemoteSettingsServerArg::Dev => RemoteSettingsServer::Dev,
119            RemoteSettingsServerArg::Stage => RemoteSettingsServer::Stage,
120            RemoteSettingsServerArg::Prod => RemoteSettingsServer::Prod,
121        }),
122        bucket_name: cli.bucket.clone(),
123        app_context: None,
124    };
125    cli_support::ensure_cli_data_dir_exists();
126    let storage_dir = cli
127        .storage_dir
128        .clone()
129        .unwrap_or_else(|| cli_support::cli_data_subdir("remote-settings-data"));
130    Ok(RemoteSettingsService::new(storage_dir, config))
131}
132
133fn sync(service: RemoteSettingsService, collections: Vec<String>) -> Result<()> {
134    // Create a bunch of clients so that sync() syncs their collections
135    let _clients = collections
136        .into_iter()
137        .map(|collection| service.make_client(collection))
138        .collect::<Vec<_>>();
139    service.sync()?;
140    Ok(())
141}
142
143fn get_records(service: RemoteSettingsService, collection: String, sync_if_empty: bool) {
144    let client = service.make_client(collection);
145    match client.get_records(sync_if_empty) {
146        Some(records) => {
147            for record in records {
148                println!("{record:?}");
149            }
150        }
151        None => println!("No cached records"),
152    }
153}
examples_remote_settings_cli/main.rs

examples_remote_settings_cli/
main.rs
