dynamically require org flag based on whether or not user is multi org member

Faris Huskovic · Faris Huskovic · commit 8258ed2afb4b · 2020-10-27T09:45:01.000-05:00
diff --git a/internal/cmd/ceapi.go b/internal/cmd/ceapi.go
@@ -83,8 +83,8 @@ func findEnv(ctx context.Context, client *coder.Client, envName, userEmail strin
 	)
 }
 
-func findImg(ctx context.Context, client *coder.Client, email, imgName string) (*coder.Image, error) {
-	userImgs, err := lookupUserImgs(ctx, client, email, imgName)
+func findImg(ctx context.Context, client *coder.Client, email, imgName, orgName string) (*coder.Image, error) {
+	userImgs, err := lookupUserImgs(ctx, client, email, imgName, orgName)
 	if err != nil {
 		return nil, err
 	}
@@ -111,8 +111,8 @@ func findImg(ctx context.Context, client *coder.Client, email, imgName string) (
 	return &userImgs[0], nil
 }
 
-func getImgs(ctx context.Context, client *coder.Client, userEmail string) ([]coder.Image, error) {
-	u, err := client.UserByEmail(ctx, userEmail)
+func getImgs(ctx context.Context, client *coder.Client, email, orgName string) ([]coder.Image, error) {
+	u, err := client.UserByEmail(ctx, email)
 	if err != nil {
 		return nil, err
 	}
@@ -123,32 +123,29 @@ func getImgs(ctx context.Context, client *coder.Client, userEmail string) ([]cod
 	}
 
 	orgs = lookupUserOrgs(u, orgs)
-	var importedImgs []coder.Image
 
-	// Get all of the imported images for all of the orgs the user belongs to.
 	for _, org := range orgs {
-		imgs, err := client.GetOrganizationImages(ctx, org.ID)
-		if err != nil {
-			return nil, err
+		if orgName == org.Name {
+			imgs, err := client.GetOrganizationImages(ctx, org.ID)
+			if err != nil {
+				return nil, err
+			}
+			return imgs, nil
 		}
-		importedImgs = append(importedImgs, imgs...)
 	}
-	return importedImgs, nil
+	return nil, xerrors.Errorf("org name %q not found", orgName)
 }
 
 // returns all images that contain imgName as a substring and have been imported by the user.
-func lookupUserImgs(ctx context.Context, client *coder.Client, email, imgName string) ([]coder.Image, error) {
+func lookupUserImgs(ctx context.Context, client *coder.Client, email, imgName, orgName string) ([]coder.Image, error) {
 	switch {
-	case client == nil:
-		return nil, xerrors.New("nil client")
 	case email == "":
 		return nil, xerrors.New("user email unset")
 	case imgName == "":
 		return nil, xerrors.New("image name unset")
 	}
 
-	// Get all imported images for each org the user belongs to.
-	imgs, err := getImgs(ctx, client, email)
+	imgs, err := getImgs(ctx, client, email, orgName)
 	if err != nil {
 		return nil, err
 	}
@@ -172,3 +169,26 @@ func lookupUserImgs(ctx context.Context, client *coder.Client, email, imgName st
 	}
 	return userImgs, nil
 }
+
+func isMultiOrgMember(email string) (*bool, error) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	client, err := newClient()
+	if err != nil {
+		return nil, err
+	}
+
+	u, err := client.UserByEmail(ctx, email)
+	if err != nil {
+		return nil, xerrors.New("email not found")
+	}
+
+	orgs, err := client.Organizations(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	isMulti := len(lookupUserOrgs(u, orgs)) > 1
+	return &isMulti, nil
+}
diff --git a/internal/cmd/envs.go b/internal/cmd/envs.go
@@ -174,7 +174,7 @@ coder envs create --cpu 4 --disk 100 --memory 8 --image 5f443b16-30652892427b955
 				return err
 			}
 
-			importedImg, err := findImg(cmd.Context(), client, *user, img)
+			importedImg, err := findImg(cmd.Context(), client, *user, img, org)
 			if err != nil {
 				return err
 			}
@@ -233,11 +233,22 @@ coder envs create --cpu 4 --disk 100 --memory 8 --image 5f443b16-30652892427b955
 	cmd.Flags().StringVarP(&img, "image", "i", "", "name of the image to base the environment off of.")
 	cmd.Flags().BoolVar(&follow, "follow", false, "follow buildlog after initiating rebuild")
 	_ = cmd.MarkFlagRequired("image")
+
+	multiOrgMember, err := isMultiOrgMember(*user)
+	if err != nil {
+		clog.Fatal(fmt.Sprintf("checking if user is mutli-org member: %s", err.Error()))
+	}
+
+	// if the user belongs to multiple organizations we need them to specify which one.
+	if *multiOrgMember {
+		_ = cmd.MarkFlagRequired("org")
+	}
 	return cmd
 }
 
 func editEnvCommand(user *string) *cobra.Command {
 	var (
+		org		 string
 		img      string
 		tag      string
 		cpuCores float32
@@ -301,13 +312,26 @@ coder envs edit back-end-env --disk 20`,
 			return nil
 		},
 	}
+	cmd.Flags().StringVarP(&org, "org", "o", defaultOrg, "name of the organization the environment should be created under.")
 	cmd.Flags().StringVarP(&img, "image", "i", "", "name of the image you wan't the environment to be based off of.")
 	cmd.Flags().StringVarP(&tag, "tag", "t", "latest", "image tag of the image you wan't to base the environment off of.")
 	cmd.Flags().Float32P("cpu", "c", cpuCores, "The number of cpu cores the environment should be provisioned with.")
 	cmd.Flags().Float32P("memory", "m", memGB, "The amount of RAM an environment should be provisioned with.")
 	cmd.Flags().IntP("disk", "d", diskGB, "The amount of disk storage an environment should be provisioned with.")
 	cmd.Flags().IntP("gpu", "g", gpus, "The amount of disk storage to provision the environment with.")
 	cmd.Flags().BoolVar(&follow, "follow", false, "follow buildlog after initiating rebuild")
+	// if the image is being changed we need to know whether or not we should require the org flag.
+	if img != "" {
+		multiOrgMember, err := isMultiOrgMember(*user)
+		if err != nil {
+			clog.Fatal(fmt.Sprintf("checking if user is mutli-org member: %s", err.Error()))
+		}
+	
+		// if the user belongs to multiple organizations we need them to specify which one.
+		if *multiOrgMember {
+			_ = cmd.MarkFlagRequired("org")
+		}
+	}
 	return cmd
 }
 
