go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/db/dbloader/loader.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

package dbloader

import (
	"context"

	"go.charczuk.com/sdk/db"
)

// LoaderFor returns a loader for a given type.
//
// Specifically it uses a `GetMany` command on the dbc to fetch
// the keys for a given batch.
//
// There are very specific constraints on the types we can use as the values here, specifically
// the types can only have (1) primary key, and the key values will be resolved
// with reflection incurring a bit of a performance hit.
func LoaderFor[K comparable, V any](dbc *db.Connection, opts ...db.InvocationOption) Loader[K, V] {
	var v V
	primaryKeys := db.TypeMetaFor(v).PrimaryKeys()
	if len(primaryKeys) == 0 {
		panic("LoaderFor provided type does not have any primary keys")
	}
	if len(primaryKeys) > 1 {
		panic("LoaderFor provided type has multiple primary keys")
	}
	primaryKey := primaryKeys[0]
	keyAccessor := func(v V) (k K) {
		rawKeyValue := primaryKey.GetValue(v)
		typedKeyValue, ok := rawKeyValue.(K)
		if !ok {
			return
		}
		return typedKeyValue
	}
	return Loader[K, V]{
		Resolver: func(ctx context.Context, keys ...K) (output map[K]V, err error) {
			// ~~ UGLY HACK ALERT ~~
			// we have to do this because []K is not assignable to []any
			// for reasons that i'm sure the go team has a patronizing
			// answer for but regardless the ergonomics are awful.
			anyKeys := make([]any, len(keys))
			for index := range keys {
				anyKeys[index] = keys[index]
			}
			var values []V
			err = dbc.Invoke(append(opts, db.OptContext(ctx))...).GetMany(&values, anyKeys...)
			if err != nil {
				return
			}
			output = make(map[K]V)
			for _, v := range values {
				k := keyAccessor(v)
				output[k] = v
			}
			return
		},
	}
}

// Loader is a specific "batch" of calls to
// the loader resolver.
type Loader[K comparable, V any] struct {
	Resolver func(context.Context, ...K) (map[K]V, error)
}

func (l Loader[K, V]) New() *Batch[K, V] {
	return &Batch[K, V]{
		loader:  &l,
		storage: make(map[K]*V),
	}
}

type Batch[K comparable, V any] struct {
	loader  *Loader[K, V]
	storage map[K]*V
}

// Load returns a promise to load a value with a given key.
func (b *Batch[K, V]) Load(k K) *V {
	if v, ok := b.storage[k]; ok {
		return v
	}
	// store ~something~ until we resolve
	// we can't store <nil> here because we
	// have to deref this pointer when we resolve
	var v V
	b.storage[k] = &v
	return &v
}

// Keys returns the "scheduled" keys to be fetched.
func (b *Batch[K, V]) Keys() (output []K) {
	output = make([]K, 0, len(b.storage))
	for k := range b.storage {
		output = append(output, k)
	}
	return
}

// Resolve resoles the "promises" for the values for given keys and
// makes the stored pointers point at valid targets.
func (b *Batch[K, V]) Resolve(ctx context.Context) error {
	values, err := b.loader.Resolver(ctx, b.Keys()...)
	if err != nil {
		return err
	}
	for k, v := range values {
		if sv, ok := b.storage[k]; ok {
			// this assignment is _very_ important
			// we have to set the _target_ of the pointer
			// to our value, not set the _variable_ to the address
			// of our value.
			*sv = v
		}
	}
	return nil
}