Support COPY query source to file (#343)

Author: KKould

src/binder/copy.rs

@@ -86,30 +86,34 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
         target: CopyTarget,
         options: &[CopyOption],
     ) -> Result<LogicalPlan, DatabaseError> {
+        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<str> = 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 <source_table> TO <dest_file>
@@ -139,6 +143,20 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
     }
 }
 
+fn copy_ext_source(target: CopyTarget, options: &[CopyOption]) -> Result<ExtSource, DatabaseError> {
+    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 {

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]
+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