Split and move ssh logic to experimental/ssh #3544

ilia-db · 2025-09-03T11:34:30Z

Changes

Previous PRs in the stack:

Move from flat structure to cmd and internal client/server/proxy/keys/workspace sub packages.

Moved a bunch of hardcoded values to constants (some hardcoded stuff remains in the keys and workspace sub packages, will be improved in a follow up)

Added unit tests for servers connection manager logic and updated proxy tests.
The coverage isn't great yet, the goal is to avoid unit tests with excessive mocking and focus on integration tests from now on.

Why

/experimental already has the databricks-bundles python project, seems like a good place to put ssh stuff too

Tests

Unit and manual

eng-dev-ecosystem-bot · 2025-09-03T11:46:03Z

Run: 17575476240

	Env	✅pass	❌FAIL	🔄flaky	🙈skip
❌	aws linux	305	1		523
❌	aws windows	304	1	2	522
✅	aws-ucws linux	420			421
✅	aws-ucws windows	421			420
✅	azure linux	308			522
✅	azure windows	309			521
🔄	azure-ucws linux	413		7	420
✅	azure-ucws windows	421			419
✅	gcp linux	307			524
✅	gcp windows	308			523

10 failing tests:

Test Name	aws linux	aws windows	azure-ucws linux
TestAccept/bundle/deploy/dashboard/generate_inplace	✅pass	✅pass	🔄flaky
TestAccept/bundle/deploy/dashboard/nested-folders	✅pass	✅pass	🔄flaky
TestAccept/bundle/deploy/dashboard/simple	✅pass	✅pass	🔄flaky
TestAccept/bundle/deploy/dashboard/simple_outside_bundle_root	✅pass	✅pass	🔄flaky
TestAccept/bundle/deploy/dashboard/simple_syncroot	✅pass	✅pass	🔄flaky
TestAccept/bundle/deployment/bind/dashboard	✅pass	✅pass	🔄flaky
TestAccept/bundle/templates/default-python/combinations/classic/DATABRICKS_CLI_DEPLOYMENT=terraform/DLT=yes/NBOOK=no/PY=no	✅pass	🔄flaky	✅pass
TestAccept/bundle/templates/default-python/combinations/classic/DATABRICKS_CLI_DEPLOYMENT=terraform/DLT=yes/NBOOK=yes/PY=no	✅pass	🔄flaky	✅pass
TestDashboardAssumptions_WorkspaceImport	✅pass	✅pass	🔄flaky
TestFetchRepositoryInfoAPI_FromRepo	❌FAIL	❌FAIL	✅pass

pietern

I haven't looked at the unit tests closely, but worry that they can run for a long time or be flaky. A good practice is to include a context.Context everywhere and add a short-lived deadline to tear down everything it spawns.

pietern · 2025-09-08T13:49:05Z

experimental/ssh/cmd/connect.go

+		err := client.RunClient(ctx, wsClient, opts)
+		if err != nil && proxy.IsNormalClosure(err) {
+			return nil
 		}


The normal closure check was added here and removed from the server case.

Can they both be omitted (or hoisted into the "run" function)?

Improved in #3569

pietern · 2025-09-08T13:49:43Z

experimental/ssh/cmd/server.go

-	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", 10*time.Minute, "Delay before shutting down after no pings from clients")
+	cmd.Flags().StringVar(&secretsScope, "secrets-scope-name", "", "Databricks secrets scope name")
+	cmd.MarkFlagRequired("secrets-scope-name")
+	cmd.Flags().StringVar(&publicKeySecretName, "client-key-name", "", "Databricks client key name")


Is this an actual key name or the secret name in the secret scope?

Secret name in the scope - you have to combine it with secrets-scope-name to get the value.

Server uses the scope not only for downloading the client key value though, but also for managing its own key-pair

pietern · 2025-09-08T13:50:18Z

experimental/ssh/cmd/server.go

 	cmd.MarkFlagRequired("cluster")
-	cmd.Flags().IntVar(&maxClients, "max-clients", 10, "Maximum number of SSH clients")
-	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", 10*time.Minute, "Delay before shutting down after no pings from clients")
+	cmd.Flags().StringVar(&secretsScope, "secrets-scope-name", "", "Databricks secrets scope name")


The product name is "secret scope" (singular).

pietern · 2025-09-08T13:54:56Z

experimental/ssh/internal/client/ssh-server-bootstrap.py

                "server",
                f"--cluster={ctx.clusterId}",
