[Feature]Refine create stage interpreter test

Signed-off-by: grapebaba <281165273@qq.com>
datafuselabs · Dec 6, 2021 · 33b14af · 33b14af
1 parent 1c0906d
commit 33b14af
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 6 deletions.
diff --git a/query/src/interpreters/interpreter_stage_create.rs b/query/src/interpreters/interpreter_stage_create.rs
@@ -51,14 +51,15 @@ impl Interpreter for CreatStageInterpreter {
         let plan = self.plan.clone();
         let user_mgr = self.ctx.get_sessions_manager().get_user_manager();
         let user_stage = plan.user_stage_info;
-        let stage_name = user_stage.stage_name.clone();
         let create_stage = user_mgr.add_stage(user_stage).await;
         if plan.if_not_exists {
-            create_stage.map_err(|_e| {
-                ErrorCode::StageAlreadyExists(format!(
-                    "Stage already exists, name [{}]",
-                    stage_name
-                ))
+            create_stage.or_else(|e| {
+                // StageAlreadyExists(4061)
+                if e.code() == 4061 {
+                    Ok(u64::MIN)
+                } else {
+                    Err(e)
+                }
             })?;
         } else {
             create_stage?;

diff --git a/query/src/interpreters/interpreter_stage_create_test.rs b/query/src/interpreters/interpreter_stage_create_test.rs
@@ -53,5 +53,51 @@ async fn test_create_stage_interpreter() -> Result<()> {
         panic!()
     }
 
+    if let PlanNode::CreateUserStage(plan) = PlanParser::parse(TEST_QUERY, ctx.clone()).await? {
+        let executor = CreatStageInterpreter::try_create(ctx.clone(), plan.clone())?;
+        assert_eq!(executor.name(), "CreatStageInterpreter");
+        let is_err = executor.execute(None).await.is_err();
+        assert!(!is_err);
+        let stage = ctx
+            .get_sessions_manager()
+            .get_user_manager()
+            .get_stage("test_stage")
+            .await?;
+
+        assert_eq!(
+            stage.file_format.clone(),
+            Some(FileFormat::Csv {
+                compression: Compression::Gzip,
+                record_delimiter: ",".to_string()
+            })
+        );
+        assert_eq!(stage.comments.clone(), String::from("test"))
+    } else {
+        panic!()
+    }
+
+    static TEST_QUERY1: &str = "CREATE STAGE test_stage url='s3://load/files/' credentials=(access_key_id='1a2b3c' secret_access_key='4x5y6z') file_format=(FORMAT=CSV compression=GZIP record_delimiter=',') comments='test'";
+    if let PlanNode::CreateUserStage(plan) = PlanParser::parse(TEST_QUERY1, ctx.clone()).await? {
+        let executor = CreatStageInterpreter::try_create(ctx.clone(), plan.clone())?;
+        assert_eq!(executor.name(), "CreatStageInterpreter");
+        let is_err = executor.execute(None).await.is_err();
+        assert!(is_err);
+        let stage = ctx
+            .get_sessions_manager()
+            .get_user_manager()
+            .get_stage("test_stage")
+            .await?;
+
+        assert_eq!(
+            stage.file_format.clone(),
+            Some(FileFormat::Csv {
+                compression: Compression::Gzip,
+                record_delimiter: ",".to_string()
+            })
+        );
+        assert_eq!(stage.comments.clone(), String::from("test"))
+    } else {
+        panic!()
+    }
     Ok(())
 }