[codex] consolidate terraform bulk execution on scheduler

cmd/ansible/ansible.go

@@ -1,7 +1,11 @@
 package ansible
 
 import (
+	"os"
+	"strings"
+
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 
 	"github.com/cloudposse/atmos/cmd/internal"
@@ -48,6 +52,17 @@ func init() {
 		panic(err)
 	}
 
+	// Normalize only the long identity flag for Ansible. The -i shorthand belongs
+	// to Ansible inventory, so Atmos identity must remain long-form here.
+	identityRegistry := flags.NewFlagRegistry()
+	identityRegistry.Register(&flags.StringFlag{
+		Name:        cfg.IdentityFlagName,
+		Default:     "",
+		Description: "Specify the identity to authenticate before running Ansible commands.",
+		NoOptDefVal: cfg.IdentityFlagSelectValue,
+	})
+	internal.RegisterCommandFlagRegistry("ansible", identityRegistry)
+
 	// Add subcommands.
 	ansibleCmd.AddCommand(playbookCmd)
 	ansibleCmd.AddCommand(versionCmd)
@@ -108,6 +123,75 @@ func ansibleGlobalFlagsHandler(cmd *cobra.Command, args []string) error {
 	return cmd.Usage()
 }
 
+func resolveAnsibleIdentity(cmd *cobra.Command, parsedIdentity string) string {
+	if parsedIdentity != "" {
+		return cfg.NormalizeIdentityValue(parsedIdentity)
+	}
+
+	if value, ok := getChangedIdentityFlagValue(cmd); ok {
+		return cfg.NormalizeIdentityValue(value)
+	}
+
+	if value, ok := getLongIdentityFromArgs(os.Args[1:]); ok {
+		return cfg.NormalizeIdentityValue(value)
+	}
+
+	if envIdentity := os.Getenv("ATMOS_IDENTITY"); envIdentity != "" {
+		return cfg.NormalizeIdentityValue(envIdentity)
+	}
+
+	return ""
+}
+
+func getChangedIdentityFlagValue(cmd *cobra.Command) (string, bool) {
+	for current := cmd; current != nil; current = current.Parent() {
+		for _, flagSet := range []*pflag.FlagSet{
+			current.Flags(),
+			current.InheritedFlags(),
+			current.PersistentFlags(),
+		} {
+			if flagSet == nil {
+				continue
+			}
+
+			flag := flagSet.Lookup(cfg.IdentityFlagName)
+			if flag == nil || !flag.Changed {
+				continue
+			}
+
+			return flag.Value.String(), true
+		}
+	}
+
+	return "", false
+}
+
+func getLongIdentityFromArgs(args []string) (string, bool) {
+	for i := 0; i < len(args); i++ {
+		arg := args[i]
+		if arg == "--" {
+			return "", false
+		}
+
+		if arg == cfg.IdentityFlag {
+			if i+1 < len(args) && !strings.HasPrefix(args[i+1], "-") {
+				return args[i+1], true
+			}
+			return cfg.IdentityFlagSelectValue, true
+		}
+
+		if strings.HasPrefix(arg, cfg.IdentityFlag+"=") {
+			value := strings.TrimPrefix(arg, cfg.IdentityFlag+"=")
+			if value == "" {
+				return cfg.IdentityFlagSelectValue, true
+			}
+			return value, true
+		}
+	}
+
+	return "", false
+}
+
 // buildConfigAndStacksInfo creates a ConfigAndStacksInfo with global flags populated.
 // This ensures config selection flags (--base-path, --config, --config-path, --profile)
 // are properly honored when initializing CLI config.
@@ -119,6 +203,7 @@ func buildConfigAndStacksInfo(cmd *cobra.Command) schema.ConfigAndStacksInfo {
 		AtmosBasePath:           globalFlags.BasePath,
 		AtmosConfigFilesFromArg: globalFlags.Config,
 		AtmosConfigDirsFromArg:  globalFlags.ConfigPath,
+		Identity:                resolveAnsibleIdentity(cmd, globalFlags.Identity.Value()),
 		ProfilesFromArg:         globalFlags.Profile,
 	}
 

cmd/ansible/ansible_test.go

@@ -4,12 +4,62 @@ import (
 	"testing"
 
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/cloudposse/atmos/cmd/internal"
+	cfg "github.com/cloudposse/atmos/pkg/config"
+	"github.com/cloudposse/atmos/pkg/flags/preprocess"
 	"github.com/cloudposse/atmos/pkg/schema"
 )
 
+func initAnsibleCommandTest(t *testing.T) {
+	t.Helper()
+
+	// cmd.NewTestKit cannot be imported here: cmd/root.go imports cmd/ansible,
+	// so importing cmd from this package would create a circular dependency.
+	// Snapshot the package-level Cobra commands directly, matching the cleanup
+	// behavior these tests need.
+	resetAnsibleCommandFlags(t, ansibleCmd)
+	resetAnsibleCommandFlags(t, playbookCmd)
+	resetAnsibleCommandFlags(t, versionCmd)
+}
+
+func resetAnsibleCommandFlags(t *testing.T, cmd *cobra.Command) {
+	t.Helper()
+
+	type flagSnapshot struct {
+		value   string
+		changed bool
+	}
+
+	snapshot := map[*pflag.Flag]flagSnapshot{}
+	for _, flagSet := range []*pflag.FlagSet{
+		cmd.Flags(),
+		cmd.PersistentFlags(),
+		cmd.InheritedFlags(),
+	} {
+		if flagSet == nil {
+			continue
+		}
+		flagSet.VisitAll(func(flag *pflag.Flag) {
+			snapshot[flag] = flagSnapshot{
+				value:   flag.Value.String(),
+				changed: flag.Changed,
+			}
+		})
+	}
+
+	t.Cleanup(func() {
+		for flag, snap := range snapshot {
+			_ = flag.Value.Set(snap.value)
+			flag.Changed = snap.changed
+		}
+	})
+}
+
 func TestAnsibleCommandProvider(t *testing.T) {
 	provider := &AnsibleCommandProvider{}
 
@@ -75,6 +125,85 @@ func TestAnsibleCommandStructure(t *testing.T) {
 	})
 }
 
