Adding upstream version 0.28.1.

Signed-off-by: Daniel Baumann <daniel@debian.org>

tools/subscriptions/broker.go

@@ -0,0 +1,65 @@
+package subscriptions
+
+import (
+	"fmt"
+
+	"github.com/pocketbase/pocketbase/tools/list"
+	"github.com/pocketbase/pocketbase/tools/store"
+)
+
+// Broker defines a struct for managing subscriptions clients.
+type Broker struct {
+	store *store.Store[string, Client]
+}
+
+// NewBroker initializes and returns a new Broker instance.
+func NewBroker() *Broker {
+	return &Broker{
+		store: store.New[string, Client](nil),
+	}
+}
+
+// Clients returns a shallow copy of all registered clients indexed
+// with their connection id.
+func (b *Broker) Clients() map[string]Client {
+	return b.store.GetAll()
+}
+
+// ChunkedClients splits the current clients into a chunked slice.
+func (b *Broker) ChunkedClients(chunkSize int) [][]Client {
+	return list.ToChunks(b.store.Values(), chunkSize)
+}
+
+// TotalClients returns the total number of registered clients.
+func (b *Broker) TotalClients() int {
+	return b.store.Length()
+}
+
+// ClientById finds a registered client by its id.
+//
+// Returns non-nil error when client with clientId is not registered.
+func (b *Broker) ClientById(clientId string) (Client, error) {
+	client, ok := b.store.GetOk(clientId)
+	if !ok {
+		return nil, fmt.Errorf("no client associated with connection ID %q", clientId)
+	}
+
+	return client, nil
+}
+
+// Register adds a new client to the broker instance.
+func (b *Broker) Register(client Client) {
+	b.store.Set(client.Id(), client)
+}
+
+// Unregister removes a single client by its id and marks it as discarded.
+//
+// If client with clientId doesn't exist, this method does nothing.
+func (b *Broker) Unregister(clientId string) {
+	client := b.store.Get(clientId)
+	if client == nil {
+		return
+	}
+	client.Discard()
+	b.store.Remove(clientId)
+}

tools/subscriptions/broker_test.go

@@ -0,0 +1,134 @@
+package subscriptions_test
+
+import (
+	"testing"
+
+	"github.com/pocketbase/pocketbase/tools/subscriptions"
+)
+
+func TestNewBroker(t *testing.T) {
+	b := subscriptions.NewBroker()
+
+	if b.Clients() == nil {
+		t.Fatal("Expected clients map to be initialized")
+	}
+}
+
+func TestClients(t *testing.T) {
+	b := subscriptions.NewBroker()
+
+	if total := len(b.Clients()); total != 0 {
+		t.Fatalf("Expected no clients, got %v", total)
+	}
+
+	b.Register(subscriptions.NewDefaultClient())
+	b.Register(subscriptions.NewDefaultClient())
+
+	// check if it is a shallow copy
+	clients := b.Clients()
+	for k := range clients {
+		delete(clients, k)
+	}
+
+	// should return a new copy
+	if total := len(b.Clients()); total != 2 {
+		t.Fatalf("Expected 2 clients, got %v", total)
+	}
+}
+
+func TestChunkedClients(t *testing.T) {
+	b := subscriptions.NewBroker()
+
+	chunks := b.ChunkedClients(2)
+	if total := len(chunks); total != 0 {
+		t.Fatalf("Expected %d chunks, got %d", 0, total)
+	}
+
+	b.Register(subscriptions.NewDefaultClient())
+	b.Register(subscriptions.NewDefaultClient())
+	b.Register(subscriptions.NewDefaultClient())
+
+	chunks = b.ChunkedClients(2)
+	if total := len(chunks); total != 2 {
+		t.Fatalf("Expected %d chunks, got %d", 2, total)
+	}
+
+	if total := len(chunks[0]); total != 2 {
+		t.Fatalf("Expected the first chunk to have 2 clients, got %d", total)
+	}
+
+	if total := len(chunks[1]); total != 1 {
+		t.Fatalf("Expected the second chunk to have 1 client, got %d", total)
+	}
+}
+
+func TestTotalClients(t *testing.T) {
+	b := subscriptions.NewBroker()
+
+	if total := b.TotalClients(); total != 0 {
+		t.Fatalf("Expected no clients, got %d", total)
+	}
+
+	b.Register(subscriptions.NewDefaultClient())
+	b.Register(subscriptions.NewDefaultClient())
+
+	if total := b.TotalClients(); total != 2 {
+		t.Fatalf("Expected %d clients, got %d", 2, total)
+	}
+}
+
+func TestClientById(t *testing.T) {
+	b := subscriptions.NewBroker()
+
+	clientA := subscriptions.NewDefaultClient()
+	clientB := subscriptions.NewDefaultClient()
+	b.Register(clientA)
+	b.Register(clientB)
+
+	resultClient, err := b.ClientById(clientA.Id())
+	if err != nil {
+		t.Fatalf("Expected client with id %s, got error %v", clientA.Id(), err)
+	}
+	if resultClient.Id() != clientA.Id() {
+		t.Fatalf("Expected client %s, got %s", clientA.Id(), resultClient.Id())
+	}
+
+	if c, err := b.ClientById("missing"); err == nil {
+		t.Fatalf("Expected error, found client %v", c)
+	}
+}
+
+func TestRegister(t *testing.T) {
+	b := subscriptions.NewBroker()
+
+	client := subscriptions.NewDefaultClient()
+	b.Register(client)
+
+	if _, err := b.ClientById(client.Id()); err != nil {
+		t.Fatalf("Expected client with id %s, got error %v", client.Id(), err)
+	}
+}
+
+func TestUnregister(t *testing.T) {
+	b := subscriptions.NewBroker()
+
+	clientA := subscriptions.NewDefaultClient()
+	clientB := subscriptions.NewDefaultClient()
+	b.Register(clientA)
+	b.Register(clientB)
+
+	if _, err := b.ClientById(clientA.Id()); err != nil {
+		t.Fatalf("Expected client with id %s, got error %v", clientA.Id(), err)
+	}
+
+	b.Unregister(clientA.Id())
+
+	if c, err := b.ClientById(clientA.Id()); err == nil {
+		t.Fatalf("Expected error, found client %v", c)
+	}
+
+	// clientB shouldn't have been removed
+	if _, err := b.ClientById(clientB.Id()); err != nil {
+		t.Fatalf("Expected client with id %s, got error %v", clientB.Id(), err)
+	}
+}

