From 507053ee50563bae098557246b4415b1b272d6ec Mon Sep 17 00:00:00 2001 From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com> Date: Wed, 1 May 2024 14:07:03 +0530 Subject: [PATCH] Annotate DLT pipelines when deployed using DABs (#1410) ## Changes This PR annotates any pipelines that were deployed using DABs to have `deployment.kind` set to "BUNDLE", mirroring the annotation for Jobs (similar PR for jobs FYI: https://github.com/databricks/cli/pull/880). Breakglass UI is not yet available for pipelines, so this annotation will just be used for revenue attribution ATM. Note: The API field has been deployed in all regions including GovCloud. ## Tests Unit tests and manually. Manually verified that the kind and metadata_file_path are being set by DABs, and are returned by a GET API to a pipeline deployed using a DAB. Example: ``` "deployment": { "kind":"BUNDLE", "metadata_file_path":"/Users/shreyas.goenka@databricks.com/.bundle/bundle-playground/default/state/metadata.json" }, ``` --- bundle/deploy/metadata/annotate_jobs.go | 3 +- bundle/deploy/metadata/annotate_pipelines.go | 34 +++++++++ .../metadata/annotate_pipelines_test.go | 72 +++++++++++++++++++ bundle/deploy/metadata/upload.go | 9 ++- bundle/phases/initialize.go | 1 + 5 files changed, 115 insertions(+), 4 deletions(-) create mode 100644 bundle/deploy/metadata/annotate_pipelines.go create mode 100644 bundle/deploy/metadata/annotate_pipelines_test.go diff --git a/bundle/deploy/metadata/annotate_jobs.go b/bundle/deploy/metadata/annotate_jobs.go index 2b03a59b..f42d4693 100644 --- a/bundle/deploy/metadata/annotate_jobs.go +++ b/bundle/deploy/metadata/annotate_jobs.go @@ -2,7 +2,6 @@ package metadata import ( "context" - "path" "github.com/databricks/cli/bundle" "github.com/databricks/cli/libs/diag" @@ -27,7 +26,7 @@ func (m *annotateJobs) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnosti job.JobSettings.Deployment = &jobs.JobDeployment{ Kind: jobs.JobDeploymentKindBundle, - MetadataFilePath: path.Join(b.Config.Workspace.StatePath, MetadataFileName), + MetadataFilePath: metadataFilePath(b), } job.JobSettings.EditMode = jobs.JobEditModeUiLocked job.JobSettings.Format = jobs.FormatMultiTask diff --git a/bundle/deploy/metadata/annotate_pipelines.go b/bundle/deploy/metadata/annotate_pipelines.go new file mode 100644 index 00000000..990f4890 --- /dev/null +++ b/bundle/deploy/metadata/annotate_pipelines.go @@ -0,0 +1,34 @@ +package metadata + +import ( + "context" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/databricks-sdk-go/service/pipelines" +) + +type annotatePipelines struct{} + +func AnnotatePipelines() bundle.Mutator { + return &annotatePipelines{} +} + +func (m *annotatePipelines) Name() string { + return "metadata.AnnotatePipelines" +} + +func (m *annotatePipelines) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics { + for _, pipeline := range b.Config.Resources.Pipelines { + if pipeline.PipelineSpec == nil { + continue + } + + pipeline.PipelineSpec.Deployment = &pipelines.PipelineDeployment{ + Kind: pipelines.DeploymentKindBundle, + MetadataFilePath: metadataFilePath(b), + } + } + + return nil +} diff --git a/bundle/deploy/metadata/annotate_pipelines_test.go b/bundle/deploy/metadata/annotate_pipelines_test.go new file mode 100644 index 00000000..448a022d --- /dev/null +++ b/bundle/deploy/metadata/annotate_pipelines_test.go @@ -0,0 +1,72 @@ +package metadata + +import ( + "context" + "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/pipelines" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAnnotatePipelinesMutator(t *testing.T) { + b := &bundle.Bundle{ + Config: config.Root{ + Workspace: config.Workspace{ + StatePath: "/a/b/c", + }, + Resources: config.Resources{ + Pipelines: map[string]*resources.Pipeline{ + "my-pipeline-1": { + PipelineSpec: &pipelines.PipelineSpec{ + Name: "My Pipeline One", + }, + }, + "my-pipeline-2": { + PipelineSpec: &pipelines.PipelineSpec{ + Name: "My Pipeline Two", + }, + }, + }, + }, + }, + } + + diags := bundle.Apply(context.Background(), b, AnnotatePipelines()) + require.NoError(t, diags.Error()) + + assert.Equal(t, + &pipelines.PipelineDeployment{ + Kind: pipelines.DeploymentKindBundle, + MetadataFilePath: "/a/b/c/metadata.json", + }, + b.Config.Resources.Pipelines["my-pipeline-1"].PipelineSpec.Deployment) + + assert.Equal(t, + &pipelines.PipelineDeployment{ + Kind: pipelines.DeploymentKindBundle, + MetadataFilePath: "/a/b/c/metadata.json", + }, + b.Config.Resources.Pipelines["my-pipeline-2"].PipelineSpec.Deployment) +} + +func TestAnnotatePipelinesMutatorPipelineWithoutASpec(t *testing.T) { + b := &bundle.Bundle{ + Config: config.Root{ + Workspace: config.Workspace{ + StatePath: "/a/b/c", + }, + Resources: config.Resources{ + Pipelines: map[string]*resources.Pipeline{ + "my-pipeline-1": {}, + }, + }, + }, + } + + diags := bundle.Apply(context.Background(), b, AnnotatePipelines()) + require.NoError(t, diags.Error()) +} diff --git a/bundle/deploy/metadata/upload.go b/bundle/deploy/metadata/upload.go index a040a0ae..ee87816d 100644 --- a/bundle/deploy/metadata/upload.go +++ b/bundle/deploy/metadata/upload.go @@ -4,13 +4,18 @@ import ( "bytes" "context" "encoding/json" + "path" "github.com/databricks/cli/bundle" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/filer" ) -const MetadataFileName = "metadata.json" +const metadataFileName = "metadata.json" + +func metadataFilePath(b *bundle.Bundle) string { + return path.Join(b.Config.Workspace.StatePath, metadataFileName) +} type upload struct{} @@ -33,5 +38,5 @@ func (m *upload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { return diag.FromErr(err) } - return diag.FromErr(f.Write(ctx, MetadataFileName, bytes.NewReader(metadata), filer.CreateParentDirectories, filer.OverwriteIfExists)) + return diag.FromErr(f.Write(ctx, metadataFileName, bytes.NewReader(metadata), filer.CreateParentDirectories, filer.OverwriteIfExists)) } diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index 2f5eab30..ded2e198 100644 --- a/bundle/phases/initialize.go +++ b/bundle/phases/initialize.go @@ -45,6 +45,7 @@ func Initialize() bundle.Mutator { permissions.ApplyBundlePermissions(), permissions.FilterCurrentUser(), metadata.AnnotateJobs(), + metadata.AnnotatePipelines(), terraform.Initialize(), scripts.Execute(config.ScriptPostInit), },