Skip to content

Commit

Permalink
refactor: add the main cli entry point struct Command{} and simplif…
Browse files Browse the repository at this point in the history 
…y `main()` (#3850)

refactor: add the main cli entrypoint struct `Command{}` to simplify `main()`
  • Loading branch information
@zyy17
zyy17 authored May 7, 2024
1 parent 71c1c7c commit a68072c
Show file tree
Hide file tree
Showing 9 changed files with 130 additions and 125 deletions.
85 changes: 49 additions & 36 deletions src/cmd/src/bin/greptime.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,24 +16,42 @@

use std::fmt;

use clap::{FromArgMatches, Parser, Subcommand};
use clap::{Parser, Subcommand};
use cmd::error::Result;
use cmd::options::{CliOptions, Options};
use cmd::{
cli, datanode, frontend, greptimedb_cli, log_versions, metasrv, standalone, start_app, App,
};
use cmd::options::{GlobalOptions, Options};
use cmd::{cli, datanode, frontend, log_versions, metasrv, standalone, start_app, App};
use common_version::{short_version, version};

#[derive(Parser)]
#[command(name = "greptime", author, version, long_version = version!(), about)]
#[command(propagate_version = true)]
pub(crate) struct Command {
#[clap(subcommand)]
pub(crate) subcmd: SubCommand,

#[clap(flatten)]
pub(crate) global_options: GlobalOptions,
}

#[derive(Subcommand)]
enum SubCommand {
/// Start datanode service.
#[clap(name = "datanode")]
Datanode(datanode::Command),

/// Start frontend service.
#[clap(name = "frontend")]
Frontend(frontend::Command),

/// Start metasrv service.
#[clap(name = "metasrv")]
Metasrv(metasrv::Command),

/// Run greptimedb as a standalone service.
#[clap(name = "standalone")]
Standalone(standalone::Command),

/// Execute the cli tools for greptimedb.
#[clap(name = "cli")]
Cli(cli::Command),
}
Expand Down Expand Up @@ -67,13 +85,13 @@ impl SubCommand {
Ok(app)
}

fn load_options(&self, cli_options: &CliOptions) -> Result<Options> {
fn load_options(&self, global_options: &GlobalOptions) -> Result<Options> {
match self {
SubCommand::Datanode(cmd) => cmd.load_options(cli_options),
SubCommand::Frontend(cmd) => cmd.load_options(cli_options),
SubCommand::Metasrv(cmd) => cmd.load_options(cli_options),
SubCommand::Standalone(cmd) => cmd.load_options(cli_options),
SubCommand::Cli(cmd) => cmd.load_options(cli_options),
SubCommand::Datanode(cmd) => cmd.load_options(global_options),
SubCommand::Frontend(cmd) => cmd.load_options(global_options),
SubCommand::Metasrv(cmd) => cmd.load_options(global_options),
SubCommand::Standalone(cmd) => cmd.load_options(global_options),
SubCommand::Cli(cmd) => cmd.load_options(global_options),
}
}
}
Expand All @@ -96,44 +114,39 @@ static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

#[tokio::main]
async fn main() -> Result<()> {
let metadata = human_panic::Metadata {
version: env!("CARGO_PKG_VERSION").into(),
name: "GreptimeDB".into(),
authors: Default::default(),
homepage: "https://github.com/GreptimeTeam/greptimedb/discussions".into(),
};
human_panic::setup_panic!(metadata);

common_telemetry::set_panic_hook();

let version = version!();
let cli = greptimedb_cli().version(version);

let cli = SubCommand::augment_subcommands(cli);

let args = cli.get_matches();
setup_human_panic();
start(Command::parse()).await
}

let subcmd = match SubCommand::from_arg_matches(&args) {
Ok(subcmd) => subcmd,
Err(e) => e.exit(),
};
async fn start(cli: Command) -> Result<()> {
let subcmd = cli.subcmd;

let app_name = subcmd.to_string();

let cli_options = CliOptions::new(&args);

let opts = subcmd.load_options(&cli_options)?;
let opts = subcmd.load_options(&cli.global_options)?;

let _guard = common_telemetry::init_global_logging(
&app_name,
opts.logging_options(),
cli_options.tracing_options(),
cli.global_options.tracing_options(),
opts.node_id(),
);

log_versions(version, short_version!());
log_versions(version!(), short_version!());

let app = subcmd.build(opts).await?;

start_app(app).await
}

fn setup_human_panic() {
let metadata = human_panic::Metadata {
version: env!("CARGO_PKG_VERSION").into(),
name: "GreptimeDB".into(),
authors: Default::default(),
homepage: "https://github.com/GreptimeTeam/greptimedb/discussions".into(),
};
human_panic::setup_panic!(metadata);

common_telemetry::set_panic_hook();
}
8 changes: 4 additions & 4 deletions src/cmd/src/cli.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ use upgrade::UpgradeCommand;

use self::export::ExportCommand;
use crate::error::Result;
use crate::options::{CliOptions, Options};
use crate::options::{GlobalOptions, Options};
use crate::App;

#[async_trait]
Expand Down Expand Up @@ -80,14 +80,14 @@ impl Command {
self.cmd.build().await
}

pub fn load_options(&self, cli_options: &CliOptions) -> Result<Options> {
pub fn load_options(&self, global_options: &GlobalOptions) -> Result<Options> {
let mut logging_opts = LoggingOptions::default();

if let Some(dir) = &cli_options.log_dir {
if let Some(dir) = &global_options.log_dir {
logging_opts.dir.clone_from(dir);
}

logging_opts.level.clone_from(&cli_options.log_level);
logging_opts.level.clone_from(&global_options.log_level);

Ok(Options::Cli(Box::new(logging_opts)))
}
Expand Down
36 changes: 19 additions & 17 deletions src/cmd/src/datanode.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ use servers::Mode;
use snafu::{OptionExt, ResultExt};

use crate::error::{MissingConfigSnafu, Result, ShutdownDatanodeSnafu, StartDatanodeSnafu};
use crate::options::{CliOptions, Options};
use crate::options::{GlobalOptions, Options};
use crate::App;

pub struct Instance {
Expand Down Expand Up @@ -82,8 +82,8 @@ impl Command {
self.subcmd.build(opts).await
}

pub fn load_options(&self, cli_options: &CliOptions) -> Result<Options> {
self.subcmd.load_options(cli_options)
pub fn load_options(&self, global_options: &GlobalOptions) -> Result<Options> {
self.subcmd.load_options(global_options)
}
}

Expand All @@ -99,9 +99,9 @@ impl SubCommand {
}
}

fn load_options(&self, cli_options: &CliOptions) -> Result<Options> {
fn load_options(&self, global_options: &GlobalOptions) -> Result<Options> {
match self {
SubCommand::Start(cmd) => cmd.load_options(cli_options),
SubCommand::Start(cmd) => cmd.load_options(global_options),
}
}
}
Expand Down Expand Up @@ -131,19 +131,19 @@ struct StartCommand {
}

impl StartCommand {
fn load_options(&self, cli_options: &CliOptions) -> Result<Options> {
fn load_options(&self, global_options: &GlobalOptions) -> Result<Options> {
let mut opts: DatanodeOptions = Options::load_layered_options(
self.config_file.as_deref(),
self.env_prefix.as_ref(),
DatanodeOptions::env_list_keys(),
)?;

if let Some(dir) = &cli_options.log_dir {
if let Some(dir) = &global_options.log_dir {
opts.logging.dir.clone_from(dir);
}

if cli_options.log_level.is_some() {
opts.logging.level.clone_from(&cli_options.log_level);
if global_options.log_level.is_some() {
opts.logging.level.clone_from(&global_options.log_level);
}

if let Some(addr) = &self.rpc_addr {
Expand Down Expand Up @@ -259,7 +259,7 @@ mod tests {
use servers::Mode;

use super::*;
use crate::options::{CliOptions, ENV_VAR_SEP};
use crate::options::{GlobalOptions, ENV_VAR_SEP};

#[test]
fn test_read_from_config_file() {
Expand Down Expand Up @@ -315,7 +315,8 @@ mod tests {
..Default::default()
};

let Options::Datanode(options) = cmd.load_options(&CliOptions::default()).unwrap() else {
let Options::Datanode(options) = cmd.load_options(&GlobalOptions::default()).unwrap()
else {
unreachable!()
};

Expand Down Expand Up @@ -377,7 +378,7 @@ mod tests {
#[test]
fn test_try_from_cmd() {
if let Options::Datanode(opt) = StartCommand::default()
.load_options(&CliOptions::default())
.load_options(&GlobalOptions::default())
.unwrap()
{
assert_eq!(Mode::Standalone, opt.mode)
Expand All @@ -388,7 +389,7 @@ mod tests {
metasrv_addr: Some(vec!["127.0.0.1:3002".to_string()]),
..Default::default()
})
.load_options(&CliOptions::default())
.load_options(&GlobalOptions::default())
.unwrap()
{
assert_eq!(Mode::Distributed, opt.mode)
Expand All @@ -398,15 +399,15 @@ mod tests {
metasrv_addr: Some(vec!["127.0.0.1:3002".to_string()]),
..Default::default()
})
.load_options(&CliOptions::default())
.load_options(&GlobalOptions::default())
.is_err());

// Providing node_id but leave metasrv_addr absent is ok since metasrv_addr has default value
assert!((StartCommand {
node_id: Some(42),
..Default::default()
})
.load_options(&CliOptions::default())
.load_options(&GlobalOptions::default())
.is_ok());
}

Expand All @@ -415,7 +416,7 @@ mod tests {
let cmd = StartCommand::default();

let options = cmd
.load_options(&CliOptions {
.load_options(&GlobalOptions {
log_dir: Some("/tmp/greptimedb/test/logs".to_string()),
log_level: Some("debug".to_string()),

Expand Down Expand Up @@ -504,7 +505,8 @@ mod tests {
..Default::default()
};

let Options::Datanode(opts) = command.load_options(&CliOptions::default()).unwrap()
let Options::Datanode(opts) =
command.load_options(&GlobalOptions::default()).unwrap()
else {
unreachable!()
};
Expand Down
29 changes: 15 additions & 14 deletions src/cmd/src/frontend.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ use servers::Mode;
use snafu::{OptionExt, ResultExt};

use crate::error::{self, InitTimezoneSnafu, MissingConfigSnafu, Result, StartFrontendSnafu};
use crate::options::{CliOptions, Options};
use crate::options::{GlobalOptions, Options};
use crate::App;

pub struct Instance {
Expand Down Expand Up @@ -90,8 +90,8 @@ impl Command {
self.subcmd.build(opts).await
}

pub fn load_options(&self, cli_options: &CliOptions) -> Result<Options> {
self.subcmd.load_options(cli_options)
pub fn load_options(&self, global_options: &GlobalOptions) -> Result<Options> {
self.subcmd.load_options(global_options)
}
}

Expand All @@ -107,9 +107,9 @@ impl SubCommand {
}
}

fn load_options(&self, cli_options: &CliOptions) -> Result<Options> {
fn load_options(&self, global_options: &GlobalOptions) -> Result<Options> {
match self {
SubCommand::Start(cmd) => cmd.load_options(cli_options),
SubCommand::Start(cmd) => cmd.load_options(global_options),
}
}
}
Expand Down Expand Up @@ -147,19 +147,19 @@ pub struct StartCommand {
}

impl StartCommand {
fn load_options(&self, cli_options: &CliOptions) -> Result<Options> {
fn load_options(&self, global_options: &GlobalOptions) -> Result<Options> {
let mut opts: FrontendOptions = Options::load_layered_options(
self.config_file.as_deref(),
self.env_prefix.as_ref(),
FrontendOptions::env_list_keys(),
)?;

if let Some(dir) = &cli_options.log_dir {
if let Some(dir) = &global_options.log_dir {
opts.logging.dir.clone_from(dir);
}

if cli_options.log_level.is_some() {
opts.logging.level.clone_from(&cli_options.log_level);
if global_options.log_level.is_some() {
opts.logging.level.clone_from(&global_options.log_level);
}

let tls_opts = TlsOption::new(
Expand Down Expand Up @@ -304,7 +304,7 @@ mod tests {
use servers::http::HttpOptions;

use super::*;
use crate::options::{CliOptions, ENV_VAR_SEP};
use crate::options::{GlobalOptions, ENV_VAR_SEP};

#[test]
fn test_try_from_start_command() {
Expand All @@ -317,7 +317,8 @@ mod tests {
..Default::default()
};

let Options::Frontend(opts) = command.load_options(&CliOptions::default()).unwrap() else {
let Options::Frontend(opts) = command.load_options(&GlobalOptions::default()).unwrap()
else {
unreachable!()
};

Expand Down Expand Up @@ -367,7 +368,7 @@ mod tests {
..Default::default()
};

let Options::Frontend(fe_opts) = command.load_options(&CliOptions::default()).unwrap()
let Options::Frontend(fe_opts) = command.load_options(&GlobalOptions::default()).unwrap()
else {
unreachable!()
};
Expand Down Expand Up @@ -414,7 +415,7 @@ mod tests {
};

let options = cmd
.load_options(&CliOptions {
.load_options(&GlobalOptions {
log_dir: Some("/tmp/greptimedb/test/logs".to_string()),
log_level: Some("debug".to_string()),

Expand Down Expand Up @@ -500,7 +501,7 @@ mod tests {
};

let Options::Frontend(fe_opts) =
command.load_options(&CliOptions::default()).unwrap()
command.load_options(&GlobalOptions::default()).unwrap()
else {
unreachable!()
};
Expand Down
Loading

0 comments on commit a68072c

Please sign in to comment.