Improve Authorize performance

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/coredata/membership.go

@@ -19,7 +19,6 @@ import (
 	"errors"
 	"fmt"
 	"maps"
-	"strings"
 	"time"
 
 	"github.com/jackc/pgx/v5"
@@ -227,81 +226,63 @@ LEFT JOIN
 	return nil
 }
 
-// LoadRoleByIdentityAndEntityID loads an identity's role by querying any entity to extract its organization_id
-func (m *Membership) LoadRoleByIdentityAndEntityID(
+func LoadRoleByIdentityAndEntityIDOnly(
 	ctx context.Context,
 	conn pg.Conn,
 	scope Scoper,
 	identityID gid.GID,
 	entityID gid.GID,
-) error {
+) (MembershipRole, error) {
 	entityType := entityID.EntityType()
 
-	// For organization, the entity ID is the organization ID
+	// For organization, the entity ID is the organization ID - optimized path
 	if entityType == OrganizationEntityType {
-		return m.LoadByIdentityAndOrg(ctx, conn, scope, identityID, entityID)
+		query := `
+SELECT role
+FROM iam_memberships
+WHERE
+    identity_id = $1
+    AND tenant_id = $2
+    AND organization_id = $3
+LIMIT 1;
+`
+		var role MembershipRole
+		err := conn.QueryRow(ctx, query, identityID, scope.GetTenantID(), entityID).Scan(&role)
+		if err != nil {
+			if errors.Is(err, pgx.ErrNoRows) {
+				return "", ErrResourceNotFound
+			}
+			return "", fmt.Errorf("cannot query role: %w", err)
+		}
+		return role, nil
 	}
 
 	tableName, ok := EntityTable(entityType)
 	if !ok {
-		return fmt.Errorf("unsupported entity type for role lookup: %d", entityType)
+		return "", fmt.Errorf("unsupported entity type for role lookup: %d", entityType)
 	}
 
-	// Build scope fragment with table alias to avoid ambiguity
-	scopeFragment := scope.SQLFragment()
-	// Replace column references with table-qualified versions
-	scopeFragment = strings.ReplaceAll(scopeFragment, "tenant_id =", "m.tenant_id =")
-
-	query := fmt.Sprintf(`
-SELECT
-    m.id,
-    m.identity_id,
-    m.organization_id,
-    m.role,
-    m.created_at,
-    m.updated_at
-FROM
-    iam_memberships m
-    INNER JOIN %s e ON e.id = @entity_id
+	query := `
+SELECT m.role
+FROM iam_memberships m
 WHERE
-    %s
-    AND m.identity_id = @identity_id
-    AND m.organization_id = e.organization_id
+    m.identity_id = $1
+    AND tenant_id = $2
+    AND m.organization_id = (SELECT organization_id FROM ` + tableName + ` WHERE id = $3)
 LIMIT 1;
-`, tableName, scopeFragment)
+`
 
-	args := pgx.NamedArgs{
-		"identity_id": identityID,
-		"entity_id":   entityID,
-	}
-	maps.Copy(args, scope.SQLArguments())
-
-	rows, err := conn.Query(ctx, query, args)
+	var role MembershipRole
+	err := conn.QueryRow(ctx, query, identityID, scope.GetTenantID(), entityID).Scan(&role)
 	if err != nil {
-		return fmt.Errorf("cannot query membership by entity: %w", err)
-	}
-	defer rows.Close()
+		if errors.Is(err, pgx.ErrNoRows) {
+			return "", ErrResourceNotFound
+		}
 
-	if !rows.Next() {
-		return ErrResourceNotFound
+		return "", fmt.Errorf("cannot query role: %w", err)
 	}
 
-	var membership Membership
-	err = rows.Scan(
-		&membership.ID,
-		&membership.IdentityID,
-		&membership.OrganizationID,
-		&membership.Role,
-		&membership.CreatedAt,
-		&membership.UpdatedAt,
-	)
-
-	if err != nil {
-		return fmt.Errorf("cannot scan membership: %w", err)
-	}
-
-	*m = membership
-	return nil
+	return role, nil
 }
 
 func (m *Membership) LoadByIdentityAndOrg(

pkg/coredata/migrations/20251223T190400Z.sql

@@ -0,0 +1,2 @@
+CREATE INDEX IF NOT EXISTS idx_iam_memberships_identity_tenant_org ON iam_memberships(identity_id, tenant_id, organization_id);
+

pkg/iam/authorizer.go

@@ -57,17 +57,18 @@ func (a *Authorizer) Authorize(ctx context.Context, params AuthorizeParams) erro
 
 	policies := a.buildPolicies(ctx, params)
 
+	// Pre-allocate Resource map with capacity for id + attributes
+	resourceAttrs := make(map[string]string, 1+len(params.ResourceAttributes))
+	resourceAttrs["id"] = params.Resource.String()
+	maps.Copy(resourceAttrs, params.ResourceAttributes)
+
 	conditionCtx := policy.ConditionContext{
 		Principal: map[string]string{
 			"id": params.Principal.String(),
 		},
-		Resource: map[string]string{
-			"id": params.Resource.String(),
-		},
+		Resource: resourceAttrs,
 	}
 
-	maps.Copy(conditionCtx.Resource, params.ResourceAttributes)
-
 	req := policy.AuthorizationRequest{
 		Principal:        params.Principal,
 		Resource:         params.Resource,
@@ -81,53 +82,44 @@ func (a *Authorizer) Authorize(ctx context.Context, params AuthorizeParams) erro
 		return nil
 	}
 
-	if result.Decision == policy.DecisionDeny {
-		return NewInsufficientPermissionsError(params.Principal, params.Resource, params.Action)
-	}
-
 	return NewInsufficientPermissionsError(params.Principal, params.Resource, params.Action)
 }
 
-// buildPolicies constructs the list of policies to evaluate.
-// This includes self-management policies and role-based policies.
 func (a *Authorizer) buildPolicies(ctx context.Context, params AuthorizeParams) []*policy.Policy {
-	// Start with self-management policies
-	policies := make([]*policy.Policy, len(a.policySet.SelfManagePolicies))
-	copy(policies, a.policySet.SelfManagePolicies)
+	selfManageCount := len(a.policySet.SelfManagePolicies)
 
-	// For organization-scoped resources, add role-based policies
+	var rolePolicies []*policy.Policy
 	if params.Resource.TenantID() != gid.NilTenant {
-		rolePolicies := a.loadRolePolicies(ctx, params.Principal, params.Resource)
-		policies = append(policies, rolePolicies...)
+		rolePolicies = a.loadRolePolicies(ctx, params.Principal, params.Resource)
 	}
 
+	totalCount := selfManageCount + len(rolePolicies)
+	policies := make([]*policy.Policy, selfManageCount, totalCount)
+	copy(policies, a.policySet.SelfManagePolicies)
+	policies = append(policies, rolePolicies...)
+
 	return policies
 }
 
-// loadRolePolicies loads the role-based policies for a user in an organization.
 func (a *Authorizer) loadRolePolicies(ctx context.Context, principalID gid.GID, resourceID gid.GID) []*policy.Policy {
-	var roleName string
+	var role coredata.MembershipRole
 
-	err := a.pg.WithConn(ctx, func(conn pg.Conn) error {
-		scope := coredata.NewScopeFromObjectID(resourceID)
-
-		var m coredata.Membership
-		if err := m.LoadRoleByIdentityAndEntityID(ctx, conn, scope, principalID, resourceID); err != nil {
+	err := a.pg.WithConn(
+		ctx,
+		func(conn pg.Conn) (err error) {
+			scope := coredata.NewScopeFromObjectID(resourceID)
+			role, err = coredata.LoadRoleByIdentityAndEntityIDOnly(ctx, conn, scope, principalID, resourceID)
 			if errors.Is(err, coredata.ErrResourceNotFound) {
 				return nil // No membership = no role-based policies
 			}
+
 			return err
-		}
+		},
+	)
 
-		roleName = m.Role.String()
-		return nil
-	})
-
-	if err != nil || roleName == "" {
-		// On error or no role, return empty policies (fail closed)
+	if err != nil || role == "" {
 		return nil
 	}
 
-	// Get policies for the user's role
-	return a.policySet.RolePolicies[roleName]
+	return a.policySet.RolePolicies[role.String()]
 }