eway/internal/storage/badger/goodsitem.go

package badger

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"unsafe"

	"git.loyso.art/frx/eway/internal/encoding/fbs"
	"git.loyso.art/frx/eway/internal/entity"

	badger "github.com/dgraph-io/badger/v4"
	"github.com/dgraph-io/ristretto/z"
	"github.com/rs/zerolog"
)

const useJSON = false

type goodsItemClient struct {
	db *badger.DB
}

func newGoodsItemClient(db *badger.DB) *goodsItemClient {
	return &goodsItemClient{
		db: db,
	}
}

func (*goodsItemClient) prefix() []byte {
	return []byte("!!goodsitem!!")
}

func (c *goodsItemClient) prefixed(key []byte) []byte {
	return append(c.prefix(), key...)
}

func (c *goodsItemClient) prefixedStr(key string) []byte {
	keyBytes := unsafe.Slice(unsafe.StringData(key), len(key))
	return c.prefixed(keyBytes)
}

func (c *goodsItemClient) prefixedIDByCartStr(key int64) []byte {
	var keyBytes [8]byte
	binary.BigEndian.PutUint64(keyBytes[:], uint64(key))
	return c.prefixedIDByCart(keyBytes[:])
}

func (*goodsItemClient) prefixedIDByCart(key []byte) []byte {
	return append([]byte("!!goodsitem_card_idx!!"), key...)
}

func (c *goodsItemClient) ListIter(
	ctx context.Context, prefetchCount int,
) (out <-chan entity.GoodsItem, err error) {
	stream := c.db.NewStream()
	stream.Prefix = c.prefix()
	stream.LogPrefix = "list_iter"

	bus := make(chan entity.GoodsItem, prefetchCount)
	stream.Send = func(buf *z.Buffer) error {
		list, err := badger.BufferToKVList(buf)
		if err != nil {
			return err
		}

		for _, kv := range list.GetKv() {
			var gooditem entity.GoodsItem

			if useJSON {
				err = json.Unmarshal(kv.GetValue(), &gooditem)
				if err != nil {
					return err
				}
			} else {
				gooditem, err = fbs.ParseGoodsItem(kv.GetValue())
				if err != nil {
					return err
				}
			}

			bus <- gooditem
		}

		return nil
	}

	go func(ctx context.Context) {
		defer close(bus)

		err := stream.Orchestrate(ctx)
		if err != nil {
			zerolog.Ctx(ctx).Warn().Err(err).Msg("unable to orchestrate")
		}
		println("finished")
	}(ctx)

	return bus, nil
}

func (c *goodsItemClient) List(
	ctx context.Context,
) (out []entity.GoodsItem, err error) {
	err = c.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchValues = true
		opts.PrefetchSize = 20
		iter := txn.NewIterator(opts)
		defer iter.Close()

		prefix := c.prefix()
		var cursor int
		for iter.Seek(prefix); iter.ValidForPrefix(prefix); iter.Next() {
			cursor++
			current := iter.Item()
			err = current.Value(func(val []byte) error {
				var goodsItem entity.GoodsItem
				if useJSON {
					err := json.Unmarshal(val, &goodsItem)
					if err != nil {
						return err
					}
				} else {
					goodsItem, err = fbs.ParseGoodsItem(val)
					if err != nil {
						return err
					}
				}
				out = append(out, goodsItem)

				return nil
			})
			if err != nil {
				return fmt.Errorf("getting value: %w", err)
			}
		}

		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("viewing: %w", err)
	}

	return out, nil
}

func (c *goodsItemClient) Get(
	ctx context.Context,
	sku string,
) (out entity.GoodsItem, err error) {
	err = c.db.View(func(txn *badger.Txn) error {
		key := unsafe.Slice(unsafe.StringData(sku), len(sku))
		out, err = c.getBySKU(key, txn)
		return err
	})
	if err != nil {
		if errors.Is(err, badger.ErrKeyNotFound) {
			return out, entity.ErrNotFound
		}

		return out, fmt.Errorf("viewing: %w", err)
	}

	return out, nil
}

func (c *goodsItemClient) GetByCart(ctx context.Context, id int64) (out entity.GoodsItem, err error) {
	err = c.db.View(func(txn *badger.Txn) error {
		var idByte [8]byte
		binary.BigEndian.PutUint64(idByte[:], uint64(id))

		item, err := txn.Get(c.prefixedIDByCart(idByte[:]))
		if err != nil {
			return fmt.Errorf("getting key: %w", err)
		}

		sku := make([]byte, item.ValueSize())
		sku, err = item.ValueCopy(sku)
		if err != nil {
			return fmt.Errorf("getting value of idx: %w", err)
		}

		sku = bytes.TrimPrefix(sku, c.prefix())

		out, err = c.getBySKU(sku, txn)
		return err
	})
	if err != nil {
		if errors.Is(err, badger.ErrKeyNotFound) {
			return out, entity.ErrNotFound
		}

		return out, fmt.Errorf("viewing: %w", err)
	}

	return out, nil
}

func (c *goodsItemClient) UpsertMany(ctx context.Context, items ...entity.GoodsItem) ([]entity.GoodsItem, error) {
	return items, c.upsertByBatch(ctx, items)
}

func (c *goodsItemClient) upsertByBatch(ctx context.Context, items []entity.GoodsItem) error {
	batch := c.db.NewWriteBatch()
	defer batch.Cancel()

	log := zerolog.Ctx(ctx)

	for _, item := range items {
		select {
		case <-ctx.Done():
			break
		default:
		}
		key := c.prefixedStr(item.Articul)
		var value []byte
		if useJSON {
			value, _ = json.Marshal(item)
		} else {
			value = fbs.MakeDomainGoodItemFinished(item)
		}

		idxValue := make([]byte, len(key))
		copy(idxValue, key)

		coreEntry := badger.NewEntry(key, value)
		if err := batch.SetEntry(coreEntry); err != nil {
			log.Warn().Err(err).Msg("unable to set item, breaking")
			break
		}

		idxKey := c.prefixedIDByCartStr(item.Cart)
		idxEntry := badger.NewEntry(idxKey, idxValue)
		if err := batch.SetEntry(idxEntry); err != nil {
			log.Warn().Err(err).Msg("unable to set idx, breaking")
			break
		}
	}

	err := batch.Flush()
	if err != nil {
		return fmt.Errorf("flushing changes: %w", err)
	}

	return nil
}

func (c *goodsItemClient) getBySKU(sku []byte, txn *badger.Txn) (out entity.GoodsItem, err error) {
	item, err := txn.Get(c.prefixed(sku))
	if err != nil {
		return out, fmt.Errorf("getting key: %w", err)
	}

	err = item.Value(func(val []byte) error {
		if useJSON {
			return json.Unmarshal(val, &out)
		}

		out, err = fbs.ParseGoodsItem(val)
		return err
	})

	if err != nil {
		return out, fmt.Errorf("reading value: %w", err)
	}

	return out, nil
}