Correctly overwrite local state if remote state is newer

bundle/deploy/terraform/filer.go

@@ -0,0 +1,21 @@
+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)
+}
+
+// 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
+	}
+}

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,40 @@
+package terraform
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/stretchr/testify/require"
+)
+
+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)
+}