Add support for datasource query template variables (#51)

src/components/VariableEditor/VariableEditor.tsx

@@ -0,0 +1,143 @@
+import React, { useCallback, useRef } from 'react';
+import { AsyncSelect, InlineField } from '@grafana/ui';
+import { DataQuery } from '@grafana/schema';
+import { getBackendSrv } from '@grafana/runtime';
+import {
+  CustomVariableSupport,
+  DataQueryRequest,
+  DataQueryResponse,
+  DataSourceApi,
+  QueryEditorProps,
+  SelectableValue,
+} from '@grafana/data';
+import { Observable } from 'rxjs';
+import invariant from 'tiny-invariant';
+
+/**
+ * The _Variables_ DataSource is defined separately from the UQL query
+ * DataSource because we use a distinct "query" (which is actually just an
+ * attributeKey)
+ *
+ * Calling this out explicitly just in case it results in some type mismatches
+ * in the future since the Grafana assumption seems to be that the editor and
+ * variables queries will overlap.
+ */
+type VariableDataSource = DataSourceApi<VariableQuery>;
+interface VariableQuery extends DataQuery {
+  attributeKey: string;
+}
+
+export class VariableEditor extends CustomVariableSupport<VariableDataSource> {
+  constructor(readonly url: string, readonly projectName: string) {
+    super();
+  }
+
+  /**
+   * Variable values fetching function that is called for each dashboard
+   * variable.
+   */
+  query = (request: DataQueryRequest<VariableQuery>): Observable<DataQueryResponse> => {
+    const { url, projectName } = this;
+    const { attributeKey } = request.targets[0];
+    invariant(typeof attributeKey === 'string', 'Invalid attribute key');
+
+    return new Observable((subscriber) => {
+      getBackendSrv()
+        .post(`${url}/projects/${projectName}/telemetry/attributes`, {
+          data: {
+            'attribute-types': ['values'],
+            'telemetry-types': ['spans', 'metrics', 'logs'],
+            'scope-to-attribute-keys': [attributeKey],
+            'oldest-time': request.range.from,
+            'youngest-time': request.range.to,
+          },
+        })
+        .then((res: AttributeRes) => {
+          subscriber.next({
+            data: res.data[attributeKey].map((v) => ({
+              text: v.value,
+            })),
+          });
+        })
+        .catch(() => {
+          // todo: analytics
+          subscriber.next({
+            data: [],
+          });
+        });
+    });
+  };
+
+  /**
+   * Component definition for the UI editor shown to users for creating the
+   * query that will be used to fetch variable options.
+   *
+   * For us, we currently don't have a "query" in the traditional UQL sense,
+   * just an attribute key that we will fetch the values for.
+   */
+  editor = ({ onChange, query, range }: QueryEditorProps<VariableDataSource, VariableQuery>) => {
+    // nb we don't have a way to scope the attribute keys request so we cache
+    // the values for performance
+    const attributeKeysCache = useRef<null | string[]>(null);
+
+    // options fetching fn called on mount and on each change of the select
+    // input
+    const loadOptions = useCallback(
+      async (val: string) => {
+        if (attributeKeysCache.current === null) {
+          const res: AttributeRes = await getBackendSrv().post(
+            `${this.url}/projects/${this.projectName}/telemetry/attributes`,
+            {
+              data: {
+                'attribute-types': ['keys'],
+                'telemetry-types': ['spans', 'metrics', 'logs'],
+                'oldest-time': range?.from,
+                'youngest-time': range?.to,
+              },
+            }
+          );
+          attributeKeysCache.current = Object.keys(res.data);
+        }
+
+        const options: Array<SelectableValue<string>> = attributeKeysCache.current
+          .filter((key) => key.includes(val))
+          .map((key) => ({
+            label: key,
+            value: key,
+          }));
+
+        return options;
+      },
+      [range]
+    );
+
+    return (
+      <div className="gf-form">
+        <InlineField
+          label="Attribute key"
+          tooltip="Cloud Observability uses this key to populate the selectable values for the variable when viewing the dashboard. Choose from any attributes currently on your logs, metics, or traces."
+        >
+          <AsyncSelect
+            defaultOptions
+            cacheOptions
+            defaultValue={query ? { label: query.attributeKey, value: query.attributeKey } : undefined}
+            loadOptions={loadOptions}
+            onChange={(v) => {
+              if (v.value) {
+                onChange({ refId: v.value, attributeKey: v.value });
+              }
+            }}
+          />
+        </InlineField>
+      </div>
+    );
+  };
+}
+
+/** attributes endpoint response shape */
+type AttributeRes = {
+  data: Record<
+    string,
+    Array<{ type: 'string' | 'int64'; value: string; telemetry_type: 'SPANS' | 'METRICS' | 'LOGS' }>
+  >;
+};

src/datasource.ts

@@ -8,6 +8,7 @@ import {
 import { config, getBackendSrv, getTemplateSrv } from '@grafana/runtime';
 import { preprocessData } from './preprocessors';
 import { LightstepDataSourceOptions, LightstepQuery } from './types';
+import { VariableEditor } from 'components/VariableEditor/VariableEditor';
 
 /**
  * THE DATASOURCE
@@ -26,6 +27,8 @@ export class DataSource extends DataSourceApi<LightstepQuery, LightstepDataSourc
     this.projectName = instanceSettings.jsonData.projectName || '';
     this.orgName = instanceSettings.jsonData.orgName || '';
     this.url = instanceSettings.url || '';
+
+    this.variables = new VariableEditor(this.url, this.defaultProjectName());
   }
 
   /**
@@ -61,12 +64,13 @@ export class DataSource extends DataSourceApi<LightstepQuery, LightstepDataSourc
         const res = await getBackendSrv().post(`${this.url}/projects/${query.projectName}/telemetry/query_timeseries`, {
           data: {
             attributes: {
-              query: getTemplateSrv().replace(query.text, request.scopedVars),
+              query: query.text,
               'input-language': query.language,
               'oldest-time': request.range.from,
               'youngest-time': request.range.to,
               // query_timeseries minimum supported output-period is 1 second
               'output-period': Math.max(1, rangeUtil.intervalToSeconds(request.interval)),
+              'template-variables': createRequestVariables(),
             },
             analytics: {
               anonymized_user: hashedEmail,
@@ -142,6 +146,35 @@ export class DataSource extends DataSourceApi<LightstepQuery, LightstepDataSourc
   }
 }
 
+/**
+ * Translates Grafana dashboard variables into a set of LS API template
+ * variables
+ */
+function createRequestVariables() {
+  return getTemplateSrv()
+    .getVariables()
+    .map((v) => {
+      if (v.type === 'query' || v.type === 'textbox' || v.type === 'custom' || v.type === 'constant') {
+        // normalize different variables values formats into request standard
+        // array of strings
+        const { value } = v.current;
+        let values = Array.isArray(value) ? value : [value];
+        if (values.length === 1 && values[0] === '$__all') {
+          values = [];
+        }
+
+        return {
+          name: v.name,
+          values,
+        };
+      }
+
+      // SKIP adhoc, datasource, system, and interval template variables
+      return false;
+    })
+    .filter(Boolean);
+}
+
 /**
  * Create an *anonymous* unique id from user email
  * @see https://developer.mozilla.org/en-US/docs/Web/API/SubtleCrypto/digest#converting_a_digest_to_a_hex_string