mirror of https://github.com/databricks/cli.git
96 lines
1.9 KiB
Go
96 lines
1.9 KiB
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 fnTemplateData func(task *jobs.Task) (map[string]any, error)
|
|
type fnCleanUp func(task *jobs.Task)
|
|
type fnTasks func(b *bundle.Bundle) []*jobs.Task
|
|
|
|
type trampoline struct {
|
|
name string
|
|
getTasks fnTasks
|
|
templateData fnTemplateData
|
|
cleanUp fnCleanUp
|
|
template string
|
|
}
|
|
|
|
func NewTrampoline(
|
|
name string,
|
|
tasks fnTasks,
|
|
templateData fnTemplateData,
|
|
cleanUp fnCleanUp,
|
|
template string,
|
|
) *trampoline {
|
|
return &trampoline{name, tasks, templateData, cleanUp, 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.getTasks(b)
|
|
for _, task := range tasks {
|
|
err := m.generateNotebookWrapper(b, task)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *trampoline) generateNotebookWrapper(b *bundle.Bundle, task *jobs.Task) error {
|
|
internalDir, err := b.InternalDir()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
notebookName := fmt.Sprintf("notebook_%s", 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.templateData(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
|
|
}
|
|
|
|
m.cleanUp(task)
|
|
remotePath := path.Join(b.Config.Workspace.FilesPath, filepath.ToSlash(internalDirRel), notebookName)
|
|
|
|
task.NotebookTask = &jobs.NotebookTask{
|
|
NotebookPath: remotePath,
|
|
}
|
|
|
|
return t.Execute(f, data)
|
|
}
|