set ReadyForQuery transaction indicater for in transaction query

Author: jennifersp

go.mod

@@ -10,7 +10,7 @@ require (
 	github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20241119094239-f4e529af734d
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
 	github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad
-	github.com/dolthub/go-mysql-server v0.20.1-0.20250520164330-094c26d756a3
+	github.com/dolthub/go-mysql-server v0.20.1-0.20250520221132-8961ce0f4cab
 	github.com/dolthub/sqllogictest/go v0.0.0-20240618184124-ca47f9354216
 	github.com/dolthub/vitess v0.0.0-20250512224608-8fb9c6ea092c
 	github.com/fatih/color v1.13.0

go.sum

@@ -266,8 +266,8 @@ github.com/dolthub/fslock v0.0.3 h1:iLMpUIvJKMKm92+N1fmHVdxJP5NdyDK5bK7z7Ba2s2U=
 github.com/dolthub/fslock v0.0.3/go.mod h1:QWql+P17oAAMLnL4HGB5tiovtDuAjdDTPbuqx7bYfa0=
 github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad h1:66ZPawHszNu37VPQckdhX1BPPVzREsGgNxQeefnlm3g=
 github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad/go.mod h1:ylU4XjUpsMcvl/BKeRRMXSH7e7WBrPXdSLvnRJYrxEA=
-github.com/dolthub/go-mysql-server v0.20.1-0.20250520164330-094c26d756a3 h1:rFkHy4L/GfVQh7KxwFPI/QmkPcBpT/PjnHLp+zde0xA=
-github.com/dolthub/go-mysql-server v0.20.1-0.20250520164330-094c26d756a3/go.mod h1:5ZdrW0fHZbz+8CngT9gksqSX4H3y+7v1pns7tJCEpu0=
+github.com/dolthub/go-mysql-server v0.20.1-0.20250520221132-8961ce0f4cab h1:gbxIHdLroWf8fdPRvs2o5CyvcHcgRJIF2D9Lw+L6Wd8=
+github.com/dolthub/go-mysql-server v0.20.1-0.20250520221132-8961ce0f4cab/go.mod h1:5ZdrW0fHZbz+8CngT9gksqSX4H3y+7v1pns7tJCEpu0=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63 h1:OAsXLAPL4du6tfbBgK0xXHZkOlos63RdKYS3Sgw/dfI=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63/go.mod h1:lV7lUeuDhH5thVGDCKXbatwKy2KW80L4rMT46n+Y2/Q=
 github.com/dolthub/ishell v0.0.0-20240701202509-2b217167d718 h1:lT7hE5k+0nkBdj/1UOSFwjWpNxf+LCApbRHgnCA17XE=

server/connection_handler.go

@@ -65,6 +65,8 @@ type ConnectionHandler struct {
 	// copyFromStdinState is set when this connection is in the COPY FROM STDIN mode, meaning it is waiting on
 	// COPY DATA messages from the client to import data into tables.
 	copyFromStdinState *copyFromStdinState
+	// inTransaction is set to true with BEGIN query and false with COMMIT query.
+	inTransaction bool
 }
 
 // Set this env var to disable panic handling in the connection, which is useful when debugging a panic
@@ -448,6 +450,10 @@ func (h *ConnectionHandler) handleQuery(message *pgproto3.Query) (endOfMessages
 // and any error that occurred while handling the query.
 func (h *ConnectionHandler) handleQueryOutsideEngine(query ConvertedQuery) (handled bool, endOfMessages bool, err error) {
 	switch stmt := query.AST.(type) {
+	case *sqlparser.Begin:
+		h.inTransaction = true
+	case *sqlparser.Commit:
+		h.inTransaction = false
 	case *sqlparser.Deallocate:
 		// TODO: handle ALL keyword
 		return true, true, h.deallocatePreparedStatement(stmt.Name, h.preparedStatements, query, h.Conn())
@@ -1032,10 +1038,15 @@ func (h *ConnectionHandler) handledPSQLCommands(statement string) (bool, error)
 // query. A nil error should be provided if this is being called naturally.
 func (h *ConnectionHandler) endOfMessages(err error) {
 	if err != nil {
+		// TODO: is ReadyForQueryTransactionIndicator_FailedTransactionBlock used here?
 		h.sendError(err)
 	}
+	ti := ReadyForQueryTransactionIndicator_Idle
+	if h.inTransaction {
+		ti = ReadyForQueryTransactionIndicator_TransactionBlock
+	}
 	if sendErr := h.send(&pgproto3.ReadyForQuery{
-		TxStatus: byte(ReadyForQueryTransactionIndicator_Idle),
+		TxStatus: byte(ti),
 	}); sendErr != nil {
 		// We panic here for the same reason as above.
 		panic(sendErr)

testing/go/prepared_statement_test.go

@@ -425,6 +425,19 @@ var preparedStatementTests = []ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "insert returning",
+		SetUpScript: []string{
+			"CREATE TABLE test (id serial, name text)",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    `INSERT INTO test (name) VALUES ($1) RETURNING id;`,
+				BindVars: []any{"test_name"},
+				Expected: []sql.Row{{1}},
+			},
+		},
+	},
 }
 
 var pgCatalogTests = []ScriptTest{