tools/subscriptions/client.go

@@ -0,0 +1,280 @@
+package subscriptions
+
+import (
+	"encoding/json"
+	"net/url"
+	"strings"
+	"sync"
+
+	"github.com/pocketbase/pocketbase/tools/inflector"
+	"github.com/pocketbase/pocketbase/tools/security"
+	"github.com/spf13/cast"
+)
+
+const optionsParam = "options"
+
+// SubscriptionOptions defines the request options (query params, headers, etc.)
+// for a single subscription topic.
+type SubscriptionOptions struct {
+	Query   map[string]string `json:"query"`
+	Headers map[string]string `json:"headers"`
+}
+
+// Client is an interface for a generic subscription client.
+type Client interface {
+	// Id Returns the unique id of the client.
+	Id() string
+
+	// Channel returns the client's communication channel.
+	//
+	// NB! The channel shouldn't be used after calling Discard().
+	Channel() chan Message
+
+	// Subscriptions returns a shallow copy of the client subscriptions matching the prefixes.
+	// If no prefix is specified, returns all subscriptions.
+	Subscriptions(prefixes ...string) map[string]SubscriptionOptions
+
+	// Subscribe subscribes the client to the provided subscriptions list.
+	//
+	// Each subscription can also have "options" (json serialized SubscriptionOptions) as query parameter.
+	//
+	// Example:
+	//
+	// 	Subscribe(
+	// 	    "subscriptionA",
+	// 	    `subscriptionB?options={"query":{"a":1},"headers":{"x_token":"abc"}}`,
+	// 	)
+	Subscribe(subs ...string)
+
+	// Unsubscribe unsubscribes the client from the provided subscriptions list.
+	Unsubscribe(subs ...string)
+
+	// HasSubscription checks if the client is subscribed to `sub`.
+	HasSubscription(sub string) bool
+
+	// Set stores any value to the client's context.
+	Set(key string, value any)
+
+	// Unset removes a single value from the client's context.
+	Unset(key string)
+
+	// Get retrieves the key value from the client's context.
+	Get(key string) any
+
+	// Discard marks the client as "discarded" (and closes its channel),
+	// meaning that it shouldn't be used anymore for sending new messages.
+	//
+	// It is safe to call Discard() multiple times.
+	Discard()
+
+	// IsDiscarded indicates whether the client has been "discarded"
+	// and should no longer be used.
+	IsDiscarded() bool
+
+	// Send sends the specified message to the client's channel (if not discarded).
+	Send(m Message)
+}
+
+// ensures that DefaultClient satisfies the Client interface
+var _ Client = (*DefaultClient)(nil)
+
+// DefaultClient defines a generic subscription client.
+type DefaultClient struct {
+	store         map[string]any
+	subscriptions map[string]SubscriptionOptions
+	channel       chan Message
+	id            string
+	mux           sync.RWMutex
+	isDiscarded   bool
+}
+
+// NewDefaultClient creates and returns a new DefaultClient instance.
+func NewDefaultClient() *DefaultClient {
+	return &DefaultClient{
+		id:            security.RandomString(40),
+		store:         map[string]any{},
+		channel:       make(chan Message),
+		subscriptions: map[string]SubscriptionOptions{},
+	}
+}
+
+// Id implements the [Client.Id] interface method.
+func (c *DefaultClient) Id() string {
+	c.mux.RLock()
+	defer c.mux.RUnlock()
+
+	return c.id
+}
+
+// Channel implements the [Client.Channel] interface method.
+func (c *DefaultClient) Channel() chan Message {
+	c.mux.RLock()
+	defer c.mux.RUnlock()
+
+	return c.channel
+}
+
+// Subscriptions implements the [Client.Subscriptions] interface method.
+//
+// It returns a shallow copy of the client subscriptions matching the prefixes.
+// If no prefix is specified, returns all subscriptions.
+func (c *DefaultClient) Subscriptions(prefixes ...string) map[string]SubscriptionOptions {
+	c.mux.RLock()
+	defer c.mux.RUnlock()
+
+	// no prefix -> return copy of all subscriptions
+	if len(prefixes) == 0 {
+		result := make(map[string]SubscriptionOptions, len(c.subscriptions))
+
+		for s, options := range c.subscriptions {
+			result[s] = options
+		}
+
+		return result
+	}
+
+	result := make(map[string]SubscriptionOptions)
+
+	for _, prefix := range prefixes {
+		for s, options := range c.subscriptions {
+			// "?" ensures that the options query start character is always there
+			// so that it can be used as an end separator when looking only for the main subscription topic
+			if strings.HasPrefix(s+"?", prefix) {
+				result[s] = options
+			}
+		}
+	}
+
+	return result
+}
+
+// Subscribe implements the [Client.Subscribe] interface method.
+//
+// Empty subscriptions (aka. "") are ignored.
+func (c *DefaultClient) Subscribe(subs ...string) {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	for _, s := range subs {
+		if s == "" {
+			continue // skip empty
+		}
+
+		// extract subscription options (if any)
+		rawOptions := struct {
+			// note: any instead of string to minimize the breaking changes with earlier versions
+			Query   map[string]any `json:"query"`
+			Headers map[string]any `json:"headers"`
+		}{}
+		u, err := url.Parse(s)
+		if err == nil {
+			raw := u.Query().Get(optionsParam)
+			if raw != "" {
+				json.Unmarshal([]byte(raw), &rawOptions)
+			}
+		}
+
+		options := SubscriptionOptions{
+			Query:   make(map[string]string, len(rawOptions.Query)),
+			Headers: make(map[string]string, len(rawOptions.Headers)),
+		}
+
+		// normalize query
+		// (currently only single string values are supported for consistency with the default routes handling)
+		for k, v := range rawOptions.Query {
+			options.Query[k] = cast.ToString(v)
+		}
+
+		// normalize headers name and values, eg. "X-Token" is converted to "x_token"
+		// (currently only single string values are supported for consistency with the default routes handling)
+		for k, v := range rawOptions.Headers {
+			options.Headers[inflector.Snakecase(k)] = cast.ToString(v)
+		}
+
+		c.subscriptions[s] = options
+	}
+}
+
+// Unsubscribe implements the [Client.Unsubscribe] interface method.
+//
+// If subs is not set, this method removes all registered client's subscriptions.
+func (c *DefaultClient) Unsubscribe(subs ...string) {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	if len(subs) > 0 {
+		for _, s := range subs {
+			delete(c.subscriptions, s)
+		}
+	} else {
+		// unsubscribe all
+		for s := range c.subscriptions {
+			delete(c.subscriptions, s)
+		}
+	}
+}
+
+// HasSubscription implements the [Client.HasSubscription] interface method.
+func (c *DefaultClient) HasSubscription(sub string) bool {
+	c.mux.RLock()
+	defer c.mux.RUnlock()
+
+	_, ok := c.subscriptions[sub]
+
+	return ok
+}
+
+// Get implements the [Client.Get] interface method.
+func (c *DefaultClient) Get(key string) any {
+	c.mux.RLock()
+	defer c.mux.RUnlock()
+
+	return c.store[key]
+}
+
+// Set implements the [Client.Set] interface method.
+func (c *DefaultClient) Set(key string, value any) {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	c.store[key] = value
+}
+
+// Unset implements the [Client.Unset] interface method.
+func (c *DefaultClient) Unset(key string) {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	delete(c.store, key)
+}
+
+// Discard implements the [Client.Discard] interface method.
+func (c *DefaultClient) Discard() {
+	c.mux.Lock()
+	defer c.mux.Unlock()
+
+	if c.isDiscarded {
+		return
+	}
+
+	close(c.channel)
+
+	c.isDiscarded = true
+}
+
+// IsDiscarded implements the [Client.IsDiscarded] interface method.
+func (c *DefaultClient) IsDiscarded() bool {
+	c.mux.RLock()
+	defer c.mux.RUnlock()
+
+	return c.isDiscarded
+}
+
+// Send sends the specified message to the client's channel (if not discarded).
+func (c *DefaultClient) Send(m Message) {
+	if c.IsDiscarded() {
+		return
+	}
+
+	c.Channel() <- m
+}

