Add basic steelthread tests for bulk import and export of relationships

internal/services/steelthreadtesting/definitions.go

@@ -327,4 +327,63 @@ var steelThreadTestCases = []steelThreadTestCase{
 			},
 		},
 	},
+	{
+		name:     "basic import export",
+		datafile: "document-with-a-few-relationships.yaml",
+		operations: []steelThreadOperationCase{
+			{
+				name:          "import export with no limit and no filter",
+				operationName: "bulkImportExportRelationships",
+				arguments: map[string]any{
+					"rels_file": "basic-import-export-relationships.txt",
+				},
+				resultsFileName: "basic-import-export-results.yaml",
+			},
+		},
+	},
+	{
+		name:     "basic import export with optional limit",
+		datafile: "document-with-a-few-relationships.yaml",
+		operations: []steelThreadOperationCase{
+			{
+				name:          "import export with limit and no filter",
+				operationName: "bulkImportExportRelationships",
+				arguments: map[string]any{
+					"rels_file":      "basic-import-export-relationships.txt",
+					"optional_limit": 3,
+				},
+				resultsFileName: "basic-import-export-results.yaml",
+			},
+		},
+	},
+	{
+		name:     "basic import export with filter",
+		datafile: "document-with-a-few-relationships.yaml",
+		operations: []steelThreadOperationCase{
+			{
+				name:          "import export with filter filter",
+				operationName: "bulkImportExportRelationships",
+				arguments: map[string]any{
+					"rels_file":     "basic-import-export-relationships.txt",
+					"resource_type": "document",
+				},
+				resultsFileName: "basic-import-export-results.yaml",
+			},
+		},
+	},
+	{
+		name:     "basic import export with object ID filter",
+		datafile: "document-with-a-few-relationships.yaml",
+		operations: []steelThreadOperationCase{
+			{
+				name:          "import export with filter filter",
+				operationName: "bulkImportExportRelationships",
+				arguments: map[string]any{
+					"rels_file":                 "basic-import-export-relationships.txt",
+					"filter_resource_id_prefix": "doc-1",
+				},
+				resultsFileName: "filtered-import-export-results.yaml",
+			},
+		},
+	},
 }

internal/services/steelthreadtesting/operations.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"sort"
 	"strings"
 	"time"
