Support Databricks WebAPI 2.1 version and Support existing_cluster_id and new_cluster options to create a Job

flyteplugins/go/tasks/plugins/webapi/databricks/integration_test.go

@@ -44,7 +44,7 @@ func TestEndToEnd(t *testing.T) {
 	plugin, err := pluginEntry.LoadPlugin(context.TODO(), newFakeSetupContext())
 	assert.NoError(t, err)
 
-	t.Run("run a databricks job", func(t *testing.T) {
+	t.Run("run a databricks job by new_cluster key", func(t *testing.T) {
 		databricksConfDict := map[string]interface{}{
 			"name": "flytekit databricks plugin example",
 			"new_cluster": map[string]string{
@@ -75,6 +75,34 @@ func TestEndToEnd(t *testing.T) {
 		phase := tests.RunPluginEndToEndTest(t, plugin, &template, inputs, nil, nil, iter)
 		assert.Equal(t, true, phase.Phase().IsSuccess())
 	})
+
+	t.Run("run a databricks job by new_cluster key", func(t *testing.T) {
+		databricksConfDict := map[string]interface{}{
+			"name":                "flytekit databricks plugin example",
+			"existing_cluster_id": "1201-my-cluster",
+			"timeout_seconds":     3600,
+			"max_retries":         1,
+		}
+		databricksConfig, err := utils.MarshalObjToStruct(databricksConfDict)
+		assert.NoError(t, err)
+		sparkJob := plugins.SparkJob{DatabricksConf: databricksConfig, DatabricksToken: "token", SparkConf: map[string]string{"spark.driver.bindAddress": "127.0.0.1"}}
+		st, err := utils.MarshalPbToStruct(&sparkJob)
+		assert.NoError(t, err)
+		inputs, _ := coreutils.MakeLiteralMap(map[string]interface{}{"x": 1})
+		template := flyteIdlCore.TaskTemplate{
+			Type:   "databricks",
+			Custom: st,
+			Target: &coreIdl.TaskTemplate_Container{
+				Container: &coreIdl.Container{
+					Command: []string{"command"},
+					Args:    []string{"pyflyte-execute"},
+				},
+			},
+		}
+
+		phase := tests.RunPluginEndToEndTest(t, plugin, &template, inputs, nil, nil, iter)
+		assert.Equal(t, true, phase.Phase().IsSuccess())
+	})
 }
 
 func newFakeDatabricksServer() *httptest.Server {

flyteplugins/go/tasks/plugins/webapi/databricks/plugin.go

@@ -29,7 +29,7 @@ const (
 	ErrSystem       errors.ErrorCode = "System"
 	post            string           = "POST"
 	get             string           = "GET"
-	databricksAPI   string           = "/api/2.0/jobs/runs"
+	databricksAPI   string           = "/api/2.1/jobs/runs"
 	newCluster      string           = "new_cluster"
 	dockerImage     string           = "docker_image"
 	sparkConfig     string           = "spark_conf"
@@ -114,10 +114,15 @@ func (p Plugin) Create(ctx context.Context, taskCtx webapi.TaskExecutionContextR
 		return nil, nil, fmt.Errorf("failed to unmarshal databricksJob: %v: %v", sparkJob.DatabricksConf, err)
 	}
 
-	if _, ok := databricksJob[newCluster]; ok {
-		databricksJob[newCluster].(map[string]interface{})[dockerImage] = map[string]string{url: container.Image}
-		if len(sparkJob.SparkConf) != 0 {
-			databricksJob[newCluster].(map[string]interface{})[sparkConfig] = sparkJob.SparkConf
+	// If "existing_cluster_id" is in databricks_job, then we don't need to set "new_cluster"
+	// Refer the docs here: https://docs.databricks.com/en/workflows/jobs/jobs-2.0-api.html#request-structure
+	if clusterConfig, ok := databricksJob[newCluster].(map[string]interface{}); ok {
+		if dockerConfig, ok := clusterConfig[dockerImage].(map[string]interface{}); !ok || dockerConfig[url] == nil {
+			clusterConfig[dockerImage] = map[string]string{url: container.Image}
+		}
+
+		if clusterConfig[sparkConfig] == nil && len(sparkJob.SparkConf) != 0 {
+			clusterConfig[sparkConfig] = sparkJob.SparkConf
 		}
 	}
 	databricksJob[sparkPythonTask] = map[string]interface{}{pythonFile: p.cfg.EntrypointFile, parameters: modifiedArgs}
@@ -222,7 +227,7 @@ func (p Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phase
 	case http.StatusAccepted:
 		return core.PhaseInfoRunning(pluginsCore.DefaultPhaseVersion, taskInfo), nil
 	case http.StatusOK:
-		if lifeCycleState == "TERMINATED" || lifeCycleState == "SKIPPED" || lifeCycleState == "INTERNAL_ERROR" {
+		if lifeCycleState == "TERMINATED" || lifeCycleState == "TERMINATING" || lifeCycleState == "INTERNAL_ERROR" {
 			if resultState == "SUCCESS" {
 				if err := writeOutput(ctx, taskCtx); err != nil {
 					pluginsCore.PhaseInfoFailure(string(rune(statusCode)), "failed to write output", taskInfo)
@@ -231,6 +236,11 @@ func (p Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phase
 			}
 			return pluginsCore.PhaseInfoFailure(string(rune(statusCode)), message, taskInfo), nil
 		}
+
+		if lifeCycleState == "PENDING" {
+			return core.PhaseInfoInitializing(time.Now(), core.DefaultPhaseVersion, message, taskInfo), nil
+		}
+
 		return core.PhaseInfoRunning(pluginsCore.DefaultPhaseVersion, taskInfo), nil
 	case http.StatusBadRequest:
 		fallthrough
@@ -278,7 +288,7 @@ func buildRequest(
 	var err error
 	if isCancel {
 		databricksURL += "/cancel"
-		data = []byte(fmt.Sprintf("{ run_id: %v }", runID))
+		data = []byte(fmt.Sprintf("{ \"run_id\": %v }", runID))
 	} else if method == post {
 		databricksURL += "/submit"
 		mJSON, err := json.Marshal(databricksJob)

flyteplugins/go/tasks/plugins/webapi/databricks/plugin_test.go

@@ -67,7 +67,7 @@ func TestBuildRequest(t *testing.T) {
 	token := "test-token"
 	runID := "019e70eb"
 	databricksEndpoint := ""
-	databricksURL := "https://" + testInstance + "/api/2.0/jobs/runs"
+	databricksURL := "https://" + testInstance + "/api/2.1/jobs/runs"
 	t.Run("build http request for submitting a databricks job", func(t *testing.T) {
 		req, err := buildRequest(post, nil, databricksEndpoint, testInstance, token, runID, false)
 		header := http.Header{}