Merge pull request #76 from atlanhq/FT-914

FT-914 : Implement UserCache and GroupCache

atlan/assets/group_cache.go

@@ -0,0 +1,141 @@
+package assets
+
+import (
+	"errors"
+	"sync"
+)
+
+// GroupCache provides a lazily-loaded cache for translating Atlan-internal groups into their IDs and names.
+type GroupCache struct {
+	groupClient  *GroupClient
+	cacheByID    map[string]AtlanGroup
+	mapIDToName  map[string]string
+	mapNameToID  map[string]string
+	mapAliasToID map[string]string
+	mutex        sync.Mutex
+}
+
+var (
+	groupCaches = make(map[string]*GroupCache)
+	groupMutex  sync.Mutex
+)
+
+// GetGroupCache retrieves the GroupCache for the default Atlan client.
+func GetGroupCache() (*GroupCache, error) {
+	groupMutex.Lock()
+	defer groupMutex.Unlock()
+
+	client := DefaultAtlanClient
+	cacheKey := generateCacheKey(client.host, client.ApiKey)
+
+	if groupCaches[cacheKey] == nil {
+		groupCaches[cacheKey] = &GroupCache{
+			groupClient:  client.GroupClient,
+			cacheByID:    make(map[string]AtlanGroup),
+			mapIDToName:  make(map[string]string),
+			mapNameToID:  make(map[string]string),
+			mapAliasToID: make(map[string]string),
+		}
+	}
+	return groupCaches[cacheKey], nil
+}
+
+// GetGroupIDForGroupName translates the provided group name to its GUID.
+func GetGroupIDForGroupName(name string) (string, error) {
+	cache, err := GetGroupCache()
+	if err != nil {
+		return "", err
+	}
+	return cache.getIDForName(name), nil
+}
+
+// GetGroupIDForAlias translates the provided group alias to its GUID.
+func GetGroupIDForAlias(alias string) (string, error) {
+	cache, err := GetGroupCache()
+	if err != nil {
+		return "", err
+	}
+	return cache.getIDForAlias(alias), nil
+}
+
+// GetGroupNameForGroupID translates the provided group GUID to its name.
+func GetGroupNameForGroupID(id string) (string, error) {
+	cache, err := GetGroupCache()
+	if err != nil {
+		return "", err
+	}
+	return cache.getNameForID(id), nil
+}
+
+// ValidateGroupAliases validates that the given group aliases are valid.
+func ValidateGroupAliases(aliases []string) error {
+	cache, err := GetGroupCache()
+	if err != nil {
+		return err
+	}
+	return cache.validateAliases(aliases)
+}
+
+func (gc *GroupCache) refreshCache() error {
+	gc.mutex.Lock()
+	defer gc.mutex.Unlock()
+
+	groups, err := gc.groupClient.GetAll(20, 0, "")
+	if err != nil {
+		return err
+	}
+
+	gc.cacheByID = make(map[string]AtlanGroup)
+	gc.mapIDToName = make(map[string]string)
+	gc.mapNameToID = make(map[string]string)
+	gc.mapAliasToID = make(map[string]string)
+
+	for _, group := range groups {
+		groupID := *group.ID
+		groupName := *group.Name
+		groupAlias := *group.Alias
+
+		gc.cacheByID[groupID] = *group
+		gc.mapIDToName[groupID] = groupName
+		gc.mapNameToID[groupName] = groupID
+		gc.mapAliasToID[groupAlias] = groupID
+	}
+
+	return nil
+}
+
+func (gc *GroupCache) getIDForName(name string) string {
+	if id, exists := gc.mapNameToID[name]; exists {
+		return id
+	}
+	gc.refreshCache()
+	return gc.mapNameToID[name]
+}
+
+func (gc *GroupCache) getIDForAlias(alias string) string {
+	if id, exists := gc.mapAliasToID[alias]; exists {
+		return id
+	}
+	gc.refreshCache()
+	return gc.mapAliasToID[alias]
+}
+
+func (gc *GroupCache) getNameForID(id string) string {
+	if name, exists := gc.mapIDToName[id]; exists {
+		return name
+	}
+	gc.refreshCache()
+	return gc.mapIDToName[id]
+}
+
+func (gc *GroupCache) validateAliases(aliases []string) error {
+	for _, alias := range aliases {
+		if _, exists := gc.mapAliasToID[alias]; !exists {
+			gc.refreshCache()
+			if _, exists := gc.mapAliasToID[alias]; !exists {
+				return errors.New("provided group alias not found in Atlan")
+			}
+		}
+	}
+	return nil
+}

atlan/assets/user_cache.go

