cloud-services/pkg/redisv2/conn.go

package redisv2

import (
	"context"
	"encoding/json"
	"sync"
	"time"

	"github.com/fiskerinc/cloud-services/pkg/logger"
	"github.com/pkg/errors"
	"github.com/redis/go-redis/v9"
)

// This is a wrapper around a redis connection that does a bunch of
// different stuff

func NewClient(redisClient *redis.Client) (client *Connection) {
	if redisClient == nil {
		redisClient = NewConnection()
	}

	client = &Connection{
		Client: redisClient,
	}
	return
}

// Client defines the function signatures associated with sending messages
//
//	and setting/getting objects
type ClientInterface interface {
	Close() error
	Ping() error

	GetClient() *redis.Client
	SetClient(*redis.Client)

	queueMessage(string, interface{}) error
	publishMessage(string, interface{}) error

	BatchQueueMessages(ids []string, messages []interface{}) error
	BatchPublishMessages(ids []string, messages []interface{}) error

	SafeQueueMessage(string, interface{}) error
	SafePublishMessage(string, interface{}) error

	// Simple redis operations
	Set(string, interface{}) error
	Get(string) (interface{}, error)
	Delete([]string) error
	GetMulti(ids []string) ([]interface{}, error)

	// Sets
	NewSet(string, interface{}, time.Duration) error
	GetSet(string, interface{}) error
	AddToSet(id string, data interface{}, expire time.Duration) error

	// Use objects when you wish to access individual fields in future
	SetObject(string, interface{}, time.Duration) error
	SetObjectField(string, string, interface{}) error
	SetObjects([]string, []interface{}, time.Duration) error
	GetObject(string, interface{}) error
	GetObjectField(string, string) (string, error)
	GetObjectMap(string) (map[string]string, error)
	GetObjectRaw(string) (map[string][]byte, error)
	GetObjectsMulti([]string, []interface{}) error
	GetObjectsMultiMap([]string) (map[string]map[string]string, error)

	// General execution
	Retrieve(command string, data interface{}) error

	// Cache functions marshal/unmarshal any data type to redis
	SetCache(string, interface{}, time.Duration) error

	// Thread-safe variations
	SafeSet(string, interface{}) error
	SafeGet(string) (interface{}, error)
	SafeDelete([]string) error

	SafeNewSet(string, interface{}, time.Duration) error
	SafeGetSet(string, interface{}) error

	SafeSetObject(string, interface{}, time.Duration) error
	SafeGetObject(string, interface{}) error

	Execute(command ...interface{}) (interface{}, error)
	SafeExecute(command ...interface{}) (interface{}, error)

	ExecuteBatch(batch *RedisBatchCommands) ([]interface{}, error)
	SafeExecuteBatch(batch *RedisBatchCommands) ([]interface{}, error)

	Keys(pattern string) (keys []string, err error)
}

// Connection holds a client to redis
//
//	The methods for connection are NOT thread safe.
type Connection struct {
	*redis.Client
	once sync.Once
	mu   sync.Mutex
}

// Retrieve implements ClientInterface.
func (*Connection) Retrieve(command string, data interface{}) error {
	panic("unimplemented")
}

var _ ClientInterface = &Connection{}

func (c *Connection) GetClient() (client *redis.Client) {
	return c.Client
}

func (c *Connection) SetClient(newClient *redis.Client) {
	c.Client = newClient
}

// Close the client if it exists
func (c *Connection) Close() error {
	if c.Client == nil {
		return nil
	}

	err := c.Client.Close()
	if err != nil {
		return errors.WithStack(err)
	}

	c.Client = nil

	return nil
}

func (c *Connection) do(commandName string, args ...interface{}) (reply interface{}, err error) {
	p := append([]interface{}{commandName}, args...)
	reply, err = c.Do(context.Background(), p...).Result()
	if err != nil {
		return reply, errors.WithStack(err)
	}
	return reply, nil
}

// func (c *Connection) send(commandName string, args ...interface{}) error {
// 	err := c.Client.Send(commandName, args...)
// 	if err != nil {
// 		return errors.WithStack(err)
// 	}

