chore: Change query package to use new schema v2 accessor methods (#2659)

Author: josephschorr

pkg/query/build_tree.go

@@ -75,7 +75,7 @@ func (b *iteratorBuilder) buildIteratorFromSchemaInternal(definitionName string,
 	// Track the position in the sentinels list before building
 	sentinelsLenBefore := len(b.recursiveSentinels)
 
-	def, ok := b.schema.Definitions()[definitionName]
+	def, ok := b.schema.GetTypeDefinition(definitionName)
 	if !ok {
 		// Remove before returning error
 		delete(b.building, id)
@@ -84,9 +84,9 @@ func (b *iteratorBuilder) buildIteratorFromSchemaInternal(definitionName string,
 
 	var result Iterator
 	var err error
-	if p, ok := def.Permissions()[relationName]; ok {
+	if p, ok := def.GetPermission(relationName); ok {
 		result, err = b.buildIteratorFromPermission(p)
-	} else if r, ok := def.Relations()[relationName]; ok {
+	} else if r, ok := def.GetRelation(relationName); ok {
 		result, err = b.buildIteratorFromRelation(r, withSubRelations)
 	} else {
 		err = fmt.Errorf("BuildIteratorFromSchema: couldn't find a relation or permission named `%s` in definition `%s`", relationName, definitionName)
@@ -160,7 +160,7 @@ func (b *iteratorBuilder) buildIteratorFromPermission(p *schema.Permission) (Ite
 func (b *iteratorBuilder) buildIteratorFromOperation(p *schema.Permission, op schema.Operation) (Iterator, error) {
 	switch perm := op.(type) {
 	case *schema.ArrowReference:
-		rel, ok := p.Parent().Relations()[perm.Left()]
+		rel, ok := p.Parent().GetRelation(perm.Left())
 		if !ok {
 			return nil, fmt.Errorf("BuildIteratorFromSchema: couldn't find left-hand relation for arrow `%s->%s` for permission `%s` in definition `%s`", perm.Left(), perm.Right(), p.Name(), p.Parent().Name())
 		}
@@ -208,7 +208,7 @@ func (b *iteratorBuilder) buildIteratorFromOperation(p *schema.Permission, op sc
 		return NewExclusion(mainIt, excludedIt), nil
 
 	case *schema.FunctionedArrowReference:
-		rel, ok := p.Parent().Relations()[perm.Left()]
+		rel, ok := p.Parent().GetRelation(perm.Left())
 		if !ok {
 			return nil, fmt.Errorf("BuildIteratorFromSchema: couldn't find arrow relation `%s` for functioned arrow `%s.%s(%s)` for permission `%s` in definition `%s`", perm.Left(), perm.Left(), functionTypeString(perm.Function()), perm.Right(), p.Name(), p.Parent().Name())
 		}

pkg/query/build_tree_test.go

@@ -189,10 +189,10 @@ func TestBuildTreeArrowOperation(t *testing.T) {
 	require.NotNil(dsSchema)
 
 	// Verify that the schema contains definitions with arrow-like operations
-	docDef, ok := dsSchema.Definitions()["document"]
+	docDef, ok := dsSchema.GetTypeDefinition("document")
 	require.True(ok, "should find document definition")
 
-	viewPerm, ok := docDef.Permissions()["view"]
+	viewPerm, ok := docDef.GetPermission("view")
 	require.True(ok, "should find view permission that contains arrow operations")
 	require.NotNil(viewPerm.Operation, "view permission should have operation defined")
 }
@@ -781,10 +781,10 @@ func TestBuildTreeWildcardIterator(t *testing.T) {
 	require.NoError(err)
 
 	// Verify the schema has the wildcard BaseRelation
-	documentDef := dsSchema.Definitions()["document"]
-	require.NotNil(documentDef)
-	viewerRelation := documentDef.Relations()["viewer"]
-	require.NotNil(viewerRelation)
+	documentDef, ok := dsSchema.GetTypeDefinition("document")
+	require.True(ok)
+	viewerRelation, ok := documentDef.GetRelation("viewer")
+	require.True(ok)
 	require.Len(viewerRelation.BaseRelations(), 1)
 	baseRel := viewerRelation.BaseRelations()[0]
 	require.True(baseRel.Wildcard(), "BaseRelation should have Wildcard: true")

pkg/query/quick_e2e_test.go

@@ -30,8 +30,11 @@ func TestCheck(t *testing.T) {
 	// by iterating through the relationships in the schema and then walking them.
 	//
 	// In this case, it's a little contrived.
-	vande := NewRelationIterator(dsSchema.Definitions()["document"].Relations()["viewer_and_editor"].BaseRelations()[0])
-	edit := NewRelationIterator(dsSchema.Definitions()["document"].Relations()["editor"].BaseRelations()[0])
+	docDef, _ := dsSchema.GetTypeDefinition("document")
+	vandeRel, _ := docDef.GetRelation("viewer_and_editor")
+	vande := NewRelationIterator(vandeRel.BaseRelations()[0])
+	editRel, _ := docDef.GetRelation("editor")
+	edit := NewRelationIterator(editRel.BaseRelations()[0])
 	it := NewIntersection()
 	it.addSubIterator(vande)
 	it.addSubIterator(edit)
@@ -63,7 +66,9 @@ func TestBaseIterSubjects(t *testing.T) {
 	dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
 	require.NoError(err)
 
-	vande := NewRelationIterator(dsSchema.Definitions()["document"].Relations()["viewer_and_editor"].BaseRelations()[0])
+	docDef, _ := dsSchema.GetTypeDefinition("document")
+	vandeRel, _ := docDef.GetRelation("viewer_and_editor")
+	vande := NewRelationIterator(vandeRel.BaseRelations()[0])
 
 	ctx := &Context{
 		Context:  t.Context(),
@@ -93,8 +98,12 @@ func TestCheckArrow(t *testing.T) {
 	require.NoError(err)
 
 	// This is effectively `permission foo = parent_folder->viewer`
-	folders := NewRelationIterator(dsSchema.Definitions()["document"].Relations()["parent"].BaseRelations()[0])
-	view := NewRelationIterator(dsSchema.Definitions()["folder"].Relations()["viewer"].BaseRelations()[0])
+	docDef, _ := dsSchema.GetTypeDefinition("document")
+	parentRel, _ := docDef.GetRelation("parent")
+	folders := NewRelationIterator(parentRel.BaseRelations()[0])
+	folderDef, _ := dsSchema.GetTypeDefinition("folder")
+	viewRel, _ := folderDef.GetRelation("viewer")
+	view := NewRelationIterator(viewRel.BaseRelations()[0])
 	it := NewArrow(folders, view)
 
 	ctx := &Context{