Pyroscope: Send start/end with profile types query (#77523)

pkg/tsdb/grafana-pyroscope-datasource/instance.go

@@ -29,7 +29,7 @@ var (
 )
 
 type ProfilingClient interface {
-	ProfileTypes(context.Context) ([]*ProfileType, error)
+	ProfileTypes(ctx context.Context, start int64, end int64) ([]*ProfileType, error)
 	LabelNames(ctx context.Context, labelSelector string, start int64, end int64) ([]string, error)
 	LabelValues(ctx context.Context, label string, labelSelector string, start int64, end int64) ([]string, error)
 	GetSeries(ctx context.Context, profileTypeID string, labelSelector string, start int64, end int64, groupBy []string, step float64) (*SeriesResponse, error)
@@ -86,7 +86,30 @@ func (d *PyroscopeDatasource) CallResource(ctx context.Context, req *backend.Cal
 
 func (d *PyroscopeDatasource) profileTypes(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
 	ctxLogger := logger.FromContext(ctx)
-	types, err := d.client.ProfileTypes(ctx)
+
+	u, err := url.Parse(req.URL)
+	if err != nil {
+		ctxLogger.Error("Failed to parse URL", "error", err, "function", logEntrypoint())
+		return err
+	}
+	query := u.Query()
+
+	var start, end int64
+	if query.Has("start") && query.Has("end") {
+		start, err = strconv.ParseInt(query.Get("start"), 10, 64)
+		if err != nil {
+			ctxLogger.Error("Failed to parse start as int", "error", err, "function", logEntrypoint())
+			return err
+		}
+
+		end, err = strconv.ParseInt(query.Get("end"), 10, 64)
+		if err != nil {
+			ctxLogger.Error("Failed to parse end as int", "error", err, "function", logEntrypoint())
+			return err
+		}
+	}
+
+	types, err := d.client.ProfileTypes(ctx, start, end)
 	if err != nil {
 		ctxLogger.Error("Received error from client", "error", err, "function", logEntrypoint())
 		return err
@@ -199,7 +222,7 @@ func (d *PyroscopeDatasource) labelValues(ctx context.Context, req *backend.Call
 // contains Frames ([]*Frame).
 func (d *PyroscopeDatasource) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
 	ctxLogger := logger.FromContext(ctx)
-	ctxLogger.Debug("Processing queries", "queryLenght", len(req.Queries), "function", logEntrypoint())
+	ctxLogger.Debug("Processing queries", "queryLength", len(req.Queries), "function", logEntrypoint())
 
 	// create response struct
 	response := backend.NewQueryDataResponse()
@@ -228,7 +251,11 @@ func (d *PyroscopeDatasource) CheckHealth(ctx context.Context, _ *backend.CheckH
 	status := backend.HealthStatusOk
 	message := "Data source is working"
 
-	if _, err := d.client.ProfileTypes(ctx); err != nil {
+	// Since this is a health check mechanism and we only care about whether the
+	// request succeeded or failed, we set the window to be small.
+	start := time.Now().Add(-5 * time.Minute).UnixMilli()
+	end := time.Now().UnixMilli()
+	if _, err := d.client.ProfileTypes(ctx, start, end); err != nil {
 		status = backend.HealthStatusError
 		message = err.Error()
 	}

pkg/tsdb/grafana-pyroscope-datasource/pyroscopeClient.go

@@ -70,10 +70,13 @@ func NewPyroscopeClient(httpClient *http.Client, url string) *PyroscopeClient {
 	}
 }
 
-func (c *PyroscopeClient) ProfileTypes(ctx context.Context) ([]*ProfileType, error) {
+func (c *PyroscopeClient) ProfileTypes(ctx context.Context, start int64, end int64) ([]*ProfileType, error) {
 	ctx, span := tracing.DefaultTracer().Start(ctx, "datasource.pyroscope.ProfileTypes")
 	defer span.End()
-	res, err := c.connectClient.ProfileTypes(ctx, connect.NewRequest(&querierv1.ProfileTypesRequest{}))
+	res, err := c.connectClient.ProfileTypes(ctx, connect.NewRequest(&querierv1.ProfileTypesRequest{
+		Start: start,
+		End:   end,
+	}))
 	if err != nil {
 		logger.Error("Received error from client", "error", err, "function", logEntrypoint())
 		span.RecordError(err)

pkg/tsdb/grafana-pyroscope-datasource/query_test.go

@@ -275,7 +275,7 @@ type FakeClient struct {
 	Args []any
 }
 
-func (f *FakeClient) ProfileTypes(ctx context.Context) ([]*ProfileType, error) {
+func (f *FakeClient) ProfileTypes(ctx context.Context, start int64, end int64) ([]*ProfileType, error) {
 	return []*ProfileType{
 		{
 			ID:    "type:1",

public/app/core/components/TraceToProfiles/TraceToProfilesSettings.tsx

@@ -55,7 +55,7 @@ export function TraceToProfilesSettings({ options, onOptionsChange }: Props) {
       supportedDataSourceTypes.includes(dataSource.type) &&
       dataSource.uid === options.jsonData.tracesToProfiles?.datasourceUid
     ) {
-      dataSource.getProfileTypes().then((profileTypes) => {
+      dataSource.getAllProfileTypes().then((profileTypes) => {
         setProfileTypes(profileTypes);
       });
     } else {

public/app/plugins/datasource/grafana-pyroscope-datasource/QueryEditor/ProfileTypesCascader.tsx

@@ -1,5 +1,6 @@
 import React, { useEffect, useMemo, useState } from 'react';
 
+import { TimeRange } from '@grafana/data';
 import { Cascader, CascaderOption } from '@grafana/ui';
 
 import { PyroscopeDataSource } from '../datasource';
@@ -70,16 +71,22 @@ function useCascaderOptions(profileTypes?: ProfileTypeMessage[]): CascaderOption
  * This is exported and not used directly in the ProfileTypesCascader component because in some case we need to know
  * the profileTypes before rendering the cascader.
  * @param datasource
+ * @param range Time range for the profile types query.
  */
-export function useProfileTypes(datasource: PyroscopeDataSource) {
+export function useProfileTypes(datasource: PyroscopeDataSource, range?: TimeRange) {
   const [profileTypes, setProfileTypes] = useState<ProfileTypeMessage[]>();
 
+  const impreciseRange = {
+    to: Math.ceil((range?.to.valueOf() || 0) / 60000) * 60000,
+    from: Math.floor((range?.from.valueOf() || 0) / 60000) * 60000,
+  };
+
   useEffect(() => {
     (async () => {
-      const profileTypes = await datasource.getProfileTypes();
+      const profileTypes = await datasource.getProfileTypes(impreciseRange.from.valueOf(), impreciseRange.to.valueOf());
       setProfileTypes(profileTypes);
     })();
-  }, [datasource]);
+  }, [datasource, impreciseRange.from, impreciseRange.to]);
 
   return profileTypes;
 }

public/app/plugins/datasource/grafana-pyroscope-datasource/QueryEditor/QueryEditor.tsx

@@ -27,7 +27,7 @@ export function QueryEditor(props: Props) {
     onRunQuery();
   }
 
-  const profileTypes = useProfileTypes(datasource);
+  const profileTypes = useProfileTypes(datasource, range);
   const { labels, getLabelValues, onLabelSelectorChange } = useLabels(range, datasource, query, onChange);
   useNormalizeQuery(query, profileTypes, onChange, app);
 

public/app/plugins/datasource/grafana-pyroscope-datasource/VariableQueryEditor.tsx

@@ -1,6 +1,6 @@
 import React, { useEffect, useState } from 'react';
 
-import { QueryEditorProps, SelectableValue } from '@grafana/data';
+import { QueryEditorProps, SelectableValue, TimeRange } from '@grafana/data';
 import { InlineField, InlineFieldRow, LoadingPlaceholder, Select } from '@grafana/ui';
 
 import { ProfileTypesCascader, useProfileTypes } from './QueryEditor/ProfileTypesCascader';
@@ -65,6 +65,7 @@ export function VariableQueryEditor(props: QueryEditorProps<PyroscopeDataSource,
               props.onChange({ ...props.query, profileTypeId: val });
             }
           }}
+          range={props.range}
         />
       )}
 
@@ -131,8 +132,9 @@ function ProfileTypeRow(props: {
   datasource: PyroscopeDataSource;
   onChange: (val: string) => void;
   initialValue?: string;
+  range?: TimeRange;
 }) {
-  const profileTypes = useProfileTypes(props.datasource);
+  const profileTypes = useProfileTypes(props.datasource, props.range);
   return (
     <InlineFieldRow>
       <InlineField

public/app/plugins/datasource/grafana-pyroscope-datasource/VariableSupport.ts

@@ -9,7 +9,7 @@ import { PyroscopeDataSource } from './datasource';
 import { ProfileTypeMessage, VariableQuery } from './types';
 
 export interface DataAPI {
-  getProfileTypes(): Promise<ProfileTypeMessage[]>;
+  getProfileTypes(start: number, end: number): Promise<ProfileTypeMessage[]>;
   getLabelNames(query: string, start: number, end: number): Promise<string[]>;
   getLabelValues(query: string, label: string, start: number, end: number): Promise<string[]>;
 }
@@ -26,7 +26,9 @@ export class VariableSupport extends CustomVariableSupport<PyroscopeDataSource>
 
   query(request: DataQueryRequest<VariableQuery>): Observable<DataQueryResponse> {
     if (request.targets[0].type === 'profileType') {
-      return from(this.dataAPI.getProfileTypes()).pipe(
+      return from(
+        this.dataAPI.getProfileTypes(this.timeSrv.timeRange().from.valueOf(), this.timeSrv.timeRange().to.valueOf())
+      ).pipe(
         map((values) => {
           return { data: values.map<MetricFindValue>((v) => ({ text: v.label, value: v.id })) };
         })

public/app/plugins/datasource/grafana-pyroscope-datasource/datasource.ts

@@ -48,7 +48,14 @@ export class PyroscopeDataSource extends DataSourceWithBackend<Query, PyroscopeD
     });
   }
 
-  async getProfileTypes(): Promise<ProfileTypeMessage[]> {
+  async getProfileTypes(start: number, end: number): Promise<ProfileTypeMessage[]> {
+    return await this.getResource('profileTypes', {
+      start,
+      end,
+    });
+  }
+
+  async getAllProfileTypes(): Promise<ProfileTypeMessage[]> {
     return await this.getResource('profileTypes');
   }