STACKITGit/cli
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add command to open codespaces in JupyterLab

Browse source 
Add command to open codespaces in JupyterLab
This commit is contained in:
Caleb Brose 2022-04-28 16:11:14 -05:00 • committed by GitHub
commit 55bce59ab7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 150 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

33
pkg/cmd/codespace/common.go
View file

@ -14,8 +14,10 @@ import (
"github.com/AlecAivazis/survey/v2"
"github.com/AlecAivazis/survey/v2/terminal"
"github.com/cli/cli/v2/internal/codespaces"
"github.com/cli/cli/v2/internal/codespaces/api"
"github.com/cli/cli/v2/pkg/iostreams"
"github.com/cli/cli/v2/pkg/liveshare"
"github.com/spf13/cobra"
"golang.org/x/term"
)
@ -58,6 +60,37 @@ func (a *App) StopProgressIndicator() {
a.io.StopProgressIndicator()
}
// Connects to a codespace using Live Share and returns that session
func startLiveShareSession(ctx context.Context, codespace *api.Codespace, a *App, debug bool, debugFile string) (session *liveshare.Session, err error) {
// While connecting, ensure in the background that the user has keys installed.
// That lets us report a more useful error message if they don't.
authkeys := make(chan error, 1)
go func() {
authkeys <- checkAuthorizedKeys(ctx, a.apiClient)
}()
liveshareLogger := noopLogger()
if debug {
debugLogger, err := newFileLogger(debugFile)
if err != nil {
return nil, fmt.Errorf("couldn't create file logger: %w", err)
}
defer safeClose(debugLogger, &err)
liveshareLogger = debugLogger.Logger
a.errLogger.Printf("Debug file located at: %s", debugLogger.Name())
}
session, err = codespaces.ConnectToLiveshare(ctx, a, liveshareLogger, a.apiClient, codespace)
if err != nil {
if authErr := <-authkeys; authErr != nil {
return nil, fmt.Errorf("failed to fetch authorization keys: %w", authErr)
}
return nil, fmt.Errorf("failed to connect to Live Share: %w", err)
}
return
}
//go:generate moq -fmt goimports -rm -skip-ensure -out mock_api.go . apiClient
type apiClient interface {
GetUser(ctx context.Context) (*api.User, error)

82
pkg/cmd/codespace/jupyter.go Normal file
View file

@ -0,0 +1,82 @@
package codespace
import (
"context"
"fmt"
"net"
"strings"
"github.com/cli/cli/v2/pkg/liveshare"
"github.com/spf13/cobra"
)
func newJupyterCmd(app *App) *cobra.Command {
var codespace string
jupyterCmd := &cobra.Command{
Use: "jupyter",
Short: "Open a codespace in JupyterLab",
Args: noArgsConstraint,
RunE: func(cmd *cobra.Command, args []string) error {
return app.Jupyter(cmd.Context(), codespace)
},
}
jupyterCmd.Flags().StringVarP(&codespace, "codespace", "c", "", "Name of the codespace")
return jupyterCmd
}
func (a *App) Jupyter(ctx context.Context, codespaceName string) (err error) {
// Ensure all child tasks (e.g. port forwarding) terminate before return.
ctx, cancel := context.WithCancel(ctx)
defer cancel()
codespace, err := getOrChooseCodespace(ctx, a.apiClient, codespaceName)
if err != nil {
return err
}
session, err := startLiveShareSession(ctx, codespace, a, false, "")
if err != nil {
return err
}
defer safeClose(session, &err)
a.StartProgressIndicatorWithLabel("Starting JupyterLab on codespace")
serverPort, serverUrl, err := session.StartJupyterServer(ctx)
a.StopProgressIndicator()
if err != nil {
return fmt.Errorf("failed to start JupyterLab server: %w", err)
}
// Pass 0 to pick a random port
listen, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", 0))
if err != nil {
return err
}
defer listen.Close()
destPort := listen.Addr().(*net.TCPAddr).Port
tunnelClosed := make(chan error, 1)
go func() {
fwd := liveshare.NewPortForwarder(session, "jupyter", serverPort, true)
tunnelClosed <- fwd.ForwardToListener(ctx, listen) // always non-nil
}()
// Server URL contains an authentication token that must be preserved
targetUrl := strings.Replace(serverUrl, fmt.Sprintf("%d", serverPort), fmt.Sprintf("%d", destPort), 1)
err = a.browser.Browse(targetUrl)
if err != nil {
return fmt.Errorf("failed to open JupyterLab in browser: %w", err)
}
fmt.Fprintln(a.io.Out, targetUrl)
select {
case err := <-tunnelClosed:
return fmt.Errorf("tunnel closed: %w", err)
case <-ctx.Done():
return nil // success
}
}

13
pkg/cmd/codespace/logs.go
View file

@ -41,20 +41,11 @@ func (a *App) Logs(ctx context.Context, codespaceName string, follow bool) (err
return err
}
authkeys := make(chan error, 1)
go func() {
authkeys <- checkAuthorizedKeys(ctx, a.apiClient)
}()
session, err := codespaces.ConnectToLiveshare(ctx, a, noopLogger(), a.apiClient, codespace)
session, err := startLiveShareSession(ctx, codespace, a, false, "")
if err != nil {
return fmt.Errorf("connecting to codespace: %w", err)
}
defer safeClose(session, &err)
if err := <-authkeys; err != nil {
return err
}
defer safeClose(session, &err)
// Ensure local port is listening before client (getPostCreateOutput) connects.
listen, err := net.Listen("tcp", "127.0.0.1:0") // arbitrary port

4
pkg/cmd/codespace/ports.go
View file

@ -55,9 +55,9 @@ func (a *App) ListPorts(ctx context.Context, codespaceName string, exporter cmdu
devContainerCh := getDevContainer(ctx, a.apiClient, codespace)
session, err := codespaces.ConnectToLiveshare(ctx, a, noopLogger(), a.apiClient, codespace)
session, err := startLiveShareSession(ctx, codespace, a, false, "")
if err != nil {
return fmt.Errorf("error connecting to codespace: %w", err)
return err
}
defer safeClose(session, &err)

1
pkg/cmd/codespace/root.go
View file

@ -14,6 +14,7 @@ func NewRootCmd(app *App) *cobra.Command {
root.AddCommand(newCreateCmd(app))
root.AddCommand(newEditCmd(app))
root.AddCommand(newDeleteCmd(app))
root.AddCommand(newJupyterCmd(app))
root.AddCommand(newListCmd(app))
root.AddCommand(newLogsCmd(app))
root.AddCommand(newPortsCmd(app))

31
pkg/cmd/codespace/ssh.go
View file

@ -115,36 +115,14 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
ctx, cancel := context.WithCancel(ctx)
defer cancel()
// While connecting, ensure in the background that the user has keys installed.
// That lets us report a more useful error message if they don't.
authkeys := make(chan error, 1)
go func() {
authkeys <- checkAuthorizedKeys(ctx, a.apiClient)
}()
codespace, err := getOrChooseCodespace(ctx, a.apiClient, opts.codespace)
if err != nil {
return err
}
liveshareLogger := noopLogger()
if opts.debug {
debugLogger, err := newFileLogger(opts.debugFile)
if err != nil {
return fmt.Errorf("error creating debug logger: %w", err)
}
defer safeClose(debugLogger, &err)
liveshareLogger = debugLogger.Logger
a.errLogger.Printf("Debug file located at: %s", debugLogger.Name())
}
session, err := codespaces.ConnectToLiveshare(ctx, a, liveshareLogger, a.apiClient, codespace)
session, err := startLiveShareSession(ctx, codespace, a, opts.debug, opts.debugFile)
if err != nil {
if authErr := <-authkeys; authErr != nil {
return authErr
}
return fmt.Errorf("error connecting to codespace: %w", err)
return err
}
defer safeClose(session, &err)
@ -208,11 +186,10 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
}
}
func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) error {
func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) (err error) {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
var err error
var csList []*api.Codespace
if opts.codespace == "" {
a.StartProgressIndicatorWithLabel("Fetching codespaces")
@ -253,7 +230,7 @@ func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) error {
if err != nil {
result.err = fmt.Errorf("error connecting to codespace: %w", err)
} else {
defer session.Close()
defer safeClose(session, &err)
_, result.user, err = session.StartSSHServer(ctx)
if err != nil {

26
pkg/liveshare/session.go
View file

@ -62,6 +62,32 @@ func (s *Session) StartSSHServer(ctx context.Context) (int, string, error) {
return port, response.User, nil
}
// StartJupyterServer starts a Juypyter server in the container and returns
// the port on which it listens and the server URL.
func (s *Session) StartJupyterServer(ctx context.Context) (int, string, error) {
var response struct {
Result bool `json:"result"`
Message string `json:"message"`
Port string `json:"port"`
ServerUrl string `json:"serverUrl"`
}
if err := s.rpc.do(ctx, "IJupyterServerHostService.getRunningServer", []string{}, &response); err != nil {
return 0, "", fmt.Errorf("failed to invoke JupyterLab RPC: %w", err)
}
if !response.Result {
return 0, "", fmt.Errorf("failed to start JupyterLab: %s", response.Message)
}
port, err := strconv.Atoi(response.Port)
if err != nil {
return 0, "", fmt.Errorf("failed to parse JupyterLab port: %w", err)
}
return port, response.ServerUrl, nil
}
// heartbeat runs until context cancellation, periodically checking whether there is a
// reason to keep the connection alive, and if so, notifying the Live Share host to do so.
// Heartbeat ensures it does not send more than one request every "interval" to ratelimit