diff --git a/src/binder/copy.rs b/src/binder/copy.rs index 5ea67653..c690c64d 100644 --- a/src/binder/copy.rs +++ b/src/binder/copy.rs @@ -86,30 +86,34 @@ impl> Binder<'_, '_, T, A> target: CopyTarget, options: &[CopyOption], ) -> Result { + let ext_source = copy_ext_source(target, options)?; + let (table_name, ..) = match source { CopySource::Table { table_name, columns, } => (table_name, columns), - CopySource::Query(_) => { - return Err(DatabaseError::UnsupportedStmt("'COPY SOURCE'".to_string())); + CopySource::Query(query) => { + if !to { + return Err(DatabaseError::UnsupportedStmt( + "'COPY FROM query'".to_string(), + )); + } + let mut input_plan = self.bind_query(&query)?; + let schema_ref = input_plan.output_schema().clone(); + return Ok(LogicalPlan::new( + Operator::CopyToFile(CopyToFileOperator { + target: ext_source, + schema_ref, + }), + Childrens::Only(Box::new(input_plan)), + )); } }; let table_name: Arc = lower_case_name(&table_name)?.into(); if let Some(table) = self.context.table(table_name.clone())? { let schema_ref = table.schema_ref().clone(); - let ext_source = ExtSource { - path: match target { - CopyTarget::File { filename } => filename.into(), - t => { - return Err(DatabaseError::UnsupportedStmt(format!( - "copy target: {t:?}" - ))) - } - }, - format: FileFormat::from_options(options), - }; if to { // COPY TO @@ -139,6 +143,20 @@ impl> Binder<'_, '_, T, A> } } +fn copy_ext_source(target: CopyTarget, options: &[CopyOption]) -> Result { + Ok(ExtSource { + path: match target { + CopyTarget::File { filename } => filename.into(), + t => { + return Err(DatabaseError::UnsupportedStmt(format!( + "copy target: {t:?}" + ))) + } + }, + format: FileFormat::from_options(options), + }) +} + impl FileFormat { /// Create from copy options. pub fn from_options(options: &[CopyOption]) -> Self { diff --git a/tests/slt/copy.slt b/tests/slt/copy.slt index c5985203..96ec43d9 100644 --- a/tests/slt/copy.slt +++ b/tests/slt/copy.slt @@ -16,4 +16,22 @@ SELECT * FROM test_copy query I COPY test_copy TO './copy.csv' ( DELIMITER ',' ); ---- -Copy To ./copy.csv [a, b, c] \ No newline at end of file +Copy To ./copy.csv [a, b, c] + +query I +COPY (SELECT a, c FROM test_copy WHERE a = 1) TO './copy_query.csv' ( DELIMITER ',', HEADER true ); +---- +Copy To ./copy_query.csv [a, c] + +statement ok +create table test_copy_query (a int primary key, c varchar(10)) + +query I +COPY test_copy_query FROM './copy_query.csv' ( DELIMITER ',', HEADER true ); +---- +1 + +query I +SELECT * FROM test_copy_query +---- +1 two