feat(core): improve entitlements performance

service/authorization/authorization.go

@@ -20,6 +20,7 @@ import (
 	"github.com/opentdf/platform/protocol/go/entityresolution"
 	"github.com/opentdf/platform/protocol/go/policy"
 	attr "github.com/opentdf/platform/protocol/go/policy/attributes"
+	"github.com/opentdf/platform/protocol/go/policy/subjectmapping"
 	otdf "github.com/opentdf/platform/sdk"
 	"github.com/opentdf/platform/service/internal/access"
 	"github.com/opentdf/platform/service/internal/entitlements"
@@ -375,40 +376,107 @@ func (as *AuthorizationService) GetDecisions(ctx context.Context, req *authoriza
 	return rsp, nil
 }
 
-func (as *AuthorizationService) GetEntitlements(ctx context.Context, req *authorization.GetEntitlementsRequest) (*authorization.GetEntitlementsResponse, error) {
-	as.logger.DebugContext(ctx, "getting entitlements")
-	request := attr.GetAttributeValuesByFqnsRequest{
-		WithValue: &policy.AttributeValueSelector{
-			WithSubjectMaps: true,
-		},
+// makeSubMapsByValLookup creates a lookup map of subject mappings by attribute value ID.
+func makeSubMapsByValLookup(subjectMappings []*policy.SubjectMapping) map[string][]*policy.SubjectMapping {
+	// map keys will be attribute value IDs
+	lookup := make(map[string][]*policy.SubjectMapping)
+	for _, sm := range subjectMappings {
+		val := sm.GetAttributeValue()
+		id := val.GetId()
+		// if attribute value ID exists
+		if id != "" {
+			// append the subject mapping to the slice of subject mappings for the attribute value ID
+			lookup[id] = append(lookup[id], sm)
+		}
 	}
-	// Lack of scope has impacts on performance
-	// https://github.com/opentdf/platform/issues/365
-	if req.GetScope() == nil {
-		// TODO: Reomve and use MatchSubjectMappings instead later in the flow
-		listAttributeResp, err := as.sdk.Attributes.ListAttributes(ctx, &attr.ListAttributesRequest{})
-		if err != nil {
-			as.logger.ErrorContext(ctx, "failed to list attributes", slog.String("error", err.Error()))
-			return nil, status.Error(codes.Internal, "failed to list attributes")
+	return lookup
+}
+
+// updateValsWithSubMaps updates the subject mappings of values using the lookup map.
+func updateValsWithSubMaps(values []*policy.Value, subMapsByVal map[string][]*policy.SubjectMapping) []*policy.Value {
+	for i, v := range values {
+		// if subject mappings exist for the value
+		if subjectMappings, ok := subMapsByVal[v.GetId()]; ok {
+			// update the subject mappings of the value
+			values[i].SubjectMappings = subjectMappings
 		}
-		var attributeFqns []string
-		for _, attr := range listAttributeResp.GetAttributes() {
-			for _, val := range attr.GetValues() {
-				attributeFqns = append(attributeFqns, val.GetFqn())
-			}
+	}
+	return values
+}
+
+// updateValsByFqnLookup updates the lookup map with attribute values by FQN.
+func updateValsByFqnLookup(attribute *policy.Attribute, scopeMap map[string]bool, fqnAttrVals map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue) map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue {
+	rule := attribute.GetRule()
+	for _, v := range attribute.GetValues() {
+		// if scope exists and current attribute value FQN is not in scope
+		if !(scopeMap == nil || scopeMap[v.GetFqn()]) {
+			// skip
+			continue
+		}
+		// trim attribute values (by default only keep single value relevant to FQN)
+		// This is key to minimizing the rego query size.
+		values := []*policy.Value{v}
+		if rule == policy.AttributeRuleTypeEnum_ATTRIBUTE_RULE_TYPE_ENUM_HIERARCHY {
+			// restore ALL attribute values if attribute rule is hierarchical
+			// This is key to honoring comprehensive hierarchy.
+			values = attribute.GetValues()
 		}
-		request.Fqns = attributeFqns
-	} else {
-		// get subject mappings
-		request.Fqns = req.GetScope().GetAttributeValueFqns()
+		// only clone relevant fields for attribute
+		a := &policy.Attribute{Rule: rule, Values: values}
+		fqnAttrVals[v.GetFqn()] = &attr.GetAttributeValuesByFqnsResponse_AttributeAndValue{Attribute: a, Value: v}
+	}
+	return fqnAttrVals
+}
+
+// makeValsByFqnsLookup creates a lookup map of attribute values by FQN.
+func makeValsByFqnsLookup(attrs []*policy.Attribute, subMapsByVal map[string][]*policy.SubjectMapping, scopeMap map[string]bool) map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue {
+	fqnAttrVals := make(map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue)
+	for i := range attrs {
+		// add subject mappings to attribute values
+		attrs[i].Values = updateValsWithSubMaps(attrs[i].GetValues(), subMapsByVal)
+		// update the lookup map with attribute values by FQN
+		fqnAttrVals = updateValsByFqnLookup(attrs[i], scopeMap, fqnAttrVals)
+	}
+	return fqnAttrVals
+}
+
+// makeScopeMap creates a map of attribute value FQNs.
+func makeScopeMap(scope *authorization.ResourceAttribute) map[string]bool {
+	// if scope not defined, return nil pointer
+	if scope == nil {
+		return nil
+	}
+	scopeMap := make(map[string]bool)
+	// add attribute value FQNs from scope to the map
+	for _, fqn := range scope.GetAttributeValueFqns() {
+		scopeMap[fqn] = true
+	}
+	return scopeMap
+}
+
+func (as *AuthorizationService) GetEntitlements(ctx context.Context, req *authorization.GetEntitlementsRequest) (*authorization.GetEntitlementsResponse, error) {
+	as.logger.DebugContext(ctx, "getting entitlements")
+	attrsRes, err := as.sdk.Attributes.ListAttributes(ctx, &attr.ListAttributesRequest{})
+	if err != nil {
+		as.logger.ErrorContext(ctx, "failed to list attributes", slog.String("error", err.Error()))
+		return nil, status.Error(codes.Internal, "failed to list attributes")
 	}
-	avf, err := as.sdk.Attributes.GetAttributeValuesByFqns(ctx, &request)
+	subMapsRes, err := as.sdk.SubjectMapping.ListSubjectMappings(ctx, &subjectmapping.ListSubjectMappingsRequest{})
 	if err != nil {
-		as.logger.ErrorContext(ctx, "failed to get attribute values by fqns", slog.String("error", err.Error()))
-		return nil, status.Error(codes.Internal, "failed to get attribute values by fqns")
+		as.logger.ErrorContext(ctx, "failed to list subject mappings", slog.String("error", err.Error()))
+		return nil, status.Error(codes.Internal, "failed to list subject mappings")
+	}
+	// create a lookup map of attribute value FQNs (based on request scope)
+	scopeMap := makeScopeMap(req.GetScope())
+	// create a lookup map of subject mappings by attribute value ID
+	subMapsByVal := makeSubMapsByValLookup(subMapsRes.GetSubjectMappings())
+	// create a lookup map of attribute values by FQN (for rego query)
+	fqnAttrVals := makeValsByFqnsLookup(attrsRes.GetAttributes(), subMapsByVal, scopeMap)
+	avf := &attr.GetAttributeValuesByFqnsResponse{
+		FqnAttributeValues: fqnAttrVals,
 	}
 	subjectMappings := avf.GetFqnAttributeValues()
-	as.logger.DebugContext(ctx, "retrieved from subject mappings service", slog.Any("subject_mappings: ", subjectMappings))
+	as.logger.DebugContext(ctx, fmt.Sprintf("retrieved %d subject mappings", len(subjectMappings)))
 	// TODO: this could probably be moved to proto validation https://github.com/opentdf/platform/issues/1057
 	if req.Entities == nil {
 		as.logger.ErrorContext(ctx, "requires entities")
@@ -431,7 +499,6 @@ func (as *AuthorizationService) GetEntitlements(ctx context.Context, req *author
 		as.logger.ErrorContext(ctx, "failed to build rego input", slog.String("error", err.Error()))
 		return nil, status.Error(codes.Internal, "failed to build rego input")
 	}
-	as.logger.DebugContext(ctx, "entitlements", "input", fmt.Sprintf("%+v", in))
 
 	results, err := as.eval.Eval(ctx,
 		rego.EvalInput(in),
@@ -463,7 +530,6 @@ func (as *AuthorizationService) GetEntitlements(ctx context.Context, req *author
 		return rsp, nil
 	}
 	as.logger.DebugContext(ctx, "opa results", "results", fmt.Sprintf("%+v", results))
-
 	for idx, entity := range req.GetEntities() {
 		// Ensure the entity has an ID
 		entityID := entity.GetId()

service/authorization/authorization_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/opentdf/platform/protocol/go/entityresolution"
 	"github.com/opentdf/platform/protocol/go/policy"
 	attr "github.com/opentdf/platform/protocol/go/policy/attributes"
+	sm "github.com/opentdf/platform/protocol/go/policy/subjectmapping"
 	otdf "github.com/opentdf/platform/sdk"
 	"github.com/opentdf/platform/service/logger"
 	"github.com/stretchr/testify/assert"
@@ -25,6 +26,7 @@ import (
 var (
 	getAttributesByValueFqnsResponse attr.GetAttributeValuesByFqnsResponse
 	listAttributeResp                attr.ListAttributesResponse
+	listSubjectMappings              sm.ListSubjectMappingsResponse
 	createEntityChainResp            entityresolution.CreateEntityChainFromJwtResponse
 	resolveEntitiesResp              entityresolution.ResolveEntitiesResponse
 	mockNamespace                    = "www.example.org"
@@ -51,6 +53,14 @@ type myERSClient struct {
 	entityresolution.EntityResolutionServiceClient
 }
 
+type mySubjectMappingClient struct {
+	sm.SubjectMappingServiceClient
+}
+
+func (*mySubjectMappingClient) ListSubjectMappings(_ context.Context, _ *sm.ListSubjectMappingsRequest, _ ...grpc.CallOption) (*sm.ListSubjectMappingsResponse, error) {
+	return &listSubjectMappings, nil
+}
+
 func (*myERSClient) CreateEntityChainFromJwt(_ context.Context, _ *entityresolution.CreateEntityChainFromJwtRequest, _ ...grpc.CallOption) (*entityresolution.CreateEntityChainFromJwtResponse, error) {
 	return &createEntityChainResp, nil
 }
@@ -214,7 +224,8 @@ func Test_GetDecisionsAllOf_Pass(t *testing.T) {
 	}}
 
 	as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
-		Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
+		SubjectMapping: &mySubjectMappingClient{},
+		Attributes:     &myAttributesClient{}, EntityResoution: &myERSClient{}},
 		tokenSource: &testTokenSource, eval: prepared}
 
 	resp, err := as.GetDecisions(ctxb, &req)
@@ -372,7 +383,8 @@ func Test_GetDecisions_AllOf_Fail(t *testing.T) {
 	require.NoError(t, err)
 
 	as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
-		Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
+		SubjectMapping: &mySubjectMappingClient{},
+		Attributes:     &myAttributesClient{}, EntityResoution: &myERSClient{}},
 		tokenSource: &testTokenSource, eval: prepared}
 
 	resp, err := as.GetDecisions(ctxb, &req)
@@ -467,7 +479,8 @@ func Test_GetDecisionsAllOfWithEnvironmental_Pass(t *testing.T) {
 	}}
 
 	as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
-		Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
+		SubjectMapping: &mySubjectMappingClient{},
+		Attributes:     &myAttributesClient{}, EntityResoution: &myERSClient{}},
 		tokenSource: &testTokenSource, eval: prepared}
 
 	resp, err := as.GetDecisions(ctxb, &req)
@@ -560,7 +573,8 @@ func Test_GetDecisionsAllOfWithEnvironmental_Fail(t *testing.T) {
 	}}
 
 	as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
-		Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
+		SubjectMapping: &mySubjectMappingClient{},
+		Attributes:     &myAttributesClient{}, EntityResoution: &myERSClient{}},
 		tokenSource: &testTokenSource, eval: prepared}
 
 	resp, err := as.GetDecisions(ctxb, &req)
@@ -634,7 +648,8 @@ func Test_GetEntitlementsSimple(t *testing.T) {
 	require.NoError(t, err)
 
 	as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
-		Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
+		SubjectMapping: &mySubjectMappingClient{},
+		Attributes:     &myAttributesClient{}, EntityResoution: &myERSClient{}},
 		tokenSource: &testTokenSource, eval: prepared}
 
 	req := authorization.GetEntitlementsRequest{
@@ -653,8 +668,6 @@ func Test_GetEntitlementsSimple(t *testing.T) {
 
 func Test_GetEntitlementsWithComprehensiveHierarchy(t *testing.T) {
 	logger := logger.CreateTestLogger()
-
-	listAttributeResp = attr.ListAttributesResponse{}
 	attrDef := policy.Attribute{
 		Name: mockAttrName,
 		Namespace: &policy.Namespace{
@@ -672,6 +685,7 @@ func Test_GetEntitlementsWithComprehensiveHierarchy(t *testing.T) {
 			},
 		},
 	}
+	listAttributeResp.Attributes = []*policy.Attribute{&attrDef}
 	getAttributesByValueFqnsResponse = attr.GetAttributeValuesByFqnsResponse{FqnAttributeValues: map[string]*attr.GetAttributeValuesByFqnsResponse_AttributeAndValue{
 		"https://www.example.org/attr/foo/value/value1": {
 			Attribute: &attrDef,
@@ -712,7 +726,8 @@ func Test_GetEntitlementsWithComprehensiveHierarchy(t *testing.T) {
 	prepared, err := rego.PrepareForEval(ctxb)
 	require.NoError(t, err)
 	as := AuthorizationService{logger: logger, sdk: &otdf.SDK{
-		Attributes: &myAttributesClient{}, EntityResoution: &myERSClient{}},
+		SubjectMapping: &mySubjectMappingClient{},
+		Attributes:     &myAttributesClient{}, EntityResoution: &myERSClient{}},
 		tokenSource: &testTokenSource, eval: prepared}
 
 	withHierarchy := true

service/policy/db/subject_mappings.go

@@ -189,15 +189,6 @@ func subjectMappingHydrateItem(row pgx.Row, logger *logger.Logger) (*policy.Subj
 		&scsJSON,
 		&attributeValueJSON,
 	)
-	logger.Debug(
-		"subjectMappingHydrateItem",
-		slog.Any("row", row),
-		slog.String("id", id),
-		slog.String("actionsJSON", string(actionsJSON)),
-		slog.String("metadataJSON", string(metadataJSON)),
-		slog.String("scsJSON", string(scsJSON)),
-		slog.String("attributeValueJSON", string(attributeValueJSON)),
-	)
 	if err != nil {
 		return nil, db.WrapIfKnownInvalidQueryErr(err)
 	}