diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 3cecb9c30e..67b20aa84d 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -1150,6 +1150,14 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC } func logSshTunnelEvent(ctx context.Context, opts ClientOptions, isSuccess, isReconnect bool, serverStartTimeMs int64) { + telemetry.Log(ctx, protos.DatabricksCliLog{ + SshTunnelEvent: buildSshTunnelEvent(opts, isSuccess, isReconnect, serverStartTimeMs), + }) +} + +// buildSshTunnelEvent maps the connection options and outcome onto the telemetry +// event. It is separated from logSshTunnelEvent so the field mapping can be unit tested. +func buildSshTunnelEvent(opts ClientOptions, isSuccess, isReconnect bool, serverStartTimeMs int64) *protos.SshTunnelEvent { computeType := protos.SshTunnelComputeTypeDedicated if opts.IsServerlessMode() { computeType = protos.SshTunnelComputeTypeServerless @@ -1165,16 +1173,15 @@ func logSshTunnelEvent(ctx context.Context, opts ClientOptions, isSuccess, isRec clientMode = protos.SshTunnelClientModeSSH } - telemetry.Log(ctx, protos.DatabricksCliLog{ - SshTunnelEvent: &protos.SshTunnelEvent{ - ComputeType: computeType, - AcceleratorType: opts.Accelerator, - IdeType: opts.IDE, - ClientMode: clientMode, - IsReconnect: isReconnect, - AutoStartCluster: opts.AutoStartCluster, - ServerStartTimeMs: serverStartTimeMs, - IsSuccess: isSuccess, - }, - }) + return &protos.SshTunnelEvent{ + ComputeType: computeType, + AcceleratorType: opts.Accelerator, + IdeType: opts.IDE, + ClientMode: clientMode, + IsReconnect: isReconnect, + AutoStartCluster: opts.AutoStartCluster, + ServerStartTimeMs: serverStartTimeMs, + IsSuccess: isSuccess, + HasBaseEnvironment: opts.BaseEnvironment != "", + } } diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index 57861eb25c..f71f38d089 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -7,6 +7,7 @@ import ( "time" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/telemetry/protos" "github.com/databricks/databricks-sdk-go/experimental/mocks" "github.com/databricks/databricks-sdk-go/service/compute" "github.com/databricks/databricks-sdk-go/service/environments" @@ -416,3 +417,68 @@ func TestTailWriterRetainsTail(t *testing.T) { assert.Equal(t, "ab", w.String()) }) } + +func TestBuildSshTunnelEvent(t *testing.T) { + tests := []struct { + name string + opts ClientOptions + want protos.SshTunnelEvent + }{ + { + name: "dedicated cluster via raw SSH client", + opts: ClientOptions{ClusterID: "abc-123", AutoStartCluster: true}, + want: protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeDedicated, + ClientMode: protos.SshTunnelClientModeSSH, + AutoStartCluster: true, + }, + }, + { + name: "serverless with accelerator", + opts: ClientOptions{ConnectionName: "my-conn", Accelerator: "GPU_1xA10"}, + want: protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeServerless, + AcceleratorType: "GPU_1xA10", + ClientMode: protos.SshTunnelClientModeSSH, + }, + }, + { + name: "proxy mode takes precedence over IDE", + opts: ClientOptions{ConnectionName: "my-conn", ProxyMode: true, IDE: "vscode"}, + want: protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeServerless, + IdeType: "vscode", + ClientMode: protos.SshTunnelClientModeProxy, + }, + }, + { + name: "IDE mode", + opts: ClientOptions{ConnectionName: "my-conn", IDE: "cursor"}, + want: protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeServerless, + IdeType: "cursor", + ClientMode: protos.SshTunnelClientModeIDE, + }, + }, + { + // The raw --base-environment value can carry PII, so only its presence is recorded. + name: "custom base environment records presence only", + opts: ClientOptions{ConnectionName: "my-conn", BaseEnvironment: "/Workspace/Users/me@example.com/env.yaml"}, + want: protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeServerless, + ClientMode: protos.SshTunnelClientModeSSH, + HasBaseEnvironment: true, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := buildSshTunnelEvent(tt.opts, true, true, 1500) + tt.want.IsSuccess = true + tt.want.IsReconnect = true + tt.want.ServerStartTimeMs = 1500 + assert.Equal(t, &tt.want, got) + }) + } +} diff --git a/libs/telemetry/protos/ssh_tunnel.go b/libs/telemetry/protos/ssh_tunnel.go index 42be8233b7..062bc0fff4 100644 --- a/libs/telemetry/protos/ssh_tunnel.go +++ b/libs/telemetry/protos/ssh_tunnel.go @@ -38,6 +38,11 @@ type SshTunnelEvent struct { // Whether the cluster was auto-started by the CLI. AutoStartCluster bool `json:"auto_start_cluster,omitempty"` + // Whether a custom base environment was set via --base-environment. + // Only the presence is recorded: the flag value can be an env.yaml path + // or display name carrying PII, so the value itself is not logged. + HasBaseEnvironment bool `json:"has_base_environment,omitempty"` + // Time in milliseconds spent starting the SSH server. // Zero if server was already running. ServerStartTimeMs int64 `json:"server_start_time_ms"`