Correctly overwrite local state if remote state is newer (#1008)

## Changes

A bug in the code that pulls the remote state could cause the local
state to be empty instead of a copy of the remote state. This happened
only if the local state was present and stale when compared to the
remote version.

We correctly checked for the state serial to see if the local state had
to be replaced but didn't seek back on the remote state before writing
it out. Because the staleness check would read the remote state in full,
copying from the same reader would immediately yield an EOF.

## Tests

* Unit tests for state pull and push mutators that rely on a mocked
filer.
* An integration test that deploys the same bundle from multiple paths,
triggering the staleness logic.

Both failed prior to the fix and now pass.

NOTICE

@@ -16,6 +16,10 @@ go-ini/ini - https://github.com/go-ini/ini
 Copyright ini authors
 License - https://github.com/go-ini/ini/blob/main/LICENSE
 
+uber-go/mock - https://go.uber.org/mock
+Copyright Google Inc.
+License - https://github.com/uber-go/mock/blob/main/LICENSE
+
 —--
 
 This software contains code from the following open source projects, licensed under the MPL 2.0 license:

bundle/deploy/terraform/filer.go

@@ -0,0 +1,14 @@
+package terraform
+
+import (
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/libs/filer"
+)
+
+// filerFunc is a function that returns a filer.Filer.
+type filerFunc func(b *bundle.Bundle) (filer.Filer, error)
+
+// stateFiler returns a filer.Filer that can be used to read/write state files.
+func stateFiler(b *bundle.Bundle) (filer.Filer, error) {
+	return filer.NewWorkspaceFilesClient(b.WorkspaceClient(), b.Config.Workspace.StatePath)
+}

bundle/deploy/terraform/state_pull.go

@@ -1,6 +1,7 @@
 package terraform
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"io"
@@ -13,14 +14,38 @@ import (
 	"github.com/databricks/cli/libs/log"
 )
 
-type statePull struct{}
+type statePull struct {
+	filerFunc
+}
 
 func (l *statePull) Name() string {
 	return "terraform:state-pull"
 }
 
+func (l *statePull) remoteState(ctx context.Context, f filer.Filer) (*bytes.Buffer, error) {
+	// Download state file from filer to local cache directory.
+	remote, err := f.Read(ctx, TerraformStateFileName)
+	if err != nil {
+		// On first deploy this state file doesn't yet exist.
+		if errors.Is(err, fs.ErrNotExist) {
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	defer remote.Close()
+
+	var buf bytes.Buffer
+	_, err = io.Copy(&buf, remote)
+	if err != nil {
+		return nil, err
+	}
+
+	return &buf, nil
+}
+
 func (l *statePull) Apply(ctx context.Context, b *bundle.Bundle) error {
-	f, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(), b.Config.Workspace.StatePath)
+	f, err := l.filerFunc(b)
 	if err != nil {
 		return err
 	}
@@ -32,15 +57,15 @@ func (l *statePull) Apply(ctx context.Context, b *bundle.Bundle) error {
 
 	// Download state file from filer to local cache directory.
 	log.Infof(ctx, "Opening remote state file")
-	remote, err := f.Read(ctx, TerraformStateFileName)
+	remote, err := l.remoteState(ctx, f)
 	if err != nil {
-		// On first deploy this state file doesn't yet exist.
-		if errors.Is(err, fs.ErrNotExist) {
-			log.Infof(ctx, "Remote state file does not exist")
-			return nil
-		}
+		log.Infof(ctx, "Unable to open remote state file: %s", err)
 		return err
 	}
+	if remote == nil {
+		log.Infof(ctx, "Remote state file does not exist")
+		return nil
+	}
 
 	// Expect the state file to live under dir.
 	local, err := os.OpenFile(filepath.Join(dir, TerraformStateFileName), os.O_CREATE|os.O_RDWR, 0600)
@@ -49,7 +74,7 @@ func (l *statePull) Apply(ctx context.Context, b *bundle.Bundle) error {
 	}
 	defer local.Close()
 
-	if !IsLocalStateStale(local, remote) {
+	if !IsLocalStateStale(local, bytes.NewReader(remote.Bytes())) {
 		log.Infof(ctx, "Local state is the same or newer, ignoring remote state")
 		return nil
 	}
@@ -60,7 +85,7 @@ func (l *statePull) Apply(ctx context.Context, b *bundle.Bundle) error {
 
 	// Write file to disk.
 	log.Infof(ctx, "Writing remote state file to local cache directory")
-	_, err = io.Copy(local, remote)
+	_, err = io.Copy(local, bytes.NewReader(remote.Bytes()))
 	if err != nil {
 		return err
 	}
@@ -69,5 +94,5 @@ func (l *statePull) Apply(ctx context.Context, b *bundle.Bundle) error {
 }
 
 func StatePull() bundle.Mutator {
-	return &statePull{}
+	return &statePull{stateFiler}
 }

bundle/deploy/terraform/state_pull_test.go

@@ -0,0 +1,128 @@
+package terraform
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"io/fs"
+	"os"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	mock "github.com/databricks/cli/internal/mocks/libs/filer"
+	"github.com/databricks/cli/libs/filer"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+)
+
+func mockStateFilerForPull(t *testing.T, contents map[string]int, merr error) filer.Filer {
+	buf, err := json.Marshal(contents)
+	require.NoError(t, err)
+
+	ctrl := gomock.NewController(t)
+	mock := mock.NewMockFiler(ctrl)
+	mock.
+		EXPECT().
+		Read(gomock.Any(), gomock.Eq(TerraformStateFileName)).
+		Return(io.NopCloser(bytes.NewReader(buf)), merr).
+		Times(1)
+	return mock
+}
+
+func statePullTestBundle(t *testing.T) *bundle.Bundle {
+	return &bundle.Bundle{
+		Config: config.Root{
+			Bundle: config.Bundle{
+				Target: "default",
+			},
+			Path: t.TempDir(),
+		},
+	}
+}
+
+func TestStatePullLocalMissingRemoteMissing(t *testing.T) {
+	m := &statePull{
+		identityFiler(mockStateFilerForPull(t, nil, os.ErrNotExist)),
+	}
+
+	ctx := context.Background()
+	b := statePullTestBundle(t)
+	err := bundle.Apply(ctx, b, m)
+	assert.NoError(t, err)
+
+	// Confirm that no local state file has been written.
+	_, err = os.Stat(localStateFile(t, ctx, b))
+	assert.ErrorIs(t, err, fs.ErrNotExist)
+}
+
+func TestStatePullLocalMissingRemotePresent(t *testing.T) {
+	m := &statePull{
+		identityFiler(mockStateFilerForPull(t, map[string]int{"serial": 5}, nil)),
+	}
+
+	ctx := context.Background()
+	b := statePullTestBundle(t)
+	err := bundle.Apply(ctx, b, m)
+	assert.NoError(t, err)
+
+	// Confirm that the local state file has been updated.
+	localState := readLocalState(t, ctx, b)
+	assert.Equal(t, map[string]int{"serial": 5}, localState)
+}
+
+func TestStatePullLocalStale(t *testing.T) {
+	m := &statePull{
+		identityFiler(mockStateFilerForPull(t, map[string]int{"serial": 5}, nil)),
+	}
+
+	ctx := context.Background()
+	b := statePullTestBundle(t)
+
+	// Write a stale local state file.
+	writeLocalState(t, ctx, b, map[string]int{"serial": 4})
+	err := bundle.Apply(ctx, b, m)
+	assert.NoError(t, err)
+
+	// Confirm that the local state file has been updated.
+	localState := readLocalState(t, ctx, b)
+	assert.Equal(t, map[string]int{"serial": 5}, localState)
+}
+
+func TestStatePullLocalEqual(t *testing.T) {
+	m := &statePull{
+		identityFiler(mockStateFilerForPull(t, map[string]int{"serial": 5, "some_other_key": 123}, nil)),
+	}
+
+	ctx := context.Background()
+	b := statePullTestBundle(t)
+
+	// Write a local state file with the same serial as the remote.
+	writeLocalState(t, ctx, b, map[string]int{"serial": 5})
+	err := bundle.Apply(ctx, b, m)
+	assert.NoError(t, err)
+
+	// Confirm that the local state file has not been updated.
+	localState := readLocalState(t, ctx, b)
+	assert.Equal(t, map[string]int{"serial": 5}, localState)
+}
+
+func TestStatePullLocalNewer(t *testing.T) {
+	m := &statePull{
+		identityFiler(mockStateFilerForPull(t, map[string]int{"serial": 5, "some_other_key": 123}, nil)),
+	}
+
+	ctx := context.Background()
+	b := statePullTestBundle(t)
+
+	// Write a local state file with a newer serial as the remote.
+	writeLocalState(t, ctx, b, map[string]int{"serial": 6})
+	err := bundle.Apply(ctx, b, m)
+	assert.NoError(t, err)
+
+	// Confirm that the local state file has not been updated.
+	localState := readLocalState(t, ctx, b)
+	assert.Equal(t, map[string]int{"serial": 6}, localState)
+}

bundle/deploy/terraform/state_push.go

@@ -10,14 +10,16 @@ import (
 	"github.com/databricks/cli/libs/log"
 )
 
-type statePush struct{}
+type statePush struct {
+	filerFunc
+}
 
 func (l *statePush) Name() string {
 	return "terraform:state-push"
 }
 
 func (l *statePush) Apply(ctx context.Context, b *bundle.Bundle) error {
-	f, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(), b.Config.Workspace.StatePath)
+	f, err := l.filerFunc(b)
 	if err != nil {
 		return err
 	}
@@ -45,5 +47,5 @@ func (l *statePush) Apply(ctx context.Context, b *bundle.Bundle) error {
 }
 
 func StatePush() bundle.Mutator {
-	return &statePush{}
+	return &statePush{stateFiler}
 }

bundle/deploy/terraform/state_push_test.go

@@ -0,0 +1,63 @@
+package terraform
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	mock "github.com/databricks/cli/internal/mocks/libs/filer"
+	"github.com/databricks/cli/libs/filer"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/mock/gomock"
+)
+
+func mockStateFilerForPush(t *testing.T, fn func(body io.Reader)) filer.Filer {
+	ctrl := gomock.NewController(t)
+	mock := mock.NewMockFiler(ctrl)
+	mock.
+		EXPECT().
+		Write(gomock.Any(), gomock.Any(), gomock.Any(), filer.CreateParentDirectories, filer.OverwriteIfExists).
+		Do(func(ctx context.Context, path string, reader io.Reader, mode ...filer.WriteMode) error {
+			fn(reader)
+			return nil
+		}).
+		Return(nil).
+		Times(1)
+	return mock
+}
+
+func statePushTestBundle(t *testing.T) *bundle.Bundle {
+	return &bundle.Bundle{
+		Config: config.Root{
+			Bundle: config.Bundle{
+				Target: "default",
+			},
+			Path: t.TempDir(),
+		},
+	}
+}
+
+func TestStatePush(t *testing.T) {
+	mock := mockStateFilerForPush(t, func(body io.Reader) {
+		dec := json.NewDecoder(body)
+		var contents map[string]int
+		err := dec.Decode(&contents)
+		assert.NoError(t, err)
+		assert.Equal(t, map[string]int{"serial": 4}, contents)
+	})
+
+	m := &statePush{
+		identityFiler(mock),
+	}
+
+	ctx := context.Background()
+	b := statePushTestBundle(t)
+
+	// Write a stale local state file.
+	writeLocalState(t, ctx, b, map[string]int{"serial": 4})
+	err := bundle.Apply(ctx, b, m)
+	assert.NoError(t, err)
+}

bundle/deploy/terraform/state_test.go

@@ -0,0 +1,48 @@
+package terraform
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/libs/filer"
+	"github.com/stretchr/testify/require"
+)
+
+// identityFiler returns a filerFunc that returns the specified filer.
+func identityFiler(f filer.Filer) filerFunc {
+	return func(_ *bundle.Bundle) (filer.Filer, error) {
+		return f, nil
+	}
+}
+
+func localStateFile(t *testing.T, ctx context.Context, b *bundle.Bundle) string {
+	dir, err := Dir(ctx, b)
+	require.NoError(t, err)
+	return filepath.Join(dir, TerraformStateFileName)
+}
+
+func readLocalState(t *testing.T, ctx context.Context, b *bundle.Bundle) map[string]int {
+	f, err := os.Open(localStateFile(t, ctx, b))
+	require.NoError(t, err)
+	defer f.Close()
+
+	var contents map[string]int
+	dec := json.NewDecoder(f)
+	err = dec.Decode(&contents)
+	require.NoError(t, err)
+	return contents
+}
+
+func writeLocalState(t *testing.T, ctx context.Context, b *bundle.Bundle, contents map[string]int) {
+	f, err := os.Create(localStateFile(t, ctx, b))
+	require.NoError(t, err)
+	defer f.Close()
+
+	enc := json.NewEncoder(f)
+	err = enc.Encode(contents)
+	require.NoError(t, err)
+}