feat: introduce DynamicTimeoutLayer (#5006)

* feat: introduce `DynamicTimeoutLayer`

* test: add unit test

* chore: apply suggestions from CR

* feat: add timeout option for cli

src/cmd/Cargo.toml

@@ -53,6 +53,7 @@ flow.workspace = true
 frontend = { workspace = true, default-features = false }
 futures.workspace = true
 human-panic = "2.0"
+humantime.workspace = true
 lazy_static.workspace = true
 meta-client.workspace = true
 meta-srv.workspace = true

src/cmd/src/cli/database.rs

@@ -12,11 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::time::Duration;
+
 use base64::engine::general_purpose;
 use base64::Engine;
 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+use humantime::format_duration;
 use serde_json::Value;
 use servers::http::greptime_result_v1::GreptimedbV1Response;
+use servers::http::header::constants::GREPTIME_DB_HEADER_TIMEOUT;
 use servers::http::GreptimeQueryOutput;
 use snafu::ResultExt;
 
@@ -26,10 +30,16 @@ pub(crate) struct DatabaseClient {
     addr: String,
     catalog: String,
     auth_header: Option<String>,
+    timeout: Option<Duration>,
 }
 
 impl DatabaseClient {
-    pub fn new(addr: String, catalog: String, auth_basic: Option<String>) -> Self {
+    pub fn new(
+        addr: String,
+        catalog: String,
+        auth_basic: Option<String>,
+        timeout: Option<Duration>,
+    ) -> Self {
         let auth_header = if let Some(basic) = auth_basic {
             let encoded = general_purpose::STANDARD.encode(basic);
             Some(format!("basic {}", encoded))
@@ -41,6 +51,7 @@ impl DatabaseClient {
             addr,
             catalog,
             auth_header,
+            timeout,
         }
     }
 
@@ -62,6 +73,12 @@ impl DatabaseClient {
         if let Some(ref auth) = self.auth_header {
             request = request.header("Authorization", auth);
         }
+        if let Some(ref timeout) = self.timeout {
+            request = request.header(
+                GREPTIME_DB_HEADER_TIMEOUT,
+                format_duration(*timeout).to_string(),
+            );
+        }
 
         let response = request.send().await.with_context(|_| HttpQuerySqlSnafu {
             reason: format!("bad url: {}", url),

src/cmd/src/cli/export.rs

@@ -15,6 +15,7 @@
 use std::collections::HashSet;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::time::Duration;
 
 use async_trait::async_trait;
 use clap::{Parser, ValueEnum};
@@ -83,14 +84,22 @@ pub struct ExportCommand {
     /// The basic authentication for connecting to the server
     #[clap(long)]
     auth_basic: Option<String>,
+
+    /// The timeout of invoking the database.
+    #[clap(long, value_parser = humantime::parse_duration)]
+    timeout: Option<Duration>,
 }
 
 impl ExportCommand {
     pub async fn build(&self, guard: Vec<WorkerGuard>) -> Result<Instance> {
         let (catalog, schema) = database::split_database(&self.database)?;
 
-        let database_client =
-            DatabaseClient::new(self.addr.clone(), catalog.clone(), self.auth_basic.clone());
+        let database_client = DatabaseClient::new(
+            self.addr.clone(),
+            catalog.clone(),
+            self.auth_basic.clone(),
+            self.timeout,
+        );
 
         Ok(Instance::new(
             Box::new(Export {

src/cmd/src/cli/import.rs

@@ -14,6 +14,7 @@
 
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::time::Duration;
 
 use async_trait::async_trait;
 use clap::{Parser, ValueEnum};
@@ -68,13 +69,21 @@ pub struct ImportCommand {
     /// The basic authentication for connecting to the server
     #[clap(long)]
     auth_basic: Option<String>,
+
+    /// The timeout of invoking the database.
+    #[clap(long, value_parser = humantime::parse_duration)]
+    timeout: Option<Duration>,
 }
 
 impl ImportCommand {
     pub async fn build(&self, guard: Vec<WorkerGuard>) -> Result<Instance> {
         let (catalog, schema) = database::split_database(&self.database)?;
-        let database_client =
-            DatabaseClient::new(self.addr.clone(), catalog.clone(), self.auth_basic.clone());
+        let database_client = DatabaseClient::new(
+            self.addr.clone(),
+            catalog.clone(),
+            self.auth_basic.clone(),
+            self.timeout,
+        );
 
         Ok(Instance::new(
             Box::new(Import {