redis_broker.go

// Copyright (c) 2019 Sick Yoon
// This file is part of gocelery which is released under MIT license.
// See file LICENSE for full license details.

package gocelery

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

	"github.com/redis/go-redis/v9"
)

// RedisCeleryBroker is celery broker for redis
type RedisCeleryBroker struct {
	redis.UniversalClient
	QueueName string
}

// NewRedisBroker creates new RedisCeleryBroker with given redis connection pool
func NewRedisBroker(conn redis.UniversalClient) *RedisCeleryBroker {
	return &RedisCeleryBroker{
		UniversalClient: conn,
		QueueName:       "celery",
	}
}

// NewRedisCeleryBroker creates new RedisCeleryBroker based on given uri
//
// Deprecated: NewRedisCeleryBroker exists for historical compatibility
// and should not be used. Use NewRedisBroker instead to create new RedisCeleryBroker.
func NewRedisCeleryBroker(uri string) *RedisCeleryBroker {
	connectOpions, err := redis.ParseURL(uri)
	if err != nil {
		return nil
	}

	client := redis.NewClient(connectOpions)

	return &RedisCeleryBroker{
		UniversalClient: client,
		QueueName:       "celery",
	}
}

// SendCeleryMessage sends CeleryMessage to redis queue
func (cb *RedisCeleryBroker) SendCeleryMessage(ctx context.Context, message *CeleryMessage) error {
	jsonBytes, err := json.Marshal(message)
	if err != nil {
		return err
	}
	resp := cb.LPush(ctx, message.Properties.DeliveryInfo.RoutingKey, jsonBytes)
	if resp.Err() != nil {
		return err
	}
	return nil
}

// GetCeleryMessage retrieves celery message from redis queue
func (cb *RedisCeleryBroker) GetCeleryMessage(ctx context.Context) (*CeleryMessage, error) {
	resp := cb.BRPop(ctx, time.Second*1, cb.QueueName)
	if resp.Err() != nil {
		return nil, resp.Err()
	}
	if resp.Val() == nil {
		return nil, fmt.Errorf("null message received from redis")
	}
	messageList := resp.Val()
	if messageList[0] != cb.QueueName {
		return nil, fmt.Errorf("not a celery message: %v", messageList[0])
	}
	var message CeleryMessage
	if err := json.Unmarshal([]byte(messageList[1]), &message); err != nil {
		return nil, err
	}
	return &message, nil
}

// GetTaskMessage retrieves task message from redis queue
func (cb *RedisCeleryBroker) GetTaskMessage(ctx context.Context) (*TaskMessage, error) {
	celeryMessage, err := cb.GetCeleryMessage(ctx)
	if err != nil {
		return nil, err
	}
	return celeryMessage.GetTaskMessage(), nil
}