Allow flink onprem commands to work with MDS auth

internal/flink/command.go

@@ -1,10 +1,13 @@
 package flink
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/spf13/cobra"
 
+	cmfsdk "github.com/confluentinc/cmf-sdk-go/v1"
+
 	pcmd "github.com/confluentinc/cli/v4/pkg/cmd"
 	"github.com/confluentinc/cli/v4/pkg/config"
 	"github.com/confluentinc/cli/v4/pkg/featureflags"
@@ -22,8 +25,10 @@ func New(cfg *config.Config, prerunner pcmd.PreRunner) *cobra.Command {
 
 	c := &command{pcmd.NewAuthenticatedCLICommand(cmd, prerunner)}
 
-	if !cfg.IsCloudLogin() {
-		// On-prem commands don't require login, so change the pre-runner to account for that.
+	// On-prem commands are able to run with or without login. Accordingly, set the pre-runner.
+	if cfg.IsOnPremLogin() {
+		c = &command{pcmd.NewAuthenticatedWithMDSCLICommand(cmd, prerunner)}
+	} else if !cfg.IsCloudLogin() {
 		cmd.PersistentPreRunE = prerunner.Anonymous(c.AuthenticatedCLICommand.CLICommand, false)
 	}
 
@@ -104,3 +109,10 @@ func addCmfFlagSet(cmd *cobra.Command) {
 	cmd.Flags().String("client-cert-path", "", `Path to client cert to be verified by Confluent Manager for Apache Flink. Include for mTLS authentication. Environment variable "CONFLUENT_CMF_CLIENT_CERT_PATH" may be set in place of this flag.`)
 	cmd.Flags().String("certificate-authority-path", "", `Path to a PEM-encoded Certificate Authority to verify the Confluent Manager for Apache Flink connection. Environment variable "CONFLUENT_CMF_CERTIFICATE_AUTHORITY_PATH" may be set in place of this flag.`)
 }
+
+func (c *command) createContext() context.Context {
+	if !c.Config.IsOnPremLogin() {
+		return context.Background()
+	}
+	return context.WithValue(context.Background(), cmfsdk.ContextAccessToken, c.Context.GetAuthToken())
+}

internal/flink/command_application_create.go

@@ -66,7 +66,7 @@ func (c *command) applicationCreate(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	outputApplication, err := client.CreateApplication(cmd.Context(), environment, application)
+	outputApplication, err := client.CreateApplication(c.createContext(), environment, application)
 	if err != nil {
 		return err
 	}

internal/flink/command_application_delete.go

@@ -37,7 +37,7 @@ func (c *command) applicationDelete(cmd *cobra.Command, args []string) error {
 	}
 
 	existenceFunc := func(name string) bool {
-		_, err := client.DescribeApplication(cmd.Context(), environment, name)
+		_, err := client.DescribeApplication(c.createContext(), environment, name)
 		return err == nil
 	}
 
@@ -50,7 +50,7 @@ func (c *command) applicationDelete(cmd *cobra.Command, args []string) error {
 	}
 
 	deleteFunc := func(name string) error {
-		return client.DeleteApplication(cmd.Context(), environment, name)
+		return client.DeleteApplication(c.createContext(), environment, name)
 	}
 
 	_, err = deletion.Delete(args, deleteFunc, resource.FlinkApplication)

internal/flink/command_application_describe.go

@@ -36,7 +36,7 @@ func (c *command) applicationDescribe(cmd *cobra.Command, args []string) error {
 	}
 
 	applicationName := args[0]
-	application, err := client.DescribeApplication(cmd.Context(), environment, applicationName)
+	application, err := client.DescribeApplication(c.createContext(), environment, applicationName)
 	if err != nil {
 		return err
 	}

internal/flink/command_application_list.go

@@ -35,7 +35,7 @@ func (c *command) applicationList(cmd *cobra.Command, _ []string) error {
 		return err
 	}
 
-	applications, err := client.ListApplications(cmd.Context(), environment)
+	applications, err := client.ListApplications(c.createContext(), environment)
 	if err != nil {
 		return err
 	}

internal/flink/command_application_update.go

@@ -66,7 +66,7 @@ func (c *command) applicationUpdate(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	outputApplication, err := client.UpdateApplication(cmd.Context(), environment, application)
+	outputApplication, err := client.UpdateApplication(c.createContext(), environment, application)
 	if err != nil {
 		return err
 	}

internal/flink/command_application_webui.go

@@ -62,7 +62,7 @@ func (c *command) applicationWebUiForward(cmd *cobra.Command, args []string) err
 
 	// Get the name of the application and check for its existence
 	applicationName := args[0]
-	_, err = cmfClient.DescribeApplication(cmd.Context(), environment, applicationName)
+	_, err = cmfClient.DescribeApplication(c.createContext(), environment, applicationName)
 	if err != nil {
 		return fmt.Errorf(`failed to forward web UI: %s`, err)
 	}
@@ -78,7 +78,7 @@ func (c *command) applicationWebUiForward(cmd *cobra.Command, args []string) err
 	}
 
 	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		handleRequest(w, r, url, environment, applicationName, cmfClient.APIClient.GetConfig().UserAgent, client)
+		c.handleRequest(w, r, url, environment, applicationName, cmfClient.APIClient.GetConfig().UserAgent, client)
 	})
 
 	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
@@ -93,7 +93,7 @@ func (c *command) applicationWebUiForward(cmd *cobra.Command, args []string) err
 	return nil
 }
 
