try to filter through objects isntead of fetching all objects (#374)

Author: adelowo

fgax/benchmark_test.go

@@ -0,0 +1,254 @@
+//go:build integration
+
+package fgax_test
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/theopenlane/iam/fgax"
+	"github.com/theopenlane/iam/fgax/testutils"
+)
+
+const benchmarkModelDSL = `
+model
+  schema 1.1
+
+type user
+
+type organization
+  relations
+    define member: [user]
+    define admin: [user]
+    define owner: [user]
+`
+
+func setupIntegrationBenchmark(b *testing.B) (*fgax.Client, context.Context, func()) {
+	b.Helper()
+
+	ctx := context.Background()
+
+	tmpModelFile := b.TempDir() + "/model.fga"
+	err := os.WriteFile(tmpModelFile, []byte(benchmarkModelDSL), 0644)
+	require.NoError(b, err)
+
+	container := testutils.NewFGATestcontainer(ctx,
+		testutils.WithStoreName("benchmark-store"),
+		testutils.WithReuse(true),
+		testutils.WithContainerName("fga-benchmark"),
+		testutils.WithModelFile(tmpModelFile),
+	)
+
+	client, err := container.NewFgaClient(ctx)
+	require.NoError(b, err)
+
+	return client, ctx, func() {
+		container.TeardownFixture()
+	}
+}
+
+func makeTuples(b *testing.B, client *fgax.Client, ctx context.Context, totalTuples, targetTuples int, targetObject string) {
+	b.Helper()
+
+	var writes []fgax.TupleKey
+
+	for i := 0; i < targetTuples; i++ {
+		writes = append(writes, fgax.TupleKey{
+			Subject: fgax.Entity{
+				Kind:       "user",
+				Identifier: fmt.Sprintf("target-user-%d", i),
+			},
+			Relation: "member",
+			Object: fgax.Entity{
+				Kind:       "organization",
+				Identifier: targetObject,
+			},
+		})
+	}
+
+	for i := 0; i < totalTuples-targetTuples; i++ {
+		writes = append(writes, fgax.TupleKey{
+			Subject: fgax.Entity{
+				Kind:       "user",
+				Identifier: fmt.Sprintf("other-user-%d", i),
+			},
+			Relation: "member",
+			Object: fgax.Entity{
+				Kind:       "organization",
+				Identifier: fmt.Sprintf("other-org-%d", i),
+			},
+		})
+	}
+
+	maxWrites := 10
+	for i := 0; i < len(writes); i += maxWrites {
+		end := i + maxWrites
+		if end > len(writes) {
+			end = len(writes)
+		}
+
+		batch := writes[i:end]
+		_, err := client.WriteTupleKeys(ctx, batch, []fgax.TupleKey{})
+		require.NoError(b, err)
+	}
+}
+
+func cleanupTuples(b *testing.B, client *fgax.Client, ctx context.Context) {
+	b.Helper()
+
+	allTuples, err := client.GetAllTuples(ctx)
+	if err != nil {
+		b.Logf("Warning: failed to get tuples for cleanup: %v", err)
+		return
+	}
+
+	if len(allTuples) == 0 {
+		return
+	}
+
+	var deletes []fgax.TupleKey
+	for _, tuple := range allTuples {
+		subject, _ := fgax.ParseEntity(tuple.Key.User)
+		object, _ := fgax.ParseEntity(tuple.Key.Object)
+
+		deletes = append(deletes, fgax.TupleKey{
+			Subject:  subject,
+			Relation: fgax.Relation(tuple.Key.Relation),
+			Object:   object,
+		})
+	}
+
+	maxWrites := 10
+	for i := 0; i < len(deletes); i += maxWrites {
+		end := i + maxWrites
+		if end > len(deletes) {
+			end = len(deletes)
+		}
+
+		batch := deletes[i:end]
+		_, err := client.WriteTupleKeys(ctx, []fgax.TupleKey{}, batch)
+		if err != nil {
+			b.Logf("Warning: failed to delete tuples in cleanup: %v", err)
+		}
+	}
+}
+
+func BenchmarkDeleteAllObjectRelations_Small(b *testing.B) {
+	client, ctx, cleanup := setupIntegrationBenchmark(b)
+	defer cleanup()
+
+	targetObject := "target-org-small"
+	totalTuples := 100
+	targetTuples := 10
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		makeTuples(b, client, ctx, totalTuples, targetTuples, targetObject)
+		b.StartTimer()
+
+		err := client.DeleteAllObjectRelations(ctx, fmt.Sprintf("organization:%s", targetObject), []string{})
+		require.NoError(b, err)
+	}
+}
+
+func BenchmarkDeleteAllObjectRelations_Medium(b *testing.B) {
+	client, ctx, cleanup := setupIntegrationBenchmark(b)
+	defer cleanup()
+
+	targetObject := "target-org-medium"
+	totalTuples := 1000
+	targetTuples := 10
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		makeTuples(b, client, ctx, totalTuples, targetTuples, targetObject)
+		b.StartTimer()
+
+		err := client.DeleteAllObjectRelations(ctx, fmt.Sprintf("organization:%s", targetObject), []string{})
+		require.NoError(b, err)
+	}
+}
+
+func BenchmarkDeleteAllObjectRelations_Large(b *testing.B) {
+	client, ctx, cleanup := setupIntegrationBenchmark(b)
+	defer cleanup()
+
+	targetObject := "target-org-large"
+	totalTuples := 5000
+	targetTuples := 20
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		makeTuples(b, client, ctx, totalTuples, targetTuples, targetObject)
+		b.StartTimer()
+
+		err := client.DeleteAllObjectRelations(ctx, fmt.Sprintf("organization:%s", targetObject), []string{})
+		require.NoError(b, err)
+	}
+}
+
+func BenchmarkGetTuplesForObject(b *testing.B) {
+	client, ctx, cleanup := setupIntegrationBenchmark(b)
+	defer cleanup()
+
+	testCases := []struct {
+		name         string
+		totalTuples  int
+		targetTuples int
+	}{
+		{"Small", 100, 10},
+		{"Medium", 1000, 10},
+		{"Large", 5000, 20},
+	}
+
+	for _, tc := range testCases {
+		b.Run(tc.name, func(b *testing.B) {
+			targetObject := fmt.Sprintf("target-org-%s", tc.name)
+
+			makeTuples(b, client, ctx, tc.totalTuples, tc.targetTuples, targetObject)
+			defer cleanupTuples(b, client, ctx)
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				_, err := client.GetTuplesForObject(ctx, fmt.Sprintf("organization:%s", targetObject))
+				require.NoError(b, err)
+			}
+		})
+	}
+}
+
+func BenchmarkGetAllTuples(b *testing.B) {
+	client, ctx, cleanup := setupIntegrationBenchmark(b)
+	defer cleanup()
+
+	testCases := []struct {
+		name        string
+		totalTuples int
+	}{
+		{"Small", 100},
+		{"Medium", 1000},
+		{"Large", 5000},
+	}
+
+	for _, tc := range testCases {
+		b.Run(tc.name, func(b *testing.B) {
+			targetObject := fmt.Sprintf("any-org-%s", tc.name)
+
+			makeTuples(b, client, ctx, tc.totalTuples, 10, targetObject)
+			defer cleanupTuples(b, client, ctx)
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				_, err := client.GetAllTuples(ctx)
+				require.NoError(b, err)
+			}
+		})
+	}
+}

