Merge pull request #8687 from youngsofun/fix

feat(download): download support CsvWithNames.
databendlabs · Nov 7, 2022 · b26309a · b26309a · vercel · Nov 7, 2022
2 parents 422b802 + b9ed2f0
commit b26309a
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 10 deletions.
diff --git a/src/query/service/src/servers/http/v1/download.rs b/src/query/service/src/servers/http/v1/download.rs
@@ -20,8 +20,8 @@ use common_catalog::plan::DataSourcePlan;
 use common_catalog::plan::PushDownInfo;
 use common_catalog::table_context::TableContext;
 use common_exception::Result;
+use common_formats::ClickhouseFormatType;
 use common_formats::FileFormatOptionsExt;
-use common_meta_types::StageFileFormatType;
 use common_storages_fuse_result::ResultTable;
 use futures::StreamExt;
 
@@ -37,7 +37,7 @@ pub trait Downloader {
     async fn download(
         &self,
         ctx: Arc<QueryContext>,
-        fmt: StageFileFormatType,
+        format: ClickhouseFormatType,
         limit: Option<usize>,
     ) -> Result<SendableVu8Stream>;
 }
@@ -47,7 +47,7 @@ impl Downloader for ResultTable {
     async fn download(
         &self,
         ctx: Arc<QueryContext>,
-        fmt: StageFileFormatType,
+        format: ClickhouseFormatType,
         limit: Option<usize>,
     ) -> Result<SendableVu8Stream> {
         let push_downs = match limit {
@@ -74,8 +74,8 @@ impl Downloader for ResultTable {
                 push_downs,
             })
             .await?;
-        let mut output_format = FileFormatOptionsExt::get_output_format_from_settings(
-            fmt,
+        let mut output_format = FileFormatOptionsExt::get_output_format_from_settings_clickhouse(
+            format,
             self.schema(),
             &ctx.get_settings(),
         )?;

diff --git a/src/query/service/src/servers/http/v1/http_query_handlers.rs b/src/query/service/src/servers/http/v1/http_query_handlers.rs
@@ -12,12 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::str::FromStr;
-
 use common_datavalues::DataSchemaRef;
 use common_exception::ErrorCode;
-use common_meta_types::StageFileFormatType;
+use common_formats::ClickhouseFormatType;
 use common_storages_fuse_result::ResultTable;
+use poem::error::BadRequest;
 use poem::error::Error as PoemError;
 use poem::error::InternalServerError;
 use poem::error::NotFound;
@@ -332,9 +331,9 @@ async fn result_download_handler(
     Query(params): Query<DownloadHandlerParams>,
 ) -> PoemResult<Body> {
     let default_format = "csv".to_string();
+    let format_name = &params.format.unwrap_or(default_format);
     let session = ctx.get_session(SessionType::HTTPQuery);
-    let format = StageFileFormatType::from_str(&params.format.unwrap_or(default_format))
-        .map_err(|e| PoemError::from_string(e, StatusCode::BAD_REQUEST))?;
+    let format = ClickhouseFormatType::parse_clickhouse_format(format_name).map_err(BadRequest)?;
 
     let http_query_manager = HttpQueryManager::instance();
     if let Some(query) = http_query_manager.get_query(&query_id).await {

diff --git a/src/query/service/tests/it/servers/http/http_query_handlers.rs b/src/query/service/tests/it/servers/http/http_query_handlers.rs
@@ -1106,6 +1106,28 @@ pub async fn download(ep: &EndpointType, query_id: &str) -> Response {
     get_uri(ep, &uri).await
 }
 
+#[tokio::test(flavor = "current_thread")]
+async fn test_download_csv_with_names() -> Result<()> {
+    let _guard = TestGlobalServices::setup(ConfigBuilder::create().build()).await?;
+
+    let ep = create_endpoint().await?;
+
+    let sql = "select number, number + 1 from numbers(2)";
+    let (status, result) = post_sql_to_endpoint_new_session(&ep, sql, 2).await?;
+    assert_eq!(status, StatusCode::OK, "{:?}", result);
+    assert_eq!(result.data.len(), 2, "{:?}", result);
+    assert_eq!(result.state, ExecuteStateKind::Succeeded, "{:?}", result);
+
+    // succeeded query
+    let uri = format!("/v1/query/{}/download?format=csvWithNames", result.id);
+    let resp = get_uri(&ep, &uri).await;
+
+    assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
+    let exp = "\"number\",\"number + 1\"\n0,1\n1,2\n";
+    assert_eq!(resp.into_body().into_string().await.unwrap(), exp);
+    Ok(())
+}
+
 #[tokio::test(flavor = "current_thread")]
 async fn test_download() -> Result<()> {
     let _guard = TestGlobalServices::setup(ConfigBuilder::create().build()).await?;