From 390461f9e140427c2e72658fe9e7a25266d3b68a Mon Sep 17 00:00:00 2001
From: Joey <90795735+joey-grafana@users.noreply.github.com>
Date: Tue, 6 Feb 2024 13:30:59 +0000
Subject: [PATCH] Tempo: Move query ref to TraceQLEditor (#81686)

Move ref to TraceQLEditor
---
 .../datasource/tempo/traceql/QueryEditor.tsx  | 16 +++----------
 .../tempo/traceql/TraceQLEditor.tsx           | 24 ++++++++++++++-----
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/public/app/plugins/datasource/tempo/traceql/QueryEditor.tsx b/public/app/plugins/datasource/tempo/traceql/QueryEditor.tsx
index 21b4cfc6cd1..7e5f7e024e8 100644
--- a/public/app/plugins/datasource/tempo/traceql/QueryEditor.tsx
+++ b/public/app/plugins/datasource/tempo/traceql/QueryEditor.tsx
@@ -1,6 +1,6 @@
 import { css } from '@emotion/css';
 import { defaults } from 'lodash';
-import React, { useRef, useState } from 'react';
+import React, { useState } from 'react';
 
 import { GrafanaTheme2, QueryEditorProps } from '@grafana/data';
 import { config, reportInteraction } from '@grafana/runtime';
@@ -27,16 +27,6 @@ export function QueryEditor(props: Props) {
     return genQuery === query.query || genQuery === '{}';
   });
 
-  // The Monaco Editor uses the first version of props.onChange in handleOnMount i.e. always has the initial
-  // value of query because underlying Monaco editor is passed `query` below in the onEditorChange callback.
-  // handleOnMount is called only once when the editor is mounted and does not get updates to query.
-  // So we need useRef to get the latest version of query in the onEditorChange callback.
-  const queryRef = useRef(query);
-  queryRef.current = query;
-  const onEditorChange = (value: string) => {
-    props.onChange({ ...queryRef.current, query: value });
-  };
-
   return (
     <>
       <InlineLabel>
@@ -73,8 +63,8 @@ export function QueryEditor(props: Props) {
       )}
       <TraceQLEditor
         placeholder="Enter a TraceQL query or trace ID (run with Shift+Enter)"
-        value={query.query || ''}
-        onChange={onEditorChange}
+        query={query}
+        onChange={props.onChange}
         datasource={props.datasource}
         onRunQuery={props.onRunQuery}
       />
diff --git a/public/app/plugins/datasource/tempo/traceql/TraceQLEditor.tsx b/public/app/plugins/datasource/tempo/traceql/TraceQLEditor.tsx
index 253a6240870..ff5e252393f 100644
--- a/public/app/plugins/datasource/tempo/traceql/TraceQLEditor.tsx
+++ b/public/app/plugins/datasource/tempo/traceql/TraceQLEditor.tsx
@@ -7,6 +7,7 @@ import { reportInteraction } from '@grafana/runtime';
 import { CodeEditor, Monaco, monacoTypes, useTheme2 } from '@grafana/ui';
 
 import { TempoDatasource } from '../datasource';
+import { TempoQuery } from '../types';
 
 import { CompletionProvider, CompletionType } from './autocomplete';
 import { getErrorNodes, setMarkers } from './highlighting';
@@ -14,8 +15,8 @@ import { languageDefinition } from './traceql';
 
 interface Props {
   placeholder: string;
-  value: string;
-  onChange: (val: string) => void;
+  query: TempoQuery;
+  onChange: (val: TempoQuery) => void;
   onRunQuery: () => void;
   datasource: TempoDatasource;
   readOnly?: boolean;
@@ -24,10 +25,21 @@ interface Props {
 export function TraceQLEditor(props: Props) {
   const [alertText, setAlertText] = useState('');
 
-  const { onChange, onRunQuery, placeholder } = props;
+  const { query, onChange, onRunQuery, placeholder } = props;
   const setupAutocompleteFn = useAutocomplete(props.datasource, setAlertText);
   const theme = useTheme2();
   const styles = getStyles(theme, placeholder);
+
+  // The Monaco Editor uses the first version of props.onChange in handleOnMount i.e. always has the initial
+  // value of query because underlying Monaco editor is passed `query` below in the onEditorChange callback.
+  // handleOnMount is called only once when the editor is mounted and does not get updates to query.
+  // So we need useRef to get the latest version of query in the onEditorChange callback.
+  const queryRef = useRef(query);
+  queryRef.current = query;
+  const onEditorChange = (value: string) => {
+    onChange({ ...queryRef.current, query: value });
+  };
+
   // work around the problem that `onEditorDidMount` is called once
   // and wouldn't get new version of onRunQuery
   const onRunQueryRef = useRef(onRunQuery);
@@ -38,10 +50,10 @@ export function TraceQLEditor(props: Props) {
   return (
     <>
       <CodeEditor
-        value={props.value}
+        value={query.query || ''}
         language={langId}
-        onBlur={onChange}
-        onChange={onChange}
+        onBlur={onEditorChange}
+        onChange={onEditorChange}
         containerStyles={styles.queryField}
         readOnly={props.readOnly}
         monacoOptions={{