From 95cf69aa3284248ae301024fdd83e0bcbf927d9f Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 23 Feb 2015 21:57:17 -0800
Subject: [PATCH] command: StateHook for continous state updates

---
 command/apply.go           | 15 ++++++++++++++-
 command/hook_state.go      | 33 +++++++++++++++++++++++++++++++++
 command/hook_state_test.go | 29 +++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)
 create mode 100644 command/hook_state.go
 create mode 100644 command/hook_state_test.go

diff --git a/command/apply.go b/command/apply.go
index a568b0431..d46b71679 100644
--- a/command/apply.go
+++ b/command/apply.go
@@ -68,7 +68,8 @@ func (c *ApplyCommand) Run(args []string) int {
 
 	// Prepare the extra hooks to count resources
 	countHook := new(CountHook)
-	c.Meta.extraHooks = []terraform.Hook{countHook}
+	stateHook := new(StateHook)
+	c.Meta.extraHooks = []terraform.Hook{countHook, stateHook}
 
 	if !c.Destroy && maybeInit {
 		// Do a detect to determine if we need to do an init + apply.
@@ -151,6 +152,18 @@ func (c *ApplyCommand) Run(args []string) int {
 		}
 	}
 
+	// Setup the state hook for continous state updates
+	{
+		state, err := c.State()
+		if err != nil {
+			c.Ui.Error(fmt.Sprintf(
+				"Error reading state: %s", err))
+			return 1
+		}
+
+		stateHook.State = state
+	}
+
 	// Start the apply in a goroutine so that we can be interrupted.
 	var state *terraform.State
 	var applyErr error
diff --git a/command/hook_state.go b/command/hook_state.go
new file mode 100644
index 000000000..ab5c47a11
--- /dev/null
+++ b/command/hook_state.go
@@ -0,0 +1,33 @@
+package command
+
+import (
+	"sync"
+
+	"github.com/hashicorp/terraform/state"
+	"github.com/hashicorp/terraform/terraform"
+)
+
+// StateHook is a hook that continuously updates the state by calling
+// WriteState on a state.State.
+type StateHook struct {
+	terraform.NilHook
+	sync.Mutex
+
+	State state.State
+}
+
+func (h *StateHook) PostStateUpdate(
+	s *terraform.State) (terraform.HookAction, error) {
+	h.Lock()
+	defer h.Unlock()
+
+	if h.State != nil {
+		// Write the new state
+		if err := h.State.WriteState(s); err != nil {
+			return terraform.HookActionHalt, err
+		}
+	}
+
+	// Continue forth
+	return terraform.HookActionContinue, nil
+}
diff --git a/command/hook_state_test.go b/command/hook_state_test.go
new file mode 100644
index 000000000..0d0fd7927
--- /dev/null
+++ b/command/hook_state_test.go
@@ -0,0 +1,29 @@
+package command
+
+import (
+	"testing"
+
+	"github.com/hashicorp/terraform/state"
+	"github.com/hashicorp/terraform/terraform"
+)
+
+func TestStateHook_impl(t *testing.T) {
+	var _ terraform.Hook = new(StateHook)
+}
+
+func TestStateHook(t *testing.T) {
+	is := &state.InmemState{}
+	var hook terraform.Hook = &StateHook{State: is}
+
+	s := state.TestStateInitial()
+	action, err := hook.PostStateUpdate(s)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if action != terraform.HookActionContinue {
+		t.Fatalf("bad: %v", action)
+	}
+	if !is.State().Equal(s) {
+		t.Fatalf("bad state: %#v", is.State())
+	}
+}