diff --git a/src/cmd/src/cli/export.rs b/src/cmd/src/cli/export.rs index cdd8a9dba2c4..d2222fa5b30c 100644 --- a/src/cmd/src/cli/export.rs +++ b/src/cmd/src/cli/export.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::collections::HashSet; -use std::path::Path; +use std::path::PathBuf; use std::sync::Arc; use async_trait::async_trait; @@ -120,6 +120,10 @@ pub struct Export { } impl Export { + fn catalog_path(&self) -> PathBuf { + PathBuf::from(&self.output_dir).join(&self.catalog) + } + async fn get_db_names(&self) -> Result> { let db_names = self.all_db_names().await?; let Some(schema) = &self.schema else { @@ -182,8 +186,11 @@ impl Export { and table_catalog = \'{catalog}\' \ and table_schema = \'{schema}\'" ); - let result = self.database_client.sql_in_public(&sql).await?; - let records = result.context(EmptyResultSnafu)?; + let records = self + .database_client + .sql_in_public(&sql) + .await? + .context(EmptyResultSnafu)?; let mut metric_physical_tables = HashSet::with_capacity(records.len()); for value in records { let mut t = Vec::with_capacity(3); @@ -203,8 +210,11 @@ impl Export { and table_catalog = \'{catalog}\' \ and table_schema = \'{schema}\'", ); - let result = self.database_client.sql_in_public(&sql).await?; - let records = result.context(EmptyResultSnafu)?; + let records = self + .database_client + .sql_in_public(&sql) + .await? + .context(EmptyResultSnafu)?; debug!("Fetched table/view list: {:?}", records); @@ -255,8 +265,11 @@ impl Export { ), None => format!(r#"SHOW CREATE {} "{}"."{}""#, show_type, catalog, schema), }; - let result = self.database_client.sql_in_public(&sql).await?; - let records = result.context(EmptyResultSnafu)?; + let records = self + .database_client + .sql_in_public(&sql) + .await? + .context(EmptyResultSnafu)?; let Value::String(create) = &records[0][1] else { unreachable!() }; @@ -269,13 +282,11 @@ impl Export { let db_names = self.get_db_names().await?; let db_count = db_names.len(); for schema in db_names { - let output_dir = Path::new(&self.output_dir) - .join(&self.catalog) - .join(format!("{schema}/")); - tokio::fs::create_dir_all(&output_dir) + let db_dir = self.catalog_path().join(format!("{schema}/")); + tokio::fs::create_dir_all(&db_dir) .await .context(FileIoSnafu)?; - let file = Path::new(&output_dir).join("create_database.sql"); + let file = db_dir.join("create_database.sql"); let mut file = File::create(file).await.context(FileIoSnafu)?; match self .show_create("DATABASE", &self.catalog, &schema, None) @@ -312,13 +323,11 @@ impl Export { self.get_table_list(&self.catalog, &schema).await?; let table_count = metric_physical_tables.len() + remaining_tables.len() + views.len(); - let output_dir = Path::new(&self.output_dir) - .join(&self.catalog) - .join(format!("{schema}/")); - tokio::fs::create_dir_all(&output_dir) + let db_dir = self.catalog_path().join(format!("{schema}/")); + tokio::fs::create_dir_all(&db_dir) .await .context(FileIoSnafu)?; - let file = Path::new(&output_dir).join("create_tables.sql"); + let file = db_dir.join("create_tables.sql"); let mut file = File::create(file).await.context(FileIoSnafu)?; for (c, s, t) in metric_physical_tables.into_iter().chain(remaining_tables) { match self.show_create("TABLE", &c, &s, Some(&t)).await { @@ -348,7 +357,7 @@ impl Export { info!( "Finished exporting {}.{schema} with {table_count} table schemas to path: {}", self.catalog, - output_dir.to_string_lossy() + db_dir.to_string_lossy() ); Ok::<(), Error>(()) @@ -383,10 +392,8 @@ impl Export { let semaphore_moved = semaphore.clone(); tasks.push(async move { let _permit = semaphore_moved.acquire().await.unwrap(); - let output_dir = Path::new(&self.output_dir) - .join(&self.catalog) - .join(format!("{schema}/")); - tokio::fs::create_dir_all(&output_dir) + let db_dir = self.catalog_path().join(format!("{schema}/")); + tokio::fs::create_dir_all(&db_dir) .await .context(FileIoSnafu)?; @@ -410,7 +417,7 @@ impl Export { r#"COPY DATABASE "{}"."{}" TO '{}' {};"#, self.catalog, schema, - output_dir.to_str().unwrap(), + db_dir.to_str().unwrap(), with_options ); @@ -421,18 +428,18 @@ impl Export { info!( "Finished exporting {}.{schema} data into path: {}", self.catalog, - output_dir.to_string_lossy() + db_dir.to_string_lossy() ); // The export copy from sql - let copy_from_file = output_dir.join("copy_from.sql"); + let copy_from_file = db_dir.join("copy_from.sql"); let mut writer = BufWriter::new(File::create(copy_from_file).await.context(FileIoSnafu)?); let copy_database_from_sql = format!( r#"COPY DATABASE "{}"."{}" FROM '{}' WITH (FORMAT='parquet');"#, self.catalog, schema, - output_dir.to_str().unwrap() + db_dir.to_str().unwrap() ); writer .write(copy_database_from_sql.as_bytes()) diff --git a/src/cmd/src/cli/import.rs b/src/cmd/src/cli/import.rs index 217c19c3fd90..3f8a6c06a9b4 100644 --- a/src/cmd/src/cli/import.rs +++ b/src/cmd/src/cli/import.rs @@ -17,6 +17,7 @@ use std::sync::Arc; use async_trait::async_trait; use clap::{Parser, ValueEnum}; +use common_catalog::consts::DEFAULT_SCHEMA_NAME; use common_telemetry::{error, info, warn}; use snafu::ResultExt; use tokio::sync::Semaphore; @@ -100,14 +101,17 @@ pub struct Import { impl Import { async fn import_create_table(&self) -> Result<()> { - self.do_sql_job("create_tables.sql").await + // Use default db to creates other dbs + self.do_sql_job("create_database.sql", Some(DEFAULT_SCHEMA_NAME)) + .await?; + self.do_sql_job("create_tables.sql", None).await } async fn import_database_data(&self) -> Result<()> { - self.do_sql_job("copy_from.sql").await + self.do_sql_job("copy_from.sql", None).await } - async fn do_sql_job(&self, filename: &str) -> Result<()> { + async fn do_sql_job(&self, filename: &str, exec_db: Option<&str>) -> Result<()> { let timer = Instant::now(); let semaphore = Arc::new(Semaphore::new(self.parallelism)); let db_names = self.get_db_names().await?; @@ -125,7 +129,8 @@ impl Import { if sql.is_empty() { info!("Empty `{filename}` {database_input_dir:?}"); } else { - self.database_client.sql(&sql, &schema).await?; + let db = exec_db.unwrap_or(&schema); + self.database_client.sql(&sql, db).await?; info!("Imported `{filename}` for database {schema}"); } diff --git a/src/servers/src/http/handler.rs b/src/servers/src/http/handler.rs index d1690e79a944..4d5ca5846159 100644 --- a/src/servers/src/http/handler.rs +++ b/src/servers/src/http/handler.rs @@ -20,6 +20,7 @@ use aide::transform::TransformOperation; use axum::extract::{Json, Query, State}; use axum::response::{IntoResponse, Response}; use axum::{Extension, Form}; +use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_plugins::GREPTIME_EXEC_WRITE_COST; @@ -76,6 +77,11 @@ pub async fn sql( ) -> HttpResponse { let start = Instant::now(); let sql_handler = &state.sql_handler; + if let Some(db) = &query_params.db.or(form_params.db) { + let (catalog, schema) = parse_catalog_and_schema_from_db_string(db); + query_ctx.set_current_catalog(&catalog); + query_ctx.set_current_schema(&schema); + } let db = query_ctx.get_db_string(); query_ctx.set_channel(Channel::Http);