Files
SugarDB/src/modules/pubsub/commands.go

215 lines
7.1 KiB
Go

// Copyright 2024 Kelvin Clement Mwinuka
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pubsub
import (
"context"
"errors"
"fmt"
"github.com/echovault/echovault/src/utils"
"net"
"strings"
)
func handleSubscribe(ctx context.Context, cmd []string, server utils.EchoVault, conn *net.Conn) ([]byte, error) {
pubsub, ok := server.GetPubSub().(*PubSub)
if !ok {
return nil, errors.New("could not load pubsub module")
}
channels := cmd[1:]
if len(channels) == 0 {
return nil, errors.New(utils.WrongArgsResponse)
}
withPattern := strings.EqualFold(cmd[0], "psubscribe")
pubsub.Subscribe(ctx, conn, channels, withPattern)
return nil, nil
}
func handleUnsubscribe(ctx context.Context, cmd []string, server utils.EchoVault, conn *net.Conn) ([]byte, error) {
pubsub, ok := server.GetPubSub().(*PubSub)
if !ok {
return nil, errors.New("could not load pubsub module")
}
channels := cmd[1:]
withPattern := strings.EqualFold(cmd[0], "punsubscribe")
return pubsub.Unsubscribe(ctx, conn, channels, withPattern), nil
}
func handlePublish(ctx context.Context, cmd []string, server utils.EchoVault, _ *net.Conn) ([]byte, error) {
pubsub, ok := server.GetPubSub().(*PubSub)
if !ok {
return nil, errors.New("could not load pubsub module")
}
if len(cmd) != 3 {
return nil, errors.New(utils.WrongArgsResponse)
}
pubsub.Publish(ctx, cmd[2], cmd[1])
return []byte(utils.OkResponse), nil
}
func handlePubSubChannels(_ context.Context, cmd []string, server utils.EchoVault, _ *net.Conn) ([]byte, error) {
if len(cmd) > 3 {
return nil, errors.New(utils.WrongArgsResponse)
}
pubsub, ok := server.GetPubSub().(*PubSub)
if !ok {
return nil, errors.New("could not load pubsub module")
}
pattern := ""
if len(cmd) == 3 {
pattern = cmd[2]
}
return pubsub.Channels(pattern), nil
}
func handlePubSubNumPat(_ context.Context, _ []string, server utils.EchoVault, _ *net.Conn) ([]byte, error) {
pubsub, ok := server.GetPubSub().(*PubSub)
if !ok {
return nil, errors.New("could not load pubsub module")
}
num := pubsub.NumPat()
return []byte(fmt.Sprintf(":%d\r\n", num)), nil
}
func handlePubSubNumSubs(_ context.Context, cmd []string, server utils.EchoVault, _ *net.Conn) ([]byte, error) {
pubsub, ok := server.GetPubSub().(*PubSub)
if !ok {
return nil, errors.New("could not load pubsub module")
}
return pubsub.NumSub(cmd[2:]), nil
}
func Commands() []utils.Command {
return []utils.Command{
{
Command: "subscribe",
Categories: []string{utils.PubSubCategory, utils.ConnectionCategory, utils.SlowCategory},
Description: "(SUBSCRIBE channel [channel ...]) Subscribe to one or more channels.",
Sync: false,
KeyExtractionFunc: func(cmd []string) ([]string, error) {
// Treat the channels as keys
if len(cmd) < 2 {
return nil, errors.New(utils.WrongArgsResponse)
}
return cmd[1:], nil
},
HandlerFunc: handleSubscribe,
},
{
Command: "psubscribe",
Categories: []string{utils.PubSubCategory, utils.ConnectionCategory, utils.SlowCategory},
Description: "(PSUBSCRIBE pattern [pattern ...]) Subscribe to one or more glob patterns.",
Sync: false,
KeyExtractionFunc: func(cmd []string) ([]string, error) {
// Treat the patterns as keys
if len(cmd) < 2 {
return nil, errors.New(utils.WrongArgsResponse)
}
return cmd[1:], nil
},
HandlerFunc: handleSubscribe,
},
{
Command: "publish",
Categories: []string{utils.PubSubCategory, utils.FastCategory},
Description: "(PUBLISH channel message) Publish a message to the specified channel.",
Sync: true,
KeyExtractionFunc: func(cmd []string) ([]string, error) {
// Treat the channel as a key
if len(cmd) != 3 {
return nil, errors.New(utils.WrongArgsResponse)
}
return []string{cmd[1]}, nil
},
HandlerFunc: handlePublish,
},
{
Command: "unsubscribe",
Categories: []string{utils.PubSubCategory, utils.ConnectionCategory, utils.SlowCategory},
Description: `(UNSUBSCRIBE [channel [channel ...]]) Unsubscribe from a list of channels.
If the channel list is not provided, then the connection will be unsubscribed from all the channels that
it's currently subscribe to.`,
Sync: false,
KeyExtractionFunc: func(cmd []string) ([]string, error) {
// Treat the channels as keys
return cmd[1:], nil
},
HandlerFunc: handleUnsubscribe,
},
{
Command: "punsubscribe",
Categories: []string{utils.PubSubCategory, utils.ConnectionCategory, utils.SlowCategory},
Description: `(PUNSUBSCRIBE [pattern [pattern ...]]) Unsubscribe from a list of channels using patterns.
If the pattern list is not provided, then the connection will be unsubscribed from all the patterns that
it's currently subscribe to.`,
Sync: false,
KeyExtractionFunc: func(cmd []string) ([]string, error) {
// Treat the channels as keys
return cmd[1:], nil
},
HandlerFunc: handleUnsubscribe,
},
{
Command: "pubsub",
Categories: []string{},
Description: "",
Sync: false,
KeyExtractionFunc: func(cmd []string) ([]string, error) { return []string{}, nil },
HandlerFunc: func(_ context.Context, _ []string, _ utils.EchoVault, _ *net.Conn) ([]byte, error) {
return nil, errors.New("provide CHANNELS, NUMPAT, or NUMSUB subcommand")
},
SubCommands: []utils.SubCommand{
{
Command: "channels",
Categories: []string{utils.PubSubCategory, utils.SlowCategory},
Description: `(PUBSUB CHANNELS [pattern]) Returns an array containing the list of channels that
match the given pattern. If no pattern is provided, all active channels are returned. Active channels are
channels with 1 or more subscribers.`,
Sync: false,
KeyExtractionFunc: func(cmd []string) ([]string, error) { return []string{}, nil },
HandlerFunc: handlePubSubChannels,
},
{
Command: "numpat",
Categories: []string{utils.PubSubCategory, utils.SlowCategory},
Description: `(PUBSUB NUMPAT) Return the number of patterns that are currently subscribed to by clients.`,
Sync: false,
KeyExtractionFunc: func(cmd []string) ([]string, error) { return []string{}, nil },
HandlerFunc: handlePubSubNumPat,
},
{
Command: "numsub",
Categories: []string{utils.PubSubCategory, utils.SlowCategory},
Description: `(PUBSUB NUMSUB [channel [channel ...]]) Return an array of arrays containing the provided
channel name and how many clients are currently subscribed to the channel.`,
Sync: false,
KeyExtractionFunc: func(cmd []string) ([]string, error) { return cmd[2:], nil },
HandlerFunc: handlePubSubNumSubs,
},
},
},
}
}