feat: limit total rows copied in COPY TABLE FROM with LIMIT segment (#3910)

* feat: limit total rows copied in COPY TABLE FROM with LIMIT segment

* fmt

* disable default limit

* fix: check parse

* fix test, add error case

* fix: forbide LIMIT in database

* fix: only support LIMIT segment

* fix: simplify

* fix

* fix

* fix

* fix

* fix: test

* fix: change error info

* fix clippy

* fix: fix error msg

* fix test

* fix: test error info

src/sql/src/parsers/copy_parser.rs

@@ -56,7 +56,14 @@ impl<'a> ParserContext<'a> {
             })?;
 
         let req = if self.parser.parse_keyword(Keyword::TO) {
-            let (with, connection, location) = self.parse_copy_parameters()?;
+            let (with, connection, location, limit) = self.parse_copy_parameters()?;
+            if limit.is_some() {
+                return error::InvalidSqlSnafu {
+                    msg: "limit is not supported",
+                }
+                .fail();
+            }
+
             let argument = CopyDatabaseArgument {
                 database_name,
                 with: with.into(),
@@ -68,7 +75,14 @@ impl<'a> ParserContext<'a> {
             self.parser
                 .expect_keyword(Keyword::FROM)
                 .context(error::SyntaxSnafu)?;
-            let (with, connection, location) = self.parse_copy_parameters()?;
+            let (with, connection, location, limit) = self.parse_copy_parameters()?;
+            if limit.is_some() {
+                return error::InvalidSqlSnafu {
+                    msg: "limit is not supported",
+                }
+                .fail();
+            }
+
             let argument = CopyDatabaseArgument {
                 database_name,
                 with: with.into(),
@@ -91,28 +105,30 @@ impl<'a> ParserContext<'a> {
         let table_name = Self::canonicalize_object_name(raw_table_name);
 
         if self.parser.parse_keyword(Keyword::TO) {
-            let (with, connection, location) = self.parse_copy_parameters()?;
+            let (with, connection, location, limit) = self.parse_copy_parameters()?;
             Ok(CopyTable::To(CopyTableArgument {
                 table_name,
                 with: with.into(),
                 connection: connection.into(),
                 location,
+                limit,
             }))
         } else {
             self.parser
                 .expect_keyword(Keyword::FROM)
                 .context(error::SyntaxSnafu)?;
-            let (with, connection, location) = self.parse_copy_parameters()?;
+            let (with, connection, location, limit) = self.parse_copy_parameters()?;
             Ok(CopyTable::From(CopyTableArgument {
                 table_name,
                 with: with.into(),
                 connection: connection.into(),
                 location,
+                limit,
             }))
         }
     }
 
-    fn parse_copy_parameters(&mut self) -> Result<(With, Connection, String)> {
+    fn parse_copy_parameters(&mut self) -> Result<(With, Connection, String, Option<u64>)> {
         let location =
             self.parser
                 .parse_literal_string()
@@ -142,7 +158,21 @@ impl<'a> ParserContext<'a> {
             .map(parse_option_string)
             .collect::<Result<Connection>>()?;
 
-        Ok((with, connection, location))
+        let limit = if self.parser.parse_keyword(Keyword::LIMIT) {
+            Some(
+                self.parser
+                    .parse_literal_uint()
+                    .with_context(|_| error::UnexpectedSnafu {
+                        sql: self.sql,
+                        expected: "the number of maximum rows",
+                        actual: self.peek_token_as_string(),
+                    })?,
+            )
+        } else {
+            None
+        };
+
+        Ok((with, connection, location, limit))
     }
 }
 

src/sql/src/statements/copy.rs

@@ -111,6 +111,7 @@ pub struct CopyTableArgument {
     pub connection: OptionMap,
     /// Copy tbl [To|From] 'location'.
     pub location: String,
+    pub limit: Option<u64>,
 }
 
 #[cfg(test)]