Add sasl-assume-role flag to more places

Author: yolken

cmd/topicctl/subcmd/apply.go

@@ -193,6 +193,7 @@ func applyTopic(
 			applyConfig.dryRun,
 			applyConfig.shared.saslUsername,
 			applyConfig.shared.saslPassword,
+			applyConfig.shared.saslAssumeRole,
 		)
 		if err != nil {
 			return err

cmd/topicctl/subcmd/bootstrap.go

@@ -75,6 +75,7 @@ func bootstrapRun(cmd *cobra.Command, args []string) error {
 		true,
 		bootstrapConfig.shared.saslUsername,
 		bootstrapConfig.shared.saslPassword,
+		bootstrapConfig.shared.saslAssumeRole,
 	)
 	if err != nil {
 		return err

cmd/topicctl/subcmd/check.go

@@ -138,6 +138,7 @@ func checkTopicFile(
 				true,
 				checkConfig.shared.saslUsername,
 				checkConfig.shared.saslPassword,
+				checkConfig.shared.saslAssumeRole,
 			)
 			if err != nil {
 				return false, err

cmd/topicctl/subcmd/shared.go

@@ -121,6 +121,7 @@ func (s sharedOptions) getAdminClient(
 			readOnly,
 			s.saslUsername,
 			s.saslPassword,
+			s.saslAssumeRole,
 		)
 	} else if s.brokerAddr != "" {
 		tlsEnabled := (s.tlsEnabled ||
@@ -300,4 +301,10 @@ func addSharedConfigOnlyFlags(cmd *cobra.Command, options *sharedOptions) {
 		os.Getenv("TOPICCTL_SASL_USERNAME"),
 		"SASL username if using SASL; will override value set in cluster config",
 	)
+	cmd.Flags().StringVar(
+		&options.saslAssumeRole,
+		"sasl-assume-role",
+		os.Getenv("TOPICCTL_SASL_ASSUME_ROLE"),
+		"SASL assume role if using SASL AWS-MSK-IAM; will override value set in cluster config",
+	)
 }

pkg/apply/apply_test.go

@@ -879,7 +879,7 @@ func TestApplyOverrides(t *testing.T) {
 		},
 	}
 
-	adminClient, err := clusterConfig.NewAdminClient(ctx, nil, false, "", "")
+	adminClient, err := clusterConfig.NewAdminClient(ctx, nil, false, "", "", "")
 	require.NoError(t, err)
 
 	applier, err := NewTopicApplier(
@@ -922,7 +922,7 @@ func testApplier(
 		},
 	}
 
-	adminClient, err := clusterConfig.NewAdminClient(ctx, nil, false, "", "")
+	adminClient, err := clusterConfig.NewAdminClient(ctx, nil, false, "", "", "")
 	require.NoError(t, err)
 
 	applier, err := NewTopicApplier(

pkg/check/check_test.go

@@ -28,7 +28,7 @@ func TestCheck(t *testing.T) {
 		},
 	}
 
-	adminClient, err := clusterConfig.NewAdminClient(ctx, nil, false, "", "")
+	adminClient, err := clusterConfig.NewAdminClient(ctx, nil, false, "", "", "")
 	require.NoError(t, err)
 
 	topicName := util.RandomString("check-topic-", 6)

pkg/config/cluster.go

@@ -194,12 +194,14 @@ func (c ClusterConfig) NewAdminClient(
 	readOnly bool,
 	usernameOverride string,
 	passwordOverride string,
+	assumeRoleOverride string,
 ) (admin.Client, error) {
 	if len(c.Spec.ZKAddrs) == 0 {
 		log.Debug("No ZK addresses provided, using broker admin client")
 
 		var saslUsername string
 		var saslPassword string
+		var saslAssumeRole string
 		if usernameOverride != "" {
 			log.Debugf("Setting SASL username from override value")
 			saslUsername = usernameOverride
@@ -214,6 +216,13 @@ func (c ClusterConfig) NewAdminClient(
 			saslPassword = c.Spec.SASL.Password
 		}
 
+		if assumeRoleOverride != "" {
+			log.Debugf("Setting SASL assume role from override value")
+			saslAssumeRole = assumeRoleOverride
+		} else {
+			saslAssumeRole = c.Spec.SASL.AssumeRole
+		}
+
 		var saslMechanism admin.SASLMechanism
 		var err error
 
@@ -242,7 +251,7 @@ func (c ClusterConfig) NewAdminClient(
 						Mechanism:  saslMechanism,
 						Username:   saslUsername,
 						Password:   saslPassword,
-						AssumeRole: c.Spec.SASL.AssumeRole,
+						AssumeRole: saslAssumeRole,
 					},
 				},
 				ExpectedClusterID: c.Spec.ClusterID,