From 899fbab1b2ccabf2bd7c71251e1b7d640e8f7b36 Mon Sep 17 00:00:00 2001
From: Ventsislav Georgiev <ventsislav.georgiev@live.com>
Date: Sun, 17 Nov 2024 21:06:23 +0200
Subject: [PATCH] feat: add state lock command

---
 commands.go                        |   6 ++
 internal/command/clistate/state.go |  10 +++
 internal/command/state_lock.go     | 120 +++++++++++++++++++++++++++++
 3 files changed, 136 insertions(+)
 create mode 100644 internal/command/state_lock.go

diff --git a/commands.go b/commands.go
index 6796091e3306..7419ce8ad849 100644
--- a/commands.go
+++ b/commands.go
@@ -372,6 +372,12 @@ func initCommands(
 			}, nil
 		},
 
+		"state lock": func() (cli.Command, error) {
+			return &command.StateLockCommand{
+				Meta: meta,
+			}, nil
+		},
+
 		"state rm": func() (cli.Command, error) {
 			return &command.StateRmCommand{
 				StateMeta: command.StateMeta{
diff --git a/internal/command/clistate/state.go b/internal/command/clistate/state.go
index 64cd574cf5e6..8c90e5fc0def 100644
--- a/internal/command/clistate/state.go
+++ b/internal/command/clistate/state.go
@@ -65,6 +65,8 @@ type Locker interface {
 
 	// Timeout returns the configured timeout duration
 	Timeout() time.Duration
+
+	LockID() string
 }
 
 type locker struct {
@@ -167,6 +169,10 @@ func (l *locker) Timeout() time.Duration {
 	return l.timeout
 }
 
+func (l *locker) LockID() string {
+	return l.lockID
+}
+
 type noopLocker struct{}
 
 // NewNoopLocker returns a valid Locker that does nothing.
@@ -191,3 +197,7 @@ func (l noopLocker) Unlock() tfdiags.Diagnostics {
 func (l noopLocker) Timeout() time.Duration {
 	return 0
 }
+
+func (l noopLocker) LockID() string {
+	return ""
+}
diff --git a/internal/command/state_lock.go b/internal/command/state_lock.go
new file mode 100644
index 000000000000..7b2bbffa2534
--- /dev/null
+++ b/internal/command/state_lock.go
@@ -0,0 +1,120 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package command
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/hashicorp/terraform/internal/command/arguments"
+	"github.com/hashicorp/terraform/internal/command/clistate"
+	"github.com/hashicorp/terraform/internal/command/views"
+	"github.com/hashicorp/terraform/internal/states/statemgr"
+
+	"github.com/hashicorp/cli"
+	"github.com/hashicorp/terraform/internal/tfdiags"
+)
+
+// StateLockCommand is a cli.Command implementation that manually locks
+// the state.
+type StateLockCommand struct {
+	Meta
+	StateMeta
+}
+
+func (c *StateLockCommand) Run(args []string) int {
+	args = c.Meta.process(args)
+	var statePath string
+	cmdFlags := c.Meta.defaultFlagSet("state lock")
+	cmdFlags.DurationVar(&c.Meta.stateLockTimeout, "lock-timeout", 0, "lock timeout")
+	cmdFlags.StringVar(&statePath, "state", "", "path")
+	if err := cmdFlags.Parse(args); err != nil {
+		c.Ui.Error(fmt.Sprintf("Error parsing command-line flags: %s\n", err.Error()))
+		return cli.RunResultHelp
+	}
+	args = cmdFlags.Args()
+
+	if statePath != "" {
+		c.Meta.statePath = statePath
+	}
+
+	// assume everything is initialized. The user can manually init if this is
+	// required.
+	configPath, err := ModulePath(args)
+	if err != nil {
+		c.Ui.Error(err.Error())
+		return 1
+	}
+
+	var diags tfdiags.Diagnostics
+
+	backendConfig, backendDiags := c.loadBackendConfig(configPath)
+	diags = diags.Append(backendDiags)
+	if diags.HasErrors() {
+		c.showDiagnostics(diags)
+		return 1
+	}
+
+	// Load the backend
+	b, backendDiags := c.Backend(&BackendOpts{
+		Config: backendConfig,
+	})
+	diags = diags.Append(backendDiags)
+	if backendDiags.HasErrors() {
+		c.showDiagnostics(diags)
+		return 1
+	}
+
+	env, err := c.Workspace()
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error selecting workspace: %s", err))
+		return 1
+	}
+	stateMgr, err := b.StateMgr(env)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Failed to load state: %s", err))
+		return 1
+	}
+
+	_, isLocal := stateMgr.(*statemgr.Filesystem)
+	if isLocal {
+		c.Ui.Error("Local state locking is redundant.")
+		return 0
+	}
+
+	stateLocker := clistate.NewLocker(c.stateLockTimeout, views.NewStateLocker(arguments.ViewHuman, c.View))
+	if diags := stateLocker.Lock(stateMgr, "state-lock"); diags.HasErrors() {
+		c.showDiagnostics(diags)
+		return 1
+	}
+
+	c.Ui.Output(c.Colorize().Color(strings.TrimSpace(fmt.Sprintf(outputStateLockSuccess, stateLocker.LockID()))))
+	return 0
+}
+
+func (c *StateLockCommand) Help() string {
+	helpText := `
+Usage: terraform [global options] state lock
+
+  Manually lock the state for the defined configuration.
+
+  This will not modify your infrastructure. This command adds a lock on the
+  state for the current workspace. The behavior of this lock is dependent
+  on the backend being used. Local state files cannot be locked.
+`
+	return strings.TrimSpace(helpText)
+}
+
+func (c *StateLockCommand) Synopsis() string {
+	return "Acquire a lock on the current workspace"
+}
+
+const outputStateLockSuccess = `
+LOCKID: %s
+
+[reset][bold][green]Terraform state has been successfully been locked![reset][green]
+
+The state has been locked, and Terraform commands should now be blocked from
+obtaining a new lock on the remote state.
+`