Restructure Cluster, ExecutionContext and Command sdk (#49)

kartikgupta-db · web-flow · commit 730584ce83c9 · 2022-08-04T10:33:40.000+01:00
diff --git a/packages/databricks-sdk-js/src/apis/executionContext.integ.ts b/packages/databricks-sdk-js/src/apis/executionContext.integ.ts
@@ -21,13 +21,13 @@ describe(__filename, function () {
         let commandsApi = new CommandsService(integSetup.client);
 
         let context = await executionContextApi.create({
-            clusterId: integSetup.clusterId,
+            clusterId: integSetup.cluster.id,
             language: "python",
         });
         //console.log("Execution context", context);
 
         let command = await commandsApi.execute({
-            clusterId: integSetup.clusterId,
+            clusterId: integSetup.cluster.id,
             contextId: context.id,
             language: "python",
             command: "print('juhu')",
@@ -38,7 +38,7 @@ describe(__filename, function () {
         while (true) {
             await sleep(3000);
             status = await commandsApi.status({
-                clusterId: integSetup.clusterId,
+                clusterId: integSetup.cluster.id,
                 contextId: context.id,
                 commandId: command.id,
             });
@@ -55,7 +55,7 @@ describe(__filename, function () {
         assert.equal(status.results.data, "juhu");
 
         await executionContextApi.destroy({
-            clusterId: integSetup.clusterId,
+            clusterId: integSetup.cluster.id,
             contextId: context.id,
         });
     });
diff --git a/packages/databricks-sdk-js/src/apis/jobs.integ.ts b/packages/databricks-sdk-js/src/apis/jobs.integ.ts
@@ -32,7 +32,7 @@ describe(__filename, function () {
             tasks: [
                 {
                     task_key: "hello_world",
-                    existing_cluster_id: integSetup.clusterId,
+                    existing_cluster_id: integSetup.cluster.id,
                     spark_python_task: {
                         python_file: `dbfs:${jobPath}`,
                     },
diff --git a/packages/databricks-sdk-js/src/services/Cluster.integ.ts b/packages/databricks-sdk-js/src/services/Cluster.integ.ts
@@ -15,34 +15,25 @@ describe(__filename, function () {
     });
 
     it("should create an execution context", async () => {
-        let cluster = await Cluster.fromClusterId(
-            integSetup.client,
-            integSetup.clusterId
-        );
-
-        await cluster.canExecute();
+        await integSetup.cluster.canExecute();
 
-        let ctx = await cluster.createExecutioncontext();
-        let command = await ctx.execute("print('hello')");
-        let result = await command.response();
+        let ctx = await integSetup.cluster.createExecutionContext();
+        let {result} = await ctx.execute("print('hello')");
 
         assert(result.results);
         assert(result.results.resultType === "text");
         assert.equal(result.results.data, "hello");
     });
 
     it("should load a cluster by name", async () => {
-        let clusterA = await Cluster.fromClusterId(
-            integSetup.client,
-            integSetup.clusterId
-        );
+        let clusterA = integSetup.cluster;
 
         let clusterB = await Cluster.fromClusterName(
             integSetup.client,
             clusterA.details.cluster_name!
         );
 
-        assert(clusterA.details.cluster_id);
-        assert.equal(clusterA.details.cluster_id, clusterB?.details.cluster_id);
+        assert(clusterA.id);
+        assert.equal(clusterA.id, clusterB?.id);
     });
 });
diff --git a/packages/databricks-sdk-js/src/services/Cluster.ts b/packages/databricks-sdk-js/src/services/Cluster.ts
@@ -100,19 +100,15 @@ export class Cluster {
         }
     }
 
-    async createExecutioncontext(): Promise<ExecutionContext> {
-        return await ExecutionContext.create(
-            this.client,
-            this.clusterDetails.cluster_id!
-        );
+    async createExecutionContext(): Promise<ExecutionContext> {
+        return await ExecutionContext.create(this.client, this);
     }
 
     async canExecute(): Promise<boolean> {
         let context: ExecutionContext | undefined;
         try {
-            context = await this.createExecutioncontext();
-            let command = await context.execute("print('hello')");
-            await command.response();
+            context = await this.createExecutionContext();
+            await context.execute("print('hello')");
             return true;
         } catch (e) {
             return false;
diff --git a/packages/databricks-sdk-js/src/services/Command.ts b/packages/databricks-sdk-js/src/services/Command.ts
@@ -1,13 +1,21 @@
 import EventEmitter from "node:events";
 import {CommandsService, CommandsStatusResponse} from "../apis/commands";
-import {Language} from "../apis/executionContext";
 import {ExecutionContext} from "./ExecutionContext";
 
+export interface CommandWithResult {
+    cmd: Command;
+    result: CommandsStatusResponse;
+}
+
+export type StatusUpdateListener = (result: CommandsStatusResponse) => void;
+
 export class Command extends EventEmitter {
     readonly context: ExecutionContext;
     readonly commandsApi: CommandsService;
     id?: string;
 
+    private static statusUpdateEvent: string = "statusUpdate";
+
     constructor(context: ExecutionContext) {
         super();
         this.context = context;
@@ -17,12 +25,12 @@ export class Command extends EventEmitter {
     async response(): Promise<CommandsStatusResponse> {
         while (true) {
             let result = await this.commandsApi.status({
-                clusterId: this.context.clusterId,
+                clusterId: this.context.cluster.id,
                 contextId: this.context.id!,
                 commandId: this.id!,
             });
 
-            this.emit("statusUpdate", result);
+            this.emit(Command.statusUpdateEvent, result);
 
             if (
                 result.status === "Cancelled" ||
@@ -41,19 +49,23 @@ export class Command extends EventEmitter {
     static async execute(
         context: ExecutionContext,
         command: string,
-        language: Language = "python"
-    ): Promise<Command> {
+        onStatusUpdate: StatusUpdateListener = () => {}
+    ): Promise<CommandWithResult> {
         //console.log(`Executing (${language}): ${command}`);
         let cmd = new Command(context);
 
+        cmd.on(Command.statusUpdateEvent, onStatusUpdate);
+
         let result = await cmd.commandsApi.execute({
-            clusterId: cmd.context.clusterId,
+            clusterId: cmd.context.cluster.id,
             contextId: cmd.context.id!,
-            language,
+            language: cmd.context.language,
             command,
         });
 
         cmd.id = result.id;
-        return cmd;
+        const executionResult = await cmd.response();
+
+        return {cmd: cmd, result: executionResult};
     }
 }
diff --git a/packages/databricks-sdk-js/src/services/ExecutionContext.integ.ts b/packages/databricks-sdk-js/src/services/ExecutionContext.integ.ts
@@ -17,20 +17,20 @@ describe(__filename, function () {
     it("should run python with high level API", async () => {
         let context = await ExecutionContext.create(
             integSetup.client,
-            integSetup.clusterId
+            integSetup.cluster
         );
 
-        var command = await context.execute("print('juhu')");
-
         let statusUpdateCalled = false;
-        command.on("statusUpdate", (e) => (statusUpdateCalled = true));
-
-        var result = await command.response();
-        // console.log(result);
+        var {cmd, result} = await context.execute(
+            "print('juhu')",
+            (e) => (statusUpdateCalled = true)
+        );
 
-        command = await context.execute("print('kinners')");
-        result = await command.response();
-        // console.log(result);
+        statusUpdateCalled = false;
+        ({cmd, result} = await context.execute(
+            "print('kinners')",
+            (e) => (statusUpdateCalled = true)
+        ));
 
         assert(statusUpdateCalled);
         assert(result.results);
diff --git a/packages/databricks-sdk-js/src/services/ExecutionContext.ts b/packages/databricks-sdk-js/src/services/ExecutionContext.ts
@@ -1,27 +1,30 @@
+import {cluster} from "..";
 import {ApiClient} from "../api-client";
 import {ExecutionContextService, Language} from "../apis/executionContext";
-import {Command} from "./Command";
+import {Cluster} from "./Cluster";
+import {Command, CommandWithResult, StatusUpdateListener} from "./Command";
 
 export class ExecutionContext {
     readonly executionContextApi: ExecutionContextService;
-    readonly clusterId: string;
-    readonly client: ApiClient;
     id?: string;
 
-    constructor(client: ApiClient, clusterId: string) {
-        this.client = client;
+    constructor(
+        readonly client: ApiClient,
+        readonly cluster: Cluster,
+        readonly language: Language
+    ) {
         this.executionContextApi = new ExecutionContextService(client);
-        this.clusterId = clusterId;
     }
 
     static async create(
         client: ApiClient,
-        clusterId: string
+        cluster: Cluster,
+        language: Language = "python"
     ): Promise<ExecutionContext> {
-        let context = new ExecutionContext(client, clusterId);
+        let context = new ExecutionContext(client, cluster, language);
         let response = await context.executionContextApi.create({
-            clusterId: context.clusterId,
-            language: "python",
+            clusterId: context.cluster.id,
+            language: context.language,
         });
 
         context.id = response.id;
@@ -30,9 +33,9 @@ export class ExecutionContext {
 
     async execute(
         command: string,
-        language: Language = "python"
-    ): Promise<Command> {
-        return await Command.execute(this, command, language);
+        onStatusUpdate: StatusUpdateListener = () => {}
+    ): Promise<CommandWithResult> {
+        return await Command.execute(this, command, onStatusUpdate);
     }
 
     async destroy() {
@@ -41,7 +44,7 @@ export class ExecutionContext {
         }
 
         await this.executionContextApi.destroy({
-            clusterId: this.clusterId,
+            clusterId: this.cluster.id,
             contextId: this.id,
         });
     }
diff --git a/packages/databricks-sdk-js/src/services/WorkflowRun.integ.ts b/packages/databricks-sdk-js/src/services/WorkflowRun.integ.ts
@@ -15,7 +15,7 @@ describe(__filename, function () {
     it("should run a python job", async () => {
         let cluster = await Cluster.fromClusterId(
             integSetup.client,
-            integSetup.clusterId
+            integSetup.cluster.id
         );
 
         let dbfsApi = new DbfsService(integSetup.client);
@@ -55,7 +55,7 @@ describe(__filename, function () {
     it("should run a notebook job", async () => {
         let cluster = await Cluster.fromClusterId(
             integSetup.client,
-            integSetup.clusterId
+            integSetup.cluster.id
         );
 
         let jobPath = `/tmp/sdk-js-integ-${integSetup.testRunId}.py`;
diff --git a/packages/databricks-sdk-js/src/test/IntegrationTestSetup.ts b/packages/databricks-sdk-js/src/test/IntegrationTestSetup.ts
@@ -3,19 +3,19 @@
 import {v4 as uuidv4} from "uuid";
 import {ApiClient} from "../api-client";
 import {ClusterService} from "../apis/cluster";
+import {Cluster} from "../services/Cluster";
 
 export class IntegrationTestSetup {
     readonly testRunId: string;
 
-    constructor(readonly client: ApiClient, readonly clusterId: string) {
+    constructor(readonly client: ApiClient, readonly cluster: Cluster) {
         this.testRunId = uuidv4();
     }
 
     private static _instance: IntegrationTestSetup;
     static async getInstance(): Promise<IntegrationTestSetup> {
         if (!this._instance) {
             let client = new ApiClient();
-            let clustersApi = new ClusterService(client);
 
             if (!process.env["TEST_DEFAULT_CLUSTER_ID"]) {
                 throw new Error(
@@ -25,20 +25,13 @@ export class IntegrationTestSetup {
 
             const clusterId =
                 process.env["TEST_DEFAULT_CLUSTER_ID"]!.split("'").join("");
-            clustersApi.start({cluster_id: clusterId});
 
+            const cluster = await Cluster.fromClusterId(client, clusterId);
+            cluster.start();
             // wait for cluster to be running
-            while (true) {
-                let cluster = await clustersApi.get({cluster_id: clusterId});
-                if (cluster.state === "RUNNING") {
-                    break;
-                }
-
-                await new Promise((resolve) => {
-                    setTimeout(resolve, 1000);
-                });
-            }
-            this._instance = new IntegrationTestSetup(client, clusterId);
+            await cluster.waitForState("RUNNING");
+
+            this._instance = new IntegrationTestSetup(client, cluster);
         }
         return this._instance;
     }

Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ describe(__filename, function () {`
`32`	`32`	`tasks: [`
`33`	`33`	`{`
`34`	`34`	`task_key: "hello_world",`
`35`		`- existing_cluster_id: integSetup.clusterId,`
	`35`	`+ existing_cluster_id: integSetup.cluster.id,`
`36`	`36`	`spark_python_task: {`
`37`	`37`	python_file: `dbfs:${jobPath}`,
`38`	`38`	`},`