Add autocomplete to LogContext editor

Author: ddelemeny

src/LogContext/LogContextProvider.ts

@@ -126,6 +126,6 @@ export class LogContextProvider {
       runContextQuery?: (() => void),
       origQuery?: ElasticsearchQuery
   ): ReactNode {
-    return ( LogContextUI({row, runContextQuery, origQuery, updateQuery: query=>{this.contextQuery=query}}))
+    return ( LogContextUI({row, runContextQuery, origQuery, updateQuery: query=>{this.contextQuery=query}, datasource:this.datasource}))
   }
 }

src/LogContext/components/LogContextQueryBuilderSidebar.tsx

@@ -1,7 +1,7 @@
 import React, { useEffect, useMemo, useState } from "react";
 // import { Field } from '@grafana/data';
 import { useTheme2, CollapsableSection, Icon } from '@grafana/ui';
-import { LogContextProps } from "./LogContextUI";
+import { LogContextProps, useSearchableFields } from "./LogContextUI";
 import { css, cx } from "@emotion/css";
 import { LuceneQuery } from "utils/lucene";
 import { useQueryBuilderContext } from 'LogContext/QueryBuilder';
@@ -20,8 +20,8 @@ const excludedFields = [
   'timestamp_nanos',
 ];
 
-function isPrimitive(val: any) {
-  return val === null || ['string', 'number', "boolean", "undefined"].includes(typeof val)
+function isPrimitive(valT: any) {
+  return ['string', 'number', "boolean", "undefined"].includes(valT)
 }
 
 type FieldContingency = { [value: string]: {
@@ -103,19 +103,18 @@ type QueryBuilderProps = {
 
 export function LogContextQueryBuilderSidebar(props: LogContextProps & QueryBuilderProps) {
   const builder = useQueryBuilderContext();
+  const searchableFields = useSearchableFields();
 
   const {row, updateQuery} = props;
   const [fields, setFields] = useState<Field[]>([]);
 
-
-  const filteredFields = useMemo(() => (
-    row.dataFrame.fields
+  const filteredFields = useMemo(() => {
+    return searchableFields
+      // exclude some low-filterability fields
+      .filter((f)=> !excludedFields.includes(f.name) && isPrimitive(f.type))
       // sort fields by name
       .sort((f1, f2)=> (f1.name>f2.name ? 1 : -1))
-      // exclude some low-filterability fields
-      .filter((f)=> !excludedFields.includes(f.name) && isPrimitive(f.values[0])
-      )
-  ), [row.dataFrame.fields]);
+  }, [searchableFields]);
 
   useEffect(() => {
     const fields = filteredFields

src/LogContext/components/LogContextQueryEditor.tsx

@@ -1,17 +1,55 @@
-import React, { useRef} from "react";
+import React, { useCallback, useEffect, useMemo, useRef} from "react";
 import { css } from "@emotion/css";
 
 import { useQueryBuilderContext } from 'LogContext/QueryBuilder';
 
 import CodeMirror, { ReactCodeMirrorRef } from '@uiw/react-codemirror';
 import {linter, Diagnostic, lintGutter} from "@codemirror/lint"
+import {autocompletion, CompletionContext, Completion} from "@codemirror/autocomplete"
+import { useDatasource } from "components/QueryEditor/ElasticsearchQueryContext";
+import { useSearchableFields } from "./LogContextUI";
 
+function useAutocompleter(){
+  const datasource = useDatasource()
+  const searchableFields = useSearchableFields();
+  
+  useEffect(()=>{
+  }, [searchableFields])
+
+  const getSuggestions = useCallback(async (word: string)=>{
+    let suggestions: {from: number, options: Completion[]} = {from: 0, options: []};
+
+    const wordIsField = word.match(/([\w\.]+):(\S*)/)
+    if (wordIsField?.length) {
+      const [_match, fieldName, _fieldValue] = wordIsField;
+      const candidateValues = await datasource.getTagValues({key:fieldName})
+      suggestions.from = fieldName.length+1 // Replace only the value part
+      suggestions.options = candidateValues.map(v=>({
+        type: 'text',
+        label: typeof v.text === 'number' ? `${v.text}` : `"${v.text}"`
+      }))
+    } else {
+      const candidateFields = searchableFields
+      suggestions.from = 0
+      suggestions.options = candidateFields.map(f=>({
+        type: 'variable',
+        label: f.name,
+        detail: `${f.type}`
+      }))
+    }
+    return suggestions
+
+  },[datasource, searchableFields])
+  return getSuggestions
+}
 
 export function LogContextQueryEditor(){
   const editorRef = useRef<ReactCodeMirrorRef|null>(null)
   const builder = useQueryBuilderContext();
   const {parsedQuery} = builder;
 
+  const autocompleter = useAutocompleter()
+
   const queryLinter =  linter( view => {
     let diagnostics: Diagnostic[] = [];
 
@@ -26,13 +64,26 @@ export function LogContextQueryEditor(){
     }
     return diagnostics
   })
+
+  const autocomplete = useMemo(()=>autocompletion({
+    override: [async (context: CompletionContext)=>{
+      let word = context.matchBefore(/\S*/);
+      if (!word){ return null }
+      const suggestions = await autocompleter(word?.text);
+      return {
+        from: word.from + suggestions.from,
+        options: suggestions.options
+      }
+    }]
+  }),[autocompleter])
+
   return (<CodeMirror 
     ref={editorRef}
     className={css`height:100%`} // XXX : need to set height for both wrapper elements
     height="100%"
     theme={'dark'} 
     value={builder.query}
     onChange={(query)=>builder.setQuery(query || '' )}
-    extensions={[queryLinter, lintGutter()]}
+    extensions={[queryLinter, lintGutter(), autocomplete]}
     />);
 }

src/LogContext/components/LogContextUI.tsx

@@ -1,12 +1,15 @@
-import React, { useEffect, useState, useCallback, useMemo } from "react";
-import { LogRowModel, } from '@grafana/data';
+import React, { useEffect, useState, useCallback, useMemo, createContext } from "react";
+import { Field, LogRowModel, } from '@grafana/data';
 import { ElasticsearchQuery as DataQuery } from '../../types';
 import { LogContextQueryEditor } from "./LogContextQueryEditor";
 
 import { css } from "@emotion/css";
 import { Button } from "@grafana/ui";
 import { useQueryBuilder, QueryBuilderContext } from 'LogContext/QueryBuilder';
 import { LogContextQueryBuilderSidebar } from "./LogContextQueryBuilderSidebar";
+import { DatasourceContext } from "components/QueryEditor/ElasticsearchQueryContext";
+import { QuickwitDataSource } from "datasource";
+import { getHook } from "utils/context";
 
 const logContextUiStyle = css`
   display: flex;
@@ -15,20 +18,39 @@ const logContextUiStyle = css`
   height: 200px;
 `
 
+
 export interface LogContextProps {
   row: LogRowModel,
   runContextQuery?: (() => void)
   origQuery?: DataQuery
 }
 export interface LogContextUIProps extends LogContextProps {
+  datasource: QuickwitDataSource,
   updateQuery: (query: string) => void
 }
+const SearchableFieldsContext = createContext<Field[]|undefined>(undefined)
+export const useSearchableFields = getHook(SearchableFieldsContext)
 
 export function LogContextUI(props: LogContextUIProps ){
   const builder = useQueryBuilder();
   const {query, parsedQuery, setQuery, setParsedQuery} = builder;
   const [canRunQuery, setCanRunQuery] = useState<boolean>(false);
   const { origQuery, updateQuery, runContextQuery } = props;
+  const datasource = props.datasource;
+  const [fields, setFields] = useState<typeof props.row.dataFrame.fields>([]);
+
+  useEffect(()=>{
+    const dfFields = props.row.dataFrame.fields
+    // Get datasource fields definitions
+    datasource.getTagKeys({searchable:true}).then((dsFields)=>{
+      const dsFieldsNames = dsFields.map(f=>f.text)
+      setFields( dfFields.filter(f=> dsFieldsNames.includes(f.name)))
+
+    });
+
+
+
+  }, [props.row, datasource, setFields])
 
   useEffect(()=>{
     setQuery(origQuery?.query || '')
@@ -55,13 +77,17 @@ export function LogContextUI(props: LogContextUIProps ){
 
   return (
     <div className={logContextUiStyle}>
-      <QueryBuilderContext.Provider value={builder}>
-        <LogContextQueryBuilderSidebar {...props} updateQuery={setParsedQuery}/>
-        <div className={css`width:100%; display:flex; flex-direction:column; gap:0.5rem; min-width:0;`}>
-          {ActionBar}
-          <LogContextQueryEditor />
-        </div>
-      </QueryBuilderContext.Provider>
+      <DatasourceContext.Provider value={props.datasource}>
+        <QueryBuilderContext.Provider value={builder}>
+          <SearchableFieldsContext.Provider value={fields}>
+            <LogContextQueryBuilderSidebar {...props} updateQuery={setParsedQuery}/>
+            <div className={css`width:100%; display:flex; flex-direction:column; gap:0.5rem; min-width:0;`}>
+              {ActionBar}
+              <LogContextQueryEditor />
+            </div>
+          </SearchableFieldsContext.Provider>
+        </QueryBuilderContext.Provider>
+      </DatasourceContext.Provider>
     </div>
   );
 }

src/components/QueryEditor/ElasticsearchQueryContext.tsx

@@ -11,13 +11,13 @@ import { reducer as metricsReducer } from './MetricAggregationsEditor/state/redu
 import { aliasPatternReducer, queryReducer, initQuery, initExploreQuery } from './state';
 import { getHook } from 'utils/context';
 
-const RangeContext = createContext<TimeRange | undefined>(undefined);
+export const RangeContext = createContext<TimeRange | undefined>(undefined);
 export const useRange = getHook(RangeContext);
 
-const QueryContext = createContext<ElasticsearchQuery | undefined>(undefined);
+export const QueryContext = createContext<ElasticsearchQuery | undefined>(undefined);
 export const useQuery = getHook(QueryContext);
 
-const DatasourceContext = createContext<ElasticDatasource | undefined>(undefined);
+export const DatasourceContext = createContext<ElasticDatasource | undefined>(undefined);
 export const useDatasource = getHook(DatasourceContext);
 
 interface Props {

src/datasource.ts

@@ -54,6 +54,13 @@ export const REF_ID_STARTER_LOG_VOLUME = 'log-volume-';
 
 export type ElasticDatasource = QuickwitDataSource;
 
+type FieldCapsSpec = {
+  aggregatable?: boolean,
+  searchable?: boolean,
+  type?: string[],
+  _range?: TimeRange
+}
+
 export class QuickwitDataSource
   extends DataSourceWithBackend<ElasticsearchQuery, QuickwitOptions>
   implements
@@ -368,21 +375,21 @@ export class QuickwitDataSource
     );
   }
 
-  getFields(aggregatable?: boolean, type?: string[], _range?: TimeRange): Observable<MetricFindValue[]> {
+  getFields(spec: FieldCapsSpec={}): Observable<MetricFindValue[]> {
     // TODO: use the time range.
     return from(this.getResource('_elastic/' + this.index + '/_field_caps')).pipe(
       map((field_capabilities_response: FieldCapabilitiesResponse) => {
         const shouldAddField = (field: any) => {
-          if (aggregatable === undefined) {
-            return true;
+          if (spec.aggregatable !== undefined && field.aggregatable !== spec.aggregatable) {
+            return false
           }
-          if (aggregatable !== undefined && field.aggregatable !== aggregatable) {
-            return false;
+          if (spec.searchable !== undefined && field.searchable !== spec.searchable){
+            return false
           }
-          if (type?.length === 0) {
-            return true;
+          if (spec.type && spec.type.length !== 0 && !(spec.type.includes(field.type) || spec.type.includes(fieldTypeMap[field.type]))) {
+            return false
           }
-          return type?.includes(field.type) || type?.includes(fieldTypeMap[field.type]);
+          return true
         };
         const fieldCapabilities = Object.entries(field_capabilities_response.fields)
           .flatMap(([field_name, field_capabilities]) => {
@@ -412,8 +419,8 @@ export class QuickwitDataSource
   /**
    * Get tag keys for adhoc filters
    */
-  getTagKeys() {
-    return lastValueFrom(this.getFields());
+  getTagKeys(spec?: FieldCapsSpec) {
+    return lastValueFrom(this.getFields(spec));
   }
 
   /**
@@ -514,7 +521,7 @@ export class QuickwitDataSource
     if (query) {
       if (parsedQuery.find === 'fields') {
         parsedQuery.type = this.interpolateLuceneQuery(parsedQuery.type);
-        return lastValueFrom(this.getFields(true, parsedQuery.type, range));
+        return lastValueFrom(this.getFields({aggregatable:true, type:parsedQuery.type, _range:range}));
       }
       if (parsedQuery.find === 'terms') {
         parsedQuery.field = this.interpolateLuceneQuery(parsedQuery.field);

src/hooks/useFields.test.tsx

@@ -11,6 +11,8 @@ import { ElasticsearchQuery, MetricAggregationType, BucketAggregationType } from
 import { useFields } from './useFields';
 import { renderHook } from '@testing-library/react-hooks';
 
+
+
 describe('useFields hook', () => {
   // TODO: If we move the field type to the configuration objects as described in the hook's source
   // we can stop testing for getField to be called with the correct parameters.
@@ -48,39 +50,39 @@ describe('useFields hook', () => {
       { wrapper, initialProps: 'cardinality' }
     );
     result.current();
-    expect(getFields).toHaveBeenLastCalledWith(true, [], timeRange);
+    expect(getFields).toHaveBeenLastCalledWith({aggregatable:true, type:[], _range:timeRange});
 
     // All other metric aggregations only work on numbers
     rerender('avg');
     result.current();
-    expect(getFields).toHaveBeenLastCalledWith(true, ['number'], timeRange);
+    expect(getFields).toHaveBeenLastCalledWith({aggregatable:true, type:['number'], _range:timeRange});
 
     //
     // BUCKET AGGREGATIONS
     //
     // Date Histrogram only works on dates
     rerender('date_histogram');
     result.current();
-    expect(getFields).toHaveBeenLastCalledWith(true, ['date'], timeRange);
+    expect(getFields).toHaveBeenLastCalledWith({aggregatable:true, type:['date'], _range:timeRange});
 
     // Histrogram only works on numbers
     rerender('histogram');
     result.current();
-    expect(getFields).toHaveBeenLastCalledWith(true, ['number'], timeRange);
+    expect(getFields).toHaveBeenLastCalledWith({aggregatable:true, type:['number'], _range:timeRange});
 
     // Geohash Grid only works on geo_point data
     rerender('geohash_grid');
     result.current();
-    expect(getFields).toHaveBeenLastCalledWith(true, ['geo_point'], timeRange);
+    expect(getFields).toHaveBeenLastCalledWith({aggregatable:true, type:['geo_point'], _range:timeRange});
 
     // All other bucket aggregation work on any kind of data
     rerender('terms');
     result.current();
-    expect(getFields).toHaveBeenLastCalledWith(true, [], timeRange);
+    expect(getFields).toHaveBeenLastCalledWith({aggregatable:true, type:[], _range:timeRange});
 
     // top_metrics work on only on numeric data in 7.7
     rerender('top_metrics');
     result.current();
-    expect(getFields).toHaveBeenLastCalledWith(true, ['number'], timeRange);
+    expect(getFields).toHaveBeenLastCalledWith({aggregatable:true, type:['number'], _range:timeRange});
   });
 });

src/hooks/useFields.ts

@@ -62,7 +62,7 @@ export const useFields = (type: AggregationType | string[]) => {
   return async (q?: string) => {
     // _mapping doesn't support filtering, we avoid sending a request everytime q changes
     if (!rawFields) {
-      rawFields = await lastValueFrom(datasource.getFields(true, filter, range));
+      rawFields = await lastValueFrom(datasource.getFields({aggregatable:true, type:filter, _range:range}));
     }
 
     return rawFields.filter(({ text }) => q === undefined || text.includes(q)).map(toSelectableValue);