Add support for context aware autocompletion

cypress/integration/smoke.spec.ts

@@ -109,22 +109,27 @@ e2e.scenario({
             e2eSelectors.ConfigEditor.table.wrapper().contains('cloudtrail_logs');
             e2eSelectors.ConfigEditor.table.input().type('cloudtrail_logs').type('{enter}');
 
-            // Verify editor suggestions
-            e2eSelectors.QueryEditor.CodeEditor.container().click({ force: true }).type(`{selectall}$__table`);
-            e2eSelectors.QueryEditor.CodeEditor.container().contains('(Macro) cloudtrail_logs');
+            // The follwing section will verify that autocompletion in behaving as expected.
+            // Throughout the composition of the SQL query, the autocompletion engine will provide appropriate suggestions.
+            // In this test the first few suggestions are accepted by hitting enter which will create a basic query.
+            // Increasing delay to allow tables names and columns names to be resolved async by the plugin
+            e2eSelectors.QueryEditor.CodeEditor.container()
+              .click({ force: true })
+              .type(`s{enter}{enter}{enter}{enter} {enter}{enter}`, { delay: 3000 });
+            e2eSelectors.QueryEditor.CodeEditor.container().contains(
+              'SELECT * FROM cloudtrail_logs GROUP BY additionaleventdata'
+            );
 
-            e2eSelectors.QueryEditor.CodeEditor.container().click({ force: true }).type(`{selectall}{enter}
-SELECT 
-    $__parseTime(eventtime, 'yyyy-MM-dd''T''HH:mm:ss''Z'), 
-    sum(cast(json_extract_scalar(additionaleventdata, '$.bytesTransferredOut') as real)) AS bytes 
-FROM 
-    $__table 
-WHERE additionaleventdata IS NOT NULL AND json_extract_scalar(additionaleventdata, '$.bytesTransferredOut') IS NOT NULL 
-AND 
-    $__timeFilter(eventtime, 'yyyy-MM-dd''T''HH:mm:ss''Z') 
+            e2eSelectors.QueryEditor.CodeEditor.container()
+              .click({ force: true })
+              .type(
+                `{selectall}
+SELECT $__parseTime(eventtime, 'yyyy-MM-dd''T''HH:mm:ss''Z'), sum(cast(json_extract_scalar(additionaleventdata, '$.bytesTransferredOut') as real)) AS bytes 
+FROM $__table WHERE additionaleventdata IS NOT NULL AND json_extract_scalar(additionaleventdata, '$.bytesTransferredOut') IS NOT NULL AND  $__timeFilter(eventtime, 'yyyy-MM-dd''T''HH:mm:ss''Z') 
 GROUP BY 1 
 ORDER BY 1 
-`);
+`
+              );
             // blur and wait for loading
             cy.get('.panel-content').click();
             cy.get('.panel-loading');

package.json

@@ -18,6 +18,7 @@
   "author": "Grafana Labs",
   "license": "Apache-2.0",
   "devDependencies": {
+    "@grafana/experimental": "^0.0.2-canary.38",
     "@grafana/aws-sdk": "0.0.37",
     "@grafana/data": "8.2.1",
     "@grafana/e2e": "8.2.1",

src/QueryEditor.test.tsx

@@ -7,6 +7,7 @@ import { select } from 'react-select-event';
 import { selectors } from 'tests/selectors';
 import { defaultKey } from 'types';
 import * as runtime from '@grafana/runtime';
+import * as experimental from '@grafana/experimental';
 
 const ds = mockDatasource;
 const q = mockQuery;
@@ -19,6 +20,13 @@ jest
 
 jest.spyOn(ds, 'getVariables').mockImplementation(mockGetVariables);
 
+jest.mock('@grafana/experimental', () => ({
+  ...jest.requireActual<typeof experimental>('@grafana/experimental'),
+  SQLEditor: function SQLEditor() {
+    return <></>;
+  },
+}));
+
 const props = {
   datasource: ds,
   query: q,

src/QueryEditor.tsx

@@ -3,10 +3,10 @@ import { QueryEditorProps, SelectableValue } from '@grafana/data';
 import { DataSource } from './datasource';
 import { AthenaDataSourceOptions, AthenaQuery, defaultQuery, SelectableFormatOptions } from './types';
 import { InlineSegmentGroup } from '@grafana/ui';
-import { FormatSelect, ResourceSelector, QueryCodeEditor } from '@grafana/aws-sdk';
+import { FormatSelect, ResourceSelector } from '@grafana/aws-sdk';
 import { selectors } from 'tests/selectors';
-import { getSuggestions } from 'Suggestions';
 import { appendTemplateVariables } from 'utils';
+import SQLEditor from 'SQLEditor';
 
 type Props = QueryEditorProps<DataSource, AthenaQuery, AthenaDataSourceOptions>;
 
@@ -133,13 +133,13 @@ export function QueryEditor(props: Props) {
             onRunQuery={props.onRunQuery}
           />
         </div>
+
         <div style={{ minWidth: '400px', marginLeft: '10px', flex: 1 }}>
-          <QueryCodeEditor
-            language="sql"
+          <SQLEditor
             query={queryWithDefaults}
-            onChange={props.onChange}
             onRunQuery={props.onRunQuery}
-            getSuggestions={getSuggestions}
+            onChange={props.onChange}
+            datasource={props.datasource}
           />
         </div>
       </InlineSegmentGroup>

src/SQLEditor.tsx

@@ -0,0 +1,54 @@
+import { SQLEditor as SQLCodeEditor } from '@grafana/experimental';
+import { DataSource } from 'datasource';
+import { getAthenaCompletionProvider } from 'language/completionItemProvider';
+import { TABLE_MACRO } from 'language/macros';
+import React, { useRef, useMemo, useCallback, useEffect } from 'react';
+import { AthenaQuery } from 'types';
+
+interface RawEditorProps {
+  query: AthenaQuery;
+  onRunQuery: () => void;
+  onChange: (q: AthenaQuery) => void;
+  datasource: DataSource;
+}
+
+export default function SQLEditor({ query, datasource, onRunQuery, onChange }: RawEditorProps) {
+  const queryRef = useRef<AthenaQuery>(query);
+  useEffect(() => {
+    queryRef.current = query;
+  }, [query]);
+
+  const getTables = useCallback(async () => {
+    const tables: string[] = await datasource.getTables(queryRef.current).catch(() => []);
+    return tables.map((table) => ({ name: table, completion: table }));
+  }, [datasource]);
+
+  const getColumns = useCallback(async (tableName?: string) => {
+    const columns: string[] = await datasource
+      .getColumns({
+        ...queryRef.current,
+        table: tableName ? tableName.replace(TABLE_MACRO, queryRef.current.table ?? '') : queryRef.current.table,
+      })
+      .catch(() => []);
+    return columns.map((column) => ({ name: column, completion: column }));
+  }, []);
+
+  const getTablesRef = useRef(getTables);
+  const getColumnsRef = useRef(getColumns);
+  const completionProvider = useMemo(
+    () => getAthenaCompletionProvider({ getTables: getTablesRef, getColumns: getColumnsRef }),
+    []
+  );
+
+  return (
+    <SQLCodeEditor
+      query={query.rawSQL}
+      onBlur={() => onRunQuery()}
+      onChange={(rawSQL) => onChange({ ...queryRef.current, rawSQL })}
+      language={{
+        id: 'sql',
+        completionProvider,
+      }}
+    ></SQLCodeEditor>
+  );
+}

src/VariableQueryEditor.tsx

@@ -1,10 +1,9 @@
 import React from 'react';
-import { QueryCodeEditor } from '@grafana/aws-sdk';
-import { getSuggestions } from 'Suggestions';
 import { AthenaQuery, AthenaDataSourceOptions } from './types';
 import { QueryEditorProps } from '@grafana/data';
 import { DataSource } from 'datasource';
+import { QueryEditor } from 'QueryEditor';
 
 export function VariableQueryCodeEditor(props: QueryEditorProps<DataSource, AthenaQuery, AthenaDataSourceOptions>) {
-  return <QueryCodeEditor {...props} language="sql" getSuggestions={getSuggestions} />;
+  return <QueryEditor {...props} />;
 }

src/language/completionItemProvider.ts

@@ -0,0 +1,32 @@
+import {
+  ColumnDefinition,
+  getStandardSQLCompletionProvider,
+  LanguageCompletionProvider,
+  TableDefinition,
+  TableIdentifier,
+} from '@grafana/experimental';
+import { MACROS } from './macros';
+
+interface CompletionProviderGetterArgs {
+  getTables: React.MutableRefObject<(d?: string) => Promise<TableDefinition[]>>;
+  getColumns: React.MutableRefObject<(table: string) => Promise<ColumnDefinition[]>>;
+}
+
+export const getAthenaCompletionProvider: (args: CompletionProviderGetterArgs) => LanguageCompletionProvider =
+  ({ getTables, getColumns }) =>
+  (monaco, language) => {
+    return {
+      // get standard SQL completion provider which will resolve functions and macros
+      ...(language && getStandardSQLCompletionProvider(monaco, language)),
+      triggerCharacters: ['.', ' ', '$', ',', '(', "'"],
+      tables: {
+        resolve: async () => {
+          return await getTables.current();
+        },
+      },
+      columns: {
+        resolve: async (t: TableIdentifier) => getColumns.current(t.table!),
+      },
+      supportedMacros: () => MACROS,
+    };
+  };