fgax/tuples.go

@@ -314,8 +314,8 @@ func (c *Client) deleteRelationshipTuple(ctx context.Context, tuples []openfga.T
 	return resp, nil
 }
 
-// getAllTuples gets all the relationship tuples in the openFGA store
-func (c *Client) getAllTuples(ctx context.Context, opts ...RequestOption) ([]openfga.Tuple, error) {
+// GetAllTuples gets all the relationship tuples in the openFGA store
+func (c *Client) GetAllTuples(ctx context.Context, opts ...RequestOption) ([]openfga.Tuple, error) {
 	var tuples []openfga.Tuple
 
 	ropts := getReadOptions(opts...)
@@ -342,6 +342,42 @@ func (c *Client) getAllTuples(ctx context.Context, opts ...RequestOption) ([]ope
 	return tuples, nil
 }
 
+// GetTuplesForObject gets all the relationship tuples for a specific object in the openFGA store
+func (c *Client) GetTuplesForObject(ctx context.Context, object string) ([]openfga.Tuple, error) {
+	var tuples []openfga.Tuple
+
+	readRequest := ofgaclient.ClientReadRequest{
+		Object: openfga.PtrString(object),
+	}
+
+	opts := ofgaclient.ClientReadOptions{
+		PageSize:    openfga.PtrInt32(defaultPageSize),
+		Consistency: &defaultConsistency,
+	}
+
+	notComplete := true
+
+	for notComplete {
+		resp, err := c.Ofga.Read(ctx).Body(readRequest).Options(opts).Execute()
+		if err != nil {
+			log.Error().Err(err).Str("object", object).Msg("error getting relationship tuples for object")
+
+			return nil, err
+		}
+
+		tuples = append(tuples, resp.GetTuples()...)
+
+		if resp.ContinuationToken == "" {
+			notComplete = false
+			continue
+		}
+
+		opts.ContinuationToken = &resp.ContinuationToken
+	}
+
+	return tuples, nil
+}
+
 // DeleteAllObjectRelations deletes all the relationship tuples for a given object
 func (c *Client) DeleteAllObjectRelations(ctx context.Context, object string, excludeRelations []string, opts ...RequestOption) error {
 	// validate object is not empty
@@ -354,28 +390,26 @@ func (c *Client) DeleteAllObjectRelations(ctx context.Context, object string, ex
 		return newInvalidEntityError(object)
 	}
 
-	tuples, err := c.getAllTuples(ctx, opts...)
+	tuples, err := c.GetTuplesForObject(ctx, object)
 	if err != nil {
 		return err
 	}
 
 	var tuplesToDelete []openfga.TupleKeyWithoutCondition
 
-	// check all the tuples for the object
+	// Filter out relations that should be excluded
 	for _, t := range tuples {
-		if t.Key.Object == object {
-			// if the relation is in the exclude list, skip it
-			if slices.Contains(excludeRelations, t.Key.Relation) {
-				continue
-			}
-
-			k := openfga.TupleKeyWithoutCondition{
-				User:     t.Key.User,
-				Relation: t.Key.Relation,
-				Object:   t.Key.Object,
-			}
-			tuplesToDelete = append(tuplesToDelete, k)
+		// if the relation is in the exclude list, skip it
+		if slices.Contains(excludeRelations, t.Key.Relation) {
+			continue
+		}
+
+		k := openfga.TupleKeyWithoutCondition{
+			User:     t.Key.User,
+			Relation: t.Key.Relation,
+			Object:   t.Key.Object,
 		}
+		tuplesToDelete = append(tuplesToDelete, k)
 	}
 
 	// delete the tuples in batches of 10, the max supported by the OpenFGA transactional write api