From a21a9b9c6cbe0319e7a90fefc37f68300f694181 Mon Sep 17 00:00:00 2001
From: Matias Chomicki <matyax@gmail.com>
Date: Thu, 6 Jun 2024 11:47:17 +0200
Subject: [PATCH] Loki: Kick start your query now applies templates to the
 current query (#87658)

* Kick start your query: keep pipe operations in the original query

* QueryPatternsModal: define keep operations from the operations list

* QueryPatternsModal: resolve deprecation

* QueryPatternsModal: use the correct import

* QueryPatternsModal: use category instead of order rank

* QueryPatternsModal: add unit test case

* Chore: change button cta
---
 .../querybuilder/components/QueryPattern.tsx  |  2 +-
 .../components/QueryPatternsModal.test.tsx    | 61 +++++++++++++------
 .../components/QueryPatternsModal.tsx         | 30 ++++++++-
 3 files changed, 69 insertions(+), 24 deletions(-)

diff --git a/public/app/plugins/datasource/loki/querybuilder/components/QueryPattern.tsx b/public/app/plugins/datasource/loki/querybuilder/components/QueryPattern.tsx
index c3d15c9711b..3407c447345 100644
--- a/public/app/plugins/datasource/loki/querybuilder/components/QueryPattern.tsx
+++ b/public/app/plugins/datasource/loki/querybuilder/components/QueryPattern.tsx
@@ -68,7 +68,7 @@ export const QueryPattern = (props: Props) => {
                 onPatternSelect(pattern);
               }}
             >
