ソースを参照

clear portals on ROLLBACK TO SAVEPOINT and reject named re-Bind

A suspended portal holds a materialized result and completion tag, so a
ROLLBACK TO SAVEPOINT that removes the rows it streamed would let a resumed
Execute replay rolled-back DataRows and a stale tag. Every COMMIT, ROLLBACK, and
ROLLBACK TO SAVEPOINT now clears the portal table, in both the extended path
(via the portal's clearsPortals flag) and the simple-query path (which previously
only cleared when the transaction went idle). Clearing on ROLLBACK TO SAVEPOINT
is stricter than PostgreSQL, which keeps portals created before the savepoint;
the client gets 34000 instead of stale rows.

Bind now matches PostgreSQL's CreatePortal rule: a named portal that already
exists fails with 42P03 and keeps its state, while the unnamed portal is
replaced on every Bind (that is what clients re-bind).

Tests:
- suspended INSERT ... RETURNING + ROLLBACK TO SAVEPOINT -> 34000, no replayed
  rows or tag, and the rolled-back rows are gone (fails without the fix with
  "Ds" instead of "E")
- COMMIT/ROLLBACK/ROLLBACK TO SAVEPOINT through a portal: tag emitted once,
  sibling suspended portals cleared, later Execute 34000
- BEGIN through a portal: tag emitted once, block starts, sibling portals kept,
  repeated Execute repeats the tag without replaying BEGIN
- named re-Bind 42P03 with the transaction aborted; unnamed re-Bind restart

Capabilities and README document the savepoint clearing and the 42P03 rule.
Danilo Fragoso 21 時間 前
親
コミット
ea3639f3f9

ファイルの差分が大きいため隠しています
+ 0 - 0
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: 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"},
+			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",
 		},

+ 1 - 1
pkg/capabilities/capabilities_test.go

@@ -340,7 +340,7 @@ func TestBoundedExecuteIsAdvertised(t *testing.T) {
 	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",
+		"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)

+ 36 - 12
pkg/pgserver/connection.go

@@ -96,10 +96,10 @@ type portalExecution struct {
 	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
+	// 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
@@ -539,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()
@@ -593,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"))
@@ -719,8 +734,7 @@ func (c *Connection) handleExecute(msg *Message) error {
 	if err := c.sendPortalRows(portal, maxRows); err != nil {
 		return err
 	}
-	if portal.execution.endsTransaction {
-		// PostgreSQL destroys the portals of the transaction that just ended.
+	if portal.execution.clearsPortals {
 		c.clearPortals()
 	}
 	return nil
@@ -769,12 +783,7 @@ func (c *Connection) materializePortal(portal *portal) *executeError {
 		// 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 == ""
-		}
+		execution.clearsPortals = resetsPortals(stmt)
 	}
 	columnTypes := wireColumnTypes(execution.result)
 	formats, formatErr := c.resolveResultFormats(portal.resultFormats, columnTypes, len(execution.result.Columns))
@@ -1153,6 +1162,21 @@ 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 {

+ 163 - 9
pkg/pgserver/portal_test.go

@@ -241,22 +241,51 @@ func TestCloseEndsSuspendedPortal(t *testing.T) {
 	}
 }
 
-// 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) {
+// 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("p", "stmt", nil, nil))
-	assertFrameTypes(t, s.send(executeRowsMessage("p", 2)), "TDDs")
+	s.send(parseMessage("stmt", portalSelectSQL), bindResultMessage("", "stmt", nil, nil))
+	assertFrameTypes(t, s.send(executeRowsMessage("", 2)), "TDDs")
 
-	frames := s.send(
-		bindResultMessage("p", "stmt", nil, nil),
-		executeRowsMessage("p", 0),
-	)
+	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) {
@@ -450,3 +479,128 @@ func TestFailedExecuteLeavesTransactionConsistent(t *testing.T) {
 	}
 	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)
+	}
+}

+ 1 - 0
pkg/pgserver/protocol.go

@@ -90,6 +90,7 @@ const (
 	ErrCodeUndefinedColumn      = "42703"
 	ErrCodeDuplicateTable       = "42P07"
 	ErrCodeDuplicateColumn      = "42701"
+	ErrCodeDuplicateCursor      = "42P03"
 	ErrCodeInvalidParameter     = "22023"
 	ErrCodeInvalidText          = "22P02"
 	ErrCodeDatatypeMismatch     = "42804"

+ 3 - 1
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 {

+ 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},

この差分においてかなりの量のファイルが変更されているため、一部のファイルを表示していません