Connecting live tail webSocket

Author: nojaf

pkg/opensearch/lucene_handler.go

@@ -44,7 +44,8 @@ func (h *luceneHandler) processQuery(q *Query) error {
 	if len(q.BucketAggs) == 0 {
 		// If no aggregations, only trace, document, and logs queries are valid
 		if q.luceneQueryType != "Traces" {
-			if len(q.Metrics) == 0 || !(q.Metrics[0].Type == rawDataType || q.Metrics[0].Type == rawDocumentType) {
+			// Add logsType to the allowed types when no bucketAggs are present
+			if len(q.Metrics) == 0 || !(q.Metrics[0].Type == rawDataType || q.Metrics[0].Type == rawDocumentType || q.Metrics[0].Type == logsType) {
 				return fmt.Errorf("invalid query, missing metrics and aggregations")
 			}
 		}
@@ -132,20 +133,6 @@ func processLogsQuery(q *Query, b *client.SearchRequestBuilder, from, to int64,
 		size = defaultLogsSize
 	}
 	b.Size(size)
-
-	// For log query, we use only date histogram aggregation
-	aggBuilder := b.Agg()
-	defaultBucketAgg := &BucketAgg{
-		Type:  dateHistType,
-		Field: defaultTimeField,
-		ID:    "1",
-		Settings: utils.NewJsonFromAny(map[string]interface{}{
-			"interval": "auto",
-		})}
-	defaultBucketAgg.Settings = utils.NewJsonFromAny(
-		defaultBucketAgg.generateSettingsForDSL(),
-	)
-	_ = addDateHistogramAgg(aggBuilder, defaultBucketAgg, from, to, defaultTimeField)
 }
 
 func (bucketAgg BucketAgg) generateSettingsForDSL() map[string]interface{} {

pkg/opensearch/opensearch.go

@@ -25,7 +25,9 @@ import (
 // This is a list of interfaces that the Service implements.
 // Go will not compile if we don't implement all of these interfaces.
 var (
-	_ backend.StreamHandler = (*OpenSearchDatasource)(nil)
+	_ backend.QueryDataHandler    = (*OpenSearchDatasource)(nil)
+	_ backend.StreamHandler       = (*OpenSearchDatasource)(nil)
+	_ backend.CallResourceHandler = (*OpenSearchDatasource)(nil)
 )
 
 type datasourceInfo struct {
@@ -51,26 +53,35 @@ func newInstanceSettings(client *http.Client) datasource.InstanceFactoryFunc {
 // OpenSearchExecutor represents a handler for handling OpenSearch datasource request
 type OpenSearchExecutor struct{}
 
+// OpenSearchDatasource is an OpenSearch data source.
 type OpenSearchDatasource struct {
-	HttpClient *http.Client
-	backend.StreamHandler
-	im     instancemgmt.InstanceManager
-	logger log.Logger
+	// instance manager just holds a pointer to the open search client
+	im instancemgmt.InstanceManager
+
+	// httpClient is the general httpClient for the datasource to use for miscellaneous calls
+	httpClient *http.Client
+	logger     log.Logger
+
+	// streamQueries stores active queries for streaming sessions, keyed by a unique ID (e.g., refId)
+	streamQueries sync.Map
 }
 
+// NewOpenSearchDatasource creates a new OpenSearchDatasource and sets up its configuration.
 func NewOpenSearchDatasource(ctx context.Context, settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
 	log.DefaultLogger.Debug("Initializing new data source instance")
 
+	// Use a default http client for now
 	httpClient, err := client.NewDatasourceHttpClient(ctx, &settings)
 	if err != nil {
 		return nil, err
 	}
 
-	return &OpenSearchDatasource{
-		HttpClient: httpClient,
+	ds := &OpenSearchDatasource{
 		im:         datasource.NewInstanceManager(newInstanceSettings(httpClient)),
+		httpClient: httpClient,
 		logger:     backend.NewLoggerWith("logger", "tsdb.opensearch"),
-	}, nil
+	}
+	return ds, nil
 }
 
 // CheckHealth handles health checks sent from Grafana to the plugin.
@@ -146,7 +157,7 @@ func (ds *OpenSearchDatasource) CheckHealth(ctx context.Context, req *backend.Ch
 		}
 		request.Header = req.GetHTTPHeaders()
 
-		response, err := ds.HttpClient.Do(request)
+		response, err := ds.httpClient.Do(request)
 		if err != nil {
 			res.Status = backend.HealthStatusError
 			res.Message = err.Error()
@@ -213,12 +224,18 @@ func (ds *OpenSearchDatasource) CheckHealth(ctx context.Context, req *backend.Ch
 // The QueryDataResponse contains a map of RefID to the response for each query, and each response
 // contains Frames ([]*Frame).
 func (ds *OpenSearchDatasource) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	ds.logger.Info("QueryData called", "numQueries", len(req.Queries))
 	if len(req.Queries) == 0 {
 		return nil, fmt.Errorf("query contains no queries")
 	}
 
+	if len(req.Queries) > 0 {
+		jsonData, _ := req.Queries[0].JSON.MarshalJSON()
+		ds.logger.Info("QueryData - First query JSON", "refId", req.Queries[0].RefID, "json", string(jsonData), "interval", req.Queries[0].Interval.String(), "timeRangeFrom", req.Queries[0].TimeRange.From.String(), "timeRangeTo", req.Queries[0].TimeRange.To.String())
+	}
+
 	timeRange := req.Queries[0].TimeRange
-	osClient, err := client.NewClient(ctx, req.PluginContext.DataSourceInstanceSettings, ds.HttpClient, &timeRange)
+	osClient, err := client.NewClient(ctx, req.PluginContext.DataSourceInstanceSettings, ds.httpClient, &timeRange)
 	if err != nil {
 		return nil, err
 	}
@@ -229,7 +246,11 @@ func (ds *OpenSearchDatasource) QueryData(ctx context.Context, req *backend.Quer
 	}
 
 	query := newQueryRequest(osClient, req.Queries, req.PluginContext.DataSourceInstanceSettings)
+	ds.logger.Info("QueryData - About to execute query")
 	response, err = wrapError(query.execute(ctx))
+	if err != nil {
+		ds.logger.Error("QueryData - Error executing query", "error", err.Error())
+	}
 	return response, err
 }
 
@@ -335,6 +356,14 @@ func extractParametersFromServiceMapFrames(resp *backend.QueryDataResponse) ([]s
 }
 
 func (ds *OpenSearchDatasource) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+	ds.logger.Info("CallResource called", "path", req.Path, "method", req.Method)
+
+	// Route for registering stream queries
+	if strings.HasPrefix(req.Path, "_stream_query_register/") {
+		return ds.handleRegisterStreamQuery(ctx, req, sender)
+	}
+
+	// Existing resource call logic
 	// allowed paths for resource calls:
 	// - empty string for fetching db version
 	// - /_mapping for fetching index mapping, e.g. requests going to `index/_mapping`
@@ -355,7 +384,7 @@ func (ds *OpenSearchDatasource) CallResource(ctx context.Context, req *backend.C
 	}
 	request.Header = req.GetHTTPHeaders()
 
-	response, err := ds.HttpClient.Do(request)
+	response, err := ds.httpClient.Do(request)
 	if err != nil {
 		return err
 	}

pkg/opensearch/resource_handlers.go

@@ -0,0 +1,77 @@
+package opensearch
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"strings"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+)
+
+// handleRegisterStreamQuery handles the HTTP request to register a query for a streaming session.
+func (ds *OpenSearchDatasource) handleRegisterStreamQuery(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+	ds.logger.Info("handleRegisterStreamQuery called", "path", req.Path, "method", req.Method)
+
+	if req.Method != http.MethodPost {
+		return sender.Send(&backend.CallResourceResponse{
+			Status: http.StatusMethodNotAllowed,
+			Body:   []byte("Method not allowed"),
+		})
+	}
+
+	// Extract refId from path, e.g., /_stream_query_register/{refId}
+	// The path in CallResourceRequest does not include the initial /resources prefix.
+	pathParts := strings.Split(strings.TrimPrefix(req.Path, "_stream_query_register/"), "/")
+	if len(pathParts) == 0 || pathParts[0] == "" {
+		ds.logger.Error("handleRegisterStreamQuery: missing refId in path", "path", req.Path)
+		return sender.Send(&backend.CallResourceResponse{
+			Status: http.StatusBadRequest,
+			Body:   []byte("Missing refId in path"),
+		})
+	}
+	refId := pathParts[0]
+
+	// req.Body is already []byte, no need for io.ReadAll
+	bodyBytes := req.Body
+	// if err != nil { // This was for io.ReadAll, not needed now
+	// 	ds.logger.Error("handleRegisterStreamQuery: failed to read request body", "refId", refId, "error", err)
+	// 	return sender.Send(&backend.CallResourceResponse{
+	// 		Status: http.StatusInternalServerError,
+	// 		Body:   []byte("Failed to read request body"),
+	// 	})
+	// }
+	// req.Body is automatically closed by the SDK's resource handling (if it were an io.Closer, which it isn't directly here)
+
+	ds.logger.Debug("handleRegisterStreamQuery: received body", "refId", refId, "body", string(bodyBytes))
+
+	var query Query // Using the Query struct from models.go
+	if err := json.Unmarshal(bodyBytes, &query); err != nil {
+		ds.logger.Error("handleRegisterStreamQuery: failed to unmarshal query", "refId", refId, "error", err, "body", string(bodyBytes))
+		return sender.Send(&backend.CallResourceResponse{
+			Status: http.StatusBadRequest,
+			Body:   []byte("Failed to unmarshal query payload: " + err.Error()),
+		})
+	}
+
+	// Store the query. streamQueries is a sync.Map declared in OpenSearchDatasource struct.
+	ds.streamQueries.Store(refId, query)
+	ds.logger.Info("handleRegisterStreamQuery: successfully stored query for streaming", "refId", refId, "queryIsLuceQueryType", query.luceneQueryType)
+
+	// Send a JSON response
+	responseBody := map[string]string{"message": "Query registered for streaming"}
+	jsonBody, err := json.Marshal(responseBody)
+	if err != nil {
+		ds.logger.Error("handleRegisterStreamQuery: failed to marshal JSON response", "refId", refId, "error", err)
+		return sender.Send(&backend.CallResourceResponse{
+			Status: http.StatusInternalServerError,
+			Body:   []byte("Error creating success response"), // Keep simple error for brevity
+		})
+	}
+
+	return sender.Send(&backend.CallResourceResponse{
+		Status:  http.StatusOK,
+		Body:    jsonBody,
+		Headers: map[string][]string{"Content-Type": {"application/json"}}, // Explicitly set Content-Type for the response
+	})
+}