Skip to content

Commit 3c03fdd

Browse files
shreyas-goenkashreyas-goenka
authored
Add top level run_as support for DLT pipelines (#3307)
## Changes This PR adds support for top-level run_as to DABs. Before DABs would error if pipelines were used with the `run_as` field set to a value. Now, we'll transparently read the run_as value and set it for all pipelines in the bundle. ## Tests New acceptance tests. Also manually tested that the run_as feature works as expected.
1 parent ae9b412 commit 3c03fdd

File tree

11 files changed

+114
-68
lines changed

11 files changed

+114
-68
lines changed

‎acceptance/bundle/run_as/databricks.yml‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
resources:
2+
jobs:
3+
job_one:
4+
name: Job 1
5+
6+
job_two:
7+
name: Job 2
8+
run_as:
9+
service_principal_name: "sp_override"
10+
11+
job_three:
12+
name: Job 3
13+
run_as:
14+
user_name: "user_override"
15+
16+
pipelines:
17+
pipeline_one:
18+
name: Pipeline 1
19+
20+
pipeline_two:
21+
name: Pipeline 2
22+
run_as:
23+
service_principal_name: "sp_override"
24+
25+
pipeline_three:
26+
name: Pipeline 3
27+
run_as:
28+
user_name: "user_override"
29+
30+
targets:
31+
no_run_as:
32+
default: true
33+
34+
with_user_run_as:
35+
run_as:
36+
user_name: "user_base"
37+
38+
with_sp_run_as:
39+
run_as:
40+
service_principal_name: "sp_base"

‎acceptance/bundle/run_as/out.test.toml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Local = true
2+
Cloud = false
3+
4+
[EnvMatrix]
5+
DATABRICKS_CLI_DEPLOYMENT = ["terraform", "direct-exp"]

‎acceptance/bundle/run_as/output.txt‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
2+
>>> [CLI] bundle validate -o json -t no_run_as
3+
jobs.job_one: null
4+
jobs.job_three: {"user_name":"user_override"}
5+
jobs.job_two: {"service_principal_name":"sp_override"}
6+
pipelines.pipeline_one: null
7+
pipelines.pipeline_three: {"user_name":"user_override"}
8+
pipelines.pipeline_two: {"service_principal_name":"sp_override"}
9+
10+
>>> [CLI] bundle validate -o json -t with_user_run_as
11+
jobs.job_one: {"user_name":"user_base"}
12+
jobs.job_three: {"user_name":"user_override"}
13+
jobs.job_two: {"service_principal_name":"sp_override"}
14+
pipelines.pipeline_one: {"user_name":"user_base"}
15+
pipelines.pipeline_three: {"user_name":"user_override"}
16+
pipelines.pipeline_two: {"service_principal_name":"sp_override"}
17+
18+
>>> [CLI] bundle validate -o json -t with_sp_run_as
19+
jobs.job_one: {"service_principal_name":"sp_base"}
20+
jobs.job_three: {"user_name":"user_override"}
21+
jobs.job_two: {"service_principal_name":"sp_override"}
22+
pipelines.pipeline_one: {"service_principal_name":"sp_base"}
23+
pipelines.pipeline_three: {"user_name":"user_override"}
24+
pipelines.pipeline_two: {"service_principal_name":"sp_override"}

‎acceptance/bundle/run_as/pipelines/_script‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
trace errcode $CLI bundle validate -t t_user_name
2-
trace musterr $CLI bundle validate -t t_user_name_different
2+
trace errcode $CLI bundle validate -t t_user_name_different
33
trace errcode $CLI bundle validate -t t_service_principal_name
4-
trace musterr $CLI bundle validate -t t_service_principal_name_different
4+
trace errcode $CLI bundle validate -t t_service_principal_name_different

‎acceptance/bundle/run_as/pipelines/regular_user/output.txt‎

Lines changed: 6 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11

2-
=== t_user_name target must succeed, the rest must fail
2+
=== run_as should succeed in all cases
33
>>> errcode [CLI] bundle validate -t t_user_name
44
Name: run_as
55
Target: t_user_name
@@ -9,50 +9,29 @@ Workspace:
99

1010
Validation OK!
1111

12-
>>> musterr [CLI] bundle validate -t t_user_name_different
13-
Error: pipelines do not support a setting a run_as user that is different from the owner.
14-
Current identity: [USERNAME]. Run as identity: [email protected].
15-
See https://docs.databricks.com/dev-tools/bundles/run-as.html to learn more about the run_as property.
16-
in databricks.yml:20:5
17-
12+
>>> errcode [CLI] bundle validate -t t_user_name_different
1813
Name: run_as
1914
Target: t_user_name_different
2015
Workspace:
2116
User: [USERNAME]
2217
Path: /Workspace/Users/[USERNAME]/.bundle/run_as/t_user_name_different
2318

24-
Found 1 error
25-
26-
Exit code (musterr): 1
19+
Validation OK!
2720

2821
>>> errcode [CLI] bundle validate -t t_service_principal_name
29-
Error: pipelines do not support a setting a run_as user that is different from the owner.
30-
Current identity: [USERNAME]. Run as identity: [UUID].
31-
See https://docs.databricks.com/dev-tools/bundles/run-as.html to learn more about the run_as property.
32-
in databricks.yml:20:5
33-
3422
Name: run_as
3523
Target: t_service_principal_name
3624
Workspace:
3725
User: [USERNAME]
3826
Path: /Workspace/Users/[USERNAME]/.bundle/run_as/t_service_principal_name
3927

40-
Found 1 error
41-
42-
Exit code: 1
43-
44-
>>> musterr [CLI] bundle validate -t t_service_principal_name_different
45-
Error: pipelines do not support a setting a run_as user that is different from the owner.
46-
Current identity: [USERNAME]. Run as identity: [UUID].
47-
See https://docs.databricks.com/dev-tools/bundles/run-as.html to learn more about the run_as property.
48-
in databricks.yml:20:5
28+
Validation OK!
4929

30+
>>> errcode [CLI] bundle validate -t t_service_principal_name_different
5031
Name: run_as
5132
Target: t_service_principal_name_different
5233
Workspace:
5334
User: [USERNAME]
5435
Path: /Workspace/Users/[USERNAME]/.bundle/run_as/t_service_principal_name_different
5536

56-
Found 1 error
57-
58-
Exit code (musterr): 1
37+
Validation OK!

‎acceptance/bundle/run_as/pipelines/regular_user/script‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
cp -r $TESTDIR/../{databricks.yml,dlt} .
2-
title "t_user_name target must succeed, the rest must fail"
2+
title "run_as should succeed in all cases"
33
source $TESTDIR/../_script

‎acceptance/bundle/run_as/pipelines/service_principal/output.txt‎

Lines changed: 6 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,22 @@
11

2-
=== t_service_principal_name must succeed, the rest must fail
2+
=== run_as should succeed in all cases
33
>>> errcode [CLI] bundle validate -t t_user_name
4-
Error: pipelines do not support a setting a run_as user that is different from the owner.
5-
Current identity: [USERNAME]. Run as identity: [email protected].
6-
See https://docs.databricks.com/dev-tools/bundles/run-as.html to learn more about the run_as property.
7-
in databricks.yml:20:5
8-
94
Name: run_as
105
Target: t_user_name
116
Workspace:
127
User: [USERNAME]
138
Path: /Workspace/Users/[USERNAME]/.bundle/run_as/t_user_name
149

15-
Found 1 error
16-
17-
Exit code: 1
18-
19-
>>> musterr [CLI] bundle validate -t t_user_name_different
20-
Error: pipelines do not support a setting a run_as user that is different from the owner.
21-
Current identity: [USERNAME]. Run as identity: [email protected].
22-
See https://docs.databricks.com/dev-tools/bundles/run-as.html to learn more about the run_as property.
23-
in databricks.yml:20:5
10+
Validation OK!
2411

12+
>>> errcode [CLI] bundle validate -t t_user_name_different
2513
Name: run_as
2614
Target: t_user_name_different
2715
Workspace:
2816
User: [USERNAME]
2917
Path: /Workspace/Users/[USERNAME]/.bundle/run_as/t_user_name_different
3018

31-
Found 1 error
32-
33-
Exit code (musterr): 1
19+
Validation OK!
3420

3521
>>> errcode [CLI] bundle validate -t t_service_principal_name
3622
Name: run_as
@@ -41,18 +27,11 @@ Workspace:
4127

4228
Validation OK!
4329

44-
>>> musterr [CLI] bundle validate -t t_service_principal_name_different
45-
Error: pipelines do not support a setting a run_as user that is different from the owner.
46-
Current identity: [USERNAME]. Run as identity: [UUID].
47-
See https://docs.databricks.com/dev-tools/bundles/run-as.html to learn more about the run_as property.
48-
in databricks.yml:20:5
49-
30+
>>> errcode [CLI] bundle validate -t t_service_principal_name_different
5031
Name: run_as
5132
Target: t_service_principal_name_different
5233
Workspace:
5334
User: [USERNAME]
5435
Path: /Workspace/Users/[USERNAME]/.bundle/run_as/t_service_principal_name_different
5536

56-
Found 1 error
57-
58-
Exit code (musterr): 1
37+
Validation OK!

‎acceptance/bundle/run_as/pipelines/service_principal/script‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
cp -r $TESTDIR/../{databricks.yml,dlt} .
2-
title "t_service_principal_name must succeed, the rest must fail"
2+
title "run_as should succeed in all cases"
33
source $TESTDIR/../_script

‎acceptance/bundle/run_as/script‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
trace $CLI bundle validate -o json -t no_run_as | jq -r '.resources | to_entries[] | .key as $resource_type | .value | to_entries[] | "\($resource_type).\(.key): \(.value.run_as)"'
2+
trace $CLI bundle validate -o json -t with_user_run_as | jq -r '.resources | to_entries[] | .key as $resource_type | .value | to_entries[] | "\($resource_type).\(.key): \(.value.run_as)"'
3+
trace $CLI bundle validate -o json -t with_sp_run_as | jq -r '.resources | to_entries[] | .key as $resource_type | .value | to_entries[] | "\($resource_type).\(.key): \(.value.run_as)"'

‎bundle/config/mutator/resourcemutator/run_as.go‎

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"github.com/databricks/cli/libs/diag"
1111
"github.com/databricks/cli/libs/dyn"
1212
"github.com/databricks/databricks-sdk-go/service/jobs"
13+
"github.com/databricks/databricks-sdk-go/service/pipelines"
1314
)
1415

