oc-scheduler/controllers/streams.go

package controllers

import (
	"context"
	"fmt"
	"net/http"
	"time"

	oclib "cloud.o-forge.io/core/oc-lib"
	"cloud.o-forge.io/core/oc-lib/dbs"
	"cloud.o-forge.io/core/oc-lib/models/booking"
	libutils "cloud.o-forge.io/core/oc-lib/models/utils"
	"cloud.o-forge.io/core/oc-lib/models/workflow_execution"
	"cloud.o-forge.io/core/oc-lib/tools"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

// streamMsg is the envelope pushed over every stream WebSocket.
type streamMsg struct {
	Type    string      `json:"type"` // "snapshot" | "update" | "delete"
	Data    interface{} `json:"data,omitempty"`
	Deleted bool        `json:"deleted,omitempty"`
}

// ---------------------------------------------------------------------------
// Booking stream
// ---------------------------------------------------------------------------

// BookingStreamHandler opens a WebSocket that:
//  1. sends an immediate snapshot of matching bookings ("snapshot")
//  2. pushes each subsequent create/update/delete as an individual "update" or
//     "delete" message.
//
// Query params (all optional):
//
//	executions_id — filter to a specific scheduling session
//	is_draft      — "true" | "false"  (omit = non-draft)
//	start_date    — YYYY-MM-DD  (expected_start_date >=)
//	end_date      — YYYY-MM-DD  (expected_start_date <=)
func BookingStreamHandler(w http.ResponseWriter, r *http.Request) {
	user, peerID, groups := oclib.ExtractTokenInfoWs(*r)

	q := r.URL.Query()
	executionID := q.Get("execution_id")
	executionsID := q.Get("executions_id")
	isDraftStr := q.Get("is_draft")
	onlyDraft := isDraftStr == "true"
	filterDraft := isDraftStr != "" // whether the caller wants draft filtering at all
	startDate, _ := time.ParseInLocation("2006-01-02", q.Get("start_date"), time.UTC)
	endDate, _ := time.ParseInLocation("2006-01-02", q.Get("end_date"), time.UTC)

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	matchesFilter := func(b *booking.Booking) bool {
		if executionID != "" && b.GetID() != executionID {
			return false
		}
		if executionsID != "" && b.ExecutionsID != executionsID {
			return false
		}
		if filterDraft && b.IsDraft != onlyDraft {
			return false
		}
		if !startDate.IsZero() && b.ExpectedStartDate.Before(startDate) {
			return false
		}
		if !endDate.IsZero() && b.ExpectedStartDate.After(endDate) {
			return false
		}
		return true
	}

	// Build snapshot filters
	andF := map[string][]dbs.Filter{}
	if executionID != "" {
		andF["id"] = []dbs.Filter{{Operator: dbs.EQUAL.String(), Value: executionID}}
	}
	if executionsID != "" {
		andF["executions_id"] = []dbs.Filter{{Operator: dbs.EQUAL.String(), Value: executionsID}}
	}
	if !startDate.IsZero() {
		andF["expected_start_date"] = append(andF["expected_start_date"],
			dbs.Filter{Operator: "gte", Value: primitive.NewDateTimeFromTime(startDate)})
	}
	if !endDate.IsZero() {
		andF["expected_start_date"] = append(andF["expected_start_date"],
			dbs.Filter{Operator: "lte", Value: primitive.NewDateTimeFromTime(endDate)})
	}
	var snapshotFilter *dbs.Filters
	if len(andF) > 0 {
		snapshotFilter = &dbs.Filters{And: andF}
	}

	snapshot := oclib.NewRequest(oclib.LibDataEnum(oclib.BOOKING), user, peerID, groups, nil).
		Search(snapshotFilter, "", onlyDraft, 0, 10000)
	if err := conn.WriteJSON(streamMsg{Type: "snapshot", Data: snapshot.Data}); err != nil {
		return
	}

	changeCh, unsub := libutils.SubscribeChanges(tools.BOOKING)
	defer unsub()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// detect client disconnect
	closeCh := make(chan struct{})
	go func() {
		defer close(closeCh)
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}()

	for {
		select {
		case evt := <-changeCh:
			b, ok := evt.Object.(*booking.Booking)
			if !ok || !matchesFilter(b) {
				continue
			}
			if evt.Deleted {
				_ = conn.WriteJSON(streamMsg{Type: "delete", Data: b, Deleted: true})
			} else {
				_ = conn.WriteJSON(streamMsg{Type: "update", Data: b})
			}
		case <-closeCh:
			return
		case <-ctx.Done():
			return
		}
	}
}

// ---------------------------------------------------------------------------
// WorkflowExecution stream
// ---------------------------------------------------------------------------

// ExecutionStreamHandler opens a WebSocket that:
//  1. sends an immediate snapshot of matching executions ("snapshot")
//  2. pushes each subsequent create/update/delete as "update" or "delete".
//
// Query params (all optional):
//
//	executions_id — filter to a specific scheduling session
//	is_draft      — "true" | "false"  (omit = non-draft)
//	start_date    — YYYY-MM-DD  (execution_date >=)
//	end_date      — YYYY-MM-DD  (execution_date <=)
func ExecutionStreamHandler(w http.ResponseWriter, r *http.Request) {
	user, peerID, groups := oclib.ExtractTokenInfoWs(*r)

	q := r.URL.Query()
	executionID := q.Get("execution_id")
	executionsID := q.Get("executions_id")
	isDraftStr := q.Get("is_draft")
	onlyDraft := isDraftStr == "true"
	filterDraft := isDraftStr != ""
	startDate, _ := time.ParseInLocation("2006-01-02", q.Get("start_date"), time.UTC)
	endDate, _ := time.ParseInLocation("2006-01-02", q.Get("end_date"), time.UTC)

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	matchesFilter := func(e *workflow_execution.WorkflowExecution) bool {
		if executionID != "" && e.GetID() != executionID {
			return false
		}
		if executionsID != "" && e.ExecutionsID != executionsID {
			return false
		}
		if filterDraft && e.IsDraft != onlyDraft {
			return false
		}
		if !startDate.IsZero() && e.ExecDate.Before(startDate) {
			return false
		}
		if !endDate.IsZero() && e.ExecDate.After(endDate) {
			return false
		}
		return true
	}

	// Build snapshot filters
	andF := map[string][]dbs.Filter{}
	if executionID != "" {
		andF["id"] = []dbs.Filter{{Operator: dbs.EQUAL.String(), Value: executionID}}
	}
	if executionsID != "" {
		andF["executions_id"] = []dbs.Filter{{Operator: dbs.EQUAL.String(), Value: executionsID}}
	}
	if !startDate.IsZero() {
		andF["execution_date"] = append(andF["execution_date"],
			dbs.Filter{Operator: "gte", Value: primitive.NewDateTimeFromTime(startDate)})
	}
	if !endDate.IsZero() {
		andF["execution_date"] = append(andF["execution_date"],
			dbs.Filter{Operator: "lte", Value: primitive.NewDateTimeFromTime(endDate)})
	}
	var snapshotFilter *dbs.Filters
	if len(andF) > 0 {
		snapshotFilter = &dbs.Filters{And: andF}
	}

	snapshot := oclib.NewRequest(oclib.LibDataEnum(oclib.WORKFLOW_EXECUTION), user, peerID, groups, nil).
		Search(snapshotFilter, "", onlyDraft, 0, 10000)
	if err := conn.WriteJSON(streamMsg{Type: "snapshot", Data: snapshot.Data}); err != nil {
		return
	}

	changeCh, unsub := libutils.SubscribeChanges(tools.WORKFLOW_EXECUTION)
	defer unsub()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	closeCh := make(chan struct{})
	go func() {
		defer close(closeCh)
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}()

	for {
		select {
		case evt := <-changeCh:
			e, ok := evt.Object.(*workflow_execution.WorkflowExecution)
			fmt.Println("CHANGE!", e, ok, matchesFilter(e))
			if !ok || !matchesFilter(e) {
				continue
			}
			if evt.Deleted {
				_ = conn.WriteJSON(streamMsg{Type: "delete", Data: e, Deleted: true})
			} else {
				_ = conn.WriteJSON(streamMsg{Type: "update", Data: e})
			}
		case <-closeCh:
			return
		case <-ctx.Done():
			return
		}
	}
}