gitlab.com/picnic-app/backend/role-api@v0.0.0-20230614140944-06a76ff3696d/internal/repo/spanner/helpers/helpers.go

package helpers

import (
	"context"
	"fmt"
	"reflect"

	"cloud.google.com/go/spanner"
	"github.com/Masterminds/squirrel"
	"google.golang.org/api/iterator"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"gitlab.com/picnic-app/backend/role-api/internal/repo/spanner/tables"
)

type Reader interface {
	Read(ctx context.Context, table string, k spanner.KeySet, cols []string) *spanner.RowIterator
}

type RowReader interface {
	ReadRow(ctx context.Context, table string, k spanner.Key, cols []string) (*spanner.Row, error)
}

type Queryer interface {
	Query(context.Context, spanner.Statement) *spanner.RowIterator
}

type Updater interface {
	Update(context.Context, spanner.Statement) (int64, error)
}

type BufferWriter interface {
	BufferWrite([]*spanner.Mutation) error
}

type DBWriter interface {
	Updater
	BufferWriter
}

type IndexReader interface {
	ReadUsingIndex(ctx context.Context, table, index string, keys spanner.KeySet, columns []string) *spanner.RowIterator
}

func DeleteByBuilder(ctx context.Context, db Updater, b squirrel.DeleteBuilder) (int64, error) {
	q, args, err := b.PlaceholderFormat(squirrel.AtP).ToSql()
	if err != nil {
		return 0, err
	}

	stmt := spanner.Statement{SQL: q, Params: ArgsToParams(args)}
	return db.Update(ctx, stmt)
}

func UpdateByBuilder(ctx context.Context, db Updater, b squirrel.UpdateBuilder) (int64, error) {
	q, args, err := b.PlaceholderFormat(squirrel.AtP).ToSql()
	if err != nil {
		return 0, err
	}

	stmt := spanner.Statement{SQL: q, Params: ArgsToParams(args)}
	return db.Update(ctx, stmt)
}

func GetResultsByBuilder[M any](ctx context.Context, db Queryer, b squirrel.SelectBuilder) ([]M, error) {
	q, args, err := b.PlaceholderFormat(squirrel.AtP).ToSql()
	if err != nil {
		return nil, err
	}

	stmt := spanner.Statement{SQL: q, Params: ArgsToParams(args)}
	return GetResults[M](db.Query(ctx, stmt))
}

func GetResultByBuilder[M any](ctx context.Context, db Queryer, b squirrel.SelectBuilder) (m M, err error) {
	q, args, err := b.PlaceholderFormat(squirrel.AtP).ToSql()
	if err != nil {
		return m, err
	}

	stmt := spanner.Statement{SQL: q, Params: ArgsToParams(args)}
	return GetResult[M](db.Query(ctx, stmt))
}

func GetByKey[M any, T tables.Table, K Key](ctx context.Context, db RowReader, k K) (m M, err error) {
	var t T
	row, err := db.ReadRow(ctx, t.TableName(), key(k), t.Columns())
	if err != nil {
		return m, err
	}
	return scanFunc[M]()(row)
}

func GetByKeys[M any, T tables.Table, K Key](ctx context.Context, db Reader, ids ...K) ([]M, error) {
	var t T
	return GetResults[M](db.Read(ctx, t.TableName(), keySet(ids...), t.Columns()))
}

func Delete[T tables.Table, K Key](db BufferWriter, keys ...K) error {
	if len(keys) == 0 {
		return nil
	}

	var t T
	m := spanner.Delete(t.TableName(), keySet(keys...))
	return db.BufferWrite([]*spanner.Mutation{m})
}

// GetResult returns the result from the iterator. Calls Stop after the iterator
// is finished.
func GetResult[M any](iter *spanner.RowIterator) (out M, err error) {
	defer iter.Stop()

	row, err := iter.Next()
	if err != nil {
		if err == iterator.Done {
			return out, status.Error(codes.NotFound, reflect.TypeOf(out).Name())
		}

		return out, err
	}

	return scanFunc[M]()(row)
}

// GetResults returns the results from the iterator. Calls Stop after the
// iterator is finished.
func GetResults[M any](iter *spanner.RowIterator) (out []M, err error) {
	f := scanFunc[M]()
	err = iter.Do(func(row *spanner.Row) error {
		m, err := f(row)
		if err != nil {
			return err
		}
		out = append(out, m)
		return nil
	})
	return out, err
}

// GetPtrResults returns the results from the iterator. Calls Stop after the
// iterator is finished.
func GetPtrResults[M any](iter *spanner.RowIterator) (out []*M, err error) {
	f := scanFunc[M]()
	err = iter.Do(
		func(row *spanner.Row) (err error) {
			m, err := f(row)
			if err == nil {
				out = append(out, &m)
			}
			return err
		},
	)
	return out, err
}

func structScanFunc[M any](row *spanner.Row) (m M, err error)    { return m, row.ToStructLenient(&m) }
func primitiveScanFunc[M any](row *spanner.Row) (m M, err error) { return m, row.Columns(&m) }
func scanFunc[M any]() func(row *spanner.Row) (M, error) {
	var m M
	if reflect.TypeOf(m).Kind() != reflect.Struct {
		return primitiveScanFunc[M]
	}
	return structScanFunc[M]
}

func ArgsToParams(args []interface{}) map[string]interface{} {
	params := make(map[string]interface{}, len(args))
	for n := 0; n < len(args); n++ {
		params[fmt.Sprintf("p%d", n+1)] = args[n]
	}

	return params
}

type Key interface {
	string | spanner.Key | spanner.KeyRange
}

func key[K Key](key K) spanner.Key {
	switch ids := (interface{})(key).(type) {
	case spanner.Key:
		return ids
	case string:
		return spanner.Key{ids}
	}
	return nil
}

func keySet[K Key](keys ...K) spanner.KeySet {
	if len(keys) == 0 {
		return spanner.AllKeys()
	}

	switch ids := (interface{})(keys).(type) {
	case []spanner.KeyRange:
		keys := make([]spanner.KeySet, len(ids))
		for i, id := range ids {
			keys[i] = id
		}
		return spanner.KeySets(keys...)
	case []spanner.Key:
		return spanner.KeySetFromKeys(ids...)
	case []string:
		keys := make([]spanner.Key, len(ids))
		for i, id := range ids {
			keys[i] = spanner.Key{id}
		}
		return spanner.KeySetFromKeys(keys...)
	}
	return nil
}