Merge pull request #6760 from planetscale/bp-6748

[Backport] Fix reserved connection in autocommit mode on DML #6748

go/test/endtoend/vtgate/setstatement/sysvar_test.go

@@ -192,6 +192,36 @@ func TestSetSystemVariableAndThenSuccessfulTx(t *testing.T) {
 	assertMatches(t, conn, "select @@sql_safe_updates", "[[INT64(1)]]")
 }
 
+func TestSetSystemVariableAndThenSuccessfulAutocommitDML(t *testing.T) {
+	vtParams := mysql.ConnParams{
+		Host: "localhost",
+		Port: clusterInstance.VtgateMySQLPort,
+	}
+
+	conn, err := mysql.Connect(context.Background(), &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+	checkedExec(t, conn, `delete from test`)
+
+	checkedExec(t, conn, `set sql_safe_updates = 1`)
+
+	checkedExec(t, conn, `insert into test (id, val1) values (80, null)`)
+	assertMatches(t, conn, `select id, val1 from test`, `[[INT64(80) NULL]]`)
+	assertMatches(t, conn, `select @@sql_safe_updates`, `[[INT64(1)]]`)
+
+	checkedExec(t, conn, `update test set val2 = 2 where val1 is null`)
+	assertMatches(t, conn, `select id, val1, val2 from test`, `[[INT64(80) NULL INT32(2)]]`)
+	assertMatches(t, conn, `select @@sql_safe_updates`, `[[INT64(1)]]`)
+
+	checkedExec(t, conn, `update test set val1 = 'text' where val1 is null`)
+	assertMatches(t, conn, `select id, val1, val2 from test`, `[[INT64(80) VARCHAR("text") INT32(2)]]`)
+	assertMatches(t, conn, `select @@sql_safe_updates`, `[[INT64(1)]]`)
+
+	checkedExec(t, conn, `delete from test where val1 = 'text'`)
+	assertMatches(t, conn, `select id, val1, val2 from test`, `[]`)
+	assertMatches(t, conn, `select @@sql_safe_updates`, `[[INT64(1)]]`)
+}
+
 func TestStartTxAndSetSystemVariableAndThenSuccessfulCommit(t *testing.T) {
 	vtParams := mysql.ConnParams{
 		Host: "localhost",

go/vt/vtgate/safe_session.go

@@ -214,7 +214,10 @@ func (session *SafeSession) AppendOrUpdate(shardSession *vtgatepb.Session_ShardS
 	session.mu.Lock()
 	defer session.mu.Unlock()
 
-	if session.autocommitState == autocommitted {
+	// additional check of transaction id is required
+	// as now in autocommit mode there can be session due to reserved connection
+	// that needs to be stored as shard session.
+	if session.autocommitState == autocommitted && shardSession.TransactionId != 0 {
 		// Should be unreachable
 		return vterrors.New(vtrpcpb.Code_INTERNAL, "BUG: SafeSession.AppendOrUpdate: unexpected autocommit state")
 	}