@@ -17,6 +18,7 @@ import (
 	"gopkg.in/yaml.v3"
 
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
+	"github.com/authzed/spicedb/pkg/tuple"
 )
 
 func lookupSubjects(parameters map[string]any, client v1.PermissionsServiceClient) (any, error) {
@@ -223,10 +225,114 @@ func cursoredLookupResources(parameters map[string]any, client v1.PermissionsSer
 	return nodeSets, nil
 }
 
+func bulkImportExportRelationships(parameters map[string]any, client v1.PermissionsServiceClient) (any, error) {
+	// Read the list of relationships to pass to the import operation.
+	importRelsFile := parameters["rels_file"].(string)
+
+	relsFile, err := os.ReadFile("testdata/" + importRelsFile)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read relationships file: %w", err)
+	}
+
+	// Run the import relationships.
+	importRels := make([]*v1.Relationship, 0)
+	for _, line := range strings.Split(string(relsFile), "\n") {
+		if line == "" {
+			continue
+		}
+
+		parsed, err := tuple.ParseV1Rel(line)
+		if err != nil {
+			return nil, err
+		}
+
+		importRels = append(importRels, parsed)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	r, err := client.ImportBulkRelationships(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, rel := range importRels {
+		err := r.Send(&v1.ImportBulkRelationshipsRequest{Relationships: []*v1.Relationship{rel}})
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	resp, err := r.CloseAndRecv()
+	if err != nil {
+		return nil, err
+	}
+
+	if len(importRels) != int(resp.NumLoaded) {
+		return nil, fmt.Errorf("expected %d relationships to be loaded, got %d", len(importRels), resp.NumLoaded)
+	}
+
+	// Run bulk export and return the results.
+	var optionalLimit uint32
+	if optionalLimitValue, ok := parameters["optional_limit"]; ok {
+		optionalLimit = uint32(optionalLimitValue.(int))
+	}
+
+	var filter *v1.RelationshipFilter
+	if resourceTypeValue, ok := parameters["resource_type"]; ok {
+		filter = &v1.RelationshipFilter{
+			ResourceType: resourceTypeValue.(string),
+		}
+	}
+
+	if filterPrefixValue, ok := parameters["filter_resource_id_prefix"]; ok {
+		filter = &v1.RelationshipFilter{
+			OptionalResourceIdPrefix: filterPrefixValue.(string),
+		}
+	}
+
+	exr, err := client.ExportBulkRelationships(ctx, &v1.ExportBulkRelationshipsRequest{
+		Consistency: &v1.Consistency{
+			Requirement: &v1.Consistency_FullyConsistent{
+				FullyConsistent: true,
+			},
+		},
+		OptionalLimit:              optionalLimit,
+		OptionalRelationshipFilter: filter,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	exportedRels := make([]yaml.Node, 0)
+	for {
+		resp, err := exr.Recv()
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+
+			return nil, err
+		}
+
+		for _, rel := range resp.Relationships {
+			exportedRels = append(exportedRels, yaml.Node{
+				Kind:  yaml.ScalarNode,
+				Value: tuple.MustV1RelString(rel),
+				Style: yaml.SingleQuotedStyle,
+			})
+		}
+	}
+
+	return exportedRels, nil
+}
+
 var operations = map[string]stOperation{
-	"lookupSubjects":          lookupSubjects,
-	"lookupResources":         lookupResources,
-	"cursoredLookupResources": cursoredLookupResources,
+	"lookupSubjects":                lookupSubjects,
+	"lookupResources":               lookupResources,
+	"cursoredLookupResources":       cursoredLookupResources,
+	"bulkImportExportRelationships": bulkImportExportRelationships,
 }
 
 func formatResolvedResource(resource *v1.LookupResourcesResponse) string {

internal/services/steelthreadtesting/steelresults/basic-import-export-results.yaml

@@ -0,0 +1,46 @@
+---
+- 'document:doc-1#viewer@user:user-0'
+- 'document:doc-1#viewer@user:user-1'
+- 'document:doc-1#viewer@user:user-2'
+- 'document:doc-1#viewer@user:user-3'
+- 'document:doc-1#viewer@user:user-4'
+- 'document:doc-1#viewer@user:user-5'
+- 'document:doc-1#viewer@user:user-6'
+- 'document:doc-1#viewer@user:user-7'
+- 'document:doc-1#viewer@user:user-8'
+- 'document:doc-17#viewer@user:user-1[somecaveat:{"somecondition":43}][expiration:2323-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-10[somecaveat:{"somecondition":52}][expiration:3223-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-2[somecaveat:{"somecondition":44}][expiration:2423-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-3[somecaveat:{"somecondition":45}][expiration:2523-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-4[somecaveat:{"somecondition":46}][expiration:2623-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-5[somecaveat:{"somecondition":47}][expiration:2723-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-6[somecaveat:{"somecondition":48}][expiration:2823-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-7[somecaveat:{"somecondition":49}][expiration:2923-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-8[somecaveat:{"somecondition":50}][expiration:3023-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-9[somecaveat:{"somecondition":51}][expiration:3123-01-02T13:14:15Z]'
+- 'document:doc-2#viewer@user:user-0'
+- 'document:doc-3#viewer@user:user-2'
+- 'document:doc-3#viewer@user:user-3'
+- 'document:doc-3#viewer@user:user-4'
+- 'document:doc-3#viewer@user:user-5'
+- 'document:doc-3#viewer@user:user-6'
+- 'document:doc-3#viewer@user:user-7'
+- 'document:doc-3#viewer@user:user-8'
+- 'document:doc-4#viewer@user:user-1[somecaveat]'
+- 'document:doc-4#viewer@user:user-2[somecaveat:{"somecondition":41}]'
+- 'document:doc-4#viewer@user:user-3[somecaveat:{"somecondition":42}]'
+- 'document:doc-5#viewer@user:user-1[expiration:2323-01-02T13:14:15Z]'
+- 'document:doc-5#viewer@user:user-2[expiration:2423-01-02T13:14:15Z]'
+- 'document:doc-6#viewer@user:user-1[somecaveat][expiration:2323-01-02T13:14:15Z]'
+- 'document:doc-6#viewer@user:user-2[somecaveat:{"somecondition":41}][expiration:2123-01-02T13:14:15Z]'
+- 'document:doc-6#viewer@user:user-3[somecaveat:{"somecondition":42}][expiration:2223-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-1[somecaveat:{"somecondition":43}][expiration:2323-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-10[somecaveat:{"somecondition":52}][expiration:3223-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-2[somecaveat:{"somecondition":44}][expiration:2423-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-3[somecaveat:{"somecondition":45}][expiration:2523-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-4[somecaveat:{"somecondition":46}][expiration:2623-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-5[somecaveat:{"somecondition":47}][expiration:2723-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-6[somecaveat:{"somecondition":48}][expiration:2823-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-7[somecaveat:{"somecondition":49}][expiration:2923-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-8[somecaveat:{"somecondition":50}][expiration:3023-01-02T13:14:15Z]'
+- 'document:doc-7#viewer@user:user-9[somecaveat:{"somecondition":51}][expiration:3123-01-02T13:14:15Z]'

internal/services/steelthreadtesting/steelresults/filtered-import-export-results.yaml

@@ -0,0 +1,20 @@
+---
+- 'document:doc-1#viewer@user:user-0'
+- 'document:doc-1#viewer@user:user-1'
+- 'document:doc-1#viewer@user:user-2'
+- 'document:doc-1#viewer@user:user-3'
+- 'document:doc-1#viewer@user:user-4'
+- 'document:doc-1#viewer@user:user-5'
+- 'document:doc-1#viewer@user:user-6'
+- 'document:doc-1#viewer@user:user-7'
+- 'document:doc-1#viewer@user:user-8'
+- 'document:doc-17#viewer@user:user-1[somecaveat:{"somecondition":43}][expiration:2323-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-10[somecaveat:{"somecondition":52}][expiration:3223-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-2[somecaveat:{"somecondition":44}][expiration:2423-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-3[somecaveat:{"somecondition":45}][expiration:2523-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-4[somecaveat:{"somecondition":46}][expiration:2623-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-5[somecaveat:{"somecondition":47}][expiration:2723-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-6[somecaveat:{"somecondition":48}][expiration:2823-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-7[somecaveat:{"somecondition":49}][expiration:2923-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-8[somecaveat:{"somecondition":50}][expiration:3023-01-02T13:14:15Z]'
+- 'document:doc-17#viewer@user:user-9[somecaveat:{"somecondition":51}][expiration:3123-01-02T13:14:15Z]'

internal/services/steelthreadtesting/testdata/basic-import-export-relationships.txt

@@ -0,0 +1,42 @@
+document:doc-1#viewer@user:user-2
+document:doc-1#viewer@user:user-3
+document:doc-1#viewer@user:user-4
+document:doc-1#viewer@user:user-5
+document:doc-1#viewer@user:user-6
+document:doc-1#viewer@user:user-7
+document:doc-1#viewer@user:user-8
+document:doc-3#viewer@user:user-2
+document:doc-3#viewer@user:user-3
+document:doc-3#viewer@user:user-4
+document:doc-3#viewer@user:user-5
+document:doc-3#viewer@user:user-6
+document:doc-3#viewer@user:user-7
+document:doc-3#viewer@user:user-8
+document:doc-4#viewer@user:user-1[somecaveat]
+document:doc-4#viewer@user:user-2[somecaveat:{"somecondition": 41}]
+document:doc-4#viewer@user:user-3[somecaveat:{"somecondition": 42}]
+document:doc-5#viewer@user:user-1[expiration:2323-01-02T13:14:15Z]
+document:doc-5#viewer@user:user-2[expiration:2423-01-02T13:14:15Z]
+document:doc-6#viewer@user:user-1[somecaveat][expiration:2323-01-02T13:14:15Z]
+document:doc-6#viewer@user:user-2[somecaveat:{"somecondition": 41}][expiration:2123-01-02T13:14:15Z]
+document:doc-6#viewer@user:user-3[somecaveat:{"somecondition": 42}][expiration:2223-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-1[somecaveat:{"somecondition": 43}][expiration:2323-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-2[somecaveat:{"somecondition": 44}][expiration:2423-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-3[somecaveat:{"somecondition": 45}][expiration:2523-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-4[somecaveat:{"somecondition": 46}][expiration:2623-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-5[somecaveat:{"somecondition": 47}][expiration:2723-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-6[somecaveat:{"somecondition": 48}][expiration:2823-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-7[somecaveat:{"somecondition": 49}][expiration:2923-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-8[somecaveat:{"somecondition": 50}][expiration:3023-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-9[somecaveat:{"somecondition": 51}][expiration:3123-01-02T13:14:15Z]
+document:doc-7#viewer@user:user-10[somecaveat:{"somecondition": 52}][expiration:3223-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-1[somecaveat:{"somecondition": 43}][expiration:2323-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-2[somecaveat:{"somecondition": 44}][expiration:2423-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-3[somecaveat:{"somecondition": 45}][expiration:2523-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-4[somecaveat:{"somecondition": 46}][expiration:2623-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-5[somecaveat:{"somecondition": 47}][expiration:2723-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-6[somecaveat:{"somecondition": 48}][expiration:2823-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-7[somecaveat:{"somecondition": 49}][expiration:2923-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-8[somecaveat:{"somecondition": 50}][expiration:3023-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-9[somecaveat:{"somecondition": 51}][expiration:3123-01-02T13:14:15Z]
+document:doc-17#viewer@user:user-10[somecaveat:{"somecondition": 52}][expiration:3223-01-02T13:14:15Z]

internal/services/steelthreadtesting/testdata/document-with-a-few-relationships.yaml

@@ -0,0 +1,19 @@
+---
+schema: |+
+  use expiration
+
+  caveat somecaveat(somecondition int) {
+    somecondition == 42
+  }
+
+  definition user {}
+
+  definition document {
+    relation viewer: user | user with expiration | user with somecaveat | user with somecaveat and expiration
+    permission view = viewer
+  }
+
+relationships: |
+  document:doc-1#viewer@user:user-0
+  document:doc-1#viewer@user:user-1
+  document:doc-2#viewer@user:user-0