From 09d1846e13a5ca875d943609c268d95da3c37b1b Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
Date: Mon, 4 Mar 2024 17:12:10 +0100
Subject: [PATCH] Return `application_id` for service principal lookups (#1245)

## Changes
Return ApplicationId for service principals lookups

Fixes #1234

## Tests
Added (regression) tests
---
 .codegen/lookup.go.tmpl                       | 12 ++++----
 .../resolve_resource_references_test.go       | 28 +++++++++++++++++++
 bundle/config/variable/lookup.go              |  2 +-
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/.codegen/lookup.go.tmpl b/.codegen/lookup.go.tmpl
index a982f151..7e643a90 100644
--- a/.codegen/lookup.go.tmpl
+++ b/.codegen/lookup.go.tmpl
@@ -18,6 +18,11 @@ package variable
 		"warehouses"
 }}
 
+{{	$customField :=
+		dict
+			"service-principals" "ApplicationId"
+}}
+
 import (
 	"context"
 	"fmt"
@@ -116,15 +121,10 @@ func allResolvers() *resolvers {
 			return "", err
 		}
 
-		return fmt.Sprint(entity{{ template "field-path" .List.NamedIdMap.IdPath }}), nil
+		return fmt.Sprint(entity.{{ getOrDefault $customField .KebabName ((index .List.NamedIdMap.IdPath 0).PascalName) }}), nil
 	}
 	{{end -}}
 	{{- end}}
 
 	return r
 }
-
-
-{{- define "field-path" -}}
-	{{- range .}}.{{.PascalName}}{{end}}
-{{- end -}}
diff --git a/bundle/config/mutator/resolve_resource_references_test.go b/bundle/config/mutator/resolve_resource_references_test.go
index 4d51285c..5f5dab31 100644
--- a/bundle/config/mutator/resolve_resource_references_test.go
+++ b/bundle/config/mutator/resolve_resource_references_test.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
 	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 )
 
 func TestResolveClusterReference(t *testing.T) {
@@ -105,3 +106,30 @@ func TestNoLookupIfVariableIsSet(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, "random value", *b.Config.Variables["my-cluster-id"].Value)
 }
+
+func TestResolveServicePrincipal(t *testing.T) {
+	spName := "Some SP name"
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Variables: map[string]*variable.Variable{
+				"my-sp": {
+					Lookup: &variable.Lookup{
+						ServicePrincipal: spName,
+					},
+				},
+			},
+		},
+	}
+
+	m := mocks.NewMockWorkspaceClient(t)
+	b.SetWorkpaceClient(m.WorkspaceClient)
+	spApi := m.GetMockServicePrincipalsAPI()
+	spApi.EXPECT().GetByDisplayName(mock.Anything, spName).Return(&iam.ServicePrincipal{
+		Id:            "1234",
+		ApplicationId: "app-1234",
+	}, nil)
+
+	err := bundle.Apply(context.Background(), b, ResolveResourceReferences())
+	require.NoError(t, err)
+	require.Equal(t, "app-1234", *b.Config.Variables["my-sp"].Value)
+}
diff --git a/bundle/config/variable/lookup.go b/bundle/config/variable/lookup.go
index 3b29783e..56d2ca81 100755
--- a/bundle/config/variable/lookup.go
+++ b/bundle/config/variable/lookup.go
@@ -297,7 +297,7 @@ func allResolvers() *resolvers {
 			return "", err
 		}
 
-		return fmt.Sprint(entity.Id), nil
+		return fmt.Sprint(entity.ApplicationId), nil
 	}
 	r.Warehouse = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
 		entity, err := w.Warehouses.GetByName(ctx, name)