Fix #4712: Make Result.Close() idempotent (#4741)

pkg/db/db_common/client.go

@@ -23,7 +23,7 @@ type Client interface {
 	AcquireSession(context.Context) *AcquireSessionResult
 
 	ExecuteSync(context.Context, string, ...any) (*pqueryresult.SyncQueryResult, error)
-	Execute(context.Context, string, ...any) (*pqueryresult.Result[queryresult.TimingResultStream], error)
+	Execute(context.Context, string, ...any) (*queryresult.Result, error)
 
 	ExecuteSyncInSession(context.Context, *DatabaseSession, string, ...any) (*pqueryresult.SyncQueryResult, error)
 	ExecuteInSession(context.Context, *DatabaseSession, func(), string, ...any) (*queryresult.Result, error)

pkg/db/db_common/execute.go

@@ -18,7 +18,7 @@ func ExecuteQuery(ctx context.Context, client Client, queryString string, args .
 		return nil, err
 	}
 	go func() {
-		resultsStreamer.StreamResult(result)
+		resultsStreamer.StreamResult(result.Result)
 		resultsStreamer.Close()
 	}()
 	return resultsStreamer, nil

pkg/interactive/interactive_client.go

@@ -401,7 +401,7 @@ func (c *InteractiveClient) executeQuery(ctx context.Context, queryCtx context.C
 			querydisplay.DisplayErrorTiming(t)
 		}
 	} else {
-		c.promptResult.Streamer.StreamResult(result)
+		c.promptResult.Streamer.StreamResult(result.Result)
 	}
 }
 

pkg/query/queryexecute/execute.go

@@ -15,7 +15,7 @@ import (
 	"github.com/turbot/pipe-fittings/v2/modconfig"
 	"github.com/turbot/pipe-fittings/v2/pipes"
 	"github.com/turbot/pipe-fittings/v2/querydisplay"
-	"github.com/turbot/pipe-fittings/v2/queryresult"
+	pqueryresult "github.com/turbot/pipe-fittings/v2/queryresult"
 	"github.com/turbot/pipe-fittings/v2/steampipeconfig"
 	"github.com/turbot/pipe-fittings/v2/utils"
 	"github.com/turbot/steampipe/v2/pkg/cmdconfig"
@@ -25,6 +25,7 @@ import (
 	"github.com/turbot/steampipe/v2/pkg/error_helpers"
 	"github.com/turbot/steampipe/v2/pkg/interactive"
 	"github.com/turbot/steampipe/v2/pkg/query"
+	"github.com/turbot/steampipe/v2/pkg/query/queryresult"
 	"github.com/turbot/steampipe/v2/pkg/snapshot"
 )
 
@@ -37,9 +38,11 @@ func RunInteractiveSession(ctx context.Context, initData *query.InitData) error
 
 	// print the data as it comes
 	for r := range result.Streamer.Results {
+		// wrap the result from pipe-fittings with our wrapper that has idempotent Close
+		wrapped := queryresult.WrapResult(r)
 		rowCount, _ := querydisplay.ShowOutput(ctx, r)
 		// show timing
-		display.DisplayTiming(r, rowCount)
+		display.DisplayTiming(wrapped, rowCount)
 		// signal to the resultStreamer that we are done with this chunk of the stream
 		result.Streamer.AllResultsRead()
 	}
@@ -123,6 +126,8 @@ func executeQuery(ctx context.Context, initData *query.InitData, resolvedQuery *
 	rowErrors := 0 // get the number of rows that returned an error
 	// print the data as it comes
 	for r := range resultsStreamer.Results {
+		// wrap the result from pipe-fittings with our wrapper that has idempotent Close
+		wrapped := queryresult.WrapResult(r)
 
 		// if the output format is snapshot or export is set or share/snapshot args are set, we need to generate a snapshot
 		if needSnapshot() {
@@ -133,7 +138,7 @@ func executeQuery(ctx context.Context, initData *query.InitData, resolvedQuery *
 
 			// re-generate the query result from the snapshot. since the row stream in the actual queryresult has been exhausted(while generating the snapshot),
 			// we need to re-generate it for other output formats
-			newQueryResult, err := snapshot.SnapshotToQueryResult[queryresult.TimingContainer](snap, initData.StartTime)
+			newQueryResult, err := snapshot.SnapshotToQueryResult[pqueryresult.TimingContainer](snap, initData.StartTime)
 			if err != nil {
 				return err, 0
 			}
@@ -177,7 +182,7 @@ func executeQuery(ctx context.Context, initData *query.InitData, resolvedQuery *
 			// if other output formats are also needed, we call the querydisplay using the re-generated query result
 			rowCount, _ := querydisplay.ShowOutput(ctx, newQueryResult)
 			// show timing
-			display.DisplayTiming(r, rowCount)
+			display.DisplayTiming(wrapped, rowCount)
 
 			// signal to the resultStreamer that we are done with this result
 			resultsStreamer.AllResultsRead()
@@ -187,7 +192,7 @@ func executeQuery(ctx context.Context, initData *query.InitData, resolvedQuery *
 		// for other output formats, we call the querydisplay code in pipe-fittings
 		rowCount, rowErrs := querydisplay.ShowOutput(ctx, r)
 		// show timing
-		display.DisplayTiming(r, rowCount)
+		display.DisplayTiming(wrapped, rowCount)
 
 		// signal to the resultStreamer that we are done with this result
 		resultsStreamer.AllResultsRead()

pkg/query/queryresult/result.go

@@ -1,12 +1,35 @@
 package queryresult
 
-import "github.com/turbot/pipe-fittings/v2/queryresult"
+import (
+	"sync"
 
-// Result is a type alias for queryresult.Result[TimingResultStream]
-type Result = queryresult.Result[TimingResultStream]
+	"github.com/turbot/pipe-fittings/v2/queryresult"
+)
+
+// Result wraps queryresult.Result[TimingResultStream] with idempotent Close()
+type Result struct {
+	*queryresult.Result[TimingResultStream]
+	closeOnce sync.Once
+}
 
 func NewResult(cols []*queryresult.ColumnDef) *Result {
-	return queryresult.NewResult[TimingResultStream](cols, NewTimingResultStream())
+	return &Result{
+		Result: queryresult.NewResult[TimingResultStream](cols, NewTimingResultStream()),
+	}
+}
+
+// Close closes the row channel in an idempotent manner
+func (r *Result) Close() {
+	r.closeOnce.Do(func() {
+		r.Result.Close()
+	})
+}
+
+// WrapResult wraps a pipe-fittings Result with our wrapper that has idempotent Close
+func WrapResult(r *queryresult.Result[TimingResultStream]) *Result {
+	return &Result{
+		Result: r,
+	}
 }
 
 // ResultStreamer is a type alias for queryresult.ResultStreamer[TimingResultStream]

pkg/query/queryresult/result_test.go

@@ -0,0 +1,25 @@
+package queryresult
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/turbot/pipe-fittings/v2/queryresult"
+)
+
+func TestResultClose_DoubleClose(t *testing.T) {
+	// Create a result with some column definitions
+	cols := []*queryresult.ColumnDef{
+		{Name: "id", DataType: "integer"},
+		{Name: "name", DataType: "text"},
+	}
+	result := NewResult(cols)
+
+	// Close the result once
+	result.Close()
+
+	// Closing again should not panic (idempotent behavior)
+	assert.NotPanics(t, func() {
+		result.Close()
+	}, "Result.Close() should be idempotent and not panic on second call")
+}