F

internal/cmd/beta/session_cmd.go

@@ -48,17 +48,17 @@ All exported variables during the session will be available to the subsequent co
 						Stderr:  cmd.ErrOrStderr(),
 					}
 
-					runmeCmd, err := cmdFactory.Build(cfg, options)
+					program, err := cmdFactory.Build(cfg, options)
 					if err != nil {
 						return err
 					}
 
-					err = runmeCmd.Start(cmd.Context())
+					err = program.Start(cmd.Context())
 					if err != nil {
 						return err
 					}
 
-					err = runmeCmd.Wait()
+					err = program.Wait()
 					if err != nil {
 						return err
 					}
@@ -68,10 +68,15 @@ All exported variables during the session will be available to the subsequent co
 						return errors.WithStack(err)
 					}
 
+					// TODO(adamb): currently, the collected env are printed out,
+					// but they could be put in a session.
+					if _, err := cmd.ErrOrStderr().Write([]byte("Collected env during the session:\n")); err != nil {
+						return errors.WithStack(err)
+					}
 					for _, env := range changed {
 						_, err := cmd.OutOrStdout().Write([]byte(env + "\n"))
 						if err != nil {
-							return err
+							return errors.WithStack(err)
 						}
 					}
 
@@ -87,6 +92,8 @@ All exported variables during the session will be available to the subsequent co
 }
 
 func sessionSetupCmd() *cobra.Command {
+	var debug bool
+
 	cmd := cobra.Command{
 		Use:    "setup",
 		Hidden: true,
@@ -106,13 +113,17 @@ func sessionSetupCmd() *cobra.Command {
 					envSetter := command.NewFileBasedEnvSetter(
 						os.Getenv(command.EnvCollectorSessionPrePathEnvName),
 						os.Getenv(command.EnvCollectorSessionPostPathEnvName),
+						debug,
 					)
-					buf := new(bytes.Buffer)
 
-					_, _ = buf.WriteString("#!/bin/sh\n")
-					_, _ = buf.WriteString("set -euxo pipefail\n")
+					buf := bytes.NewBufferString("#!/bin/sh\n")
+					if debug {
+						_, _ = buf.WriteString("set -euxo pipefail\n")
+					}
 					_ = envSetter.SetOnShell(buf)
-					_, _ = buf.WriteString("set +euxo pipefail\n")
+					if debug {
+						_, _ = buf.WriteString("set +euxo pipefail\n")
+					}
 
 					_, err := cmd.OutOrStdout().Write(buf.Bytes())
 					return errors.WithStack(err)
@@ -121,6 +132,8 @@ func sessionSetupCmd() *cobra.Command {
 		},
 	}
 
+	cmd.Flags().BoolVar(&debug, "debug", false, "Enable debug mode.")
+
 	return &cmd
 }
 

internal/command/env_collector_fifo_unix.go

@@ -125,7 +125,7 @@ func (c *envCollectorFifo) ExtraEnv() []string {
 }
 
 func (c *envCollectorFifo) SetOnShell(shell io.Writer) error {
-	return setOnShell(shell, envDumpCommand, true, c.prePath(), c.postPath())
+	return setOnShell(shell, envDumpCommand, true, false, false, c.prePath(), c.postPath())
 }
 
 func (c *envCollectorFifo) prePath() string {

internal/command/env_collector_file.go

@@ -73,7 +73,7 @@ func (c *envCollectorFile) ExtraEnv() []string {
 }
 
 func (c *envCollectorFile) SetOnShell(shell io.Writer) error {
-	return setOnShell(shell, envDumpCommand, true, c.prePath(), c.postPath())
+	return setOnShell(shell, envDumpCommand, true, false, false, c.prePath(), c.postPath())
 }
 
 func (c *envCollectorFile) prePath() string {

internal/command/env_shell.go

@@ -14,44 +14,61 @@ func createEnv(key, value string) string {
 	return key + "=" + value
 }
 
-type FileBasedEnvSetter struct {
+// ScriptEnvSetter returns a shell script that installs itself and
+// collects environment variables to provided pre- and post-paths.
+type ScriptEnvSetter struct {
+	debug       bool
 	dumpCommand string
 	prePath     string
 	postPath    string
 }
 
-func NewFileBasedEnvSetter(prePath, postPath string) *FileBasedEnvSetter {
-	return &FileBasedEnvSetter{
+func NewFileBasedEnvSetter(prePath, postPath string, debug bool) *ScriptEnvSetter {
+	return &ScriptEnvSetter{
+		debug:       debug,
 		dumpCommand: envDumpCommand,
 		prePath:     prePath,
 		postPath:    postPath,
 	}
 }
 
-func (s *FileBasedEnvSetter) SetOnShell(shell io.Writer) error {
-	return setOnShell(shell, s.dumpCommand, true, s.prePath, s.postPath)
+func (s *ScriptEnvSetter) SetOnShell(shell io.Writer) error {
+	return setOnShell(shell, s.dumpCommand, false, true, s.debug, s.prePath, s.postPath)
 }
 
 func setOnShell(
 	shell io.Writer,
 	dumpCommand string,
 	skipShellHistory bool,
+	asFile bool,
+	debug bool,
 	prePath string,
 	postPath string,
 ) error {
 	prefix := ""
 	if skipShellHistory {
-		// Prefix commands with a space to avoid polluting the shell history.
-		prefix = " "
+		prefix = " " // space avoids polluting the shell history
 	}
 
 	w := bulkWriter{Writer: shell}
 
-	// First, dump all env at the beginning, so that a diff can be calculated.
-	w.Write([]byte(prefix + dumpCommand + " > " + prePath + "\n"))
+	if asFile {
+		w.WriteString("#!/bin/sh\n")
+	}
+
+	if debug {
+		w.WriteString("set -euxo pipefail\n")
+	}
+
+	// Dump all env at the beginning, so that a diff can be calculated.
+	w.WriteString(prefix + dumpCommand + " > " + prePath + "\n")
 	// Then, set a trap on EXIT to dump all env at the end.
-	w.Write([]byte(prefix + "__cleanup() {\nrv=$?\n" + (envDumpCommand + " > " + postPath) + "\nexit $rv\n}\n"))
-	w.Write([]byte(prefix + "trap -- \"__cleanup\" EXIT\n"))
+	w.WriteString(prefix + "__cleanup() {\nrv=$?\n" + (envDumpCommand + " > " + postPath) + "\nexit $rv\n}\n")
+	w.WriteString(prefix + "trap -- \"__cleanup\" EXIT\n")
+
+	if debug {
+		w.WriteString("set +euxo pipefail\n")
+	}
 
 	_, err := w.Done()
 	return err

internal/command/env_shell_test.go

@@ -15,7 +15,7 @@ func TestSetOnShell(t *testing.T) {
 
 	buf := new(bytes.Buffer)
 
-	err := setOnShell(buf, envDumpCommand, false, "prePath", "postPath")
+	err := setOnShell(buf, envDumpCommand, false, false, false, "prePath", "postPath")
 	require.NoError(t, err)
 
 	expected := (envDumpCommand + " > prePath\n" +
@@ -33,7 +33,7 @@ func TestSetOnShell_SkipShellHistory(t *testing.T) {
 
 	buf := new(bytes.Buffer)
 
-	err := setOnShell(buf, envDumpCommand, true, "prePath", "postPath")
+	err := setOnShell(buf, envDumpCommand, true, false, false, "prePath", "postPath")
 	require.NoError(t, err)
 
 	expected := (" " + envDumpCommand + " > prePath\n" +