feat(cli): handle org selection client-side (#1669)

Signed-off-by: Miguel Martinez <[email protected]>

Author: migmartri

app/cli/cmd/config.go

@@ -26,7 +26,7 @@ import (
 
 // Map of all the possible configuration options that we expect viper to handle
 var confOptions = struct {
-	authToken, controlplaneAPI, CASAPI, controlplaneCA, CASCA, insecure *confOpt
+	authToken, controlplaneAPI, CASAPI, controlplaneCA, CASCA, insecure, organization *confOpt
 }{
 	insecure: &confOpt{
 		viperKey: "api-insecure",
@@ -50,6 +50,10 @@ var confOptions = struct {
 		viperKey: "artifact-cas.api-ca",
 		flagName: "artifact-cas-ca",
 	},
+	organization: &confOpt{
+		viperKey: "organization",
+		flagName: "org",
+	},
 }
 
 type confOpt struct {

app/cli/cmd/organization_create.go

@@ -17,6 +17,7 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/chainloop-dev/chainloop/app/cli/internal/action"
 	"github.com/spf13/cobra"
@@ -34,6 +35,11 @@ func newOrganizationCreateCmd() *cobra.Command {
 				return err
 			}
 
+			// set the local state
+			if err := setLocalOrganization(name); err != nil {
+				return fmt.Errorf("writing config file: %w", err)
+			}
+
 			logger.Info().Msgf("Organization %q created!", org.Name)
 			return nil
 		},

app/cli/cmd/organization_leave.go

@@ -66,6 +66,11 @@ func newOrganizationLeaveCmd() *cobra.Command {
 				return fmt.Errorf("deleting membership: %w", err)
 			}
 
+			// set the local state
+			if err := setLocalOrganization(""); err != nil {
+				return fmt.Errorf("writing config file: %w", err)
+			}
+
 			logger.Info().Msg("Membership deleted")
 			return nil
 		},

app/cli/cmd/organization_list.go

@@ -22,6 +22,7 @@ import (
 	"github.com/chainloop-dev/chainloop/app/cli/internal/action"
 	"github.com/jedib0t/go-pretty/v6/table"
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
 )
 
 const UserWithNoOrganizationMsg = "you are not part of any organization, please run \"chainloop organization create --name ORG_NAME\" to create one"
@@ -50,11 +51,15 @@ func orgMembershipTableOutput(items []*action.MembershipItem) error {
 		return nil
 	}
 
+	// Get the current organization from viper configuration
+	currentOrg := viper.GetString(confOptions.organization.viperKey)
+
 	t := newTableWriter()
-	t.AppendHeader(table.Row{"Name", "Current", "Role", "Joined At"})
+	t.AppendHeader(table.Row{"Name", "Current", "Default", "Role", "Joined At"})
 
 	for _, i := range items {
-		t.AppendRow(table.Row{i.Org.Name, i.Current, i.Role, i.CreatedAt.Format(time.RFC822)})
+		current := i.Org.Name == currentOrg
+		t.AppendRow(table.Row{i.Org.Name, current, i.Default, i.Role, i.CreatedAt.Format(time.RFC822)})
 		t.AppendSeparator()
 	}
 

app/cli/cmd/organization_set.go

@@ -24,10 +24,18 @@ import (
 
 func newOrganizationSet() *cobra.Command {
 	var orgName string
+	var setDefault bool
 
 	cmd := &cobra.Command{
 		Use:   "set",
-		Short: "Set the current organization associated with this user",
+		Short: "Set the current organization to be used by this CLI",
+		Example: `
+  # Set the current organization to be used by this CLI
+  $ chainloop org set --name my-org
+
+  # Optionally set the organization as the default one for all clients by storing the preference server-side
+  $ chainloop org set --name my-org --default
+		`,
 		RunE: func(cmd *cobra.Command, _ []string) error {
 			ctx := cmd.Context()
 			// To find the membership ID, we need to iterate and filter by org
@@ -38,17 +46,27 @@ func newOrganizationSet() *cobra.Command {
 				return fmt.Errorf("organization %s not found", orgName)
 			}
 
-			m, err := action.NewMembershipSet(actionOpts).Run(ctx, membership.ID)
-			if err != nil {
-				return err
+			// set the local state
+			if err := setLocalOrganization(orgName); err != nil {
+				return fmt.Errorf("writing config file: %w", err)
+			}
+
+			// change the state server side
+			if setDefault {
+				var err error
+				membership, err = action.NewMembershipSet(actionOpts).Run(ctx, membership.ID)
+				if err != nil {
+					return err
+				}
 			}
 
 			logger.Info().Msg("Organization switched!")
-			return encodeOutput([]*action.MembershipItem{m}, orgMembershipTableOutput)
+			return encodeOutput([]*action.MembershipItem{membership}, orgMembershipTableOutput)
 		},
 	}
 
 	cmd.Flags().StringVar(&orgName, "name", "", "organization name to make the switch")
+	cmd.Flags().BoolVar(&setDefault, "default", false, "set this organization as the default one for all clients")
 	cobra.CheckErr(cmd.MarkFlagRequired("name"))
 
 	return cmd

app/cli/cmd/root.go

@@ -30,6 +30,7 @@ import (
 	"github.com/chainloop-dev/chainloop/app/cli/internal/action"
 	"github.com/chainloop-dev/chainloop/app/cli/internal/telemetry"
 	"github.com/chainloop-dev/chainloop/app/cli/internal/telemetry/posthog"
+	v1 "github.com/chainloop-dev/chainloop/app/controlplane/api/controlplane/v1"
 	"github.com/chainloop-dev/chainloop/internal/grpcconn"
 	"github.com/golang-jwt/jwt/v4"
 	"github.com/rs/zerolog"
@@ -72,13 +73,29 @@ type parsedToken struct {
 // Environment variable prefix for vipers
 const envPrefix = "CHAINLOOP"
 
+func Execute(l zerolog.Logger) error {
+	rootCmd := NewRootCmd(l)
+	if err := rootCmd.Execute(); err != nil {
+		// The local file is pointing to the wrong organization, we remove it
+		if v1.IsUserNotMemberOfOrgErrorNotInOrg(err) {
+			if err := setLocalOrganization(""); err != nil {
+				logger.Debug().Err(err).Msg("failed to remove organization from config")
+			}
+		}
+
+		return err
+	}
+
+	return nil
+}
+
 func NewRootCmd(l zerolog.Logger) *cobra.Command {
 	rootCmd := &cobra.Command{
 		Use:           appName,
 		Short:         "Chainloop Command Line Interface",
 		SilenceErrors: true,
 		SilenceUsage:  true,
-		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+		PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {
 			var err error
 			logger, err = initLogger(l)
 			if err != nil {
@@ -105,11 +122,31 @@ func NewRootCmd(l zerolog.Logger) *cobra.Command {
 			}
 
 			controlplaneURL := viper.GetString(confOptions.controlplaneAPI.viperKey)
+
+			// If no organization is set in local configuration, we load it from server and save it
+			if viper.GetString(confOptions.organization.viperKey) == "" {
+				conn, err := grpcconn.New(controlplaneURL, apiToken, opts...)
+				if err != nil {
+					return err
+				}
+
+				currentContext, err := action.NewConfigCurrentContext(newActionOpts(logger, conn)).Run()
+				if err == nil && currentContext.CurrentMembership != nil {
+					if err := setLocalOrganization(currentContext.CurrentMembership.Org.Name); err != nil {
+						return fmt.Errorf("writing config file: %w", err)
+					}
+				}
+			}
+
+			// reload the connection now that we have the org name
+			if orgName := viper.GetString(confOptions.organization.viperKey); orgName != "" {
+				opts = append(opts, grpcconn.WithOrgName(orgName))
+			}
+
 			conn, err := grpcconn.New(controlplaneURL, apiToken, opts...)
 			if err != nil {
 				return err
 			}
-
 			actionOpts = newActionOpts(logger, conn)
 
 			if !isTelemetryDisabled() {
@@ -152,7 +189,7 @@ func NewRootCmd(l zerolog.Logger) *cobra.Command {
 
 			return nil
 		},
-		PersistentPostRunE: func(cmd *cobra.Command, args []string) error {
+		PersistentPostRunE: func(_ *cobra.Command, _ []string) error {
 			return cleanup(actionOpts.CPConnection)
 		},
 	}
@@ -189,6 +226,9 @@ func NewRootCmd(l zerolog.Logger) *cobra.Command {
 	// Override the oauth authentication requirement for the CLI by providing an API token
 	rootCmd.PersistentFlags().StringVarP(&apiToken, "token", "t", "", fmt.Sprintf("API token. NOTE: Alternatively use the env variable %s", tokenEnvVarName))
 
+	rootCmd.PersistentFlags().StringP(confOptions.organization.flagName, "n", "", "organization name")
+	cobra.CheckErr(viper.BindPFlag(confOptions.organization.viperKey, rootCmd.PersistentFlags().Lookup(confOptions.organization.flagName)))
+
 	rootCmd.AddCommand(newWorkflowCmd(), newAuthCmd(), NewVersionCmd(),
 		newAttestationCmd(), newArtifactCmd(), newConfigCmd(),
 		newIntegrationCmd(), newOrganizationCmd(), newCASBackendCmd(),
@@ -453,3 +493,9 @@ func hashControlPlaneURL() string {
 func apiInsecure() bool {
 	return viper.GetBool(confOptions.insecure.viperKey)
 }
+
+// setLocalOrganization updates the local organization configuration
+func setLocalOrganization(orgName string) error {
+	viper.Set(confOptions.organization.viperKey, orgName)
+	return viper.WriteConfig()
+}

app/cli/internal/action/membership_list.go

@@ -34,7 +34,7 @@ type OrgItem struct {
 
 type MembershipItem struct {
 	ID        string     `json:"id"`
-	Current   bool       `json:"current"`
+	Default   bool       `json:"current"`
 	CreatedAt *time.Time `json:"joinedAt"`
 	UpdatedAt *time.Time `json:"updatedAt"`
 	Org       *OrgItem
@@ -96,7 +96,7 @@ func pbMembershipItemToAction(in *pb.OrgMembershipItem) *MembershipItem {
 		CreatedAt: toTimePtr(in.GetCreatedAt().AsTime()),
 		UpdatedAt: toTimePtr(in.GetCreatedAt().AsTime()),
 		Org:       pbOrgItemToAction(in.Org),
-		Current:   in.Current,
+		Default:   in.Current,
 		Role:      pbRoleToString(in.Role),
 		User:      pbUserItemToAction(in.User),
 	}

app/cli/main.go

@@ -34,10 +34,9 @@ import (
 func main() {
 	// Couldn't find an easier way to disable the timestamp
 	logger := zerolog.New(zerolog.ConsoleWriter{Out: os.Stderr, FormatTimestamp: func(interface{}) string { return "" }})
-	rootCmd := cmd.NewRootCmd(logger)
 
 	// Run the command
-	if err := rootCmd.Execute(); err != nil {
+	if err := cmd.Execute(logger); err != nil {
 		msg, exitCode := errorInfo(err, logger)
 		logger.Error().Msg(msg)
 		os.Exit(exitCode)
@@ -90,8 +89,8 @@ func errorInfo(err error, logger zerolog.Logger) (string, int) {
 		msg = "your authentication token has expired, please run chainloop auth login again"
 	case isWrappedErr(st, jwtMiddleware.ErrMissingJwtToken):
 		msg = "authentication required, please run \"chainloop auth login\""
-	case v1.IsUserWithNoMembershipErrorNotInOrg(err):
-		msg = cmd.UserWithNoOrganizationMsg
+	case v1.IsUserNotMemberOfOrgErrorNotInOrg(err), v1.IsUserWithNoMembershipErrorNotInOrg(err):
+		msg = "the organization you are trying to access does not exist, please run \"chainloop auth login\""
 	case errors.As(err, &cmd.GracefulError{}):
 		// Graceful recovery if the flag is set and the received error is marked as recoverable
 		if cmd.GracefulExit {