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

@@ -17,6 +17,9 @@
   },
   "author": "Grafana Labs",
   "license": "Apache-2.0",
+  "dependencies": {
+    "@grafana/experimental": "^0.0.2-canary.39"
+  },
   "devDependencies": {
     "@grafana/aws-sdk": "0.0.37",
     "@grafana/data": "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,
@@ -153,4 +161,10 @@ describe('QueryEditor', () => {
     });
     expect(onRunQuery).toHaveBeenCalled();
   });
+
+  it('should display query options by default', async () => {
+    render(<QueryEditor {...props} />);
+    const selectEl = screen.getByLabelText('Format as');
+    expect(selectEl).toBeInTheDocument();
+  });
 });

src/QueryEditor.tsx

@@ -3,12 +3,12 @@ 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>;
+type Props = QueryEditorProps<DataSource, AthenaQuery, AthenaDataSourceOptions> & { hideOptions?: boolean };
 
 type QueryProperties = 'regions' | 'catalogs' | 'databases' | 'tables' | 'columns';
 
@@ -125,21 +125,25 @@ export function QueryEditor(props: Props) {
             labelWidth={11}
             className="width-12"
           />
-          <h6>Frames</h6>
-          <FormatSelect
-            query={props.query}
-            options={SelectableFormatOptions}
-            onChange={props.onChange}
-            onRunQuery={props.onRunQuery}
-          />
+          {!props.hideOptions && (
+            <>
+              <h6>Frames</h6>
+              <FormatSelect
+                query={props.query}
+                options={SelectableFormatOptions}
+                onChange={props.onChange}
+                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,57 @@
+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 }));
+    },
+    [datasource]
+  );
+
+  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.test.tsx

@@ -0,0 +1,45 @@
+import React from 'react';
+import { render, screen } from '@testing-library/react';
+import { VariableQueryCodeEditor } from './VariableQueryEditor';
+import { mockDatasource, mockQuery } from './__mocks__/datasource';
+import '@testing-library/jest-dom';
+import * as runtime from '@grafana/runtime';
+import * as experimental from '@grafana/experimental';
+
+const ds = mockDatasource;
+const q = mockQuery;
+
+const mockGetVariables = jest.fn().mockReturnValue([]);
+
+jest
+  .spyOn(runtime, 'getTemplateSrv')
+  .mockImplementation(() => ({ getVariables: mockGetVariables, replace: jest.fn() }));
+
+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,
+  onChange: jest.fn(),
+  onRunQuery: jest.fn(),
+};
+
+beforeEach(() => {
+  ds.getResource = jest.fn().mockResolvedValue([]);
+  ds.postResource = jest.fn().mockResolvedValue([]);
+});
+
+describe('VariableQueryEditor', () => {
+  it('should not display query options', async () => {
+    render(<VariableQueryCodeEditor {...props} />);
+    const selectEl = screen.queryByLabelText('Format as');
+    expect(selectEl).toBeNull();
+  });
+});

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} hideOptions={true} />;
 }

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,
+    };
+  };