introduce a new method mysqlConn.Reset() allowing to reset the mysql connection (implemented as "return mc.sendNoArgsCommandWithResultOK(ctx, comConnReset)" where comConnReset=31); add tests

Author: zenovich

connection.go

@@ -685,6 +685,11 @@ func (mc *mysqlConn) startWatcher() {
 	}()
 }
 
+// Reset resets the MySQL connection.
+func (mc *mysqlConn) Reset(ctx context.Context) (err error) {
+	return mc.sendNoArgsCommandWithResultOK(ctx, comConnReset)
+}
+
 func (mc *mysqlConn) CheckNamedValue(nv *driver.NamedValue) (err error) {
 	nv.Value, err = converter{}.ConvertValue(nv.Value)
 	return

connection_test.go

@@ -129,65 +129,98 @@ func TestCheckNamedValue(t *testing.T) {
 	}
 }
 
-// TestCleanCancel tests passed context is cancelled at start.
+// TestNoArgsCommandCleanCancel tests passed context is cancelled at start.
 // No packet should be sent.  Connection should keep current status.
-func TestCleanCancel(t *testing.T) {
-	mc := &mysqlConn{
-		closech: make(chan struct{}),
-	}
-	mc.startWatcher()
-	defer mc.cleanup()
-
-	ctx, cancel := context.WithCancel(context.Background())
-	cancel()
-
-	for range 3 { // Repeat same behavior
-		err := mc.Ping(ctx)
-		if err != context.Canceled {
-			t.Errorf("expected context.Canceled, got %#v", err)
-		}
-
-		if mc.closed.Load() {
-			t.Error("expected mc is not closed, closed actually")
-		}
-
-		if mc.watching {
-			t.Error("expected watching is false, but true")
-		}
+func TestNoArgsCommandCleanCancel(t *testing.T) {
+	for _, test := range []struct {
+		name string
+		funcToCall func(ctx context.Context, mc *mysqlConn) error
+	} {
+		{name: "Ping", funcToCall: func(ctx context.Context, mc *mysqlConn) error { return mc.Ping(ctx) }},
+		{name: "Reset", funcToCall: func(ctx context.Context, mc *mysqlConn) error { return mc.Reset(ctx) }},
+	} {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			mc := &mysqlConn{
+				closech: make(chan struct{}),
+			}
+			mc.startWatcher()
+			defer mc.cleanup()
+
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel()
+
+			for range 3 { // Repeat same behavior
+				err := test.funcToCall(ctx, mc)
+				if err != context.Canceled {
+					t.Errorf("expected context.Canceled, got %#v", err)
+				}
+
+				if mc.closed.Load() {
+					t.Error("expected mc is not closed, closed actually")
+				}
+
+				if mc.watching {
+					t.Error("expected watching is false, but true")
+				}
+			}
+		})
 	}
 }
 
-func TestPingMarkBadConnection(t *testing.T) {
-	nc := badConnection{err: errors.New("boom")}
-	mc := &mysqlConn{
-		netConn:          nc,
-		buf:              newBuffer(),
-		maxAllowedPacket: defaultMaxAllowedPacket,
-		closech:          make(chan struct{}),
-		cfg:              NewConfig(),
-	}
-
-	err := mc.Ping(context.Background())
-
-	if err != driver.ErrBadConn {
-		t.Errorf("expected driver.ErrBadConn, got  %#v", err)
+func TestNoArgsCommandMarkBadConnection(t *testing.T) {
+	for _, test := range []struct {
+		name string
+		funcToCall func(mc *mysqlConn) error
+	} {
+		{name: "Ping", funcToCall: func(mc *mysqlConn) error { return mc.Ping(context.Background()) }},
+		{name: "Reset", funcToCall: func(mc *mysqlConn) error { return mc.Reset(context.Background()) }},
+	} {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			nc := badConnection{err: errors.New("boom")}
+			mc := &mysqlConn{
+				netConn:          nc,
+				buf:              newBuffer(),
+				maxAllowedPacket: defaultMaxAllowedPacket,
+				closech:          make(chan struct{}),
+				cfg:              NewConfig(),
+			}
+
+			err := test.funcToCall(mc)
+
+			if err != driver.ErrBadConn {
+				t.Errorf("expected driver.ErrBadConn, got  %#v", err)
+			}
+		})
 	}
 }
 
