diff --git a/bundle/config/experimental.go b/bundle/config/experimental.go index be0e7d8f..62d1ae73 100644 --- a/bundle/config/experimental.go +++ b/bundle/config/experimental.go @@ -2,6 +2,14 @@ package config type Experimental struct { Scripts map[ScriptHook]Command `json:"scripts,omitempty"` + + // By default Python wheel tasks deployed as is to Databricks platform. + // If notebook wrapper required (for example, used in DBR < 13.1 or other configuration differences), users can provide a following experimental setting + // experimental: + // python_wheel_wrapper: true + // In this case the configured wheel task will be deployed as a notebook task which install defined wheel in runtime and executes it. + // For more details see https://github.com/databricks/cli/pull/797 and https://github.com/databricks/cli/pull/635 + PythonWheelWrapper bool `json:"python_wheel_wrapper,omitempty"` } type Command string diff --git a/bundle/config/mutator/if.go b/bundle/config/mutator/if.go new file mode 100644 index 00000000..462d8f00 --- /dev/null +++ b/bundle/config/mutator/if.go @@ -0,0 +1,35 @@ +package mutator + +import ( + "context" + + "github.com/databricks/cli/bundle" +) + +type ifMutator struct { + condition func(*bundle.Bundle) bool + onTrueMutator bundle.Mutator + onFalseMutator bundle.Mutator +} + +func If( + condition func(*bundle.Bundle) bool, + onTrueMutator bundle.Mutator, + onFalseMutator bundle.Mutator, +) bundle.Mutator { + return &ifMutator{ + condition, onTrueMutator, onFalseMutator, + } +} + +func (m *ifMutator) Apply(ctx context.Context, b *bundle.Bundle) error { + if m.condition(b) { + return bundle.Apply(ctx, b, m.onTrueMutator) + } else { + return bundle.Apply(ctx, b, m.onFalseMutator) + } +} + +func (m *ifMutator) Name() string { + return "If" +} diff --git a/bundle/config/mutator/noop.go b/bundle/config/mutator/noop.go new file mode 100644 index 00000000..91c16385 --- /dev/null +++ b/bundle/config/mutator/noop.go @@ -0,0 +1,21 @@ +package mutator + +import ( + "context" + + "github.com/databricks/cli/bundle" +) + +type noop struct{} + +func (*noop) Apply(context.Context, *bundle.Bundle) error { + return nil +} + +func (*noop) Name() string { + return "NoOp" +} + +func NoOp() bundle.Mutator { + return &noop{} +} diff --git a/bundle/python/conditional_transform_test.go b/bundle/python/conditional_transform_test.go new file mode 100644 index 00000000..5bf33721 --- /dev/null +++ b/bundle/python/conditional_transform_test.go @@ -0,0 +1,114 @@ +package python + +import ( + "context" + "path" + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/require" +) + +func TestNoTransformByDefault(t *testing.T) { + tmpDir := t.TempDir() + + b := &bundle.Bundle{ + Config: config.Root{ + Path: tmpDir, + Bundle: config.Bundle{ + Target: "development", + }, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "job1": { + JobSettings: &jobs.JobSettings{ + Tasks: []jobs.Task{ + { + TaskKey: "key1", + PythonWheelTask: &jobs.PythonWheelTask{ + PackageName: "test_package", + EntryPoint: "main", + }, + Libraries: []compute.Library{ + {Whl: "/Workspace/Users/test@test.com/bundle/dist/test.whl"}, + }, + }, + }, + }, + }, + }, + }, + }, + } + + trampoline := TransformWheelTask() + err := bundle.Apply(context.Background(), b, trampoline) + require.NoError(t, err) + + task := b.Config.Resources.Jobs["job1"].Tasks[0] + require.NotNil(t, task.PythonWheelTask) + require.Equal(t, "test_package", task.PythonWheelTask.PackageName) + require.Equal(t, "main", task.PythonWheelTask.EntryPoint) + require.Equal(t, "/Workspace/Users/test@test.com/bundle/dist/test.whl", task.Libraries[0].Whl) + + require.Nil(t, task.NotebookTask) +} + +func TestTransformWithExperimentalSettingSetToTrue(t *testing.T) { + tmpDir := t.TempDir() + + b := &bundle.Bundle{ + Config: config.Root{ + Path: tmpDir, + Bundle: config.Bundle{ + Target: "development", + }, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "job1": { + JobSettings: &jobs.JobSettings{ + Tasks: []jobs.Task{ + { + TaskKey: "key1", + PythonWheelTask: &jobs.PythonWheelTask{ + PackageName: "test_package", + EntryPoint: "main", + }, + Libraries: []compute.Library{ + {Whl: "/Workspace/Users/test@test.com/bundle/dist/test.whl"}, + }, + }, + }, + }, + }, + }, + }, + Experimental: &config.Experimental{ + PythonWheelWrapper: true, + }, + }, + } + + trampoline := TransformWheelTask() + err := bundle.Apply(context.Background(), b, trampoline) + require.NoError(t, err) + + task := b.Config.Resources.Jobs["job1"].Tasks[0] + require.Nil(t, task.PythonWheelTask) + require.NotNil(t, task.NotebookTask) + + dir, err := b.InternalDir(context.Background()) + require.NoError(t, err) + + internalDirRel, err := filepath.Rel(b.Config.Path, dir) + require.NoError(t, err) + + require.Equal(t, path.Join(filepath.ToSlash(internalDirRel), "notebook_job1_key1"), task.NotebookTask.NotebookPath) + + require.Empty(t, task.Libraries) +} diff --git a/bundle/python/transform.go b/bundle/python/transform.go index d8eb33f5..f6207a59 100644 --- a/bundle/python/transform.go +++ b/bundle/python/transform.go @@ -49,10 +49,16 @@ dbutils.notebook.exit(s) // which installs uploaded wheels using %pip and then calling corresponding // entry point. func TransformWheelTask() bundle.Mutator { - return mutator.NewTrampoline( - "python_wheel", - &pythonTrampoline{}, - NOTEBOOK_TEMPLATE, + return mutator.If( + func(b *bundle.Bundle) bool { + return b.Config.Experimental != nil && b.Config.Experimental.PythonWheelWrapper + }, + mutator.NewTrampoline( + "python_wheel", + &pythonTrampoline{}, + NOTEBOOK_TEMPLATE, + ), + mutator.NoOp(), ) } @@ -113,7 +119,7 @@ func (t *pythonTrampoline) generateParameters(task *jobs.PythonWheelTask) (strin if task.Parameters != nil && task.NamedParameters != nil { return "", fmt.Errorf("not allowed to pass both paramaters and named_parameters") } - params := append([]string{"python"}, task.Parameters...) + params := append([]string{task.PackageName}, task.Parameters...) for k, v := range task.NamedParameters { params = append(params, fmt.Sprintf("%s=%s", k, v)) } diff --git a/bundle/python/transform_test.go b/bundle/python/transform_test.go index a7448f23..1ccdba56 100644 --- a/bundle/python/transform_test.go +++ b/bundle/python/transform_test.go @@ -25,26 +25,26 @@ type testCaseNamed struct { } var paramsTestCases []testCase = []testCase{ - {[]string{}, `"python"`}, - {[]string{"a"}, `"python", "a"`}, - {[]string{"a", "b"}, `"python", "a", "b"`}, - {[]string{"123!@#$%^&*()-="}, `"python", "123!@#$%^&*()-="`}, - {[]string{`{"a": 1}`}, `"python", "{\"a\": 1}"`}, + {[]string{}, `"my_test_code"`}, + {[]string{"a"}, `"my_test_code", "a"`}, + {[]string{"a", "b"}, `"my_test_code", "a", "b"`}, + {[]string{"123!@#$%^&*()-="}, `"my_test_code", "123!@#$%^&*()-="`}, + {[]string{`{"a": 1}`}, `"my_test_code", "{\"a\": 1}"`}, } var paramsTestCasesNamed []testCaseNamed = []testCaseNamed{ - {map[string]string{}, `"python"`}, - {map[string]string{"a": "1"}, `"python", "a=1"`}, - {map[string]string{"a": "'1'"}, `"python", "a='1'"`}, - {map[string]string{"a": `"1"`}, `"python", "a=\"1\""`}, - {map[string]string{"a": "1", "b": "2"}, `"python", "a=1", "b=2"`}, - {map[string]string{"data": `{"a": 1}`}, `"python", "data={\"a\": 1}"`}, + {map[string]string{}, `"my_test_code"`}, + {map[string]string{"a": "1"}, `"my_test_code", "a=1"`}, + {map[string]string{"a": "'1'"}, `"my_test_code", "a='1'"`}, + {map[string]string{"a": `"1"`}, `"my_test_code", "a=\"1\""`}, + {map[string]string{"a": "1", "b": "2"}, `"my_test_code", "a=1", "b=2"`}, + {map[string]string{"data": `{"a": 1}`}, `"my_test_code", "data={\"a\": 1}"`}, } func TestGenerateParameters(t *testing.T) { trampoline := pythonTrampoline{} for _, c := range paramsTestCases { - task := &jobs.PythonWheelTask{Parameters: c.Actual} + task := &jobs.PythonWheelTask{PackageName: "my_test_code", Parameters: c.Actual} result, err := trampoline.generateParameters(task) require.NoError(t, err) require.Equal(t, c.Expected, result) @@ -54,7 +54,7 @@ func TestGenerateParameters(t *testing.T) { func TestGenerateNamedParameters(t *testing.T) { trampoline := pythonTrampoline{} for _, c := range paramsTestCasesNamed { - task := &jobs.PythonWheelTask{NamedParameters: c.Actual} + task := &jobs.PythonWheelTask{PackageName: "my_test_code", NamedParameters: c.Actual} result, err := trampoline.generateParameters(task) require.NoError(t, err) diff --git a/internal/bundle/bundles/python_wheel_task/databricks_template_schema.json b/internal/bundle/bundles/python_wheel_task/databricks_template_schema.json index f7f4b634..0695eb2b 100644 --- a/internal/bundle/bundles/python_wheel_task/databricks_template_schema.json +++ b/internal/bundle/bundles/python_wheel_task/databricks_template_schema.json @@ -16,6 +16,10 @@ "unique_id": { "type": "string", "description": "Unique ID for job name" + }, + "python_wheel_wrapper": { + "type": "boolean", + "description": "Whether or not to enable python wheel wrapper" } } } diff --git a/internal/bundle/bundles/python_wheel_task/template/databricks.yml.tmpl b/internal/bundle/bundles/python_wheel_task/template/databricks.yml.tmpl index e715cdf1..8729dcba 100644 --- a/internal/bundle/bundles/python_wheel_task/template/databricks.yml.tmpl +++ b/internal/bundle/bundles/python_wheel_task/template/databricks.yml.tmpl @@ -4,6 +4,11 @@ bundle: workspace: root_path: "~/.bundle/{{.unique_id}}" +{{if .python_wheel_wrapper}} +experimental: + python_wheel_wrapper: true +{{end}} + resources: jobs: some_other_job: @@ -14,6 +19,7 @@ resources: num_workers: 1 spark_version: "{{.spark_version}}" node_type_id: "{{.node_type_id}}" + data_security_mode: USER_ISOLATION python_wheel_task: package_name: my_test_code entry_point: run diff --git a/internal/bundle/python_wheel_test.go b/internal/bundle/python_wheel_test.go index ee5d897d..fd5c9acc 100644 --- a/internal/bundle/python_wheel_test.go +++ b/internal/bundle/python_wheel_test.go @@ -8,7 +8,7 @@ import ( "github.com/stretchr/testify/require" ) -func TestAccPythonWheelTaskDeployAndRun(t *testing.T) { +func runPythonWheelTest(t *testing.T, pythonWheelWrapper bool) { env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") t.Log(env) @@ -22,9 +22,10 @@ func TestAccPythonWheelTaskDeployAndRun(t *testing.T) { } bundleRoot, err := initTestTemplate(t, "python_wheel_task", map[string]any{ - "node_type_id": nodeTypeId, - "unique_id": uuid.New().String(), - "spark_version": "13.2.x-snapshot-scala2.12", + "node_type_id": nodeTypeId, + "unique_id": uuid.New().String(), + "spark_version": "13.2.x-snapshot-scala2.12", + "python_wheel_wrapper": pythonWheelWrapper, }) require.NoError(t, err) @@ -39,5 +40,13 @@ func TestAccPythonWheelTaskDeployAndRun(t *testing.T) { require.NoError(t, err) require.Contains(t, out, "Hello from my func") require.Contains(t, out, "Got arguments:") - require.Contains(t, out, "['python', 'one', 'two']") + require.Contains(t, out, "['my_test_code', 'one', 'two']") +} + +func TestAccPythonWheelTaskDeployAndRunWithoutWrapper(t *testing.T) { + runPythonWheelTest(t, false) +} + +func TestAccPythonWheelTaskDeployAndRunWithWrapper(t *testing.T) { + runPythonWheelTest(t, true) }