feat: handle Ctrl-C command in MySQL client (#6320)

* feat/answer-ctrl-c-in-mysql:
 ## Implement Connection ID-based Query Killing

 ### Key Changes:
 - **Connection ID Management:**
   - Added `connection_id` to `Session` and `QueryContext` in `src/session/src/lib.rs` and `src/session/src/context.rs`.
   - Updated `MysqlInstanceShim` and `MysqlServer` to handle `connection_id` in `src/servers/src/mysql/handler.rs` and `src/servers/src/mysql/server.rs`.

 - **KILL Statement Enhancements:**
   - Introduced `Kill` enum to handle both `ProcessId` and `ConnectionId` in `src/sql/src/statements/kill.rs`.
   - Updated `ParserContext` to parse `KILL QUERY <connection_id>` in `src/sql/src/parser.rs`.
   - Modified `StatementExecutor` to support killing queries by `connection_id` in `src/operator/src/statement/kill.rs`.

 - **Process Management:**
   - Refactored `ProcessManager` to include `connection_id` in `src/catalog/src/process_manager.rs`.
   - Added `kill_local_process` method for local query termination.

 - **Testing:**
   - Added tests for `KILL` statement parsing and execution in `src/sql/src/parser.rs`.

 ### Affected Files:
 - `Cargo.lock`, `Cargo.toml`
 - `src/catalog/src/process_manager.rs`
 - `src/frontend/src/instance.rs`
 - `src/frontend/src/stream_wrapper.rs`
 - `src/operator/src/statement.rs`
 - `src/operator/src/statement/kill.rs`
 - `src/servers/src/mysql/federated.rs`
 - `src/servers/src/mysql/handler.rs`
 - `src/servers/src/mysql/server.rs`
 - `src/servers/src/postgres.rs`
 - `src/session/src/context.rs`
 - `src/session/src/lib.rs`
 - `src/sql/src/parser.rs`
 - `src/sql/src/statements.rs`
 - `src/sql/src/statements/kill.rs`
 - `src/sql/src/statements/statement.rs`

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>


 Conflicts:

	Cargo.lock

	Cargo.toml

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* feat/answer-ctrl-c-in-mysql:
 ### Enhance Process Management and Execution

 - **`process_manager.rs`**: Added a new method `find_processes_by_connection_id` to filter processes by connection ID, improving process management capabilities.
 - **`kill.rs`**: Refactored the process killing logic to utilize the new `find_processes_by_connection_id` method, streamlining the execution flow and reducing redundant checks.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* feat/answer-ctrl-c-in-mysql:
 ## Commit Message

 ### Update Process ID Type and Refactor Code

 - **Change Process ID Type**: Updated the process ID type from `u64` to `u32` across multiple files to optimize memory usage. Affected files include `process_manager.rs`, `lib.rs`, `database.rs`, `instance.rs`, `server.rs`, `stream_wrapper.rs`, `kill.rs`, `federated.rs`, `handler.rs`, `server.rs`,
 `postgres.rs`, `mysql_server_test.rs`, `context.rs`, `lib.rs`, and `test_util.rs`.

 - **Remove Connection ID**: Removed the `connection_id` field and related logic from `process_manager.rs`, `lib.rs`, `instance.rs`, `server.rs`, `stream_wrapper.rs`, `kill.rs`, `federated.rs`, `handler.rs`, `server.rs`, `postgres.rs`, `mysql_server_test.rs`, `context.rs`, `lib.rs`, and `test_util.rs` to
 simplify the codebase.

 - **Refactor Process Management**: Refactored process management logic to improve clarity and maintainability in `process_manager.rs`, `kill.rs`, and `handler.rs`.

 - **Enhance MySQL Server Handling**: Improved MySQL server handling by integrating process management in `server.rs` and `mysql_server_test.rs`.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* feat/answer-ctrl-c-in-mysql:
 ### Add Process Manager to Postgres Server

 - **`src/frontend/src/server.rs`**: Updated server initialization to include `process_manager`.
 - **`src/servers/src/postgres.rs`**: Modified `MakePostgresServerHandler` to accept `process_id` for session creation.
 - **`src/servers/src/postgres/server.rs`**: Integrated `process_manager` into `PostgresServer` for generating `process_id` during connection handling.
 - **`src/servers/tests/postgres/mod.rs`** and **`tests-integration/src/test_util.rs`**: Adjusted test server setup to accommodate optional `process_manager`.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* feat/answer-ctrl-c-in-mysql:
 Update `greptime-proto` Dependency

 - Updated the `greptime-proto` dependency to a new revision in both `Cargo.lock` and `Cargo.toml`.
   - `Cargo.lock`: Changed source revision from `d75a56e05a87594fe31ad5c48525e9b2124149ba` to `fdcbe5f1c7c467634c90a1fd1a00a784b92a4e80`.
   - `Cargo.toml`: Updated the `greptime-proto` git revision to match the new commit.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

---------

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

src/sql/src/parser.rs

@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::str::FromStr;
+
 use snafu::ResultExt;
-use sqlparser::ast::{Ident, Query};
+use sqlparser::ast::{Ident, Query, Value};
 use sqlparser::dialect::Dialect;
 use sqlparser::keywords::Keyword;
 use sqlparser::parser::{Parser, ParserError, ParserOptions};
@@ -22,6 +24,7 @@ use sqlparser::tokenizer::{Token, TokenWithSpan};
 use crate::ast::{Expr, ObjectName};
 use crate::error::{self, Result, SyntaxSnafu};
 use crate::parsers::tql_parser;
+use crate::statements::kill::Kill;
 use crate::statements::statement::Statement;
 use crate::statements::transform_statements;
 
@@ -190,14 +193,43 @@ impl ParserContext<'_> {
 
                 Keyword::KILL => {
                     let _ = self.parser.next_token();
-                    let process_id_ident =
-                        self.parser.parse_literal_string().with_context(|_| {
-                            error::UnexpectedSnafu {
-                                expected: "process id string literal",
-                                actual: self.peek_token_as_string(),
+                    let kill = if self.parser.parse_keyword(Keyword::QUERY) {
+                        // MySQL KILL QUERY <connection id> statements
+                        let connection_id_exp =
+                            self.parser.parse_number_value().with_context(|_| {
+                                error::UnexpectedSnafu {
+                                    expected: "MySQL numeric connection id",
+                                    actual: self.peek_token_as_string(),
+                                }
+                            })?;
+                        let Value::Number(s, _) = connection_id_exp else {
+                            return error::UnexpectedTokenSnafu {
+                                expected: "MySQL numeric connection id",
+                                actual: connection_id_exp.to_string(),
                             }
+                            .fail();
+                        };
+
+                        let connection_id = u32::from_str(&s).map_err(|_| {
+                            error::UnexpectedTokenSnafu {
+                                expected: "MySQL numeric connection id",
+                                actual: s,
+                            }
+                            .build()
                         })?;
-                    Ok(Statement::Kill(process_id_ident))
+                        Kill::ConnectionId(connection_id)
+                    } else {
+                        let process_id_ident =
+                            self.parser.parse_literal_string().with_context(|_| {
+                                error::UnexpectedSnafu {
+                                    expected: "process id string literal",
+                                    actual: self.peek_token_as_string(),
+                                }
+                            })?;
+                        Kill::ProcessId(process_id_ident)
+                    };
+
+                    Ok(Statement::Kill(kill))
                 }
 
                 _ => self.unsupported(self.peek_token_as_string()),
@@ -440,4 +472,191 @@ mod tests {
         let stmt_name = ParserContext::parse_mysql_deallocate_stmt(sql, &MySqlDialect {}).unwrap();
         assert_eq!(stmt_name, "stmt2");
     }
+
+    #[test]
+    pub fn test_parse_kill_query_statement() {
+        use crate::statements::kill::Kill;
+
+        // Test MySQL-style KILL QUERY with connection ID
+        let sql = "KILL QUERY 123";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ConnectionId(connection_id)) => {
+                assert_eq!(*connection_id, 123);
+            }
+            _ => panic!("Expected Kill::ConnectionId statement"),
+        }
+
+        // Test with larger connection ID
+        let sql = "KILL QUERY 999999";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ConnectionId(connection_id)) => {
+                assert_eq!(*connection_id, 999999);
+            }
+            _ => panic!("Expected Kill::ConnectionId statement"),
+        }
+    }
+
+    #[test]
+    pub fn test_parse_kill_process_statement() {
+        use crate::statements::kill::Kill;
+
+        // Test KILL with process ID string
+        let sql = "KILL 'process-123'";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ProcessId(process_id)) => {
+                assert_eq!(process_id, "process-123");
+            }
+            _ => panic!("Expected Kill::ProcessId statement"),
+        }
+
+        // Test with double quotes
+        let sql = "KILL \"process-456\"";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ProcessId(process_id)) => {
+                assert_eq!(process_id, "process-456");
+            }
+            _ => panic!("Expected Kill::ProcessId statement"),
+        }
+
+        // Test with UUID-like process ID
+        let sql = "KILL 'f47ac10b-58cc-4372-a567-0e02b2c3d479'";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ProcessId(process_id)) => {
+                assert_eq!(process_id, "f47ac10b-58cc-4372-a567-0e02b2c3d479");
+            }
+            _ => panic!("Expected Kill::ProcessId statement"),
+        }
+    }
+
+    #[test]
+    pub fn test_parse_kill_statement_errors() {
+        // Test KILL QUERY without connection ID
+        let sql = "KILL QUERY";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(result.is_err());
+
+        // Test KILL QUERY with non-numeric connection ID
+        let sql = "KILL QUERY 'not-a-number'";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(result.is_err());
+
+        // Test KILL without any argument
+        let sql = "KILL";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(result.is_err());
+
+        // Test KILL QUERY with connection ID that's too large for u32
+        let sql = "KILL QUERY 4294967296"; // u32::MAX + 1
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(result.is_err());
+    }
+
+    #[test]
+    pub fn test_parse_kill_statement_edge_cases() {
+        use crate::statements::kill::Kill;
+
+        // Test KILL QUERY with zero connection ID
+        let sql = "KILL QUERY 0";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ConnectionId(connection_id)) => {
+                assert_eq!(*connection_id, 0);
+            }
+            _ => panic!("Expected Kill::ConnectionId statement"),
+        }
+
+        // Test KILL QUERY with maximum u32 value
+        let sql = "KILL QUERY 4294967295"; // u32::MAX
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ConnectionId(connection_id)) => {
+                assert_eq!(*connection_id, 4294967295);
+            }
+            _ => panic!("Expected Kill::ConnectionId statement"),
+        }
+
+        // Test KILL with empty string process ID
+        let sql = "KILL ''";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ProcessId(process_id)) => {
+                assert_eq!(process_id, "");
+            }
+            _ => panic!("Expected Kill::ProcessId statement"),
+        }
+    }
+
+    #[test]
+    pub fn test_parse_kill_statement_case_insensitive() {
+        use crate::statements::kill::Kill;
+
+        // Test lowercase
+        let sql = "kill query 123";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ConnectionId(connection_id)) => {
+                assert_eq!(*connection_id, 123);
+            }
+            _ => panic!("Expected Kill::ConnectionId statement"),
+        }
+
+        // Test mixed case
+        let sql = "Kill Query 456";
+        let statements =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        assert_eq!(statements.len(), 1);
+        match &statements[0] {
+            Statement::Kill(Kill::ConnectionId(connection_id)) => {
+                assert_eq!(*connection_id, 456);
+            }
+            _ => panic!("Expected Kill::ConnectionId statement"),
+        }
+    }
 }

