Skip to content

Commit e474948

Browse files
andrewnesterandrewnester
authored
Generate correct YAML if custom_tags or spark_conf is used for pipeline or job cluster configuration (#1210)
These fields (key and values) needs to be double quoted in order for yaml loader to read, parse and unmarshal it into Go struct correctly because these fields are `map[string]string` type. ## Tests Added regression unit and E2E tests
1 parent 299e9b5 commit e474948

7 files changed

Lines changed: 292 additions & 39 deletions

File tree

‎cmd/bundle/generate/generate_test.go‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import (
1212
"github.com/databricks/cli/bundle"
1313
"github.com/databricks/cli/bundle/config"
1414
"github.com/databricks/databricks-sdk-go/experimental/mocks"
15+
"github.com/databricks/databricks-sdk-go/service/compute"
16+
"github.com/databricks/databricks-sdk-go/service/jobs"
1517
"github.com/databricks/databricks-sdk-go/service/pipelines"
1618
"github.com/databricks/databricks-sdk-go/service/workspace"
1719
"github.com/stretchr/testify/mock"
@@ -36,6 +38,18 @@ func TestGeneratePipelineCommand(t *testing.T) {
3638
Name: "test-pipeline",
3739
Spec: &pipelines.PipelineSpec{
3840
Name: "test-pipeline",
41+
Clusters: []pipelines.PipelineCluster{
42+
{
43+
CustomTags: map[string]string{
44+
"Tag1": "24X7-1234",
45+
},
46+
},
47+
{
48+
SparkConf: map[string]string{
49+
"spark.databricks.delta.preview.enabled": "true",
50+
},
51+
},
52+
},
3953
Libraries: []pipelines.PipelineLibrary{
4054
{Notebook: &pipelines.NotebookLibrary{
4155
Path: "/test/notebook",
@@ -85,6 +99,11 @@ func TestGeneratePipelineCommand(t *testing.T) {
8599
pipelines:
86100
test_pipeline:
87101
name: test-pipeline
102+
clusters:
103+
- custom_tags:
104+
"Tag1": "24X7-1234"
105+
- spark_conf:
106+
"spark.databricks.delta.preview.enabled": "true"
88107
libraries:
89108
- notebook:
90109
path: %s
@@ -100,3 +119,93 @@ func TestGeneratePipelineCommand(t *testing.T) {
100119
require.NoError(t, err)
101120
require.Equal(t, "Py content", string(data))
102121
}
122+
123+
func TestGenerateJobCommand(t *testing.T) {
124+
cmd := NewGenerateJobCommand()
125+
126+
root := t.TempDir()
127+
b := &bundle.Bundle{
128+
Config: config.Root{
129+
Path: root,
130+
},
131+
}
132+
133+
m := mocks.NewMockWorkspaceClient(t)
134+
b.SetWorkpaceClient(m.WorkspaceClient)
135+
136+
jobsApi := m.GetMockJobsAPI()
137+
jobsApi.EXPECT().Get(mock.Anything, jobs.GetJobRequest{JobId: 1234}).Return(&jobs.Job{
138+
Settings: &jobs.JobSettings{
139+
Name: "test-job",
140+
JobClusters: []jobs.JobCluster{
141+
{NewCluster: &compute.ClusterSpec{
142+
CustomTags: map[string]string{
143+
"Tag1": "24X7-1234",
144+
},
145+
}},
146+
{NewCluster: &compute.ClusterSpec{
147+
SparkConf: map[string]string{
148+
"spark.databricks.delta.preview.enabled": "true",
149+
},
150+
}},
151+
},
152+
Tasks: []jobs.Task{
153+
{
154+
TaskKey: "notebook_task",
155+
NotebookTask: &jobs.NotebookTask{
156+
NotebookPath: "/test/notebook",
157+
},
158+
},
159+
},
160+
},
161+
}, nil)
162+
163+
workspaceApi := m.GetMockWorkspaceAPI()
164+
workspaceApi.EXPECT().GetStatusByPath(mock.Anything, "/test/notebook").Return(&workspace.ObjectInfo{
165+
ObjectType: workspace.ObjectTypeNotebook,
166+
Language: workspace.LanguagePython,
167+
Path: "/test/notebook",
168+
}, nil)
169+
170+
notebookContent := io.NopCloser(bytes.NewBufferString("# Databricks notebook source\nNotebook content"))
171+
workspaceApi.EXPECT().Download(mock.Anything, "/test/notebook", mock.Anything).Return(notebookContent, nil)
172+
173+
cmd.SetContext(bundle.Context(context.Background(), b))
174+
cmd.Flag("existing-job-id").Value.Set("1234")
175+
176+
configDir := filepath.Join(root, "resources")
177+
cmd.Flag("config-dir").Value.Set(configDir)
178+
179+
srcDir := filepath.Join(root, "src")
180+
cmd.Flag("source-dir").Value.Set(srcDir)
181+
182+
var key string
183+
cmd.Flags().StringVar(&key, "key", "test_job", "")
184+
185+
err := cmd.RunE(cmd, []string{})
186+
require.NoError(t, err)
187+
188+
data, err := os.ReadFile(filepath.Join(configDir, "test_job.yml"))
189+
require.NoError(t, err)
190+
191+
require.Equal(t, fmt.Sprintf(`resources:
192+
jobs:
193+
test_job:
194+
name: test-job
195+
job_clusters:
196+
- new_cluster:
197+
custom_tags:
198+
"Tag1": "24X7-1234"
199+
- new_cluster:
200+
spark_conf:
201+
"spark.databricks.delta.preview.enabled": "true"
202+
tasks:
203+
- task_key: notebook_task
204+
notebook_task:
205+
notebook_path: %s
206+
`, filepath.Join("..", "src", "notebook.py")), string(data))
207+
208+
data, err = os.ReadFile(filepath.Join(srcDir, "notebook.py"))
209+
require.NoError(t, err)
210+
require.Equal(t, "# Databricks notebook source\nNotebook content", string(data))
211+
}

‎cmd/bundle/generate/job.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"github.com/databricks/cli/libs/textutil"
1515
"github.com/databricks/databricks-sdk-go/service/jobs"
1616
"github.com/spf13/cobra"
17+
"gopkg.in/yaml.v3"
1718
)
1819

1920
func NewGenerateJobCommand() *cobra.Command {
@@ -82,7 +83,13 @@ func NewGenerateJobCommand() *cobra.Command {
8283
}
8384

8485
filename := filepath.Join(configDir, fmt.Sprintf("%s.yml", jobKey))
85-
err = yamlsaver.SaveAsYAML(result, filename, force)
86+
saver := yamlsaver.NewSaverWithStyle(map[string]yaml.Style{
87+
// Including all JobSettings and nested fields which are map[string]string type
88+
"spark_conf": yaml.DoubleQuotedStyle,
89+
"custom_tags": yaml.DoubleQuotedStyle,
90+
"tags": yaml.DoubleQuotedStyle,
91+
})
92+
err = saver.SaveAsYAML(result, filename, force)
8693
if err != nil {
8794
return err
8895
}

‎cmd/bundle/generate/pipeline.go‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"github.com/databricks/cli/libs/textutil"
1515
"github.com/databricks/databricks-sdk-go/service/pipelines"
1616
"github.com/spf13/cobra"
17+
"gopkg.in/yaml.v3"
1718
)
1819

1920
func NewGeneratePipelineCommand() *cobra.Command {
@@ -82,7 +83,15 @@ func NewGeneratePipelineCommand() *cobra.Command {
8283
}
8384

8485
filename := filepath.Join(configDir, fmt.Sprintf("%s.yml", pipelineKey))
85-
err = yamlsaver.SaveAsYAML(result, filename, force)
86+
saver := yamlsaver.NewSaverWithStyle(
87+
// Including all PipelineSpec and nested fields which are map[string]string type
88+
map[string]yaml.Style{
89+
"spark_conf": yaml.DoubleQuotedStyle,
90+
"custom_tags": yaml.DoubleQuotedStyle,
91+
"configuration": yaml.DoubleQuotedStyle,
92+
},
93+
)
94+
err = saver.SaveAsYAML(result, filename, force)
8695
if err != nil {
8796
return err
8897
}

‎internal/bundle/generate_job_test.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,11 @@ func (gt *generateJobTest) createTestJob(ctx context.Context) int64 {
103103
SparkVersion: "13.3.x-scala2.12",
104104
NumWorkers: 1,
105105
NodeTypeId: nodeTypeId,
106+
SparkConf: map[string]string{
107+
"spark.databricks.enableWsfs": "true",
108+
"spark.databricks.hive.metastore.glueCatalog.enabled": "true",
109+
"spark.databricks.pip.ignoreSSL": "true",
110+
},
106111
},
107112
NotebookTask: &jobs.NotebookTask{
108113
NotebookPath: path.Join(tmpdir, "test"),

‎internal/bundle/generate_pipeline_test.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,9 @@ func (gt *generatePipelineTest) createTestPipeline(ctx context.Context) (string,
9494
err = f.Write(ctx, "test.py", strings.NewReader("print('Hello!')"))
9595
require.NoError(t, err)
9696

97+
env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV")
98+
nodeTypeId := internal.GetNodeTypeId(env)
99+
97100
name := internal.RandomName("generated-pipeline-")
98101
resp, err := w.Pipelines.Create(ctx, pipelines.CreatePipeline{
99102
Name: name,
@@ -109,6 +112,22 @@ func (gt *generatePipelineTest) createTestPipeline(ctx context.Context) (string,
109112
},
110113
},
111114
},
115+
Clusters: []pipelines.PipelineCluster{
116+
{
117+
CustomTags: map[string]string{
118+
"Tag1": "Yes",
119+
"Tag2": "24X7",
120+
"Tag3": "APP-1234",
121+
},
122+
NodeTypeId: nodeTypeId,
123+
NumWorkers: 2,
124+
SparkConf: map[string]string{
125+
"spark.databricks.enableWsfs": "true",
126+
"spark.databricks.hive.metastore.glueCatalog.enabled": "true",
127+
"spark.databricks.pip.ignoreSSL": "true",
128+
},
129+
},
130+
},
112131
})
113132
require.NoError(t, err)
114133

‎libs/dyn/yamlsaver/saver.go‎

Lines changed: 47 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,21 @@ import (
1313
"gopkg.in/yaml.v3"
1414
)
1515

16-
func SaveAsYAML(data any, filename string, force bool) error {
16+
type saver struct {
17+
nodesWithStyle map[string]yaml.Style
18+
}
19+
20+
func NewSaver() *saver {
21+
return &saver{}
22+
}
23+
24+
func NewSaverWithStyle(nodesWithStyle map[string]yaml.Style) *saver {
25+
return &saver{
26+
nodesWithStyle: nodesWithStyle,
27+
}
28+
}
29+
30+
func (s *saver) SaveAsYAML(data any, filename string, force bool) error {
1731
err := os.MkdirAll(filepath.Dir(filename), 0755)
1832
if err != nil {
1933
return err
@@ -36,15 +50,15 @@ func SaveAsYAML(data any, filename string, force bool) error {
3650
}
3751
defer file.Close()
3852

39-
err = encode(data, file)
53+
err = s.encode(data, file)
4054
if err != nil {
4155
return err
4256
}
4357
return nil
4458
}
4559

46-
func encode(data any, w io.Writer) error {
47-
yamlNode, err := ToYamlNode(dyn.V(data))
60+
func (s *saver) encode(data any, w io.Writer) error {
61+
yamlNode, err := s.toYamlNode(dyn.V(data))
4862
if err != nil {
4963
return err
5064
}
@@ -53,7 +67,11 @@ func encode(data any, w io.Writer) error {
5367
return enc.Encode(yamlNode)
5468
}
5569

56-
func ToYamlNode(v dyn.Value) (*yaml.Node, error) {
70+
func (s *saver) toYamlNode(v dyn.Value) (*yaml.Node, error) {
71+
return s.toYamlNodeWithStyle(v, yaml.Style(0))
72+
}
73+
74+
func (s *saver) toYamlNodeWithStyle(v dyn.Value, style yaml.Style) (*yaml.Node, error) {
5775
switch v.Kind() {
5876
case dyn.KindMap:
5977
m, _ := v.AsMap()
@@ -68,49 +86,60 @@ func ToYamlNode(v dyn.Value) (*yaml.Node, error) {
6886
content := make([]*yaml.Node, 0)
6987
for _, k := range keys {
7088
item := m[k]
71-
node := yaml.Node{Kind: yaml.ScalarNode, Value: k}
72-
c, err := ToYamlNode(item)
89+
node := yaml.Node{Kind: yaml.ScalarNode, Value: k, Style: style}
90+
var nestedNodeStyle yaml.Style
91+
if customStyle, ok := s.hasStyle(k); ok {
92+
nestedNodeStyle = customStyle
93+
} else {
94+
nestedNodeStyle = style
95+
}
96+
c, err := s.toYamlNodeWithStyle(item, nestedNodeStyle)
7397
if err != nil {
7498
return nil, err
7599
}
76100
content = append(content, &node)
77101
content = append(content, c)
78102
}
79103

80-
return &yaml.Node{Kind: yaml.MappingNode, Content: content}, nil
104+
return &yaml.Node{Kind: yaml.MappingNode, Content: content, Style: style}, nil
81105
case dyn.KindSequence:
82-
s, _ := v.AsSequence()
106+
seq, _ := v.AsSequence()
83107
content := make([]*yaml.Node, 0)
84-
for _, item := range s {
85-
node, err := ToYamlNode(item)
108+
for _, item := range seq {
109+
node, err := s.toYamlNodeWithStyle(item, style)
86110
if err != nil {
87111
return nil, err
88112
}
89113
content = append(content, node)
90114
}
91-
return &yaml.Node{Kind: yaml.SequenceNode, Content: content}, nil
115+
return &yaml.Node{Kind: yaml.SequenceNode, Content: content, Style: style}, nil
92116
case dyn.KindNil:
93-
return &yaml.Node{Kind: yaml.ScalarNode, Value: "null"}, nil
117+
return &yaml.Node{Kind: yaml.ScalarNode, Value: "null", Style: style}, nil
94118
case dyn.KindString:
95119
// If the string is a scalar value (bool, int, float and etc.), we want to quote it.
96120
if isScalarValueInString(v) {
97121
return &yaml.Node{Kind: yaml.ScalarNode, Value: v.MustString(), Style: yaml.DoubleQuotedStyle}, nil
98122
}
99-
return &yaml.Node{Kind: yaml.ScalarNode, Value: v.MustString()}, nil
123+
return &yaml.Node{Kind: yaml.ScalarNode, Value: v.MustString(), Style: style}, nil
100124
case dyn.KindBool:
101-
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustBool())}, nil
125+
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustBool()), Style: style}, nil
102126
case dyn.KindInt:
103-
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustInt())}, nil
127+
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustInt()), Style: style}, nil
104128
case dyn.KindFloat:
105-
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustFloat())}, nil
129+
return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprint(v.MustFloat()), Style: style}, nil
106130
case dyn.KindTime:
107-
return &yaml.Node{Kind: yaml.ScalarNode, Value: v.MustTime().UTC().String()}, nil
131+
return &yaml.Node{Kind: yaml.ScalarNode, Value: v.MustTime().UTC().String(), Style: style}, nil
108132
default:
109133
// Panic because we only want to deal with known types.
110134
panic(fmt.Sprintf("invalid kind: %d", v.Kind()))
111135
}
112136
}
113137

138+
func (s *saver) hasStyle(key string) (yaml.Style, bool) {
139+
style, ok := s.nodesWithStyle[key]
140+
return style, ok
141+
}
142+
114143
func isScalarValueInString(v dyn.Value) bool {
115144
if v.Kind() != dyn.KindString {
116145
return false

0 commit comments

Comments
 (0)