From df751c38b4014641da0cf59a9e13ab5d4d44355b Mon Sep 17 00:00:00 2001 From: LFC Date: Mon, 27 Feb 2023 11:00:15 +0800 Subject: [PATCH] feat: a simple REPL for debugging purpose (#1048) * feat: a simple REPL for debugging purpose * fix: rebase develop --- Cargo.lock | 26 +++++ src/cmd/Cargo.toml | 7 ++ src/cmd/src/bin/greptime.rs | 6 +- src/cmd/src/cli.rs | 62 +++++++++++ src/cmd/src/cli/cmd.rs | 154 ++++++++++++++++++++++++++++ src/cmd/src/cli/helper.rs | 112 ++++++++++++++++++++ src/cmd/src/cli/repl.rs | 199 ++++++++++++++++++++++++++++++++++++ src/cmd/src/error.rs | 44 +++++++- src/cmd/src/lib.rs | 1 + src/cmd/tests/cli.rs | 145 ++++++++++++++++++++++++++ 10 files changed, 754 insertions(+), 2 deletions(-) create mode 100644 src/cmd/src/cli.rs create mode 100644 src/cmd/src/cli/cmd.rs create mode 100644 src/cmd/src/cli/helper.rs create mode 100644 src/cmd/src/cli/repl.rs create mode 100644 src/cmd/tests/cli.rs diff --git a/Cargo.lock b/Cargo.lock index bd28e140c1..aa2a385182 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1350,14 +1350,21 @@ dependencies = [ "anymap", "build-data", "clap 3.2.23", + "client", "common-base", "common-error", + "common-query", + "common-recordbatch", "common-telemetry", "datanode", + "either", "frontend", "futures", "meta-client", "meta-srv", + "nu-ansi-term", + "rexpect", + "rustyline", "serde", "servers", "snafu", @@ -1387,6 +1394,12 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "comma" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55b672471b4e9f9e95499ea597ff64941a309b2cdbffcc46f2cc5e2d971fd335" + [[package]] name = "common-base" version = "0.1.0" @@ -5904,6 +5917,19 @@ dependencies = [ "syn-ext", ] +[[package]] +name = "rexpect" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01ff60778f96fb5a48adbe421d21bf6578ed58c0872d712e7e08593c195adff8" +dependencies = [ + "comma", + "nix 0.25.1", + "regex", + "tempfile", + "thiserror", +] + [[package]] name = "ring" version = "0.16.20" diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index b55960817f..2a0c891d45 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -12,16 +12,22 @@ path = "src/bin/greptime.rs" [dependencies] anymap = "1.0.0-beta.2" clap = { version = "3.1", features = ["derive"] } +client = { path = "../client" } common-base = { path = "../common/base" } common-error = { path = "../common/error" } +common-query = { path = "../common/query" } +common-recordbatch = { path = "../common/recordbatch" } common-telemetry = { path = "../common/telemetry", features = [ "deadlock_detection", ] } datanode = { path = "../datanode" } +either = "1.8" frontend = { path = "../frontend" } futures.workspace = true meta-client = { path = "../meta-client" } meta-srv = { path = "../meta-srv" } +nu-ansi-term = "0.46" +rustyline = "10.1" serde.workspace = true servers = { path = "../servers" } snafu.workspace = true @@ -29,6 +35,7 @@ tokio.workspace = true toml = "0.5" [dev-dependencies] +rexpect = "0.5" serde.workspace = true tempdir = "0.3" diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index 6ee20faf44..b1dec90ce8 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -16,7 +16,7 @@ use std::fmt; use clap::Parser; use cmd::error::Result; -use cmd::{datanode, frontend, metasrv, standalone}; +use cmd::{cli, datanode, frontend, metasrv, standalone}; use common_telemetry::logging::{error, info}; #[derive(Parser)] @@ -46,6 +46,8 @@ enum SubCommand { Metasrv(metasrv::Command), #[clap(name = "standalone")] Standalone(standalone::Command), + #[clap(name = "cli")] + Cli(cli::Command), } impl SubCommand { @@ -55,6 +57,7 @@ impl SubCommand { SubCommand::Frontend(cmd) => cmd.run().await, SubCommand::Metasrv(cmd) => cmd.run().await, SubCommand::Standalone(cmd) => cmd.run().await, + SubCommand::Cli(cmd) => cmd.run().await, } } } @@ -66,6 +69,7 @@ impl fmt::Display for SubCommand { SubCommand::Frontend(..) => write!(f, "greptime-frontend"), SubCommand::Metasrv(..) => write!(f, "greptime-metasrv"), SubCommand::Standalone(..) => write!(f, "greptime-standalone"), + SubCommand::Cli(_) => write!(f, "greptime-cli"), } } } diff --git a/src/cmd/src/cli.rs b/src/cmd/src/cli.rs new file mode 100644 index 0000000000..6de7a91a39 --- /dev/null +++ b/src/cmd/src/cli.rs @@ -0,0 +1,62 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod cmd; +mod helper; +mod repl; + +use clap::Parser; +use repl::Repl; + +use crate::error::Result; + +#[derive(Parser)] +pub struct Command { + #[clap(subcommand)] + cmd: SubCommand, +} + +impl Command { + pub async fn run(self) -> Result<()> { + self.cmd.run().await + } +} + +#[derive(Parser)] +enum SubCommand { + Attach(AttachCommand), +} + +impl SubCommand { + async fn run(self) -> Result<()> { + match self { + SubCommand::Attach(cmd) => cmd.run().await, + } + } +} + +#[derive(Debug, Parser)] +pub(crate) struct AttachCommand { + #[clap(long)] + pub(crate) grpc_addr: String, + #[clap(long, action)] + pub(crate) disable_helper: bool, +} + +impl AttachCommand { + async fn run(self) -> Result<()> { + let mut repl = Repl::try_new(&self)?; + repl.run().await + } +} diff --git a/src/cmd/src/cli/cmd.rs b/src/cmd/src/cli/cmd.rs new file mode 100644 index 0000000000..557a02b385 --- /dev/null +++ b/src/cmd/src/cli/cmd.rs @@ -0,0 +1,154 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::error::{Error, InvalidReplCommandSnafu, Result}; + +/// Represents the parsed command from the user (which may be over many lines) +#[derive(Debug, PartialEq)] +pub(crate) enum ReplCommand { + Help, + UseDatabase { db_name: String }, + Sql { sql: String }, + Exit, +} + +impl TryFrom<&str> for ReplCommand { + type Error = Error; + + fn try_from(input: &str) -> Result { + let input = input.trim(); + if input.is_empty() { + return InvalidReplCommandSnafu { + reason: "No command specified".to_string(), + } + .fail(); + } + + // If line ends with ';', it must be treated as a complete input. + // However, the opposite is not true. + let input_is_completed = input.ends_with(';'); + + let input = input.strip_suffix(';').map(|x| x.trim()).unwrap_or(input); + let lowercase = input.to_lowercase(); + match lowercase.as_str() { + "help" => Ok(Self::Help), + "exit" | "quit" => Ok(Self::Exit), + _ => match input.split_once(' ') { + Some((maybe_use, database)) if maybe_use.to_lowercase() == "use" => { + Ok(Self::UseDatabase { + db_name: database.trim().to_string(), + }) + } + // Any valid SQL must contains at least one whitespace. + Some(_) if input_is_completed => Ok(Self::Sql { + sql: input.to_string(), + }), + _ => InvalidReplCommandSnafu { + reason: format!("unknown command '{input}', maybe input is not completed"), + } + .fail(), + }, + } + } +} + +impl ReplCommand { + pub fn help() -> &'static str { + r#" +Available commands (case insensitive): +- 'help': print this help +- 'exit' or 'quit': exit the REPL +- 'use ': switch to another database/schema context +- Other typed in text will be treated as SQL. + You can enter new line while typing, just remember to end it with ';'. +"# + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::Error::InvalidReplCommand; + + #[test] + fn test_from_str() { + fn test_ok(s: &str, expected: ReplCommand) { + let actual: ReplCommand = s.try_into().unwrap(); + assert_eq!(expected, actual, "'{}'", s); + } + + fn test_err(s: &str) { + let result: Result = s.try_into(); + assert!(matches!(result, Err(InvalidReplCommand { .. }))) + } + + test_err(""); + test_err(" "); + test_err("\t"); + + test_ok("help", ReplCommand::Help); + test_ok("help", ReplCommand::Help); + test_ok(" help", ReplCommand::Help); + test_ok(" help ", ReplCommand::Help); + test_ok(" HELP ", ReplCommand::Help); + test_ok(" Help; ", ReplCommand::Help); + test_ok(" help ; ", ReplCommand::Help); + + test_ok("exit", ReplCommand::Exit); + test_ok("exit;", ReplCommand::Exit); + test_ok("exit ;", ReplCommand::Exit); + test_ok("EXIT", ReplCommand::Exit); + + test_ok("quit", ReplCommand::Exit); + test_ok("quit;", ReplCommand::Exit); + test_ok("quit ;", ReplCommand::Exit); + test_ok("QUIT", ReplCommand::Exit); + + test_ok( + "use Foo", + ReplCommand::UseDatabase { + db_name: "Foo".to_string(), + }, + ); + test_ok( + " use Foo ; ", + ReplCommand::UseDatabase { + db_name: "Foo".to_string(), + }, + ); + // ensure that database name is case sensitive + test_ok( + " use FOO ; ", + ReplCommand::UseDatabase { + db_name: "FOO".to_string(), + }, + ); + + // ensure that we aren't messing with capitalization + test_ok( + "SELECT * from foo;", + ReplCommand::Sql { + sql: "SELECT * from foo".to_string(), + }, + ); + // Input line (that don't belong to any other cases above) must ends with ';' to make it a valid SQL. + test_err("insert blah"); + test_ok( + "insert blah;", + ReplCommand::Sql { + sql: "insert blah".to_string(), + }, + ); + } +} diff --git a/src/cmd/src/cli/helper.rs b/src/cmd/src/cli/helper.rs new file mode 100644 index 0000000000..08b1259514 --- /dev/null +++ b/src/cmd/src/cli/helper.rs @@ -0,0 +1,112 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::borrow::Cow; + +use rustyline::completion::Completer; +use rustyline::highlight::{Highlighter, MatchingBracketHighlighter}; +use rustyline::hint::{Hinter, HistoryHinter}; +use rustyline::validate::{ValidationContext, ValidationResult, Validator}; + +use crate::cli::cmd::ReplCommand; + +pub(crate) struct RustylineHelper { + hinter: HistoryHinter, + highlighter: MatchingBracketHighlighter, +} + +impl Default for RustylineHelper { + fn default() -> Self { + Self { + hinter: HistoryHinter {}, + highlighter: MatchingBracketHighlighter::default(), + } + } +} + +impl rustyline::Helper for RustylineHelper {} + +impl Validator for RustylineHelper { + fn validate(&self, ctx: &mut ValidationContext<'_>) -> rustyline::Result { + let input = ctx.input(); + match ReplCommand::try_from(input) { + Ok(_) => Ok(ValidationResult::Valid(None)), + Err(e) => { + if input.trim_end().ends_with(';') { + // If line ends with ';', it HAS to be a valid command. + Ok(ValidationResult::Invalid(Some(e.to_string()))) + } else { + Ok(ValidationResult::Incomplete) + } + } + } + } +} + +impl Hinter for RustylineHelper { + type Hint = String; + + fn hint(&self, line: &str, pos: usize, ctx: &rustyline::Context<'_>) -> Option { + self.hinter.hint(line, pos, ctx) + } +} + +impl Highlighter for RustylineHelper { + fn highlight<'l>(&self, line: &'l str, pos: usize) -> Cow<'l, str> { + self.highlighter.highlight(line, pos) + } + + fn highlight_prompt<'b, 's: 'b, 'p: 'b>( + &'s self, + prompt: &'p str, + default: bool, + ) -> Cow<'b, str> { + self.highlighter.highlight_prompt(prompt, default) + } + + fn highlight_hint<'h>(&self, hint: &'h str) -> Cow<'h, str> { + use nu_ansi_term::Style; + Cow::Owned(Style::new().dimmed().paint(hint).to_string()) + } + + fn highlight_candidate<'c>( + &self, + candidate: &'c str, + completion: rustyline::CompletionType, + ) -> Cow<'c, str> { + self.highlighter.highlight_candidate(candidate, completion) + } + + fn highlight_char(&self, line: &str, pos: usize) -> bool { + self.highlighter.highlight_char(line, pos) + } +} + +impl Completer for RustylineHelper { + type Candidate = String; + + fn complete( + &self, + line: &str, + pos: usize, + ctx: &rustyline::Context<'_>, + ) -> rustyline::Result<(usize, Vec)> { + // If there is a hint, use that as the auto-complete when user hits `tab` + if let Some(hint) = self.hinter.hint(line, pos, ctx) { + Ok((pos, vec![hint])) + } else { + Ok((0, vec![])) + } + } +} diff --git a/src/cmd/src/cli/repl.rs b/src/cmd/src/cli/repl.rs new file mode 100644 index 0000000000..ae0e4d0c46 --- /dev/null +++ b/src/cmd/src/cli/repl.rs @@ -0,0 +1,199 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::path::PathBuf; +use std::time::Instant; + +use client::{Client, Database, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_error::prelude::ErrorExt; +use common_query::Output; +use common_recordbatch::RecordBatches; +use common_telemetry::logging; +use either::Either; +use rustyline::error::ReadlineError; +use rustyline::Editor; +use snafu::{ErrorCompat, ResultExt}; + +use crate::cli::cmd::ReplCommand; +use crate::cli::helper::RustylineHelper; +use crate::cli::AttachCommand; +use crate::error::{ + CollectRecordBatchesSnafu, PrettyPrintRecordBatchesSnafu, ReadlineSnafu, ReplCreationSnafu, + RequestDatabaseSnafu, Result, +}; + +/// Captures the state of the repl, gathers commands and executes them one by one +pub(crate) struct Repl { + /// Rustyline editor for interacting with user on command line + rl: Editor, + + /// Current prompt + prompt: String, + + /// Client for interacting with GreptimeDB + database: Database, +} + +#[allow(clippy::print_stdout)] +impl Repl { + fn print_help(&self) { + println!("{}", ReplCommand::help()) + } + + pub(crate) fn try_new(cmd: &AttachCommand) -> Result { + let mut rl = Editor::new().context(ReplCreationSnafu)?; + + if !cmd.disable_helper { + rl.set_helper(Some(RustylineHelper::default())); + + let history_file = history_file(); + if let Err(e) = rl.load_history(&history_file) { + logging::debug!( + "failed to load history file on {}, error: {e}", + history_file.display() + ); + } + } + + let client = Client::with_urls([&cmd.grpc_addr]); + let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client); + + Ok(Self { + rl, + prompt: "> ".to_string(), + database, + }) + } + + /// Parse the next command + fn next_command(&mut self) -> Result { + match self.rl.readline(&self.prompt) { + Ok(ref line) => { + let request = line.trim(); + + self.rl.add_history_entry(request.to_string()); + + request.try_into() + } + Err(ReadlineError::Eof) | Err(ReadlineError::Interrupted) => Ok(ReplCommand::Exit), + // Some sort of real underlying error + Err(e) => Err(e).context(ReadlineSnafu), + } + } + + /// Read Evaluate Print Loop (interactive command line) for GreptimeDB + /// + /// Inspired / based on repl.rs from InfluxDB IOX + pub(crate) async fn run(&mut self) -> Result<()> { + println!("Ready for commands. (Hint: try 'help')"); + + loop { + match self.next_command()? { + ReplCommand::Help => { + self.print_help(); + } + ReplCommand::UseDatabase { db_name } => { + if self.execute_sql(format!("USE {db_name}")).await { + println!("Using {db_name}"); + self.database.set_schema(&db_name); + self.prompt = format!("[{db_name}] > "); + } + } + ReplCommand::Sql { sql } => { + self.execute_sql(sql).await; + } + ReplCommand::Exit => { + return Ok(()); + } + } + } + } + + async fn execute_sql(&self, sql: String) -> bool { + self.do_execute_sql(sql) + .await + .map_err(|e| { + let status_code = e.status_code(); + let root_cause = e.iter_chain().last().unwrap(); + println!("Error: {}({status_code}), {root_cause}", status_code as u32) + }) + .is_ok() + } + + async fn do_execute_sql(&self, sql: String) -> Result<()> { + let start = Instant::now(); + + let output = self + .database + .sql(&sql) + .await + .context(RequestDatabaseSnafu { sql: &sql })?; + + let either = match output { + Output::Stream(s) => { + let x = RecordBatches::try_collect(s) + .await + .context(CollectRecordBatchesSnafu)?; + Either::Left(x) + } + Output::RecordBatches(x) => Either::Left(x), + Output::AffectedRows(rows) => Either::Right(rows), + }; + + let end = Instant::now(); + + match either { + Either::Left(recordbatches) => { + let total_rows: usize = recordbatches.iter().map(|x| x.num_rows()).sum(); + if total_rows > 0 { + println!( + "{}", + recordbatches + .pretty_print() + .context(PrettyPrintRecordBatchesSnafu)? + ); + } + println!("Total Rows: {total_rows}") + } + Either::Right(rows) => println!("Affected Rows: {rows}"), + }; + + println!("Cost {} ms", (end - start).as_millis()); + Ok(()) + } +} + +impl Drop for Repl { + fn drop(&mut self) { + if self.rl.helper().is_some() { + let history_file = history_file(); + if let Err(e) = self.rl.save_history(&history_file) { + logging::debug!( + "failed to save history file on {}, error: {e}", + history_file.display() + ); + } + } + } +} + +/// Return the location of the history file (defaults to $HOME/".greptimedb_cli_history") +fn history_file() -> PathBuf { + let mut buf = match std::env::var("HOME") { + Ok(home) => PathBuf::from(home), + Err(_) => PathBuf::new(), + }; + buf.push(".greptimedb_cli_history"); + buf +} diff --git a/src/cmd/src/error.rs b/src/cmd/src/error.rs index fd2eb7d1ae..209f41b1a1 100644 --- a/src/cmd/src/error.rs +++ b/src/cmd/src/error.rs @@ -15,6 +15,7 @@ use std::any::Any; use common_error::prelude::*; +use rustyline::error::ReadlineError; #[derive(Debug, Snafu)] #[snafu(visibility(pub))] @@ -68,6 +69,40 @@ pub enum Error { #[snafu(backtrace)] source: meta_srv::error::Error, }, + + #[snafu(display("Invalid REPL command: {reason}"))] + InvalidReplCommand { reason: String }, + + #[snafu(display("Cannot create REPL: {}", source))] + ReplCreation { + source: ReadlineError, + backtrace: Backtrace, + }, + + #[snafu(display("Error reading command: {}", source))] + Readline { + source: ReadlineError, + backtrace: Backtrace, + }, + + #[snafu(display("Failed to request database, sql: {sql}, source: {source}"))] + RequestDatabase { + sql: String, + #[snafu(backtrace)] + source: client::Error, + }, + + #[snafu(display("Failed to collect RecordBatches, source: {source}"))] + CollectRecordBatches { + #[snafu(backtrace)] + source: common_recordbatch::error::Error, + }, + + #[snafu(display("Failed to pretty print Recordbatches, source: {source}"))] + PrettyPrintRecordBatches { + #[snafu(backtrace)] + source: common_recordbatch::error::Error, + }, } pub type Result = std::result::Result; @@ -82,8 +117,15 @@ impl ErrorExt for Error { Error::ReadConfig { .. } | Error::ParseConfig { .. } | Error::MissingConfig { .. } => { StatusCode::InvalidArguments } - Error::IllegalConfig { .. } => StatusCode::InvalidArguments, + Error::IllegalConfig { .. } | Error::InvalidReplCommand { .. } => { + StatusCode::InvalidArguments + } Error::IllegalAuthConfig { .. } => StatusCode::InvalidArguments, + Error::ReplCreation { .. } | Error::Readline { .. } => StatusCode::Internal, + Error::RequestDatabase { source, .. } => source.status_code(), + Error::CollectRecordBatches { source } | Error::PrettyPrintRecordBatches { source } => { + source.status_code() + } } } diff --git a/src/cmd/src/lib.rs b/src/cmd/src/lib.rs index 61d694e4ae..157e4853f1 100644 --- a/src/cmd/src/lib.rs +++ b/src/cmd/src/lib.rs @@ -14,6 +14,7 @@ #![feature(assert_matches)] +pub mod cli; pub mod datanode; pub mod error; pub mod frontend; diff --git a/src/cmd/tests/cli.rs b/src/cmd/tests/cli.rs new file mode 100644 index 0000000000..905191e13b --- /dev/null +++ b/src/cmd/tests/cli.rs @@ -0,0 +1,145 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#[cfg(target_os = "macos")] +mod tests { + use std::path::PathBuf; + use std::process::{Command, Stdio}; + use std::time::Duration; + + use rexpect::session::PtyReplSession; + use tempdir::TempDir; + + struct Repl { + repl: PtyReplSession, + } + + impl Repl { + fn send_line(&mut self, line: &str) { + self.repl.send_line(line).unwrap(); + + // read a line to consume the prompt + self.read_line(); + } + + fn read_line(&mut self) -> String { + self.repl.read_line().unwrap() + } + + fn read_expect(&mut self, expect: &str) { + assert_eq!(self.read_line(), expect); + } + + fn read_contains(&mut self, pat: &str) { + assert!(self.read_line().contains(pat)); + } + } + + #[test] + fn test_repl() { + let data_dir = TempDir::new_in("/tmp", "data").unwrap(); + let wal_dir = TempDir::new_in("/tmp", "wal").unwrap(); + + let mut bin_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + bin_path.push("../../target/debug"); + let bin_path = bin_path.to_str().unwrap(); + + let mut datanode = Command::new("./greptime") + .current_dir(bin_path) + .args([ + "datanode", + "start", + "--rpc-addr=0.0.0.0:4321", + "--node-id=1", + &format!("--data-dir={}", data_dir.path().display()), + &format!("--wal-dir={}", wal_dir.path().display()), + ]) + .stdout(Stdio::null()) + .spawn() + .unwrap(); + + // wait for Datanode actually started + std::thread::sleep(Duration::from_secs(3)); + + let mut repl_cmd = Command::new("./greptime"); + repl_cmd.current_dir(bin_path).args([ + "--log-level=off", + "cli", + "attach", + "--grpc-addr=0.0.0.0:4321", + // history commands can sneaky into stdout and mess up our tests, so disable it + "--disable-helper", + ]); + let pty_session = rexpect::session::spawn_command(repl_cmd, Some(5_000)).unwrap(); + let repl = PtyReplSession { + prompt: "> ".to_string(), + pty_session, + quit_command: None, + echo_on: false, + }; + let repl = &mut Repl { repl }; + repl.read_expect("Ready for commands. (Hint: try 'help')"); + + test_create_database(repl); + + test_use_database(repl); + + test_create_table(repl); + + test_insert(repl); + + test_select(repl); + + datanode.kill().unwrap(); + datanode.wait().unwrap(); + } + + fn test_create_database(repl: &mut Repl) { + repl.send_line("CREATE DATABASE db;"); + repl.read_expect("Affected Rows: 1"); + repl.read_contains("Cost"); + } + + fn test_use_database(repl: &mut Repl) { + repl.send_line("USE db"); + repl.read_expect("Total Rows: 0"); + repl.read_contains("Cost"); + repl.read_expect("Using db"); + } + + fn test_create_table(repl: &mut Repl) { + repl.send_line("CREATE TABLE t(x STRING, ts TIMESTAMP TIME INDEX);"); + repl.read_expect("Affected Rows: 0"); + repl.read_contains("Cost"); + } + + fn test_insert(repl: &mut Repl) { + repl.send_line("INSERT INTO t(x, ts) VALUES ('hello', 1676895812239);"); + repl.read_expect("Affected Rows: 1"); + repl.read_contains("Cost"); + } + + fn test_select(repl: &mut Repl) { + repl.send_line("SELECT * FROM t;"); + + repl.read_expect("+-------+-------------------------+"); + repl.read_expect("| x | ts |"); + repl.read_expect("+-------+-------------------------+"); + repl.read_expect("| hello | 2023-02-20T12:23:32.239 |"); + repl.read_expect("+-------+-------------------------+"); + repl.read_expect("Total Rows: 1"); + + repl.read_contains("Cost"); + } +}