allow token to come from env

replicate · Jan 30, 2023 · dd516e2 · dd516e2
1 parent c1a1bcb
commit dd516e2
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/pkg/cli/affix.go b/pkg/cli/affix.go
@@ -28,7 +28,6 @@ func newAffixCommand() *cobra.Command {
 	}
 
 	cmd.Flags().StringVarP(&sToken, "token", "t", "", "replicate cog token")
-	cmd.MarkFlagRequired("token")
 	cmd.Flags().StringVarP(&sRegistry, "registry", "r", "r8.im", "registry host")
 	cmd.Flags().StringVarP(&baseRef, "base", "b", "", "base image reference - include tag: r8.im/username/modelname@sha256:hexdigest")
 	cmd.MarkFlagRequired("base")
@@ -42,6 +41,9 @@ func newAffixCommand() *cobra.Command {
 }
 
 func affixCommmand(cmd *cobra.Command, args []string) error {
+	if sToken == "" {
+		sToken = os.Getenv("COG_TOKEN")
+	}
 
 	u, err := auth.VerifyCogToken(sRegistry, sToken)
 	if err != nil {

diff --git a/pkg/cli/layers.go b/pkg/cli/layers.go
@@ -11,10 +11,6 @@ import (
 	"github.com/replicate/r8/pkg/images"
 )
 
-// var (
-// 	sToken string
-// )
-
 func newLayerCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:    "layers [image]",
@@ -26,12 +22,14 @@ func newLayerCommand() *cobra.Command {
 	}
 
 	cmd.Flags().StringVarP(&sToken, "token", "t", "", "replicate cog token")
-	cmd.MarkFlagRequired("token")
 
 	return cmd
 }
 
 func layersCommmand(cmd *cobra.Command, args []string) error {
+	if sToken == "" {
+		sToken = os.Getenv("COG_TOKEN")
+	}
 
 	if len(args) == 0 {
 		return nil

diff --git a/pkg/cli/remix.go b/pkg/cli/remix.go
@@ -25,12 +25,14 @@ func newRemixCommand() *cobra.Command {
 	}
 
 	cmd.Flags().StringVarP(&sToken, "token", "t", "", "replicate cog token")
-	cmd.MarkFlagRequired("token")
 
 	return cmd
 }
 
 func remixCommmand(cmd *cobra.Command, args []string) error {
+	if sToken == "" {
+		sToken = os.Getenv("COG_TOKEN")
+	}
 
 	u, err := auth.VerifyCogToken(sRegistry, sToken)
 	if err != nil {