+                f"--secrets-scope-name={secrets_scope}",
+                f"--client-key-name={public_key_secret_name}",


I recommend being more verbose to make it unambiguous what's being referred to here:

Idea:

--keypair-secret-scope-name

--keypair-secret-key-name

Or pubkey if the secret scope only stores the public key.

The secret scope has 3 keys - public and private keys of the server, and the public key of the client (to be added to AuthorizedKeysFile).

Maybe --keys-secret-scope-name and --authorized-key-secret-name?

Technically we don't even need to pass these options to the server, as it can re-create the names using the same constants and logic the client uses. Since the client enforces the same server version, it will be safe. It doesn't hurt to keep server a bit more isolated though.

If the goal is to keep these isolated, then all of them can specify the scope name.

--server-pubkey-secret-name
--server-privkey-secret-name
--client-pubkey-secret-name

With the format being scope/key, like https://docs.databricks.com/aws/en/security/secrets/secrets-spark-conf-env-var#reference-a-secret-with-a-spark-configuration-property.

That removes all implicits.

pietern · 2025-09-08T13:56:06Z

experimental/ssh/cmd/connect.go

 			ClientPublicKeyName: defaultClientPublicKeyName,
+			ServerTimeout:       serverTimeout,
+		}
+		err := client.RunClient(ctx, wsClient, opts)


Can be shortened to client.Run.

Changed it in #3569

pietern · 2025-09-08T13:56:14Z

experimental/ssh/cmd/server.go

+			PortRange:            serverPortRange,
 		}
-		return err
+		return server.RunServer(ctx, client, opts)


Can be shortened to server.Run.

The client variable here is unfortunate.

Common variable names for the workspace client include w or wsc.

Changed it in #3569

pietern · 2025-09-08T13:57:31Z

experimental/ssh/internal/server/connections.go

+	shutdownTimerMu sync.Mutex
+	connections     map[string]*proxy.ProxyConnection
+	connectionsMu   sync.Mutex
+	TimedOut        chan bool


Why exported?

internal/server package listens to this channel (and terminates the process)

Oh wait, it indeed doesn't make sense in this PR, but in the follow up I've moved connections to the internal/proxy, so the export is needed there

pietern · 2025-09-08T14:11:33Z

libs/ssh/ssh-server-bootstrap.py

Wasn't this implemented in the previous PR?

It was, git didn't register this as a rename for some reason. The main thing I've changed is removing databricks_ssh_key env var, since secrets logic was moved into the server itself

pietern · 2025-09-08T14:15:54Z

experimental/ssh/internal/server/server.go

+	SSHDConfigPath string
+}
+
+func newSSHHandler(ctx context.Context, connections *connectionsManager, sshDConfigPath string) *sshHandler {


Can be a separate PR:

There is mixed casing for ssh/sshd. You can stick to just lowercase and only upcase the first char when it's the second word in an unexported variable, or the first in an exported variable.

In the follow up the pure http.Handler part of the logic was moved to the internal/proxy, and sshD* mixed case is fixed there.

If you mean that I sometimes use ssh and sometimes SSH - it seems to be in line with go camel case + acronyms style? And the real reason is because Ssh throws me off :-)

Also sshD, and the newSSHHandler also looks weird because of the double H.

They can be just sshd (referring to the process name), and newSshHandler to stick to Go patterns.

Or alternatively, move everything into an ssh package and call it handler, configPath, etc.

Move from flat structure to cmd and internal client/server/proxy/keys/workspace sub packages

ilia-db · 2025-09-09T09:41:37Z

I haven't looked at the unit tests closely, but worry that they can run for a long time or be flaky. A good practice is to include a context.Context everywhere and add a short-lived deadline to tear down everything it spawns.

Moved the tests to use the new testing/synctest std lib package in the follow up PR #3569

pietern · 2025-09-10T09:15:23Z

experimental/ssh/internal/client/ssh-server-bootstrap.py

                "server",
                f"--cluster={ctx.clusterId}",
+                f"--secrets-scope-name={secrets_scope}",
+                f"--client-key-name={public_key_secret_name}",


If the goal is to keep these isolated, then all of them can specify the scope name.

--server-pubkey-secret-name
--server-privkey-secret-name
--client-pubkey-secret-name

With the format being scope/key, like https://docs.databricks.com/aws/en/security/secrets/secrets-spark-conf-env-var#reference-a-secret-with-a-spark-configuration-property.

