Add support for query source in COPY .. TO statement (#858)

* Add support for query source in COPY .. TO statement

* Fix compile error

src/ast/mod.rs

@@ -1173,14 +1173,11 @@ pub enum Statement {
         source: Box<Query>,
     },
     Copy {
-        /// TABLE
-        #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
-        table_name: ObjectName,
-        /// COLUMNS
-        columns: Vec<Ident>,
+        /// The source of 'COPY TO', or the target of 'COPY FROM'
+        source: CopySource,
         /// If true, is a 'COPY TO' statement. If false is a 'COPY FROM'
         to: bool,
-        /// The source of 'COPY FROM', or the target of 'COPY TO'
+        /// The target of 'COPY TO', or the source of 'COPY FROM'
         target: CopyTarget,
         /// WITH options (from PostgreSQL version 9.0)
         options: Vec<CopyOption>,
@@ -1926,17 +1923,25 @@ impl fmt::Display for Statement {
             }
 
             Statement::Copy {
-                table_name,
-                columns,
+                source,
                 to,
                 target,
                 options,
                 legacy_options,
                 values,
             } => {
-                write!(f, "COPY {table_name}")?;
-                if !columns.is_empty() {
-                    write!(f, " ({})", display_comma_separated(columns))?;
+                write!(f, "COPY")?;
+                match source {
+                    CopySource::Query(query) => write!(f, " ({query})")?,
+                    CopySource::Table {
+                        table_name,
+                        columns,
+                    } => {
+                        write!(f, " {table_name}")?;
+                        if !columns.is_empty() {
+                            write!(f, " ({})", display_comma_separated(columns))?;
+                        }
+                    }
                 }
                 write!(f, " {} {}", if *to { "TO" } else { "FROM" }, target)?;
                 if !options.is_empty() {
@@ -3750,6 +3755,20 @@ impl fmt::Display for SqliteOnConflict {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum CopySource {
+    Table {
+        /// The name of the table to copy from.
+        table_name: ObjectName,
+        /// A list of column names to copy. Empty list means that all columns
+        /// are copied.
+        columns: Vec<Ident>,
+    },
+    Query(Box<Query>),
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/parser.rs

@@ -4015,13 +4015,32 @@ impl<'a> Parser<'a> {
 
     /// Parse a copy statement
     pub fn parse_copy(&mut self) -> Result<Statement, ParserError> {
-        let table_name = self.parse_object_name()?;
-        let columns = self.parse_parenthesized_column_list(Optional, false)?;
+        let source;
+        if self.consume_token(&Token::LParen) {
+            source = CopySource::Query(Box::new(self.parse_query()?));
+            self.expect_token(&Token::RParen)?;
+        } else {
+            let table_name = self.parse_object_name()?;
+            let columns = self.parse_parenthesized_column_list(Optional, false)?;
+            source = CopySource::Table {
+                table_name,
+                columns,
+            };
+        }
         let to = match self.parse_one_of_keywords(&[Keyword::FROM, Keyword::TO]) {
             Some(Keyword::FROM) => false,
             Some(Keyword::TO) => true,
             _ => self.expected("FROM or TO", self.peek_token())?,
         };
+        if !to {
+            // Use a separate if statement to prevent Rust compiler from complaining about
+            // "if statement in this position is unstable: https://github.com/rust-lang/rust/issues/53667"
+            if let CopySource::Query(_) = source {
+                return Err(ParserError::ParserError(
+                    "COPY ... FROM does not support query as a source".to_string(),
+                ));
+            }
+        }
         let target = if self.parse_keyword(Keyword::STDIN) {
             CopyTarget::Stdin
         } else if self.parse_keyword(Keyword::STDOUT) {
@@ -4052,8 +4071,7 @@ impl<'a> Parser<'a> {
             vec![]
         };
         Ok(Statement::Copy {
-            table_name,
-            columns,
+            source,
             to,
             target,
             options,