feat: import cli tool

src/cmd/src/cli.rs

@@ -12,18 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod bench;
-
-// Wait for https://github.com/GreptimeTeam/greptimedb/issues/2373
-#[allow(unused)]
-mod cmd;
-mod export;
-mod helper;
-
-// Wait for https://github.com/GreptimeTeam/greptimedb/issues/2373
-#[allow(unused)]
-mod repl;
-
 use async_trait::async_trait;
 use bench::BenchTableMetadataCommand;
 use clap::Parser;
@@ -32,10 +20,25 @@ pub use repl::Repl;
 use tracing_appender::non_blocking::WorkerGuard;
 
 use self::export::ExportCommand;
+use crate::cli::import::ImportCommand;
 use crate::error::Result;
 use crate::options::GlobalOptions;
 use crate::App;
 
+mod bench;
+
+// Wait for https://github.com/GreptimeTeam/greptimedb/issues/2373
+#[allow(unused)]
+mod cmd;
+mod export;
+mod helper;
+
+// Wait for https://github.com/GreptimeTeam/greptimedb/issues/2373
+mod database;
+mod import;
+#[allow(unused)]
+mod repl;
+
 pub const APP_NAME: &str = "greptime-cli";
 
 #[async_trait]
@@ -114,6 +117,7 @@ enum SubCommand {
     // Attach(AttachCommand),
     Bench(BenchTableMetadataCommand),
     Export(ExportCommand),
+    Import(ImportCommand),
 }
 
 impl SubCommand {
@@ -122,6 +126,7 @@ impl SubCommand {
             // SubCommand::Attach(cmd) => cmd.build().await,
             SubCommand::Bench(cmd) => cmd.build(guard).await,
             SubCommand::Export(cmd) => cmd.build(guard).await,
+            SubCommand::Import(cmd) => cmd.build(guard).await,
         }
     }
 }

src/cmd/src/cli/database.rs

@@ -0,0 +1,119 @@
+// 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 base64::engine::general_purpose;
+use base64::Engine;
+use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+use serde_json::Value;
+use servers::http::greptime_result_v1::GreptimedbV1Response;
+use servers::http::GreptimeQueryOutput;
+use snafu::ResultExt;
+
+use crate::error::{HttpQuerySqlSnafu, Result, SerdeJsonSnafu};
+
+pub(crate) struct DatabaseClient {
+    addr: String,
+    catalog: String,
+    auth_header: Option<String>,
+}
+
+impl DatabaseClient {
+    pub fn new(addr: String, catalog: String, auth_basic: Option<String>) -> Self {
+        let auth_header = if let Some(basic) = auth_basic {
+            let encoded = general_purpose::STANDARD.encode(basic);
+            Some(format!("basic {}", encoded))
+        } else {
+            None
+        };
+
+        Self {
+            addr,
+            catalog,
+            auth_header,
+        }
+    }
+
+    pub async fn sql_in_public(&self, sql: &str) -> Result<Option<Vec<Vec<Value>>>> {
+        self.sql(sql, DEFAULT_SCHEMA_NAME).await
+    }
+
+    /// Execute sql query.
+    pub async fn sql(&self, sql: &str, schema: &str) -> Result<Option<Vec<Vec<Value>>>> {
+        let url = format!("http://{}/v1/sql", self.addr);
+        let params = [
+            ("db", format!("{}-{}", self.catalog, schema)),
+            ("sql", sql.to_string()),
+        ];
+        let mut request = reqwest::Client::new()
+            .post(&url)
+            .form(&params)
+            .header("Content-Type", "application/x-www-form-urlencoded");
+        if let Some(ref auth) = self.auth_header {
+            request = request.header("Authorization", auth);
+        }
+
+        let response = request.send().await.with_context(|_| HttpQuerySqlSnafu {
+            reason: format!("bad url: {}", url),
+        })?;
+        let response = response
+            .error_for_status()
+            .with_context(|_| HttpQuerySqlSnafu {
+                reason: format!("query failed: {}", sql),
+            })?;
+
+        let text = response.text().await.with_context(|_| HttpQuerySqlSnafu {
+            reason: "cannot get response text".to_string(),
+        })?;
+
+        let body = serde_json::from_str::<GreptimedbV1Response>(&text).context(SerdeJsonSnafu)?;
+        Ok(body.output().first().and_then(|output| match output {
+            GreptimeQueryOutput::Records(records) => Some(records.rows().clone()),
+            GreptimeQueryOutput::AffectedRows(_) => None,
+        }))
+    }
+}
+
+/// Split at `-`.
+pub(crate) fn split_database(database: &str) -> Result<(String, Option<String>)> {
+    let (catalog, schema) = match database.split_once('-') {
+        Some((catalog, schema)) => (catalog, schema),
+        None => (DEFAULT_CATALOG_NAME, database),
+    };
+
+    if schema == "*" {
+        Ok((catalog.to_string(), None))
+    } else {
+        Ok((catalog.to_string(), Some(schema.to_string())))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_split_database() {
+        let result = split_database("catalog-schema").unwrap();
+        assert_eq!(result, ("catalog".to_string(), Some("schema".to_string())));
+
+        let result = split_database("schema").unwrap();
+        assert_eq!(result, ("greptime".to_string(), Some("schema".to_string())));
+
+        let result = split_database("catalog-*").unwrap();
+        assert_eq!(result, ("catalog".to_string(), None));
+
+        let result = split_database("*").unwrap();
+        assert_eq!(result, ("greptime".to_string(), None));
+    }
+}