// 	return nil
// }

func (c *Connection) Ping() error {
	return c.Client.Ping(context.Background()).Err()
}

// QueueMessage writes a message to the corresponding list
// follows the format "queue:<id>"
// Messages are guaranteed to be delivered upon websocket connection
func (c *Connection) queueMessage(id string, message interface{}) error {
	data, err := json.Marshal(message)
	if err != nil {
		return errors.WithStack(err)
	}

	_, err = c.do("RPUSH", QueueKey(id), data)
	if err != nil {
		logger.At(logger.Error(), ChannelKey(id), "redis").
			Str("msg", string(data)).Err(err).Send()
	}

	c.do("EXPIRE", QueueKey(id), 3600) // 3600 = 1hr

	return err
}

// PublishMessage writes a message to the corresponding channel
// follows the format "channel:<id>"
// This is a fire and forget mechanism
func (c *Connection) publishMessage(id string, message interface{}) error {
	data, err := json.Marshal(message)
	if err != nil {
		return errors.WithStack(err)
	}

	_, err = c.do("PUBLISH", ChannelKey(id), data)
	if err != nil {
		return err
	}

	logger.At(logger.Debug(), ChannelKey(id), "redis").
		Str("msg", string(data)).
		Msgf("sent redis msg to %s", id)

	return nil
}

// BatchQueueMessages is the same as QueueMessage except performs a batch call
func (c *Connection) BatchQueueMessages(ids []string, messages []interface{}) error {
	if len(ids) != len(messages) {
		return errors.Errorf(
			"mismatch number of ids and messages. have %d ids and %d messages",
			len(ids),
			len(messages),
		)
	}

	batch := NewRedisBatchCommands()
	for i := 0; i < len(ids); i++ {
		data, err := json.Marshal(messages[i])
		if err != nil {
			return errors.WithStack(err)
		}
		batch.Add(Command{
			Command:   "RPUSH",
			Arguments: []interface{}{QueueKey(ids[i]), data},
		})
	}

	_, err := c.ExecuteBatch(batch)
	if err != nil {
		return err
	}

	return nil
}

// BatchPublishMessages is the same as PublishMessage except performs a batch call
func (c *Connection) BatchPublishMessages(ids []string, messages []interface{}) error {
	if len(ids) != len(messages) {
		return errors.Errorf(
			"mismatch number of ids and messages. have %d ids and %d messages",
			len(ids),
			len(messages),
		)
	}

	batch := NewRedisBatchCommands()
	for i := 0; i < len(ids); i++ {
		data, err := json.Marshal(messages[i])
		if err != nil {
			return errors.WithStack(err)
		}
		batch.Add(Command{
			Command:   "PUBLISH",
			Arguments: []interface{}{ChannelKey(ids[i]), data},
		})
	}

	_, err := c.ExecuteBatch(batch)
	if err != nil {
		return err
	}

	return nil
}

// SafeQueueMessage is the thread-safe implementation of QueueMessage
func (c *Connection) SafeQueueMessage(id string, message interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.queueMessage(id, message)
}

// SafePublishMessage is the thread-safe implementation of PublishMessage
func (c *Connection) SafePublishMessage(id string, message interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.publishMessage(id, message)
}

// Set replicates redis "SET"
func (c *Connection) Set(id string, data interface{}) error {
	if _, err := c.do("SET", id, data); err != nil {
		return err
	}

	return nil
}

// Get replicates redis "GET", must properly unpack interface returned
func (c *Connection) Get(id string) (interface{}, error) {
	data, err := c.do("GET", id)
	return data, err
}

// Delete removes all ids inputted
func (c *Connection) Delete(id []string) error {
	numDeleted64, err := c.Del(context.Background(), id...).Result()
	if err != nil {
		return err
	}
	numDeleted := int(numDeleted64)
	if numDeleted != len(id) {
		return errors.Errorf(
			"tried to delete %v (total: %v), however only %v were deleted",
			id, len(id), numDeleted,
		)
	}

	return nil
}

