Fix #90

Author: kitta65

src/parser.rs

@@ -2907,7 +2907,7 @@ impl Parser {
         let mut load = self.construct_node(NodeType::LoadStatement)?;
         self.next_token()?; // -> DATA
         load.push_node("data", self.construct_node(NodeType::Keyword)?);
-        self.next_token()?; // -> INTO
+        self.next_token()?; // -> INTO | OVERWRITE
         load.push_node("into", self.construct_node(NodeType::Keyword)?);
         self.next_token()?; // -> ident
         load.push_node("ident", self.parse_identifier()?);
@@ -2935,12 +2935,28 @@ impl Parser {
         print!("{:?}", self.construct_node(NodeType::Keyword)?);
         load.push_node("from_files", self.parse_grouped_exprs(false)?);
         print!("{:?}", self.construct_node(NodeType::Keyword)?);
-        self.next_token()?; // -> WITH
-        load.push_node("with", self.construct_node(NodeType::Keyword)?);
-        self.next_token()?; // -> CONNECTION
-        load.push_node("connection", self.construct_node(NodeType::Keyword)?);
-        self.next_token()?; // -> connection_name
-        load.push_node("connection_name", self.parse_identifier()?);
+        if self.get_token(1)?.is("WITH") && self.get_token(2)?.is("PARTITION") {
+            self.next_token()?; // -> WITH
+            let mut with = self.construct_node(NodeType::WithPartitionColumnsClause)?;
+            self.next_token()?; // -> PARTITION
+            with.push_node_vec("partition_columns", self.parse_n_keywords(2)?);
+            if self.get_token(1)?.is("(") {
+                self.next_token()?; // -> (
+                with.push_node(
+                    "column_schema_group",
+                    self.parse_grouped_type_declarations(false)?,
+                );
+            }
+            load.push_node("with_partition_columns", with);
+        }
+        if self.get_token(1)?.is("WITH") && self.get_token(2)?.is("CONNECTION") {
+            self.next_token()?; // -> WITH
+            load.push_node("with", self.construct_node(NodeType::Keyword)?);
+            self.next_token()?; // -> CONNECTION
+            load.push_node("connection", self.construct_node(NodeType::Keyword)?);
+            self.next_token()?; // -> connection_name
+            load.push_node("connection_name", self.parse_identifier()?);
+        }
 
         if self.get_token(1)?.is(";") && semicolon {
             self.next_token()?; // -> ;

src/parser/tests/tests_other.rs

@@ -171,6 +171,50 @@ partitionby:
   - self: dt (Identifier)
 with:
   self: WITH (Keyword)
+",
+            0,
+        )),
+        Box::new(SuccessTestCase::new(
+            "\
+LOAD DATA OVERWRITE ident
+FROM FILES (dummy = 'dummy')
+WITH PARTITION COLUMNS (x STRING)
+",
+            "\
+self: LOAD (LoadStatement)
+data:
+  self: DATA (Keyword)
+files:
+  self: FILES (Keyword)
+from:
+  self: FROM (Keyword)
+from_files:
+  self: ( (GroupedExprs)
+  exprs:
+  - self: = (BinaryOperator)
+    left:
+      self: dummy (Identifier)
+    right:
+      self: 'dummy' (StringLiteral)
+  rparen:
+    self: ) (Symbol)
+ident:
+  self: ident (Identifier)
+into:
+  self: OVERWRITE (Keyword)
+with_partition_columns:
+  self: WITH (WithPartitionColumnsClause)
+  column_schema_group:
+    self: ( (GroupedTypeDeclarations)
+    declarations:
+    - self: x (TypeDeclaration)
+      type:
+        self: STRING (Type)
+    rparen:
+      self: ) (Symbol)
+  partition_columns:
+  - self: PARTITION (Keyword)
+  - self: COLUMNS (Keyword)
 ",
             0,
         )),

src/types.rs

@@ -921,9 +921,10 @@ export type LoadStatement = XXXStatement & {
     from: NodeChild;
     files: NodeChild;
     from_files: NodeChild;
-    with: NodeChild;
-    connection: NodeChild;
-    connection_name: NodeChild;
+    with_partition_columns?: NodeChild;
+    with?: NodeChild;
+    connection?: NodeChild;
+    connection_name?: NodeChild;
   };
 };