Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[NOREF] Dataloader sort and error utilities #1402

Merged
merged 7 commits into from
Oct 11, 2024
110 changes: 110 additions & 0 deletions pkg/helpers/mapping.go
Original file line number Diff line number Diff line change
@@ -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
}
6 changes: 6 additions & 0 deletions pkg/helpers/pointer.go
Original file line number Diff line number Diff line change
@@ -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
}
3 changes: 3 additions & 0 deletions pkg/storage/loaders/errors.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
39 changes: 39 additions & 0 deletions pkg/storage/loaders/helpers.go
Original file line number Diff line number Diff line change
@@ -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
}
34 changes: 6 additions & 28 deletions pkg/storage/loaders/model_plan_loader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)
}
Original file line number Diff line number Diff line change
Expand Up @@ -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)

}
34 changes: 6 additions & 28 deletions pkg/storage/loaders/operational_solution_loader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)

}
Loading