Loki: Add the ability to prettify logql queries (grafana#64337)

* pushed to get help of a genius

* fix: error response is not json

* feat: make request on click

* refactor: remove print statement

* refactor: remove unnecessary code

* feat: convert grafana variables to value for API request

* use the parser to interpolate and recover the original query (grafana#64591)

* Prettify query: use the parser to interpolate and recover the original query

* Fix typo

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>

* Fix typo

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>

* fix: reverse transformation not working

---------

Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>
Co-authored-by: Gareth Dawson <gwdawson.work@gmail.com>

* fix: bugs created from merge

* refactor: move prettify code out of monaco editor

* fix: variables with the same value get converted back to the incorect variable

* refactor

* use consistent styling with bigquery

* fix: only allow text/plain and application/json

* fix: only make the request if the query is valid

* endpoint now returns application/json

* prettify from js

* WIP: not all cases are handles, code still needs cleaning up

* WIP

* large refactor, finished support for all pipeline expressions

* add tests for all format functions

* idk why these files changed

* add support for range aggregation expr & refactor

* add support for vector aggregation expressions

* add support for bin op expression

* add support for literal and vector expressions

* add tests and fix some bugs

* add support for distinct and decolorize

* feat: update variable replace and return

* fix: lezer throws an errow when using a range variable

* remove api implementation

* remove api implementation

* remove type assertions

* add feature flag

* update naming

* fix: bug incorrectly formatting unwrap with labelfilter

* support label replace expr

* remove duplicate code (after migration)

* add more tests

* validate query before formatting

* move tests to lezer repo

* add feature tracking

* populate feature tracking with some data

* upgrade lezer version to 0.1.7

* bump lezer to 0.1.8

* add tests

---------

Co-authored-by: Matias Chomicki <matyax@gmail.com>
Co-authored-by: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>

Author: 3 people

docs/sources/setup-grafana/configure-grafana/feature-toggles/index.md

@@ -111,6 +111,7 @@ Experimental features might be changed or removed without prior notice.
 | `pluginsFrontendSandbox`                    | Enables the plugins frontend sandbox                                                                         |
 | `dashboardEmbed`                            | Allow embedding dashboard for external use in Code editors                                                   |
 | `frontendSandboxMonitorOnly`                | Enables monitor only in the plugin frontend sandbox (if enabled)                                             |
+| `lokiFormatQuery`                           | Enables the ability to format Loki queries                                                                   |
 | `cloudWatchLogsMonacoEditor`                | Enables the Monaco editor for CloudWatch Logs queries                                                        |
 | `exploreScrollableLogsContainer`            | Improves the scrolling behavior of logs in Explore                                                           |
 | `recordedQueriesMulti`                      | Enables writing multiple items from a single query within Recorded Queries                                   |

package.json

@@ -265,7 +265,7 @@
     "@grafana/faro-core": "1.1.0",
     "@grafana/faro-web-sdk": "1.1.0",
     "@grafana/google-sdk": "0.1.1",
-    "@grafana/lezer-logql": "0.1.5",
+    "@grafana/lezer-logql": "0.1.8",
     "@grafana/monaco-logql": "^0.0.7",
     "@grafana/runtime": "workspace:*",
     "@grafana/scenes": "0.22.0",

packages/grafana-data/src/types/featureToggles.gen.ts

@@ -98,6 +98,7 @@ export interface FeatureToggles {
   dashboardEmbed?: boolean;
   frontendSandboxMonitorOnly?: boolean;
   sqlDatasourceDatabaseSelection?: boolean;
+  lokiFormatQuery?: boolean;
   cloudWatchLogsMonacoEditor?: boolean;
   exploreScrollableLogsContainer?: boolean;
   recordedQueriesMulti?: boolean;

pkg/services/featuremgmt/registry.go

@@ -546,6 +546,13 @@ var (
 			Stage:        FeatureStagePublicPreview,
 			Owner:        grafanaBiSquad,
 		},
+		{
+			Name:         "lokiFormatQuery",
+			Description:  "Enables the ability to format Loki queries",
+			FrontendOnly: true,
+			Stage:        FeatureStageExperimental,
+			Owner:        grafanaObservabilityLogsSquad,
+		},
 		{
 			Name:         "cloudWatchLogsMonacoEditor",
 			Description:  "Enables the Monaco editor for CloudWatch Logs queries",

pkg/services/featuremgmt/toggles_gen.csv

@@ -79,6 +79,7 @@ pluginsFrontendSandbox,experimental,@grafana/plugins-platform-backend,false,fals
 dashboardEmbed,experimental,@grafana/grafana-as-code,false,false,false,true
 frontendSandboxMonitorOnly,experimental,@grafana/plugins-platform-backend,false,false,false,true
 sqlDatasourceDatabaseSelection,preview,@grafana/grafana-bi-squad,false,false,false,true
+lokiFormatQuery,experimental,@grafana/observability-logs,false,false,false,true
 cloudWatchLogsMonacoEditor,experimental,@grafana/aws-datasources,false,false,false,true
 exploreScrollableLogsContainer,experimental,@grafana/observability-logs,false,false,false,true
 recordedQueriesMulti,experimental,@grafana/observability-metrics,false,false,false,false

pkg/services/featuremgmt/toggles_gen.go

@@ -327,6 +327,10 @@ const (
 	// Enables previous SQL data source dataset dropdown behavior
 	FlagSqlDatasourceDatabaseSelection = "sqlDatasourceDatabaseSelection"
 
+	// FlagLokiFormatQuery
+	// Enables the ability to format Loki queries
+	FlagLokiFormatQuery = "lokiFormatQuery"
+
 	// FlagCloudWatchLogsMonacoEditor
 	// Enables the Monaco editor for CloudWatch Logs queries
 	FlagCloudWatchLogsMonacoEditor = "cloudWatchLogsMonacoEditor"

public/app/plugins/datasource/loki/components/LokiQueryField.tsx

@@ -75,7 +75,7 @@ export class LokiQueryField extends React.PureComponent<LokiQueryFieldProps, Lok
           className="gf-form-inline gf-form-inline--xs-view-flex-column flex-grow-1"
           data-testid={this.props['data-testid']}
         >
-          <div className="gf-form gf-form--grow flex-shrink-1 min-width-15">
+          <div className="gf-form--grow flex-shrink-1 min-width-15">
             <MonacoQueryFieldWrapper
               datasource={datasource}
               history={history ?? []}

public/app/plugins/datasource/loki/queryUtils.test.ts

@@ -1,5 +1,6 @@
 import { String } from '@grafana/lezer-logql';
 
+import { createLokiDatasource } from './mocks';
 import {
   getHighlighterExpressionsFromQuery,
   getLokiQueryType,
@@ -17,6 +18,7 @@ import {
   getLogQueryFromMetricsQuery,
   getNormalizedLokiQuery,
   getNodePositionsFromQuery,
+  formatLogqlQuery,
 } from './queryUtils';
 import { LokiQuery, LokiQueryType } from './types';
 
@@ -440,3 +442,27 @@ describe('getNodePositionsFromQuery', () => {
     expect(nodePositions.length).toBe(0);
   });
 });
+
+describe('formatLogqlQuery', () => {
+  const ds = createLokiDatasource();
+
+  it('formats a logs query', () => {
+    expect(formatLogqlQuery('{job="grafana"}', ds)).toBe('{job="grafana"}');
+  });
+
+  it('formats a metrics query', () => {
+    expect(formatLogqlQuery('count_over_time({job="grafana"}[1m])', ds)).toBe(
+      'count_over_time(\n  {job="grafana"}\n  [1m]\n)'
+    );
+  });
+
+  it('formats a metrics query with variables', () => {
+    // mock the interpolateString return value so it passes the isValid check
+    ds.interpolateString = jest.fn(() => 'rate({job="grafana"}[1s])');
+
+    expect(formatLogqlQuery('rate({job="grafana"}[$__range])', ds)).toBe('rate(\n  {job="grafana"}\n  [$__range]\n)');
+    expect(formatLogqlQuery('rate({job="grafana"}[$__interval])', ds)).toBe(
+      'rate(\n  {job="grafana"}\n  [$__interval]\n)'
+    );
+  });
+});

public/app/plugins/datasource/loki/queryUtils.ts

@@ -19,11 +19,15 @@ import {
   Identifier,
   Distinct,
   Range,
+  formatLokiQuery,
 } from '@grafana/lezer-logql';
+import { reportInteraction } from '@grafana/runtime';
 import { DataQuery } from '@grafana/schema';
 
-import { ErrorId } from '../prometheus/querybuilder/shared/parsingUtils';
+import { ErrorId, replaceVariables, returnVariables } from '../prometheus/querybuilder/shared/parsingUtils';
 
+import { placeHolderScopedVars } from './components/monaco-query-field/monaco-completion-provider/validation';
+import { LokiDatasource } from './datasource';
 import { getStreamSelectorPositions, NodePosition } from './modifyQuery';
 import { LokiQuery, LokiQueryType } from './types';
 
@@ -293,3 +297,39 @@ export const getLokiQueryFromDataQuery = (query?: DataQuery): LokiQuery | undefi
 
   return query;
 };
+
+export function formatLogqlQuery(query: string, datasource: LokiDatasource) {
+  const isInvalid = isQueryWithError(datasource.interpolateString(query, placeHolderScopedVars));
+
+  reportInteraction('grafana_loki_format_query_clicked', {
+    is_invalid: isInvalid,
+    query_type: isLogsQuery(query) ? 'logs' : 'metric',
+  });
+
+  if (isInvalid) {
+    return query;
+  }
+
+  let transformedQuery = replaceVariables(query);
+  const transformationMatches = [];
+  const tree = parser.parse(transformedQuery);
+
+  // Variables are considered errors inside of the parser, so we need to remove them before formatting
+  // We replace all variables with [0s] and keep track of the replaced variables
+  // After formatting we replace [0s] with the original variable
+  if (tree.topNode.firstChild?.firstChild?.type.id === MetricExpr) {
+    const pattern = /\[__V_[0-2]__\w+__V__\]/g;
+    transformationMatches.push(...transformedQuery.matchAll(pattern));
+    transformedQuery = transformedQuery.replace(pattern, '[0s]');
+  }
+
+  let formatted = formatLokiQuery(transformedQuery);
+
+  if (tree.topNode.firstChild?.firstChild?.type.id === MetricExpr) {
+    transformationMatches.forEach((match) => {
+      formatted = formatted.replace('[0s]', match[0]);
+    });
+  }
+
+  return returnVariables(formatted);
+}

public/app/plugins/datasource/loki/querybuilder/components/LokiQueryCodeEditor.tsx

@@ -2,12 +2,15 @@ import { css } from '@emotion/css';
 import React from 'react';
 
 import { GrafanaTheme2 } from '@grafana/data';
-import { useStyles2 } from '@grafana/ui';
+import { config } from '@grafana/runtime';
+import { useStyles2, HorizontalGroup, IconButton, Tooltip, Icon } from '@grafana/ui';
+import { getModKey } from 'app/core/utils/browser';
 
 import { testIds } from '../../components/LokiQueryEditor';
 import { LokiQueryField } from '../../components/LokiQueryField';
 import { getStats } from '../../components/stats';
 import { LokiQueryEditorProps } from '../../components/types';
+import { formatLogqlQuery } from '../../queryUtils';
 import { QueryStats } from '../../types';
 
 import { LokiQueryBuilderExplained } from './LokiQueryBuilderExplained';
@@ -31,6 +34,9 @@ export function LokiQueryCodeEditor({
 }: Props) {
   const styles = useStyles2(getStyles);
 
+  const lokiFormatQuery = config.featureToggles.lokiFormatQuery;
+  const onClickFormatQueryButton = async () => onChange({ ...query, expr: formatLogqlQuery(query.expr, datasource) });
+
   return (
     <div className={styles.wrapper}>
       <LokiQueryField
@@ -47,6 +53,27 @@ export function LokiQueryCodeEditor({
           const stats = await getStats(datasource, query);
           setQueryStats(stats);
         }}
+        ExtraFieldElement={
+          <>
+            {lokiFormatQuery && (
+              <div className={styles.buttonGroup}>
+                <div>
+                  <HorizontalGroup spacing="sm">
+                    <IconButton
+                      onClick={onClickFormatQueryButton}
+                      name="brackets-curly"
+                      size="xs"
+                      tooltip="Format query"
+                    />
+                    <Tooltip content={`Use ${getModKey()}+z to undo`}>
+                      <Icon className={styles.hint} name="keyboard" />
+                    </Tooltip>
+                  </HorizontalGroup>
+                </div>
+              </div>
+            )}
+          </>
+        }
       />
       {showExplain && <LokiQueryBuilderExplained query={query.expr} />}
     </div>
@@ -61,5 +88,20 @@ const getStyles = (theme: GrafanaTheme2) => {
         margin-bottom: 0.5;
       }
     `,
+    buttonGroup: css`
+      border: 1px solid ${theme.colors.border.medium};
+      border-top: none;
+      padding: ${theme.spacing(0.5, 0.5, 0.5, 0.5)};
+      margin-bottom: ${theme.spacing(0.5)};
+      display: flex;
+      flex-grow: 1;
+      justify-content: end;
+      font-size: ${theme.typography.bodySmall.fontSize};
+    `,
+    hint: css`
+      color: ${theme.colors.text.disabled};
+      white-space: nowrap;
+      cursor: help;
+    `,
   };
 };