Elasticsearch: Fix query initialization logic & query transformation from Promethous/Loki (#31322)

* Elasticsearch: Fix query initialization logic

* Only import prometheus & loki queries as log queries

public/app/plugins/datasource/elasticsearch/components/QueryEditor/BucketAggregationsEditor/state/reducer.ts

@@ -16,7 +16,7 @@ import { bucketAggregationConfig } from '../utils';
 import { removeEmpty } from '../../../../utils';
 
 export const reducer = (
-  state: BucketAggregation[],
+  state: ElasticsearchQuery['bucketAggs'],
   action: BucketAggregationAction | ChangeMetricTypeAction | InitAction
 ): ElasticsearchQuery['bucketAggs'] => {
   switch (action.type) {
@@ -28,18 +28,18 @@ export const reducer = (
       };
 
       // If the last bucket aggregation is a `date_histogram` we add the new one before it.
-      const lastAgg = state[state.length - 1];
+      const lastAgg = state![state!.length - 1];
       if (lastAgg?.type === 'date_histogram') {
-        return [...state.slice(0, state.length - 1), newAgg, lastAgg];
+        return [...state!.slice(0, state!.length - 1), newAgg, lastAgg];
       }
 
-      return [...state, newAgg];
+      return [...state!, newAgg];
 
     case REMOVE_BUCKET_AGG:
-      return state.filter((bucketAgg) => bucketAgg.id !== action.payload.id);
+      return state!.filter((bucketAgg) => bucketAgg.id !== action.payload.id);
 
     case CHANGE_BUCKET_AGG_TYPE:
-      return state.map((bucketAgg) => {
+      return state!.map((bucketAgg) => {
         if (bucketAgg.id !== action.payload.id) {
           return bucketAgg;
         }
@@ -58,7 +58,7 @@ export const reducer = (
       });
 
     case CHANGE_BUCKET_AGG_FIELD:
-      return state.map((bucketAgg) => {
+      return state!.map((bucketAgg) => {
         if (bucketAgg.id !== action.payload.id) {
           return bucketAgg;
         }
@@ -74,7 +74,7 @@ export const reducer = (
       // we remove all of them.
       if (metricAggregationConfig[action.payload.type].isSingleMetric) {
         return [];
-      } else if (state.length === 0) {
+      } else if (state!.length === 0) {
         // Else, if there are no bucket aggregations we restore a default one.
         // This happens when switching from a metric that requires the absence of bucket aggregations to
         // one that requires it.
@@ -83,7 +83,7 @@ export const reducer = (
       return state;
 
     case CHANGE_BUCKET_AGG_SETTING:
-      return state.map((bucketAgg) => {
+      return state!.map((bucketAgg) => {
         if (bucketAgg.id !== action.payload.bucketAgg.id) {
           return bucketAgg;
         }
@@ -102,6 +102,9 @@ export const reducer = (
       });
 
     case INIT:
+      if (state?.length || 0 > 0) {
+        return state;
+      }
       return [defaultBucketAgg('2')];
 
     default:

public/app/plugins/datasource/elasticsearch/components/QueryEditor/ElasticsearchQueryContext.test.tsx

@@ -7,6 +7,7 @@ import { ElasticDatasource } from '../../datasource';
 
 const query: ElasticsearchQuery = {
   refId: 'A',
+  query: '',
   metrics: [{ id: '1', type: 'count' }],
   bucketAggs: [{ type: 'date_histogram', id: '2' }],
 };

public/app/plugins/datasource/elasticsearch/components/QueryEditor/ElasticsearchQueryContext.tsx

@@ -48,7 +48,7 @@ export const ElasticsearchProvider: FunctionComponent<Props> = ({
 
   // This initializes the query by dispatching an init action to each reducer.
   // useStatelessReducer will then call `onChange` with the newly generated query
-  if (!query.metrics && !query.bucketAggs) {
+  if (!query.metrics || !query.bucketAggs || query.query === undefined) {
     dispatch(initQuery());
 
     return null;

public/app/plugins/datasource/elasticsearch/components/QueryEditor/MetricAggregationsEditor/SettingsEditor/index.test.tsx

@@ -12,6 +12,7 @@ describe('Settings Editor', () => {
       const initialSize = '500';
       const query: ElasticsearchQuery = {
         refId: 'A',
+        query: '',
         metrics: [
           {
             id: metricId,
@@ -21,6 +22,7 @@ describe('Settings Editor', () => {
             },
           },
         ],
+        bucketAggs: [],
       };
 
       const onChange = jest.fn();

public/app/plugins/datasource/elasticsearch/components/QueryEditor/MetricAggregationsEditor/state/reducer.ts

@@ -22,24 +22,26 @@ import {
 } from './types';
 
 export const reducer = (
-  state: MetricAggregation[],
+  state: ElasticsearchQuery['metrics'],
   action: MetricAggregationAction | InitAction
 ): ElasticsearchQuery['metrics'] => {
   switch (action.type) {
     case ADD_METRIC:
-      return [...state, defaultMetricAgg(action.payload.id)];
+      return [...state!, defaultMetricAgg(action.payload.id)];
 
     case REMOVE_METRIC:
-      const metricToRemove = state.find((m) => m.id === action.payload.id)!;
-      const metricsToRemove = [metricToRemove, ...getChildren(metricToRemove, state)];
-      const resultingMetrics = state.filter((metric) => !metricsToRemove.some((toRemove) => toRemove.id === metric.id));
+      const metricToRemove = state!.find((m) => m.id === action.payload.id)!;
+      const metricsToRemove = [metricToRemove, ...getChildren(metricToRemove, state!)];
+      const resultingMetrics = state!.filter(
+        (metric) => !metricsToRemove.some((toRemove) => toRemove.id === metric.id)
+      );
       if (resultingMetrics.length === 0) {
         return [defaultMetricAgg('1')];
       }
       return resultingMetrics;
 
     case CHANGE_METRIC_TYPE:
-      return state
+      return state!
         .filter((metric) =>
           // When the new metric type is `isSingleMetric` we remove all other metrics from the query
           // leaving only the current one.
@@ -64,7 +66,7 @@ export const reducer = (
         });
 
     case CHANGE_METRIC_FIELD:
-      return state.map((metric) => {
+      return state!.map((metric) => {
         if (metric.id !== action.payload.id) {
           return metric;
         }
@@ -82,7 +84,7 @@ export const reducer = (
       });
 
     case TOGGLE_METRIC_VISIBILITY:
-      return state.map((metric) => {
+      return state!.map((metric) => {
         if (metric.id !== action.payload.id) {
           return metric;
         }
@@ -94,7 +96,7 @@ export const reducer = (
       });
 
     case CHANGE_METRIC_SETTING:
-      return state.map((metric) => {
+      return state!.map((metric) => {
         if (metric.id !== action.payload.metric.id) {
           return metric;
         }
@@ -119,7 +121,7 @@ export const reducer = (
       });
 
     case CHANGE_METRIC_META:
-      return state.map((metric) => {
+      return state!.map((metric) => {
         if (metric.id !== action.payload.metric.id) {
           return metric;
         }
@@ -140,7 +142,7 @@ export const reducer = (
       });
 
     case CHANGE_METRIC_ATTRIBUTE:
-      return state.map((metric) => {
+      return state!.map((metric) => {
         if (metric.id !== action.payload.metric.id) {
           return metric;
         }
@@ -152,6 +154,9 @@ export const reducer = (
       });
 
     case INIT:
+      if (state?.length || 0 > 0) {
+        return state;
+      }
       return [defaultMetricAgg('1')];
 
     default:

public/app/plugins/datasource/elasticsearch/components/QueryEditor/index.test.tsx

@@ -10,6 +10,7 @@ describe('QueryEditor', () => {
       const alias = '{{metric}}';
       const query: ElasticsearchQuery = {
         refId: 'A',
+        query: '',
         alias,
         metrics: [
           {

public/app/plugins/datasource/elasticsearch/components/QueryEditor/state.test.ts

@@ -1,8 +1,29 @@
 import { reducerTester } from 'test/core/redux/reducerTester';
 import { ElasticsearchQuery } from '../../types';
-import { aliasPatternReducer, changeAliasPattern, changeQuery, queryReducer } from './state';
+import { aliasPatternReducer, changeAliasPattern, changeQuery, initQuery, queryReducer } from './state';
 
 describe('Query Reducer', () => {
+  describe('On Init', () => {
+    it('Should maintain the previous `query` if present', () => {
+      const initialQuery: ElasticsearchQuery['query'] = 'Some lucene query';
+
+      reducerTester()
+        .givenReducer(queryReducer, initialQuery)
+        .whenActionIsDispatched(initQuery())
+        .thenStateShouldEqual(initialQuery);
+    });
+
+    it('Should set an empty `query` if it is not already set', () => {
+      const initialQuery: ElasticsearchQuery['query'] = undefined;
+      const expectedQuery = '';
+
+      reducerTester()
+        .givenReducer(queryReducer, initialQuery)
+        .whenActionIsDispatched(initQuery())
+        .thenStateShouldEqual(expectedQuery);
+    });
+  });
+
   it('Should correctly set `query`', () => {
     const expectedQuery: ElasticsearchQuery['query'] = 'Some lucene query';
 

public/app/plugins/datasource/elasticsearch/components/QueryEditor/state.ts

@@ -1,4 +1,5 @@
 import { Action } from '../../hooks/useStatelessReducer';
+import { ElasticsearchQuery } from '../../types';
 
 export const INIT = 'init';
 const CHANGE_QUERY = 'change_query';
@@ -18,6 +19,10 @@ interface ChangeAliasPatternAction extends Action<typeof CHANGE_ALIAS_PATTERN> {
   };
 }
 
+/**
+ * When the `initQuery` Action is dispatched, the query gets populated with default values where values are not present.
+ * This means it won't override any existing value in place, but just ensure the query is in a "runnable" state.
+ */
 export const initQuery = (): InitAction => ({ type: INIT });
 
 export const changeQuery = (query: string): ChangeQueryAction => ({
@@ -34,26 +39,29 @@ export const changeAliasPattern = (aliasPattern: string): ChangeAliasPatternActi
   },
 });
 
-export const queryReducer = (prevQuery: string, action: ChangeQueryAction | InitAction) => {
+export const queryReducer = (prevQuery: ElasticsearchQuery['query'], action: ChangeQueryAction | InitAction) => {
   switch (action.type) {
     case CHANGE_QUERY:
       return action.payload.query;
 
     case INIT:
-      return '';
+      return prevQuery || '';
 
     default:
       return prevQuery;
   }
 };
 
-export const aliasPatternReducer = (prevAliasPattern: string, action: ChangeAliasPatternAction | InitAction) => {
+export const aliasPatternReducer = (
+  prevAliasPattern: ElasticsearchQuery['alias'],
+  action: ChangeAliasPatternAction | InitAction
+) => {
   switch (action.type) {
     case CHANGE_ALIAS_PATTERN:
       return action.payload.aliasPattern;
 
     case INIT:
-      return '';
+      return prevAliasPattern || '';
 
     default:
       return prevAliasPattern;

public/app/plugins/datasource/elasticsearch/hooks/useNextId.test.tsx

@@ -8,6 +8,7 @@ describe('useNextId', () => {
   it('Should return the next available id', () => {
     const query: ElasticsearchQuery = {
       refId: 'A',
+      query: '',
       metrics: [{ id: '1', type: 'avg' }],
       bucketAggs: [{ id: '2', type: 'date_histogram' }],
     };

public/app/plugins/datasource/elasticsearch/language_provider.test.ts

@@ -2,8 +2,10 @@ import LanguageProvider from './language_provider';
 import { PromQuery } from '../prometheus/types';
 import { ElasticDatasource } from './datasource';
 import { DataSourceInstanceSettings } from '@grafana/data';
-import { ElasticsearchOptions } from './types';
+import { ElasticsearchOptions, ElasticsearchQuery } from './types';
 import { TemplateSrv } from '../../../features/templating/template_srv';
+import { defaultBucketAgg } from './query_def';
+import { DateHistogram } from './components/QueryEditor/BucketAggregationsEditor/aggregations';
 
 const templateSrvStub = {
   getAdhocFilters: jest.fn(() => [] as any[]),
@@ -22,81 +24,152 @@ const dataSource = new ElasticDatasource(
   } as DataSourceInstanceSettings<ElasticsearchOptions>,
   templateSrvStub as TemplateSrv
 );
+
+const baseLogsQuery: Partial<ElasticsearchQuery> = {
+  isLogsQuery: true,
+  metrics: [{ type: 'logs', id: '1' }],
+  bucketAggs: [{ ...defaultBucketAgg('2'), field: dataSource.timeField } as DateHistogram],
+};
+
 describe('transform prometheus query to elasticsearch query', () => {
-  it('Prometheus query with exact equals labels ( 2 labels ) and metric __name__', () => {
+  it('With exact equals labels ( 2 labels ) and metric __name__', () => {
     const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: 'my_metric{label1="value1",label2="value2"}' };
+    const promQuery: PromQuery = { refId: 'bar', expr: 'my_metric{label1="value1",label2="value2"}' };
     const result = instance.importQueries([promQuery], 'prometheus');
+
     expect(result).toEqual([
-      { isLogsQuery: true, query: '__name__:"my_metric" AND label1:"value1" AND label2:"value2"', refId: 'bar' },
+      {
+        ...baseLogsQuery,
+        query: '__name__:"my_metric" AND label1:"value1" AND label2:"value2"',
+        refId: promQuery.refId,
+      },
     ]);
   });
-  it('Prometheus query with exact equals labels ( 1 labels ) and metric __name__', () => {
+
+  it('With exact equals labels ( 1 labels ) and metric __name__', () => {
     const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: 'my_metric{label1="value1"}' };
-    const result = instance.importQueries([promQuery], 'prometheus');
-    expect(result).toEqual([{ isLogsQuery: true, query: '__name__:"my_metric" AND label1:"value1"', refId: 'bar' }]);
-  });
-  it('Prometheus query with exact equals labels ( 1 labels )', () => {
-    const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: '{label1="value1"}' };
-    const result = instance.importQueries([promQuery], 'prometheus');
-    expect(result).toEqual([{ isLogsQuery: true, query: 'label1:"value1"', refId: 'bar' }]);
-  });
-  it('Prometheus query with no label and metric __name__', () => {
-    const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: 'my_metric{}' };
-    const result = instance.importQueries([promQuery], 'prometheus');
-    expect(result).toEqual([{ isLogsQuery: true, query: '__name__:"my_metric"', refId: 'bar' }]);
-  });
-  it('Prometheus query with no label and metric __name__ without bracket', () => {
-    const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: 'my_metric' };
-    const result = instance.importQueries([promQuery], 'prometheus');
-    expect(result).toEqual([{ isLogsQuery: true, query: '__name__:"my_metric"', refId: 'bar' }]);
-  });
-  it('Prometheus query with rate function and exact equals labels ( 2 labels ) and metric __name__', () => {
-    const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: 'rate(my_metric{label1="value1",label2="value2"}[5m])' };
+    const promQuery: PromQuery = { refId: 'bar', expr: 'my_metric{label1="value1"}' };
     const result = instance.importQueries([promQuery], 'prometheus');
+
     expect(result).toEqual([
-      { isLogsQuery: true, query: '__name__:"my_metric" AND label1:"value1" AND label2:"value2"', refId: 'bar' },
+      {
+        ...baseLogsQuery,
+        query: '__name__:"my_metric" AND label1:"value1"',
+        refId: promQuery.refId,
+      },
     ]);
   });
-  it('Prometheus query with rate function and exact equals labels not equals labels regex and not regex labels and metric __name__', () => {
+
+  it('With exact equals labels ( 1 labels )', () => {
     const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = {
+    const promQuery: PromQuery = { refId: 'bar', expr: '{label1="value1"}' };
+    const result = instance.importQueries([promQuery], 'prometheus');
+
+    expect(result).toEqual([
+      {
+        ...baseLogsQuery,
+        query: 'label1:"value1"',
+        refId: promQuery.refId,
+      },
+    ]);
+  });
+
+  it('With no label and metric __name__', () => {
+    const instance = new LanguageProvider(dataSource);
+    const promQuery: PromQuery = { refId: 'bar', expr: 'my_metric{}' };
+    const result = instance.importQueries([promQuery], 'prometheus');
+
+    expect(result).toEqual([
+      {
+        ...baseLogsQuery,
+        query: '__name__:"my_metric"',
+        refId: promQuery.refId,
+      },
+    ]);
+  });
+
+  it('With no label and metric __name__ without bracket', () => {
+    const instance = new LanguageProvider(dataSource);
+    const promQuery: PromQuery = { refId: 'bar', expr: 'my_metric' };
+    const result = instance.importQueries([promQuery], 'prometheus');
+
+    expect(result).toEqual([
+      {
+        ...baseLogsQuery,
+        query: '__name__:"my_metric"',
+        refId: promQuery.refId,
+      },
+    ]);
+  });
+
+  it('With rate function and exact equals labels ( 2 labels ) and metric __name__', () => {
+    const instance = new LanguageProvider(dataSource);
+    const promQuery: PromQuery = { refId: 'bar', expr: 'rate(my_metric{label1="value1",label2="value2"}[5m])' };
+    const result = instance.importQueries([promQuery], 'prometheus');
+
+    expect(result).toEqual([
+      {
+        ...baseLogsQuery,
+        query: '__name__:"my_metric" AND label1:"value1" AND label2:"value2"',
+        refId: promQuery.refId,
+      },
+    ]);
+  });
+
+  it('With rate function and exact equals labels not equals labels regex and not regex labels and metric __name__', () => {
+    const instance = new LanguageProvider(dataSource);
+    const promQuery: PromQuery = {
       refId: 'bar',
       expr: 'rate(my_metric{label1="value1",label2!="value2",label3=~"value.+",label4!~".*tothemoon"}[5m])',
     };
     const result = instance.importQueries([promQuery], 'prometheus');
+
     expect(result).toEqual([
       {
-        isLogsQuery: true,
+        ...baseLogsQuery,
         query:
           '__name__:"my_metric" AND label1:"value1" AND NOT label2:"value2" AND label3:/value.+/ AND NOT label4:/.*tothemoon/',
-        refId: 'bar',
+        refId: promQuery.refId,
       },
     ]);
   });
 });
-describe('transform prometheus query to elasticsearch query errors', () => {
-  it('bad prometheus query with only bracket', () => {
+
+describe('transform malformed prometheus query to elasticsearch query', () => {
+  it('With only bracket', () => {
     const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: '{' };
+    const promQuery: PromQuery = { refId: 'bar', expr: '{' };
     const result = instance.importQueries([promQuery], 'prometheus');
-    expect(result).toEqual([{ isLogsQuery: true, query: '', refId: 'bar' }]);
+
+    expect(result).toEqual([
+      {
+        ...baseLogsQuery,
+        query: '',
+        refId: promQuery.refId,
+      },
+    ]);
   });
-  it('bad prometheus empty query', async () => {
+
+  it('Empty query', async () => {
     const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: '' };
+    const promQuery: PromQuery = { refId: 'bar', expr: '' };
     const result = instance.importQueries([promQuery], 'prometheus');
-    expect(result).toEqual([{ isLogsQuery: true, query: '', refId: 'bar' }]);
-  });
-  it('graphite query not handle', async () => {
-    const instance = new LanguageProvider(dataSource);
-    var promQuery: PromQuery = { refId: 'bar', expr: '' };
-    const result = instance.importQueries([promQuery], 'graphite');
-    expect(result).toEqual([{ isLogsQuery: true, query: '', refId: 'bar' }]);
+
+    expect(result).toEqual([
+      {
+        ...baseLogsQuery,
+        query: '',
+        refId: promQuery.refId,
+      },
+    ]);
+  });
+});
+
+describe('Unsupportated datasources', () => {
+  it('Generates a default query', async () => {
+    const instance = new LanguageProvider(dataSource);
+    const someQuery = { refId: 'bar' };
+    const result = instance.importQueries([someQuery], 'THIS DATASOURCE TYPE DOESNT EXIST');
+    expect(result).toEqual([{ refId: someQuery.refId }]);
   });
 });

public/app/plugins/datasource/elasticsearch/language_provider.ts

@@ -7,6 +7,7 @@ import { PromQuery } from '../prometheus/types';
 
 import Prism, { Token } from 'prismjs';
 import grammar from '../prometheus/promql';
+import { defaultBucketAgg } from './query_def';
 
 function getNameLabelValue(promQuery: string, tokens: any): string {
   let nameLabelValue = '';
@@ -104,13 +105,25 @@ export default class ElasticsearchLanguageProvider extends LanguageProvider {
     Object.assign(this, initialValues);
   }
 
+  /**
+   * The current implementation only supports switching from Prometheus/Loki queries.
+   * For them we transform the query to an ES Logs query since it's the behaviour most users expect.
+   * For every other datasource we just copy the refId and let the query editor initialize a default query.
+   * */
   importQueries(queries: DataQuery[], datasourceType: string): ElasticsearchQuery[] {
     if (datasourceType === 'prometheus' || datasourceType === 'loki') {
       return queries.map((query) => {
-        let prometheusQuery: PromQuery = query as PromQuery;
+        let prometheusQuery = query as PromQuery;
         const expr = getElasticsearchQuery(extractPrometheusLabels(prometheusQuery.expr));
         return {
           isLogsQuery: true,
+          metrics: [
+            {
+              id: '1',
+              type: 'logs',
+            },
+          ],
+          bucketAggs: [{ ...defaultBucketAgg('2'), field: this.datasource.timeField }],
           query: expr,
           refId: query.refId,
         };
@@ -118,8 +131,6 @@ export default class ElasticsearchLanguageProvider extends LanguageProvider {
     }
     return queries.map((query) => {
       return {
-        isLogsQuery: true,
-        query: '',
         refId: query.refId,
       };
     });