Sfoglia il codice sorgente

Merge pull request #42 from database-pizza/engine/c-portals

add resumable portals and bounded Execute
Danilo Fragoso 1 giorno fa
parent
commit
123f1fa671

File diff suppressed because it is too large
+ 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'), the end of its transaction (Sync outside an explicit block, COMMIT/ROLLBACK inside one), or a ROLLBACK TO SAVEPOINT, which also drops portals created before the savepoint (stricter than PostgreSQL), after which Execute fails with 34000 (inside an aborted transaction, 25P02); a Bind replaces the unnamed portal but re-binding a named portal that already exists fails with 42P03; 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')", "ROLLBACK TO SAVEPOINT", "42P03", "34000", "25P02", "08P01", "RETURNING",
+	} {
+		if !strings.Contains(bounded.Notes, claim) {
+			t.Errorf("bounded_execute notes do not mention %q", claim)
+		}
+	}
+}

+ 257 - 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
+	// clearsPortals reports that this statement invalidated every open portal
+	// (COMMIT, ROLLBACK, or ROLLBACK TO SAVEPOINT); see resetsPortals.
+	clearsPortals 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
@@ -509,6 +539,13 @@ func (c *Connection) handleQuery(msg *Message) error {
 				return err
 			}
 		}
+		if resetsPortals(stmt) {
+			// Same rule as the extended path: a transaction end or a savepoint
+			// rollback invalidates every open portal. sendReadyForQuery covers
+			// the idle case; ROLLBACK TO SAVEPOINT keeps the block open and
+			// needs this explicit clear.
+			c.clearPortals()
+		}
 	}
 
 	return c.sendReadyForQuery()
@@ -563,6 +600,14 @@ func (c *Connection) handleBind(msg *Message) error {
 	if !ok {
 		return c.failExtended(ErrCodeInvalidParameter, fmt.Errorf("prepared statement %q does not exist", statementName))
 	}
+	// Named portals cannot be replaced: PostgreSQL reports 42P03 and leaves the
+	// existing portal in place. Only the unnamed portal is re-bound freely,
+	// because the new Bind replaces it.
+	if portalName != "" {
+		if _, exists := c.portals[portalName]; exists {
+			return c.failExtended(ErrCodeDuplicateCursor, fmt.Errorf("portal %q already exists", portalName))
+		}
+	}
 	formats, pos, err := readInt16List(msg.Data, pos)
 	if err != nil || pos+2 > len(msg.Data) {
 		return c.failExtended(ErrCodeProtocolViolation, fmt.Errorf("invalid Bind format list"))
@@ -669,14 +714,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 +726,196 @@ 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)
+		}
 	}
-	portal.executed = true
+	if err := c.sendPortalRows(portal, maxRows); err != nil {
+		return err
+	}
+	if portal.execution.clearsPortals {
+		c.clearPortals()
+	}
+	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)
 		}
-		return c.failExtended(code, fmt.Errorf("syntax error: %w", 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)
+		execution.clearsPortals = resetsPortals(stmt)
 	}
-	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
+		}
+	} 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
 		}
-		if uniqueErr != nil {
-			return c.failExtended(code, fmt.Errorf("execution error: UNIQUE constraint failed: %s", uniqueErr.Constraint))
+	} 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
 	}