-func TestPingErrInvalidConn(t *testing.T) {
-	nc := badConnection{err: errors.New("failed to write"), n: 10}
-	mc := &mysqlConn{
-		netConn:          nc,
-		buf:              newBuffer(),
-		maxAllowedPacket: defaultMaxAllowedPacket,
-		closech:          make(chan struct{}),
-		cfg:              NewConfig(),
-	}
-
-	err := mc.Ping(context.Background())
-
-	if err != nc.err {
-		t.Errorf("expected %#v, got  %#v", nc.err, err)
+func TestNoArgsCommandErrInvalidConn(t *testing.T) {
+	for _, test := range []struct {
+		name string
+		funcToCall func(mc *mysqlConn) error
+	} {
+		{name: "Ping", funcToCall: func(mc *mysqlConn) error { return mc.Ping(context.Background()) }},
+		{name: "Reset", funcToCall: func(mc *mysqlConn) error { return mc.Reset(context.Background()) }},
+	} {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			nc := badConnection{err: errors.New("failed to write"), n: 10}
+			mc := &mysqlConn{
+				netConn:          nc,
+				buf:              newBuffer(),
+				maxAllowedPacket: defaultMaxAllowedPacket,
+				closech:          make(chan struct{}),
+				cfg:              NewConfig(),
+			}
+
+			err := test.funcToCall(mc)
+
+			if err != nc.err {
+				t.Errorf("expected %#v, got  %#v", nc.err, err)
+			}
+		})
 	}
 }
 

const.go

@@ -115,6 +115,7 @@ const (
 	comStmtReset
 	comSetOption
 	comStmtFetch
+	comConnReset = 31
 )
 
 // https://dev.mysql.com/doc/internals/en/com-query-response.html#packet-Protocol::ColumnType

driver_test.go

@@ -2322,13 +2322,69 @@ func TestRejectReadOnly(t *testing.T) {
 }
 
 func TestPing(t *testing.T) {
-	ctx := context.Background()
 	runTests(t, dsn, func(dbt *DBTest) {
 		if err := dbt.db.Ping(); err != nil {
 			dbt.fail("Ping", "Ping", err)
 		}
 	})
+}
+
+func TestNoArgsCommand(t *testing.T) {
+	ctx := context.Background()
+	for _, test := range []struct{
+		method string
+		query string
+		funcToCall func(ctx context.Context, mc *mysqlConn) error
+	} {
+		{method: "Ping", query: "Ping", funcToCall: func(ctx context.Context, mc *mysqlConn) error {return mc.Ping(ctx)}},
+		{method: "Conn", query: "Reset", funcToCall: func(ctx context.Context, mc *mysqlConn) error {return mc.Reset(ctx)}},
+	} {
+		test := test
+		t.Run(test.method+"_"+test.query, func(t *testing.T) {
+			runTests(t, dsn, func(dbt *DBTest) {
+				conn, err := dbt.db.Conn(ctx)
+				if err != nil {
+					dbt.fail("db", "Conn", err)
+				}
+
+				// Check that affectedRows and insertIds are cleared after each call.
+				conn.Raw(func(conn any) error {
+					c := conn.(*mysqlConn)
+
+					// Issue a query that sets affectedRows and insertIds.
+					q, err := c.Query(`SELECT 1`, nil)
+					if err != nil {
+						dbt.fail("Conn", "Query", err)
+					}
+					if got, want := c.result.affectedRows, []int64{0}; !reflect.DeepEqual(got, want) {
+						dbt.Fatalf("bad affectedRows: got %v, want=%v", got, want)
+					}
+					if got, want := c.result.insertIds, []int64{0}; !reflect.DeepEqual(got, want) {
+						dbt.Fatalf("bad insertIds: got %v, want=%v", got, want)
+					}
+					q.Close()
+
+					// Verify that Ping()/Reset() clears both fields.
+					for range 2 {
+						if err := test.funcToCall(ctx, c); err != nil {
+							dbt.fail(test.method, test.query, err)
+						}
+						if got, want := c.result.affectedRows, []int64(nil); !reflect.DeepEqual(got, want) {
+							t.Errorf("bad affectedRows: got %v, want=%v", got, want)
+						}
+						if got, want := c.result.insertIds, []int64(nil); !reflect.DeepEqual(got, want) {
+							t.Errorf("bad affectedRows: got %v, want=%v", got, want)
+						}
+					}
+					return nil
+				})
+			})
+		})
+	}
+}
 
