From e788b645020a26a213b17217ec15a1cd11b9bd08 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Wed, 11 Mar 2026 15:02:54 +0100 Subject: [PATCH 1/2] Add telemetry events for SSH tunnel connections Track SSH tunnel usage by logging an SshTunnelEvent at the end of each connection attempt. Captures compute type (dedicated/serverless), accelerator, IDE, client mode, cluster ID, server start time, and errors. Co-Authored-By: Claude Opus 4.6 --- experimental/ssh/internal/client/client.go | 51 ++++++++++++++ .../ssh/internal/client/client_test.go | 69 +++++++++++++++++++ libs/telemetry/protos/frontend_log.go | 1 + libs/telemetry/protos/ssh_tunnel_event.go | 52 ++++++++++++++ 4 files changed, 173 insertions(+) create mode 100644 libs/telemetry/protos/ssh_tunnel_event.go diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 5d9ca6cd24..d731bc79ac 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -27,6 +27,8 @@ import ( "github.com/databricks/cli/internal/build" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/telemetry" + "github.com/databricks/cli/libs/telemetry/protos" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/retries" "github.com/databricks/databricks-sdk-go/service/compute" @@ -203,6 +205,50 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt cancel() }() + event := BuildTelemetryEvent(opts) + + runErr := runConnect(ctx, client, opts, event) + if runErr != nil { + event.ErrorMessage = runErr.Error() + } + + telemetry.Log(ctx, protos.DatabricksCliLog{ + SshTunnelEvent: event, + }) + + return runErr +} + +// BuildTelemetryEvent creates an SshTunnelEvent pre-populated with data from client options. +func BuildTelemetryEvent(opts ClientOptions) *protos.SshTunnelEvent { + event := &protos.SshTunnelEvent{ + AcceleratorType: opts.Accelerator, + IdeType: opts.IDE, + AutoStartCluster: opts.AutoStartCluster, + } + + if opts.IsServerlessMode() { + event.ComputeType = protos.SshTunnelComputeTypeServerless + } else { + event.ComputeType = protos.SshTunnelComputeTypeDedicated + } + + switch { + case opts.ProxyMode: + event.ClientMode = protos.SshTunnelClientModeProxy + case opts.IDE != "": + event.ClientMode = protos.SshTunnelClientModeIDE + default: + event.ClientMode = protos.SshTunnelClientModeSSH + } + + // If metadata is provided, the server is already running — this is a reconnect from ProxyCommand. + event.IsReconnect = opts.ServerMetadata != "" + + return event +} + +func runConnect(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions, event *protos.SshTunnelEvent) error { // For serverless without explicit --name: auto-generate or reconnect to existing session. if opts.IsServerlessMode() && opts.ConnectionName == "" && !opts.ProxyMode { err := resolveServerlessSession(ctx, client, &opts) @@ -296,10 +342,13 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt return fmt.Errorf("failed to upload ssh-tunnel binaries: %w", err) } sp.Close() + + serverStartTime := time.Now() userName, serverPort, clusterID, err = ensureSSHServerIsRunning(ctx, client, version, secretScopeName, opts) if err != nil { return fmt.Errorf("failed to ensure that ssh server is running: %w", err) } + event.ServerStartTimeMs = time.Since(serverStartTime).Milliseconds() } else { // Metadata format: ",," metadata := strings.Split(opts.ServerMetadata, ",") @@ -326,6 +375,8 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt return errors.New("cluster ID is required for serverless connections but was not found in metadata") } + event.ClusterID = clusterID + log.Infof(ctx, "Remote user name: %s", userName) log.Infof(ctx, "Server port: %d", serverPort) if opts.IsServerlessMode() { diff --git a/experimental/ssh/internal/client/client_test.go b/experimental/ssh/internal/client/client_test.go index 5821ab7b6c..3a361e464c 100644 --- a/experimental/ssh/internal/client/client_test.go +++ b/experimental/ssh/internal/client/client_test.go @@ -7,6 +7,7 @@ import ( "time" "github.com/databricks/cli/experimental/ssh/internal/client" + "github.com/databricks/cli/libs/telemetry/protos" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -169,3 +170,71 @@ func TestToProxyCommand(t *testing.T) { }) } } + +func TestBuildTelemetryEvent(t *testing.T) { + tests := []struct { + name string + opts client.ClientOptions + want *protos.SshTunnelEvent + }{ + { + name: "dedicated cluster with SSH client", + opts: client.ClientOptions{ + ClusterID: "abc-123", + AutoStartCluster: true, + }, + want: &protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeDedicated, + ClientMode: protos.SshTunnelClientModeSSH, + AutoStartCluster: true, + }, + }, + { + name: "serverless with IDE", + opts: client.ClientOptions{ + ConnectionName: "my-conn", + Accelerator: "GPU_1xA10", + IDE: "vscode", + }, + want: &protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeServerless, + ClientMode: protos.SshTunnelClientModeIDE, + AcceleratorType: "GPU_1xA10", + IdeType: "vscode", + }, + }, + { + name: "proxy mode with metadata (reconnect)", + opts: client.ClientOptions{ + ClusterID: "abc-123", + ProxyMode: true, + ServerMetadata: "user,2222,abc-123", + }, + want: &protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeDedicated, + ClientMode: protos.SshTunnelClientModeProxy, + IsReconnect: true, + }, + }, + { + name: "serverless proxy mode", + opts: client.ClientOptions{ + ConnectionName: "my-conn", + Accelerator: "GPU_8xH100", + ProxyMode: true, + }, + want: &protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeServerless, + ClientMode: protos.SshTunnelClientModeProxy, + AcceleratorType: "GPU_8xH100", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := client.BuildTelemetryEvent(tt.opts) + assert.Equal(t, tt.want, got) + }) + } +} diff --git a/libs/telemetry/protos/frontend_log.go b/libs/telemetry/protos/frontend_log.go index 7e6ab1012b..816297a8ee 100644 --- a/libs/telemetry/protos/frontend_log.go +++ b/libs/telemetry/protos/frontend_log.go @@ -19,4 +19,5 @@ type DatabricksCliLog struct { CliTestEvent *CliTestEvent `json:"cli_test_event,omitempty"` BundleInitEvent *BundleInitEvent `json:"bundle_init_event,omitempty"` BundleDeployEvent *BundleDeployEvent `json:"bundle_deploy_event,omitempty"` + SshTunnelEvent *SshTunnelEvent `json:"ssh_tunnel_event,omitempty"` } diff --git a/libs/telemetry/protos/ssh_tunnel_event.go b/libs/telemetry/protos/ssh_tunnel_event.go new file mode 100644 index 0000000000..51d6c245c9 --- /dev/null +++ b/libs/telemetry/protos/ssh_tunnel_event.go @@ -0,0 +1,52 @@ +package protos + +// SshTunnelComputeType represents the type of compute used for SSH tunnel. +type SshTunnelComputeType string + +const ( + SshTunnelComputeTypeUnspecified SshTunnelComputeType = "TYPE_UNSPECIFIED" + SshTunnelComputeTypeDedicated SshTunnelComputeType = "DEDICATED" + SshTunnelComputeTypeServerless SshTunnelComputeType = "SERVERLESS" +) + +// SshTunnelClientMode represents how the SSH tunnel client is used. +type SshTunnelClientMode string + +const ( + SshTunnelClientModeUnspecified SshTunnelClientMode = "TYPE_UNSPECIFIED" + SshTunnelClientModeSSH SshTunnelClientMode = "SSH_CLIENT" + SshTunnelClientModeProxy SshTunnelClientMode = "PROXY" + SshTunnelClientModeIDE SshTunnelClientMode = "IDE" +) + +// SshTunnelEvent tracks SSH tunnel connection lifecycle and usage. +type SshTunnelEvent struct { + // Type of compute: dedicated cluster or serverless. + ComputeType SshTunnelComputeType `json:"compute_type,omitempty"` + + // GPU accelerator type for serverless compute (e.g., "GPU_1xA10", "GPU_8xH100"). + AcceleratorType string `json:"accelerator_type,omitempty"` + + // IDE used for the connection (e.g., "vscode", "cursor"), empty if none. + IdeType string `json:"ide_type,omitempty"` + + // How the client is used: SSH client, proxy mode, or IDE mode. + ClientMode SshTunnelClientMode `json:"client_mode,omitempty"` + + // Cluster ID of the compute resource (for dedicated clusters this is the input, + // for serverless it is discovered from the server metadata via Driver Proxy). + ClusterID string `json:"cluster_id,omitempty"` + + // Whether this is a reconnection to an existing session. + IsReconnect bool `json:"is_reconnect,omitempty"` + + // Whether the cluster was auto-started by the CLI. + AutoStartCluster bool `json:"auto_start_cluster,omitempty"` + + // Time in milliseconds spent starting the SSH server (including job submission + // and waiting for the server to become ready). Zero if server was already running. + ServerStartTimeMs int64 `json:"server_start_time_ms"` + + // Error message if the connection failed. Empty on success. + ErrorMessage string `json:"error_message,omitempty"` +} From bd0048b0c6ce97a55bf51ec1e8cc15549520c397 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Wed, 11 Mar 2026 17:16:26 +0100 Subject: [PATCH 2/2] change fields --- experimental/ssh/internal/client/client.go | 6 +++--- libs/telemetry/protos/ssh_tunnel_event.go | 8 ++------ 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index d731bc79ac..112b5bddc1 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -209,7 +209,9 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt runErr := runConnect(ctx, client, opts, event) if runErr != nil { - event.ErrorMessage = runErr.Error() + event.IsSuccess = false + } else { + event.IsSuccess = true } telemetry.Log(ctx, protos.DatabricksCliLog{ @@ -375,8 +377,6 @@ func runConnect(ctx context.Context, client *databricks.WorkspaceClient, opts Cl return errors.New("cluster ID is required for serverless connections but was not found in metadata") } - event.ClusterID = clusterID - log.Infof(ctx, "Remote user name: %s", userName) log.Infof(ctx, "Server port: %d", serverPort) if opts.IsServerlessMode() { diff --git a/libs/telemetry/protos/ssh_tunnel_event.go b/libs/telemetry/protos/ssh_tunnel_event.go index 51d6c245c9..cd7abf1ecd 100644 --- a/libs/telemetry/protos/ssh_tunnel_event.go +++ b/libs/telemetry/protos/ssh_tunnel_event.go @@ -33,10 +33,6 @@ type SshTunnelEvent struct { // How the client is used: SSH client, proxy mode, or IDE mode. ClientMode SshTunnelClientMode `json:"client_mode,omitempty"` - // Cluster ID of the compute resource (for dedicated clusters this is the input, - // for serverless it is discovered from the server metadata via Driver Proxy). - ClusterID string `json:"cluster_id,omitempty"` - // Whether this is a reconnection to an existing session. IsReconnect bool `json:"is_reconnect,omitempty"` @@ -47,6 +43,6 @@ type SshTunnelEvent struct { // and waiting for the server to become ready). Zero if server was already running. ServerStartTimeMs int64 `json:"server_start_time_ms"` - // Error message if the connection failed. Empty on success. - ErrorMessage string `json:"error_message,omitempty"` + // Flag to indicate if the connection was successful + IsSuccess bool `json:"is_success,omitempty"` }