-	if err := c.sendResult(result, stmt, columnTypes, formats); err != nil {
+	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()
+	}
+	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 +1154,29 @@ 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)
+}
+
+// resetsPortals reports whether stmt invalidates every open portal. COMMIT and
+// ROLLBACK end the transaction, and ROLLBACK TO SAVEPOINT undoes writes a
+// suspended portal may already have materialized: resuming it would stream rows
+// and a completion tag from a removed state. Clearing every portal on ROLLBACK
+// TO SAVEPOINT is stricter than PostgreSQL, which keeps portals created before
+// the savepoint; the coarse rule is what keeps a failed or rewound statement
+// from leaking into a resumed stream.
+func resetsPortals(stmt parser.Statement) bool {
+	switch stmt.(type) {
+	case *parser.CommitStmt, *parser.RollbackStmt:
+		return true
+	}
+	return false
+}
+
 func (c *Connection) failExtended(code string, err error) error {
 	c.extendedFailed = true
 	if c.txStatus == TxStatusInBlock {
@@ -1408,6 +1553,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 +1682,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 {

+ 606 - 0
pkg/pgserver/portal_test.go

@@ -0,0 +1,606 @@
+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")
+	}
+}
+
+// The unnamed portal is replaced by every Bind, which restarts its materialized
+// result, row cursor, and row metadata.
+func TestUnnamedRebindRestartsPortal(t *testing.T) {
+	s := newPortalSession(t)
+	createNumberedTable(t, s, 5)
+	s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("", "stmt", nil, nil))
+	assertFrameTypes(t, s.send(executeRowsMessage("", 2)), "TDDs")
+
+	frames := s.send(bindResultMessage("", "stmt", nil, nil), executeRowsMessage("", 0))
+	assertFrameTypes(t, frames, "2TDDDDDC")
+	assertRowIDs(t, frames, "1", "2", "3", "4", "5")
+}
+
+// A named portal cannot be replaced by another Bind: PostgreSQL reports 42P03,
+// and the rejected Bind fails the cycle, so an explicit transaction is aborted
+// and the portal goes with the ROLLBACK.
+func TestNamedRebindIsRejected(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)
+	}
+	s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
+	assertFrameTypes(t, s.send(executeRowsMessage("p", 1)), "TDs")
+
+	frames := s.send(bindResultMessage("p", "stmt", nil, nil))
+	assertFrameTypes(t, frames, "E")
+	if code := errorCode(findError(frames)); code != ErrCodeDuplicateCursor {
+		t.Fatalf("SQLSTATE = %q, want %q", code, ErrCodeDuplicateCursor)
+	}
+	if status := readyStatus(t, s.send(syncMessage())); status != TxStatusFailed {
+		t.Fatalf("status after rejected re-Bind = %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("p", 1))
+	assertFrameTypes(t, frames, "E")
+	if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
+		t.Fatalf("SQLSTATE after ROLLBACK = %q, want %q", code, ErrCodeUndefinedCursor)
+	}
+	s.send(syncMessage())
+}
+
+// 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())
+}
+
+// A ROLLBACK TO SAVEPOINT can undo writes a suspended portal already
+// materialized, so it drops every portal. Resuming one would stream rows and a
+// completion tag from a removed state; the client gets 34000 instead. Clearing
+// on ROLLBACK TO SAVEPOINT is stricter than PostgreSQL, which keeps portals
+// created before the savepoint.
+func TestRollbackToSavepointClearsSuspendedPortals(t *testing.T) {
+	s := newPortalSession(t)
+	s.setup("CREATE TABLE t (id INTEGER PRIMARY KEY)")
+	if status := readyStatus(t, s.send(queryMessageData("BEGIN"))); status != TxStatusInBlock {
+		t.Fatalf("BEGIN status = %c, want %c", status, TxStatusInBlock)
+	}
+	if status := readyStatus(t, s.send(queryMessageData("SAVEPOINT s"))); status != TxStatusInBlock {
+		t.Fatalf("SAVEPOINT status = %c, want %c", status, TxStatusInBlock)
+	}
+	s.send(
+		parseMessage("stmt", "INSERT INTO t (id) VALUES (1), (2) RETURNING id"),
+		bindResultMessage("p", "stmt", nil, nil),
+	)
+	frames := s.send(executeRowsMessage("p", 1))
+	assertFrameTypes(t, frames, "TDs")
+	assertRowIDs(t, frames, "1")
+
+	if status := readyStatus(t, s.send(queryMessageData("ROLLBACK TO SAVEPOINT s"))); status != TxStatusInBlock {
+		t.Fatalf("ROLLBACK TO SAVEPOINT status = %c, want %c", status, TxStatusInBlock)
+	}
+
+	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())
+
+	if rows := s.countRows("t"); rows != 0 {
+		t.Fatalf("table has %d rows after ROLLBACK TO SAVEPOINT, want 0", rows)
+	}
+	if status := readyStatus(t, s.send(queryMessageData("ROLLBACK"))); status != TxStatusIdle {
+		t.Fatalf("ROLLBACK status = %c, want %c", status, TxStatusIdle)
+	}
+}
+
+// Transaction commands executed through a portal emit their tag once and drop
+// every open portal, so a suspended sibling can never replay rows or a
+// completion tag from a transaction that ended or was rewound.
+func TestTransactionCommandPortalClearsSiblings(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		command    string
+		wantTag    string
+		wantStatus byte
+	}{
+		{name: "commit", command: "COMMIT", wantTag: "COMMIT", wantStatus: TxStatusIdle},
+		{name: "rollback", command: "ROLLBACK", wantTag: "ROLLBACK", wantStatus: TxStatusIdle},
+		{name: "rollback_to_savepoint", command: "ROLLBACK TO SAVEPOINT s", wantTag: "ROLLBACK", wantStatus: TxStatusInBlock},
+	} {
+		t.Run(tc.name, 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)
+			}
+			if status := readyStatus(t, s.send(queryMessageData("SAVEPOINT s"))); status != TxStatusInBlock {
+				t.Fatalf("SAVEPOINT status = %c, want %c", status, TxStatusInBlock)
+			}
+			s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("first", "stmt", nil, nil))
+			assertFrameTypes(t, s.send(executeRowsMessage("first", 1)), "TDs")
+			s.send(bindResultMessage("second", "stmt", nil, nil))
+			assertFrameTypes(t, s.send(executeRowsMessage("second", 1)), "TDs")
+
+			s.send(parseMessage("txn", tc.command), bindResultMessage("txn_portal", "txn", nil, nil))
+			frames := s.send(executeRowsMessage("txn_portal", 0))
+			assertFrameTypes(t, frames, "C")
+			if tag := commandTagOf(frames[0]); tag != tc.wantTag {
+				t.Fatalf("tag = %q, want %q", tag, tc.wantTag)
+			}
+			if status := readyStatus(t, s.send(syncMessage())); status != tc.wantStatus {
+				t.Fatalf("status after %s = %c, want %c", tc.command, status, tc.wantStatus)
+			}
+
+			for _, name := range []string{"first", "second", "txn_portal"} {
+				frames = s.send(executeRowsMessage(name, 1))
+				assertFrameTypes(t, frames, "E")
+				if code := errorCode(findError(frames)); code != ErrCodeUndefinedCursor {
+					t.Fatalf("Execute %q after %s: SQLSTATE = %q, want %q", name, tc.command, code, ErrCodeUndefinedCursor)
+				}
+				s.send(syncMessage())
+			}
+		})
+	}
+}
+
+// BEGIN through a portal starts the block and emits its tag. It does not end a
+// transaction, so sibling portals are kept, and a repeated Execute repeats the
+// stored tag without replaying BEGIN (the engine rejects a second BEGIN, so a
+// replay would surface as an error).
+func TestBeginPortalStartsTransaction(t *testing.T) {
+	s := newPortalSession(t)
+	createNumberedTable(t, s, 3)
+	s.send(parseMessage("begin_stmt", "BEGIN"), bindResultMessage("begin_portal", "begin_stmt", nil, nil))
+	frames := s.send(executeRowsMessage("begin_portal", 0))
+	assertFrameTypes(t, frames, "C")
+	if tag := commandTagOf(frames[0]); tag != "BEGIN" {
+		t.Fatalf("tag = %q, want %q", tag, "BEGIN")
+	}
+	if status := readyStatus(t, s.send(syncMessage())); status != TxStatusInBlock {
+		t.Fatalf("status after BEGIN = %c, want %c", status, TxStatusInBlock)
+	}
+
+	s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("p", "stmt", nil, nil))
+	assertFrameTypes(t, s.send(executeRowsMessage("p", 1)), "TDs")
+
+	frames = s.send(executeRowsMessage("begin_portal", 0))
+	assertFrameTypes(t, frames, "C")
+	if tag := commandTagOf(frames[0]); tag != "BEGIN" {
+		t.Fatalf("repeated tag = %q, want %q", tag, "BEGIN")
+	}
+
+	frames = s.send(executeRowsMessage("p", 0))
+	assertFrameTypes(t, frames, "DDC")
+	assertRowIDs(t, frames, "2", "3")
+	if status := readyStatus(t, s.send(queryMessageData("ROLLBACK"))); status != TxStatusIdle {
+		t.Fatalf("ROLLBACK status = %c, want %c", status, TxStatusIdle)
+	}
+}

