Move utility functions dealing with IAM to libs/iamutil (#1820)

## Changes

The two functions `GetShortUserName` and `IsServicePrincipal` are
unrelated to auth or the purpose of the auth package. This change moves
them into their own package and updates `IsServicePrincipal` to take an
`*iam.User` argument instead of a string username.

## Tests

Tests pass.
This commit is contained in:
Pieter Noordhuis 2024-10-10 15:02:25 +02:00 committed by GitHub
parent e885794722
commit 3270afaff4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 35 additions and 28 deletions

View File

@ -5,8 +5,8 @@ import (
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/libs/auth"
"github.com/databricks/cli/libs/diag"
"github.com/databricks/cli/libs/iamutil"
"github.com/databricks/cli/libs/tags"
)
@ -33,7 +33,7 @@ func (m *populateCurrentUser) Apply(ctx context.Context, b *bundle.Bundle) diag.
}
b.Config.Workspace.CurrentUser = &config.User{
ShortName: auth.GetShortUserName(me),
ShortName: iamutil.GetShortUserName(me),
User: me,
}

View File

@ -6,9 +6,9 @@ import (
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/libs/auth"
"github.com/databricks/cli/libs/diag"
"github.com/databricks/cli/libs/dyn"
"github.com/databricks/cli/libs/iamutil"
"github.com/databricks/cli/libs/log"
)
@ -174,7 +174,7 @@ func (m *processTargetMode) Apply(ctx context.Context, b *bundle.Bundle) diag.Di
transformDevelopmentMode(ctx, b)
return diags
case config.Production:
isPrincipal := auth.IsServicePrincipal(b.Config.Workspace.CurrentUser.UserName)
isPrincipal := iamutil.IsServicePrincipal(b.Config.Workspace.CurrentUser.User)
return validateProductionMode(ctx, b, isPrincipal)
case "":
// No action

View File

@ -5,8 +5,8 @@ import (
"fmt"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/libs/auth"
"github.com/databricks/cli/libs/diag"
"github.com/databricks/cli/libs/iamutil"
"github.com/databricks/cli/libs/log"
)
@ -17,9 +17,10 @@ import (
func ReportPossiblePermissionDenied(ctx context.Context, b *bundle.Bundle, path string) diag.Diagnostics {
log.Errorf(ctx, "Failed to update, encountered possible permission error: %v", path)
user := b.Config.Workspace.CurrentUser.UserName
if auth.IsServicePrincipal(user) {
user = b.Config.Workspace.CurrentUser.DisplayName
me := b.Config.Workspace.CurrentUser.User
userName := me.UserName
if iamutil.IsServicePrincipal(me) {
userName = me.DisplayName
}
canManageBundle, assistance := analyzeBundlePermissions(b)
@ -30,7 +31,7 @@ func ReportPossiblePermissionDenied(ctx context.Context, b *bundle.Bundle, path
"%s\n"+
"They may need to redeploy the bundle to apply the new permissions.\n"+
"Please refer to https://docs.databricks.com/dev-tools/bundles/permissions.html for more on managing permissions.",
path, user, assistance),
path, userName, assistance),
Severity: diag.Error,
ID: diag.PathPermissionDenied,
}}
@ -44,7 +45,7 @@ func ReportPossiblePermissionDenied(ctx context.Context, b *bundle.Bundle, path
"%s\n"+
"They can redeploy the project to apply the latest set of permissions.\n"+
"Please refer to https://docs.databricks.com/dev-tools/bundles/permissions.html for more on managing permissions.",
path, user, assistance),
path, userName, assistance),
Severity: diag.Error,
ID: diag.CannotChangePathPermissions,
}}

View File

@ -12,7 +12,7 @@ import (
"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/internal/testutil"
"github.com/databricks/cli/libs/auth"
"github.com/databricks/cli/libs/iamutil"
"github.com/databricks/databricks-sdk-go"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@ -126,7 +126,7 @@ func TestAccBundleInitHelpers(t *testing.T) {
}{
{
funcName: "{{short_name}}",
expected: auth.GetShortUserName(me),
expected: iamutil.GetShortUserName(me),
},
{
funcName: "{{user_name}}",
@ -138,7 +138,7 @@ func TestAccBundleInitHelpers(t *testing.T) {
},
{
funcName: "{{is_service_principal}}",
expected: strconv.FormatBool(auth.IsServicePrincipal(me.UserName)),
expected: strconv.FormatBool(iamutil.IsServicePrincipal(me)),
},
{
funcName: "{{smallest_node_type}}",

View File

@ -1,15 +1,16 @@
package auth
package iamutil
import (
"github.com/databricks/databricks-sdk-go/service/iam"
"github.com/google/uuid"
)
// Determines whether a given user name is a service principal.
// Determines whether a given user is a service principal.
// This function uses a heuristic: if the user name 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(userName string) bool {
_, err := uuid.Parse(userName)
func IsServicePrincipal(user *iam.User) bool {
_, err := uuid.Parse(user.UserName)
return err == nil
}

View File

@ -1,19 +1,24 @@
package auth
package iamutil
import (
"testing"
"github.com/databricks/databricks-sdk-go/service/iam"
"github.com/stretchr/testify/assert"
)
func TestIsServicePrincipal_ValidUUID(t *testing.T) {
userId := "8b948b2e-d2b5-4b9e-8274-11b596f3b652"
isSP := IsServicePrincipal(userId)
user := &iam.User{
UserName: "8b948b2e-d2b5-4b9e-8274-11b596f3b652",
}
isSP := IsServicePrincipal(user)
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)
user := &iam.User{
UserName: "invalid",
}
isSP := IsServicePrincipal(user)
assert.False(t, isSP, "Expected user ID to not be recognized as a service principal")
}

View File

@ -1,4 +1,4 @@
package auth
package iamutil
import (
"strings"
@ -12,7 +12,7 @@ import (
// including dots, which are not supported in e.g. experiment names.
func GetShortUserName(user *iam.User) string {
name := user.UserName
if IsServicePrincipal(user.UserName) && user.DisplayName != "" {
if IsServicePrincipal(user) && user.DisplayName != "" {
name = user.DisplayName
}
local, _, _ := strings.Cut(name, "@")

View File

@ -1,4 +1,4 @@
package auth
package iamutil
import (
"testing"

View File

@ -11,7 +11,7 @@ import (
"text/template"
"github.com/databricks/cli/cmd/root"
"github.com/databricks/cli/libs/auth"
"github.com/databricks/cli/libs/iamutil"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/service/iam"
@ -119,7 +119,7 @@ func loadHelpers(ctx context.Context) template.FuncMap {
return "", err
}
}
return auth.GetShortUserName(cachedUser), nil
return iamutil.GetShortUserName(cachedUser), nil
},
// Get the default workspace catalog. If there is no default, or if
// Unity Catalog is not enabled, return an empty string.
@ -151,7 +151,7 @@ func loadHelpers(ctx context.Context) template.FuncMap {
return false, err
}
}
result := auth.IsServicePrincipal(cachedUser.UserName)
result := iamutil.IsServicePrincipal(cachedUser)
cachedIsServicePrincipal = &result
return result, nil
},