That removes all implicits.

pietern · 2025-09-10T09:19:57Z

experimental/ssh/internal/server/server.go

+	SSHDConfigPath string
+}
+
+func newSSHHandler(ctx context.Context, connections *connectionsManager, sshDConfigPath string) *sshHandler {


Also sshD, and the newSSHHandler also looks weird because of the double H.

They can be just sshd (referring to the process name), and newSshHandler to stick to Go patterns.

Or alternatively, move everything into an ssh package and call it handler, configPath, etc.

…3569) ## Changes Based on #3544 Move pure client and server logic to the proxy package and test it. Now `proxyConnection` is completely private, and `proxy` package exposes only 2 public functions: - `RunClientProxy` function - `NewProxyServer` constructor Meanwhile the internal logic doesn't depend on any ssh stuff and can be unit tested. `internal/client` and `internal/server` packages are now much smaller and contain mainly ssh/jobs/cluster glue. Also adding `bench.sh` until we have proper e2e tests. ## Why Tests ## Tests Unit and manual tests

ilia-db requested review from pietern, andrewnester, shreyas-goenka, denik and anton-107 as code owners September 3, 2025 11:34

ilia-db temporarily deployed to test-trigger-is September 3, 2025 11:34 — with GitHub Actions Inactive

ilia-db mentioned this pull request Sep 3, 2025

Add "ssh connect" command #3471

Merged

ilia-db force-pushed the ssh-tunnel-exp branch from b3a35d3 to e86af23 Compare September 3, 2025 12:16

ilia-db temporarily deployed to test-trigger-is September 3, 2025 12:16 — with GitHub Actions Inactive

ilia-db force-pushed the ssh-tunnel-exp branch from e86af23 to 78e037d Compare September 3, 2025 12:24

ilia-db force-pushed the ssh-tunnel-server branch from 60fc8a3 to 8be1123 Compare September 3, 2025 12:24

ilia-db temporarily deployed to test-trigger-is September 3, 2025 12:24 — with GitHub Actions Inactive

ilia-db mentioned this pull request Sep 3, 2025

Add "ssh server" command #3475

Merged

ilia-db force-pushed the ssh-tunnel-exp branch from 78e037d to bf3373e Compare September 5, 2025 10:31

ilia-db temporarily deployed to test-trigger-is September 5, 2025 10:31 — with GitHub Actions Inactive

ilia-db force-pushed the ssh-tunnel-server branch from 8be1123 to 7ebfe7e Compare September 8, 2025 12:54

ilia-db force-pushed the ssh-tunnel-exp branch from bf3373e to 2222cb9 Compare September 8, 2025 12:54

ilia-db temporarily deployed to test-trigger-is September 8, 2025 12:54 — with GitHub Actions Inactive

ilia-db force-pushed the ssh-tunnel-exp branch from 2222cb9 to 94c2ae2 Compare September 8, 2025 13:55

ilia-db temporarily deployed to test-trigger-is September 8, 2025 13:55 — with GitHub Actions Inactive

ilia-db mentioned this pull request Sep 8, 2025

Move pure ssh client and server logic to proxy package and test it #3569

Merged

Base automatically changed from ssh-tunnel-server to main September 8, 2025 14:05

pietern reviewed Sep 8, 2025

View reviewed changes

Split and move ssh logic to experimental/ssh

8f69376

Move from flat structure to cmd and internal client/server/proxy/keys/workspace sub packages

ilia-db force-pushed the ssh-tunnel-exp branch from 94c2ae2 to 8f69376 Compare September 9, 2025 07:40

ilia-db temporarily deployed to test-trigger-is September 9, 2025 07:41 — with GitHub Actions Inactive

pietern approved these changes Sep 10, 2025

View reviewed changes

ilia-db added this pull request to the merge queue Sep 10, 2025

Merged via the queue into main with commit c2e3322 Sep 10, 2025
13 checks passed

ilia-db deleted the ssh-tunnel-exp branch September 10, 2025 10:38

Split and move ssh logic to experimental/ssh #3544

Split and move ssh logic to experimental/ssh #3544

Uh oh!

Conversation

ilia-db commented Sep 3, 2025

Changes

Why

Tests

Uh oh!

eng-dev-ecosystem-bot commented Sep 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pietern left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ilia-db commented Sep 9, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

eng-dev-ecosystem-bot commented Sep 3, 2025 •

edited

Loading