fix: fix create or replace stage bug, only remove dir when success (#…

…14766) * fix: fix create or replace stage bug, only remove dir when success * fix: fix create or replace stage bug, only remove dir when success
databendlabs · Feb 27, 2024 · ea87365 · ea87365
1 parent e88992f
commit ea87365
Showing 1 changed file with 22 additions and 19 deletions.
diff --git a/src/query/service/src/interpreters/interpreter_user_stage_create.rs b/src/query/service/src/interpreters/interpreter_user_stage_create.rs
@@ -82,20 +82,30 @@ impl Interpreter for CreateUserStageInterpreter {
             )));
         };
 
-        // when create or replace stage, if old stage is not External stage, remove stage files
-        if let CreateOption::CreateOrReplace = plan.create_option {
-            if let Ok(stage) = user_mgr
+        let old_stage = match plan.create_option {
+            CreateOption::CreateOrReplace => user_mgr
                 .get_stage(&plan.tenant, &user_stage.stage_name)
                 .await
-            {
-                if stage.stage_type != StageType::External {
-                    let op = StageTable::get_op(&stage)?;
-                    op.remove_all("/").await?;
-                    info!(
-                        "create or replace stage {:?} with all objects removed in stage",
-                        user_stage.stage_name
-                    );
-                }
+                .ok(),
+            _ => None,
+        };
+
+        let mut user_stage = user_stage;
+        user_stage.creator = Some(self.ctx.get_current_user()?.identity());
+        user_stage.created_on = Utc::now();
+        let _ = user_mgr
+            .add_stage(&plan.tenant, user_stage.clone(), &plan.create_option)
+            .await?;
+
+        // when create or replace stage success, if old stage is not External stage, remove stage files
+        if let Some(stage) = old_stage {
+            if stage.stage_type != StageType::External {
+                let op = StageTable::get_op(&stage)?;
+                op.remove_all("/").await?;
+                info!(
+                    "create or replace stage {:?} with all objects removed in stage",
+                    user_stage.stage_name
+                );
             }
         }
 
@@ -105,13 +115,6 @@ impl Interpreter for CreateUserStageInterpreter {
             op.create_dir(&user_stage.stage_prefix()).await?
         }
 
-        let mut user_stage = user_stage;
-        user_stage.creator = Some(self.ctx.get_current_user()?.identity());
-        user_stage.created_on = Utc::now();
-        let _ = user_mgr
-            .add_stage(&plan.tenant, user_stage, &plan.create_option)
-            .await?;
-
         // Grant ownership as the current role
         let tenant = self.ctx.get_tenant();
         let role_api = UserApiProvider::instance().get_role_api_client(&tenant)?;