-func handleRequest(userResponseWriter http.ResponseWriter, userRequest *http.Request, url, environmentName, applicationName, userAgent string, client *http.Client) {
+func (c *command) handleRequest(userResponseWriter http.ResponseWriter, userRequest *http.Request, url, environmentName, applicationName, userAgent string, client *http.Client) {
 	body, err := io.ReadAll(userRequest.Body)
 	if err != nil {
 		http.Error(userResponseWriter, fmt.Sprintf("Failed to read request body: %s", err), http.StatusInternalServerError)
@@ -109,6 +109,11 @@ func handleRequest(userResponseWriter http.ResponseWriter, userRequest *http.Req
 	reqToCmf.Header = userRequest.Header
 	reqToCmf.Header.Set("x-confluent-cli-version", userAgent)
 
+	if c.Config.IsOnPremLogin() {
+		accessToken := c.Context.GetAuthToken()
+		reqToCmf.Header.Set("Authorization", fmt.Sprintf("Bearer %s", accessToken))
+	}
+
 	resFromCmf, err := client.Do(reqToCmf)
 	if err != nil {
 		http.Error(userResponseWriter, fmt.Sprintf("failed to forward the request: %s", err), http.StatusInternalServerError)

internal/flink/command_environment_create.go

@@ -86,7 +86,7 @@ func (c *command) environmentCreate(cmd *cobra.Command, args []string) error {
 	}
 	postEnvironment.KubernetesNamespace = kubernetesNamespace
 
-	outputEnvironment, err := client.CreateEnvironment(cmd.Context(), postEnvironment)
+	outputEnvironment, err := client.CreateEnvironment(c.createContext(), postEnvironment)
 	if err != nil {
 		return err
 	}

internal/flink/command_environment_delete.go

@@ -31,7 +31,7 @@ func (c *command) environmentDelete(cmd *cobra.Command, args []string) error {
 	}
 
 	existenceFunc := func(name string) bool {
-		_, err := client.DescribeEnvironment(cmd.Context(), name)
+		_, err := client.DescribeEnvironment(c.createContext(), name)
 		return err == nil
 	}
 
@@ -44,7 +44,7 @@ func (c *command) environmentDelete(cmd *cobra.Command, args []string) error {
 	}
 
 	deleteFunc := func(name string) error {
-		return client.DeleteEnvironment(cmd.Context(), name)
+		return client.DeleteEnvironment(c.createContext(), name)
 	}
 
 	_, err = deletion.Delete(args, deleteFunc, resource.FlinkEnvironment)

internal/flink/command_environment_describe.go

@@ -33,7 +33,7 @@ func (c *command) environmentDescribe(cmd *cobra.Command, args []string) error {
 
 	// Get the name of the environment to be retrieved
 	environmentName := args[0]
-	environment, err := client.DescribeEnvironment(cmd.Context(), environmentName)
+	environment, err := client.DescribeEnvironment(c.createContext(), environmentName)
 	if err != nil {
 		return err
 	}

internal/flink/command_environment_list.go

@@ -28,7 +28,7 @@ func (c *command) environmentList(cmd *cobra.Command, _ []string) error {
 		return err
 	}
 
-	environments, err := client.ListEnvironments(cmd.Context())
+	environments, err := client.ListEnvironments(c.createContext())
 	if err != nil {
 		return err
 	}

internal/flink/command_environment_update.go

@@ -75,7 +75,7 @@ func (c *command) environmentUpdate(cmd *cobra.Command, args []string) error {
 		environment.FlinkApplicationDefaults = defaultsParsed
 	}
 
-	outputEnvironment, err := client.UpdateEnvironment(cmd.Context(), environment)
+	outputEnvironment, err := client.UpdateEnvironment(c.createContext(), environment)
 	if err != nil {
 		return err
 	}

test/fixtures/output/flink/environment/list-cloud.golden

@@ -0,0 +1,5 @@
+Error: you must log out of Confluent Cloud to use this command
+
+Suggestions:
+    Log out with `confluent logout`.
+

test/flink_onprem_test.go

@@ -6,6 +6,20 @@ import (
 	"github.com/confluentinc/cli/v4/pkg/auth"
 )
 
+// Runs the integration test with login = "" and login = "onprem"
+func runIntegrationTestsWithMultipleAuth(s *CLITestSuite, tests []CLITest) {
+	for _, test := range tests {
+		test.login = ""
+		s.T().Setenv("LOGIN_TYPE", "")
+		s.runIntegrationTest(test)
+
+		test.name = test.args + "-onprem"
+		test.login = "onprem"
+		s.T().Setenv("LOGIN_TYPE", "onprem")
+		s.runIntegrationTest(test)
+	}
+}
+
 func (s *CLITestSuite) TestFlinkApplicationList() {
 	tests := []CLITest{
 		// failure scenarios
@@ -17,9 +31,7 @@ func (s *CLITestSuite) TestFlinkApplicationList() {
 		{args: "flink application list --environment default  --output human", fixture: "flink/application/list-human.golden"},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkApplicationDelete() {
@@ -35,9 +47,7 @@ func (s *CLITestSuite) TestFlinkApplicationDelete() {
 		{args: "flink application delete --environment delete-test default-application-1 non-existent --force", fixture: "flink/application/delete-mixed.golden", exitCode: 1},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkEnvironmentList() {
@@ -47,9 +57,7 @@ func (s *CLITestSuite) TestFlinkEnvironmentList() {
 		{args: "flink environment list --output human", fixture: "flink/environment/list-human.golden"},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkEnvironmentDelete() {
@@ -63,9 +71,7 @@ func (s *CLITestSuite) TestFlinkEnvironmentDelete() {
 		{args: "flink environment delete default non-existent --force", fixture: "flink/environment/delete-mixed.golden", exitCode: 1},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkApplicationCreate() {
@@ -81,9 +87,7 @@ func (s *CLITestSuite) TestFlinkApplicationCreate() {
 		{args: "flink application create --environment default test/fixtures/input/flink/application/create-new.json --output human", fixture: "flink/application/create-with-human.golden"},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkApplicationUpdate() {
@@ -100,9 +104,7 @@ func (s *CLITestSuite) TestFlinkApplicationUpdate() {
 		{args: "flink application update --environment default test/fixtures/input/flink/application/update-successful.json --output human", fixture: "flink/application/update-with-human.golden"},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkEnvironmentCreate() {
@@ -118,9 +120,7 @@ func (s *CLITestSuite) TestFlinkEnvironmentCreate() {
 		{args: "flink environment create default", fixture: "flink/environment/create-no-namespace.golden", exitCode: 1},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkEnvironmentUpdate() {
@@ -135,9 +135,7 @@ func (s *CLITestSuite) TestFlinkEnvironmentUpdate() {
 		{args: "flink environment update get-failure", fixture: "flink/environment/update-get-failure.golden", exitCode: 1},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkEnvironmentDescribe() {
@@ -151,9 +149,7 @@ func (s *CLITestSuite) TestFlinkEnvironmentDescribe() {
 		{args: "flink environment describe", fixture: "flink/environment/describe-no-environment.golden", exitCode: 1},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkApplicationDescribe() {
@@ -169,9 +165,7 @@ func (s *CLITestSuite) TestFlinkApplicationDescribe() {
 		{args: "flink application describe default-application-1", fixture: "flink/application/describe-no-environment.golden", exitCode: 1},
 	}
 
-	for _, test := range tests {
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
 }
 
 func (s *CLITestSuite) TestFlinkApplicationWebUiForward() {
@@ -182,14 +176,17 @@ func (s *CLITestSuite) TestFlinkApplicationWebUiForward() {
 		{args: "flink --url dummy-url application web-ui-forward non-existent --environment default", fixture: "flink/application/forward-nonexistent-application.golden", exitCode: 1},
 		{args: "flink --url dummy-url application web-ui-forward default-application-1 --environment non-existent", fixture: "flink/application/forward-nonexistent-environment.golden", exitCode: 1},
 		{args: "flink --url dummy-url application web-ui-forward get-failure --environment default", fixture: "flink/application/forward-get-failure.golden", exitCode: 1},
-		{name: "no-url-set", args: "flink application web-ui-forward --environment does-not-matter missing-applications", fixture: "flink/application/url-missing.golden", exitCode: 1},
 	}
 
-	for _, test := range tests {
-		if test.name == "no-url-set" {
-			// unset the environment variable
-			os.Unsetenv(auth.ConfluentPlatformCmfURL)
-		}
-		s.runIntegrationTest(test)
-	}
+	runIntegrationTestsWithMultipleAuth(s, tests)
+
+	noUrlSetTest := CLITest{name: "no-url-set", args: "flink application web-ui-forward --environment does-not-matter missing-applications", fixture: "flink/application/url-missing.golden", exitCode: 1}
+	// unset the environment variable
+	os.Unsetenv(auth.ConfluentPlatformCmfURL)
+	s.runIntegrationTest(noUrlSetTest)
+}
+
+func (s *CLITestSuite) TestFlinkOnPremWithCloudLogin() {
+	test := CLITest{args: "flink environment list --output json", fixture: "flink/environment/list-cloud.golden", login: "cloud", exitCode: 1}
+	s.runIntegrationTest(test)
 }