From bd0ad1fb85c763dbf02d23d29cadc4e20a4196d2 Mon Sep 17 00:00:00 2001
From: rektdeckard <friedtm@gmail.com>
Date: Wed, 11 Dec 2024 11:17:20 -0700
Subject: [PATCH] feat(cmd): expand project env from anywhere

---
 cmd/lk/app.go              | 37 +++++++++++++------------
 pkg/bootstrap/bootstrap.go | 56 ++++++++++++++++++++------------------
 2 files changed, 49 insertions(+), 44 deletions(-)

diff --git a/cmd/lk/app.go b/cmd/lk/app.go
index 9a30240a..2ffb427e 100644
--- a/cmd/lk/app.go
+++ b/cmd/lk/app.go
@@ -102,7 +102,7 @@ var (
 				},
 				{
 					Name:  "env",
-					Usage: "Print project environment variables expanded from a .env.example file",
+					Usage: "Expand environment variables from the current project, and if present, the .env.example file",
 					Flags: []cli.Flag{
 						&cli.BoolFlag{
 							Name:    "w",
@@ -112,22 +112,7 @@ var (
 					},
 					ArgsUsage: "[DIR] location of the project directory (default: current directory)",
 					Before:    requireProject,
-					Action: func(ctx context.Context, cmd *cli.Command) error {
-						rootDir := cmd.Args().First()
-						if rootDir == "" {
-							rootDir = "."
-						}
-
-						env, err := instantiateEnv(ctx, cmd, rootDir, nil)
-						if err != nil {
-							return err
-						}
-						if cmd.Bool("write") {
-							return bootstrap.WriteDotEnv(rootDir, env)
-						} else {
-							return bootstrap.PrintDotEnv(env)
-						}
-					},
+					Action:    manageEnv,
 				},
 			},
 		},
@@ -363,6 +348,24 @@ func cleanupTemplate(ctx context.Context, cmd *cli.Command, appName string) erro
 	return bootstrap.CleanupTemplate(appName)
 }
 
+func manageEnv(ctx context.Context, cmd *cli.Command) error {
+	rootDir := cmd.Args().First()
+	if rootDir == "" {
+		rootDir = "."
+	}
+
+	env, err := instantiateEnv(ctx, cmd, rootDir, nil)
+	if err != nil {
+		return err
+	}
+
+	if cmd.Bool("write") {
+		return bootstrap.WriteDotEnv(rootDir, env)
+	} else {
+		return bootstrap.PrintDotEnv(env)
+	}
+}
+
 func instantiateEnv(ctx context.Context, cmd *cli.Command, rootPath string, addlEnv *map[string]string) (map[string]string, error) {
 	env := map[string]string{
 		"LIVEKIT_API_KEY":    project.APIKey,
diff --git a/pkg/bootstrap/bootstrap.go b/pkg/bootstrap/bootstrap.go
index 8bcf41c8..ab4d8971 100644
--- a/pkg/bootstrap/bootstrap.go
+++ b/pkg/bootstrap/bootstrap.go
@@ -192,40 +192,42 @@ type PromptFunc func(key string, value string) (string, error)
 // prompting for others, and returning the result as a map.
 func InstantiateDotEnv(ctx context.Context, rootDir string, substitutions map[string]string, verbose bool, prompt PromptFunc) (map[string]string, error) {
 	promptedVars := map[string]string{}
-
 	envExamplePath := path.Join(rootDir, EnvExampleFile)
+
 	stat, err := os.Stat(envExamplePath)
-	if err != nil {
+	if err != nil && !errors.Is(err, fs.ErrNotExist) {
 		return nil, err
-	}
-	if stat.IsDir() {
-		return nil, errors.New("env.example file is a directory")
-	}
+	} else if stat != nil {
+		if stat.IsDir() {
+			return nil, errors.New("env.example file is a directory")
+		}
 
-	envMap, err := godotenv.Read(envExamplePath)
-	if err != nil {
-		return nil, err
-	}
+		envMap, err := godotenv.Read(envExamplePath)
+		if err != nil {
+			return nil, err
+		}
 
-	for key, oldValue := range envMap {
-		// if key is a substitution, replace it
-		if value, ok := substitutions[key]; ok {
-			envMap[key] = value
-			// if key was already promped, use that value
-		} else if alreadyPromptedValue, ok := promptedVars[key]; ok {
-			envMap[key] = alreadyPromptedValue
-		} else {
-			// prompt for value
-			newValue, err := prompt(key, oldValue)
-			if err != nil {
-				return nil, err
+		for key, oldValue := range envMap {
+			// if key is a substitution, replace it
+			if value, ok := substitutions[key]; ok {
+				envMap[key] = value
+				// if key was already promped, use that value
+			} else if alreadyPromptedValue, ok := promptedVars[key]; ok {
+				envMap[key] = alreadyPromptedValue
+			} else {
+				// prompt for value
+				newValue, err := prompt(key, oldValue)
+				if err != nil {
+					return nil, err
+				}
+				envMap[key] = newValue
+				promptedVars[key] = newValue
 			}
-			envMap[key] = newValue
-			promptedVars[key] = newValue
 		}
+		return envMap, nil
+	} else {
+		return substitutions, nil
 	}
-
-	return envMap, nil
 }
 
 func PrintDotEnv(envMap map[string]string) error {
@@ -243,7 +245,7 @@ func WriteDotEnv(rootDir string, envMap map[string]string) error {
 		return err
 	}
 	envLocalPath := path.Join(rootDir, EnvLocalFile)
-	return os.WriteFile(envLocalPath, []byte(envContents), 0700)
+	return os.WriteFile(envLocalPath, []byte(envContents+"\n"), 0700)
 }
 
 func CloneTemplate(url, dir string) (string, string, error) {