tools/subscriptions/client_test.go

@@ -0,0 +1,244 @@
+package subscriptions_test
+
+import (
+	"encoding/json"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/pocketbase/pocketbase/tools/subscriptions"
+)
+
+func TestNewDefaultClient(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	if c.Channel() == nil {
+		t.Errorf("Expected channel to be initialized")
+	}
+
+	if c.Subscriptions() == nil {
+		t.Errorf("Expected subscriptions map to be initialized")
+	}
+
+	if c.Id() == "" {
+		t.Errorf("Expected unique id to be set")
+	}
+}
+
+func TestId(t *testing.T) {
+	clients := []*subscriptions.DefaultClient{
+		subscriptions.NewDefaultClient(),
+		subscriptions.NewDefaultClient(),
+		subscriptions.NewDefaultClient(),
+		subscriptions.NewDefaultClient(),
+	}
+
+	ids := map[string]struct{}{}
+	for i, c := range clients {
+		// check uniqueness
+		if _, ok := ids[c.Id()]; ok {
+			t.Errorf("(%d) Expected unique id, got %v", i, c.Id())
+		} else {
+			ids[c.Id()] = struct{}{}
+		}
+
+		// check length
+		if len(c.Id()) != 40 {
+			t.Errorf("(%d) Expected unique id to have 40 chars length, got %v", i, c.Id())
+		}
+	}
+}
+
+func TestChannel(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	if c.Channel() == nil {
+		t.Fatalf("Expected channel to be initialized, got")
+	}
+}
+
+func TestSubscriptions(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	if len(c.Subscriptions()) != 0 {
+		t.Fatalf("Expected subscriptions to be empty")
+	}
+
+	c.Subscribe("sub1", "sub11", "sub2")
+
+	scenarios := []struct {
+		prefixes []string
+		expected []string
+	}{
+		{nil, []string{"sub1", "sub11", "sub2"}},
+		{[]string{"missing"}, nil},
+		{[]string{"sub1"}, []string{"sub1", "sub11"}},
+		{[]string{"sub2"}, []string{"sub2"}}, // with extra query start char
+	}
+
+	for _, s := range scenarios {
+		t.Run(strings.Join(s.prefixes, ","), func(t *testing.T) {
+			subs := c.Subscriptions(s.prefixes...)
+
+			if len(subs) != len(s.expected) {
+				t.Fatalf("Expected %d subscriptions, got %d", len(s.expected), len(subs))
+			}
+
+			for _, s := range s.expected {
+				if _, ok := subs[s]; !ok {
+					t.Fatalf("Missing subscription %q in \n%v", s, subs)
+				}
+			}
+		})
+	}
+}
+
+func TestSubscribe(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	subs := []string{"", "sub1", "sub2", "sub3"}
+	expected := []string{"sub1", "sub2", "sub3"}
+
+	c.Subscribe(subs...) // empty string should be skipped
+
+	if len(c.Subscriptions()) != 3 {
+		t.Fatalf("Expected 3 subscriptions, got %v", c.Subscriptions())
+	}
+
+	for i, s := range expected {
+		if !c.HasSubscription(s) {
+			t.Errorf("(%d) Expected sub %s", i, s)
+		}
+	}
+}
+
+func TestSubscribeOptions(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	sub1 := "test1"
+	sub2 := `test2?options={"query":{"name":123},"headers":{"X-Token":456}}`
+
+	c.Subscribe(sub1, sub2)
+
+	subs := c.Subscriptions()
+
+	scenarios := []struct {
+		name            string
+		expectedOptions string
+	}{
+		{sub1, `{"query":{},"headers":{}}`},
+		{sub2, `{"query":{"name":"123"},"headers":{"x_token":"456"}}`},
+	}
+
+	for _, s := range scenarios {
+		t.Run(s.name, func(t *testing.T) {
+			options, ok := subs[s.name]
+			if !ok {
+				t.Fatalf("Missing subscription \n%q \nin \n%v", s.name, subs)
+			}
+
+			rawBytes, err := json.Marshal(options)
+			if err != nil {
+				t.Fatal(err)
+			}
+			rawStr := string(rawBytes)
+
+			if rawStr != s.expectedOptions {
+				t.Fatalf("Expected options\n%v\ngot\n%v", s.expectedOptions, rawStr)
+			}
+		})
+	}
+}
+
+func TestUnsubscribe(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	c.Subscribe("sub1", "sub2", "sub3")
+
+	c.Unsubscribe("sub1")
+
+	if c.HasSubscription("sub1") {
+		t.Fatalf("Expected sub1 to be removed")
+	}
+
+	c.Unsubscribe( /* all */ )
+	if len(c.Subscriptions()) != 0 {
+		t.Fatalf("Expected all subscriptions to be removed, got %v", c.Subscriptions())
+	}
+}
+
+func TestHasSubscription(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	if c.HasSubscription("missing") {
+		t.Error("Expected false, got true")
+	}
+
+	c.Subscribe("sub")
+
+	if !c.HasSubscription("sub") {
+		t.Error("Expected true, got false")
+	}
+}
+
+func TestSetAndGet(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	c.Set("demo", 1)
+
+	result, _ := c.Get("demo").(int)
+
+	if result != 1 {
+		t.Errorf("Expected 1, got %v", result)
+	}
+}
+
+func TestDiscard(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	if v := c.IsDiscarded(); v {
+		t.Fatal("Expected false, got true")
+	}
+
+	c.Discard()
+
+	if v := c.IsDiscarded(); !v {
+		t.Fatal("Expected true, got false")
+	}
+}
+
+func TestSend(t *testing.T) {
+	c := subscriptions.NewDefaultClient()
+
+	received := []string{}
+	go func() {
+		for m := range c.Channel() {
+			received = append(received, m.Name)
+		}
+	}()
+
+	c.Send(subscriptions.Message{Name: "m1"})
+	c.Send(subscriptions.Message{Name: "m2"})
+	c.Discard()
+	c.Send(subscriptions.Message{Name: "m3"})
+	c.Send(subscriptions.Message{Name: "m4"})
+	time.Sleep(5 * time.Millisecond)
+
+	expected := []string{"m1", "m2"}
+
+	if len(received) != len(expected) {
+		t.Fatalf("Expected %d messages, got %d", len(expected), len(received))
+	}
+	for _, name := range expected {
+		var exists bool
+		for _, n := range received {
+			if n == name {
+				exists = true
+				break
+			}
+		}
+		if !exists {
+			t.Fatalf("Missing expected %q message, got %v", name, received)
+		}
+	}
+}

