query.go

package backend

import (
	"context"

	opentracing "github.com/opentracing/opentracing-go"

	"github.com/pulumi/pulumi/pkg/v2/backend/display"
	"github.com/pulumi/pulumi/pkg/v2/engine"
	"github.com/pulumi/pulumi/sdk/v2/go/common/util/result"
)

type MakeQuery func(context.Context, QueryOperation) (engine.QueryInfo, error)

// RunQuery executes a query program against the resource outputs of a locally hosted stack.
func RunQuery(ctx context.Context, b Backend, op QueryOperation,
	callerEventsOpt chan<- engine.Event, newQuery MakeQuery) result.Result {
	q, err := newQuery(ctx, op)
	if err != nil {
		return result.FromError(err)
	}

	// Render query output to CLI.
	displayEvents := make(chan engine.Event)
	displayDone := make(chan bool)
	go display.ShowQueryEvents("running query", displayEvents, displayDone, op.Opts.Display)

	// The engineEvents channel receives all events from the engine, which we then forward onto other
	// channels for actual processing. (displayEvents and callerEventsOpt.)
	engineEvents := make(chan engine.Event)
	eventsDone := make(chan bool)
	go func() {
		for e := range engineEvents {
			displayEvents <- e
			if callerEventsOpt != nil {
				callerEventsOpt <- e
			}
		}

		close(eventsDone)
	}()

	// Depending on the action, kick off the relevant engine activity.  Note that we don't immediately check and
	// return error conditions, because we will do so below after waiting for the display channels to close.
	cancellationScope := op.Scopes.NewScope(engineEvents, true /*dryRun*/)
	engineCtx := &engine.Context{
		Cancel:        cancellationScope.Context(),
		Events:        engineEvents,
		BackendClient: NewBackendClient(b),
	}
	if parentSpan := opentracing.SpanFromContext(ctx); parentSpan != nil {
		engineCtx.ParentSpan = parentSpan.Context()
	}

	res := engine.Query(engineCtx, q, op.Opts.Engine)

	// Wait for dependent channels to finish processing engineEvents before closing.
	<-displayDone
	cancellationScope.Close() // Don't take any cancellations anymore, we're shutting down.
	close(engineEvents)

	// Make sure that the goroutine writing to displayEvents and callerEventsOpt
	// has exited before proceeding
	<-eventsDone
	close(displayEvents)

	return res
}