adds Logs method to QueryRuns

Author: brandonc

helper_test.go

@@ -1252,6 +1252,27 @@ func createRunWaitForAnyStatuses(t *testing.T, client *Client, w *Workspace, sta
 	}
 }
 
+func createQueryRunWaitForAnyStatuses(t *testing.T, client *Client, w *Workspace, statuses []QueryRunStatus) (*QueryRun, func()) {
+	ctx := context.Background()
+	qr := createQueryRun(t, client, w)
+
+	timeout := 2 * time.Minute
+
+	ctxPollQueryRunReady, cancelPollQueryRunReady := context.WithTimeout(ctx, timeout)
+
+	run := pollQueryRunStatus(
+		t,
+		client,
+		ctxPollQueryRunReady,
+		qr,
+		append(statuses, QueryRunErrored),
+	)
+
+	return run, func() {
+		cancelPollQueryRunReady()
+	}
+}
+
 func applyableStatuses(r *Run) []RunStatus {
 	if len(r.PolicyChecks) > 0 {
 		return []RunStatus{
@@ -1298,6 +1319,39 @@ func pollRunStatus(t *testing.T, client *Client, ctx context.Context, r *Run, rs
 	return r
 }
 
+// pollQueryRunStatus will poll the given query run until its status matches one of the given run statuses or the given context
+// times out.
+func pollQueryRunStatus(t *testing.T, client *Client, ctx context.Context, q *QueryRun, rss []QueryRunStatus) *QueryRun {
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		t.Logf("No deadline was set to poll query run %q which could result in an infinite loop", q.ID)
+	}
+
+	t.Logf("Polling query run %q for status included in %q with deadline of %s", q.ID, rss, deadline)
+
+	ticker := time.NewTicker(2 * time.Second)
+	defer ticker.Stop()
+
+	for finished := false; !finished; {
+		t.Log("...")
+		select {
+		case <-ctx.Done():
+			t.Fatalf("Run %q had status %q at deadline", q.ID, q.Status)
+		case <-ticker.C:
+			q = readQueryRun(t, client, ctx, q)
+			t.Logf("Query Run %q had status %q", q.ID, q.Status)
+			for _, rs := range rss {
+				if rs == q.Status {
+					finished = true
+					break
+				}
+			}
+		}
+	}
+
+	return q
+}
+
 // pollStateVersionStatus will poll the given state version until its status
 // matches one of the given statuses or the given context times out.
 func pollStateVersionStatus(t *testing.T, client *Client, ctx context.Context, sv *StateVersion, statuses []StateVersionStatus) *StateVersion {
@@ -1347,6 +1401,18 @@ func readRun(t *testing.T, client *Client, ctx context.Context, r *Run) *Run {
 	return rr
 }
 
+// readQueryRun will re-read the given query run.
+func readQueryRun(t *testing.T, client *Client, ctx context.Context, r *QueryRun) *QueryRun {
+	t.Logf("Reading query run %q", r.ID)
+
+	qr, err := client.QueryRuns.Read(ctx, r.ID)
+	if err != nil {
+		t.Fatalf("Could not read run %q: %s", r.ID, err)
+	}
+
+	return qr
+}
+
 // applyRun will apply the given run.
 func applyRun(t *testing.T, client *Client, ctx context.Context, r *Run) {
 	t.Logf("Applying run %q", r.ID)

query_runs.go

@@ -6,6 +6,7 @@ package tfe
 import (
 	"context"
 	"fmt"
+	"io"
 	"net/url"
 	"time"
 )
@@ -31,6 +32,9 @@ type QueryRuns interface {
 	// ReadWithOptions reads a query run by its ID using the options supplied
 	ReadWithOptions(ctx context.Context, queryRunID string, options *QueryRunReadOptions) (*QueryRun, error)
 
+	// Logs retrieves the logs of a query run.
+	Logs(ctx context.Context, queryRunID string) (io.Reader, error)
+
 	// Cancel a query run by its ID.
 	Cancel(ctx context.Context, runID string) error
 
@@ -46,18 +50,17 @@ type QueryRunCreateOptions struct {
 	// https://jsonapi.org/format/#crud-creating
 	Type string `jsonapi:"primary,queries"`
 
-	// TerraformVersion specifies the Terraform version to use in this run.
-	// Only valid for plan-only runs; must be a valid Terraform version available to the organization.
+	// TerraformVersion specifies the Terraform version to use in this query run.
 	TerraformVersion *string `jsonapi:"attr,terraform-version,omitempty"`
 
 	Source QueryRunSource `jsonapi:"attr,source"`
 
-	// Specifies the configuration version to use for this run. If the
+	// Specifies the configuration version to use for this query run. If the
 	// configuration version object is omitted, the run will be created using the
 	// workspace's latest configuration version.
 	ConfigurationVersion *ConfigurationVersion `jsonapi:"relation,configuration-version"`
 
-	// Specifies the workspace where the run will be executed.
+	// Specifies the workspace where the query run will be executed.
 	Workspace *Workspace `jsonapi:"relation,workspace"`
 
 	// Variables allows you to specify terraform input variables for
@@ -82,6 +85,19 @@ type QueryRunIncludeOpt string
 // QueryRunSource represents the available sources for query runs.
 type QueryRunSource string
 
+// QueryRunStatus is the query run state
+type QueryRunStatus string
+
+// List all available run statuses.
+const (
+	QueryRunCanceled QueryRunStatus = "canceled"
+	QueryRunErrored  QueryRunStatus = "errored"
+	QueryRunPending  QueryRunStatus = "pending"
+	QueryRunQueuing  QueryRunStatus = "queued"
+	QueryRunRunning  QueryRunStatus = "running"
+	QueryRunFinished QueryRunStatus = "queuing_apply"
+)
+
 // List all available run sources.
 const (
 	QueryRunSourceAPI QueryRunSource = "tfe-api"
@@ -116,13 +132,14 @@ type QueryRunReadOptions struct {
 
 // QueryRun represents a Terraform Enterprise query run.
 type QueryRun struct {
-	ID               string               `jsonapi:"primary,queries"`
-	CreatedAt        time.Time            `jsonapi:"attr,created-at,iso8601"`
-	Source           RunSource            `jsonapi:"attr,source"`
-	Status           RunStatus            `jsonapi:"attr,status"`
-	StatusTimestamps *RunStatusTimestamps `jsonapi:"attr,status-timestamps"`
-	TerraformVersion string               `jsonapi:"attr,terraform-version"`
-	Variables        []*RunVariableAttr   `jsonapi:"attr,variables"`
+	ID               string                    `jsonapi:"primary,queries"`
+	CreatedAt        time.Time                 `jsonapi:"attr,created-at,iso8601"`
+	Source           QueryRunSource            `jsonapi:"attr,source"`
+	Status           QueryRunStatus            `jsonapi:"attr,status"`
+	StatusTimestamps *QueryRunStatusTimestamps `jsonapi:"attr,status-timestamps"`
+	TerraformVersion string                    `jsonapi:"attr,terraform-version"`
+	Variables        []*RunVariableAttr        `jsonapi:"attr,variables"`
+	LogReadURL       string                    `jsonapi:"attr,log-read-url"`
 
 	// Relations
 	ConfigurationVersion *ConfigurationVersion `jsonapi:"relation,configuration-version"`
@@ -206,6 +223,49 @@ func (r *queryRuns) ReadWithOptions(ctx context.Context, queryRunID string, opti
 	return &run, nil
 }
 
+func (r *queryRuns) Logs(ctx context.Context, queryRunID string) (io.Reader, error) {
+	if !validStringID(&queryRunID) {
+		return nil, ErrInvalidQueryRunID
+	}
+
+	// Get the query to make sure it exists.
+	q, err := r.Read(ctx, queryRunID)
+	if err != nil {
+		return nil, err
+	}
+
+	// Return an error if the log URL is empty.
+	if q.LogReadURL == "" {
+		return nil, fmt.Errorf("query %s does not have a log URL", queryRunID)
+	}
+
+	u, err := url.Parse(q.LogReadURL)
+	if err != nil {
+		return nil, fmt.Errorf("invalid log URL: %w", err)
+	}
+
+	done := func() (bool, error) {
+		p, err := r.Read(ctx, q.ID)
+		if err != nil {
+			return false, err
+		}
+
+		switch p.Status {
+		case QueryRunCanceled, QueryRunErrored, QueryRunFinished:
+			return true, nil
+		default:
+			return false, nil
+		}
+	}
+
+	return &LogReader{
+		client: r.client,
+		ctx:    ctx,
+		done:   done,
+		logURL: u,
+	}, nil
+}
+
 func (r *queryRuns) Cancel(ctx context.Context, queryRunID string) error {
 	if queryRunID == "" {
 		return ErrInvalidQueryRunID

query_runs_integration_test.go

@@ -5,6 +5,7 @@ package tfe
 
 import (
 	"context"
+	"io"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -166,11 +167,12 @@ func TestQueryRunsCreate(t *testing.T) {
 		assert.Equal(t, len(vars), len(qr.Variables))
 
 		for _, v := range qr.Variables {
-			if v.Key == "test_foo" {
+			switch v.Key {
+			case "test_foo":
 				assert.Equal(t, v.Value, "Hello, Foo!")
-			} else if v.Key == "test_variable" {
+			case "test_variable":
 				assert.Equal(t, v.Value, "Hello, World!")
-			} else {
+			default:
 				t.Fatalf("Unexpected variable key: %s", v.Key)
 			}
 		}
@@ -263,6 +265,28 @@ func TestQueryRunsCancel(t *testing.T) {
 	})
 }
 
+func TestQueryRunsLogs(t *testing.T) {
+	skipUnlessBeta(t)
+	client := testClient(t)
+	ctx := context.Background()
+
+	wTest, wTestCleanup := createWorkspace(t, client, nil)
+	defer wTestCleanup()
+
+	qr, cleanup := createQueryRunWaitForAnyStatuses(t, client, wTest, []QueryRunStatus{QueryRunErrored, QueryRunFinished})
+	t.Cleanup(cleanup)
+
+	t.Run("when the query run exists", func(t *testing.T) {
+		// We assume the second query run is in a state that can be canceled.
+		reader, err := client.QueryRuns.Logs(ctx, qr.ID)
+		require.NoError(t, err)
+
+		logs, err := io.ReadAll(reader)
+		require.NoError(t, err)
+		assert.NotEmpty(t, logs, "some logs should be returned")
+	})
+}
+
 func TestQueryRunsForceCancel(t *testing.T) {
 	skipUnlessBeta(t)
 	client := testClient(t)