sql: parse EXECUTE UNIT TEST

src/adapter/src/coord/command_handler.rs

@@ -1307,7 +1307,8 @@ impl Coordinator {
                     | Statement::RevokeRole(_)
                     | Statement::Update(_)
                     | Statement::ValidateConnection(_)
-                    | Statement::Comment(_) => {
+                    | Statement::Comment(_)
+                    | Statement::ExecuteUnitTest(_) => {
                         let txn_status = ctx.session_mut().transaction_mut();
 
                         // If we're not in an implicit transaction and we could generate exactly one

src/sql-lexer/src/keywords.txt

@@ -289,6 +289,7 @@ Message
 Metadata
 Minute
 Minutes
+Mock
 Mode
 Month
 Months
@@ -472,6 +473,7 @@ Task
 Tasks
 Temp
 Temporary
+Test
 Text
 Then
 Tick
@@ -499,6 +501,7 @@ Unbounded
 Uncommitted
 Union
 Unique
+Unit
 Unknown
 Unnest
 Until

src/sql-parser/src/ast/defs/statement.rs

@@ -105,6 +105,7 @@ pub enum Statement<T: AstInfo> {
     Close(CloseStatement),
     Prepare(PrepareStatement<T>),
     Execute(ExecuteStatement<T>),
+    ExecuteUnitTest(ExecuteUnitTestStatement<T>),
     Deallocate(DeallocateStatement),
     Raise(RaiseStatement),
     GrantRole(GrantRoleStatement<T>),
@@ -184,6 +185,7 @@ impl<T: AstInfo> AstDisplay for Statement<T> {
             Statement::Fetch(stmt) => f.write_node(stmt),
             Statement::Prepare(stmt) => f.write_node(stmt),
             Statement::Execute(stmt) => f.write_node(stmt),
+            Statement::ExecuteUnitTest(stmt) => f.write_node(stmt),
             Statement::Deallocate(stmt) => f.write_node(stmt),
             Statement::Raise(stmt) => f.write_node(stmt),
             Statement::GrantRole(stmt) => f.write_node(stmt),
@@ -268,6 +270,7 @@ pub fn statement_kind_label_value(kind: StatementKind) -> &'static str {
         StatementKind::Close => "close",
         StatementKind::Prepare => "prepare",
         StatementKind::Execute => "execute",
+        StatementKind::ExecuteUnitTest => "execute_unit_test",
         StatementKind::Deallocate => "deallocate",
         StatementKind::Raise => "raise",
         StatementKind::GrantRole => "grant_role",
@@ -5092,6 +5095,74 @@ impl<T: AstInfo> AstDisplay for ExecuteStatement<T> {
 }
 impl_display_t!(ExecuteStatement);
 
+/// `EXECUTE UNIT TEST ...`
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct ExecuteUnitTestStatement<T: AstInfo> {
+    pub name: Ident,
+    pub target: T::ItemName,
+    pub at_time: Option<Expr<T>>,
+    pub mocks: Vec<MockViewDef<T>>,
+    pub expected: ExpectedResultDef<T>,
+}
+
+impl<T: AstInfo> AstDisplay for ExecuteUnitTestStatement<T> {
+    fn fmt<W: fmt::Write>(&self, f: &mut AstFormatter<W>) {
+        f.write_str("EXECUTE UNIT TEST ");
+        f.write_node(&self.name);
+        f.write_str(" FOR ");
+        f.write_node(&self.target);
+        if let Some(at_time) = &self.at_time {
+            f.write_str(" AT TIME ");
+            f.write_node(at_time);
+        }
+        for (i, mock) in self.mocks.iter().enumerate() {
+            f.write_str(if i == 0 { " MOCK " } else { ", MOCK " });
+            f.write_node(mock);
+        }
+        f.write_str(" EXPECTED ");
+        f.write_node(&self.expected);
+    }
+}
+impl_display_t!(ExecuteUnitTestStatement);
+
+/// Mock view definition for EXECUTE UNIT TEST
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct MockViewDef<T: AstInfo> {
+    pub name: T::ItemName,
+    pub columns: Vec<ColumnDef<T>>,
+    pub query: Query<T>,
+}
+
+impl<T: AstInfo> AstDisplay for MockViewDef<T> {
+    fn fmt<W: fmt::Write>(&self, f: &mut AstFormatter<W>) {
+        f.write_node(&self.name);
+        f.write_str("(");
+        f.write_node(&display::comma_separated(&self.columns));
+        f.write_str(") AS (");
+        f.write_node(&self.query);
+        f.write_str(")");
+    }
+}
+impl_display_t!(MockViewDef);
+
+/// Expected result definition for EXECUTE UNIT TEST
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct ExpectedResultDef<T: AstInfo> {
+    pub columns: Vec<ColumnDef<T>>,
+    pub query: Query<T>,
+}
+
+impl<T: AstInfo> AstDisplay for ExpectedResultDef<T> {
+    fn fmt<W: fmt::Write>(&self, f: &mut AstFormatter<W>) {
+        f.write_str("(");
+        f.write_node(&display::comma_separated(&self.columns));
+        f.write_str(") AS (");
+        f.write_node(&self.query);
+        f.write_str(")");
+    }
+}
+impl_display_t!(ExpectedResultDef);
+
 /// `DEALLOCATE ...`
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct DeallocateStatement {

src/sql-parser/src/parser.rs

@@ -9396,6 +9396,12 @@ impl<'a> Parser<'a> {
     /// Parse a `EXECUTE` statement, assuming that the `EXECUTE` token
     /// has already been consumed.
     fn parse_execute(&mut self) -> Result<Statement<Raw>, ParserError> {
+        // Check if this is EXECUTE UNIT TEST
+        if self.parse_keywords(&[UNIT, TEST]) {
+            return self.parse_execute_unit_test();
+        }
+
+        // Otherwise parse as regular EXECUTE (prepared statement)
         let name = self.parse_identifier()?;
         let params = if self.consume_token(&Token::LParen) {
             let params = self.parse_comma_separated(Parser::parse_expr)?;
@@ -9407,6 +9413,97 @@ impl<'a> Parser<'a> {
         Ok(Statement::Execute(ExecuteStatement { name, params }))
     }
 
+    /// Parse an `EXECUTE UNIT TEST` statement, assuming that the
+    /// `EXECUTE UNIT TEST` tokens have already been consumed.
+    fn parse_execute_unit_test(&mut self) -> Result<Statement<Raw>, ParserError> {
+        // Parse test name
+        let name = self.parse_identifier()?;
+
+        // Expect FOR keyword
+        self.expect_keyword(FOR)?;
+
+        // Parse target view name
+        let target = self.parse_raw_name()?;
+
+        // Parse optional AT TIME clause
+        let at_time = if self.parse_keywords(&[AT, TIME]) {
+            Some(self.parse_expr()?)
+        } else {
+            None
+        };
+
+        // Parse MOCK definitions (0 or more). Like CTEs, MOCK clauses are
+        // comma-separated; no comma appears before the trailing `EXPECTED`
+        // clause.
+        let mut mocks = Vec::new();
+        if self.parse_keyword(MOCK) {
+            mocks.push(self.parse_mock_view_def()?);
+            while self.consume_token(&Token::Comma) {
+                self.expect_keyword(MOCK)?;
+                mocks.push(self.parse_mock_view_def()?);
+            }
+        }
+
+        self.expect_keyword(EXPECTED)?;
+
+        self.expect_token(&Token::LParen)?;
+        let expected_columns = self.parse_comma_separated(|parser| {
+            Ok(ColumnDef {
+                name: parser.parse_identifier()?,
+                data_type: parser.parse_data_type()?,
+                collation: None,
+                options: vec![],
+            })
+        })?;
+        self.expect_token(&Token::RParen)?;
+
+        self.expect_keyword(AS)?;
+        self.expect_token(&Token::LParen)?;
+        let expected_query = self.parse_query()?;
+        self.expect_token(&Token::RParen)?;
+
+        let expected = ExpectedResultDef {
+            columns: expected_columns,
+            query: expected_query,
+        };
+
+        Ok(Statement::ExecuteUnitTest(ExecuteUnitTestStatement {
+            name,
+            target,
+            at_time,
+            mocks,
+            expected,
+        }))
+    }
+
+    /// Parse a single `MOCK <name> (<cols>) AS (<query>)` definition,
+    /// assuming the leading `MOCK` keyword has already been consumed.
+    fn parse_mock_view_def(&mut self) -> Result<MockViewDef<Raw>, ParserError> {
+        let name = self.parse_raw_name()?;
+
+        self.expect_token(&Token::LParen)?;
+        let columns = self.parse_comma_separated(|parser| {
+            Ok(ColumnDef {
+                name: parser.parse_identifier()?,
+                data_type: parser.parse_data_type()?,
+                collation: None,
+                options: vec![],
+            })
+        })?;
+        self.expect_token(&Token::RParen)?;
+
+        self.expect_keyword(AS)?;
+        self.expect_token(&Token::LParen)?;
+        let query = self.parse_query()?;
+        self.expect_token(&Token::RParen)?;
+
+        Ok(MockViewDef {
+            name,
+            columns,
+            query,
+        })
+    }
+
     /// Parse a `DEALLOCATE` statement, assuming that the `DEALLOCATE` token
     /// has already been consumed.
     fn parse_deallocate(&mut self) -> Result<Statement<Raw>, ParserError> {