dataRT/data/mongo/event.go

package mongo

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

	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

const (
	dbevent string = "cl"
	tbevent string = "events"
)

const (
	EventStatusHappen = iota
	EventStatusDataAt
	EventStatusReport
	EventStatusConfirm
	EventStatusPersist
	EventStatusClose
)

const (
	EventActionHappen  = "happen"
	EventActionDataAt  = "data_attach"
	EventActionReport  = "report"
	EventActionConfirm = "confirm"
	EventActionPersist = "persist"
	EventActionClose   = "close"
)

var EventStatusAction = []string{
	EventStatusHappen:  EventActionHappen,
	EventStatusDataAt:  EventActionDataAt,
	EventStatusReport:  EventActionReport,
	EventStatusConfirm: EventActionConfirm,
	EventStatusPersist: EventActionPersist,
	EventStatusClose:   EventActionClose,
}

type operation struct {
	Action string `bson:"action" json:"action"`
	OP     string `bson:"op" json:"op"`
	TS     int64  `bson:"ts" json:"ts"`
}

func GenOperation(action, op string) operation {
	return operation{
		Action: action,
		OP:     op,
		TS:     time.Now().UnixMilli(),
	}
}

type Event struct {
	Event      string       `bson:"event" json:"event"`
	EventUUID  string       `bson:"event_uuid" json:"event_uuid"`
	Type       int          `bson:"type" json:"type"`
	Priority   int          `bson:"priority" json:"priority"` // 0~9
	Status     int          `bson:"status" json:"status"`
	Timestamp  int64        `bson:"timestamp" json:"timestamp"`
	From       string       `bson:"from" json:"from"`
	Operations []*operation `bson:"operations" json:"operations"`
	// others
	Alarm *Alarm `bson:"alarm" json:"alarm"`
}

func (e *Event) Marshall() ([]byte, error) {
	return json.Marshal(e)
}

func InsertEvent(ctx context.Context, doc any) error {
	_, err := getCollection(dbevent, tbevent).InsertOne(ctx, doc)
	return err
}

func DeleteEvent[T bson.M | bson.D](ctx context.Context, filter T) error {
	_, err := getCollection(dbevent, tbevent).DeleteOne(ctx, filter)
	return err
}

func UpdateEvent[T bson.M | bson.D](ctx context.Context, filter T, update T) error {
	_, err := getCollection(dbevent, tbevent).UpdateOne(ctx, filter, update)
	return err
}

func FindEventsWithPageLimit[T bson.M | bson.D](ctx context.Context, filter T,
	sort int, page int64, limit int64) ([]*Event, error) {

	opt := options.Find()
	if sort == 1 || sort == -1 {
		opt.SetSort(bson.D{{Key: "timestamp", Value: sort}})
	} else {
		opt.SetSort(bson.D{{Key: "_id", Value: 1}})
	}

	if page > 0 && limit > 0 {
		opt.SetSkip(limit * (page - 1)).SetLimit(limit)
	}

	cursor, err := getCollection(dbevent, tbevent).Find(ctx, filter, opt)
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var docs []*Event
	for cursor.Next(ctx) {
		doc := new(Event)
		if err := cursor.Decode(doc); err != nil {
			return nil, err
		}
		docs = append(docs, doc)
	}
	if err := cursor.Err(); err != nil {
		return docs, err
	}

	return docs, nil
}

func BulkWriteEventsWithUUID(ctx context.Context, curd byte, events []map[string]any) ([]string, error) {
	length := len(events)
	if length <= 0 {
		return nil, errors.New("no event")
	}

	models := make([]mongo.WriteModel, 0, length)
	idx2UUID := make(map[int]string, length)
	for i, event := range events {
		uuid, ok := event["event_uuid"].(string)
		if !ok || uuid == "" {
			return nil, fmt.Errorf("invalid uuid at index %d", i)
		}

		switch curd {
		case 'c':
			models = append(models,
				mongo.NewInsertOneModel().
					SetDocument(event))

		case 'u':
			filter := bson.M{"event_uuid": uuid}
			status, ok := event["status"].(int)
			if !ok {
				return nil, fmt.Errorf("invalid status at index %d", i)
			}
			operation, ok := event["operation"].(operation)
			if !ok {
				return nil, fmt.Errorf("invalid operation at index %d", i)
			}
			update := bson.M{"$set": bson.M{"status": status},
				"$push": bson.M{"operations": operation}}
			models = append(models,
				mongo.NewUpdateOneModel().
					SetFilter(filter).
					SetUpdate(update))

		case 'd':
			filter := bson.M{"event_uuid": uuid}
			models = append(models,
				mongo.NewDeleteOneModel().
					SetFilter(filter))

		default:
			return nil, fmt.Errorf("invalid curd")
		}

		idx2UUID[i] = uuid
	}

	opts := options.BulkWrite().SetOrdered(false)
	_, err := getCollection(dbevent, tbevent).BulkWrite(ctx, models, opts)

	sUUIDs := []string{}
	if err != nil {
		if bulkErr, ok := err.(mongo.BulkWriteException); ok {
			idxExist := map[int]bool{}
			for _, writeErr := range bulkErr.WriteErrors {
				idxExist[writeErr.Index] = true
			}

			for idx, uuid := range idx2UUID {
				if !idxExist[idx] {
					sUUIDs = append(sUUIDs, uuid)
				}
			}
		}
	}

	return sUUIDs, err
}

func genEventType(sys int, level int) int {
	return sys + level*3
}