+func TestAnsibleIdentityRegistryNormalizesLongIdentityOnly(t *testing.T) {
+	initAnsibleCommandTest(t)
+
+	registry := internal.GetCommandFlagRegistry("ansible")
+	require.NotNil(t, registry)
+	allFlags := registry.All()
+	flagInfos := make([]preprocess.FlagInfo, len(allFlags))
+	for i, f := range allFlags {
+		flagInfos[i] = f
+	}
+	preprocessor := preprocess.NewNoOptDefValPreprocessor(flagInfos)
+
+	assert.Equal(
+		t,
+		[]string{"playbook", "--identity=terraform", "pg-auto-failover"},
+		preprocessor.Preprocess([]string{"playbook", "--identity", "terraform", "pg-auto-failover"}),
+	)
+
+	assert.Equal(
+		t,
+		[]string{"playbook", "-i", "localhost,", "pg-auto-failover"},
+		preprocessor.Preprocess([]string{"playbook", "-i", "localhost,", "pg-auto-failover"}),
+	)
+}
+
+func TestGetLongIdentityFromArgs(t *testing.T) {
+	initAnsibleCommandTest(t)
+
+	tests := []struct {
+		name       string
+		args       []string
+		expected   string
+		expectedOK bool
+	}{
+		{
+			name:       "equals form",
+			args:       []string{"ansible", "playbook", "component", "--identity=terraform"},
+			expected:   "terraform",
+			expectedOK: true,
+		},
+		{
+			name:       "space form",
+			args:       []string{"ansible", "playbook", "component", "--identity", "terraform"},
+			expected:   "terraform",
+			expectedOK: true,
+		},
+		{
+			name:       "without value selects interactively",
+			args:       []string{"ansible", "playbook", "component", "--identity"},
+			expected:   cfg.IdentityFlagSelectValue,
+			expectedOK: true,
+		},
+		{
+			name:       "does not consume ansible inventory shorthand",
+			args:       []string{"ansible", "playbook", "component", "-i", "localhost,"},
+			expectedOK: false,
+		},
+		{
+			name:       "stops at separator",
+			args:       []string{"ansible", "playbook", "component", "--", "--identity=terraform"},
+			expectedOK: false,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			actual, ok := getLongIdentityFromArgs(tc.args)
+			assert.Equal(t, tc.expectedOK, ok)
+			assert.Equal(t, tc.expected, actual)
+		})
+	}
+}
+
+func TestResolveAnsibleIdentity_NormalizesParsedIdentity(t *testing.T) {
+	initAnsibleCommandTest(t)
+
+	assert.Equal(t, cfg.IdentityFlagDisabledValue, resolveAnsibleIdentity(&cobra.Command{}, "false"))
+}
+
 func TestPlaybookCommandStructure(t *testing.T) {
 	t.Run("playbook command has correct properties", func(t *testing.T) {
 		assert.Equal(t, "playbook", playbookCmd.Use)
@@ -97,6 +226,8 @@ func TestVersionCommandStructure(t *testing.T) {
 }
 
 func TestBuildConfigAndStacksInfo(t *testing.T) {
+	initAnsibleCommandTest(t)
+
 	t.Run("returns empty info when no stack flag", func(t *testing.T) {
 		cmd := &cobra.Command{Use: "test"}
 		info := buildConfigAndStacksInfo(cmd)
@@ -112,6 +243,41 @@ func TestBuildConfigAndStacksInfo(t *testing.T) {
 		info := buildConfigAndStacksInfo(cmd)
 		assert.Equal(t, "dev-us-east-1", info.Stack)
 	})
+
+	t.Run("returns info with identity when global identity flag is set", func(t *testing.T) {
+		viper.Set("identity", "terraform")
+		t.Cleanup(func() { viper.Set("identity", "") })
+
+		cmd := &cobra.Command{Use: "test"}
+		cmd.Flags().String("identity", "", "identity name")
+		err := cmd.Flags().Set("identity", "terraform")
+		require.NoError(t, err)
+
+		info := buildConfigAndStacksInfo(cmd)
+		assert.Equal(t, "terraform", info.Identity)
+	})
+
+	t.Run("returns info with identity from changed root persistent flag", func(t *testing.T) {
+		rootCmd := &cobra.Command{Use: "atmos"}
+		rootCmd.PersistentFlags().String("identity", "", "identity name")
+		cmd := &cobra.Command{Use: "playbook"}
+		rootCmd.AddCommand(cmd)
+
+		err := rootCmd.PersistentFlags().Set("identity", "terraform")
+		require.NoError(t, err)
+
+		info := buildConfigAndStacksInfo(cmd)
+		assert.Equal(t, "terraform", info.Identity)
+	})
+
+	t.Run("returns info with identity from environment", func(t *testing.T) {
+		t.Setenv("ATMOS_IDENTITY", "terraform")
+
+		cmd := &cobra.Command{Use: "test"}
+
+		info := buildConfigAndStacksInfo(cmd)
+		assert.Equal(t, "terraform", info.Identity)
+	})
 }
 
 func TestProcessArgs(t *testing.T) {

cmd/identity_helpers.go

@@ -95,8 +95,8 @@ func extractIdentityFromArgs(args []string) string {
 		}
 
 		// Handle -i=value format (short flag).
-		if strings.HasPrefix(arg, "-i=") {
-			value := strings.TrimPrefix(arg, "-i=")
+		if strings.HasPrefix(arg, cfg.IdentityFlagShort+"=") {
+			value := strings.TrimPrefix(arg, cfg.IdentityFlagShort+"=")
 			if value == "" {
 				// -i= (empty value) -> interactive selection.
 				return cfg.IdentityFlagSelectValue
@@ -105,7 +105,7 @@ func extractIdentityFromArgs(args []string) string {
 		}
 
 		// Handle --identity value format (space-separated).
-		if arg == cfg.IdentityFlag || arg == "-i" {
+		if arg == cfg.IdentityFlag || arg == cfg.IdentityFlagShort {
 			// Check if next arg exists and is not another flag.
 			if i+1 < len(args) && !strings.HasPrefix(args[i+1], "-") {
 				// Has value: --identity <value>.

cmd/terraform/flags.go

@@ -1,6 +1,7 @@
 package terraform
 
 import (
+	cfg "github.com/cloudposse/atmos/pkg/config"
 	"github.com/cloudposse/atmos/pkg/flags"
 	"github.com/cloudposse/atmos/pkg/perf"
 )
@@ -21,12 +22,12 @@ func TerraformFlags() *flags.FlagRegistry {
 // registerIdentityFlags adds identity and authentication related flags.
 func registerIdentityFlags(registry *flags.FlagRegistry) {
 	registry.Register(&flags.StringFlag{
-		Name:        "identity",
-		Shorthand:   "i",
+		Name:        cfg.IdentityFlagName,
+		Shorthand:   cfg.IdentityFlagShortName,
 		Default:     "",
 		Description: "Specify the identity to authenticate to before running Terraform commands. Use without value to interactively select.",
 		EnvVars:     []string{"ATMOS_IDENTITY"},
-		NoOptDefVal: "__SELECT__",
+		NoOptDefVal: cfg.IdentityFlagSelectValue,
 	})
 }