refactor: Separate the code logic for COPY INTO LOCATION/TABLE. (#18060)

refactor: Separate the code for COPY INTO LOCATION and TABLE.

Author: youngsofun

Cargo.lock

@@ -17,7 +17,6 @@ use std::fmt::Display;
 use std::fmt::Formatter;
 use std::sync::Arc;
 
-use databend_common_ast::ast::CopyIntoLocationOptions;
 use databend_common_ast::ast::CopyIntoTableOptions;
 use databend_common_exception::Result;
 use databend_common_expression::RemoteDefaultExpr;
@@ -48,10 +47,6 @@ pub struct StageTableInfo {
     pub is_select: bool,
     pub copy_into_table_options: CopyIntoTableOptions,
     pub is_variant: bool,
-
-    // copy into location only
-    pub copy_into_location_ordered: bool,
-    pub copy_into_location_options: CopyIntoLocationOptions,
 }
 
 impl StageTableInfo {

src/query/catalog/src/plan/datasource/datasource_info/stage.rs

@@ -116,6 +116,7 @@ databend-storages-common-cache = { workspace = true }
 databend-storages-common-index = { workspace = true }
 databend-storages-common-io = { workspace = true }
 databend-storages-common-session = { workspace = true }
+databend-storages-common-stage = { workspace = true }
 databend-storages-common-table-meta = { workspace = true }
 derive-visitor = { workspace = true }
 ethnum = { workspace = true }

src/query/service/Cargo.toml

@@ -1244,7 +1244,7 @@ impl AccessChecker for PrivilegeAccess {
                 }
             }
             Plan::CopyIntoLocation(plan) => {
-                self.validate_stage_access(&plan.stage, UserPrivilegeType::Write).await?;
+                self.validate_stage_access(&plan.info.stage, UserPrivilegeType::Write).await?;
                 let from = plan.from.clone();
                 return self.check(ctx, &from).await;
             }

src/query/service/src/interpreters/access/privilege_access.rs

@@ -14,17 +14,14 @@
 
 use std::sync::Arc;
 
-use databend_common_ast::ast::CopyIntoLocationOptions;
 use databend_common_base::runtime::GlobalIORuntime;
-use databend_common_catalog::plan::StageTableInfo;
 use databend_common_exception::Result;
 use databend_common_expression::infer_table_schema;
-use databend_common_meta_app::principal::StageInfo;
 use databend_common_meta_app::schema::UpdateStreamMetaReq;
 use databend_common_pipeline_core::ExecutionInfo;
 use databend_common_sql::executor::physical_plans::CopyIntoLocation;
 use databend_common_sql::executor::PhysicalPlan;
-use databend_common_storage::StageFilesInfo;
+use databend_storages_common_stage::CopyIntoLocationInfo;
 use log::debug;
 use log::info;
 
@@ -84,10 +81,8 @@ impl CopyIntoLocationInterpreter {
     #[async_backtrace::framed]
     async fn build_local_copy_into_stage_pipeline(
         &self,
-        stage: &StageInfo,
-        path: &str,
         query: &Plan,
-        options: &CopyIntoLocationOptions,
+        info: &CopyIntoLocationInfo,
     ) -> Result<(PipelineBuildResult, Vec<UpdateStreamMetaReq>)> {
         let (query_interpreter, update_stream_meta_req) = self.build_query(query).await?;
         let query_physical_plan = query_interpreter.build_physical_plan().await?;
@@ -98,25 +93,9 @@ impl CopyIntoLocationInterpreter {
             plan_id: 0,
             input: Box::new(query_physical_plan),
             project_columns: query_interpreter.get_result_columns(),
-            input_schema: query_result_schema,
-            to_stage_info: StageTableInfo {
-                schema: table_schema,
-                stage_info: stage.clone(),
-                files_info: StageFilesInfo {
-                    path: path.to_string(),
-                    files: None,
-                    pattern: None,
-                },
-                files_to_copy: None,
-                duplicated_files_detected: vec![],
-                is_select: false,
-                default_exprs: None,
-                copy_into_location_options: options.clone(),
-                copy_into_table_options: Default::default(),
-                stage_root: "".to_string(),
-                copy_into_location_ordered: self.plan.is_ordered,
-                is_variant: false,
-            },
+            input_data_schema: query_result_schema,
+            input_table_schema: table_schema,
+            info: info.clone(),
         }));
 
         let mut next_plan_id = 0;
@@ -148,12 +127,7 @@ impl Interpreter for CopyIntoLocationInterpreter {
         }
 
         let (mut pipeline_build_result, update_stream_reqs) = self
-            .build_local_copy_into_stage_pipeline(
-                &self.plan.stage,
-                &self.plan.path,
-                &self.plan.from,
-                &self.plan.options,
-            )
+            .build_local_copy_into_stage_pipeline(&self.plan.from, &self.plan.info)
             .await?;
 
         // We are going to consuming streams, which are all of the default catalog

src/query/service/src/interpreters/interpreter_copy_into_location.rs

@@ -15,7 +15,7 @@
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_sql::executor::physical_plans::CopyIntoLocation;
-use databend_common_storages_stage::StageTable;
+use databend_common_storages_stage::StageSinkTable;
 
 use crate::pipelines::PipelineBuilder;
 
@@ -32,12 +32,12 @@ impl PipelineBuilder {
             false,
         )?;
 
-        let to_table = StageTable::try_create(copy.to_stage_info.clone())?;
+        let to_table = StageSinkTable::create(copy.info.clone(), copy.input_table_schema.clone())?;
         PipelineBuilder::build_append2table_with_commit_pipeline(
             self.ctx.clone(),
             &mut self.main_pipeline,
             to_table,
-            copy.input_schema.clone(),
+            copy.input_data_schema.clone(),
             None,
             vec![],
             false,

src/query/service/src/pipelines/builders/builder_copy_into_location.rs

@@ -1701,10 +1701,8 @@ impl TableContext for QueryContext {
                         duplicated_files_detected: vec![],
                         is_select: true,
                         default_exprs: None,
-                        copy_into_location_options: Default::default(),
                         copy_into_table_options: Default::default(),
                         stage_root,
-                        copy_into_location_ordered: false,
                         is_variant: true,
                     };
                     StageTable::try_create(info)
@@ -1720,10 +1718,8 @@ impl TableContext for QueryContext {
                     duplicated_files_detected: vec![],
                     is_select: true,
                     default_exprs: None,
-                    copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
                     stage_root,
-                    copy_into_location_ordered: false,
                     is_variant: false,
                 };
                 OrcTable::try_create(info).await
@@ -1741,10 +1737,8 @@ impl TableContext for QueryContext {
                     duplicated_files_detected: vec![],
                     is_select: true,
                     default_exprs: None,
-                    copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
                     stage_root,
-                    copy_into_location_ordered: false,
                     is_variant: true,
                 };
                 StageTable::try_create(info)
@@ -1780,10 +1774,8 @@ impl TableContext for QueryContext {
                     duplicated_files_detected: vec![],
                     is_select: true,
                     default_exprs: None,
-                    copy_into_location_options: Default::default(),
                     copy_into_table_options: Default::default(),
                     stage_root,
-                    copy_into_location_ordered: false,
                     is_variant: false,
                 };
                 StageTable::try_create(info)

src/query/service/src/sessions/query_ctx.rs

@@ -34,6 +34,7 @@ databend-common-users = { workspace = true }
 databend-common-version = { workspace = true }
 databend-enterprise-data-mask-feature = { workspace = true }
 databend-storages-common-cache = { workspace = true }
+databend-storages-common-stage = { workspace = true }
 databend-storages-common-table-meta = { workspace = true }
 
 ahash = { workspace = true, features = ["no-rng"] }

src/query/sql/Cargo.toml

@@ -425,8 +425,9 @@ pub trait PhysicalPlanReplacer {
             plan_id: plan.plan_id,
             input: Box::new(input),
             project_columns: plan.project_columns.clone(),
-            input_schema: plan.input_schema.clone(),
-            to_stage_info: plan.to_stage_info.clone(),
+            input_data_schema: plan.input_data_schema.clone(),
+            input_table_schema: plan.input_table_schema.clone(),
+            info: plan.info.clone(),
         })))
     }
 

src/query/sql/src/executor/physical_plan_visitor.rs

@@ -12,13 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use databend_common_catalog::plan::StageTableInfo;
 use databend_common_exception::Result;
 use databend_common_expression::types::DataType;
 use databend_common_expression::types::NumberDataType;
 use databend_common_expression::DataField;
 use databend_common_expression::DataSchemaRef;
 use databend_common_expression::DataSchemaRefExt;
+use databend_common_expression::TableSchemaRef;
+use databend_storages_common_stage::CopyIntoLocationInfo;
 
 use crate::executor::PhysicalPlan;
 use crate::ColumnBinding;
@@ -28,8 +29,9 @@ pub struct CopyIntoLocation {
     pub plan_id: u32,
     pub input: Box<PhysicalPlan>,
     pub project_columns: Vec<ColumnBinding>,
-    pub input_schema: DataSchemaRef,
-    pub to_stage_info: StageTableInfo,
+    pub input_data_schema: DataSchemaRef,
+    pub input_table_schema: TableSchemaRef,
+    pub info: CopyIntoLocationInfo,
 }
 
 impl CopyIntoLocation {