Support for More Dashboard Variables (#50)

Signed-off-by: Taras Priadka <[email protected]>

pkg/pixie_plugin.go

@@ -88,8 +88,45 @@ func createClient(ctx context.Context, apiKey string, cloudAddr string) (*pxapi.
 type QueryType string
 
 const (
-	RunScript   QueryType = "run-script"
-	GetClusters QueryType = "get-clusters"
+	RunScript     QueryType = "run-script"
+	GetClusters   QueryType = "get-clusters"
+	GetPods       QueryType = "get-pods"
+	GetServices   QueryType = "get-services"
+	GetNamespaces QueryType = "get-namespaces"
+	GetNodes      QueryType = "get-nodes"
+)
+
+const (
+	getPodsScript string = `
+import px
+df = px.DataFrame(table='process_stats', start_time=__time_from)
+df.pod = df.ctx['pod_name']
+df = df[df.pod != '']
+df = df.groupby('pod').agg()
+px.display(df)
+`
+	getServicesScript string = `
+import px
+df = px.DataFrame(table='process_stats', start_time=__time_from)
+df.service = df.ctx['service']
+df = df[df.service != '']
+df = df.groupby('service').agg()
+px.display(df)
+`
+	getNamespacesScript string = `
+import px
+df = px.DataFrame(table='process_stats', start_time=__time_from)
+df.namespace = df.ctx['namespace']
+df = df[df.namespace != '']
+px.display(df.groupby('namespace').agg())
+`
+	getNodesScript string = `
+import px
+df = px.DataFrame(table='process_stats', start_time=__time_from)
+df.node = df.ctx['node_name']
+df = df[df.node != '']
+px.display(df.groupby('node').agg())
+`
 )
 
 type queryBody struct {
@@ -127,16 +164,24 @@ func (td *PixieDatasource) query(ctx context.Context, query backend.DataQuery,
 	}
 
 	if qm.QueryType != GetClusters && len(qm.QueryBody.ClusterID) == 0 {
-		return nil, fmt.Errorf("no clusterID present in the request")
+		return nil, fmt.Errorf("no clusterID present in the request. Please set `pixieCluster` dashboard variable to `Pixie Datasource`->`Clusters`")
 	}
 
 	clusterID := qm.QueryBody.ClusterID
 
 	switch qm.QueryType {
 	case RunScript:
-		return qp.queryScript(ctx, qm.QueryBody, query, clusterID)
+		return qp.queryScript(ctx, qm.QueryBody.PxlScript, query, clusterID)
 	case GetClusters:
-		return qp.queryClusters(ctx, apiToken)
+		return qp.queryClusters(ctx)
+	case GetPods:
+		return qp.queryScript(ctx, getPodsScript, query, clusterID)
+	case GetServices:
+		return qp.queryScript(ctx, getServicesScript, query, clusterID)
+	case GetNamespaces:
+		return qp.queryScript(ctx, getNamespacesScript, query, clusterID)
+	case GetNodes:
+		return qp.queryScript(ctx, getNodesScript, query, clusterID)
 	default:
 		return nil, fmt.Errorf("unknown query type: %v", qm.QueryType)
 	}

pkg/pixie_queries.go

@@ -68,7 +68,7 @@ type PixieQueryProcessor struct {
 // queryScript sends a request to Pixie with pxlScript and returns DataResponse about the current cluster
 func (qp PixieQueryProcessor) queryScript(
 	ctx context.Context,
-	qm queryBody,
+	pxlScript string,
 	query backend.DataQuery,
 	clusterID string,
 ) (*backend.DataResponse, error) {
@@ -83,7 +83,6 @@ func (qp PixieQueryProcessor) queryScript(
 
 	// Create TableMuxer to accept results table.
 	tm := &PixieToGrafanaTableMux{}
-	pxlScript := qm.PxlScript
 	// Update macros in query text.
 	pxlScript = replaceTimeMacroInQueryText(pxlScript, timeFromMacro,
 		query.TimeRange.From)
@@ -128,7 +127,7 @@ func (qp PixieQueryProcessor) queryScript(
 }
 
 // queryClusters sends a request to Pixie, and returns a DataResponse with healthy clusters
-func (qp PixieQueryProcessor) queryClusters(ctx context.Context, apiToken string) (*backend.DataResponse, error) {
+func (qp PixieQueryProcessor) queryClusters(ctx context.Context) (*backend.DataResponse, error) {
 	response := &backend.DataResponse{}
 	viziers, err := qp.client.ListViziers(ctx)
 

src/datasource.ts

@@ -16,9 +16,23 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { DataFrame, DataSourceInstanceSettings, ScopedVars, toDataFrame, VariableModel } from '@grafana/data';
-import { BackendSrv, DataSourceWithBackend, FetchResponse, getBackendSrv, getTemplateSrv } from '@grafana/runtime';
-import { PixieDataSourceOptions, PixieDataQuery, PixieVariableQuery } from './types';
+import {
+  DataFrame,
+  DataSourceInstanceSettings,
+  MetricFindValue,
+  ScopedVars,
+  toDataFrame,
+  VariableModel,
+} from '@grafana/data';
+import { DataSourceWithBackend, getTemplateSrv, FetchResponse, getBackendSrv, BackendSrv } from '@grafana/runtime';
+import {
+  PixieDataSourceOptions,
+  PixieDataQuery,
+  PixieVariableQuery,
+  CLUSTER_VARIABLE_NAME as CLUSTER_VARIABLE_NAME,
+  QueryType,
+  checkExhaustive,
+} from './types';
 import { getColumnsScript } from './column_filtering';
 
 const timeVars = [
@@ -42,6 +56,16 @@ export class DataSource extends DataSourceWithBackend<PixieDataQuery, PixieDataS
     this.backendSrv = getBackendSrv();
   }
 
+  getClusterId(): string {
+    const dashboardVariables: VariableModel[] = getTemplateSrv().getVariables();
+
+    // find cluster variable and convert it to any since the variable value field is not exposed
+    const pixieClusterIdVariable = dashboardVariables.find(
+      (variable) => variable.name === CLUSTER_VARIABLE_NAME
+    ) as any;
+    return pixieClusterIdVariable?.current?.value ?? '';
+  }
+
   applyTemplateVariables(query: PixieDataQuery, scopedVars: ScopedVars) {
     let pxlScript = query.queryBody?.pxlScript ?? '';
 
@@ -58,16 +82,11 @@ export class DataSource extends DataSourceWithBackend<PixieDataQuery, PixieDataS
       );
     }
 
-    const dashboardVariables: VariableModel[] = getTemplateSrv().getVariables();
-
-    // find cluster variable and convert it to any since the variable value field is not exposed
-    const pixieClusterID = dashboardVariables.find((variable) => variable.name === 'pixieCluster') as any;
-
     return {
       ...query,
       queryBody: {
         ...query.queryBody,
-        clusterID: pixieClusterID?.current?.value ?? '',
+        clusterID: this.getClusterId(),
         pxlScript: pxlScript
           ? getTemplateSrv().replace(pxlScript, {
               ...scopedVars,
@@ -81,12 +100,12 @@ export class DataSource extends DataSourceWithBackend<PixieDataQuery, PixieDataS
     const refId = options?.variable?.name ?? 'tempvar';
 
     const interpolatedQuery = {
+      ...query,
       refId,
       datasource: {
         type: this.type,
         uid: this.uid,
       },
-      queryType: query.queryType,
     };
 
     options = {
@@ -124,21 +143,64 @@ export class DataSource extends DataSourceWithBackend<PixieDataQuery, PixieDataS
     return zipped;
   }
 
-  async metricFindQuery(query: PixieVariableQuery, options?: any) {
+  /**
+   * Converts zipped output into dashboard ingestible data.
+   *
+   * @param data zipped data
+   * @param textField field to use for dashboard variable name.
+   * Set to `undefined` if textField should be the same as valueField in the output.
+   *
+   * @param valueField field to use for dashboard variable value
+   */
+  convertData(data: any[], textField: string | undefined, valueField: string): MetricFindValue[] {
+    const output = data.flatMap((entry: any) => {
+      let values: string[] = [entry[valueField] as string];
+      //check if the value is in array form
+      if (values[0].includes(',')) {
+        //expand and clean values
+        values = JSON.parse(values[0]);
+      }
+      return values.map((value) => ({
+        // if textField undefined use value for the text label
+        text: textField ? entry[textField] : value,
+        value: value,
+      }));
+    });
+    return output;
+  }
+
+  async metricFindQuery(query: PixieVariableQuery, options?: any): Promise<MetricFindValue[]> {
     const variableName: string = options.variable.name;
-    //Make sure the query is not empty. Variable query editor will send empty string if user haven't clicked on dropdown menu
+    //Make sure the query is not empty. Variable query editor will send empty query if user haven't clicked on dropdown menu
     query = query || { queryType: 'get-clusters' as const };
 
+    if (query.queryType !== 'get-clusters' && query.queryBody?.clusterID === `\$${CLUSTER_VARIABLE_NAME}`) {
+      const interpolatedClusterId = getTemplateSrv().replace(query.queryBody?.clusterID, options.scopedVars);
+      query = { ...query, queryBody: { clusterID: interpolatedClusterId } };
+    }
     // Fetch variables from the backend
     const response = await this.fetchMetricNames(query, options);
     //Convert the response to a DataFrame
-    const vizierFrame: DataFrame = toDataFrame(response!.data.results[variableName].frames[0]);
-    //Convert DataFrame to an array of objects containing fields same as column names of the DataFrame
-    const clusterData: ClusterMeta[] = this.zipGrafanaDataFrame(vizierFrame);
+    const frame: DataFrame = toDataFrame(response!.data.results[variableName].frames[0]);
 
-    return clusterData.map((entry) => ({
-      text: entry.name,
-      value: entry.id,
-    }));
+    //Convert DataFrame to an array of objects containing fields same as column names of the DataFrame
+    const flatData: ClusterMeta[] = this.zipGrafanaDataFrame(frame);
+
+    switch (query.queryType) {
+      case QueryType.GetClusters:
+        return this.convertData(flatData, 'name', 'id');
+      case QueryType.GetPods:
+        return this.convertData(flatData, undefined, 'pod');
+      case QueryType.GetServices:
+        return this.convertData(flatData, undefined, 'service');
+      case QueryType.GetNamespaces:
+        return this.convertData(flatData, undefined, 'namespace');
+      case QueryType.GetNodes:
+        return this.convertData(flatData, undefined, 'node');
+      case QueryType.RunScript:
+        return Promise.resolve([]);
+      default:
+        checkExhaustive(query.queryType);
+    }
   }
 }

src/module.ts

@@ -21,7 +21,7 @@ import { DataSource } from './datasource';
 import { ConfigEditor } from './config_editor';
 import { QueryEditor } from './query_editor';
 import { PixieDataQuery, PixieDataSourceOptions } from './types';
-import { VariableQueryEditor } from 'variable_query_editor';
+import { VariableQueryEditor } from './variable_query_editor';
 
 export const plugin = new DataSourcePlugin<DataSource, PixieDataQuery, PixieDataSourceOptions>(DataSource)
   .setConfigEditor(ConfigEditor)

src/query_editor.tsx

@@ -24,11 +24,11 @@ import Editor from 'react-simple-code-editor';
 import { highlight, languages } from 'prismjs';
 import 'prismjs/components/prism-python';
 import 'prism-themes/themes/prism-vsc-dark-plus.css';
-import './query_editor.css';
 
+import './styles.css';
 import { DataSource } from './datasource';
-import { scriptOptions, Script } from 'pxl_scripts';
-import { defaultQuery, PixieDataSourceOptions, PixieDataQuery } from './types';
+import { scriptOptions, Script } from './pxl_scripts';
+import { defaultQuery, PixieDataSourceOptions, PixieDataQuery, QueryType } from './types';
 
 type Props = QueryEditorProps<DataSource, PixieDataQuery, PixieDataSourceOptions>;
 
@@ -47,7 +47,7 @@ export class QueryEditor extends PureComponent<Props> {
     const { onChange, query } = this.props;
     onChange({
       ...query,
-      queryType: 'run-script' as const,
+      queryType: QueryType.RunScript,
       queryBody: { pxlScript: event },
     });
   }
@@ -58,7 +58,7 @@ export class QueryEditor extends PureComponent<Props> {
 
       onChange({
         ...query,
-        queryType: 'run-script' as const,
+        queryType: QueryType.RunScript,
         queryBody: { pxlScript: option?.value.script ?? '' },
         queryMeta: {
           isTabular: option.value.isTabular || false,

src/query_editor.css → src/styles.css

@@ -1,8 +1,10 @@
+/* General */
 .m-2 {
     margin-left: 5px;
     margin-right: 5px;
 }
 
+/* Code Editor */
 .code-editor {
     counter-reset: line;
   }

src/types.ts

@@ -17,14 +17,27 @@
  */
 
 import { DataQuery, DataSourceJsonData, SelectableValue } from '@grafana/data';
-import { scriptOptions } from 'pxl_scripts';
+import { scriptOptions } from './pxl_scripts';
 
 // Types of available queries to the backend
-export type QueryType = 'run-script' | 'get-clusters';
+export const enum QueryType {
+  RunScript = 'run-script',
+  GetClusters = 'get-clusters',
+  GetPods = 'get-pods',
+  GetServices = 'get-services',
+  GetNamespaces = 'get-namespaces',
+  GetNodes = 'get-nodes',
+}
+
+// predefined global dashboard variable name for cluster variable
+export const CLUSTER_VARIABLE_NAME = 'pixieCluster';
 
 // Describes variable query to be sent to the backend.
 export interface PixieVariableQuery {
   queryType: QueryType;
+  queryBody?: {
+    clusterID?: string;
+  };
 }
 
 // PixieDataQuery is the interface representing a query in Pixie.
@@ -33,6 +46,7 @@ export interface PixieDataQuery extends DataQuery {
   queryType: QueryType;
   clusterID?: string;
   queryBody?: {
+    clusterID?: string;
     pxlScript?: string;
   };
   // queryMeta is used for UI-Rendering
@@ -44,7 +58,7 @@ export interface PixieDataQuery extends DataQuery {
 }
 
 export const defaultQuery: Partial<PixieDataQuery> = {
-  queryType: 'run-script' as const,
+  queryType: QueryType.RunScript,
   queryBody: {
     pxlScript: scriptOptions[0].value?.script ?? '',
   },
@@ -58,3 +72,7 @@ export interface PixieSecureDataSourceOptions {
   // Address of Pixie cloud.
   cloudAddr?: string;
 }
+
+export function checkExhaustive(val: never): never {
+  throw new Error(`Unexpected value: ${JSON.stringify(val)}`);
+}