Capture profiles from Flight DoGet

duckdbservice/flight_handler.go

@@ -721,10 +721,6 @@ func (h *FlightSQLHandler) GetFlightInfoStatement(ctx context.Context, cmd fligh
 		return nil, status.Errorf(codes.InvalidArgument, "failed to prepare query: %v", err)
 	}
 
-	// Send DuckDB profiling output as gRPC trailing metadata so the
-	// control plane can attach it to the trace span.
-	sendProfilingMetadata(ctx, session)
-
 	handleID := fmt.Sprintf("query-%d", session.handleCounter.Add(1))
 	ticketBytes, err := flightsql.CreateStatementQueryTicket([]byte(handleID))
 	if err != nil {
@@ -822,6 +818,7 @@ func (h *FlightSQLHandler) DoGetStatement(ctx context.Context, ticket flightsql.
 
 		inTxn := tx != nil || session.sqlTxActive.Load()
 		var closeRows func() error
+		execStartedAt := clearProfilingOutput()
 		queryFn := func() (*sql.Rows, error) {
 			rows, closer, err := session.queryRows(ctx, tx, handle.Query)
 			if err != nil {
@@ -862,6 +859,7 @@ func (h *FlightSQLHandler) DoGetStatement(ctx context.Context, ticket flightsql.
 		}
 		defer func() {
 			_ = closeRows()
+			sendProfilingMetadataSince(ctx, execStartedAt)
 		}()
 
 		for {
@@ -937,6 +935,7 @@ func (h *FlightSQLHandler) DoPutCommandStatementUpdate(ctx context.Context,
 	defer session.progress.queryActive.Store(false)
 	endTxnWork := ttx.beginWork()
 	defer endTxnWork()
+	execStartedAt := clearProfilingOutput()
 
 	execFn := func() (sql.Result, error) {
 		return session.exec(ctx, tx, query)
@@ -1003,7 +1002,7 @@ func (h *FlightSQLHandler) DoPutCommandStatementUpdate(ctx context.Context,
 		return 0, status.Errorf(codes.InvalidArgument, "failed to execute update: %v", execErr)
 	}
 
-	sendProfilingMetadata(ctx, session)
+	sendProfilingMetadataSince(ctx, execStartedAt)
 
 	affected, err := result.RowsAffected()
 	if err != nil {

duckdbservice/profiling.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"os"
+	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
@@ -17,20 +18,38 @@ const profilingMetadataKey = "x-duckgres-profiling"
 // profilingOutputPath is the fixed file path where DuckDB writes profiling
 // output. Only one query runs per worker at a time (control plane enforces
 // this), so a single file is safe.
-const profilingOutputPath = "/tmp/duckgres-profiling.json"
+var profilingOutputPath = "/tmp/duckgres-profiling.json"
 
-// sendProfilingMetadata reads the profiling output file written by DuckDB
-// and sends it as gRPC trailing metadata so the control plane can attach
-// it to the trace span.
-func sendProfilingMetadata(ctx context.Context, session *Session) {
+// clearProfilingOutput removes the previous statement's profile before a
+// statement begins. The returned timestamp is used to reject a profile that
+// was not produced by this execution.
+func clearProfilingOutput() time.Time {
+	_ = os.Remove(profilingOutputPath)
+	return time.Now()
+}
+
+func profilingMetadataSince(startedAt time.Time) string {
+	info, err := os.Stat(profilingOutputPath)
+	if err != nil || !info.ModTime().After(startedAt) {
+		return ""
+	}
 	data, err := os.ReadFile(profilingOutputPath)
 	if err != nil || len(data) == 0 {
-		return
+		return ""
 	}
 	// Compact JSON to a single line — gRPC metadata values cannot contain newlines.
 	var compact bytes.Buffer
 	if json.Compact(&compact, data) != nil {
-		return
+		return ""
+	}
+	return compact.String()
+}
+
+// sendProfilingMetadataSince sends the profile written by the execution that
+// began at startedAt. It deliberately ignores absent, stale, and malformed
+// files so canceled or failed statements cannot reuse an earlier profile.
+func sendProfilingMetadataSince(ctx context.Context, startedAt time.Time) {
+	if profile := profilingMetadataSince(startedAt); profile != "" {
+		_ = grpc.SetTrailer(ctx, metadata.Pairs(profilingMetadataKey, profile))
 	}
-	_ = grpc.SetTrailer(ctx, metadata.Pairs(profilingMetadataKey, compact.String()))
 }

duckdbservice/profiling_test.go

@@ -4,8 +4,10 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"os"
 	"strings"
 	"testing"
+	"time"
 )
 
 // validateGRPCMetadataValue checks that a string is safe for use as a gRPC
@@ -113,3 +115,28 @@ func TestProfilingOutputGRPCSafety(t *testing.T) {
 		t.Errorf("query_name mangled: %v", parsed["query_name"])
 	}
 }
+
+func TestSendProfilingMetadataSinceRejectsStaleOutput(t *testing.T) {
+	previousPath := profilingOutputPath
+	profilingOutputPath = t.TempDir() + "/profiling.json"
+	t.Cleanup(func() { profilingOutputPath = previousPath })
+
+	if err := os.WriteFile(profilingOutputPath, []byte(`{"latency":0.1}`), 0o600); err != nil {
+		t.Fatalf("write profile: %v", err)
+	}
+	if got := profilingMetadataSince(time.Now()); got != "" {
+		t.Fatalf("stale profile metadata = %q, want empty", got)
+	}
+
+	startedAt := clearProfilingOutput()
+	if err := os.WriteFile(profilingOutputPath, []byte(`{"latency":0.2}`), 0o600); err != nil {
+		t.Fatalf("write completed profile: %v", err)
+	}
+	freshAt := startedAt.Add(time.Second)
+	if err := os.Chtimes(profilingOutputPath, freshAt, freshAt); err != nil {
+		t.Fatalf("set completed profile mtime: %v", err)
+	}
+	if got := profilingMetadataSince(startedAt); got != `{"latency":0.2}` {
+		t.Fatalf("completed profile metadata = %q", got)
+	}
+}

server/conn.go

@@ -129,6 +129,9 @@ type portalExec struct {
 	originalQuery  string
 	convertedQuery string
 	start          time.Time // first Execute leg's start, so the query log spans all legs
+	// finishProfiling runs after rows.Close has allowed a Flight DoGet trailer
+	// to arrive. A suspended portal retains it until its terminal Execute.
+	finishProfiling func()
 }
 
 // closeExec releases a suspended portal's open rowset (if any). Must be
@@ -139,6 +142,9 @@ func (p *portal) closeExec() {
 		return
 	}
 	_ = p.exec.rows.Close()
+	if p.exec.finishProfiling != nil {
+		p.exec.finishProfiling()
+	}
 	p.exec = nil
 }
 

server/conn_extended_query.go

@@ -912,9 +912,9 @@ func (c *clientConn) handleExecute(body []byte) {
 		}
 		rows, err = runQuery()
 	}
-	c.lastProfilingSummary = observe.EnrichSpanWithProfiling(execCtx, execSpan, execStart, c.executor, c.orgID)
-	execSpan.End()
 	if err != nil {
+		c.lastProfilingSummary = observe.EnrichSpanWithProfiling(execCtx, execSpan, execStart, c.executor, c.orgID)
+		execSpan.End()
 		queryFinalErr = err
 		errCode := classifyErrorCode(err)
 		errMsg := err.Error()
@@ -929,10 +929,26 @@ func (c *clientConn) handleExecute(body []byte) {
 		return
 	}
 	keepRowsOpen := false
-	defer func() {
-		if !keepRowsOpen {
-			_ = rows.Close()
+	rowsFinished := false
+	profilingFinished := false
+	finishProfiling := func() {
+		if profilingFinished {
+			return
 		}
+		profilingFinished = true
+		c.lastProfilingSummary = observe.EnrichSpanWithProfiling(execCtx, execSpan, execStart, c.executor, c.orgID)
+		execSpan.End()
+	}
+	finishRows := func() {
+		if keepRowsOpen || rowsFinished {
+			return
+		}
+		rowsFinished = true
+		_ = rows.Close()
+		finishProfiling()
+	}
+	defer func() {
+		finishRows()
 	}()
 
 	cols, err := rows.Columns()
@@ -941,6 +957,7 @@ func (c *clientConn) handleExecute(body []byte) {
 		c.logger().Error("Columns error.", "error", err)
 		c.sendError("ERROR", "42000", err.Error())
 		c.setTxError()
+		finishRows()
 		c.logQuery(start, originalQuery, convertedQuery, cmdType, 0, 0, "42000", err.Error(), "extended")
 		return
 	}
@@ -991,12 +1008,8 @@ func (c *clientConn) handleExecute(body []byte) {
 		} else {
 			// The retried rowset replaces the original everywhere below —
 			// including as the rowset a suspension keeps open.
+			rows = retryRows
 			activeRows = retryRows
-			defer func() {
-				if !keepRowsOpen {
-					_ = retryRows.Close()
-				}
-			}()
 			stream = c.streamSelectRows(retryRows, cols, colTypes, typeOIDs, false, p.resultFormats, maxRows)
 		}
 	}
@@ -1005,6 +1018,7 @@ func (c *clientConn) handleExecute(body []byte) {
 		queryFinalErr = stream.scanErr
 		c.sendError("ERROR", "42000", stream.scanErr.Error())
 		c.setTxError()
+		finishRows()
 		c.logQuery(start, originalQuery, convertedQuery, cmdType, 0, 0, "42000", stream.scanErr.Error(), "extended")
 		return
 	}
@@ -1029,6 +1043,7 @@ func (c *clientConn) handleExecute(body []byte) {
 			c.sendError("ERROR", errCode, errMsg)
 		}
 		c.setTxError()
+		finishRows()
 		c.logQuery(start, originalQuery, convertedQuery, cmdType, 0, 0, errCode, errMsg, "extended")
 		return
 	}
@@ -1041,14 +1056,15 @@ func (c *clientConn) handleExecute(body []byte) {
 		// entry is written by the leg that completes the portal.
 		keepRowsOpen = true
 		p.exec = &portalExec{
-			rows:           activeRows,
-			cols:           cols,
-			typeOIDs:       typeOIDs,
-			cmdType:        cmdType,
-			rowCount:       int64(rowCount),
-			originalQuery:  originalQuery,
-			convertedQuery: convertedQuery,
-			start:          start,
+			rows:            activeRows,
+			cols:            cols,
+			typeOIDs:        typeOIDs,
+			cmdType:         cmdType,
+			rowCount:        int64(rowCount),
+			originalQuery:   originalQuery,
+			convertedQuery:  convertedQuery,
+			start:           start,
+			finishProfiling: finishProfiling,
 		}
 		_ = wire.WritePortalSuspended(c.writer)
 		return
@@ -1057,6 +1073,7 @@ func (c *clientConn) handleExecute(body []byte) {
 	c.updateTxStatus(cmdType)
 	tag := buildCommandTagFromRowCount(cmdType, int64(rowCount))
 	_ = c.writeCommandComplete(tag)
+	finishRows()
 	c.logQuery(start, originalQuery, convertedQuery, cmdType, int64(rowCount), 0, "", "", "extended")
 }
 

server/conn_query_exec.go

@@ -204,9 +204,9 @@ func (c *clientConn) executeSelectQuery(query string, cmdType string, workerStat
 		}
 		rows, err = runQuery()
 	}
-	c.lastProfilingSummary = observe.EnrichSpanWithProfiling(execCtx, execSpan, execStart, c.executor, c.orgID)
-	execSpan.End()
 	if err != nil {
+		c.lastProfilingSummary = observe.EnrichSpanWithProfiling(execCtx, execSpan, execStart, c.executor, c.orgID)
+		execSpan.End()
 		queryFinalErr = err
 		errCode := classifyErrorCode(err)
 		errMsg := err.Error()
@@ -221,7 +221,11 @@ func (c *clientConn) executeSelectQuery(query string, cmdType string, workerStat
 		_ = c.flushWriter()
 		return 0, errCode, errMsg, nil
 	}
-	defer func() { _ = rows.Close() }()
+	defer func() {
+		_ = rows.Close()
+		c.lastProfilingSummary = observe.EnrichSpanWithProfiling(execCtx, execSpan, execStart, c.executor, c.orgID)
+		execSpan.End()
+	}()
 
 	cols, err := rows.Columns()
 	if err != nil {

server/conn_querylog_lifecycle_test.go

@@ -349,3 +349,79 @@ func TestLifecyclePairFiresOnHandleExecuteQuery(t *testing.T) {
 	c.handleExecute(body)
 	assertLifecyclePair(t, buf, "handleExecute-Query")
 }
+
+func TestHandleExecuteLogsProfileAfterResultStreamCloses(t *testing.T) {
+	c, cleanup := newLifecycleClientConn(t)
+	defer cleanup()
+
+	profile := `{"cpu_time":2.5,"system_peak_buffer_memory":8192,"children":[]}`
+	exec := &closeProfileExecutor{profile: profile}
+	exec.rows = &closeProfileRowSet{onClose: func() { exec.ready = true }}
+	c.executor = exec
+	c.portals["p1"] = &portal{stmt: &preparedStmt{
+		query:          "SELECT 1",
+		convertedQuery: "SELECT 1",
+	}}
+
+	c.handleExecute(append([]byte("p1\x00"), 0, 0, 0, 0))
+
+	select {
+	case entry := <-c.server.queryLogger.ch:
+		if entry.CPUTimeSeconds != 2.5 {
+			t.Fatalf("CPUTimeSeconds = %f, want 2.5", entry.CPUTimeSeconds)
+		}
+		if entry.PeakBufferMemoryBytes != 8192 {
+			t.Fatalf("PeakBufferMemoryBytes = %d, want 8192", entry.PeakBufferMemoryBytes)
+		}
+	default:
+		t.Fatal("expected terminal query-log entry")
+	}
+}
+
+type closeProfileExecutor struct {
+	rows    RowSet
+	profile string
+	ready   bool
+}
+
+func (e *closeProfileExecutor) QueryContext(context.Context, string, ...any) (RowSet, error) {
+	return e.rows, nil
+}
+func (e *closeProfileExecutor) ExecContext(context.Context, string, ...any) (ExecResult, error) {
+	return nil, errors.New("not implemented")
+}
+func (e *closeProfileExecutor) Query(string, ...any) (RowSet, error) { return e.rows, nil }
+func (e *closeProfileExecutor) Exec(string, ...any) (ExecResult, error) {
+	return nil, errors.New("not implemented")
+}
+func (e *closeProfileExecutor) ConnContext(context.Context) (RawConn, error) {
+	return nil, errors.New("not implemented")
+}
+func (e *closeProfileExecutor) PingContext(context.Context) error { return nil }
+func (e *closeProfileExecutor) Close() error                      { return nil }
+func (e *closeProfileExecutor) LastProfilingOutput() string {
+	if !e.ready {
+		return ""
+	}
+	return e.profile
+}
+
+type closeProfileRowSet struct {
+	closed  bool
+	onClose func()
+}
+
+func (*closeProfileRowSet) Columns() ([]string, error) { return []string{"x"}, nil }
+func (*closeProfileRowSet) ColumnTypes() ([]ColumnTyper, error) {
+	return []ColumnTyper{stringColumnTyper{}}, nil
+}
+func (*closeProfileRowSet) Next() bool        { return false }
+func (*closeProfileRowSet) Scan(...any) error { return nil }
+func (r *closeProfileRowSet) Close() error {
+	if !r.closed {
+		r.closed = true
+		r.onClose()
+	}
+	return nil
+}
+func (*closeProfileRowSet) Err() error { return nil }