collection_watcher.go

package mongo_watcher

import (
	"context"
	"github.com/techpro-studio/gomongo"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
	"log"
	"time"
)

type MongoSchema[T any] interface {
	GetId() bson.ObjectID
	gomongo.ModelConverted[T]
}

type MongoEvent string

const MongoEventInsert MongoEvent = "insert"
const MongoEventUpdate MongoEvent = "update"
const MongoEventDelete MongoEvent = "delete"
const MongoEventReplace MongoEvent = "replace"

type Event[T any] struct {
	Type                     MongoEvent
	FullDocument             *T
	FullDocumentBeforeChange *T
	Key                      string
}

type EventHandler[T any] interface {
	Setup(ctx context.Context, collection *mongo.Collection)
	HandleEvent(ctx context.Context, event *Event[T]) error
}

type eventHandlerFunction[T any] func(ctx context.Context, event *Event[T]) error

func FunctionEventHandler[T any](function func(ctx context.Context, event *Event[T]) error) EventHandler[T] {
	return eventHandlerFunction[T](function)
}

func (s eventHandlerFunction[T]) Setup(ctx context.Context, collection *mongo.Collection) {}
func (s eventHandlerFunction[T]) HandleEvent(ctx context.Context, event *Event[T]) error {
	return s(ctx, event)
}

type CollectionWatcher[T any, M MongoSchema[T]] struct {
	collection    *mongo.Collection
	eventHandlers []EventHandler[T]
	allowedEvents []MongoEvent
}

func NewCollectionWatcher[T any, M MongoSchema[T]](
	collection *mongo.Collection,
	eventHandlers []EventHandler[T],
	allowedEvents []MongoEvent,
) *CollectionWatcher[T, M] {
	if len(allowedEvents) == 0 {
		allowedEvents = []MongoEvent{MongoEventUpdate, MongoEventReplace, MongoEventInsert, MongoEventDelete}
	}
	for _, handler := range eventHandlers {
		handler.Setup(context.Background(), collection)
	}
	return &CollectionWatcher[T, M]{
		collection:    collection,
		eventHandlers: eventHandlers,
		allowedEvents: allowedEvents,
	}
}

func (c *CollectionWatcher[T, M]) Watch(ctx context.Context) {
	var resumeToken bson.Raw

	pipeline := mongo.Pipeline{
		{{"$match", bson.D{
			{"operationType", bson.D{{"$in", c.allowedEvents}}},
		}}},
	}

	for {
		select {
		case <-ctx.Done():
			log.Printf("Watcher for %s stopping: context cancelled", c.collection.Name())
			return
		default:
		}

		err := c.runStream(ctx, pipeline, &resumeToken)
		if err != nil {
			log.Printf("Change stream error on %s: %v", c.collection.Name(), err)
		}

		// If context was cancelled, stop retrying
		if ctx.Err() != nil {
			log.Printf("Watcher for %s stopping: context cancelled", c.collection.Name())
			return
		}

		log.Printf("Reconnecting watcher for %s in 3s...", c.collection.Name())
		select {
		case <-ctx.Done():
			return
		case <-time.After(3 * time.Second):
		}
	}
}

func (c *CollectionWatcher[T, M]) runStream(ctx context.Context, pipeline mongo.Pipeline, resumeToken *bson.Raw) error {
	opts := options.ChangeStream().
		SetFullDocument(options.UpdateLookup).
		SetFullDocumentBeforeChange(options.WhenAvailable)

	// Resume from where we left off if we have a token
	if *resumeToken != nil {
		opts.SetResumeAfter(*resumeToken)
	}

	changeStream, err := c.collection.Watch(ctx, pipeline, opts)
	if err != nil {
		return err
	}
	defer func() {
		closeCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if cerr := changeStream.Close(closeCtx); cerr != nil {
			log.Printf("Error closing change stream: %v\n", cerr)
		}
	}()

	log.Printf("Watching %s (resuming: %v)", c.collection.Name(), *resumeToken != nil)

	for changeStream.Next(ctx) {
		*resumeToken = changeStream.ResumeToken()

		var raw struct {
			FullDocument             M      `bson:"fullDocument"`
			FullDocumentBeforeChange *M     `bson:"fullDocumentBeforeChange"`
			Type                     string `bson:"operationType"`
			DocumentKey              struct {
				ID bson.ObjectID `bson:"_id"`
			} `bson:"documentKey"`
		}

		if err := changeStream.Decode(&raw); err != nil {
			log.Printf("error decoding change stream: %v\n", err)
			continue
		}

		event := c.buildEvent(raw.Type, raw.DocumentKey.ID.Hex(), raw.FullDocument, raw.FullDocumentBeforeChange)
		if event == nil {
			continue
		}

		for _, handler := range c.eventHandlers {
			if err := handler.HandleEvent(ctx, event); err != nil {
				log.Printf("Error handling event: %v\n", err)
			}
		}
	}

	return changeStream.Err()
}

func (c *CollectionWatcher[T, M]) buildEvent(opType string, key string, doc M, beforeDoc *M) *Event[T] {
	switch MongoEvent(opType) {
	case MongoEventInsert, MongoEventUpdate, MongoEventReplace:
		data := *doc.ToModel()
		var before *T
		if beforeDoc != nil {
			before = (*beforeDoc).ToModel()
		}
		return &Event[T]{
			Type:                     MongoEvent(opType),
			FullDocument:             &data,
			FullDocumentBeforeChange: before,
			Key:                      key,
		}
	case MongoEventDelete:
		return &Event[T]{
			Type: MongoEventDelete,
			Key:  key,
		}
	default:
		log.Printf("Unhandled operationType: %s", opType)
		return nil
	}
}