Allow specifying a target container for pilot VMs (#4218)

* Allow specifying a target container for pilot VMs * Error handling * format
superfly · Feb 13, 2025 · 49f10cd · 49f10cd
1 parent ef6c170
commit 49f10cd
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 8 deletions.
diff --git a/internal/command/console/console.go b/internal/command/console/console.go
@@ -56,6 +56,10 @@ func New() *cobra.Command {
 			Description: "Select the machine on which to execute the console from a list.",
 			Default:     false,
 		},
+		flag.String{
+			Name:        "container",
+			Description: "Container to connect to",
+		},
 		flag.String{
 			Name:        "user",
 			Shorthand:   "u",
@@ -221,6 +225,7 @@ func runConsole(ctx context.Context) error {
 		Dialer:         dialer,
 		Username:       flag.GetString(ctx, "user"),
 		DisableSpinner: false,
+		Container:      flag.GetString(ctx, "container"),
 		AppNames:       []string{app.Name},
 	}
 	sshClient, err := ssh.Connect(params, machine.PrivateIP)
@@ -234,7 +239,7 @@ func runConsole(ctx context.Context) error {
 		consoleCommand = flag.GetString(ctx, "command")
 	}
 
-	return ssh.Console(ctx, sshClient, consoleCommand, true)
+	return ssh.Console(ctx, sshClient, consoleCommand, true, params.Container)
 }
 
 func selectMachine(ctx context.Context, app *fly.AppCompact, appConfig *appconfig.Config) (*fly.Machine, func(), error) {

diff --git a/internal/command/machine/run.go b/internal/command/machine/run.go
@@ -479,7 +479,7 @@ func runMachineRun(ctx context.Context) error {
 			return err
 		}
 
-		err = ssh.Console(ctx, sshClient, flag.GetString(ctx, "command"), true)
+		err = ssh.Console(ctx, sshClient, flag.GetString(ctx, "command"), true, "")
 		if destroy {
 			err = soManyErrors("console", err, "destroy machine", Destroy(ctx, app, machine, true))
 		}

diff --git a/internal/command/postgres/barman.go b/internal/command/postgres/barman.go
@@ -492,7 +492,7 @@ func runConsole(ctx context.Context, cmd string) error {
 		return err
 	}
 
-	if err := ssh.Console(ctx, sshc, cmd, false); err != nil {
+	if err := ssh.Console(ctx, sshc, cmd, false, ""); err != nil {
 		captureError(ctx, err, app)
 		return err
 	}

diff --git a/internal/command/ssh/connect.go b/internal/command/ssh/connect.go
@@ -56,6 +56,7 @@ type ConnectParams struct {
 	Username       string
 	Dialer         agent.Dialer
 	DisableSpinner bool
+	Container      string
 	AppNames       []string
 }
 

diff --git a/internal/command/ssh/console.go b/internal/command/ssh/console.go
@@ -61,6 +61,10 @@ func stdArgsSSH(cmd *cobra.Command) {
 			Shorthand:   "A",
 			Description: "Address of VM to connect to",
 		},
+		flag.String{
+			Name:        "container",
+			Description: "Container to connect to",
+		},
 		flag.Bool{
 			Name:        "pty",
 			Description: "Allocate a pseudo-terminal (default: on when no command is provided)",
@@ -177,6 +181,7 @@ func runConsole(ctx context.Context) error {
 		Dialer:         dialer,
 		Username:       flag.GetString(ctx, "user"),
 		DisableSpinner: quiet(ctx),
+		Container:      flag.GetString(ctx, "container"),
 		AppNames:       []string{app.Name},
 	}
 	sshc, err := Connect(params, addr)
@@ -185,15 +190,15 @@ func runConsole(ctx context.Context) error {
 		return err
 	}
 
-	if err := Console(ctx, sshc, cmd, allocPTY); err != nil {
+	if err := Console(ctx, sshc, cmd, allocPTY, params.Container); err != nil {
 		captureError(ctx, err, app)
 		return err
 	}
 
 	return nil
 }
 
-func Console(ctx context.Context, sshClient *ssh.Client, cmd string, allocPTY bool) error {
+func Console(ctx context.Context, sshClient *ssh.Client, cmd string, allocPTY bool, container string) error {
 	currentStdin, currentStdout, currentStderr, err := setupConsole()
 	defer func() error {
 		if err := cleanupConsole(currentStdin, currentStdout, currentStderr); err != nil {
@@ -214,7 +219,7 @@ func Console(ctx context.Context, sshClient *ssh.Client, cmd string, allocPTY bo
 		TermEnv:  determineTermEnv(),
 	}
 
-	if err := sshClient.Shell(ctx, sessIO, cmd); err != nil {
+	if err := sshClient.Shell(ctx, sessIO, cmd, container); err != nil {
 		return errors.Wrap(err, "ssh shell")
 	}
 

diff --git a/internal/command/ssh/ssh_terminal.go b/internal/command/ssh/ssh_terminal.go
@@ -115,7 +115,7 @@ func SSHConnect(p *SSHParams, addr string) error {
 		TermEnv:  "xterm",
 	}
 
-	if err := sshClient.Shell(context.Background(), sessIO, p.Cmd); err != nil {
+	if err := sshClient.Shell(context.Background(), sessIO, p.Cmd, ""); err != nil {
 		return errors.Wrap(err, "ssh shell")
 	}
 

diff --git a/ssh/client.go b/ssh/client.go
@@ -104,17 +104,26 @@ func (c *Client) Connect(ctx context.Context) error {
 	}
 }
 
-func (c *Client) Shell(ctx context.Context, sessIO *SessionIO, cmd string) error {
+func (c *Client) Shell(ctx context.Context, sessIO *SessionIO, cmd string, container string) error {
 	if c.Client == nil {
 		if err := c.Connect(ctx); err != nil {
 			return err
 		}
 	}
 
 	sess, err := c.Client.NewSession()
+
 	if err != nil {
 		return err
 	}
+
+	if container != "" {
+		err = sess.Setenv("FLY_SSH_CONTAINER", container)
+		if err != nil {
+			return err
+		}
+	}
+
 	defer sess.Close()
 
 	return sessIO.attach(ctx, sess, cmd)