// Deprecated: NewSet adds items to a set in redis
func (c *Connection) NewSet(id string, data interface{}, expire time.Duration) error {
	var err error

	pipe := c.TxPipeline()

	err = pipe.Del(context.Background(), id).Err()
	if err != nil {
		pipe.Discard()
		return err
	}

	err = pipe.SAdd(context.Background(), id, data).Err()
	if err != nil {
		pipe.Discard()
		return err
	}

	if expire > 0 {
		err = pipe.Expire(context.Background(), id, expire).Err()
		if err != nil {
			pipe.Discard()
			return err
		}
	}

	_, err = pipe.Exec(context.Background())
	err = errors.WithStack(err)

	return err
}

// AddToSet adds item to a set in redis
func (c *Connection) AddToSet(id string, data interface{}, expire time.Duration) error {
	var err error

	pipe := c.TxPipeline()

	err = pipe.SAdd(context.Background(), id, data).Err()
	if err != nil {
		pipe.Discard()
		return err
	}

	if expire > 0 {
		err = pipe.Expire(context.Background(), id, expire).Err()
		if err != nil {
			pipe.Discard()
			return err
		}
	}

	_, err = pipe.Exec(context.Background())
	err = errors.WithStack(err)

	return err
}

// Deprecated: GetSet retrieves items from a set in redis
func (c *Connection) GetSet(id string, data interface{}) (err error) {
	res := c.SMembers(context.Background(), id)
	if res.Err() != nil {
		return err
	}

	err = res.ScanSlice(data)
	return err
}

// Deprecated: SetObject assigns the hash key id to the data object
// data can be of any type. Expire is in seconds, use -1 for no expire
func (c *Connection) SetObject(id string, data interface{}, expire time.Duration) error {
	var err error

	if expire > 0 {
		err = c.setExpiringObject(id, data, expire)
	} else {
		err = c.setPersistentObject(id, data)
	}

	return err
}

func (c *Connection) setPersistentObject(id string, data interface{}) error {
	err := c.HSet(context.Background(), id, data).Err()
	return err
}

func (c *Connection) setExpiringObject(id string, data interface{}, expire time.Duration) error {
	var err error

	pipe := c.TxPipeline()

	err = pipe.HSet(context.Background(), id, data).Err()
	if err != nil {
		pipe.Discard()
		return err
	}

	err = pipe.Expire(context.Background(), id, time.Duration(expire)).Err()
	if err != nil {
		pipe.Discard()
		return err
	}

	_, err = pipe.Exec(context.Background())

	return err
}

// Deprecated: SetObjectField sets a specific key of an object
func (c *Connection) SetObjectField(id string, key string, data interface{}) error {
	err := c.Do(context.Background(), "HSET", id, key, data).Err()
	return err
}

// Deprecated: SetObjects provides the same functionality as SetObject for multiple objects in one call to redis
func (c *Connection) SetObjects(ids []string, data []interface{}, expire time.Duration) error {
	var err error

	if len(ids) <= 0 || len(data) <= 0 || len(ids) != len(data) {
		return errors.Errorf("invalid lengths entered, lengths must match and be > 0. ids length: %v data length: %v",
			len(ids),
			len(data),
		)
	}

	pipe := c.TxPipeline()

	for i := range ids {
		err = pipe.Do(context.Background(), "HSET", ids[i], data[i]).Err()
		if err != nil {
			pipe.Discard()
			return err
		}
	}

	if expire > 0 {
		for _, id := range ids {
			err = pipe.Expire(context.Background(), id, time.Duration(expire)).Err()
			if err != nil {
				pipe.Discard()
				return err
			}
		}
	}

	_, err = pipe.Exec(context.Background())
	return err
}