@@ -0,0 +1,161 @@
+package assets
+
+import (
+	"errors"
+	"sync"
+)
+
+// UserCache provides a lazily-loaded cache for translating Atlan-internal users into their IDs and names.
+type UserCache struct {
+	userClient   *UserClient
+	tokenClient  *TokenClient
+	mapIDToName  map[string]string
+	mapNameToID  map[string]string
+	mapEmailToID map[string]string
+	mutex        sync.Mutex
+}
+
+var (
+	userCaches = make(map[string]*UserCache)
+	userMutex  sync.Mutex
+)
+
+// GetUserCache retrieves the UserCache for the default Atlan client.
+func GetUserCache() (*UserCache, error) {
+	userMutex.Lock()
+	defer userMutex.Unlock()
+
+	client := DefaultAtlanClient
+	cacheKey := generateCacheKey(client.host, client.ApiKey)
+
+	if userCaches[cacheKey] == nil {
+		userCaches[cacheKey] = &UserCache{
+			userClient:   client.UserClient,
+			tokenClient:  client.TokenClient,
+			mapIDToName:  make(map[string]string),
+			mapNameToID:  make(map[string]string),
+			mapEmailToID: make(map[string]string),
+		}
+	}
+	return userCaches[cacheKey], nil
+}
+
+// GetUserIDForName translates the provided human-readable username to its GUID.
+func GetUserIDForName(name string) (string, error) {
+	cache, err := GetUserCache()
+	if err != nil {
+		return "", err
+	}
+	return cache.getIDForName(name)
+}
+
+// GetUserIDForEmail translates the provided email to its GUID.
+func GetUserIDForEmail(email string) (string, error) {
+	cache, err := GetUserCache()
+	if err != nil {
+		return "", err
+	}
+	return cache.getIDForEmail(email)
+}
+
+// GetUserNameForID translates the provided user GUID to the human-readable username.
+func GetUserNameForID(id string) (string, error) {
+	cache, err := GetUserCache()
+	if err != nil {
+		return "", err
+	}
+	return cache.getNameForID(id)
+}
+
+// ValidateUserNames validates that the given human-readable usernames are valid.
+func ValidateUserNames(names []string) error {
+	cache, err := GetUserCache()
+	if err != nil {
+		return err
+	}
+	return cache.validateNames(names)
+}
+
+func (uc *UserCache) refreshCache() error {
+	uc.mutex.Lock()
+	defer uc.mutex.Unlock()
+
+	users, err := uc.userClient.GetAll(20, 0, "")
+	if err != nil {
+		return err
+	}
+
+	uc.mapIDToName = make(map[string]string)
+	uc.mapNameToID = make(map[string]string)
+	uc.mapEmailToID = make(map[string]string)
+
+	for _, user := range users {
+		userID := user.ID
+		userName := user.Username
+		userEmail := user.Email
+
+		uc.mapIDToName[userID] = *userName
+		uc.mapNameToID[*userName] = userID
+		uc.mapEmailToID[userEmail] = userID
+	}
+
+	return nil
+}
+
+func (uc *UserCache) getIDForName(name string) (string, error) {
+	if id, exists := uc.mapNameToID[name]; exists {
+		return id, nil
+	}
+	// If the name is an API token, try fetching it directly
+	if isServiceAccount(name) {
+		token, err := uc.tokenClient.GetByID(name)
+		if err != nil {
+			return "", err
+		}
+		if token != nil && token.GUID != nil {
+			uc.mapNameToID[name] = *token.GUID
+			return *token.GUID, nil
+		}
+		return "", errors.New("API token not found by name")
+	}
+	uc.refreshCache()
+	return uc.mapNameToID[name], nil
+}
+
+func (uc *UserCache) getIDForEmail(email string) (string, error) {
+	if id, exists := uc.mapEmailToID[email]; exists {
+		return id, nil
+	}
+	uc.refreshCache()
+	return uc.mapEmailToID[email], nil
+}
+
+func (uc *UserCache) getNameForID(id string) (string, error) {
+	if name, exists := uc.mapIDToName[id]; exists {
+		return name, nil
+	}
+	// If the ID is an API token, try fetching it directly
+	token, err := uc.tokenClient.GetByGUID(id)
+	if err != nil {
+		return "", err
+	}
+	if token != nil && token.ClientID != nil {
+		return *token.ClientID, nil
+	}
+	uc.refreshCache()
+	return uc.mapIDToName[id], nil
+}
+
+func (uc *UserCache) validateNames(names []string) error {
+	for _, name := range names {
+		if _, err := uc.getIDForName(name); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Helper function to check if a name is a service account
+func isServiceAccount(name string) bool {
+	return len(name) > len("service-account-") && name[:len("service-account-")] == "service-account-"
+}

atlan/model/structs/user.go

@@ -86,7 +86,7 @@ type UserRequest struct {
 func (r *UserRequest) QueryParams() map[string]interface{} {
 	qp := make(map[string]interface{})
 
-	if r.PostFilter != nil {
+	if r.PostFilter != nil && *r.PostFilter != "" {
 		qp["filter"] = *r.PostFilter
 	}
 	if r.Sort != nil && *r.Sort != "" {

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	_ "github.com/atlanhq/atlan-go/atlan"
 	"github.com/atlanhq/atlan-go/atlan/assets"
 	_ "github.com/atlanhq/atlan-go/atlan/model/structs"
@@ -11,6 +12,13 @@ func main() {
 	ctx := assets.NewContext()
 	ctx.EnableLogging("debug")
 
+	// Test User Cache
+	UserId, err := assets.GetGroupNameForGroupID("58d547d8-3f4d-4b9e-9666-39980f140661")
+	if err != nil {
+		fmt.Println("Error:", err)
+	}
+	fmt.Println(UserId)
+
 	/*
 		// Delete an API Token
 		err := ctx.TokenClient.Purge("a853f1d5-f1f4-4cdb-b86d-c61df3ecade6")