2
0
Эх сурвалжийг харах

add resumable portals and bounded Execute

Execute with maxRows > 0 now returns at most that many rows and then
PortalSuspended ('s'), and a later Execute on the same portal resumes at the
next row instead of re-running the statement; maxRows = 0 keeps returning all
rows in one response.

A portal materializes its result once on the first Execute - executor result,
per-column result formats, column OIDs, completion tag - and every later
Execute slices it, so DML ... RETURNING runs exactly once and released rows are
dropped on completion. RowDescription is sent by the first Execute only, and
per-portal result formats apply on every Execute.

Portal lifetime follows PostgreSQL: reusable until Close('P'), a Bind of the
same name, or the end of the transaction that created it (Sync outside an
explicit block, COMMIT/ROLLBACK inside one), after which Execute fails with
34000 (25P02 while the transaction is aborted). An unknown or closed portal is
now 34000 instead of 22023; a negative row limit stays 08P01.

pkg/pgserver/portal_test.go covers suspend/resume across Executes, exact row
counts and frame sequences, maxRows = 0, exact-limit suspension, Close('P'),
re-Bind, Sync and transaction end, DML RETURNING executed once, binary formats
across resumes, and the error paths. pgx v5.11.0 has no bounded Execute API, so
the raw protocol is the only path that can request a suspended portal.

Capabilities and README now describe the supported semantics and limits.
Danilo Fragoso 1 өдөр өмнө
parent
commit
89253b9bc6

Файлын зөрүү хэтэрхий том тул дарагдсан байна
+ 0 - 1
README.md


+ 1 - 1
pkg/capabilities/capabilities.go

@@ -162,7 +162,7 @@ func Current() *Document {
 			ExtendedQuery:  true,
 			Describe:       Feature{Supported: true, Status: "compatibility_no_data", Notes: "statement Describe returns ParameterDescription with the negotiated parameter types (a client Parse OID wins; an OID of 0 is inferred from casts, comparisons, assignment targets, and ANY/ALL operands, and falls back to TEXT rather than 0) then NoData; portal Describe returns NoData without a ParameterDescription, matching PostgreSQL; Execute still emits the real RowDescription; a placeholder number above the 65535-parameter protocol limit fails Parse with 54000 before any allocation"},
 			BinaryResults:  Feature{Supported: true, Status: "supported", Notes: "Bind accepts per-column result formats (0=text, 1=binary): an empty list is all text, a single entry applies to every column, and any other length must equal the result column count (08P01 otherwise); binary is encodable for bool (16), bytea (17), int8 (20), int2 (21), int4 (23), text (25), float4 (700), float8 (701), varchar (1043), numeric (1700) only as the element type of numeric[] (1231), jsonb (3802, version byte 0x01 plus canonical JSON text), and one-dimensional arrays (1000, 1005, 1007, 1009, 1015, 1016, 1021, 1022, 1231) with NULL elements and empty arrays preserved; a binary request for any other reported OID (timestamptz 1184, date 1082) fails with 0A000; formats are resolved when Execute produces the result because Describe returns NoData, so for INSERT/UPDATE/DELETE ... RETURNING the statement runs before a bad format or OID is rejected (no rows are sent and no rollback is attempted); an unrepresentable stored value (for example text in an INTEGER column, or a float4 overflow) fails with 22P02 without emitting the failing DataRow; Describe and simple-query results stay text, and numeric array elements are encoded exactly from their stored decimal text"},
-			BoundedExecute: Feature{Supported: false, Status: "rejected", SQLState: "0A000", Notes: "Execute maxRows must be zero; portals are one-shot, not resumable"},
+			BoundedExecute: Feature{Supported: true, Status: "supported", SQLState: "34000", Notes: "Execute maxRows > 0 returns at most maxRows rows and then PortalSuspended ('s'); a later Execute on the same portal resumes at the next row without re-running the statement, and maxRows = 0 returns all remaining rows and CommandComplete; the executor result is materialized once on the first Execute and later Executes slice it, so memory is bounded by that result and there is no scroll/backward fetch and no WITH HOLD; RowDescription is sent by the first Execute only, and a portal repeats its completion tag without new rows once exhausted; results are per portal, so per-column result formats apply to every Execute; a portal survives until Close('P'), a Bind of the same name, or the end of its transaction (Sync outside an explicit block, COMMIT/ROLLBACK inside one), after which Execute fails with 34000 (inside an aborted transaction, 25P02); maxRows < 0 is a protocol error (08P01) and DML ... RETURNING runs exactly once"},
 			ProxyErrors:    true,
 			TLS:            "terminated by the managed pgproxy, not the engine listener",
 		},

+ 24 - 0
pkg/capabilities/capabilities_test.go

@@ -323,3 +323,27 @@ func TestBinaryResultsAreAdvertised(t *testing.T) {
 		}
 	}
 }
+
+// Bounded Execute is advertised with its resume semantics, portal lifetime, and
+// limits. pkg/pgserver covers each claim with raw-protocol portal tests.
+func TestBoundedExecuteIsAdvertised(t *testing.T) {
+	bounded := Current().Protocol.BoundedExecute
+	if !bounded.Supported {
+		t.Fatalf("bounded_execute = %+v, want supported", bounded)
+	}
+	if bounded.Status != "supported" {
+		t.Fatalf("bounded_execute status = %q, want supported", bounded.Status)
+	}
+	if bounded.SQLState != "34000" {
+		t.Fatalf("bounded_execute sqlstate = %q, want 34000", bounded.SQLState)
+	}
+	for _, claim := range []string{
+		"PortalSuspended", "maxRows = 0", "materialized once", "no scroll",
+		"WITH HOLD", "RowDescription is sent by the first Execute only",
+		"Close('P')", "34000", "25P02", "08P01", "RETURNING",
+	} {
+		if !strings.Contains(bounded.Notes, claim) {
+			t.Errorf("bounded_execute notes do not mention %q", claim)
+		}
+	}
+}