+func TestReset(t *testing.T) {
+	ctx := context.Background()
 	runTests(t, dsn, func(dbt *DBTest) {
 		conn, err := dbt.db.Conn(ctx)
 		if err != nil {
@@ -2339,31 +2395,49 @@ func TestPing(t *testing.T) {
 		conn.Raw(func(conn any) error {
 			c := conn.(*mysqlConn)
 
-			// Issue a query that sets affectedRows and insertIds.
-			q, err := c.Query(`SELECT 1`, nil)
+			_, err = c.ExecContext(ctx, "SET @a := 1", nil)
+			if err != nil {
+				dbt.fail("Conn", "ExecContext", err)
+			}
+			var rows driver.Rows
+			rows, err = c.QueryContext(ctx, "SELECT @a", nil)
 			if err != nil {
-				dbt.fail("Conn", "Query", err)
+				dbt.fail("Conn", "QueryContext", err)
 			}
-			if got, want := c.result.affectedRows, []int64{0}; !reflect.DeepEqual(got, want) {
-				dbt.Fatalf("bad affectedRows: got %v, want=%v", got, want)
+			result := []driver.Value{0}
+			err = rows.Next(result)
+			if err != nil {
+				dbt.fail("Rows", "Next", err)
 			}
-			if got, want := c.result.insertIds, []int64{0}; !reflect.DeepEqual(got, want) {
-				dbt.Fatalf("bad insertIds: got %v, want=%v", got, want)
+			err = rows.Close()
+			if err != nil {
+				dbt.fail("Rows", "Close", err)
+			}
+			if !reflect.DeepEqual([]driver.Value{int64(1)}, result) {
+				dbt.Fatalf("failed to set @a to 1 with SET: got %v, want=%v", result, []driver.Value{int64(1)})
 			}
-			q.Close()
 
-			// Verify that Ping() clears both fields.
-			for range 2 {
-				if err := c.Ping(ctx); err != nil {
-					dbt.fail("Pinger", "Ping", err)
-				}
-				if got, want := c.result.affectedRows, []int64(nil); !reflect.DeepEqual(got, want) {
-					t.Errorf("bad affectedRows: got %v, want=%v", got, want)
-				}
-				if got, want := c.result.insertIds, []int64(nil); !reflect.DeepEqual(got, want) {
-					t.Errorf("bad affectedRows: got %v, want=%v", got, want)
-				}
+			err = c.Reset(ctx)
+			if err != nil {
+				dbt.fail("Conn", "Reset", err)
+			}
+
+			rows, err = c.QueryContext(ctx, "SELECT @a", nil)
+			if err != nil {
+				dbt.fail("Conn", "QueryContext", err)
+			}
+			err = rows.Next(result)
+			if err != nil {
+				dbt.fail("Rows", "Next", err)
+			}
+			err = rows.Close()
+			if err != nil {
+				dbt.fail("Rows", "Close", err)
 			}
+			if !reflect.DeepEqual([]driver.Value{nil}, result) {
+				dbt.Fatalf("Reset did not reset the session (@a is still set): got %v, want=%v", result, []driver.Value{nil})
+			}
+
 			return nil
 		})
 	})