Tempo: Integrate context aware autocomplete API (#67845)

* Send query in search tag values call

* Make sure to send the full query when using the code editor

* Fix merge conflicts

* Remove unused params

public/app/plugins/datasource/tempo/SearchTraceQLEditor/SearchField.test.tsx

@@ -183,6 +183,7 @@ const renderSearchField = (
       }}
       tags={tags || []}
       hideTag={hideTag}
+      query={'{}'}
     />
   );
 };

public/app/plugins/datasource/tempo/SearchTraceQLEditor/SearchField.tsx

@@ -35,6 +35,7 @@ interface Props {
   hideTag?: boolean;
   hideValue?: boolean;
   allowDelete?: boolean;
+  query: string;
 }
 const SearchField = ({
   filter,
@@ -48,6 +49,7 @@ const SearchField = ({
   hideTag,
   hideValue,
   allowDelete,
+  query,
 }: Props) => {
   const styles = useStyles2(getStyles);
   const languageProvider = useMemo(() => new TempoLanguageProvider(datasource), [datasource]);
@@ -60,7 +62,7 @@ const SearchField = ({
 
   const updateOptions = async () => {
     try {
-      return filter.tag ? await languageProvider.getOptionsV2(scopedTag) : [];
+      return filter.tag ? await languageProvider.getOptionsV2(scopedTag, query) : [];
     } catch (error) {
       // Display message if Tempo is connected but search 404's
       if (isFetchError(error) && error?.status === 404) {
@@ -72,7 +74,12 @@ const SearchField = ({
     return [];
   };
 
-  const { loading: isLoadingValues, value: options } = useAsync(updateOptions, [scopedTag, languageProvider, setError]);
+  const { loading: isLoadingValues, value: options } = useAsync(updateOptions, [
+    scopedTag,
+    languageProvider,
+    setError,
+    query,
+  ]);
 
   useEffect(() => {
     if (Array.isArray(filter.value) && filter.value.length > 1 && filter.operator !== '=~') {

public/app/plugins/datasource/tempo/SearchTraceQLEditor/TagsInput.test.tsx

@@ -112,6 +112,7 @@ describe('TagsInput', () => {
         }}
         staticTags={[]}
         isTagsLoading={false}
+        query={''}
       />
     );
   };

public/app/plugins/datasource/tempo/SearchTraceQLEditor/TagsInput.tsx

@@ -34,6 +34,7 @@ interface Props {
   staticTags: Array<string | undefined>;
   isTagsLoading: boolean;
   hideValues?: boolean;
+  query: string;
 }
 const TagsInput = ({
   updateFilter,
@@ -44,6 +45,7 @@ const TagsInput = ({
   staticTags,
   isTagsLoading,
   hideValues,
+  query,
 }: Props) => {
   const styles = useStyles2(getStyles);
   const generateId = () => uuidv4().slice(0, 8);
@@ -77,6 +79,7 @@ const TagsInput = ({
             deleteFilter={deleteFilter}
             allowDelete={true}
             hideValue={hideValues}
+            query={query}
           />
           {i === filters.length - 1 && (
             <AccessoryButton variant={'secondary'} icon={'plus'} onClick={handleOnAdd} title={'Add tag'} />

public/app/plugins/datasource/tempo/SearchTraceQLEditor/TraceQLSearch.tsx

@@ -117,6 +117,7 @@ const TraceQLSearch = ({ datasource, query, onChange }: Props) => {
                 tags={[]}
                 hideScope={true}
                 hideTag={true}
+                query={traceQlQuery}
               />
             </InlineSearchField>
           ))}
@@ -160,6 +161,7 @@ const TraceQLSearch = ({ datasource, query, onChange }: Props) => {
               deleteFilter={deleteFilter}
               staticTags={staticTags}
               isTagsLoading={isTagsLoading}
+              query={traceQlQuery}
             />
           </InlineSearchField>
         </div>

public/app/plugins/datasource/tempo/configuration/TraceQLSearchTags.tsx

@@ -89,6 +89,7 @@ export function TraceQLSearchTags({ options, onOptionsChange, datasource }: Prop
           staticTags={staticTags}
           isTagsLoading={loading}
           hideValues={true}
+          query={'{}'}
         />
       ) : (
         <div>Invalid data source, please create a valid data source and try again</div>

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

@@ -117,8 +117,8 @@ export default class TempoLanguageProvider extends LanguageProvider {
     return options;
   }
 
-  async getOptionsV2(tag: string): Promise<Array<SelectableValue<string>>> {
-    const response = await this.request(`/api/v2/search/tag/${tag}/values`);
+  async getOptionsV2(tag: string, query: string): Promise<Array<SelectableValue<string>>> {
+    const response = await this.request(`/api/v2/search/tag/${tag}/values`, query ? { q: query } : {});
     let options: Array<SelectableValue<string>> = [];
     if (response && response.tagValues) {
       options = response.tagValues.map((v: { type: string; value: string }) => ({

public/app/plugins/datasource/tempo/traceql/autocomplete.ts

@@ -88,13 +88,13 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
     this.registerInteractionCommandId = id;
   }
 
-  private async getTagValues(tagName: string): Promise<Array<SelectableValue<string>>> {
+  private async getTagValues(tagName: string, query: string): Promise<Array<SelectableValue<string>>> {
     let tagValues: Array<SelectableValue<string>>;
 
     if (this.cachedValues.hasOwnProperty(tagName)) {
       tagValues = this.cachedValues[tagName];
     } else {
-      tagValues = await this.languageProvider.getOptionsV2(tagName);
+      tagValues = await this.languageProvider.getOptionsV2(tagName, query);
       this.cachedValues[tagName] = tagValues;
     }
     return tagValues;
@@ -134,7 +134,7 @@ export class CompletionProvider implements monacoTypes.languages.CompletionItemP
       case 'SPANSET_IN_VALUE':
         let tagValues;
         try {
-          tagValues = await this.getTagValues(situation.tagName);
+          tagValues = await this.getTagValues(situation.tagName, situation.query);
         } catch (error) {
           if (isFetchError(error)) {
             dispatch(notifyApp(createErrorNotification(error.data.error, new Error(error.data.message))));

public/app/plugins/datasource/tempo/traceql/situation.test.ts

@@ -1,4 +1,4 @@
-import { getSituation, Situation } from './situation';
+import { getSituation, SituationType } from './situation';
 
 jest.mock('@grafana/runtime', () => ({
   ...jest.requireActual('@grafana/runtime'),
@@ -7,7 +7,7 @@ jest.mock('@grafana/runtime', () => ({
 interface SituationTest {
   query: string;
   cursorPos: number;
-  expected: Situation;
+  expected: SituationType;
 }
 
 describe('situation', () => {
@@ -62,7 +62,7 @@ describe('situation', () => {
   tests.forEach((test) => {
     it(`${test.query} at ${test.cursorPos} is ${test.expected.type}`, async () => {
       const sit = getSituation(test.query, test.cursorPos);
-      expect(sit).toEqual(test.expected);
+      expect(sit).toEqual({ ...test.expected, query: test.query });
     });
   });
 });

public/app/plugins/datasource/tempo/traceql/situation.ts

@@ -1,16 +1,13 @@
-// we find the first error-node in the tree that is at the cursor-position.
-// NOTE: this might be too slow, might need to optimize it
-// (ideas: we do not need to go into every subtree, based on from/to)
-// also, only go to places that are in the sub-tree of the node found
-// by default by lezer. problem is, `next()` will go upward too,
-// and we do not want to go higher than our node
 import { SyntaxNode, Tree } from '@lezer/common';
 
 import { AttributeField, FieldExpression, FieldOp, parser, SpansetFilter } from '@grafana/lezer-traceql';
 
 type Direction = 'parent' | 'firstChild' | 'lastChild' | 'nextSibling' | 'prevSibling';
 type NodeType = number;
-export type Situation =
+
+export type Situation = { query: string } & SituationType;
+
+export type SituationType =
   | {
       type: 'UNKNOWN';
     }
@@ -46,7 +43,7 @@ type Path = Array<[Direction, NodeType[]]>;
 
 type Resolver = {
   path: NodeType[];
-  fun: (node: SyntaxNode, text: string, pos: number) => Situation | null;
+  fun: (node: SyntaxNode, text: string, pos: number) => SituationType | null;
 };
 
 function getErrorNode(tree: Tree, cursorPos: number): SyntaxNode | null {
@@ -101,6 +98,7 @@ export function getSituation(text: string, offset: number): Situation | null {
   // so we handle that case first
   if (text === '') {
     return {
+      query: text,
       type: 'EMPTY',
     };
   }
@@ -127,13 +125,14 @@ export function getSituation(text: string, offset: number): Situation | null {
     ids.push(cur.type.id);
   }
 
+  let situationType: SituationType | null = null;
   for (let resolver of RESOLVERS) {
     if (isPathMatch(resolver.path, ids)) {
-      return resolver.fun(currentNode, text, offset);
+      situationType = resolver.fun(currentNode, text, offset);
     }
   }
 
-  return null;
+  return { query: text, ...(situationType ?? { type: 'UNKNOWN' }) };
 }
 
 const ERROR_NODE_ID = 0;
@@ -157,7 +156,7 @@ const RESOLVERS: Resolver[] = [
   },
 ];
 
-function resolveSpanset(node: SyntaxNode, text: string, pos: number): Situation {
+function resolveSpanset(node: SyntaxNode): SituationType {
   const lastFieldExpression = walk(node, [['lastChild', [FieldExpression]]]);
   if (lastFieldExpression) {
     return {
@@ -170,7 +169,7 @@ function resolveSpanset(node: SyntaxNode, text: string, pos: number): Situation
   };
 }
 
-function resolveAttribute(node: SyntaxNode, text: string, pos: number): Situation {
+function resolveAttribute(node: SyntaxNode, text: string): SituationType {
   const attributeFieldParent = walk(node, [['parent', [AttributeField]]]);
   const attributeFieldParentText = attributeFieldParent ? getNodeText(attributeFieldParent, text) : '';
 
@@ -194,7 +193,7 @@ function resolveAttribute(node: SyntaxNode, text: string, pos: number): Situatio
   };
 }
 
-function resolveExpression(node: SyntaxNode, text: string, pos: number): Situation {
+function resolveExpression(node: SyntaxNode, text: string): SituationType {
   if (node.prevSibling?.type.id === FieldOp) {
     let attributeField = node.prevSibling.prevSibling;
     if (attributeField) {
@@ -210,7 +209,7 @@ function resolveExpression(node: SyntaxNode, text: string, pos: number): Situati
   };
 }
 
-function resolveErrorInFilterRoot(node: SyntaxNode, text: string, pos: number): Situation {
+function resolveErrorInFilterRoot(): SituationType {
   return {
     type: 'SPANSET_IN_NAME',
   };