From d17f352662603e0be914f86bb533fa49895285a6 Mon Sep 17 00:00:00 2001 From: "posthog[bot]" <206114724+posthog[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:22:39 +0000 Subject: [PATCH 1/2] Implement PostgreSQL PortalSuspended for Execute's max_rows cap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A client Execute with a non-zero max_rows (any BI tool/driver that sets a fetch size over the extended protocol) hit the cap in streamSelectRows and was reported to the client as a complete CommandComplete SELECT instead of PortalSuspended — silent truncation with no error. The row already fetched from the DB when the cap tripped was also dropped, since the loop called rows.Next() before checking the cap. Fix: check the cap before advancing the RowSet (so no row is lost), send PortalSuspended when it trips, and keep the portal's RowSet open so the next Execute resumes the same query instead of restarting it (which would have duplicated rows). Branch: posthog-self-driving/fixserver-implement-portalsuspended-so-a3b3dc Generated-By: PostHog Code Task-Id: b830fd36-c3cf-4d21-81e1-235cfc4afd40 --- docs/postgres-compatibility.md | 1 + server/conn.go | 37 ++- server/conn_extended_query.go | 235 +++++++++++------- ...conn_extended_query_portal_suspend_test.go | 140 +++++++++++ server/conn_query_exec.go | 31 ++- server/conn_results.go | 11 + server/wire/protocol.go | 8 + 7 files changed, 367 insertions(+), 96 deletions(-) create mode 100644 server/conn_extended_query_portal_suspend_test.go diff --git a/docs/postgres-compatibility.md b/docs/postgres-compatibility.md index 591a6744d..1e16b6c62 100644 --- a/docs/postgres-compatibility.md +++ b/docs/postgres-compatibility.md @@ -164,6 +164,7 @@ Operational guidance: allow the import to complete, then manually deduplicate af | Simple query | ✅ | `protocol_test.go::TestProtocolSimpleQuery` | | | Extended query (Parse/Bind/Describe/Execute/Sync) | ✅ | `protocol_test.go::TestProtocolExtendedQuery`; server `conn_bind_test.go`, `conn_describe_test.go` | | | Extended-query error recovery (skip-until-Sync, pipelining) | ✅ | server `conn_skip_until_sync_test.go`; clients `clients_test.go::TestExtendedQueryErrorHandling` | #718 | +| Execute row limit (`max_rows`) / `PortalSuspended` | ✅ | server `conn_extended_query_portal_suspend_test.go::TestHandleExecutePortalSuspendedResumesWithoutLosingOrDuplicatingRows` | A client Execute with a non-zero `max_rows` (a BI tool / driver fetch size) sends `PortalSuspended` when the cap is hit before the result set is exhausted, and the portal's `RowSet` stays open so the next Execute resumes rather than restarting. Previously the server silently capped the stream and reported `CommandComplete`, truncating results with no error. | | Prepared statements (PREPARE/EXECUTE, reuse, NULL, 20+ params) | ✅ | `edge_cases_test.go::TestPreparedStatementEdgeCases`; clients `::TestPreparedStatements` | | | Binary vs text result format | ✅ | `protocol_test.go::TestProtocolDataTypes`; clients `::TestPgxBinaryFormatResults`; server `types_test.go` encode/decode | | | Row description metadata | ✅ | `protocol_test.go::TestProtocolRowDescription` | | diff --git a/server/conn.go b/server/conn.go index a985df7ab..6d8673ace 100644 --- a/server/conn.go +++ b/server/conn.go @@ -110,6 +110,37 @@ type portal struct { paramFormats []int16 // 0=text, 1=binary for each parameter resultFormats []int16 described bool // true if Describe was called on this portal + + // openRows is the live result set of a SELECT portal that was suspended + // by a previous Execute's max_rows cap (PortalSuspended sent instead of + // CommandComplete). Non-nil means "resume this on the next Execute" + // rather than re-running the query, which would restart from row 0 and + // duplicate everything already sent. Cleared (and the RowSet closed) once + // the result set drains, errors, or the portal/statement is Closed. + openRows RowSet + openCols []string + openColTypes []ColumnTyper + openTypeOIDs []int32 + // openRowsSent is the cumulative row count sent across all Executes of + // this portal, for the final CommandComplete tag. + openRowsSent int64 +} + +// closeOpenRows releases a suspended portal's live result set, if any. Safe +// to call on a portal that was never suspended. +func (p *portal) closeOpenRows() { + if p.openRows != nil { + _ = p.openRows.Close() + p.openRows = nil + } +} + +// closeAllOpenPortalRows releases every suspended portal's live result set on +// connection teardown, mirroring closeAllCursors. +func (c *clientConn) closeAllOpenPortalRows() { + for _, p := range c.portals { + p.closeOpenRows() + } } // decodeParams converts raw parameter bytes to Go values based on format codes. @@ -960,14 +991,16 @@ func (c *clientConn) serve() error { stopRefresh = StartCredentialRefresh(db, c.server.cfg.DuckLake) } } - // Defers run LIFO: close cursors first (they hold open RowSets), then stop - // credential refresh, then clean up the database connection. + // Defers run LIFO: close cursors and suspended portals first (they hold + // open RowSets), then stop credential refresh, then clean up the database + // connection. defer func() { if c.executor != nil { c.safeCleanupDB() } }() defer c.closeAllCursors() + defer c.closeAllOpenPortalRows() defer func() { if stopRefresh != nil { stopRefresh() diff --git a/server/conn_extended_query.go b/server/conn_extended_query.go index 47a486bed..1231e585d 100644 --- a/server/conn_extended_query.go +++ b/server/conn_extended_query.go @@ -865,95 +865,140 @@ func (c *clientConn) handleExecute(body []byte) { return } - // Result-returning query: use Query with converted query - runQuery := func() (RowSet, error) { - return c.executor.Query(convertedQuery, args...) - } - - execStart := time.Now() - execCtx, execSpan := observe.Tracer().Start(queryCtx, "duckgres.execute") - rows, err := runQuery() - if err != nil && c.txStatus == txStatusIdle && isDuckLakeTransactionConflict(err) { - ducklakeConflictTotal.Inc() - rows, err = retryOnConflict(runQuery) - } - if err != nil { - rows, err, _ = recoverAbortedTransaction( - err, - c.txStatus == txStatusIdle, - func() error { - _, rollbackErr := c.executor.ExecContext(context.Background(), "ROLLBACK") - return rollbackErr - }, - runQuery, - ) - } - // Exploratory tier: a read that blew the small worker's memory_limit is - // transparently re-executed on a normal-size worker. Prepare phase only — - // nothing has been sent to the client yet, so the retry is invisible. Never - // inside a transaction: the new worker has none of its accumulated state. - // runQuery reads c.executor at call time, so it targets the new worker. - // Same contract as executeSelectQuery; a failed escalation surfaces the - // ORIGINAL query error as FATAL and terminates the connection. - if err != nil && c.onExploratoryWorker && isWorkerOutOfMemoryError(err) && c.txStatus == txStatusIdle { - if escErr := c.escalateWorker(queryCtx, escalateReasonOOM); escErr != nil { - queryFinalErr = err - execSpan.End() - _ = c.failEscalation(convertedQuery, escErr, err.Error()) + // Result-returning query: resume a suspended portal's RowSet (from a + // previous Execute that hit max_rows), or run a fresh Query. Only a fresh + // (non-resumed) execution goes through the conflict-retry / OOM-retry / + // RowDescription machinery below — a resumed portal already has all of + // that settled from its first Execute. + resuming := p.openRows != nil + + var rows RowSet + var cols []string + var colTypes []ColumnTyper + var typeOIDs []int32 + if resuming { + rows = p.openRows + cols = p.openCols + colTypes = p.openColTypes + typeOIDs = p.openTypeOIDs + } + + // suspended is set below once streamSelectRows reports the max_rows cap + // was hit with more of the result set left. The deferred cleanup decides, + // on every exit path, whether to hand the still-open RowSet back to the + // portal for the next Execute (suspended) or close it (drained/errored) — + // this is what makes a suspended portal resumable instead of the old + // "close rows and hope the client never comes back for more" behavior. + var suspended bool + defer func() { + if suspended { + p.openRows = rows + p.openCols = cols + p.openColTypes = colTypes + p.openTypeOIDs = typeOIDs return } - rows, err = runQuery() - } - c.lastProfilingSummary = observe.EnrichSpanWithProfiling(execCtx, execSpan, execStart, c.executor, c.orgID) - execSpan.End() - if err != nil { - queryFinalErr = err - errCode := classifyErrorCode(err) - errMsg := err.Error() - if c.isCallerCancellation(err) { - errMsg = "canceling statement due to user request" - } else { - c.logQueryError(convertedQuery, err) + if rows != nil { + _ = rows.Close() } - c.sendError("ERROR", errCode, errMsg) - c.setTxError() - c.logQuery(start, originalQuery, convertedQuery, cmdType, 0, 0, errCode, errMsg, "extended") - return - } - defer func() { _ = rows.Close() }() + p.openRows = nil + }() - cols, err := rows.Columns() - if err != nil { - queryFinalErr = err - c.logger().Error("Columns error.", "error", err) - c.sendError("ERROR", "42000", err.Error()) - c.setTxError() - c.logQuery(start, originalQuery, convertedQuery, cmdType, 0, 0, "42000", err.Error(), "extended") - return - } + if !resuming { + runQuery := func() (RowSet, error) { + return c.executor.Query(convertedQuery, args...) + } - // Get column types for binary encoding - colTypes, _ := rows.ColumnTypes() - typeOIDs := make([]int32, len(cols)) - for i, ct := range colTypes { - typeOIDs[i] = getTypeInfo(ct).OID - } + execStart := time.Now() + execCtx, execSpan := observe.Tracer().Start(queryCtx, "duckgres.execute") + var err error + rows, err = runQuery() + if err != nil && c.txStatus == txStatusIdle && isDuckLakeTransactionConflict(err) { + ducklakeConflictTotal.Inc() + rows, err = retryOnConflict(runQuery) + } + if err != nil { + rows, err, _ = recoverAbortedTransaction( + err, + c.txStatus == txStatusIdle, + func() error { + _, rollbackErr := c.executor.ExecContext(context.Background(), "ROLLBACK") + return rollbackErr + }, + runQuery, + ) + } + // Exploratory tier: a read that blew the small worker's memory_limit is + // transparently re-executed on a normal-size worker. Prepare phase only — + // nothing has been sent to the client yet, so the retry is invisible. Never + // inside a transaction: the new worker has none of its accumulated state. + // runQuery reads c.executor at call time, so it targets the new worker. + // Same contract as executeSelectQuery; a failed escalation surfaces the + // ORIGINAL query error as FATAL and terminates the connection. + if err != nil && c.onExploratoryWorker && isWorkerOutOfMemoryError(err) && c.txStatus == txStatusIdle { + if escErr := c.escalateWorker(queryCtx, escalateReasonOOM); escErr != nil { + queryFinalErr = err + execSpan.End() + _ = c.failEscalation(convertedQuery, escErr, err.Error()) + return + } + rows, err = runQuery() + } + c.lastProfilingSummary = observe.EnrichSpanWithProfiling(execCtx, execSpan, execStart, c.executor, c.orgID) + execSpan.End() + if err != nil { + queryFinalErr = err + errCode := classifyErrorCode(err) + errMsg := err.Error() + if c.isCallerCancellation(err) { + errMsg = "canceling statement due to user request" + } else { + c.logQueryError(convertedQuery, err) + } + c.sendError("ERROR", errCode, errMsg) + c.setTxError() + c.logQuery(start, originalQuery, convertedQuery, cmdType, 0, 0, errCode, errMsg, "extended") + return + } - // Send RowDescription if Describe wasn't called before Execute. - // Some clients skip Describe and go straight to Execute, but still - // need the column metadata before receiving data rows. - // Skip if there are no columns - queries that return 0 columns (like - // DDL accidentally routed here) don't need RowDescription. - if !p.described && len(cols) > 0 { - if err := c.sendRowDescriptionWithFormats(cols, colTypes, p.resultFormats); err != nil { + cols, err = rows.Columns() + if err != nil { + queryFinalErr = err + c.logger().Error("Columns error.", "error", err) + c.sendError("ERROR", "42000", err.Error()) + c.setTxError() + c.logQuery(start, originalQuery, convertedQuery, cmdType, 0, 0, "42000", err.Error(), "extended") return } + + // Get column types for binary encoding + colTypes, _ = rows.ColumnTypes() + typeOIDs = make([]int32, len(cols)) + for i, ct := range colTypes { + typeOIDs[i] = getTypeInfo(ct).OID + } + + // Send RowDescription if Describe wasn't called before Execute. + // Some clients skip Describe and go straight to Execute, but still + // need the column metadata before receiving data rows. + // Skip if there are no columns - queries that return 0 columns (like + // DDL accidentally routed here) don't need RowDescription. + if !p.described && len(cols) > 0 { + if err := c.sendRowDescriptionWithFormats(cols, colTypes, p.resultFormats); err != nil { + return + } + } + // A later Execute on this portal (resuming, or after a plain restart + // without Describe) must never resend RowDescription. + p.described = true } // Send rows with the format codes from Bind. Shared with the simple-query // path so the exploratory tier's zero-row retry below behaves identically - // on both protocols; maxRows still caps the DataRows (portal suspension is - // not implemented) and the RowDescription was already handled above. + // on both protocols. maxRows caps the DataRows sent per Execute; if the + // cap is hit before the result set is exhausted, streamSelectRows leaves + // rows open and positioned to continue, and stream.suspended tells us to + // report PortalSuspended below instead of CommandComplete. stream := c.streamSelectRows(rows, cols, colTypes, typeOIDs, false, p.resultFormats, maxRows) // Exploratory tier: an OOM raised before a SINGLE DataRow reached the @@ -961,20 +1006,24 @@ func (c *clientConn) handleExecute(body []byte) { // has seen is the RowDescription, which the identical query on the same // engine reproduces exactly, so it is deliberately NOT resent. Once rows // are out the door the error must surface: a retry cannot un-send them. - if stream.rowsErr != nil && stream.rowsSent == 0 && + // Only applies to a fresh execution: a resumed portal has already sent + // rows from an earlier Execute, so the client has committed to this + // result set and a transparent restart would duplicate them. + if !resuming && stream.rowsErr != nil && stream.rowsSent == 0 && c.onExploratoryWorker && isWorkerOutOfMemoryError(stream.rowsErr) && c.txStatus == txStatusIdle { oomErr := stream.rowsErr _ = rows.Close() + rows = nil if escErr := c.escalateWorker(queryCtx, escalateReasonOOM); escErr != nil { queryFinalErr = oomErr _ = c.failEscalation(convertedQuery, escErr, oomErr.Error()) return } - retryRows, retryErr := runQuery() + retryRows, retryErr := c.executor.Query(convertedQuery, args...) if retryErr != nil { stream.rowsErr = retryErr } else { - defer func() { _ = retryRows.Close() }() + rows = retryRows stream = c.streamSelectRows(retryRows, cols, colTypes, typeOIDs, false, p.resultFormats, maxRows) } } @@ -991,8 +1040,17 @@ func (c *clientConn) handleExecute(body []byte) { return } - rowCount := stream.rowsSent - queryRowsAff = int64(rowCount) + if stream.suspended { + suspended = true + p.openRowsSent += int64(stream.rowsSent) + queryRowsAff = p.openRowsSent + _ = c.writePortalSuspended() + c.logQuery(start, originalQuery, convertedQuery, cmdType, p.openRowsSent, 0, "", "", "extended") + return + } + + rowCount := p.openRowsSent + int64(stream.rowsSent) + queryRowsAff = rowCount if err := stream.rowsErr; err != nil { queryFinalErr = err @@ -1011,10 +1069,11 @@ func (c *clientConn) handleExecute(body []byte) { return } + p.openRowsSent = 0 c.updateTxStatus(cmdType) - tag := buildCommandTagFromRowCount(cmdType, int64(rowCount)) + tag := buildCommandTagFromRowCount(cmdType, rowCount) _ = c.writeCommandComplete(tag) - c.logQuery(start, originalQuery, convertedQuery, cmdType, int64(rowCount), 0, "", "", "extended") + c.logQuery(start, originalQuery, convertedQuery, cmdType, rowCount, 0, "", "", "extended") } func (c *clientConn) handleClose(body []byte) { @@ -1034,6 +1093,9 @@ func (c *clientConn) handleClose(body []byte) { case 'S': delete(c.stmts, name) case 'P': + if p, ok := c.portals[name]; ok { + p.closeOpenRows() + } delete(c.portals, name) } @@ -1220,7 +1282,10 @@ func (c *clientConn) handleBind(body []byte) { } } - // Close existing portal with same name + // Close existing portal with same name, releasing any suspended result set. + if old, ok := c.portals[portalName]; ok { + old.closeOpenRows() + } delete(c.portals, portalName) c.portals[portalName] = &portal{ diff --git a/server/conn_extended_query_portal_suspend_test.go b/server/conn_extended_query_portal_suspend_test.go new file mode 100644 index 000000000..599db499e --- /dev/null +++ b/server/conn_extended_query_portal_suspend_test.go @@ -0,0 +1,140 @@ +package server + +import ( + "bufio" + "bytes" + "context" + "encoding/binary" + "net" + "testing" + + "github.com/posthog/duckgres/server/wire" +) + +// buildExecuteBody encodes an Execute message body: portal name followed by +// the int32 max_rows field. +func buildExecuteBody(portalName string, maxRows int32) []byte { + var body bytes.Buffer + body.WriteString(portalName) + body.WriteByte(0) + _ = binary.Write(&body, binary.BigEndian, maxRows) + return body.Bytes() +} + +// countMessages scans a pgwire backend message stream and returns how many +// times each message type byte appears, plus the payload of the last +// CommandComplete ('C') message seen (its tag, sans trailing NUL). +func countMessages(t *testing.T, buf []byte) (counts map[byte]int, lastCommandTag string) { + t.Helper() + counts = map[byte]int{} + r := bytes.NewReader(buf) + for { + msgType, data, err := wire.ReadMessage(r) + if err != nil { + break + } + counts[msgType]++ + if msgType == wire.MsgCommandComplete { + lastCommandTag = string(bytes.TrimRight(data, "\x00")) + } + } + return counts, lastCommandTag +} + +// TestHandleExecutePortalSuspendedResumesWithoutLosingOrDuplicatingRows is the +// regression test for the silent-truncation bug: an Execute whose max_rows +// cap is hit below the result set's true size must send PortalSuspended (not +// CommandComplete) and keep the portal's RowSet open so a subsequent Execute +// resumes and delivers every remaining row exactly once — no row dropped at +// the boundary, none duplicated on resume. +func TestHandleExecutePortalSuspendedResumesWithoutLosingOrDuplicatingRows(t *testing.T) { + clientSide, serverSide := net.Pipe() + defer func() { _ = clientSide.Close() }() + defer func() { _ = serverSide.Close() }() + + rows := &streamingRowSet{ + rows: [][]any{{int64(1)}, {int64(2)}, {int64(3)}}, + cols: []string{"n"}, + colTypers: []ColumnTyper{stringColumnTyper{}}, + } + executor := &lifecycleExecutor{queryRows: rows} + + var out bytes.Buffer + c := &clientConn{ + server: &Server{activeQueries: make(map[BackendKey]context.CancelFunc)}, + conn: clientSide, + reader: bufio.NewReader(clientSide), + writer: bufio.NewWriter(&out), + ctx: context.Background(), + cancel: func() {}, + txStatus: txStatusIdle, + executor: executor, + portals: map[string]*portal{ + "p": { + stmt: &preparedStmt{ + query: "SELECT n FROM t", + convertedQuery: "SELECT n FROM t", + }, + }, + }, + } + + // First Execute caps at 1 row out of 3. + c.handleExecute(buildExecuteBody("p", 1)) + _ = c.writer.Flush() + + if executor.queryCalls.Load() != 1 { + t.Fatalf("expected exactly one Query call (no restart), got %d", executor.queryCalls.Load()) + } + if rows.idx != 1 { + t.Fatalf("expected the RowSet to be positioned after exactly 1 row (no dropped lookahead row), got idx=%d", rows.idx) + } + if rows.closed { + t.Fatalf("expected the RowSet to stay open across a suspended Execute") + } + p := c.portals["p"] + if p.openRows == nil { + t.Fatalf("expected the portal to retain the open RowSet for resumption") + } + if p.openRowsSent != 1 { + t.Fatalf("expected openRowsSent=1 after the first Execute, got %d", p.openRowsSent) + } + + counts1, tag1 := countMessages(t, out.Bytes()) + if counts1[wire.MsgCommandComplete] != 0 { + t.Fatalf("expected NO CommandComplete on a suspended Execute, got tag %q", tag1) + } + if counts1[wire.MsgPortalSuspended] != 1 { + t.Fatalf("expected exactly one PortalSuspended message, got %d", counts1[wire.MsgPortalSuspended]) + } + if counts1[wire.MsgDataRow] != 1 { + t.Fatalf("expected exactly 1 DataRow in the first Execute, got %d", counts1[wire.MsgDataRow]) + } + + // Second Execute (no cap) resumes and must deliver the remaining 2 rows, + // then a CommandComplete with the row count across BOTH Executes. + out.Reset() + c.handleExecute(buildExecuteBody("p", 0)) + _ = c.writer.Flush() + + if executor.queryCalls.Load() != 1 { + t.Fatalf("expected the resumed Execute to reuse the existing RowSet, not re-run Query; got %d Query calls", executor.queryCalls.Load()) + } + if !rows.closed { + t.Fatalf("expected the RowSet to be closed once the portal drains") + } + if p.openRows != nil { + t.Fatalf("expected the portal's openRows to be cleared once drained") + } + + counts2, tag2 := countMessages(t, out.Bytes()) + if counts2[wire.MsgDataRow] != 2 { + t.Fatalf("expected exactly 2 DataRows on resume (rows 2 and 3, none dropped/duplicated), got %d", counts2[wire.MsgDataRow]) + } + if counts2[wire.MsgCommandComplete] != 1 { + t.Fatalf("expected a CommandComplete on the draining Execute, got %d", counts2[wire.MsgCommandComplete]) + } + if tag2 != "SELECT 3" { + t.Fatalf("expected cumulative command tag %q, got %q", "SELECT 3", tag2) + } +} diff --git a/server/conn_query_exec.go b/server/conn_query_exec.go index 73d713f04..b9e0046ac 100644 --- a/server/conn_query_exec.go +++ b/server/conn_query_exec.go @@ -357,18 +357,25 @@ type selectStream struct { // can reproduce the exact log wording. writeErr error writeStage string + // suspended is true when maxRows capped the stream before rows was + // exhausted. rows is left open and positioned so a later Next() continues + // exactly where this call stopped — the caller must NOT close it and + // must send PortalSuspended instead of CommandComplete. + suspended bool } -// streamSelectRows sends (optionally) the RowDescription and then every DataRow -// of rows. Extracted from executeSelectQuery so the exploratory tier can retry -// a zero-row OOM stream on the escalated worker WITHOUT resending -// RowDescription — pass sendRowDesc=false on such a retry. +// streamSelectRows sends (optionally) the RowDescription and then DataRows of +// rows, up to maxRows. Extracted from executeSelectQuery so the exploratory +// tier can retry a zero-row OOM stream on the escalated worker WITHOUT +// resending RowDescription — pass sendRowDesc=false on such a retry. // // formats are the Bind result-format codes (nil = all text, the simple-query // case). maxRows > 0 caps the DataRows sent, for the extended protocol's -// Execute row limit; portal suspension is not implemented, so the row already -// fetched when the cap is reached is dropped, and rows.Err() is still consulted -// by the caller exactly as on a fully drained stream. +// Execute row limit. The cap is checked BEFORE calling rows.Next() for the +// next row, so a capped stream never advances past (and drops) a row it +// didn't send — rows stays valid and resumable via a later Execute on the +// same portal. rows.Err() is only meaningful once the stream actually +// finished (not suspended), so it is left unset on a suspended stream. func (c *clientConn) streamSelectRows(rows RowSet, cols []string, colTypes []ColumnTyper, typeOIDs []int32, sendRowDesc bool, formats []int16, maxRows int32) selectStream { if sendRowDesc { if err := c.sendRowDescription(cols, colTypes); err != nil { @@ -377,8 +384,12 @@ func (c *clientConn) streamSelectRows(rows RowSet, cols []string, colTypes []Col } var out selectStream - for rows.Next() { + for { if maxRows > 0 && int32(out.rowsSent) >= maxRows { + out.suspended = true + break + } + if !rows.Next() { break } @@ -400,7 +411,9 @@ func (c *clientConn) streamSelectRows(rows RowSet, cols []string, colTypes []Col } out.rowsSent++ } - out.rowsErr = rows.Err() + if !out.suspended { + out.rowsErr = rows.Err() + } return out } diff --git a/server/conn_results.go b/server/conn_results.go index 0e8e74031..d658fb9d3 100644 --- a/server/conn_results.go +++ b/server/conn_results.go @@ -588,6 +588,17 @@ func (c *clientConn) writeCommandComplete(tag string) error { return nil } +func (c *clientConn) writePortalSuspended() error { + if err := wire.WritePortalSuspended(c.writer); err != nil { + c.markActiveQueryMetricsError(err) + return err + } + if c.activeQueryMetrics != nil { + return c.flushWriter() + } + return nil +} + func (c *clientConn) writeReadyForQuery(txStatus byte) error { if err := wire.WriteReadyForQuery(c.writer, txStatus); err != nil { c.markActiveQueryMetricsError(err) diff --git a/server/wire/protocol.go b/server/wire/protocol.go index 1c25b4c87..e203241c4 100644 --- a/server/wire/protocol.go +++ b/server/wire/protocol.go @@ -36,6 +36,7 @@ const ( MsgBindComplete = '2' MsgCloseComplete = '3' MsgNoData = 'n' + MsgPortalSuspended = 's' // COPY messages (both directions) MsgCopyData = 'd' // Contains COPY data @@ -340,6 +341,13 @@ func WriteNoData(w io.Writer) error { return WriteMessage(w, MsgNoData, nil) } +// WritePortalSuspended tells the client that Execute's row limit was hit +// before the result set was exhausted. The portal stays open — a later +// Execute on the same portal resumes rather than restarting the query. +func WritePortalSuspended(w io.Writer) error { + return WriteMessage(w, MsgPortalSuspended, nil) +} + // writeCopyOutResponse tells client we're about to send COPY data // Format: overall format (0=text, 1=binary), num columns, format per column func WriteCopyOutResponse(w io.Writer, numColumns int16, textFormat bool) error { From bbefa386b064376d31a0532a675990cea9ed84d7 Mon Sep 17 00:00:00 2001 From: "posthog[bot]" <206114724+posthog[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:47:30 +0000 Subject: [PATCH 2/2] Fix conn_tier_exec_test to expect PortalSuspended instead of the old truncation behavior CI caught that TestExtendedExecuteMaxRowsUnchanged asserted the pre-fix bug (CommandComplete on a capped Execute) as expected behavior. Replaced with TestExtendedExecuteMaxRowsSuspendsAndResumes, which asserts PortalSuspended is sent, the RowSet stays open with no row dropped at the cap boundary, and a follow-up Execute on the same portal resumes and delivers the remaining rows exactly once. Generated-By: PostHog Code Task-Id: b830fd36-c3cf-4d21-81e1-235cfc4afd40 --- server/conn_tier_exec_test.go | 57 ++++++++++++++++++++++++++++++----- 1 file changed, 49 insertions(+), 8 deletions(-) diff --git a/server/conn_tier_exec_test.go b/server/conn_tier_exec_test.go index d32d8ac04..5a2e98265 100644 --- a/server/conn_tier_exec_test.go +++ b/server/conn_tier_exec_test.go @@ -1340,13 +1340,16 @@ func TestExtendedExecuteMidStreamOOMAfterRowsSurfaces(t *testing.T) { } } -// TestExtendedExecuteMaxRowsUnchanged pins the non-tier behavior of the -// Execute row loop through the shared streaming helper: maxRows still caps the -// DataRows sent (portal suspension is not implemented — the surplus row is -// fetched and dropped, as before). -func TestExtendedExecuteMaxRowsUnchanged(t *testing.T) { +// TestExtendedExecuteMaxRowsSuspendsAndResumes pins the non-tier behavior of +// the Execute row loop through the shared streaming helper: maxRows caps the +// DataRows sent, but the capped Execute now reports PortalSuspended (never +// CommandComplete) and leaves the portal's RowSet open so a later Execute on +// the same portal resumes and delivers the rest — no row dropped at the +// boundary, none duplicated on resume. +func TestExtendedExecuteMaxRowsSuspendsAndResumes(t *testing.T) { + rows := &tierRowSet{rows: []int64{1, 2, 3}} exec := &tierExecutor{name: "std", queryFn: func(int, string) (RowSet, error) { - return &tierRowSet{rows: []int64{1, 2, 3}}, nil + return rows, nil }} c, out := newBufferedConn(exec) c.stmts = make(map[string]*preparedStmt) @@ -1365,8 +1368,46 @@ func TestExtendedExecuteMaxRowsUnchanged(t *testing.T) { if got := countMsgs(msgs, 'D'); got != 2 { t.Fatalf("DataRow count = %d, want 2 (maxRows): %s", got, describeMsgs(msgs)) } - if !commandCompleteWith(msgs, "SELECT 2") { - t.Fatalf("want CommandComplete 'SELECT 2': %s", describeMsgs(msgs)) + if got := countMsgs(msgs, 's'); got != 1 { + t.Fatalf("PortalSuspended count = %d, want 1: %s", got, describeMsgs(msgs)) + } + if commandCompleteWith(msgs, "SELECT 2") || countMsgs(msgs, 'C') != 0 { + t.Fatalf("expected NO CommandComplete on a suspended Execute: %s", describeMsgs(msgs)) + } + if len(exec.queryCalls) != 1 { + t.Fatalf("expected exactly one Query call, got %d", len(exec.queryCalls)) + } + if rows.idx != 2 { + t.Fatalf("expected the RowSet positioned after exactly 2 rows (no dropped lookahead row), got idx=%d", rows.idx) + } + if rows.closed != 0 { + t.Fatalf("expected the RowSet to stay open across a suspended Execute, closed=%d", rows.closed) + } + if c.portals["s1"].openRows == nil { + t.Fatalf("expected the portal to retain the open RowSet for resumption") + } + + // A second Execute on the same portal must resume, not restart: it should + // deliver only the remaining row and report the cumulative row count. + out.Reset() + if err := extExecute(c, "s1", 0); err != nil { + t.Fatalf("Execute (resume): %v", err) + } + msgs = parseWireMsgs(t, out.Bytes()) + if got := countMsgs(msgs, 'D'); got != 1 { + t.Fatalf("resumed DataRow count = %d, want 1 (row 3, not re-sending 1-2): %s", got, describeMsgs(msgs)) + } + if !commandCompleteWith(msgs, "SELECT 3") { + t.Fatalf("want cumulative CommandComplete 'SELECT 3': %s", describeMsgs(msgs)) + } + if len(exec.queryCalls) != 1 { + t.Fatalf("expected the resumed Execute to reuse the existing RowSet, not re-run Query; got %d Query calls", len(exec.queryCalls)) + } + if rows.closed != 1 { + t.Fatalf("expected the RowSet to be closed exactly once after draining, closed=%d", rows.closed) + } + if c.portals["s1"].openRows != nil { + t.Fatalf("expected the portal's openRows to be cleared once drained") } }