databricks-cli/bundle/config/mutator/trampoline.go

package mutator

import (
	"context"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"text/template"

	"github.com/databricks/cli/bundle"
	"github.com/databricks/databricks-sdk-go/service/jobs"
)

type TaskWithJobKey struct {
	Task   *jobs.Task
	JobKey string
}

type TrampolineFunctions interface {
	GetTemplateData(task *jobs.Task) (map[string]any, error)
	GetTasks(b *bundle.Bundle) []TaskWithJobKey
	CleanUp(task *jobs.Task) error
}
type trampoline struct {
	name      string
	functions TrampolineFunctions
	template  string
}

func NewTrampoline(
	name string,
	functions TrampolineFunctions,
	template string,
) *trampoline {
	return &trampoline{name, functions, template}
}

func (m *trampoline) Name() string {
	return fmt.Sprintf("trampoline(%s)", m.name)
}

func (m *trampoline) Apply(ctx context.Context, b *bundle.Bundle) error {
	tasks := m.functions.GetTasks(b)
	for _, task := range tasks {
		err := m.generateNotebookWrapper(ctx, b, task)
		if err != nil {
			return err
		}
	}
	return nil
}

func (m *trampoline) generateNotebookWrapper(ctx context.Context, b *bundle.Bundle, task TaskWithJobKey) error {
	internalDir, err := b.InternalDir(ctx)
	if err != nil {
		return err
	}

	notebookName := fmt.Sprintf("notebook_%s_%s", task.JobKey, task.Task.TaskKey)
	localNotebookPath := filepath.Join(internalDir, notebookName+".py")

	err = os.MkdirAll(filepath.Dir(localNotebookPath), 0755)
	if err != nil {
		return err
	}

	f, err := os.Create(localNotebookPath)
	if err != nil {
		return err
	}
	defer f.Close()

	data, err := m.functions.GetTemplateData(task.Task)
	if err != nil {
		return err
	}

	t, err := template.New(notebookName).Parse(m.template)
	if err != nil {
		return err
	}

	internalDirRel, err := filepath.Rel(b.Config.Path, internalDir)
	if err != nil {
		return err
	}

	err = m.functions.CleanUp(task.Task)
	if err != nil {
		return err
	}
	remotePath := path.Join(b.Config.Workspace.FilesPath, filepath.ToSlash(internalDirRel), notebookName)

	task.Task.NotebookTask = &jobs.NotebookTask{
		NotebookPath: remotePath,
	}

	return t.Execute(f, data)
}