// Deprecated: GetObject retrieves an object based off the hash key id
// and "unmarshals" it to struct pointer given
func (c *Connection) GetObject(id string, dest interface{}) error {
	result := c.HGetAll(context.Background(), id)
	err := result.Err()
	if err != nil {
		return err
	}

	err = result.Scan(&dest)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Deprecated: GetObjectMap retrieves an object based off the hash key id
// and returns it as a map[string]interface{}. GetObject()
// is preferred if you know the object being retrieved
func (c *Connection) GetObjectMap(id string) (map[string]string, error) {
	result := c.HGetAll(context.Background(), id)
	err := result.Err()
	if result.Err() != nil {
		return nil, err
	}
	object := result.Val()
	return object, err
}

// Deprecated: GetObjectRaw retrieves an object based off the hash key id
// and returns it as a map[string][]byte. Use this method when you have
// a hash of objects that you want to unmarshal.
func (c *Connection) GetObjectRaw(id string) (map[string][]byte, error) {
	var m map[string][]byte

	result := c.HGetAll(context.Background(), id)
	values, err := result.Result()
	if err != nil {
		return m, err
	}

	m = make(map[string][]byte, len(values))
	for keyS, valueS := range values {
		value := []byte(valueS)
		m[keyS] = value
	}

	return m, nil
}

// Deprecated: GetObjectField retrieves the value associated with the id and key of a hash map
func (c *Connection) GetObjectField(id string, key string) (string, error) {
	res := c.HGet(context.Background(), id, key)
	value, err := res.Result()
	return value, err
}

// Deprecated: GetObjectsMulti retrieves an array of objects based off
// the hash key ids given and returns the data as a
// []interface{}. Use this function if you know
// you need to get multiple objects from redis (one call to server).
func (c *Connection) GetObjectsMulti(ids []string, data []interface{}) error {
	var err error

	if len(ids) != len(data) {
		return errors.Errorf(
			"number of ids given %v does not match size of data slice %v",
			len(ids),
			len(data),
		)
	}

	pipe := c.TxPipeline()

	results := make([]*redis.MapStringStringCmd, 0, len(ids))
	for _, id := range ids {
		res := pipe.HGetAll(context.Background(), id)
		err := res.Err()
		if err != nil {
			pipe.Discard()
			return err
		}
		results = append(results, res)
	}

	_, err = pipe.Exec(context.Background())
	if err != nil {
		return err
	}

	for i, value := range results {
		err = value.Scan(&data[i])
		if err != nil {
			err = errors.WithStack(err)
		}
	}

	return nil
}

// Deprecated: GetObjectsMultiMap retrieves an array of objects based off
// the hash key ids given and returns the data as a
// map[string]interface{}. Use this function if you know
// you need to get multiple objects from redis (one call to server).
func (c *Connection) GetObjectsMultiMap(ids []string) (map[string]map[string]string, error) {
	var err error

	objects := make(map[string]map[string]string)

	pipe := c.TxPipeline()

	results := make([]*redis.MapStringStringCmd, 0, len(ids))
	for _, id := range ids {
		res := pipe.HGetAll(context.Background(), id)
		err := res.Err()
		if err != nil {
			pipe.Discard()
			return nil, err
		}
		results = append(results, res)
	}

	_, err = pipe.Exec(context.Background())
	if err != nil {
		return nil, err
	}

	for i, value := range results {
		err = value.Err()
		if err != nil {
			return objects, errors.WithStack(err)
		}

		o := value.Val()
		objects[ids[i]] = o
	}

	return objects, nil
}

func (c *Connection) makeKeys(ids []string) []interface{} {
	keys := make([]interface{}, len(ids))

	for i, id := range ids {
		keys[i] = id
	}

	return keys
}

// Deprecated: GetMulti
func (c *Connection) GetMulti(ids []string) ([]interface{}, error) {
	if len(ids) == 0 {
		return nil, errors.WithStack(errors.New("cannot call redis MGET with no keys"))
	}

	res := c.MGet(context.Background(), ids...)
	result, err := res.Result()

	return result, err
}

// Deprecated: SetCache marshals object and inserts into redis based on key id
// sets expiration to expire time.Duration
func (c *Connection) SetCache(id string, data interface{}, expire time.Duration) error {
	var err error
	pipe := c.TxPipeline()

	serialized, err := json.Marshal(data)
	if err != nil {
		return errors.WithStack(err)
	}

	err = pipe.Set(context.Background(), id, serialized, time.Duration(0)).Err()
	if err != nil {
		pipe.Discard()
		return err
	}

	if expire > 0 {
		err = pipe.Expire(context.Background(), id, time.Duration(expire)).Err()
		if err != nil {
			pipe.Discard()
			return err
		}
	}

	_, err = pipe.Exec(context.Background())
	err = errors.WithStack(err)

	return err
}

// Deprecated: SafeSet is the thread-safe version of Set()
func (c *Connection) SafeSet(id string, data interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.Set(id, data)
}

// Deprecated: SafeGet is the thread-safe version of Get()
func (c *Connection) SafeGet(id string) (interface{}, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.Get(id)
}

// Deprecated: SafeDelete is the thread-safe version of Delete()
func (c *Connection) SafeDelete(ids []string) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.Delete(ids)
}

