From 15ca7fe62d522c5efc5fae73d911cbde03e116b2 Mon Sep 17 00:00:00 2001
From: Arpit Jasapara <87999496+arpitjasa-db@users.noreply.github.com>
Date: Fri, 19 Jul 2024 04:38:20 -0700
Subject: [PATCH] Add UUID function to bundle template functions  (#1612)

## Changes

Add support for google/uuid.New() to DAB templates.

This is needed to generate UUIDs in downstream templates like MLOps
Stacks.

## Tests

Unit tests.
---
 libs/template/helpers.go                        |  6 ++++++
 libs/template/helpers_test.go                   | 17 +++++++++++++++++
 libs/template/testdata/uuid/template/hello.tmpl |  1 +
 3 files changed, 24 insertions(+)
 create mode 100644 libs/template/testdata/uuid/template/hello.tmpl

diff --git a/libs/template/helpers.go b/libs/template/helpers.go
index b3dea329..1dfe74d7 100644
--- a/libs/template/helpers.go
+++ b/libs/template/helpers.go
@@ -14,6 +14,8 @@ import (
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/service/iam"
+
+	"github.com/google/uuid"
 )
 
 type ErrFail struct {
@@ -51,6 +53,10 @@ func loadHelpers(ctx context.Context) template.FuncMap {
 		"random_int": func(n int) int {
 			return rand.Intn(n)
 		},
+		// Alias for https://pkg.go.dev/github.com/google/uuid#New. Returns, as a string, a UUID which is a 128 bit (16 byte) Universal Unique IDentifier as defined in RFC 4122.
+		"uuid": func() string {
+			return uuid.New().String()
+		},
 		// 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 c0848c8d..8cc7b928 100644
--- a/libs/template/helpers_test.go
+++ b/libs/template/helpers_test.go
@@ -69,6 +69,23 @@ func TestTemplateRandIntFunction(t *testing.T) {
 	assert.Empty(t, err)
 }
 
+func TestTemplateUuidFunction(t *testing.T) {
+	ctx := context.Background()
+	tmpDir := t.TempDir()
+
+	ctx = root.SetWorkspaceClient(ctx, nil)
+	helpers := loadHelpers(ctx)
+	r, err := newRenderer(ctx, nil, helpers, "./testdata/uuid/template", "./testdata/uuid/library", tmpDir)
+	require.NoError(t, err)
+
+	err = r.walk()
+	assert.NoError(t, err)
+
+	assert.Len(t, r.files, 1)
+	uuid := strings.TrimSpace(string(r.files[0].(*inMemoryFile).content))
+	assert.Regexp(t, "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", uuid)
+}
+
 func TestTemplateUrlFunction(t *testing.T) {
 	ctx := context.Background()
 	tmpDir := t.TempDir()
diff --git a/libs/template/testdata/uuid/template/hello.tmpl b/libs/template/testdata/uuid/template/hello.tmpl
new file mode 100644
index 00000000..178c2a9c
--- /dev/null
+++ b/libs/template/testdata/uuid/template/hello.tmpl
@@ -0,0 +1 @@
+{{print (uuid)}}