Add telemetry for serverless and workflow/dbconnect runs (#1543)

## Changes

Add telemetry for serverless and workflow/dbconnect runs

## Tests

Manually

Author: ilia-db

packages/databricks-vscode/src/configuration/ConnectionManager.ts

@@ -137,6 +137,11 @@ export class ConnectionManager implements Disposable {
             ) {
                 this.cli.setClusterId(clusterId);
             }
+            if (this.cluster) {
+                this.telemetry.recordEvent(Events.COMPUTE_SELECTED, {
+                    type: "cluster",
+                });
+            }
             this.onDidChangeClusterEmitter.fire(this.cluster);
         } catch (e) {
             this.configModel.set("clusterId", undefined);
@@ -462,6 +467,9 @@ export class ConnectionManager implements Disposable {
             await this.detachCluster();
             this.customWhenContext.setServerless(true);
             this.onDidChangeClusterEmitter.fire(undefined);
+            this.telemetry.recordEvent(Events.COMPUTE_SELECTED, {
+                type: "serverless",
+            });
         }
     }
 

packages/databricks-vscode/src/extension.ts

@@ -768,7 +768,8 @@ export async function activate(
         pythonExtensionWrapper,
         featureManager,
         context,
-        customWhenContext
+        customWhenContext,
+        telemetry
     );
     const debugFactory = new DatabricksDebugAdapterFactory(
         connectionManager,
@@ -780,7 +781,8 @@ export async function activate(
         connectionManager,
         configModel,
         context,
-        bundleCommands
+        bundleCommands,
+        telemetry
     );
 
     context.subscriptions.push(

packages/databricks-vscode/src/run/DatabricksWorkflowDebugAdapter.ts

@@ -31,6 +31,7 @@ import {LocalUri} from "../sync/SyncDestination";
 import {FileUtils} from "../utils";
 import {BundleCommands} from "../ui/bundle-resource-explorer/BundleCommands";
 import {ConfigModel} from "../configuration/models/ConfigModel";
+import {Telemetry} from "../telemetry";
 
 /**
  * This interface describes the mock-debug specific launch attributes
@@ -58,12 +59,14 @@ export class DatabricksWorkflowDebugAdapterFactory
         private connection: ConnectionManager,
         private configModel: ConfigModel,
         context: ExtensionContext,
-        bundleCommands: BundleCommands
+        bundleCommands: BundleCommands,
+        telemetry: Telemetry
     ) {
         this.workflowRunner = new WorkflowRunner(
             context,
             bundleCommands,
-            connection
+            connection,
+            telemetry
         );
     }
 

packages/databricks-vscode/src/run/RunCommands.ts

@@ -12,6 +12,7 @@ import {
 } from "../utils/shellUtils";
 import {CustomWhenContext} from "../vscode-objs/CustomWhenContext";
 import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
+import {Events, Telemetry} from "../telemetry";
 
 /**
  * Run related commands
@@ -23,7 +24,8 @@ export class RunCommands {
         private readonly pythonExtension: MsPythonExtensionWrapper,
         private readonly featureManager: FeatureManager,
         private readonly context: ExtensionContext,
-        private readonly customWhenContext: CustomWhenContext
+        private readonly customWhenContext: CustomWhenContext,
+        private readonly telemetry: Telemetry
     ) {
         this.context.subscriptions.push(
             window.onDidChangeActiveTextEditor(async () =>
@@ -151,6 +153,11 @@ export class RunCommands {
             this.workspaceFolderManager.activeWorkspaceFolder,
             config
         );
+
+        this.telemetry.recordEvent(Events.DBCONNECT_RUN, {
+            launchType: "debug",
+            computeType: this.connection.serverless ? "serverless" : "cluster",
+        });
     }
 
     async runFileUsingDbconnect(resource?: Uri) {
@@ -180,5 +187,10 @@ export class RunCommands {
                 bootstrapPath
             )} ${escapePathArgument(targetResource.fsPath)}`
         );
+
+        this.telemetry.recordEvent(Events.DBCONNECT_RUN, {
+            launchType: "run",
+            computeType: this.connection.serverless ? "serverless" : "cluster",
+        });
     }
 }

packages/databricks-vscode/src/run/WorkflowRunner.ts

@@ -17,6 +17,8 @@ import Convert from "ansi-to-html";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {WorkspaceFsWorkflowWrapper} from "../workspace-fs/WorkspaceFsWorkflowWrapper";
 import {BundleCommands} from "../ui/bundle-resource-explorer/BundleCommands";
+import {Events, Telemetry} from "../telemetry";
+import {ComputeType, WorkflowTaskType} from "../telemetry/constants";
 
 export class WorkflowRunner implements Disposable {
     private panels = new Map<string, WorkflowOutputPanel>();
@@ -25,7 +27,8 @@ export class WorkflowRunner implements Disposable {
     constructor(
         private context: ExtensionContext,
         private bundleCommands: BundleCommands,
-        private readonly connectionManager: ConnectionManager
+        private readonly connectionManager: ConnectionManager,
+        private readonly telemetry: Telemetry
     ) {}
 
     dispose() {
@@ -133,9 +136,14 @@ export class WorkflowRunner implements Disposable {
             }
         });
 
+        let taskType: WorkflowTaskType = "unknown";
+        const computeType: ComputeType =
+            cluster === undefined ? "serverless" : "cluster";
+        const recordRun = this.telemetry.start(Events.WORKFLOW_RUN);
         try {
             const notebookType = await FileUtils.isNotebook(program);
             if (notebookType) {
+                taskType = "notebook";
                 let remoteFilePath: string =
                     syncDestinationMapper.localToRemoteNotebook(program).path;
                 if (syncDestinationMapper.remoteUri.type === "workspace") {
@@ -168,7 +176,9 @@ export class WorkflowRunner implements Disposable {
                         token: panelCancellation.token,
                     })
                 );
+                recordRun({success: true, taskType, computeType});
             } else {
+                taskType = "python";
                 const originalFileUri =
                     syncDestinationMapper.localToRemote(program);
                 const wrappedFile =
@@ -196,6 +206,7 @@ export class WorkflowRunner implements Disposable {
                 });
                 //TODO: Respone logs will contain bootstrap code path in the error stack trace. Remove it.
                 panel.showStdoutResult(response.logs || "");
+                recordRun({success: true, taskType, computeType});
             }
         } catch (e: unknown) {
             if (e instanceof ApiError) {
@@ -212,6 +223,7 @@ export class WorkflowRunner implements Disposable {
                     message: (e as any).message,
                 });
             }
+            recordRun({success: false, taskType, computeType});
         }
     }
 }

packages/databricks-vscode/src/telemetry/constants.ts

@@ -20,6 +20,9 @@ export enum Events {
     BUNDLE_INIT = "bundleInit",
     BUNDLE_SUB_PROJECTS = "bundleSubProjects",
     CONNECTION_STATE_CHANGED = "connectionStateChanged",
+    COMPUTE_SELECTED = "computeSelected",
+    WORKFLOW_RUN = "workflowRun",
+    DBCONNECT_RUN = "dbconnectRun",
 }
 /* eslint-enable @typescript-eslint/naming-convention */
 
@@ -35,6 +38,9 @@ export type BundleRunType =
     | "validate"
     | "partial-refresh"
     | "manual-input";
+export type WorkflowTaskType = "python" | "notebook" | "unknown";
+export type LaunchType = "run" | "debug";
+export type ComputeType = "cluster" | "serverless";
 
 /** Documentation about all of the properties and metrics of the event. */
 type EventDescription<T> = {[K in keyof T]?: {comment?: string}};
@@ -162,6 +168,38 @@ export class EventTypes {
             comment: "The new state of the connection",
         },
     };
+    [Events.COMPUTE_SELECTED]: EventType<{
+        type: ComputeType;
+    }> = {
+        comment: "A compute was selected",
+        type: {
+            comment: "The type of the compute",
+        },
+    };
+    [Events.WORKFLOW_RUN]: EventType<
+        {
+            success: boolean;
+            taskType: WorkflowTaskType;
+            computeType: ComputeType;
+        } & DurationMeasurement
+    > = {
+        comment: "A workflow task was run",
+        taskType: {
+            comment: "The type of the workflow task",
+        },
+        computeType: {
+            comment: "The type of the compute",
+        },
+    };
+    [Events.DBCONNECT_RUN]: EventType<{
+        launchType: LaunchType;
+        computeType: ComputeType;
+    }> = {
+        comment: "A Databricks Connect debug run",
+        computeType: {
+            comment: "The type of the compute",
+        },
+    };
 }
 
 /**