Remove API timeout limit for execution context runs (#487)

Fixes #482

---------

Co-authored-by: Fabian Jakobs <fabian.jakobs@databricks.com>

Author: kartikgupta-db

packages/databricks-sdk-js/src/services/Command.ts

@@ -1,8 +1,9 @@
 import {EventEmitter} from "events";
-import retry, {RetriableError} from "../retries/retries";
+import retry, {DEFAULT_MAX_TIMEOUT, RetriableError} from "../retries/retries";
 import {ExecutionContext} from "./ExecutionContext";
 import {CancellationToken} from "../types";
 import {CommandExecutionService, CommandStatusResponse} from "../apis/commands";
+import Time from "../retries/Time";
 
 interface CommandErrorParams {
     commandId: string;
@@ -109,9 +110,11 @@ export class Command extends EventEmitter {
     }
 
     async response(
-        cancellationToken?: CancellationToken
+        cancellationToken?: CancellationToken,
+        timeout: Time = DEFAULT_MAX_TIMEOUT
     ): Promise<CommandStatusResponse> {
         await retry({
+            timeout: timeout,
             fn: async () => {
                 await this.refresh();
 
@@ -143,7 +146,8 @@ export class Command extends EventEmitter {
         context: ExecutionContext,
         command: string,
         onStatusUpdate: StatusUpdateListener = () => {},
-        cancellationToken?: CancellationToken
+        cancellationToken?: CancellationToken,
+        timeout: Time = DEFAULT_MAX_TIMEOUT
     ): Promise<CommandWithResult> {
         const cmd = new Command(context);
 
@@ -158,7 +162,7 @@ export class Command extends EventEmitter {
 
         cmd.id = executeApiResponse.id;
 
-        const executionResult = await cmd.response(cancellationToken);
+        const executionResult = await cmd.response(cancellationToken, timeout);
 
         return {cmd: cmd, result: executionResult};
     }

packages/databricks-sdk-js/src/services/ExecutionContext.integ.ts

@@ -6,6 +6,7 @@ import assert from "assert";
 import {IntegrationTestSetup} from "../test/IntegrationTestSetup";
 import {mock, when, instance} from "ts-mockito";
 import {TokenFixture} from "../test/fixtures/TokenFixtures";
+import {DEFAULT_MAX_TIMEOUT} from "../retries/retries";
 
 describe(__filename, function () {
     let integSetup: IntegrationTestSetup;
@@ -56,7 +57,8 @@ describe(__filename, function () {
         const {cmd, result} = await context.execute(
             "while True: pass",
             undefined,
-            instance(token)
+            instance(token),
+            DEFAULT_MAX_TIMEOUT
         );
         // The API surfaces an exception when a command is cancelled
         // The cancellation itself proceeds as expected, but the status

packages/databricks-sdk-js/src/services/ExecutionContext.ts

@@ -1,6 +1,7 @@
-import {CancellationToken} from "..";
+import {CancellationToken, Time} from "..";
 import {ApiClient} from "../api-client";
 import {CommandExecutionService, Language} from "../apis/commands";
+import {DEFAULT_MAX_TIMEOUT} from "../retries/retries";
 import {Cluster} from "./Cluster";
 import {Command, CommandWithResult, StatusUpdateListener} from "./Command";
 
@@ -34,9 +35,16 @@ export class ExecutionContext {
     async execute(
         command: string,
         onStatusUpdate: StatusUpdateListener = () => {},
-        token?: CancellationToken
+        token?: CancellationToken,
+        timeout: Time = DEFAULT_MAX_TIMEOUT
     ): Promise<CommandWithResult> {
-        return await Command.execute(this, command, onStatusUpdate, token);
+        return await Command.execute(
+            this,
+            command,
+            onStatusUpdate,
+            token,
+            timeout
+        );
     }
 
     async destroy() {

packages/databricks-vscode/src/run/DatabricksRuntime.test.ts

@@ -28,7 +28,12 @@ describe(__filename, () => {
 
         executionContextMock = mock(ExecutionContext);
         when(
-            executionContextMock.execute(anything(), anything(), anything())
+            executionContextMock.execute(
+                anything(),
+                anything(),
+                anything(),
+                anything()
+            )
         ).thenResolve({
             cmd: mock(Command),
             result: {
@@ -122,7 +127,12 @@ describe(__filename, () => {
 
     it("should handle failed executions", async () => {
         when(
-            executionContextMock.execute(anything(), anything(), anything())
+            executionContextMock.execute(
+                anything(),
+                anything(),
+                anything(),
+                anything()
+            )
         ).thenResolve({
             cmd: mock(Command),
             result: {

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

@@ -25,6 +25,7 @@ import {CodeSynchronizer} from "../sync/CodeSynchronizer";
 import * as fs from "node:fs/promises";
 import {parseErrorResult} from "./ErrorParser";
 import path from "node:path";
+import {Time, TimeUnits} from "@databricks/databricks-sdk";
 
 export interface OutputEvent {
     type: "prio" | "out" | "err";
@@ -151,6 +152,7 @@ export class DatabricksRuntime implements Disposable {
             // We wait for sync to complete so that the local files are consistant
             // with the remote repo files
             await this.codeSynchronizer.waitForSyncComplete();
+            await commands.executeCommand("workbench.panel.repl.view.focus");
 
             log(
                 `Running ${syncDestination.localUri.relativePath(
@@ -166,7 +168,8 @@ export class DatabricksRuntime implements Disposable {
                     envVars
                 ),
                 undefined,
-                this.token
+                this.token,
+                new Time(240, TimeUnits.hours)
             );
             const result = response.result;