Skip to content

Commit ead9e8c

Browse files
ilia-dbpietern
ilia-db
and
pietern
authored
Add "ssh" and "ssh setup" commands (#3470)
## Changes This is the first PR that introduces the ssh-tunnel functionality into the CLI See README for an overview of the SSH tunnel feature and its usage. This PR only adds `ssh setup` command, `connect` and `server` are here: - #3471 - #3475 The whole `ssh` command group is hidden (not available in `--help` output, and in this PR it's not even added to the root cmd), and the usage text has private preview disclaimer. Hidden flag will be removed when we move to public preview phase (no timeline yet). `setup` command updates local ssh config (`~/.ssh/config`). The config entry relies on `ProxyCommand` option calling `databricks ssh connect`, which is implemented in a follow up PR. ## Why We've agreed to move the implementation form the separate repo to the main databricks CLI ## Tests Manually and unit tests. --------- Co-authored-by: Pieter Noordhuis <[email protected]>
1 parent f534b13 commit ead9e8c

File tree

6 files changed

+828
-0
lines changed

6 files changed

+828
-0
lines changed

‎cmd/ssh/README.md‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
## SSH Tunnel for Databricks
2+
3+
The SSH tunnel lets customers connect any IDE to Databricks compute to run and debug all code - including non-Spark/ML - with environment parity, and simple setup.
4+
5+
## Compute Requirements
6+
- Dedicated (single user) access mode if you want to use Remote Development tools in IDEs
7+
- Dedicated or standard access mode for terminal SSH connections
8+
9+
## Usage
10+
A. With local ssh config setup:
11+
```shell
12+
databricks ssh setup --name=hello --cluster=id # one time only
13+
ssh hello # use system SSH client to create a session
14+
```
15+
B. Spawn an ssh session directly:
16+
```shell
17+
databricks ssh connect --cluster=id
18+
```
19+
20+
## Development
21+
```shell
22+
make build snapshot-release
23+
./cli ssh connect --cluster=<id> --releases-dir=./dist --debug # or modify ssh config accordingly
24+
```
25+
26+
## Design
27+
28+
High level:
29+
```mermaid
30+
---
31+
config:
32+
theme: redux
33+
layout: dagre
34+
---
35+
flowchart TD
36+
n1(["Client A"])
37+
subgraph s1["Control Plane"]
38+
n3["Jobs API"]
39+
n2["Driver Proxy API"]
40+
n11["Workspace API"]
41+
end
42+
subgraph s3["Spark User A or root"]
43+
n4["SSH Server A"]
44+
end
45+
subgraph s4["Spark User B or root"]
46+
n6["SSH Server B"]
47+
end
48+
subgraph s2["Cluster"]
49+
s3
50+
s4
51+
n12["Workspace Filesystem"]
52+
end
53+
n1 -. "1 - start an ssh server job" .-> n3
54+
n3 -. "2 - start ssh server" .-> n4
55+
n4 <-. "3 - save the ssh server port number" .-> n12
56+
n1 <-. "4 - get ssh server port number" .-> n11
57+
n1 <-. "6 - websocket connection" .-> n2
58+
n2 <-. "7 - websocket connection" .-> n4
59+
n6 <-.-> n12
60+
style s2 stroke:#757575
61+
style s1 stroke:#757575
62+
style s4 stroke-dasharray: 5 5
63+
style n6 stroke-dasharray: 5 5
64+
```
65+
66+
Connection flow:
67+
```mermaid
68+
---
69+
config:
70+
theme: base
71+
---
72+
sequenceDiagram
73+
autonumber
74+
participant P1 as databricks ssh connect
75+
participant P2 as ssh client
76+
participant P3 as databricks ssh connect --proxy
77+
participant P4 as wsfs
78+
participant P6 as databricks ssh server
79+
participant P7 as sshd
80+
Note over P1,P6: Try to get a port and a remote user name of an existing server<br/> ($v is databricks CLI version, $cluster is supplied by the user)
81+
activate P1
82+
P1 ->> P4: GET ~/.ssh/$v/$cluster/metadata.json
83+
P4 -->> P1: {port: xxxx} or error
84+
P1 ->> P6: GET /driver-proxy-api/$cluster/$port/metadata
85+
P6 -->> P1: {user: spark-xxxx} or {user: root} or error
86+
Note over P1,P6: Start the new server in the case of an error
87+
opt
88+
P1 -->> P1: generate<br/>key pair
89+
P1 -->> P4: PUT ~/.ssh/$v/bin/databricks, unless it's already there
90+
P1 ->> P4: PUT ~/.ssh/$v/$cluster/start-server-with-pub-key.ipynb
91+
P1 ->> P6: jobs/runs/submit start-server-with-pub-key.ipynb $cluster
92+
activate P6
93+
P6 ->> P6: start self-kill-timeout<br/>generate server key pair<br/>create custom sshd config<br/>listen for /ssh and /metadata on a free port
94+
P6 ->> P4: PUT ~/.ssh/$v/$cluster/metadata.json<br/>{port: xxxx}
95+
loop unil successful or timed out
96+
P1 -> P6: Get port and remote user name of the server (sequence 1 - 4 above)
97+
end
98+
end
99+
Note over P1,P7: We know the port and the user, spawn "ssh"
100+
P1 ->> P2: ssh -l $user -i $key<br/> -o ProxyCommand="databricks ssh connect --proxy $cluster $user $port"
101+
activate P2
102+
P2 ->> P3: exec ProxyCommand
103+
activate P3
104+
P3 ->> P6: wss:/dirver-proxy-api/$cluster/$port/ssh
105+
P6 ->> P6: stop self-kill-timeout
106+
P6 ->> P7: /usr/sbin/sshd -i -f config
107+
activate P7
108+
P2 -> P7: pubkey auth
109+
loop until the connection is closed<br/>by ssh client, sshd, or driver-proxy
110+
P2 -> P7: stdin and stdout
111+
P1 -> P2: stdin, stdout, and stderr
112+
deactivate P7
113+
deactivate P3
114+
deactivate P2
115+
deactivate P1
116+
end
117+
break when the last ws connection drops
118+
P6 ->> P6: start self-kill-timeout
119+
P6 ->> P4: DELETE ~/.ssh/$v/$cluster/metadata.json
120+
deactivate P6
121+
end
122+
```

‎cmd/ssh/setup.go‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package ssh
2+
3+
import (
4+
"time"
5+
6+
"github.com/databricks/cli/cmd/root"
7+
"github.com/databricks/cli/libs/cmdctx"
8+
"github.com/databricks/cli/libs/ssh"
9+
"github.com/spf13/cobra"
10+
)
11+
12+
func newSetupCommand() *cobra.Command {
13+
cmd := &cobra.Command{
14+
Use: "setup",
15+
Short: "Setup SSH configuration for Databricks compute",
16+
Long: `Setup SSH configuration for Databricks compute.
17+
18+
This command configures SSH to connect to Databricks compute by adding
19+
an SSH host configuration to your SSH config file.
20+
21+
` + disclaimer,
22+
}
23+
24+
var hostName string
25+
var clusterID string
26+
var sshConfigPath string
27+
var shutdownDelay time.Duration
28+
29+
cmd.Flags().StringVar(&hostName, "name", "", "Host name to use in SSH config")
30+
cmd.MarkFlagRequired("name")
31+
cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID")
32+
cmd.MarkFlagRequired("cluster")
33+
cmd.Flags().StringVar(&sshConfigPath, "ssh-config", "", "Path to SSH config file (default ~/.ssh/config)")
34+
cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", 10*time.Minute, "SSH server will terminate after this delay if there are no active connections")
35+
36+
cmd.PreRunE = root.MustWorkspaceClient
37+
cmd.RunE = func(cmd *cobra.Command, args []string) error {
38+
ctx := cmd.Context()
39+
client := cmdctx.WorkspaceClient(ctx)
40+
opts := ssh.SetupOptions{
41+
HostName: hostName,
42+
ClusterID: clusterID,
43+
SSHConfigPath: sshConfigPath,
44+
ShutdownDelay: shutdownDelay,
45+
Profile: client.Config.Profile,
46+
}
47+
return ssh.Setup(ctx, client, opts)
48+
}
49+
50+
return cmd
51+
}

‎cmd/ssh/ssh.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package ssh
2+
3+
import (
4+
"github.com/spf13/cobra"
5+
)
6+
7+
const disclaimer = `WARNING! This is an experimental feature:
8+
- The product is in preview and not intended to be used in production;
9+
- The product may change or may never be released;
10+
- While we will not charge separately for this product right now, we may charge for it in the future. You will still incur charges for DBUs;
11+
- There's no formal support or SLAs for the preview - so please reach out to your account or other contact with any questions or feedback;`
12+
13+
func New() *cobra.Command {
14+
cmd := &cobra.Command{
15+
Use: "ssh",
16+
Short: "Connect to Databricks compute with ssh",
17+
Hidden: true,
18+
Long: `Connect to Databricks compute with ssh.
19+
20+
SSH commands let you setup and establish ssh connections to Databricks compute.
21+
22+
Common workflows:
23+
databricks ssh connect --cluster=<cluster-id> --profile=<profile-name> # connect to a cluster without any setup
24+
databricks ssh setup --name=my-compute --cluster=<cluster-id> # update local ssh config
25+
ssh my-compute # connect to the compute using ssh client
26+
27+
` + disclaimer,
28+
}
29+
30+
cmd.AddCommand(newSetupCommand())
31+
32+
return cmd
33+
}

‎libs/ssh/keys.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package ssh
2+
3+
import (
4+
"fmt"
5+
"os"
6+
"path/filepath"
7+
)
8+
9+
// We use different client keys for each cluster as a good practice for better isolation and control.
10+
func getLocalSSHKeyPath(clusterID, keysDir string) (string, error) {
11+
if keysDir == "" {
12+
homeDir, err := os.UserHomeDir()
13+
if err != nil {
14+
return "", fmt.Errorf("failed to get home directory: %w", err)
15+
}
16+
keysDir = filepath.Join(homeDir, ".databricks", "ssh-tunnel-keys")
17+
}
18+
return filepath.Join(keysDir, clusterID), nil
19+
}

0 commit comments

Comments
 (0)