feat: Support KILL statement (#479)

src/ast/mod.rs

@@ -1005,6 +1005,13 @@ pub enum Statement {
         data_types: Vec<DataType>,
         statement: Box<Statement>,
     },
+    /// See <https://clickhouse.com/docs/ru/sql-reference/statements/kill/>
+    /// See <https://dev.mysql.com/doc/refman/8.0/en/kill.html>
+    Kill {
+        modifier: Option<KillType>,
+        // processlist_id
+        id: u64,
+    },
     /// EXPLAIN TABLE
     /// Note: this is a MySQL-specific statement. See <https://dev.mysql.com/doc/refman/8.0/en/explain.html>
     ExplainTable {
@@ -1047,6 +1054,15 @@ impl fmt::Display for Statement {
     #[allow(clippy::cognitive_complexity)]
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
+            Statement::Kill { modifier, id } => {
+                write!(f, "KILL ")?;
+
+                if let Some(m) = modifier {
+                    write!(f, "{} ", m)?;
+                }
+
+                write!(f, "{}", id)
+            }
             Statement::ExplainTable {
                 describe_alias,
                 table_name,
@@ -2097,6 +2113,26 @@ impl fmt::Display for ObjectType {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum KillType {
+    Connection,
+    Query,
+    Mutation,
+}
+
+impl fmt::Display for KillType {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            // MySQL
+            KillType::Connection => "CONNECTION",
+            KillType::Query => "QUERY",
+            // Clickhouse supports Mutation
+            KillType::Mutation => "MUTATION",
+        })
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub enum HiveDistributionStyle {

src/keywords.rs

@@ -136,6 +136,7 @@ define_keywords!(
     COMPUTE,
     CONDITION,
     CONNECT,
+    CONNECTION,
     CONSTRAINT,
     CONTAINS,
     CONVERT,
@@ -279,6 +280,7 @@ define_keywords!(
     JSONFILE,
     JULIAN,
     KEY,
+    KILL,
     LAG,
     LANGUAGE,
     LARGE,
@@ -319,6 +321,7 @@ define_keywords!(
     MONTH,
     MSCK,
     MULTISET,
+    MUTATION,
     NATIONAL,
     NATURAL,
     NCHAR,
@@ -384,6 +387,7 @@ define_keywords!(
     PURGE,
     QUALIFY,
     QUARTER,
+    QUERY,
     QUOTE,
     RANGE,
     RANK,

src/parser.rs

@@ -154,6 +154,7 @@ impl<'a> Parser<'a> {
     pub fn parse_statement(&mut self) -> Result<Statement, ParserError> {
         match self.next_token() {
             Token::Word(w) => match w.keyword {
+                Keyword::KILL => Ok(self.parse_kill()?),
                 Keyword::DESCRIBE => Ok(self.parse_explain(true)?),
                 Keyword::EXPLAIN => Ok(self.parse_explain(false)?),
                 Keyword::ANALYZE => Ok(self.parse_analyze()?),
@@ -2878,6 +2879,32 @@ impl<'a> Parser<'a> {
         })
     }
 
+    // KILL [CONNECTION | QUERY] processlist_id
+    pub fn parse_kill(&mut self) -> Result<Statement, ParserError> {
+        let modifier_keyword =
+            self.parse_one_of_keywords(&[Keyword::CONNECTION, Keyword::QUERY, Keyword::MUTATION]);
+
+        let id = self.parse_literal_uint()?;
+
+        let modifier = match modifier_keyword {
+            Some(Keyword::CONNECTION) => Some(KillType::Connection),
+            Some(Keyword::QUERY) => Some(KillType::Query),
+            Some(Keyword::MUTATION) => {
+                if dialect_of!(self is ClickHouseDialect | GenericDialect) {
+                    Some(KillType::Mutation)
+                } else {
+                    self.expected(
+                        "Unsupported type for KILL, allowed: CONNECTION | QUERY",
+                        self.peek_token(),
+                    )?
+                }
+            }
+            _ => None,
+        };
+
+        Ok(Statement::Kill { modifier, id })
+    }
+
     pub fn parse_explain(&mut self, describe_alias: bool) -> Result<Statement, ParserError> {
         let analyze = self.parse_keyword(Keyword::ANALYZE);
         let verbose = self.parse_keyword(Keyword::VERBOSE);

tests/sqlparser_mysql.rs

@@ -770,6 +770,36 @@ fn parse_substring_in_select() {
     }
 }
 
+#[test]
+fn parse_kill() {
+    let stmt = mysql_and_generic().verified_stmt("KILL CONNECTION 5");
+    assert_eq!(
+        stmt,
+        Statement::Kill {
+            modifier: Some(KillType::Connection),
+            id: 5,
+        }
+    );
+
+    let stmt = mysql_and_generic().verified_stmt("KILL QUERY 5");
+    assert_eq!(
+        stmt,
+        Statement::Kill {
+            modifier: Some(KillType::Query),
+            id: 5,
+        }
+    );
+
+    let stmt = mysql_and_generic().verified_stmt("KILL 5");
+    assert_eq!(
+        stmt,
+        Statement::Kill {
+            modifier: None,
+            id: 5,
+        }
+    );
+}
+
 fn mysql() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(MySqlDialect {})],

tests/sqpparser_clickhouse.rs

@@ -119,6 +119,18 @@ fn parse_array_expr() {
     )
 }
 
+#[test]
+fn parse_kill() {
+    let stmt = clickhouse().verified_stmt("KILL MUTATION 5");
+    assert_eq!(
+        stmt,
+        Statement::Kill {
+            modifier: Some(KillType::Mutation),
+            id: 5,
+        }
+    );
+}
+
 fn clickhouse() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(ClickHouseDialect {})],