src/sql/src/statements.rs

@@ -22,6 +22,7 @@ pub mod describe;
 pub mod drop;
 pub mod explain;
 pub mod insert;
+pub mod kill;
 mod option_map;
 pub mod query;
 pub mod set_variables;

src/sql/src/statements/kill.rs

@@ -0,0 +1,40 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt::{Display, Formatter};
+
+use serde::Serialize;
+use sqlparser_derive::{Visit, VisitMut};
+
+/// Arguments of `KILL` statements.
+#[derive(Debug, Clone, Eq, PartialEq, Visit, VisitMut, Serialize)]
+pub enum Kill {
+    /// Kill a remote process id.
+    ProcessId(String),
+    /// Kill MySQL connection id.
+    ConnectionId(u32),
+}
+
+impl Display for Kill {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Kill::ProcessId(id) => {
+                write!(f, "KILL {}", id)
+            }
+            Kill::ConnectionId(id) => {
+                write!(f, "KILL QUERY {}", id)
+            }
+        }
+    }
+}

src/sql/src/statements/statement.rs

@@ -32,6 +32,7 @@ use crate::statements::describe::DescribeTable;
 use crate::statements::drop::{DropDatabase, DropFlow, DropTable, DropView};
 use crate::statements::explain::Explain;
 use crate::statements::insert::Insert;
+use crate::statements::kill::Kill;
 use crate::statements::query::Query;
 use crate::statements::set_variables::SetVariables;
 use crate::statements::show::{
@@ -139,7 +140,7 @@ pub enum Statement {
     // CLOSE
     CloseCursor(CloseCursor),
     // KILL <process>
-    Kill(String),
+    Kill(Kill),
 }
 
 impl Display for Statement {