[NOREF] Dataloader sort and error utilities (#1402)

* feat: pull in  helper functions from other branch

* feat: use the sort and error functions for model plan and operational solution functions

* feat: renamed mapper --> mapping, and add comment. Remove unused mapper interface

pkg/helpers/mapping.go

@@ -0,0 +1,110 @@
+package helpers
+
+// OneToOneFunc takes a list of keys and a list of values which map one-to-one (key-to-value).
+// it relies on the transformOutput func to return the result in expected format
+// Example:
+//
+//	type User struct {
+//	    ID   string
+//	    Name string
+//	}
+//
+//	keys := []string{"1", "2"}
+//	users := []User{
+//	    {ID: "1", Name: "Alice"},
+//	    {ID: "2", Name: "Bob"},
+//	}
+//
+//	getKeyFunc := func(user User) string {
+//	    return user.ID
+//	}
+//
+//	transformOutputFunc := func(user User, exists bool) string {
+//	    if !exists {
+//	        return "Unknown"
+//	    }
+//	    return user.Name
+//	}
+//
+//	result := OneToOneFunc(keys, users, getKeyFunc, transformOutputFunc)
+//	// result: []string{"Alice", "Bob"}
+func OneToOneFunc[K comparable, V any, Output any](keys []K, vals []V, getKey func(V) K, transformOutput func(V, bool) Output) []Output {
+	store := map[K]V{}
+
+	for _, val := range vals {
+		id := getKey(val)
+		store[id] = val
+	}
+	output := make([]Output, len(keys))
+
+	for index, key := range keys {
+		data, ok := store[key]
+		output[index] = transformOutput(data, ok)
+	}
+
+	return output
+}
+
+// OneToManyFunc takes a list of keys and a list of values which map one-to-many (key-to-value)
+// ex: vals could be a list of collaborators where more than one collaborator exists for the same model plan id
+// getKey: this function takes a value, and gets mapKey key for the intermediate map of data. This is needed for cases where you can't directly infer a key from a value
+// getRes: this function takes an original key, and the intermediate response and returns a value and bool if hte value exists
+// transform output lets you cast a data result to final expected data type returned
+/* EXAMPLE
+type SolutionAndPossibleKey struct {
+	// OperationalNeedID is the ID of the need that a solution is associated with
+	OperationalNeedID uuid.UUID `json:"operational_need_id"`
+	// IncludeNotNeeded specifies if the query should return solutions with a status of not needed, or if possible solutions (not discrete records) should be included
+	IncludeNotNeeded bool `json:"include_not_needed"`
+}
+keys []storage.SolutionAndPossibleKey
+data := []*models.OperationalSolution {lots of operational solutions}
+
+	getKeyFunc := func(data *models.OperationalSolution) uuid.UUID {
+		// The key for the loader is more complicated, but can't be inferred by each individual value. Instead, we can use the get res function to further transform the data as needed
+		return data.OperationalNeedID
+	}
+
+	getResFunc := func(key storage.SolutionAndPossibleKey, resMap map[uuid.UUID][]*models.OperationalSolution) ([]*models.OperationalSolution, bool) {
+	res, ok := resMap[key.OperationalNeedID]
+		if !key.IncludeNotNeeded {
+			lo.Filter(res, func(sol *models.OperationalSolution, _ int) bool {
+				if sol.Needed == nil{
+					return false
+				}
+				return *sol.Needed
+			})
+		// }
+
+	return res, ok
+
+	}
+
+	transformFunc := func transformToDataLoaderResultAllowNils[V any](val V, valueFound bool) *dataloader.Result[V] {
+	return &dataloader.Result[V]{Data: val, Error: nil}
+}
+
+	return OneToManyFunc(keys, sols, getKeyFunc, getResFunc, transformFunc)
+
+
+*/
+func OneToManyFunc[K comparable, V any, mapKey comparable, Output any](keys []K, vals []V, getKey func(V) mapKey, getRes func(K, map[mapKey][]V) ([]V, bool), transformOutput func([]V, bool) Output) []Output {
+	// create a map to store values grouped by key (of type K)
+	// each key will map to a slice of values (of type V)
+	store := map[mapKey][]V{}
+
+	for _, val := range vals {
+		id := getKey(val)
+		if _, ok := store[id]; !ok {
+			store[id] = []V{}
+		}
+		store[id] = append(store[id], val)
+	}
+	output := make([]Output, len(keys))
+	for index, key := range keys {
+		data, ok := getRes(key, store)
+		output[index] = transformOutput(data, ok)
+	}
+
+	return output
+}

pkg/helpers/pointer.go

@@ -0,0 +1,6 @@
+package helpers
+
+// PointerTo takes in any item and returns a pointer to that item
+func PointerTo[T any](val T) *T {
+	return &val
+}

pkg/storage/loaders/errors.go

@@ -10,3 +10,6 @@ var ErrLoaderOfWrongType = errors.New("dataLoader: dataloader returned from the
 
 // ErrLoaderIsNotInstantiated is returned when a dataloader hasn't been instantiated properly, and is nil
 var ErrLoaderIsNotInstantiated = errors.New("dataLoader: dataloader has not been instantiated. make sure to assign the loader to the config on instantiation")
+
+// ErrRecordNotFoundForKey is returned when a dataloader result doesn't have a record for a given key
+var ErrRecordNotFoundForKey = errors.New("dataLoader: record not found for given key")

pkg/storage/loaders/helpers.go

@@ -0,0 +1,39 @@
+package loaders
+
+import (
+	"fmt"
+
+	"github.com/graph-gophers/dataloader/v7"
+
+	"github.com/cms-enterprise/mint-app/pkg/helpers"
+)
+
+func transformToDataLoaderResult[V any](val V, valueFound bool) *dataloader.Result[V] {
+	if valueFound {
+		return &dataloader.Result[V]{Data: val, Error: nil}
+	}
+
+	return &dataloader.Result[V]{Data: val, Error: fmt.Errorf("issue getting result for type %T, err: %w", val, ErrRecordNotFoundForKey)}
+}
+
+// transformToDataLoaderResultAllowNils transforms an output to a dataloader result. It doesn't error if there is not a value for the given key.
+func transformToDataLoaderResultAllowNils[V any](val V, valueFound bool) *dataloader.Result[V] {
+	return &dataloader.Result[V]{Data: val, Error: nil}
+}
+
+func oneToOneDataLoaderFunc[K comparable, V any](keys []K, values []V, getKey func(V) K) []*dataloader.Result[V] {
+
+	return helpers.OneToOneFunc(keys, values, getKey, transformToDataLoaderResult)
+}
+func oneToManyDataLoaderFunc[K comparable, V any, mapKey comparable](keys []K, values []V, getKey func(V) mapKey, getRes func(K, map[mapKey][]V) ([]V, bool)) []*dataloader.Result[[]V] {
+	return helpers.OneToManyFunc(keys, values, getKey, getRes, transformToDataLoaderResultAllowNils)
+}
+
+func errorPerEachKey[K comparable, V any](keys []K, err error) []*dataloader.Result[V] {
+	var empty V
+	output := make([]*dataloader.Result[V], len(keys))
+	for index := range keys {
+		output[index] = &dataloader.Result[V]{Data: empty, Error: err}
+	}
+	return output
+}

pkg/storage/loaders/model_plan_loader.go

@@ -2,10 +2,8 @@ package loaders
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/google/uuid"
-	"github.com/samber/lo"
 
 	"github.com/cms-enterprise/mint-app/pkg/appcontext"
 	"github.com/cms-enterprise/mint-app/pkg/models"
@@ -27,38 +25,18 @@ var ModelPlan = &modelPlanLoaders{
 
 func batchModelPlanByModelPlanID(ctx context.Context, modelPlanIDs []uuid.UUID) []*dataloader.Result[*models.ModelPlan] {
 	logger := appcontext.ZLogger(ctx)
-	output := make([]*dataloader.Result[*models.ModelPlan], len(modelPlanIDs))
 	loaders, err := Loaders(ctx)
 	if err != nil {
-		//TODO: (loaders) make this a helper function to return an error per result
-		for index := range modelPlanIDs {
-			output[index] = &dataloader.Result[*models.ModelPlan]{Data: nil, Error: err}
-		}
-		return output
+		return errorPerEachKey[uuid.UUID, *models.ModelPlan](modelPlanIDs, err)
 	}
 
 	data, err := storage.ModelPlansGetByModePlanIDsLOADER(loaders.DataReader.Store, logger, modelPlanIDs)
 	if err != nil {
-		//TODO: (loaders) make this a helper function to return an error per result
-		for index := range modelPlanIDs {
-			output[index] = &dataloader.Result[*models.ModelPlan]{Data: nil, Error: err}
-		}
-		return output
+		return errorPerEachKey[uuid.UUID, *models.ModelPlan](modelPlanIDs, err)
 	}
-	planByID := lo.Associate(data, func(plan *models.ModelPlan) (uuid.UUID, *models.ModelPlan) {
-		return plan.ID, plan
-	})
-
-	// RETURN IN THE SAME ORDER REQUESTED
-	for index, id := range modelPlanIDs {
-
-		plan, ok := planByID[id]
-		if ok {
-			output[index] = &dataloader.Result[*models.ModelPlan]{Data: plan, Error: nil}
-		} else {
-			err2 := fmt.Errorf("model plan not found for modelPlanID id %s", id)
-			output[index] = &dataloader.Result[*models.ModelPlan]{Data: nil, Error: err2}
-		}
+	getKeyFunc := func(modelPlan *models.ModelPlan) uuid.UUID {
+		return modelPlan.ID
 	}
-	return output
+
+	return oneToOneDataLoaderFunc(modelPlanIDs, data, getKeyFunc)
 }

pkg/storage/loaders/operational_solution_and_possible_collection_loader.go

@@ -14,44 +14,35 @@ import (
 
 func batchOperationalSolutionAndPossibleCollectionGetByOperationalNeedID(ctx context.Context, keys []storage.SolutionAndPossibleKey) []*dataloader.Result[[]*models.OperationalSolution] {
 	logger := appcontext.ZLogger(ctx)
-	output := make([]*dataloader.Result[[]*models.OperationalSolution], len(keys))
+
 	loaders, err := Loaders(ctx)
 	if err != nil {
-		for index := range keys {
-			output[index] = &dataloader.Result[[]*models.OperationalSolution]{Data: nil, Error: err}
-		}
-		return output
+		return errorPerEachKey[storage.SolutionAndPossibleKey, []*models.OperationalSolution](keys, err)
+
 	}
 
 	sols, loadErr := storage.OperationalSolutionAndPossibleCollectionGetByOperationalNeedIDLOADER(loaders.DataReader.Store, logger, keys)
-
 	if loadErr != nil {
-		for index := range keys {
-			output[index] = &dataloader.Result[[]*models.OperationalSolution]{Data: nil, Error: loadErr}
-		}
-		return output
-
-	}
-	solsByID := map[uuid.UUID][]*models.OperationalSolution{}
-	for _, sol := range sols {
-		slice, ok := solsByID[sol.OperationalNeedID]
-		if ok {
-			slice = append(slice, sol) //Add to existing slice
-			solsByID[sol.OperationalNeedID] = slice
-			continue
-		}
-		solsByID[sol.OperationalNeedID] = []*models.OperationalSolution{sol}
+		return errorPerEachKey[storage.SolutionAndPossibleKey, []*models.OperationalSolution](keys, loadErr)
 	}
 
-	for index, key := range keys {
-		//Note we aren't verifying the not needed when we return the result.... We should be including needed / not needed programmatically.
-		//  This is an edge case when a need is queried twice, once with needed once without. In practice, this doesn't happen. As this is getting refactored, we will leave it as is.
-
-		sols := solsByID[key.OperationalNeedID] //Any Solutions not found will return a zero state result eg empty array
-
-		output[index] = &dataloader.Result[[]*models.OperationalSolution]{Data: sols, Error: nil}
-
+	// We use just operational need ID for the intermediate map
+	getKeyFunc := func(data *models.OperationalSolution) uuid.UUID {
+		return data.OperationalNeedID
+	}
+	getResFunc := func(key storage.SolutionAndPossibleKey, resMap map[uuid.UUID][]*models.OperationalSolution) ([]*models.OperationalSolution, bool) {
+		res, ok := resMap[key.OperationalNeedID]
+		// NOTE: we don't filter out the not needed for this loader, as it isn't possible to request it from the resolver
+		// if !key.IncludeNotNeeded {
+		// 	lo.Filter(res, func(sol *models.OperationalSolution, _ int) bool {
+		// 		if sol.Needed == nil{
+		// 			return false
+		// 		}
+		// 		return *sol.Needed
+		// 	})
+		// }
+		return res, ok
 	}
-	return output
+	return oneToManyDataLoaderFunc(keys, sols, getKeyFunc, getResFunc)
 
 }

pkg/storage/loaders/operational_solution_loader.go

@@ -2,10 +2,8 @@ package loaders
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/google/uuid"
-	"github.com/samber/lo"
 
 	"github.com/cms-enterprise/mint-app/pkg/appcontext"
 	"github.com/cms-enterprise/mint-app/pkg/models"
@@ -34,39 +32,19 @@ func operationalSolutionGetByIDBatch(
 	ids []uuid.UUID,
 ) []*dataloader.Result[*models.OperationalSolution] {
 	logger := appcontext.ZLogger(ctx)
-	output := make([]*dataloader.Result[*models.OperationalSolution], len(ids))
 	loaders, err := Loaders(ctx)
 	if err != nil {
-		for index := range ids {
-			output[index] = &dataloader.Result[*models.OperationalSolution]{Data: nil, Error: ErrNoLoaderOnContext}
-		}
-		return output
+		return errorPerEachKey[uuid.UUID, *models.OperationalSolution](ids, err)
 	}
 
 	opSols, loadErr := storage.OperationalSolutionGetByIDLOADER(loaders.DataReader.Store, logger, ids)
 	if loadErr != nil {
-		for index := range ids {
-			output[index] = &dataloader.Result[*models.OperationalSolution]{Data: nil, Error: loadErr}
-		}
-		return output
+		return errorPerEachKey[uuid.UUID, *models.OperationalSolution](ids, loadErr)
 	}
-
-	opSolsByID := lo.Associate(opSols, func(gc *models.OperationalSolution) (uuid.UUID, *models.OperationalSolution) {
-		return gc.ID, gc
-	})
-
-	// RETURN IN THE SAME ORDER REQUESTED
-	for index, key := range ids {
-
-		opSol, ok := opSolsByID[key]
-		if ok {
-			output[index] = &dataloader.Result[*models.OperationalSolution]{Data: opSol, Error: nil}
-		} else {
-			err := fmt.Errorf("operational solution not found for id %s", key)
-			output[index] = &dataloader.Result[*models.OperationalSolution]{Data: nil, Error: err}
-		}
-
+	getKeyFunc := func(data *models.OperationalSolution) uuid.UUID {
+		return data.ID
 	}
 
-	return output
+	return oneToOneDataLoaderFunc(ids, opSols, getKeyFunc)
+
 }