194 sso lambda deletes then recreates users (#203)

* Improve error handling for failed google api calls *When the google api appears to return zero groups or users, then an error is thrown to terminate execution.
awslabs · Jun 20, 2024 · a6b6541 · a6b6541
1 parent a5bb2ae
commit a6b6541
Showing 1 changed file with 33 additions and 20 deletions.
diff --git a/internal/google/client.go b/internal/google/client.go
@@ -18,6 +18,7 @@ package google
 import (
 	"context"
 	"strings"
+	"errors"
 
 	"golang.org/x/oauth2/google"
 	admin "google.golang.org/api/admin/directory/v1"
@@ -113,18 +114,24 @@ func (c *client) GetUsers(query string) ([]*admin.User, error) {
                         return nil
                 })
 		return u, err
-        }
-
-	// The Google api doesn't support multi-part queries, but we do so we need to split into an array of query strings
-	queries := strings.Split(query, ",")
-
-	// Then call the api one query at a time, appending to our list
-	for _, subQuery := range queries {
-		err = c.service.Users.List().Query(subQuery).Customer("my_customer").Pages(c.ctx, func(users *admin.Users) error {
-			u = append(u, users.Users...)
-			return nil
-		})
+        } else {
+
+	        // The Google api doesn't support multi-part queries, but we do so we need to split into an array of query strings
+		queries := strings.Split(query, ",")
+
+		// Then call the api one query at a time, appending to our list
+		for _, subQuery := range queries {
+			err = c.service.Users.List().Query(subQuery).Customer("my_customer").Pages(c.ctx, func(users *admin.Users) error {
+				u = append(u, users.Users...)
+				return nil
+			})
+		}
 	}
+
+	// Check we've got some users otherwise something is wrong.
+        if len(u) == 0 {
+                return u, errors.New("google api returned 0 users?")
+        } 
 	return u, err
 
 
@@ -159,17 +166,23 @@ func (c *client) GetGroups(query string) ([]*admin.Group, error) {
                         return nil
                 })
 		return g, err
+        } else {
+
+        	// The Google api doesn't support multi-part queries, but we do so we need to split into an array of query strings
+        	queries := strings.Split(query, ",")
+
+        	// Then call the api one query at a time, appending to our list
+        	for _, subQuery := range queries {
+			err = c.service.Groups.List().Customer("my_customer").Query(subQuery).Pages(context.TODO(), func(groups *admin.Groups) error {
+				g = append(g, groups.Groups...)
+				return nil
+			})
+		}
 	}
 
-        // The Google api doesn't support multi-part queries, but we do so we need to split into an array of query strings
-        queries := strings.Split(query, ",")
-
-        // Then call the api one query at a time, appending to our list
-        for _, subQuery := range queries {
-		err = c.service.Groups.List().Customer("my_customer").Query(subQuery).Pages(context.TODO(), func(groups *admin.Groups) error {
-			g = append(g, groups.Groups...)
-			return nil
-		})
+	// Check we've got some users otherwise something is wrong.
+	if len(g) == 0 {
+		return g, errors.New("google api return 0 groups?")
 	}
 	return g, err
 }