// Deprecated: SafeNewSet is the thread-safe version of NewSet()
func (c *Connection) SafeNewSet(id string, data interface{}, expire time.Duration) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.NewSet(id, data, expire)
}

// Deprecated: SafeGetSet is the thread-safe version of GetSet()
func (c *Connection) SafeGetSet(id string, data interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.GetSet(id, data)
}

// Deprecated: SafeSetObject provides a thread-safe SetObject()
func (c *Connection) SafeSetObject(id string, data interface{}, expire time.Duration) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.SetObject(id, data, expire)
}

// Deprecated: SafeGetObject provides a thread-safe GetObject()
func (c *Connection) SafeGetObject(id string, data interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.GetObject(id, data)
}

// Execute provides an abstraction over redigo's Do method
// use this method over specialized methods within this library
func (c *Connection) Execute(command ...interface{}) (interface{}, error) {
	var reply interface{}
	if len(command) < 2 {
		return reply, ErrInvalidCommand
	}

	reply, err := c.do(command[0].(string), command[1:]...)
	if err != nil {
		return reply, err
	}

	return reply, nil
}

// SafeExecute provides a thread-safe Execute method
func (c *Connection) SafeExecute(command ...interface{}) (interface{}, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.Execute(command)
}

// ExecuteBatch sends all commands stored to Redis
// removes all commands regardless of success or failure
func (c *Connection) ExecuteBatch(batch *RedisBatchCommands) ([]interface{}, error) {
	if batch.IsEmpty() {
		return nil, nil
	}

	return c.executeBatch(batch)
}

func (c *Connection) SafeExecuteBatch(batch *RedisBatchCommands) ([]interface{}, error) {
	if batch.IsEmpty() {
		return nil, nil
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	return c.executeBatch(batch)
}

func (c *Connection) executeBatch(batch *RedisBatchCommands) ([]interface{}, error) {
	var err error
	defer batch.Clear()

	pipe := c.TxPipeline()
	responses := make([]interface{}, 0, len(batch.Commands))
	for _, command := range batch.Commands {
		keyword := command.Command
		// There has to be a slightly nicer way to do this
		commandList := []interface{}{keyword}
		commandList = append(commandList, command.Arguments...)
		
		res := pipe.Do(context.Background(), commandList...)
		if res.Err() != nil {
			pipe.Discard()
			return responses, err
		}
		// This will yield no responses I think
		responses = append(responses, res.Val())
	}

	_, err = pipe.Exec(context.Background())
	return responses, err
}

func (c *Connection) Keys(pattern string) (keys []string, err error) {
	keys, err = c.Keys(pattern)
	return
}

// // redisArgs is a helper function to generate redis args
// func redisArgs(id string, data interface{}) redis.Args {
// 	return redis.Args{}.Add(id).AddFlat(data)
// }

// // redisArgsMulti is a helper function to generate redis args for hash map fields and arrays
// func redisArgsMulti(data ...interface{}) redis.Args {
// 	return redis.Args{}.Add(data...)
// }