+ 233 - 99
pkg/pgserver/connection.go

@@ -69,13 +69,43 @@ type preparedStatement struct {
 	paramOIDs []int32
 }
 
+// portal is one Bind result. Its lifetime follows PostgreSQL: a portal is
+// reusable until Close('P'), a Bind of the same name, or the end of the
+// transaction it was created in (see Connection.clearPortals).
 type portal struct {
-	query    string
-	executed bool
+	query string
 	// resultFormats is the Bind result format list exactly as the client sent
 	// it: empty (all text), one entry that repeats for every column, or one
-	// entry per column. It is resolved against the produced result at Execute.
+	// entry per column. It is resolved once, when the first Execute produces the
+	// result.
 	resultFormats []int16
+	// execution is nil until the first Execute. Materializing once lets a
+	// resumed portal stream rows from execution.result instead of running the
+	// statement again.
+	execution *portalExecution
+}
+
+// portalExecution is the result a portal materialized on its first Execute,
+// plus the stream cursor a later Execute continues from.
+type portalExecution struct {
+	result *executor.Result
+	// columnTypes, oids, and formats are the result metadata resolved once for
+	// every Execute of this portal.
+	columnTypes []string
+	oids        []int32
+	formats     []int16
+	// commandTag is the completion tag sent when the rows are exhausted.
+	commandTag string
+	// endsTransaction reports that this statement ended the explicit
+	// transaction, so every portal is destroyed with it.
+	endsTransaction bool
+	rowOffset       int
+	// described is true once the RowDescription was sent, so a resumed Execute
+	// sends only DataRows.
+	described bool
+	// completed is true once CommandComplete was sent; a further Execute
+	// repeats the tag without re-running the statement.
+	completed bool
 }
 
 // NewConnection creates a new connection handler
@@ -669,14 +699,11 @@ func (c *Connection) handleExecute(msg *Message) error {
 	}
 	portal, ok := c.portals[name]
 	if !ok {
-		return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("portal %q does not exist", name))
-	}
-	if portal.executed {
-		return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("portal can only be executed once"))
+		return c.failExtended(ErrCodeUndefinedCursor, fmt.Errorf("portal %q does not exist", name))
 	}
 	maxRows := int32(binary.BigEndian.Uint32(msg.Data[pos : pos+4]))
 	if maxRows < 0 {
-		return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Execute row limit"))
+		return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Execute row limit: %d", maxRows))
 	}
 	// Check the transaction state before catalog emulation. Catalog queries do
 	// not all parse as regular PizzaSQL statements, but must still return 25P02
