Merge pull request #50 from ConductorOne/jallers/fix-user-sync

Improve pagination when listing users

Author: johnallers

pkg/connector/connector.go

@@ -53,7 +53,7 @@ func (d *Connector) Metadata(ctx context.Context) (*v2.ConnectorMetadata, error)
 // Validate is called to ensure that the connector is properly configured. It should exercise any API credentials
 // to be sure that they are valid.
 func (d *Connector) Validate(ctx context.Context) (annotations.Annotations, error) {
-	users, _, err := d.Client.ListUsers(ctx, 0, 1)
+	users, _, err := d.Client.ListUsers(ctx, "", 1)
 	if err != nil {
 		return nil, err
 	}

pkg/connector/pagination.go

@@ -30,6 +30,23 @@ func parseOffsetFromToken(i string, resourceID *v2.ResourceId) (*pagination.Bag,
 	return b, offset, nil
 }
 
+func parseCursorFromToken(i string, resourceID *v2.ResourceId) (*pagination.Bag, string, error) {
+	b := &pagination.Bag{}
+	err := b.Unmarshal(i)
+	if err != nil {
+		return nil, "", err
+	}
+
+	if b.Current() == nil {
+		b.Push(pagination.PageState{
+			ResourceTypeID: resourceID.ResourceType,
+			ResourceID:     resourceID.Resource,
+		})
+	}
+
+	return b, b.PageToken(), nil
+}
+
 func convertPageToken(token string) (int, error) {
 	if token == "" {
 		return 0, nil

pkg/connector/users.go

@@ -97,12 +97,12 @@ func getUserDetailedStatus(user *snowflake.User) string {
 // List returns all the users from the database as resource objects.
 // Users include a UserTrait because they are the 'shape' of a standard user.
 func (o *userBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId, pToken *pagination.Token) ([]*v2.Resource, string, annotations.Annotations, error) {
-	bag, offset, err := parseOffsetFromToken(pToken.Token, &v2.ResourceId{ResourceType: o.resourceType.Id})
+	bag, cursor, err := parseCursorFromToken(pToken.Token, &v2.ResourceId{ResourceType: o.resourceType.Id})
 	if err != nil {
-		return nil, "", nil, wrapError(err, "failed to get next page offset")
+		return nil, "", nil, wrapError(err, "failed to get next page cursor")
 	}
 
-	users, _, err := o.client.ListUsers(ctx, offset, resourcePageSize)
+	users, _, err := o.client.ListUsers(ctx, cursor, resourcePageSize)
 	if err != nil {
 		return nil, "", nil, wrapError(err, "failed to list users")
 	}
@@ -125,12 +125,12 @@ func (o *userBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId,
 		return resources, "", nil, nil
 	}
 
-	nextPage, err := handleNextPage(bag, offset+resourcePageSize)
+	nextCursor, err := bag.NextToken(users[len(users)-1].Username)
 	if err != nil {
 		return nil, "", nil, wrapError(err, "failed to create next page cursor")
 	}
 
-	return resources, nextPage, nil, nil
+	return resources, nextCursor, nil, nil
 }
 
 // Entitlements always returns an empty slice for users.

pkg/snowflake/user.go

@@ -181,10 +181,12 @@ func (r *GetUserRawResponse) GetValueByColumnName(columnName string) (string, bo
 	return "", false
 }
 
-func (c *Client) ListUsers(ctx context.Context, offset, limit int) ([]User, *http.Response, error) {
-	queries := []string{
-		"SHOW USERS;",
-		c.paginateLastQuery(offset, limit),
+func (c *Client) ListUsers(ctx context.Context, cursor string, limit int) ([]User, *http.Response, error) {
+	queries := []string{}
+	if cursor != "" {
+		queries = append(queries, fmt.Sprintf("SHOW USERS LIMIT %d FROM '%s';", limit, cursor))
+	} else {
+		queries = append(queries, fmt.Sprintf("SHOW USERS LIMIT %d;", limit))
 	}
 
 	req, err := c.PostStatementRequest(ctx, queries)
@@ -198,11 +200,7 @@ func (c *Client) ListUsers(ctx context.Context, offset, limit int) ([]User, *htt
 		return nil, resp, err
 	}
 
-	if len(response.StatementHandles) < 2 {
-		return nil, resp, nil
-	}
-
-	req, err = c.GetStatementResponse(ctx, response.StatementHandles[1])
+	req, err = c.GetStatementResponse(ctx, response.StatementHandle)
 	if err != nil {
 		return nil, resp, err
 	}