diff --git a/docs/postgres-compatibility.md b/docs/postgres-compatibility.md index 591a6744..1e16b6c6 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 a985df7a..6d8673ac 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 47a486be..1231e585 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 00000000..599db499 --- /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 73d713f0..b9e0046a 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 0e8e7403..d658fb9d 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/conn_tier_exec_test.go b/server/conn_tier_exec_test.go index d32d8ac0..5a2e9826 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") } } diff --git a/server/wire/protocol.go b/server/wire/protocol.go index 1c25b4c8..e203241c 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 {