entity.go

package grpcdatasource

import (
	"errors"
	"fmt"

	"github.com/tidwall/gjson"

	"github.com/wundergraph/astjson"
)

// entityIndexMap maps positions in the typed gRPC response back to positions
// in the original representations array. The slice index is the response
// position; the value is the representation index. It is built per call by
// recording the position of every representation whose __typename matches
// the requested entity type.
type entityIndexMap []int

// newEntityIndexMap builds the index map for a single entity call by collecting
// the positions of representations whose __typename matches the requested type.
// A single pass over representations populates the slice.
func newEntityIndexMap(requestedEntityType string, representations []gjson.Result) entityIndexMap {
	indexMap := make(entityIndexMap, 0, len(representations))
	for i, representation := range representations {
		if representation.Get(typenameFieldName).String() == requestedEntityType {
			indexMap = append(indexMap, i)
		}
	}
	return indexMap
}

// getRepresentations gets the representations from the variables.
// If no representations are found, it returns nil.
func getRepresentations(variables gjson.Result) []gjson.Result {
	r := variables.Get("representations")
	if !r.Exists() {
		return nil
	}

	return r.Array()
}

// validateEntityResponse verifies that the number of entities returned by the
// subgraph matches the number of representations of the requested type.
// Callers should subsequently build an entityIndexMap via newEntityIndexMap to
// merge the response — mergeEntities relies on the invariant that
// len(response entities) == len(indexMap), which this function establishes.
func validateEntityResponse(data *astjson.Value, requestedEntityType string, representations []gjson.Result) error {
	if data == nil {
		return errors.New("validateEntityResponse: subgraph response data is nil")
	}

	if requestedEntityType == "" {
		return errors.New("validateEntityResponse: requested entity type is empty; the entity RPC plan is missing a RequestedEntityType")
	}

	if len(representations) == 0 {
		return errors.New("validateEntityResponse: no entity representations provided in the request variables")
	}

	expected := 0
	for _, representation := range representations {
		if representation.Get(typenameFieldName).String() == requestedEntityType {
			expected++
		}
	}

	entities := data.Get(entityPath).GetArray()
	if len(entities) != expected {
		return fmt.Errorf("entity type %s received %d entities in the subgraph response, but %d are expected", requestedEntityType, len(entities), expected)
	}

	return nil
}