PostgreSQL: Add variable query editor support

- Introduced a new feature toggle for the PostgreSQL variable query editor `postgresVariableQueryEditor`.
2025-12-08 23:18:00 +01:00
parent 15c93100ab
commit d5215a5be2
16 changed files with 400 additions and 41 deletions
@@ -1161,6 +1161,10 @@ export interface FeatureToggles {
  */
  jaegerEnableGrpcEndpoint?: boolean;
  /**
+  * Enable the new variable query editor for the PostgreSQL data source
+  */
+  postgresVariableQueryEditor?: boolean;
+  /**
  * Load plugins on store service startup instead of wire provider, and call RegisterFixedRoles after all plugins are loaded
  * @default false
  */
@@ -21,6 +21,7 @@
    "@grafana/i18n": "12.4.0-pre",
    "@grafana/plugin-ui": "^0.11.0",
    "@grafana/runtime": "12.4.0-pre",
+    "@grafana/schema": "12.4.0-pre",
    "@grafana/ui": "12.4.0-pre",
    "@react-awesome-query-builder/ui": "6.6.15",
    "immutable": "5.1.4",
@@ -15,7 +15,7 @@ import { RawEditor } from './query-editor-raw/RawEditor';
 import { VisualEditor } from './visual-query-builder/VisualEditor';

 export interface SqlQueryEditorProps extends QueryEditorProps<SqlDatasource, SQLQuery, SQLOptions> {
-  queryHeaderProps?: Pick<QueryHeaderProps, 'dialect'>;
+  queryHeaderProps?: Pick<QueryHeaderProps, 'dialect' | 'hideRunButton' | 'hideFormatSelector'>;
 }

 export default function SqlQueryEditor({
@@ -99,6 +99,8 @@ export default function SqlQueryEditor({
        query={queryWithDefaults}
        isQueryRunnable={isQueryRunnable}
        dialect={dialect}
+        hideRunButton={queryHeaderProps?.hideRunButton}
+        hideFormatSelector={queryHeaderProps?.hideFormatSelector}
      />

      <Space v={0.5} />
@@ -25,6 +25,8 @@ export interface QueryHeaderProps {
  preconfiguredDataset: string;
  query: QueryWithDefaults;
  queryRowFilter: QueryRowFilter;
+  hideRunButton?: boolean;
+  hideFormatSelector?: boolean;
 }

 export function QueryHeader({
@@ -37,6 +39,8 @@ export function QueryHeader({
  preconfiguredDataset,
  query,
  queryRowFilter,
+  hideRunButton,
+  hideFormatSelector,
 }: QueryHeaderProps) {
  const { editorMode } = query;
  const [_, copyToClipboard] = useCopyToClipboard();
@@ -123,14 +127,16 @@ export function QueryHeader({
  return (
    <>
      <EditorHeader>
-        <InlineSelect
-          label={t('grafana-sql.components.query-header.label-format', 'Format')}
-          value={query.format}
-          placeholder={t('grafana-sql.components.query-header.placeholder-select-format', 'Select format')}
-          menuShouldPortal
-          onChange={onFormatChange}
-          options={QUERY_FORMAT_OPTIONS}
-        />
+        {!hideFormatSelector && (
+          <InlineSelect
+            label={t('grafana-sql.components.query-header.label-format', 'Format')}
+            value={query.format}
+            placeholder={t('grafana-sql.components.query-header.placeholder-select-format', 'Select format')}
+            menuShouldPortal
+            onChange={onFormatChange}
+            options={QUERY_FORMAT_OPTIONS}
+          />
+        )}

        {editorMode === EditorMode.Builder && (
          <>
@@ -222,26 +228,27 @@ export function QueryHeader({

        <FlexItem grow={1} />

-        {isQueryRunnable ? (
-          <Button icon="play" variant="primary" size="sm" onClick={() => onRunQuery()}>
-            <Trans i18nKey="grafana-sql.components.query-header.run-query">Run query</Trans>
-          </Button>
-        ) : (
-          <Tooltip
-            theme="error"
-            content={
-              <Trans i18nKey="grafana-sql.components.query-header.content-invalid-query">
-                Your query is invalid. Check below for details. <br />
-                However, you can still run this query.
-              </Trans>
-            }
-            placement="top"
-          >
-            <Button icon="exclamation-triangle" variant="secondary" size="sm" onClick={() => onRunQuery()}>
+        {!hideRunButton &&
+          (isQueryRunnable ? (
+            <Button icon="play" variant="primary" size="sm" onClick={() => onRunQuery()}>
              <Trans i18nKey="grafana-sql.components.query-header.run-query">Run query</Trans>
            </Button>
-          </Tooltip>
-        )}
+          ) : (
+            <Tooltip
+              theme="error"
+              content={
+                <Trans i18nKey="grafana-sql.components.query-header.content-invalid-query">
+                  Your query is invalid. Check below for details. <br />
+                  However, you can still run this query.
+                </Trans>
+              }
+              placement="top"
+            >
+              <Button icon="exclamation-triangle" variant="secondary" size="sm" onClick={() => onRunQuery()}>
+                <Trans i18nKey="grafana-sql.components.query-header.run-query">Run query</Trans>
+              </Button>
+            </Tooltip>
+          ))}

        <RadioButtonGroup options={editorModes} size="sm" value={editorMode} onChange={onEditorModeChange} />

@@ -2,20 +2,19 @@ import { lastValueFrom, Observable, throwError } from 'rxjs';
 import { map } from 'rxjs/operators';

 import {
-  getDefaultTimeRange,
+  CoreApp,
  DataFrame,
  DataFrameView,
-  DataQuery,
  DataQueryRequest,
  DataQueryResponse,
  DataSourceInstanceSettings,
-  MetricFindValue,
-  ScopedVars,
-  CoreApp,
+  getDefaultTimeRange,
  getSearchFilterScopedVar,
  LegacyMetricFindQueryOptions,
-  VariableWithMultiSupport,
+  MetricFindValue,
+  ScopedVars,
  TimeRange,
+  VariableWithMultiSupport,
 } from '@grafana/data';
 import { EditorMode } from '@grafana/plugin-ui';
 import {
@@ -24,15 +23,16 @@ import {
  FetchResponse,
  getBackendSrv,
  getTemplateSrv,
-  toDataQueryResponse,
-  TemplateSrv,
  reportInteraction,
+  TemplateSrv,
+  toDataQueryResponse,
 } from '@grafana/runtime';
+import { DataQuery } from '@grafana/schema';

 import { ResponseParser } from '../ResponseParser';
 import { SqlQueryEditorLazy } from '../components/QueryEditorLazy';
 import { MACRO_NAMES } from '../constants';
-import { DB, SQLQuery, SQLOptions, SqlQueryModel, QueryFormat } from '../types';
+import { DB, QueryFormat, SQLOptions, SQLQuery, SqlQueryModel } from '../types';
 import migrateAnnotation from '../utils/migration';

 export abstract class SqlDatasource extends DataSourceWithBackend<SQLQuery, SQLOptions> {
@@ -182,7 +182,7 @@ export abstract class SqlDatasource extends DataSourceWithBackend<SQLQuery, SQLO
    return;
  }

-  async metricFindQuery(query: string, options?: LegacyMetricFindQueryOptions): Promise<MetricFindValue[]> {
+  async metricFindQuery(query: SQLQuery | string, options?: LegacyMetricFindQueryOptions): Promise<MetricFindValue[]> {
    const range = options?.range;
    if (range == null) {
      // i cannot create a scenario where this happens, we handle it just to be sure.
@@ -194,12 +194,17 @@ export abstract class SqlDatasource extends DataSourceWithBackend<SQLQuery, SQLO
      refId = options.variable.name;
    }

+    const queryString = typeof query === 'string' ? query : query.rawSql;
+    if (!queryString) {
+      return [];
+    }
+
    const scopedVars = {
      ...options?.scopedVars,
-      ...getSearchFilterScopedVar({ query, wildcardChar: '%', options }),
+      ...getSearchFilterScopedVar({ query: queryString, wildcardChar: '%', options }),
    };

-    const rawSql = this.templateSrv.replace(query, scopedVars, this.interpolateVariable);
+    const rawSql = this.templateSrv.replace(queryString, scopedVars, this.interpolateVariable);

    const interpolatedQuery: SQLQuery = {
      refId: refId,
@@ -1913,6 +1913,13 @@ var (
 			Stage:       FeatureStageExperimental,
 			Owner:       grafanaOSSBigTent,
 		},
+		{
+			Name:         "postgresVariableQueryEditor",
+			Description:  "Enable the new variable query editor for the PostgreSQL data source",
+			Stage:        FeatureStageExperimental,
+			Owner:        grafanaOSSBigTent,
+			FrontendOnly: true,
+		},
 		{
 			Name:         "pluginStoreServiceLoading",
 			Description:  "Load plugins on store service startup instead of wire provider, and call RegisterFixedRoles after all plugins are loaded",
@@ -260,6 +260,7 @@ newVizSuggestions,preview,@grafana/dataviz-squad,false,false,true
 externalVizSuggestions,experimental,@grafana/dataviz-squad,false,false,true
 preventPanelChromeOverflow,preview,@grafana/grafana-frontend-platform,false,false,true
 jaegerEnableGrpcEndpoint,experimental,@grafana/oss-big-tent,false,false,false
+postgresVariableQueryEditor,experimental,@grafana/oss-big-tent,false,false,true
 pluginStoreServiceLoading,experimental,@grafana/plugins-platform-backend,false,false,false
 newPanelPadding,preview,@grafana/dashboards-squad,false,false,true
 onlyStoreActionSets,GA,@grafana/identity-access-team,false,false,false
@@ -2692,6 +2692,19 @@
        "expression": "false"
      }
    },
+    {
+      "metadata": {
+        "name": "postgresVariableQueryEditor",
+        "resourceVersion": "1765231394616",
+        "creationTimestamp": "2025-12-08T22:03:14Z"
+      },
+      "spec": {
+        "description": "Enable the new variable query editor for the PostgreSQL data source",
+        "stage": "experimental",
+        "codeowner": "@grafana/oss-big-tent",
+        "frontend": true
+      }
+    },
    {
      "metadata": {
        "name": "preferLibraryPanelTitle",
@@ -0,0 +1,20 @@
+import { QueryEditorProps } from '@grafana/data';
+import { QueryHeaderProps, SQLOptions, SQLQuery, SqlQueryEditorLazy } from '@grafana/sql';
+
+import { PostgresDatasource } from './datasource';
+import { migrateVariableQuery } from './migrations';
+
+const queryHeaderProps: Pick<QueryHeaderProps, 'dialect' | 'hideRunButton' | 'hideFormatSelector'> = {
+  dialect: 'postgres',
+  hideRunButton: true,
+  hideFormatSelector: true,
+};
+
+export function VariableQueryEditor(props: QueryEditorProps<PostgresDatasource, SQLQuery, SQLOptions>) {
+  const newProps = {
+    ...props,
+    query: migrateVariableQuery(props.query),
+    queryHeaderProps,
+  };
+  return <SqlQueryEditorLazy {...newProps} />;
+}
@@ -2,7 +2,7 @@ import { v4 as uuidv4 } from 'uuid';

 import { DataSourceInstanceSettings, ScopedVars, VariableWithMultiSupport } from '@grafana/data';
 import { LanguageDefinition } from '@grafana/plugin-ui';
-import { TemplateSrv } from '@grafana/runtime';
+import { config, TemplateSrv } from '@grafana/runtime';
 import {
  COMMON_FNS,
  DB,
@@ -16,15 +16,23 @@ import {

 import { PostgresQueryModel } from './PostgresQueryModel';
 import { getSchema, getTimescaleDBVersion, getVersion, showTables } from './postgresMetaQuery';
+import { transformMetricFindResponse } from './responseParser';
 import { fetchColumns, fetchTables, getSqlCompletionProvider } from './sqlCompletionProvider';
 import { getFieldConfig, toRawSql } from './sqlUtil';
 import { PostgresOptions } from './types';
+import { SQLVariableSupport } from './variables';

 export class PostgresDatasource extends SqlDatasource {
  sqlLanguageDefinition: LanguageDefinition | undefined = undefined;

  constructor(instanceSettings: DataSourceInstanceSettings<PostgresOptions>) {
    super(instanceSettings);
+    if (config.featureToggles.postgresVariableQueryEditor) {
+      this.variables = new SQLVariableSupport(this);
+      this.responseParser = {
+        transformMetricFindResponse: transformMetricFindResponse,
+      };
+    }
  }

  getQueryModel(target?: SQLQuery, templateSrv?: TemplateSrv, scopedVars?: ScopedVars): PostgresQueryModel {
@@ -0,0 +1,20 @@
+import { applyQueryDefaults, type SQLQuery } from '@grafana/sql';
+
+import type { VariableQuery } from './types';
+
+export function migrateVariableQuery(rawQuery: string | SQLQuery): VariableQuery {
+  if (typeof rawQuery !== 'string') {
+    return {
+      ...rawQuery,
+      query: rawQuery.rawSql || '',
+    };
+  }
+
+  return {
+    ...applyQueryDefaults({
+      refId: 'SQLVariableQueryEditor-VariableQuery',
+      rawSql: rawQuery,
+    }),
+    query: rawQuery,
+  };
+}
@@ -0,0 +1,177 @@
+import { FieldType, DataFrame } from '@grafana/data';
+
+import { transformMetricFindResponse } from './responseParser';
+
+describe('transformMetricFindResponse function', () => {
+  it('should handle big arrays', () => {
+    const stringValues = new Array(150_000).fill('a');
+    const numberValues = new Array(150_000).fill(1);
+
+    const frame: DataFrame = {
+      fields: [
+        { name: 'name', type: FieldType.string, config: {}, values: stringValues },
+        { name: 'value', type: FieldType.number, config: {}, values: numberValues },
+      ],
+      length: stringValues.length,
+    };
+
+    const result = transformMetricFindResponse(frame);
+
+    // With the new logic, first field is text/value and additional fields are properties
+    // After deduplication, we get 1 unique item since all 'a' values deduplicate
+    // Fields named 'value' are not added to properties to avoid conflicts
+    expect(result).toHaveLength(1);
+    expect(result[0]).toEqual({
+      text: 'a',
+      value: 'a',
+    });
+  });
+
+  it('should use first field as text and value with additional fields as properties', () => {
+    const frame: DataFrame = {
+      fields: [
+        { name: 'id', type: FieldType.string, config: {}, values: ['user1', 'user2', 'user3'] },
+        {
+          name: 'email',
+          type: FieldType.string,
+          config: {},
+          values: ['user1@test.com', 'user2@test.com', 'user3@test.com'],
+        },
+        { name: 'role', type: FieldType.string, config: {}, values: ['admin', 'user', 'guest'] },
+      ],
+      length: 3,
+    };
+
+    const result = transformMetricFindResponse(frame);
+
+    expect(result).toHaveLength(3);
+    expect(result[0]).toEqual({
+      text: 'user1',
+      value: 'user1',
+      properties: {
+        email: 'user1@test.com',
+        role: 'admin',
+      },
+    });
+    expect(result[1]).toEqual({
+      text: 'user2',
+      value: 'user2',
+      properties: {
+        email: 'user2@test.com',
+        role: 'user',
+      },
+    });
+    expect(result[2]).toEqual({
+      text: 'user3',
+      value: 'user3',
+      properties: {
+        email: 'user3@test.com',
+        role: 'guest',
+      },
+    });
+  });
+
+  it('should handle single field without properties', () => {
+    const frame: DataFrame = {
+      fields: [{ name: 'name', type: FieldType.string, config: {}, values: ['value1', 'value2'] }],
+      length: 2,
+    };
+
+    const result = transformMetricFindResponse(frame);
+
+    expect(result).toHaveLength(2);
+    expect(result[0]).toEqual({
+      text: 'value1',
+      value: 'value1',
+    });
+    expect(result[1]).toEqual({
+      text: 'value2',
+      value: 'value2',
+    });
+  });
+
+  it('should still handle __text and __value fields', () => {
+    const frame: DataFrame = {
+      fields: [
+        { name: '__text', type: FieldType.string, config: {}, values: ['Display 1', 'Display 2'] },
+        { name: '__value', type: FieldType.string, config: {}, values: ['val1', 'val2'] },
+      ],
+      length: 2,
+    };
+
+    const result = transformMetricFindResponse(frame);
+
+    expect(result).toHaveLength(2);
+    expect(result[0]).toEqual({
+      text: 'Display 1',
+      value: 'val1',
+    });
+    expect(result[1]).toEqual({
+      text: 'Display 2',
+      value: 'val2',
+    });
+  });
+
+  it('should not add fields named "text" or "value" to properties', () => {
+    const frame: DataFrame = {
+      fields: [
+        { name: 'id', type: FieldType.string, config: {}, values: ['item1', 'item2'] },
+        { name: 'text', type: FieldType.string, config: {}, values: ['Text 1', 'Text 2'] },
+        { name: 'value', type: FieldType.string, config: {}, values: ['Value 1', 'Value 2'] },
+        { name: 'description', type: FieldType.string, config: {}, values: ['Desc 1', 'Desc 2'] },
+      ],
+      length: 2,
+    };
+
+    const result = transformMetricFindResponse(frame);
+
+    expect(result).toHaveLength(2);
+    // Fields named 'text' and 'value' should not be in properties
+    expect(result[0]).toEqual({
+      text: 'item1',
+      value: 'item1',
+      properties: {
+        description: 'Desc 1',
+      },
+    });
+    expect(result[1]).toEqual({
+      text: 'item2',
+      value: 'item2',
+      properties: {
+        description: 'Desc 2',
+      },
+    });
+  });
+
+  it('should add additional fields as properties when __text and __value are present', () => {
+    const frame: DataFrame = {
+      fields: [
+        { name: '__text', type: FieldType.string, config: {}, values: ['Display 1', 'Display 2'] },
+        { name: '__value', type: FieldType.string, config: {}, values: ['val1', 'val2'] },
+        { name: 'category', type: FieldType.string, config: {}, values: ['cat1', 'cat2'] },
+        { name: 'priority', type: FieldType.number, config: {}, values: [1, 2] },
+      ],
+      length: 2,
+    };
+
+    const result = transformMetricFindResponse(frame);
+
+    expect(result).toHaveLength(2);
+    expect(result[0]).toEqual({
+      text: 'Display 1',
+      value: 'val1',
+      properties: {
+        category: 'cat1',
+        priority: '1',
+      },
+    });
+    expect(result[1]).toEqual({
+      text: 'Display 2',
+      value: 'val2',
+      properties: {
+        category: 'cat2',
+        priority: '2',
+      },
+    });
+  });
+});
@@ -0,0 +1,62 @@
+import { uniqBy } from 'lodash';
+
+import { DataFrame, Field, MetricFindValue } from '@grafana/data';
+
+export function transformMetricFindResponse(frame: DataFrame): MetricFindValue[] {
+  const textField = frame.fields.find((f) => f.name === '__text');
+  const valueField = frame.fields.find((f) => f.name === '__value');
+
+  let values: MetricFindValue[];
+
+  if (textField && valueField) {
+    const additionalFields = frame.fields.filter((f) => f.name !== '__text' && f.name !== '__value');
+    values = buildMetricFindValues(textField, valueField, additionalFields);
+  } else if (frame.fields.length > 0) {
+    // Support multiple fields by first field as text/value and additional fields as properties
+    const firstField = frame.fields[0];
+    const additionalFields = frame.fields.slice(1);
+    values = buildMetricFindValues(firstField, firstField, additionalFields);
+  } else {
+    values = [];
+  }
+
+  return uniqBy(values, 'text');
+}
+
+function buildMetricFindValues(textField: Field, valueField: Field, additionalFields: Field[]): MetricFindValue[] {
+  const values: MetricFindValue[] = [];
+
+  for (let i = 0; i < textField.values.length; i++) {
+    const item: MetricFindValue = {
+      text: '' + textField.values[i],
+      value: '' + valueField.values[i],
+    };
+
+    const properties = buildProperties(additionalFields, i);
+    if (properties) {
+      item.properties = properties;
+    }
+
+    values.push(item);
+  }
+
+  return values;
+}
+
+function buildProperties(fields: Field[], rowIndex: number): Record<string, string> | undefined {
+  if (fields.length === 0) {
+    return undefined;
+  }
+
+  const properties: Record<string, string> = {};
+
+  for (const field of fields) {
+    // Skip fields named 'text' or 'value' to avoid conflicts with top-level fields
+    if (field.name !== 'text' && field.name !== 'value') {
+      properties[field.name] = '' + field.values[rowIndex];
+    }
+  }
+
+  // Only return properties object if there are actual properties
+  return Object.keys(properties).length > 0 ? properties : undefined;
+}
@@ -1,4 +1,4 @@
-import { SQLOptions } from '@grafana/sql';
+import { SQLOptions, SQLQuery } from '@grafana/sql';

 export enum PostgresTLSModes {
  disable = 'disable',
@@ -25,3 +25,7 @@ export interface PostgresOptions extends SQLOptions {
 export interface SecureJsonData {
  password?: string;
 }
+
+export interface VariableQuery extends SQLQuery {
+  query: string;
+}
@@ -0,0 +1,27 @@
+import { from, map, Observable } from 'rxjs';
+
+import { CustomVariableSupport, DataQueryRequest, MetricFindValue } from '@grafana/data';
+import { applyQueryDefaults, SQLQuery } from '@grafana/sql';
+
+import { VariableQueryEditor } from './VariableQueryEditor';
+import { PostgresDatasource } from './datasource';
+import { migrateVariableQuery } from './migrations';
+
+export class SQLVariableSupport extends CustomVariableSupport<PostgresDatasource, SQLQuery> {
+  constructor(private readonly datasource: PostgresDatasource) {
+    super();
+  }
+
+  editor = VariableQueryEditor;
+
+  query(request: DataQueryRequest<SQLQuery>): Observable<{ data: MetricFindValue[] }> {
+    const queryObj = migrateVariableQuery(request.targets[0]);
+    const result = this.datasource.metricFindQuery(queryObj, { scopedVars: request.scopedVars, range: request.range });
+
+    return from(result).pipe(map((data) => ({ data })));
+  }
+
+  getDefaultQuery(): Partial<SQLQuery> {
+    return applyQueryDefaults({ refId: 'SQLVariableQueryEditor-VariableQuery' });
+  }
+}
@@ -3702,6 +3702,7 @@ __metadata:
    "@grafana/i18n": "npm:12.4.0-pre"
    "@grafana/plugin-ui": "npm:^0.11.0"
    "@grafana/runtime": "npm:12.4.0-pre"
+    "@grafana/schema": "npm:12.4.0-pre"
    "@grafana/ui": "npm:12.4.0-pre"
    "@react-awesome-query-builder/ui": "npm:6.6.15"
    "@testing-library/dom": "npm:10.4.1"