SSE support for global-query (server implementation)

cmd/global-query/pkg/distributed/query.go

@@ -20,6 +20,8 @@ import (
 type QueryRunner struct {
 	resolver hosts.Resolver
 	querier  Querier
+
+	onResult func(*results.Result) error
 }
 
 // QueryOption configures the query runner
@@ -37,6 +39,18 @@ func NewQueryRunner(resolver hosts.Resolver, querier Querier, opts ...QueryOptio
 	return
 }
 
+// SetResultReceivedFn registers a callback to be executed for every results that is
+// read off the results channel
+func (q *QueryRunner) SetResultReceivedFn(f func(*results.Result) error) *QueryRunner {
+	q.onResult = f
+	return q
+}
+
+// ResultReceived calls the result callback with res and
+func (q *QueryRunner) ResultReceived(res *results.Result) error {
+	return q.onResult(res)
+}
+
 // Run executes / runs the query and creates the final result structure
 func (q *QueryRunner) Run(ctx context.Context, args *query.Args) (*results.Result, error) {
 	// use a copy of the arguments, since some fields are modified by the querier
@@ -74,7 +88,7 @@ func (q *QueryRunner) Run(ctx context.Context, args *query.Args) (*results.Resul
 	logger.Info("reading query results from querier")
 
 	finalResult := aggregateResults(ctx, stmt,
-		q.querier.Query(ctx, hostList, &queryArgs),
+		q.querier.Query(ctx, hostList, &queryArgs), q.onResult,
 	)
 
 	finalResult.End()
@@ -115,7 +129,7 @@ func (q *QueryRunner) prepareHostList(ctx context.Context, queryHosts string) (h
 
 // aggregateResults takes finished query workloads from the workloads channel, aggregates the result by merging the rows and summaries,
 // and returns the final result. The `tracker` variable provides information about potential Run failures for individual hosts
-func aggregateResults(ctx context.Context, stmt *query.Statement, queryResults <-chan *results.Result) (finalResult *results.Result) {
+func aggregateResults(ctx context.Context, stmt *query.Statement, queryResults <-chan *results.Result, onResult func(*results.Result) error) (finalResult *results.Result) {
 	ctx, span := tracing.Start(ctx, "aggregateResults")
 	defer span.End()
 
@@ -193,6 +207,11 @@ func aggregateResults(ctx context.Context, stmt *query.Statement, queryResults <
 			// take the total from the query result. Since there may be overlap between the queries of two
 			// different systems, the overlap has to be deducted from the total
 			finalResult.Summary.Hits.Total += res.Summary.Hits.Total - merged
+
+			err := onResult(finalResult)
+			if err != nil {
+				logger.With("error", err).Error("failed to call results callback")
+			}
 		}
 	}
 }

pkg/api/api.go

@@ -17,4 +17,7 @@ const (
 
 	// ValidationRoute is the route to validate a goquery query
 	ValidationRoute = QueryRoute + "/validate"
+
+	// SSEQueryRoute runs a goquery query with a return channel for partial results
+	SSEQueryRoute = QueryRoute + "/sse"
 )

pkg/api/query.go

@@ -3,6 +3,8 @@ package api
 import (
 	"context"
 
+	"github.com/danielgtaylor/huma/v2/sse"
+	"github.com/els0r/goProbe/cmd/global-query/pkg/distributed"
 	"github.com/els0r/goProbe/pkg/query"
 	"github.com/els0r/goProbe/pkg/results"
 	"github.com/els0r/telemetry/logging"
@@ -22,6 +24,21 @@ func getBodyQueryRunnerHandler(caller string, querier query.Runner) func(context
 	}
 }
 
+func getSSEBodyQueryRunnerHandler(caller string, querier *distributed.QueryRunner) func(context.Context, *ArgsInput, sse.Sender) {
+	return func(ctx context.Context, input *ArgsInput, send sse.Sender) {
+		querier.SetResultReceivedFn(func(res *results.Result) error {
+			return send.Data(&PartialResult{res})
+		})
+
+		res, err := runQuery(ctx, caller, input.Body, querier)
+		if err != nil {
+			send.Data(err)
+			return
+		}
+		_ = send.Data(&FinalResult{res})
+	}
+}
+
 func runQuery(ctx context.Context, caller string, args *query.Args, querier query.Runner) (*results.Result, error) {
 	// make sure all defaults are available if they weren't set explicitly
 	args.SetDefaults()

pkg/api/query_api_ops.go

@@ -4,6 +4,8 @@ import (
 	"net/http"
 
 	"github.com/danielgtaylor/huma/v2"
+	"github.com/danielgtaylor/huma/v2/sse"
+	"github.com/els0r/goProbe/cmd/global-query/pkg/distributed"
 	"github.com/els0r/goProbe/pkg/query"
 	"github.com/els0r/goProbe/pkg/results"
 )
@@ -48,6 +50,32 @@ func RegisterQueryAPI(a huma.API, caller string, querier query.Runner, middlewar
 		},
 		getBodyQueryRunnerHandler(caller, querier),
 	)
+
+	// register routes specific to distributed querying
+	dqr, ok := querier.(*distributed.QueryRunner)
+	if ok {
+		registerDistributedQueryAPI(a, caller, dqr, middlewares)
+	}
+}
+
+func registerDistributedQueryAPI(a huma.API, caller string, qr *distributed.QueryRunner, middlewares huma.Middlewares) {
+	sse.Register(a,
+		huma.Operation{
+			OperationID: "query-post-run-sse",
+			Method:      http.MethodPost,
+			Path:        SSEQueryRoute,
+			Summary:     "Run query with server sent events (SSE)",
+			Description: "Runs a query based on the parameters provided in the body. Pushes back partial results via SSE",
+			Middlewares: middlewares,
+			Tags:        queryTags,
+		},
+		map[string]any{
+			"queryError":    &query.DetailError{},
+			"partialResult": &PartialResult{},
+			"finalResult":   &FinalResult{},
+		},
+		getSSEBodyQueryRunnerHandler(caller, qr),
+	)
 }
 
 // ArgsBodyInput stores the query args to be validated in the body
@@ -66,3 +94,13 @@ type ArgsParamsInput struct {
 type QueryResultOutput struct {
 	Body *results.Result
 }
+
+// PartialResult represents an update to the results structure. It SHOULD only be used if the
+// results.Result object will be further modified / aggregated. This data structure is relevant
+// only in the context of SSE
+type PartialResult struct{ *results.Result }
+
+// FinalResult represents the result which is sent after all aggregation of partial results has
+// completed. It SHOULD only be sent at the end of a streaming operation. This data structure is relevant
+// only in the context of SSE
+type FinalResult struct{ *results.Result }