@@ -684,116 +711,202 @@ func (c *Connection) handleExecute(msg *Message) error {
 	if c.txStatus == TxStatusFailed && !isRollbackStatement(portal.query) {
 		return c.failExtended(ErrCodeTransactionAborted, fmt.Errorf("current transaction is aborted, commands ignored until end of transaction block"))
 	}
-	if maxRows != 0 {
-		return c.failExtended(ErrCodeFeatureNotSupported, fmt.Errorf("bounded Execute is not supported; use maxRows=0"))
+	if portal.execution == nil {
+		if failure := c.materializePortal(portal); failure != nil {
+			return c.failExtended(failure.code, failure.err)
+		}
+	}
+	if err := c.sendPortalRows(portal, maxRows); err != nil {
+		return err
+	}
+	if portal.execution.endsTransaction {
+		// PostgreSQL destroys the portals of the transaction that just ended.
+		c.clearPortals()
 	}
-	portal.executed = true
+	return nil
+}
+
+// executeError is a rejected Execute. Its transaction-state effects have already
+// been applied by the time it is built, so handleExecute only reports it.
+type executeError struct {
+	code string
+	err  error
+}
+
+// materializePortal runs a portal's statement once and stores everything the
+// Execute responses need. The executor already produces a complete result, so a
+// resumed Execute stays a slice of that result and never runs the SQL again;
+// column metadata and the completion tag are also computed once. A nil return
+// means the portal is ready to stream.
+func (c *Connection) materializePortal(portal *portal) *executeError {
+	execution := &portalExecution{}
 	if result, handled, catalogErr := c.catalogResult(portal.query); handled {
 		if catalogErr != nil {
-			return c.failExtended(ErrCodeInternalError, catalogErr)
+			return &executeError{code: ErrCodeInternalError, err: catalogErr}
 		}
-		return c.sendTabularResult(result, "SELECT", portal.resultFormats)
-	}
-	l := lexer.New(portal.query)
-	stmt, err := parser.New(l).Parse()
-	if err != nil {
-		code := ErrCodeSyntaxError
-		if errors.Is(err, pgerror.ErrFeatureNotSupported) {
-			code = ErrCodeFeatureNotSupported
-			if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
+		execution.result = result
+		execution.commandTag = fmt.Sprintf("SELECT %d", len(result.Rows))
+	} else {
+		stmt, err := parser.New(lexer.New(portal.query)).Parse()
+		if err != nil {
+			code := ErrCodeSyntaxError
+			if errors.Is(err, pgerror.ErrFeatureNotSupported) {
+				code = ErrCodeFeatureNotSupported
+				if c.txStatus == TxStatusInBlock {
+					c.txStatus = TxStatusFailed
+				}
 			}
+			return &executeError{code: code, err: fmt.Errorf("syntax error: %w", err)}
+		}
+		result, err := c.executor.Execute(stmt)
+		if err != nil && isIdleRollback(stmt, c.txStatus, err) {
+			result, err = executor.NewResult("ROLLBACK"), nil
+		}
+		if err != nil {
+			return c.classifyExecuteError(stmt, err)
+		}
+		execution.result = result
+		// getCommandTag applies the statement's transaction transition, so it
+		// belongs here: a repeated Execute must not replay BEGIN/COMMIT/ROLLBACK.
+		execution.commandTag = c.getCommandTag(stmt, result)
+		switch s := stmt.(type) {
+		case *parser.CommitStmt:
+			execution.endsTransaction = true
+		case *parser.RollbackStmt:
+			execution.endsTransaction = s.Savepoint == ""
 		}
-		return c.failExtended(code, fmt.Errorf("syntax error: %w", err))
 	}
-	result, err := c.executor.Execute(stmt)
-	if err != nil && isIdleRollback(stmt, c.txStatus, err) {
-		result, err = executor.NewResult("ROLLBACK"), nil
+	columnTypes := wireColumnTypes(execution.result)
+	formats, formatErr := c.resolveResultFormats(portal.resultFormats, columnTypes, len(execution.result.Columns))
+	if formatErr != nil {
+		return &executeError{code: resultFormatSQLState(formatErr), err: formatErr}
 	}
-	if err != nil {
-		code := ErrCodeInternalError
-		var primaryKeyErr *storage.PrimaryKeyViolationError
-		var uniqueErr *storage.UniqueViolationError
-		var notNullErr *storage.NotNullViolationError
-		var checkErr *storage.CheckViolationError
-		var foreignKeyErr *storage.ForeignKeyViolationError
-		if errors.Is(err, storage.ErrSerialization) {
-			code = ErrCodeSerializationFailure
+	execution.columnTypes = columnTypes
+	execution.oids = c.resultColumnOIDs(columnTypes, len(execution.result.Columns))
+	execution.formats = formats
+	portal.execution = execution
+	return nil
+}
+
+// classifyExecuteError maps a statement execution error to the SQLSTATE the
+// client receives and leaves the transaction state consistent.
+func (c *Connection) classifyExecuteError(stmt parser.Statement, err error) *executeError {
+	code := ErrCodeInternalError
+	var primaryKeyErr *storage.PrimaryKeyViolationError
+	var uniqueErr *storage.UniqueViolationError
+	var notNullErr *storage.NotNullViolationError
+	var checkErr *storage.CheckViolationError
+	var foreignKeyErr *storage.ForeignKeyViolationError
+	if errors.Is(err, storage.ErrSerialization) {
+		code = ErrCodeSerializationFailure
+		c.txStatus = TxStatusIdle
+	} else if errors.Is(err, pgerror.ErrInvalidText) {
+		code = ErrCodeInvalidText
+		if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
+		}
+	} else if errors.Is(err, pgerror.ErrInvalidParameter) {
+		code = ErrCodeInvalidParameter
+		if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
+		}
+	} else if errors.Is(err, pgerror.ErrDatatypeMismatch) {
+		code = ErrCodeDatatypeMismatch
+		if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
+		}
+	} else if analysisTypeMismatch(err) {
+		code = ErrCodeDatatypeMismatch
+		if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
+		}
+	} else if errors.As(err, &notNullErr) {
+		code = "23502"
+		if _, commit := stmt.(*parser.CommitStmt); commit {
 			c.txStatus = TxStatusIdle
-		} else if errors.Is(err, pgerror.ErrInvalidText) {
-			code = ErrCodeInvalidText
-			if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if errors.Is(err, pgerror.ErrInvalidParameter) {
-			code = ErrCodeInvalidParameter
-			if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if errors.Is(err, pgerror.ErrDatatypeMismatch) {
-			code = ErrCodeDatatypeMismatch
-			if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if analysisTypeMismatch(err) {
-			code = ErrCodeDatatypeMismatch
-			if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if errors.As(err, &notNullErr) {
-			code = "23502"
-			if _, commit := stmt.(*parser.CommitStmt); commit {
-				c.txStatus = TxStatusIdle
-			} else if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if errors.As(err, &checkErr) {
-			code = ErrCodeCheckViolation
-			if _, commit := stmt.(*parser.CommitStmt); commit {
-				c.txStatus = TxStatusIdle
-			} else if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if errors.As(err, &foreignKeyErr) {
-			code = ErrCodeForeignKeyViolation
-			if _, commit := stmt.(*parser.CommitStmt); commit {
-				c.txStatus = TxStatusIdle
-			} else if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if errors.As(err, &uniqueErr) || errors.As(err, &primaryKeyErr) {
-			code = ErrCodeUniqueViolation
-			if _, commit := stmt.(*parser.CommitStmt); commit {
-				c.txStatus = TxStatusIdle
-			} else if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if errors.Is(err, pgerror.ErrFeatureNotSupported) || errors.Is(err, executor.ErrUnsupportedFeature) {
-			code = ErrCodeFeatureNotSupported
-			if c.txStatus == TxStatusInBlock {
-				c.txStatus = TxStatusFailed
-			}
-		} else if _, commit := stmt.(*parser.CommitStmt); commit {
+		} else if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
+		}
+	} else if errors.As(err, &checkErr) {
+		code = ErrCodeCheckViolation
+		if _, commit := stmt.(*parser.CommitStmt); commit {
+			c.txStatus = TxStatusIdle
+		} else if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
+		}
+	} else if errors.As(err, &foreignKeyErr) {
+		code = ErrCodeForeignKeyViolation
+		if _, commit := stmt.(*parser.CommitStmt); commit {
 			c.txStatus = TxStatusIdle
 		} else if c.txStatus == TxStatusInBlock {
 			c.txStatus = TxStatusFailed
 		}
-		if uniqueErr != nil {
-			return c.failExtended(code, fmt.Errorf("execution error: UNIQUE constraint failed: %s", uniqueErr.Constraint))
+	} else if errors.As(err, &uniqueErr) || errors.As(err, &primaryKeyErr) {
+		code = ErrCodeUniqueViolation
+		if _, commit := stmt.(*parser.CommitStmt); commit {
+			c.txStatus = TxStatusIdle
+		} else if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
+		}
+	} else if errors.Is(err, pgerror.ErrFeatureNotSupported) || errors.Is(err, executor.ErrUnsupportedFeature) {
+		code = ErrCodeFeatureNotSupported
+		if c.txStatus == TxStatusInBlock {
+			c.txStatus = TxStatusFailed
 		}
-		return c.failExtended(code, fmt.Errorf("execution error: %w", err))
+	} else if _, commit := stmt.(*parser.CommitStmt); commit {
+		c.txStatus = TxStatusIdle
+	} else if c.txStatus == TxStatusInBlock {
+		c.txStatus = TxStatusFailed
 	}
-	columnTypes := wireColumnTypes(result)
-	formats, formatErr := c.resolveResultFormats(portal.resultFormats, columnTypes, len(result.Columns))
-	if formatErr != nil {
-		return c.failResultFormat(formatErr)
+	if uniqueErr != nil {
+		return &executeError{code: code, err: fmt.Errorf("execution error: UNIQUE constraint failed: %s", uniqueErr.Constraint)}
+	}
+	return &executeError{code: code, err: fmt.Errorf("execution error: %w", err)}
+}
+
+// sendPortalRows writes one Execute response from the portal's materialized
+// result. The RowDescription is sent once, by the first Execute that has
+// columns, so a resumed Execute emits only DataRows. With maxRows > 0 the
+// response is PortalSuspended once the limit was reached; PostgreSQL reports
+// PortalSuspended even when the last row falls exactly on the limit and
+// completes on the next Execute, which is what this does too.
+func (c *Connection) sendPortalRows(portal *portal, maxRows int32) error {
+	execution := portal.execution
+	if len(execution.result.Columns) > 0 && !execution.described {
+		if err := c.sendRowDescriptionFormats(execution.result.Columns, execution.columnTypes, execution.oids, execution.formats); err != nil {
+			return err
+		}
+		execution.described = true
+	}
+	remaining := len(execution.result.Rows) - execution.rowOffset
+	limit := remaining
+	if maxRows > 0 && int(maxRows) < remaining {
+		limit = int(maxRows)
+	}
+	for i := 0; i < limit; i++ {
+		if err := c.sendDataRowFormats(execution.result.Rows[execution.rowOffset+i], execution.oids, execution.formats); err != nil {
+			return c.failBinaryEncode(err)
+		}
+	}
+	execution.rowOffset += limit
+	if maxRows > 0 && limit == int(maxRows) {
+		return c.sendPortalSuspended()
 	}
-	if err := c.sendResult(result, stmt, columnTypes, formats); err != nil {
+	firstCompletion := !execution.completed
+	execution.completed = true
+	if err := c.sendCommandComplete(execution.commandTag); err != nil {
 		return err
 	}
-	if result.LSN > 0 {
-		return c.sendLSN(result.LSN)
+	if firstCompletion && execution.result.LSN > 0 {
+		if err := c.sendLSN(execution.result.LSN); err != nil {
+			return err
+		}
 	}
+	// The portal is exhausted: keep its metadata for a repeated Execute, but
+	// release the row data so a transaction that leaves portals open does not
+	// retain every result.
+	execution.result.Rows = nil
+	execution.rowOffset = 0
 	return nil
 }
 
@@ -1032,6 +1145,14 @@ func (c *Connection) handleClose(msg *Message) error {
 	return c.writeMessage(MsgCloseComplete, nil)
 }
 
+// clearPortals drops every portal. PostgreSQL binds a portal to the transaction
+// that created it: Sync ends the implicit transaction outside an explicit
+// block, and COMMIT/ROLLBACK end an explicit one. Prepared statements are not
+// affected, so a client re-binds after the transaction ends.
+func (c *Connection) clearPortals() {
+	clear(c.portals)
+}
+
 func (c *Connection) failExtended(code string, err error) error {
 	c.extendedFailed = true
 	if c.txStatus == TxStatusInBlock {
@@ -1408,6 +1529,13 @@ func isIdleRollback(stmt parser.Statement, status byte, err error) bool {
 
 // sendReadyForQuery sends ready for query message
 func (c *Connection) sendReadyForQuery() error {
+	if c.txStatus == TxStatusIdle {
+		// The transaction this cycle ran in ends here, so its portals do too.
+		// A portal inside an explicit block survives (status T) until
+		// COMMIT/ROLLBACK, and a failed block (status E) keeps its portals so
+		// Execute still reports 25P02 until the transaction is rolled back.
+		c.clearPortals()
+	}
 	mb := NewMessageBuilder()
 	mb.AppendByte(c.txStatus)
 	if err := c.writeMessage(MsgReadyForQuery, mb.Bytes()); err != nil {
@@ -1530,6 +1658,12 @@ func (c *Connection) sendCommandComplete(tag string) error {
 	return c.writeMessage(MsgCommandComplete, mb.Bytes())
 }
 
+// sendPortalSuspended reports that a bounded Execute stopped at its row limit
+// with rows left to fetch.
+func (c *Connection) sendPortalSuspended() error {
+	return c.writeMessage(MsgPortalSuspended, nil)
+}
+
 // sendError sends an error response
 func (c *Connection) sendError(severity, code, message string) error {
 	if severity == "ERROR" && c.txStatus == TxStatusInBlock {

+ 452 - 0
pkg/pgserver/portal_test.go

@@ -0,0 +1,452 @@
+package pgserver
+
+import (
+	"bufio"
+	"bytes"
+	"strconv"
+	"testing"
+)
+
+// Bounded Execute is only reachable over the raw protocol. pgx v5.11.0 sends
+// maxRows = 0 for every Execute (pgconn.ExecParams, ExecPrepared, and the
+// Pipeline equivalents have no row-limit argument), and database/sql drives pgx,
+// so neither a pgx nor a database/sql test can request a suspended portal. These
+// tests therefore drive the protocol messages directly, one round trip per
+// Execute.
+
+// portalSession sends extended-protocol messages to a connection backed by a
+// real executor and returns the frames the client would receive for that round
+// trip.
+type portalSession struct {
+	t    *testing.T
+	c    *Connection
+	conn *countingConn
+}
+
+func newPortalSession(t *testing.T) *portalSession {
+	t.Helper()
+	c, _ := newDataConnection(t)
+	conn := &countingConn{}
+	c.conn = conn
+	c.writer = bufio.NewWriter(conn)
+	return &portalSession{t: t, c: c, conn: conn}
+}
+
+// send handles the messages, flushes buffered output the way a client's Flush
+// message would, and returns the frames produced since the previous send.
+func (s *portalSession) send(messages ...*Message) []*Message {
+	s.t.Helper()
+	for _, message := range messages {
+		if err := s.c.handleMessage(message); err != nil {
+			s.t.Fatalf("handle %c: %v", message.Type, err)
+		}
+	}
+	if err := s.c.handleMessage(&Message{Type: MsgFlush}); err != nil {
+		s.t.Fatalf("flush: %v", err)
+	}
+	frames := readAllMessages(s.t, s.conn.bytes())
+	s.conn.reset()
+	return frames
+}
+
+// setup runs SQL outside the protocol; the statements are the test fixture, not
+// the behavior under test.
+func (s *portalSession) setup(statements ...string) {
+	s.t.Helper()
+	for _, sql := range statements {
+		if _, err := s.c.executor.Execute(parseStmt(s.t, sql)); err != nil {
+			s.t.Fatalf("%s: %v", sql, err)
+		}
+	}
+}
+
+// countRows reports how many rows a table holds, so a test can prove a DML
+// portal ran once.
+func (s *portalSession) countRows(table string) int64 {
+	s.t.Helper()
+	result, err := s.c.executor.Execute(parseStmt(s.t, "SELECT COUNT(*) FROM "+table))
+	if err != nil {
+		s.t.Fatalf("count %s: %v", table, err)
+	}
+	return result.Rows[0][0].(int64)
+}
+
+func executeRowsMessage(portal string, maxRows int32) *Message {
+	mb := NewMessageBuilder()
+	mb.WriteString(portal)
+	mb.WriteInt32(maxRows)
+	return &Message{Type: MsgExecute, Data: mb.Bytes()}
+}
+
+func closePortalMessage(portal string) *Message {
+	return &Message{Type: MsgClose, Data: append([]byte("P"+portal), 0)}
+}
+
+func queryMessageData(sql string) *Message {
+	return &Message{Type: MsgQuery, Data: append([]byte(sql), 0)}
+}
+
+func syncMessage() *Message {
+	return &Message{Type: MsgSync}
+}
+
+func commandTagOf(message *Message) string {
+	if message == nil || message.Type != MsgCommandComplete {
+		return ""
+	}
+	if end := bytes.IndexByte(message.Data, 0); end >= 0 {
+		return string(message.Data[:end])
+	}
+	return string(message.Data)
+}
+
+func assertFrameTypes(t *testing.T, frames []*Message, want string) {
+	t.Helper()
+	if got := messageTypes(frames); got != want {
+		t.Fatalf("frames = %q, want %q", got, want)
+	}
+}
+
+func assertRowIDs(t *testing.T, frames []*Message, want ...string) {
+	t.Helper()
+	rows := dataRowsRaw(t, frames)
+	if len(rows) != len(want) {
+		t.Fatalf("got %d rows, want %d", len(rows), len(want))
+	}
+	for i, id := range want {
+		if len(rows[i]) != 1 || string(rows[i][0]) != id {
+			t.Fatalf("row %d = %q, want %q", i, rows[i], id)
+		}
+	}
+}
+
+func readyStatus(t *testing.T, frames []*Message) byte {
+	t.Helper()
+	for _, message := range frames {
+		if message.Type == MsgReadyForQuery {
+			if len(message.Data) != 1 {
+				t.Fatalf("ReadyForQuery payload = %q", message.Data)
+			}
+			return message.Data[0]
+		}
+	}
+	t.Fatalf("no ReadyForQuery in %q", messageTypes(frames))
+	return 0
+}
+
+const portalSelectSQL = "SELECT id FROM t ORDER BY id"
+
+func createNumberedTable(t *testing.T, s *portalSession, rows int) {
+	t.Helper()
+	s.setup("CREATE TABLE t (id INTEGER PRIMARY KEY)")
+	for i := 1; i <= rows; i++ {
+		s.setup("INSERT INTO t (id) VALUES (" + strconv.Itoa(i) + ")")
+	}
+}
+
+// A bounded Execute returns at most maxRows rows and suspends; each later
+// Execute resumes where the previous one stopped. The RowDescription is sent by
+// the first Execute only, and Describe('P') still answers NoData without
+// disturbing that sequence.
+func TestBoundedExecuteSuspendsAndResumes(t *testing.T) {
+	s := newPortalSession(t)
+	createNumberedTable(t, s, 5)
+
+	frames := s.send(
+		parseMessage("stmt", portalSelectSQL),
+		bindResultMessage("p", "stmt", nil, nil),
+		&Message{Type: MsgDescribe, Data: []byte("Pp\x00")},
+	)
+	assertFrameTypes(t, frames, "12n")
+
+	first := s.send(executeRowsMessage("p", 2))
+	assertFrameTypes(t, first, "TDDs")
+	assertRowIDs(t, first, "1", "2")
+
+	second := s.send(executeRowsMessage("p", 2))
+	assertFrameTypes(t, second, "DDs")
+	assertRowIDs(t, second, "3", "4")
+
+	third := s.send(executeRowsMessage("p", 2))
+	assertFrameTypes(t, third, "DC")
+	assertRowIDs(t, third, "5")
+	if tag := commandTagOf(third[len(third)-1]); tag != "SELECT 5" {
+		t.Fatalf("completion tag = %q, want %q", tag, "SELECT 5")
+	}
+
+	// A completed portal repeats its tag without re-running the statement.
+	fourth := s.send(executeRowsMessage("p", 0))
+	assertFrameTypes(t, fourth, "C")
+	if tag := commandTagOf(fourth[0]); tag != "SELECT 5" {
+		t.Fatalf("repeated completion tag = %q, want %q", tag, "SELECT 5")
+	}
+	// Completion drops the row data, so an exhausted portal only keeps the
+	// metadata a repeated Execute answers with.
+	if rows := s.c.portals["p"].execution.result.Rows; len(rows) != 0 {
+		t.Fatalf("completed portal still holds %d rows", len(rows))
+	}
+}
+
+// maxRows = 0 keeps the pre-existing meaning: every remaining row in one
+// Execute followed by CommandComplete.
+func TestExecuteUnboundedStreamsAllRows(t *testing.T) {
+	s := newPortalSession(t)
+	createNumberedTable(t, s, 3)
+	s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
+
+	frames := s.send(executeRowsMessage("p", 0))
+	assertFrameTypes(t, frames, "TDDDC")
+	assertRowIDs(t, frames, "1", "2", "3")
+	if tag := commandTagOf(frames[len(frames)-1]); tag != "SELECT 3" {
+		t.Fatalf("completion tag = %q, want %q", tag, "SELECT 3")
+	}
+}
+
+// A limit that falls exactly on the last row still suspends, as PostgreSQL
+// does; the following Execute reports completion with no further rows.
+func TestBoundedExecuteExactLimitSuspends(t *testing.T) {
+	s := newPortalSession(t)
+	createNumberedTable(t, s, 3)
+	s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
+
+	frames := s.send(executeRowsMessage("p", 3))
+	assertFrameTypes(t, frames, "TDDDs")
+	assertRowIDs(t, frames, "1", "2", "3")
+
+	frames = s.send(executeRowsMessage("p", 0))
+	assertFrameTypes(t, frames, "C")
+	if tag := commandTagOf(frames[0]); tag != "SELECT 3" {
+		t.Fatalf("completion tag = %q, want %q", tag, "SELECT 3")
+	}
+}
+
+// Close('P') destroys a suspended portal, and a later Execute reports 34000
+// until Sync recovers the cycle.
+func TestCloseEndsSuspendedPortal(t *testing.T) {
+	s := newPortalSession(t)
+	createNumberedTable(t, s, 3)
+	s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
+	assertFrameTypes(t, s.send(executeRowsMessage("p", 1)), "TDs")
+
+	assertFrameTypes(t, s.send(closePortalMessage("p")), "3")
+
+	frames := s.send(executeRowsMessage("p", 1))
+	assertFrameTypes(t, frames, "E")
+	if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
+		t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeUndefinedCursor)
+	}
+	assertFrameTypes(t, s.send(syncMessage()), "Z")
+	if s.c.extendedFailed {
+		t.Fatal("Sync did not clear the extended-failure state")
+	}
+}
+
+// Re-binding a portal name replaces the materialized result, the row cursor,
+// and the row metadata: the stream starts over.
+func TestRebindRestartsPortal(t *testing.T) {
+	s := newPortalSession(t)
+	createNumberedTable(t, s, 5)
+	s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
+	assertFrameTypes(t, s.send(executeRowsMessage("p", 2)), "TDDs")
+
+	frames := s.send(
+		bindResultMessage("p", "stmt", nil, nil),
+		executeRowsMessage("p", 0),
+	)
+	assertFrameTypes(t, frames, "2TDDDDDC")
+	assertRowIDs(t, frames, "1", "2", "3", "4", "5")
+}
+
+// Inside an explicit transaction a portal survives Sync and keeps resuming;
+// COMMIT and ROLLBACK end the transaction and with it every portal.
+func TestPortalLivesUntilTransactionEnd(t *testing.T) {
+	for _, ending := range []string{"COMMIT", "ROLLBACK"} {
+		t.Run(ending, func(t *testing.T) {
+			s := newPortalSession(t)
+			createNumberedTable(t, s, 3)
+
+			if status := readyStatus(t, s.send(queryMessageData("BEGIN"))); status != TxStatusInBlock {
+				t.Fatalf("BEGIN status = %c, want %c", status, TxStatusInBlock)
+			}
+			frames := s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
+			assertFrameTypes(t, frames, "12")
+
+			frames = s.send(executeRowsMessage("p", 2), syncMessage())
+			assertFrameTypes(t, frames, "TDDsZ")
+			if status := readyStatus(t, frames); status != TxStatusInBlock {
+				t.Fatalf("Sync inside the transaction reported status %c, want %c", status, TxStatusInBlock)
+			}
+
+			frames = s.send(executeRowsMessage("p", 0), syncMessage())
+			assertFrameTypes(t, frames, "DCZ")
+			assertRowIDs(t, frames, "3")
+
+			if status := readyStatus(t, s.send(queryMessageData(ending))); status != TxStatusIdle {
+				t.Fatalf("%s status = %c, want %c", ending, status, TxStatusIdle)
+			}
+			frames = s.send(executeRowsMessage("p", 1))
+			assertFrameTypes(t, frames, "E")
+			if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
+				t.Fatalf("SQLSTATE after %s = %q, want %q", ending, code, ErrCodeUndefinedCursor)
+			}
+			s.send(syncMessage())
+		})
+	}
+}
+
+// Outside an explicit block the portal belongs to the implicit transaction that
+// Sync ends, so Execute afterwards reports 34000 (PostgreSQL behaves the same);
+// the prepared statement survives and can be bound again.
+func TestPortalEndsAtSyncOutsideTransaction(t *testing.T) {
+	s := newPortalSession(t)
+	createNumberedTable(t, s, 3)
+	frames := s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
+	assertFrameTypes(t, frames, "12")
+
+	frames = s.send(executeRowsMessage("p", 1), syncMessage())
+	assertFrameTypes(t, frames, "TDsZ")
+
+	frames = s.send(executeRowsMessage("p", 1))
+	assertFrameTypes(t, frames, "E")
+	if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
+		t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeUndefinedCursor)
+	}
+	s.send(syncMessage())
+
+	frames = s.send(bindResultMessage("p", "stmt", nil, nil), executeRowsMessage("p", 0))
+	assertFrameTypes(t, frames, "2TDDDC")
+	assertRowIDs(t, frames, "1", "2", "3")
+}
+
+// A suspended INSERT ... RETURNING runs the statement exactly once: the resume
+// streams the prefetched rows, and a completed portal never runs it again.
+func TestBoundedExecuteDMLEvaluatesOnce(t *testing.T) {
+	s := newPortalSession(t)
+	s.setup("CREATE TABLE t (id INTEGER PRIMARY KEY)")
+	frames := s.send(
+		parseMessage("stmt", "INSERT INTO t (id) VALUES (1), (2), (3) RETURNING id"),
+		bindResultMessage("p", "stmt", nil, nil),
+	)
+	assertFrameTypes(t, frames, "12")
+
+	frames = s.send(executeRowsMessage("p", 2))
+	assertFrameTypes(t, frames, "TDDs")
+	assertRowIDs(t, frames, "1", "2")
+
+	frames = s.send(executeRowsMessage("p", 0))
+	// The autocommit INSERT publishes its commit LSN with the completion.
+	assertFrameTypes(t, frames, "DCS")
+	assertRowIDs(t, frames, "3")
+	if tag := commandTagOf(frames[len(frames)-2]); tag != "INSERT 0 3" {
+		t.Fatalf("completion tag = %q, want %q", tag, "INSERT 0 3")
+	}
+
+	frames = s.send(executeRowsMessage("p", 0))
+	assertFrameTypes(t, frames, "C")
+	if rows := s.countRows("t"); rows != 3 {
+		t.Fatalf("table has %d rows after three Executes, want 3", rows)
+	}
+}
+
+// Per-portal result formats keep applying on every Execute of the same portal,
+// and only the first Execute carries the RowDescription.
+func TestPortalResumeKeepsBinaryResultFormats(t *testing.T) {
+	s := newPortalSession(t)
+	s.setup(
+		"CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)",
+		"INSERT INTO t VALUES (1, 'one')",
+		"INSERT INTO t VALUES (2, 'two')",
+	)
+	frames := s.send(
+		parseMessage("stmt", "SELECT id, name FROM t ORDER BY id"),
+		bindResultMessage("p", "stmt", nil, []int16{1}),
+	)
+	assertFrameTypes(t, frames, "12")
+
+	frames = s.send(executeRowsMessage("p", 1))
+	assertFrameTypes(t, frames, "TDs")
+	descriptions := 0
+	for _, message := range frames {
+		if message.Type == MsgRowDescription {
+			descriptions++
+			formats := rowDescriptionFormats(t, message.Data)
+			if len(formats) != 2 || formats[0] != 1 || formats[1] != 1 {
+				t.Fatalf("RowDescription formats = %v, want binary for both columns", formats)
+			}
+		}
+	}
+	if descriptions != 1 {
+		t.Fatalf("got %d RowDescriptions, want 1", descriptions)
+	}
+	rows := dataRowsRaw(t, frames)
+	assertBytes(t, "first id", rows[0][0], []byte{0, 0, 0, 1})
+
+	frames = s.send(executeRowsMessage("p", 1))
+	assertFrameTypes(t, frames, "Ds")
+	rows = dataRowsRaw(t, frames)
+	assertBytes(t, "second id", rows[0][0], []byte{0, 0, 0, 2})
+
+	frames = s.send(executeRowsMessage("p", 1))
+	assertFrameTypes(t, frames, "C")
+	if tag := commandTagOf(frames[0]); tag != "SELECT 2" {
+		t.Fatalf("completion tag = %q, want %q", tag, "SELECT 2")
+	}
+}
+
+// Execute of a portal that was never bound is 34000, like PostgreSQL's
+// "portal does not exist".
+func TestExecuteUnknownPortalIsUndefinedCursor(t *testing.T) {
+	s := newPortalSession(t)
+	frames := s.send(executeRowsMessage("missing", 0))
+	assertFrameTypes(t, frames, "E")
+	if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
+		t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeUndefinedCursor)
+	}
+	s.send(syncMessage())
+}
+
+// A failed Execute leaves the transaction aborted: a later Execute yields 25P02
+// until ROLLBACK, which also destroys the portals.
+func TestFailedExecuteLeavesTransactionConsistent(t *testing.T) {
+	s := newPortalSession(t)
+	s.setup("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT NOT NULL)")
+	if status := readyStatus(t, s.send(queryMessageData("BEGIN"))); status != TxStatusInBlock {
+		t.Fatalf("BEGIN status = %c, want %c", status, TxStatusInBlock)
+	}
+
+	s.send(
+		parseMessage("failing", "INSERT INTO t (id, name) VALUES (1, NULL)"),
+		bindResultMessage("bad", "failing", nil, nil),
+	)
+	frames := s.send(executeRowsMessage("bad", 0))
+	assertFrameTypes(t, frames, "E")
+	if code := errorCode(findError(frames)); code != "23502" {
+		t.Fatalf("SQLSTATE = %q, want 23502", code)
+	}
+	if status := readyStatus(t, s.send(syncMessage())); status != TxStatusFailed {
+		t.Fatalf("status after failure = %c, want %c", status, TxStatusFailed)
+	}
+
+	s.send(
+		parseMessage("reading", portalSelectSQL),
+		bindResultMessage("ok", "reading", nil, nil),
+	)
+	frames = s.send(executeRowsMessage("ok", 0))
+	assertFrameTypes(t, frames, "E")
+	if code := errorCode(findError(frames)); code != ErrCodeTransactionAborted {
+		t.Fatalf("SQLSTATE in the aborted transaction = %q, want %q", code, ErrCodeTransactionAborted)
+	}
+	if status := readyStatus(t, s.send(syncMessage())); status != TxStatusFailed {
+		t.Fatalf("status after aborted Execute = %c, want %c", status, TxStatusFailed)
+	}
+
+	if status := readyStatus(t, s.send(queryMessageData("ROLLBACK"))); status != TxStatusIdle {
+		t.Fatalf("ROLLBACK status = %c, want %c", status, TxStatusIdle)
+	}
+	frames = s.send(executeRowsMessage("ok", 0))
+	assertFrameTypes(t, frames, "E")
+	if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
+		t.Fatalf("SQLSTATE after ROLLBACK = %q, want %q", code, ErrCodeUndefinedCursor)
+	}
+	s.send(syncMessage())
+}

+ 2 - 0
pkg/pgserver/protocol.go

@@ -46,6 +46,7 @@ const (
 	MsgNoticeResponse       = 'N' // Notice response
 	MsgParameterDescription = 't' // Parameter description
 	MsgParameterStatus      = 'S' // Parameter status
+	MsgPortalSuspended      = 's' // Portal suspended
 	MsgParseComplete        = '1' // Parse complete
 	MsgReadyForQuery        = 'Z' // Ready for query
 	MsgRowDescription       = 'T' // Row description
@@ -96,6 +97,7 @@ const (
 	ErrCodeConnectionFailure    = "08006"
 	ErrCodeProtocolViolation    = "08P01"
 	ErrCodeFeatureNotSupported  = "0A000"
+	ErrCodeUndefinedCursor      = "34000"
 	ErrCodeTransactionAborted   = "25P02"
 	ErrCodeSerializationFailure = "40001"
 	ErrCodeUniqueViolation      = "23505"

+ 0 - 1
pkg/pgserver/protocol_guards_test.go

@@ -42,7 +42,6 @@ func TestUnsupportedExtendedRequestsFailBeforeSideEffects(t *testing.T) {
 		{"invalid_result_format", &Message{Type: MsgBind, Data: guardBind(nil, []int16{2})}, ErrCodeProtocolViolation},
 		{"invalid_parameter_format", &Message{Type: MsgBind, Data: guardBind([]int16{2}, nil)}, ErrCodeProtocolViolation},
 		{"parameter_format_count", &Message{Type: MsgBind, Data: guardBind([]int16{0, 0}, nil)}, ErrCodeProtocolViolation},
-		{"bounded_execute", &Message{Type: MsgExecute, Data: []byte{'p', 0, 0, 0, 0, 1}}, ErrCodeFeatureNotSupported},
 		{"negative_execute", &Message{Type: MsgExecute, Data: []byte{'p', 0, 255, 255, 255, 255}}, ErrCodeProtocolViolation},
 	} {
 		t.Run(tc.name, func(t *testing.T) {

+ 7 - 3
pkg/pgserver/result_formats.go

@@ -75,12 +75,16 @@ func (c *Connection) resolveResultFormats(bindFormats []int16, columnTypes []str
 // connection sends an ErrorResponse, discards the rest of the cycle until Sync,
 // and marks an explicit transaction block as failed.
 func (c *Connection) failResultFormat(err error) error {
-	code := ErrCodeProtocolViolation
+	return c.failExtended(resultFormatSQLState(err), err)
+}
+
+// resultFormatSQLState reports the SQLSTATE a result-format rejection carries.
+func resultFormatSQLState(err error) string {
 	var formatErr *resultFormatError
 	if errors.As(err, &formatErr) {
-		code = formatErr.code
+		return formatErr.code
 	}
-	return c.failExtended(code, err)
+	return ErrCodeProtocolViolation
 }
 
 // typeNameAt returns the declared type name of a result column, or "" when the

Энэ ялгаанд хэт олон файл өөрчлөгдсөн тул зарим файлыг харуулаагүй болно