From 7ef75747c5995166c64befbc06853c12b5c0c909 Mon Sep 17 00:00:00 2001
From: Lennart Kats <lennart.kats@databricks.com>
Date: Tue, 5 Sep 2023 11:28:40 +0200
Subject: [PATCH] Fix isServicePrincipal() only working for workspace admins

---
 bundle/config/mutator/process_target_mode.go |  5 +----
 libs/auth/service_principal.go               | 21 ++++++++------------
 libs/auth/service_principal_test.go          | 19 ++++++++++++++++++
 libs/template/helpers.go                     |  5 +----
 4 files changed, 29 insertions(+), 21 deletions(-)
 create mode 100644 libs/auth/service_principal_test.go

diff --git a/bundle/config/mutator/process_target_mode.go b/bundle/config/mutator/process_target_mode.go
index be93512bbf..06ae7b858f 100644
--- a/bundle/config/mutator/process_target_mode.go
+++ b/bundle/config/mutator/process_target_mode.go
@@ -160,10 +160,7 @@ func (m *processTargetMode) Apply(ctx context.Context, b *bundle.Bundle) error {
 		}
 		return transformDevelopmentMode(b)
 	case config.Production:
-		isPrincipal, err := auth.IsServicePrincipal(ctx, b.WorkspaceClient(), b.Config.Workspace.CurrentUser.Id)
-		if err != nil {
-			return err
-		}
+		isPrincipal := auth.IsServicePrincipal(b.Config.Workspace.CurrentUser.Id)
 		return validateProductionMode(ctx, b, isPrincipal)
 	case "":
 		// No action
diff --git a/libs/auth/service_principal.go b/libs/auth/service_principal.go
index a6740b5037..cb488d16e5 100644
--- a/libs/auth/service_principal.go
+++ b/libs/auth/service_principal.go
@@ -1,20 +1,15 @@
 package auth
 
 import (
-	"context"
-
-	"github.com/databricks/databricks-sdk-go"
-	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/google/uuid"
 )
 
 // Determines whether a given user id is a service principal.
-// This function uses a heuristic: if no user exists with this id, we assume
-// it's a service principal. Unfortunately, the standard service principal API is too
-// slow for our purposes.
-func IsServicePrincipal(ctx context.Context, ws *databricks.WorkspaceClient, userId string) (bool, error) {
-	_, err := ws.Users.GetById(ctx, userId)
-	if apierr.IsMissing(err) {
-		return true, nil
-	}
-	return false, err
+// This function uses a heuristic: if the user id is a UUID, then we assume
+// it's a service principal. Unfortunately, the service principal listing API is too
+// slow for our purposes. And the "users" and "service principals get" APIs
+// only allow access by workspace admins.
+func IsServicePrincipal(userId string) bool {
+	_, err := uuid.Parse(userId)
+	return err == nil
 }
diff --git a/libs/auth/service_principal_test.go b/libs/auth/service_principal_test.go
new file mode 100644
index 0000000000..95e8ab5cb4
--- /dev/null
+++ b/libs/auth/service_principal_test.go
@@ -0,0 +1,19 @@
+package auth
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsServicePrincipal_ValidUUID(t *testing.T) {
+	userId := "8b948b2e-d2b5-4b9e-8274-11b596f3b652"
+	isSP := IsServicePrincipal(userId)
+	assert.True(t, isSP, "Expected user ID to be recognized as a service principal")
+}
+
+func TestIsServicePrincipal_InvalidUUID(t *testing.T) {
+	userId := "invalid"
+	isSP := IsServicePrincipal(userId)
+	assert.False(t, isSP, "Expected user ID to not be recognized as a service principal")
+}
diff --git a/libs/template/helpers.go b/libs/template/helpers.go
index f947d9ba8c..29abbe21ca 100644
--- a/libs/template/helpers.go
+++ b/libs/template/helpers.go
@@ -104,10 +104,7 @@ func loadHelpers(ctx context.Context) template.FuncMap {
 					return false, err
 				}
 			}
-			result, err := auth.IsServicePrincipal(ctx, w, user.Id)
-			if err != nil {
-				return false, err
-			}
+			result := auth.IsServicePrincipal(user.Id)
 			is_service_principal = &result
 			return result, nil
 		},