From 04e77102c90bffd49de389b2d390cbe392f39f7f Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 30 Mar 2023 12:01:09 +0200
Subject: [PATCH] Add mutators to pull and push Terraform state (#288)

## Changes

Pull state before deploying and push state after deploying.

Note: the run command was missing mutators to initialize Terraform. This
is necessary if the cache directory is removed between running "deploy"
and "run" (which is valid now that we synchronize state).

## Tests

Manually.
---
 bundle/deploy/terraform/load.go       | 12 ++++++
 bundle/deploy/terraform/pkg.go        |  3 ++
 bundle/deploy/terraform/state_pull.go | 62 +++++++++++++++++++++++++++
 bundle/deploy/terraform/state_push.go | 48 +++++++++++++++++++++
 bundle/phases/deploy.go               |  2 +
 cmd/bundle/run.go                     |  3 ++
 6 files changed, 130 insertions(+)
 create mode 100644 bundle/deploy/terraform/pkg.go
 create mode 100644 bundle/deploy/terraform/state_pull.go
 create mode 100644 bundle/deploy/terraform/state_push.go

diff --git a/bundle/deploy/terraform/load.go b/bundle/deploy/terraform/load.go
index efd329f4..7f350abf 100644
--- a/bundle/deploy/terraform/load.go
+++ b/bundle/deploy/terraform/load.go
@@ -2,8 +2,10 @@ package terraform
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/databricks/bricks/bundle"
+	"github.com/hashicorp/terraform-exec/tfexec"
 )
 
 type load struct{}
@@ -13,6 +15,16 @@ func (l *load) Name() string {
 }
 
 func (l *load) Apply(ctx context.Context, b *bundle.Bundle) ([]bundle.Mutator, error) {
+	tf := b.Terraform
+	if tf == nil {
+		return nil, fmt.Errorf("terraform not initialized")
+	}
+
+	err := tf.Init(ctx, tfexec.Upgrade(true))
+	if err != nil {
+		return nil, fmt.Errorf("terraform init: %w", err)
+	}
+
 	state, err := b.Terraform.Show(ctx)
 	if err != nil {
 		return nil, err
diff --git a/bundle/deploy/terraform/pkg.go b/bundle/deploy/terraform/pkg.go
new file mode 100644
index 00000000..5e3807be
--- /dev/null
+++ b/bundle/deploy/terraform/pkg.go
@@ -0,0 +1,3 @@
+package terraform
+
+const TerraformStateFileName = "terraform.tfstate"
diff --git a/bundle/deploy/terraform/state_pull.go b/bundle/deploy/terraform/state_pull.go
new file mode 100644
index 00000000..04a6b2eb
--- /dev/null
+++ b/bundle/deploy/terraform/state_pull.go
@@ -0,0 +1,62 @@
+package terraform
+
+import (
+	"context"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/databricks/bricks/bundle"
+	"github.com/databricks/bricks/libs/filer"
+	"github.com/databricks/bricks/libs/log"
+	"github.com/databricks/databricks-sdk-go/apierr"
+)
+
+type statePull struct{}
+
+func (l *statePull) Name() string {
+	return "terraform:state-pull"
+}
+
+func (l *statePull) Apply(ctx context.Context, b *bundle.Bundle) ([]bundle.Mutator, error) {
+	f, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(), b.Config.Workspace.StatePath.Workspace)
+	if err != nil {
+		return nil, err
+	}
+
+	dir, err := Dir(b)
+	if err != nil {
+		return nil, err
+	}
+
+	// Download state file from filer to local cache directory.
+	log.Infof(ctx, "Opening remote state file")
+	remote, err := f.Read(ctx, TerraformStateFileName)
+	if err != nil {
+		// On first deploy this state file doesn't yet exist.
+		if apierr.IsMissing(err) {
+			log.Infof(ctx, "Remote state file does not exist")
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	// Expect the state file to live under dir.
+	local, err := os.OpenFile(filepath.Join(dir, TerraformStateFileName), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	if err != nil {
+		return nil, err
+	}
+
+	// Write file to disk.
+	log.Infof(ctx, "Writing remote state file to local cache directory")
+	_, err = io.Copy(local, remote)
+	if err != nil {
+		return nil, err
+	}
+
+	return nil, nil
+}
+
+func StatePull() bundle.Mutator {
+	return &statePull{}
+}
diff --git a/bundle/deploy/terraform/state_push.go b/bundle/deploy/terraform/state_push.go
new file mode 100644
index 00000000..03c1bc4c
--- /dev/null
+++ b/bundle/deploy/terraform/state_push.go
@@ -0,0 +1,48 @@
+package terraform
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+
+	"github.com/databricks/bricks/bundle"
+	"github.com/databricks/bricks/libs/filer"
+	"github.com/databricks/bricks/libs/log"
+)
+
+type statePush struct{}
+
+func (l *statePush) Name() string {
+	return "terraform:state-push"
+}
+
+func (l *statePush) Apply(ctx context.Context, b *bundle.Bundle) ([]bundle.Mutator, error) {
+	f, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(), b.Config.Workspace.StatePath.Workspace)
+	if err != nil {
+		return nil, err
+	}
+
+	dir, err := Dir(b)
+	if err != nil {
+		return nil, err
+	}
+
+	// Expect the state file to live under dir.
+	local, err := os.Open(filepath.Join(dir, TerraformStateFileName))
+	if err != nil {
+		return nil, err
+	}
+
+	// Upload state file from local cache directory to filer.
+	log.Infof(ctx, "Writing local state file to remote state directory")
+	err = f.Write(ctx, TerraformStateFileName, local, filer.CreateParentDirectories, filer.OverwriteIfExists)
+	if err != nil {
+		return nil, err
+	}
+
+	return nil, nil
+}
+
+func StatePush() bundle.Mutator {
+	return &statePush{}
+}
diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go
index 299cd3eb..2b656c8d 100644
--- a/bundle/phases/deploy.go
+++ b/bundle/phases/deploy.go
@@ -18,7 +18,9 @@ func Deploy() bundle.Mutator {
 			artifacts.UploadAll(),
 			terraform.Interpolate(),
 			terraform.Write(),
+			terraform.StatePull(),
 			terraform.Apply(),
+			terraform.StatePush(),
 			lock.Release(),
 		},
 	)
diff --git a/cmd/bundle/run.go b/cmd/bundle/run.go
index 51d81e11..431290d7 100644
--- a/cmd/bundle/run.go
+++ b/cmd/bundle/run.go
@@ -25,6 +25,9 @@ var runCmd = &cobra.Command{
 		b := bundle.Get(cmd.Context())
 		err := bundle.Apply(cmd.Context(), b, []bundle.Mutator{
 			phases.Initialize(),
+			terraform.Interpolate(),
+			terraform.Write(),
+			terraform.StatePull(),
 			terraform.Load(),
 		})
 		if err != nil {