Repository: matrixorigin/matrixone
Pull request #13896 (2023-12-28 07:38:19)

[postfix] proxy: fix the prepare wait condition from execute to close
## What type of PR is this?

- [ ] API-change
- [x] BUG
- [ ] Improvement
- [ ] Documentation
- [ ] Feature
- [ ] Test and CI
- [ ] Code Refactoring

## Which issue(s) this PR fixes:

issue #13778 

## What this PR does / why we need it:
fix the prepare wait condition from execute to close.
Also, add support "use db" plain text back.

matrix-meow: @volgariver6 Thanks for your contributions!

Review:

Title: [postfix] proxy: fix the prepare wait condition from execute to close

Body: The pull request fixes the prepare wait condition from execute to close. It also adds support for the "use db" plain text.

Changes:
1. In `pkg/proxy/client_conn.go`, the `handleInitDB` function is modified to append the internal statement with the command type `cmdInitDB` and the string "use " + e.db.
2. In `pkg/proxy/event.go`, the `makeEvent` function is modified to handle the `Use` command and create an `initDBEvent` with the name of the database.
3. In `pkg/proxy/event.go`, the condition for setting the prepared flag is changed from `isCmdStmtExecute` to `isCmdStmtClose`.
4. In `pkg/proxy/mysql_conn_buf.go`, the constant `cmdStmtExecute` is renamed to `cmdStmtClose`.
5. In `pkg/proxy/util.go`, the function `isCmdStmtExecute` is renamed to `isCmdStmtClose`.
6. In `pkg/proxy/util_test.go`, the function `TestIsCmdStmtExecute` is renamed to `TestIsCmdStmtClose` and the test cases are updated accordingly.

Overall, the changes in the pull request seem to address the issue mentioned in the title and body. However, there are a few suggestions for improvement:

1. The commit message could be more descriptive, providing more context about the issue being fixed and the reason behind the fix.
2. It would be helpful to include test cases for the added functionality to ensure that it works as expected.
3. The code could benefit from more comments to explain the purpose and functionality of the modified functions and constants.
4. It's important to ensure that the changes made in this pull request do not introduce any security vulnerabilities. A thorough security review should be conducted to identify any potential issues.

In terms of optimization, the changes seem to be minimal and focused on fixing the specific issue. Therefore, there may not be significant opportunities for optimization in this pull request.

Patch:
diff --git a/pkg/proxy/client_conn.go b/pkg/proxy/client_conn.go
index b98b09f5d..019db3495 100644
--- a/pkg/proxy/client_conn.go
+++ b/pkg/proxy/client_conn.go
@@ -363,7 +363,7 @@ func (c *clientConn) handlePrepare(e *prepareEvent) error {
 
 // handleUse handles the use event.
 func (c *clientConn) handleInitDB(e *initDBEvent) error {
-	c.redoStmts = append(c.redoStmts, internalStmt{cmdType: cmdInitDB, s: e.db})
+	c.redoStmts = append(c.redoStmts, internalStmt{cmdType: cmdInitDB, s: "use " + e.db})
 	return nil
 }
 
diff --git a/pkg/proxy/event.go b/pkg/proxy/event.go
index 175fab889..26cabee4d 100644
--- a/pkg/proxy/event.go
+++ b/pkg/proxy/event.go
@@ -104,6 +104,8 @@ func makeEvent(msg []byte, b *msgBuf) (IEvent, bool) {
 			return makeSetVarEvent(sql), false
 		case *tree.PrepareString:
 			return makePrepareEvent(sql), false
+		case *tree.Use:
+			return makeInitDBEvent(s.Name.Origin()), false
 		default:
 			return nil, false
 		}
@@ -111,7 +113,7 @@ func makeEvent(msg []byte, b *msgBuf) (IEvent, bool) {
 		return makeInitDBEvent(getStatement(msg)), false
 	} else if isCmdStmtPrepare(msg) && b != nil {
 		b.setPrepared(true)
-	} else if isCmdStmtExecute(msg) && b != nil {
+	} else if isCmdStmtClose(msg) && b != nil {
 		b.setPrepared(false)
 	}
 	return nil, false
diff --git a/pkg/proxy/mysql_conn_buf.go b/pkg/proxy/mysql_conn_buf.go
index 498cef39e..7f1ec4b67 100644
--- a/pkg/proxy/mysql_conn_buf.go
+++ b/pkg/proxy/mysql_conn_buf.go
@@ -48,7 +48,7 @@ const (
 	cmdInitDB MySQLCmd = 0x02
 	// For stmt prepare and execute cmd from JDBC.
 	cmdStmtPrepare MySQLCmd = 0x16
-	cmdStmtExecute MySQLCmd = 0x17
+	cmdStmtClose   MySQLCmd = 0x19
 )
 
 // MySQLConn contains a buffer to save data which may be only part
diff --git a/pkg/proxy/util.go b/pkg/proxy/util.go
index a6d9c2758..9193f1984 100644
--- a/pkg/proxy/util.go
+++ b/pkg/proxy/util.go
@@ -56,8 +56,8 @@ func isCmdStmtPrepare(p []byte) bool {
 	return false
 }
 
-func isCmdStmtExecute(p []byte) bool {
-	if len(p) > 4 && p[4] == byte(cmdStmtExecute) {
+func isCmdStmtClose(p []byte) bool {
+	if len(p) > 4 && p[4] == byte(cmdStmtClose) {
 		return true
 	}
 	return false


Test patch:
diff --git a/pkg/proxy/util_test.go b/pkg/proxy/util_test.go
index 69abeb952..4297a8dfa 100644
--- a/pkg/proxy/util_test.go
+++ b/pkg/proxy/util_test.go
@@ -156,17 +156,17 @@ func TestIsCmdStmtPrepare(t *testing.T) {
 	require.True(t, ret)
 }
 
-func TestIsCmdStmtExecute(t *testing.T) {
+func TestIsCmdStmtClose(t *testing.T) {
 	var data []byte
-	ret := isCmdStmtExecute(data)
+	ret := isCmdStmtClose(data)
 	require.False(t, ret)
 
 	data = []byte{0, 0, 0, 0, 20, 0}
-	ret = isCmdStmtExecute(data)
+	ret = isCmdStmtClose(data)
 	require.False(t, ret)
 
-	data = []byte{0, 0, 0, 0, byte(cmdStmtExecute), 0}
-	ret = isCmdStmtExecute(data)
+	data = []byte{0, 0, 0, 0, byte(cmdStmtClose), 0}
+	ret = isCmdStmtClose(data)
 	require.True(t, ret)
 }