+ 3 - 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
@@ -89,6 +90,7 @@ const (
 	ErrCodeUndefinedColumn      = "42703"
 	ErrCodeDuplicateTable       = "42P07"
 	ErrCodeDuplicateColumn      = "42701"
+	ErrCodeDuplicateCursor      = "42P03"
 	ErrCodeInvalidParameter     = "22023"
 	ErrCodeInvalidText          = "22P02"
 	ErrCodeDatatypeMismatch     = "42804"
@@ -96,6 +98,7 @@ const (
 	ErrCodeConnectionFailure    = "08006"
 	ErrCodeProtocolViolation    = "08P01"
 	ErrCodeFeatureNotSupported  = "0A000"
+	ErrCodeUndefinedCursor      = "34000"
 	ErrCodeTransactionAborted   = "25P02"
 	ErrCodeSerializationFailure = "40001"
 	ErrCodeUniqueViolation      = "23505"

+ 3 - 2
pkg/pgserver/protocol_guards_test.go

@@ -17,9 +17,11 @@ func guardConnection(t *testing.T) (*Connection, *countingConn) {
 	return c, conn
 }
 
+// guardBind targets a portal the cases do not pre-create, so a malformed Bind
+// fails on its own validation rather than with 42P03.
 func guardBind(parameterFormats, resultFormats []int16) []byte {
 	b := NewMessageBuilder()
-	b.WriteString("p")
+	b.WriteString("new_portal")
 	b.WriteString("q")
 	b.WriteInt16(int16(len(parameterFormats)))
 	for _, format := range parameterFormats {
@@ -42,7 +44,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) {

+ 3 - 2
pkg/pgserver/result_binary_test.go

@@ -327,8 +327,8 @@ func TestBinaryResultsMixedPerColumnFormats(t *testing.T) {
 }
 
 // Formats belong to the portal, not the statement: two portals over one
-// prepared statement can use different formats, and re-binding a portal
-// replaces its list.
+// prepared statement can use different formats, and re-binding a closed portal
+// starts from the new list.
 func TestBinaryResultsArePerPortal(t *testing.T) {
 	c, conn := extendedResultConnection(t)
 	prepareBinaryResultTable(t, c)
@@ -339,6 +339,7 @@ func TestBinaryResultsArePerPortal(t *testing.T) {
 		bindResultMessage("binary_portal", "s", nil, []int16{1}),
 		executeMessage("text_portal"),
 		executeMessage("binary_portal"),
+		&Message{Type: MsgClose, Data: append([]byte("Ptext_portal"), 0)},
 		bindResultMessage("text_portal", "s", nil, []int16{0}),
 		executeMessage("text_portal"),
 		&Message{Type: MsgSync},

+ 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

Some files were not shown because too many files changed in this diff