Merge pull request #328 from kitta65/feature/undelete

support undrop statement
kitta65 · Mar 22, 2024 · 7ccbe10 · 7ccbe10
2 parents 34b1a6b + 4b592e2
commit 7ccbe10
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 0 deletions.
diff --git a/src/cst.rs b/src/cst.rs
@@ -128,6 +128,7 @@ pub enum NodeType {
     TypeDeclaration, // x INT64
     UnaryOperator,   // - | + | TIMESTAMP | ...
     Unknown,
+    UndropStatement,
     UnpivotOperator,
     UnpivotConfig, // ((c1, c2) FOR v IN ((v1, v2) 1, (v3, v4) 3))
     UpdateStatement,

diff --git a/src/parser.rs b/src/parser.rs
@@ -1011,6 +1011,7 @@ impl Parser {
                     self.parse_drop_statement_general(semicolon)?
                 }
             }
+            "UNDROP" => self.parse_undrop_statement(semicolon)?,
             // DCL
             "GRANT" => self.parse_grant_statement(semicolon)?,
             "REVOKE" => self.parse_revoke_statement(semicolon)?,
@@ -2984,6 +2985,22 @@ impl Parser {
         }
         Ok(drop)
     }
+    fn parse_undrop_statement(&mut self, semicolon: bool) -> BQ2CSTResult<Node> {
+        let mut undrop = self.construct_node(NodeType::UndropStatement)?;
+        self.next_token()?; // -> SCHEMA
+        undrop.push_node("what", self.construct_node(NodeType::Keyword)?);
+        if self.get_token(1)?.is("IF") {
+            self.next_token()?; // -> IF
+            undrop.push_node_vec("if_not_exists", self.parse_n_keywords(3)?);
+        }
+        self.next_token()?; // -> ident
+        undrop.push_node("ident", self.parse_identifier()?);
+        if self.get_token(1)?.is(";") && semicolon {
+            self.next_token()?; // -> ;
+            undrop.push_node("semicolon", self.construct_node(NodeType::Symbol)?);
+        }
+        Ok(undrop)
+    }
     // ----- DCL -----
     fn parse_grant_statement(&mut self, semicolon: bool) -> BQ2CSTResult<Node> {
         let mut grant = self.construct_node(NodeType::GrantStatement)?;

diff --git a/src/parser/tests/tests_ddl.rs b/src/parser/tests/tests_ddl.rs
@@ -3196,6 +3196,43 @@ if_exists:
 - self: EXISTS (Keyword)
 what:
   self: ASSIGNMENT (Keyword)
+",
+            0,
+        )),
+        // ----- UNDROP statement -----
+        Box::new(SuccessTestCase::new(
+            "\
+UNDROP SCHEMA datasetname;
+",
+            "\
+self: UNDROP (UndropStatement)
+ident:
+  self: datasetname (Identifier)
+semicolon:
+  self: ; (Symbol)
+what:
+  self: SCHEMA (Keyword)
+",
+            0,
+        )),
+        Box::new(SuccessTestCase::new(
+            "\
+UNDROP SCHEMA IF NOT EXISTS projectname.datasetname
+",
+            "\
+self: UNDROP (UndropStatement)
+ident:
+  self: . (DotOperator)
+  left:
+    self: projectname (Identifier)
+  right:
+    self: datasetname (Identifier)
+if_not_exists:
+- self: IF (Keyword)
+- self: NOT (Keyword)
+- self: EXISTS (Keyword)
+what:
+  self: SCHEMA (Keyword)
 ",
             0,
         )),

diff --git a/src/types.rs b/src/types.rs
@@ -119,6 +119,7 @@ export type UnknownNode =
   | Type
   | TypeDeclaration
   | UnaryOperator
+  | UndropStatement
   | UnpivotConfig
   | UnpivotOperator
   | UpdateStatement
@@ -1350,6 +1351,15 @@ export type UnaryOperator = Expr & {
   };
 };
 
+export type UndropStatement = XXXStatement & {
+  node_type: "UndropStatement";
+  children: {
+    what: NodeChild;
+    if_not_exists?: NodeVecChild;
+    ident: NodeChild;
+  };
+};
+
 export type UnpivotConfig = BaseNode & {
   token: Token;
   node_type: "UnpivotConfig";