1516
type setRunAs struct{}
@@ -79,16 +80,6 @@ func validateRunAs(b *bundle.Bundle) diag.Diagnostics {
7980
return diags
8081
}
8182

82-
// DLT pipelines do not support run_as in the API.
83-
if len(b.Config.Resources.Pipelines) > 0 {
84-
diags = diags.Extend(reportRunAsNotSupported(
85-
"pipelines",
86-
b.Config.GetLocation("resources.pipelines"),
87-
b.Config.Workspace.CurrentUser.UserName,
88-
identity,
89-
))
90-
}
91-
9283
// Model serving endpoints do not support run_as in the API.
9384
if len(b.Config.Resources.ModelServingEndpoints) > 0 {
9485
diags = diags.Extend(reportRunAsNotSupported(
@@ -150,6 +141,24 @@ func setRunAsForJobs(b *bundle.Bundle) {
150141
}
151142
}
152143

144+
func setRunAsForPipelines(b *bundle.Bundle) {
145+
runAs := b.Config.RunAs
146+
if runAs == nil {
147+
return
148+
}
149+
150+
for i := range b.Config.Resources.Pipelines {
151+
pipeline := b.Config.Resources.Pipelines[i]
152+
if pipeline.RunAs != nil {
153+
continue
154+
}
155+
pipeline.RunAs = &pipelines.RunAs{
156+
ServicePrincipalName: runAs.ServicePrincipalName,
157+
UserName: runAs.UserName,
158+
}
159+
}
160+
}
161+
153162
// Legacy behavior of run_as for DLT pipelines. Available under the experimental.use_run_as_legacy flag.
154163
// Only available to unblock customers stuck due to breaking changes in https://github.com/databricks/cli/pull/1233
155164
func setPipelineOwnersToRunAsIdentity(b *bundle.Bundle) {
@@ -189,6 +198,8 @@ func (m *setRunAs) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics {
189198
return nil
190199
}
191200

201+
// User has opted to use the legacy behavior of run_as with the
202+
// experimental.use_legacy_run_as flag.
192203
if b.Config.Experimental != nil && b.Config.Experimental.UseLegacyRunAs {
193204
setPipelineOwnersToRunAsIdentity(b)
194205
setRunAsForJobs(b)
@@ -209,5 +220,6 @@ func (m *setRunAs) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics {
209220
}
210221

211222
setRunAsForJobs(b)
223+
setRunAsForPipelines(b)
212224
return nil
213225
}

0 commit comments

Comments
 (0)