backend/remote: add the run ID to associate state

backend/remote/backend.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net/url"
+	"os"
 	"sort"
 	"strings"
 	"sync"
@@ -278,6 +279,9 @@ func (b *Remote) State(workspace string) (state.State, error) {
 		client:       b.client,
 		organization: b.organization,
 		workspace:    workspace,
+
+		// This is optionally set during Terraform Enterprise runs.
+		runID: os.Getenv("TFE_RUN_ID"),
 	}
 
 	return &remote.State{Client: client}, nil

backend/remote/backend_mock.go

@@ -343,8 +343,13 @@ func (m *mockStateVersions) List(ctx context.Context, options tfe.StateVersionLi
 
 func (m *mockStateVersions) Create(ctx context.Context, workspaceID string, options tfe.StateVersionCreateOptions) (*tfe.StateVersion, error) {
 	id := generateID("sv-")
+	runID := os.Getenv("TFE_RUN_ID")
 	url := fmt.Sprintf("https://app.terraform.io/_archivist/%s", id)
 
+	if runID != "" && (options.Run == nil || runID != options.Run.ID) {
+		return nil, fmt.Errorf("option.Run.ID does not contain the ID exported by TFE_RUN_ID")
+	}
+
 	sv := &tfe.StateVersion{
 		ID:          id,
 		DownloadURL: url,

backend/remote/backend_state.go

@@ -15,6 +15,7 @@ import (
 type remoteClient struct {
 	client       *tfe.Client
 	organization string
+	runID        string
 	workspace    string
 }
 
@@ -70,7 +71,7 @@ func (r *remoteClient) Put(state []byte) error {
 		return fmt.Errorf("Error retrieving workspace: %v", err)
 	}
 
-	//  the state into a buffer.
+	// Read the raw state into a Terraform state.
 	tfState, err := terraform.ReadState(bytes.NewReader(state))
 	if err != nil {
 		return fmt.Errorf("Error reading state: %s", err)
@@ -83,6 +84,12 @@ func (r *remoteClient) Put(state []byte) error {
 		State:   tfe.String(base64.StdEncoding.EncodeToString(state)),
 	}
 
+	// If we have a run ID, make sure to add it to the options
+	// so the state will be properly associated with the run.
+	if r.runID != "" {
+		options.Run = &tfe.Run{ID: r.runID}
+	}
+
 	// Create the new state.
 	_, err = r.client.StateVersions.Create(ctx, w.ID, options)
 	if err != nil {

backend/remote/backend_state_test.go

@@ -1,9 +1,12 @@
 package remote
 
 import (
+	"bytes"
+	"os"
 	"testing"
 
 	"github.com/hashicorp/terraform/state/remote"
+	"github.com/hashicorp/terraform/terraform"
 )
 
 func TestRemoteClient_impl(t *testing.T) {
@@ -14,3 +17,25 @@ func TestRemoteClient(t *testing.T) {
 	client := testRemoteClient(t)
 	remote.TestClient(t, client)
 }
+
+func TestRemoteClient_withRunID(t *testing.T) {
+	// Set the TFE_RUN_ID environment variable before creating the client!
+	if err := os.Setenv("TFE_RUN_ID", generateID("run-")); err != nil {
+		t.Fatalf("error setting env var TFE_RUN_ID: %v", err)
+	}
+
+	// Create a new test client.
+	client := testRemoteClient(t)
+
+	// Create a new empty state.
+	state := bytes.NewBuffer(nil)
+	if err := terraform.WriteState(terraform.NewState(), state); err != nil {
+		t.Fatalf("expected no error, got: %v", err)
+	}
+
+	// Store the new state to verify (this will be done
+	// by the mock that is used) that the run ID is set.
+	if err := client.Put(state.Bytes()); err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+}