From e86a949d9908dea65117419d3e2d55e8858cb5c0 Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Mon, 2 Dec 2024 15:59:29 +0530
Subject: [PATCH] Add the `bundle_uuid` helper function for templates (#1947)

## Changes
This PR adds the `bundle_uuid` helper function that'll return a stable
identifier for the bundle for the duration of the `bundle init` command.

This is also the UUID that'll be set in the telemetry event sent during
`databricks bundle init` and would be used to correlate revenue from
bundle init with resource deployments.

Template authors should add the uuid field to their `databricks.yml`
file they generate:
```
bundle:
  # A stable identified for your DAB project. We use this UUID in the Databricks backend
  # to correlate and identify multiple deployments of the same DAB project.
  uuid: {{ bundle_uuid }}
```

## Tests
Unit test
---
 bundle/config/bundle.go                        |  4 ++++
 libs/template/helpers.go                       | 10 ++++++++++
 libs/template/helpers_test.go                  | 18 ++++++++++++++++++
 .../testdata/bundle-uuid/library/.gitkeep      |  0
 .../testdata/bundle-uuid/template/bar.txt.tmpl |  1 +
 .../testdata/bundle-uuid/template/foo.txt.tmpl |  1 +
 6 files changed, 34 insertions(+)
 create mode 100644 libs/template/testdata/bundle-uuid/library/.gitkeep
 create mode 100644 libs/template/testdata/bundle-uuid/template/bar.txt.tmpl
 create mode 100644 libs/template/testdata/bundle-uuid/template/foo.txt.tmpl

diff --git a/bundle/config/bundle.go b/bundle/config/bundle.go
index f533c4d1..af9e1f4b 100644
--- a/bundle/config/bundle.go
+++ b/bundle/config/bundle.go
@@ -49,4 +49,8 @@ type Bundle struct {
 
 	// Databricks CLI version constraints required to run the bundle.
 	DatabricksCliVersion string `json:"databricks_cli_version,omitempty"`
+
+	// A stable generated UUID for the bundle. This is normally serialized by
+	// Databricks first party template when a user runs bundle init.
+	Uuid string `json:"uuid,omitempty"`
 }
diff --git a/libs/template/helpers.go b/libs/template/helpers.go
index f25cbee4..7f7acbd2 100644
--- a/libs/template/helpers.go
+++ b/libs/template/helpers.go
@@ -35,6 +35,13 @@ var cachedUser *iam.User
 var cachedIsServicePrincipal *bool
 var cachedCatalog *string
 
+// UUID that is stable for the duration of the template execution. This can be used
+// to populate the `bundle.uuid` field in databricks.yml by template authors.
+//
+// It's automatically logged in our telemetry logs when `databricks bundle init`
+// is run and can be used to attribute DBU revenue to bundle templates.
+var bundleUuid = uuid.New().String()
+
 func loadHelpers(ctx context.Context) template.FuncMap {
 	w := root.WorkspaceClient(ctx)
 	return template.FuncMap{
@@ -57,6 +64,9 @@ func loadHelpers(ctx context.Context) template.FuncMap {
 		"uuid": func() string {
 			return uuid.New().String()
 		},
+		"bundle_uuid": func() string {
+			return bundleUuid
+		},
 		// A key value pair. This is used with the map function to generate maps
 		// to use inside a template
 		"pair": func(k string, v any) pair {
diff --git a/libs/template/helpers_test.go b/libs/template/helpers_test.go
index 9f5804c0..802a8f6d 100644
--- a/libs/template/helpers_test.go
+++ b/libs/template/helpers_test.go
@@ -32,6 +32,24 @@ func TestTemplatePrintStringWithoutProcessing(t *testing.T) {
 	assert.Equal(t, `{{ fail "abc" }}`, cleanContent)
 }
 
+func TestTemplateBundleUuidFunction(t *testing.T) {
+	ctx := context.Background()
+
+	ctx = root.SetWorkspaceClient(ctx, nil)
+	helpers := loadHelpers(ctx)
+	r, err := newRenderer(ctx, nil, helpers, os.DirFS("."), "./testdata/bundle-uuid/template", "./testdata/bundle-uuid/library")
+	require.NoError(t, err)
+
+	err = r.walk()
+	assert.NoError(t, err)
+
+	assert.Len(t, r.files, 2)
+	c1 := strings.Trim(string(r.files[0].(*inMemoryFile).content), "\n\r")
+	assert.Equal(t, strings.Repeat(bundleUuid, 3), c1)
+	c2 := strings.Trim(string(r.files[1].(*inMemoryFile).content), "\n\r")
+	assert.Equal(t, strings.Repeat(bundleUuid, 5), c2)
+}
+
 func TestTemplateRegexpCompileFunction(t *testing.T) {
 	ctx := context.Background()
 
diff --git a/libs/template/testdata/bundle-uuid/library/.gitkeep b/libs/template/testdata/bundle-uuid/library/.gitkeep
new file mode 100644
index 00000000..e69de29b
diff --git a/libs/template/testdata/bundle-uuid/template/bar.txt.tmpl b/libs/template/testdata/bundle-uuid/template/bar.txt.tmpl
new file mode 100644
index 00000000..824412df
--- /dev/null
+++ b/libs/template/testdata/bundle-uuid/template/bar.txt.tmpl
@@ -0,0 +1 @@
+{{bundle_uuid}}{{bundle_uuid}}{{bundle_uuid}}
diff --git a/libs/template/testdata/bundle-uuid/template/foo.txt.tmpl b/libs/template/testdata/bundle-uuid/template/foo.txt.tmpl
new file mode 100644
index 00000000..64150e42
--- /dev/null
+++ b/libs/template/testdata/bundle-uuid/template/foo.txt.tmpl
@@ -0,0 +1 @@
+{{bundle_uuid}}{{bundle_uuid}}{{bundle_uuid}}{{bundle_uuid}}{{bundle_uuid}}