-              Replace query
+              Apply to query
             </Button>
             {hasNewQueryOption && (
               <Button
diff --git a/public/app/plugins/datasource/loki/querybuilder/components/QueryPatternsModal.test.tsx b/public/app/plugins/datasource/loki/querybuilder/components/QueryPatternsModal.test.tsx
index 5aef594b667..b6e986c1079 100644
--- a/public/app/plugins/datasource/loki/querybuilder/components/QueryPatternsModal.test.tsx
+++ b/public/app/plugins/datasource/loki/querybuilder/components/QueryPatternsModal.test.tsx
@@ -13,22 +13,24 @@ jest.mock('@grafana/runtime', () => ({
   reportInteraction: jest.fn(),
 }));
 
-const defaultProps = {
-  isOpen: true,
-  onClose: jest.fn(),
-  onChange: jest.fn(),
-  onAddQuery: jest.fn(),
-  query: {
-    refId: 'A',
-    expr: '{label1="foo", label2="bar"} |= "baz" |~ "qux"',
-  },
-  queries: [
-    {
+function getDefaultProps() {
+  return {
+    isOpen: true,
+    onClose: jest.fn(),
+    onChange: jest.fn(),
+    onAddQuery: jest.fn(),
+    query: {
       refId: 'A',
-      expr: '{label1="foo", label2="bar"}',
+      expr: '{label1="foo", label2="bar"} |= "baz" |~ "qux"',
     },
-  ],
-};
+    queries: [
+      {
+        refId: 'A',
+        expr: '{label1="foo", label2="bar"}',
+      },
+    ],
+  };
+}
 
 const queryPatterns = {
   logQueryPatterns: lokiQueryModeller.getQueryPatterns().filter((pattern) => pattern.type === LokiQueryPatternType.Log),
@@ -39,17 +41,17 @@ const queryPatterns = {
 
 describe('QueryPatternsModal', () => {
   it('renders the modal', () => {
-    render(<QueryPatternsModal {...defaultProps} />);
+    render(<QueryPatternsModal {...getDefaultProps()} />);
     expect(screen.getByText('Kick start your query')).toBeInTheDocument();
   });
   it('renders collapsible elements with all query pattern types', () => {
-    render(<QueryPatternsModal {...defaultProps} />);
+    render(<QueryPatternsModal {...getDefaultProps()} />);
     Object.values(LokiQueryPatternType).forEach((pattern) => {
       expect(screen.getByText(new RegExp(`${pattern} query starters`, 'i'))).toBeInTheDocument();
     });
   });
   it('can open and close query patterns section', async () => {
-    render(<QueryPatternsModal {...defaultProps} />);
+    render(<QueryPatternsModal {...getDefaultProps()} />);
     await userEvent.click(screen.getByText('Log query starters'));
     expect(screen.getByText(queryPatterns.logQueryPatterns[0].name)).toBeInTheDocument();
 
@@ -58,7 +60,7 @@ describe('QueryPatternsModal', () => {
   });
 
   it('can open and close multiple query patterns section', async () => {
-    render(<QueryPatternsModal {...defaultProps} />);
+    render(<QueryPatternsModal {...getDefaultProps()} />);
     await userEvent.click(screen.getByText('Log query starters'));
     expect(screen.getByText(queryPatterns.logQueryPatterns[0].name)).toBeInTheDocument();
 
@@ -73,6 +75,7 @@ describe('QueryPatternsModal', () => {
   });
 
   it('uses pattern if there is no existing query', async () => {
+    const defaultProps = getDefaultProps();
     render(<QueryPatternsModal {...defaultProps} query={{ expr: '{job="grafana"}', refId: 'A' }} />);
     await userEvent.click(screen.getByText('Log query starters'));
     expect(screen.getByText(queryPatterns.logQueryPatterns[0].name)).toBeInTheDocument();
@@ -87,7 +90,7 @@ describe('QueryPatternsModal', () => {
   });
 
   it('gives warning when selecting pattern if there is already existing query', async () => {
-    render(<QueryPatternsModal {...defaultProps} />);
+    render(<QueryPatternsModal {...getDefaultProps()} />);
     await userEvent.click(screen.getByText('Log query starters'));
     expect(screen.getByText(queryPatterns.logQueryPatterns[0].name)).toBeInTheDocument();
     const firstUseQueryButton = screen.getAllByRole('button', { name: 'Use this query' })[0];
@@ -96,6 +99,7 @@ describe('QueryPatternsModal', () => {
   });
 
   it('can use create new query when selecting pattern if there is already existing query', async () => {
+    const defaultProps = getDefaultProps();
     render(<QueryPatternsModal {...defaultProps} />);
     await userEvent.click(screen.getByText('Log query starters'));
     expect(screen.getByText(queryPatterns.logQueryPatterns[0].name)).toBeInTheDocument();
@@ -113,7 +117,7 @@ describe('QueryPatternsModal', () => {
   });
 
   it('does not show create new query option if onAddQuery function is not provided ', async () => {
-    render(<QueryPatternsModal {...defaultProps} onAddQuery={undefined} />);
+    render(<QueryPatternsModal {...getDefaultProps()} onAddQuery={undefined} />);
     await userEvent.click(screen.getByText('Log query starters'));
     expect(screen.getByText(queryPatterns.logQueryPatterns[0].name)).toBeInTheDocument();
     const useQueryButton = screen.getAllByRole('button', { name: 'Use this query' })[0];
@@ -121,4 +125,21 @@ describe('QueryPatternsModal', () => {
     expect(screen.queryByRole('button', { name: 'Create new query' })).not.toBeInTheDocument();
     expect(screen.getByText(/your current query will be replaced/)).toBeInTheDocument();
   });
+
+  it('applies a metric query on top of the existing log query', async () => {
+    const defaultProps = getDefaultProps();
+    render(<QueryPatternsModal {...defaultProps} />);
+    await userEvent.click(screen.getByText('Metric query starters'));
+    expect(screen.getByText(queryPatterns.metricQueryPatterns[0].name)).toBeInTheDocument();
+    const firstUseQueryButton = screen.getAllByRole('button', { name: 'Use this query' })[0];
+    await userEvent.click(firstUseQueryButton);
+    const createNewQueryButton = screen.getByRole('button', { name: 'Apply to query' });
+    await userEvent.click(createNewQueryButton);
+    await waitFor(() => {
+      expect(defaultProps.onChange).toHaveBeenCalledWith({
+        expr: 'sum(sum_over_time({label1="foo", label2="bar"} |= `baz` |~ `qux` | logfmt | __error__=`` | unwrap  | __error__=`` [$__auto]))',
+        refId: 'A',
+      });
+    });
+  });
 });
diff --git a/public/app/plugins/datasource/loki/querybuilder/components/QueryPatternsModal.tsx b/public/app/plugins/datasource/loki/querybuilder/components/QueryPatternsModal.tsx
index 385b032b05e..a8e0748cf22 100644
--- a/public/app/plugins/datasource/loki/querybuilder/components/QueryPatternsModal.tsx
+++ b/public/app/plugins/datasource/loki/querybuilder/components/QueryPatternsModal.tsx
@@ -2,14 +2,16 @@ import { css } from '@emotion/css';
 import { capitalize } from 'lodash';
 import React, { useMemo, useState } from 'react';
 
-import { CoreApp, DataQuery, GrafanaTheme2, getNextRefId } from '@grafana/data';
+import { CoreApp, GrafanaTheme2, getNextRefId } from '@grafana/data';
 import { reportInteraction } from '@grafana/runtime';
+import { DataQuery } from '@grafana/schema';
 import { Button, Collapse, Modal, useStyles2 } from '@grafana/ui';
 
 import { LokiQuery } from '../../types';
 import { lokiQueryModeller } from '../LokiQueryModeller';
+import { operationDefinitions } from '../operations';
 import { buildVisualQueryFromString } from '../parsing';
-import { LokiQueryPattern, LokiQueryPatternType } from '../types';
+import { LokiOperationId, LokiQueryPattern, LokiQueryPatternType, LokiVisualQueryOperationCategory } from '../types';
 
 import { QueryPattern } from './QueryPattern';
 
@@ -23,6 +25,21 @@ type Props = {
   onAddQuery?: (query: LokiQuery) => void;
 };
 
+const keepOperationCategories: string[] = [
+  LokiVisualQueryOperationCategory.Formats,
+  LokiVisualQueryOperationCategory.LineFilters,
+  LokiVisualQueryOperationCategory.LabelFilters,
+];
+const excludeOperationIds: string[] = [LokiOperationId.Unwrap];
+const keepOperations = operationDefinitions
+  .filter(
+    (operation) =>
+      operation.category &&
+      keepOperationCategories.includes(operation.category) &&
+      !excludeOperationIds.includes(operation.id)
+  )
+  .map((operation) => operation.id);
+
 export const QueryPatternsModal = (props: Props) => {
   const { isOpen, onClose, onChange, onAddQuery, query, queries, app } = props;
   const [openTabs, setOpenTabs] = useState<string[]>([]);
@@ -47,7 +64,14 @@ export const QueryPatternsModal = (props: Props) => {
       createNewQuery: hasNewQueryOption && selectAsNewQuery,
     });
 
-    visualQuery.query.operations = pattern.operations;
+    // Filter operations in the original query except those we configured to keep
+    visualQuery.query.operations = visualQuery.query.operations.filter((op) => keepOperations.includes(op.id));
+    // Filter operations in the pattern that are present in the original query
+    const patternOperations = pattern.operations.filter(
+      (patternOp) => visualQuery.query.operations.findIndex((op) => op.id === patternOp.id) < 0
+    );
+    visualQuery.query.operations = [...visualQuery.query.operations, ...patternOperations];
+
     if (hasNewQueryOption && selectAsNewQuery) {
       onAddQuery({
         ...query,