tools/subscriptions/message.go

@@ -0,0 +1,37 @@
+package subscriptions
+
+import (
+	"io"
+)
+
+// Message defines a client's channel data.
+type Message struct {
+	Name string `json:"name"`
+	Data []byte `json:"data"`
+}
+
+// WriteSSE writes the current message in a SSE format into the provided writer.
+//
+// For example, writing to a router.Event:
+//
+//	m := Message{Name: "users/create", Data: []byte{...}}
+//	m.Write(e.Response, "yourEventId")
+//	e.Flush()
+func (m *Message) WriteSSE(w io.Writer, eventId string) error {
+	parts := [][]byte{
+		[]byte("id:" + eventId + "\n"),
+		[]byte("event:" + m.Name + "\n"),
+		[]byte("data:"),
+		m.Data,
+		[]byte("\n\n"),
+	}
+
+	for _, part := range parts {
+		_, err := w.Write(part)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

tools/subscriptions/message_test.go

@@ -0,0 +1,25 @@
+package subscriptions_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/pocketbase/pocketbase/tools/subscriptions"
+)
+
+func TestMessageWrite(t *testing.T) {
+	m := subscriptions.Message{
+		Name: "test_name",
+		Data: []byte("test_data"),
+	}
+
+	var sb strings.Builder
+
+	m.WriteSSE(&sb, "test_id")
+
+	expected := "id:test_id\nevent:test_name\ndata:test_data\n\n"
+
+	if v := sb.String(); v != expected {
+		t.Fatalf("Expected writer content\n%q\ngot\n%q", expected, v)
+	}
+}