SQL: Fix code editor for SQL datasources (#58116)

* SQL: Fix code editor for sql datasources

* Fix: mysql completion with defaultdb

public/app/features/plugins/sql/components/query-editor-raw/QueryEditorRaw.tsx

@@ -1,9 +1,8 @@
 import React, { useCallback, useEffect, useRef } from 'react';
 
-import { SQLEditor } from '@grafana/experimental';
+import { LanguageDefinition, SQLEditor } from '@grafana/experimental';
 
-import { LanguageCompletionProvider, SQLQuery } from '../../types';
-import { formatSQL } from '../../utils/formatSQL';
+import { SQLQuery } from '../../types';
 
 type Props = {
   query: SQLQuery;
@@ -11,10 +10,10 @@ type Props = {
   children?: (props: { formatQuery: () => void }) => React.ReactNode;
   width?: number;
   height?: number;
-  completionProvider: LanguageCompletionProvider;
+  editorLanguageDefinition: LanguageDefinition;
 };
 
-export function QueryEditorRaw({ children, onChange, query, width, height, completionProvider }: Props) {
+export function QueryEditorRaw({ children, onChange, query, width, height, editorLanguageDefinition }: Props) {
   // We need to pass query via ref to SQLEditor as onChange is executed via monacoEditor.onDidChangeModelContent callback, not onChange property
   const queryRef = useRef<SQLQuery>(query);
   useEffect(() => {
@@ -39,7 +38,7 @@ export function QueryEditorRaw({ children, onChange, query, width, height, compl
       height={height}
       query={query.rawSql!}
       onChange={onRawQueryChange}
-      language={{ id: 'sql', completionProvider, formatter: formatSQL }}
+      language={editorLanguageDefinition}
     >
       {children}
     </SQLEditor>

public/app/features/plugins/sql/components/query-editor-raw/RawEditor.tsx

@@ -25,12 +25,12 @@ export function RawEditor({ db, query, onChange, onRunQuery, onValidate, queryTo
   const [toolboxRef, toolboxMeasure] = useMeasure<HTMLDivElement>();
   const [editorRef, editorMeasure] = useMeasure<HTMLDivElement>();
 
-  const completionProvider = useMemo(() => db.getSqlCompletionProvider(), [db]);
+  const editorLanguageDefinition = useMemo(() => db.getEditorLanguageDefinition(), [db]);
 
   const renderQueryEditor = (width?: number, height?: number) => {
     return (
       <QueryEditorRaw
-        completionProvider={completionProvider}
+        editorLanguageDefinition={editorLanguageDefinition}
         query={query}
         width={width}
         height={height ? height - toolboxMeasure.height : undefined}

public/app/features/plugins/sql/components/visual-query-builder/SQLSelectRow.tsx

@@ -1,8 +1,8 @@
 import React from 'react';
-import { useAsync } from 'react-use';
 
 import { SelectableValue, toOption } from '@grafana/data';
 
+import { COMMON_AGGREGATE_FNS } from '../../constants';
 import { QueryWithDefaults } from '../../defaults';
 import { DB, SQLQuery } from '../../types';
 import { useSqlChange } from '../../utils/useSqlChange';
@@ -18,18 +18,14 @@ interface SQLSelectRowProps {
 
 export function SQLSelectRow({ fields, query, onQueryChange, db }: SQLSelectRowProps) {
   const { onSqlChange } = useSqlChange({ query, onQueryChange, db });
-
-  const state = useAsync(async () => {
-    const functions = await db.functions();
-    return functions.map((f) => toOption(f.name));
-  }, [db]);
+  const functions = [...COMMON_AGGREGATE_FNS, ...(db.functions?.() || [])].map(toOption);
 
   return (
     <SelectRow
       columns={fields}
       sql={query.sql!}
       format={query.format}
-      functions={state.value}
+      functions={functions}
       onSqlChange={onSqlChange}
     />
   );

public/app/features/plugins/sql/constants.ts

@@ -1,115 +1,4 @@
-import { OperatorType } from './types';
-
-export const AGGREGATE_FNS = [
-  {
-    id: 'AVG',
-    name: 'AVG',
-    description: `AVG(
-    [DISTINCT]
-    expression
-  )
-  [OVER (...)]
-
-  Returns the average of non-NULL input values, or NaN if the input contains a NaN.`,
-  },
-  {
-    id: 'COUNT',
-    name: 'COUNT',
-    description: `COUNT(*)  [OVER (...)]
-  Returns the number of rows in the input.
-
-  COUNT(
-    [DISTINCT]
-    expression
-  )
-  [OVER (...)]
-
-  Returns the number of rows with expression evaluated to any value other than NULL.
-  `,
-  },
-  {
-    id: 'MAX',
-    name: 'MAX',
-    description: `MAX(
-    expression
-  )
-  [OVER (...)]
-
-  Returns the maximum value of non-NULL expressions. Returns NULL if there are zero input rows or expression evaluates to NULL for all rows. Returns NaN if the input contains a NaN.
-  `,
-  },
-  {
-    id: 'MIN',
-    name: 'MIN',
-    description: `MIN(
-    expression
-  )
-  [OVER (...)]
-
-  Returns the minimum value of non-NULL expressions. Returns NULL if there are zero input rows or expression evaluates to NULL for all rows. Returns NaN if the input contains a NaN.
-  `,
-  },
-  {
-    id: 'SUM',
-    name: 'SUM',
-    description: `SUM(
-    [DISTINCT]
-    expression
-  )
-  [OVER (...)]
-
-  Returns the sum of non-null values.
-
-  If the expression is a floating point value, the sum is non-deterministic, which means you might receive a different result each time you use this function.
-  `,
-  },
-];
-
-export const OPERATORS = [
-  { type: OperatorType.Comparison, id: 'LESS_THAN', operator: '<', description: 'Returns TRUE if X is less than Y.' },
-  {
-    type: OperatorType.Comparison,
-    id: 'LESS_THAN_EQUAL',
-    operator: '<=',
-    description: 'Returns TRUE if X is less than or equal to Y.',
-  },
-  {
-    type: OperatorType.Comparison,
-    id: 'GREATER_THAN',
-    operator: '>',
-    description: 'Returns TRUE if X is greater than Y.',
-  },
-  {
-    type: OperatorType.Comparison,
-    id: 'GREATER_THAN_EQUAL',
-    operator: '>=',
-    description: 'Returns TRUE if X is greater than or equal to Y.',
-  },
-  { type: OperatorType.Comparison, id: 'EQUAL', operator: '=', description: 'Returns TRUE if X is equal to Y.' },
-  {
-    type: OperatorType.Comparison,
-    id: 'NOT_EQUAL',
-    operator: '!=',
-    description: 'Returns TRUE if X is not equal to Y.',
-  },
-  {
-    type: OperatorType.Comparison,
-    id: 'NOT_EQUAL_ALT',
-    operator: '<>',
-    description: 'Returns TRUE if X is not equal to Y.',
-  },
-  {
-    type: OperatorType.Comparison,
-    id: 'LIKE',
-    operator: 'LIKE',
-    description: `Checks if the STRING in the first operand X matches a pattern specified by the second operand Y. Expressions can contain these characters:
-- A percent sign "%" matches any number of characters or bytes
-- An underscore "_" matches a single character or byte
-- You can escape "\", "_", or "%" using two backslashes. For example, "\\%". If you are using raw strings, only a single backslash is required. For example, r"\%".`,
-  },
-  { type: OperatorType.Logical, id: 'AND', operator: 'AND' },
-  { type: OperatorType.Logical, id: 'OR', operator: 'OR' },
-];
+export const COMMON_AGGREGATE_FNS = ['AVG', 'COUNT', 'MAX', 'MIN', 'SUM'];
 
 export const MACRO_NAMES = [
   '$__time',

public/app/features/plugins/sql/types.ts

@@ -9,7 +9,7 @@ import {
   TimeRange,
   toOption as toOptionFromData,
 } from '@grafana/data';
-import { CompletionItemKind, EditorMode, LanguageCompletionProvider } from '@grafana/experimental';
+import { CompletionItemKind, EditorMode, LanguageDefinition } from '@grafana/experimental';
 
 import { QueryWithDefaults } from './defaults';
 import {
@@ -122,12 +122,6 @@ export interface SQLSelectableValue extends SelectableValue {
   raqbFieldType?: RAQBFieldTypes;
 }
 
-export interface Aggregate {
-  id: string;
-  name: string;
-  description?: string;
-}
-
 export interface DB {
   init?: (datasourceId?: string) => Promise<boolean>;
   datasets: () => Promise<string[]>;
@@ -136,10 +130,10 @@ export interface DB {
   validateQuery: (query: SQLQuery, range?: TimeRange) => Promise<ValidationResults>;
   dsID: () => number;
   dispose?: (dsID?: string) => void;
-  lookup: (path?: string) => Promise<Array<{ name: string; completion: string }>>;
-  getSqlCompletionProvider: () => LanguageCompletionProvider;
+  lookup?: (path?: string) => Promise<Array<{ name: string; completion: string }>>;
+  getEditorLanguageDefinition: () => LanguageDefinition;
   toRawSql?: (query: SQLQuery) => string;
-  functions: () => Promise<Aggregate[]>;
+  functions?: () => string[];
 }
 
 export interface QueryEditorProps {
@@ -173,18 +167,3 @@ export interface MetaDefinition {
   completion?: string;
   kind: CompletionItemKind;
 }
-
-export {
-  CompletionItemKind,
-  LanguageCompletionProvider,
-  LinkedToken,
-  ColumnDefinition,
-  CompletionItemPriority,
-  StatementPlacementProvider,
-  SuggestionKindProvider,
-  TableDefinition,
-  TokenType,
-  OperatorType,
-  StatementPosition,
-  PositionContext,
-} from '@grafana/experimental';

public/app/plugins/datasource/cloudwatch/dynamic-labels/CompletionItemProvider.test.ts

@@ -1,5 +1,5 @@
+import { CompletionItemPriority } from '@grafana/experimental';
 import { Monaco, monacoTypes } from '@grafana/ui';
-import { CompletionItemPriority } from 'app/features/plugins/sql';
 
 import { afterLabelValue, insideLabelValue } from '../__mocks__/dynamic-label-test-data';
 import MonacoMock from '../__mocks__/monarch/Monaco';

public/app/plugins/datasource/mssql/datasource.ts

@@ -1,8 +1,9 @@
 import { DataSourceInstanceSettings, ScopedVars } from '@grafana/data';
+import { LanguageDefinition } from '@grafana/experimental';
 import { TemplateSrv } from '@grafana/runtime';
-import { AGGREGATE_FNS } from 'app/features/plugins/sql/constants';
 import { SqlDatasource } from 'app/features/plugins/sql/datasource/SqlDatasource';
-import { DB, LanguageCompletionProvider, SQLQuery, SQLSelectableValue } from 'app/features/plugins/sql/types';
+import { DB, SQLQuery, SQLSelectableValue } from 'app/features/plugins/sql/types';
+import { formatSQL } from 'app/features/plugins/sql/utils/formatSQL';
 
 import { getSchema, showDatabases, getSchemaAndName } from './MSSqlMetaQuery';
 import { MSSqlQueryModel } from './MSSqlQueryModel';
@@ -11,7 +12,7 @@ import { getIcon, getRAQBType, toRawSql } from './sqlUtil';
 import { MssqlOptions } from './types';
 
 export class MssqlDatasource extends SqlDatasource {
-  completionProvider: LanguageCompletionProvider | undefined = undefined;
+  sqlLanguageDefinition: LanguageDefinition | undefined = undefined;
   constructor(instanceSettings: DataSourceInstanceSettings<MssqlOptions>) {
     super(instanceSettings);
   }
@@ -48,16 +49,20 @@ export class MssqlDatasource extends SqlDatasource {
     return result;
   }
 
-  getSqlCompletionProvider(db: DB): LanguageCompletionProvider {
-    if (this.completionProvider !== undefined) {
-      return this.completionProvider;
+  getSqlLanguageDefinition(db: DB): LanguageDefinition {
+    if (this.sqlLanguageDefinition !== undefined) {
+      return this.sqlLanguageDefinition;
     }
     const args = {
       getColumns: { current: (query: SQLQuery) => fetchColumns(db, query) },
       getTables: { current: (dataset?: string) => fetchTables(db, dataset) },
     };
-    this.completionProvider = getSqlCompletionProvider(args);
-    return this.completionProvider;
+    this.sqlLanguageDefinition = {
+      id: 'sql',
+      completionProvider: getSqlCompletionProvider(args),
+      formatter: formatSQL,
+    };
+    return this.sqlLanguageDefinition;
   }
 
   getDB(): DB {
@@ -68,7 +73,7 @@ export class MssqlDatasource extends SqlDatasource {
       init: () => Promise.resolve(true),
       datasets: () => this.fetchDatasets(),
       tables: (dataset?: string) => this.fetchTables(dataset),
-      getSqlCompletionProvider: () => this.getSqlCompletionProvider(this.db),
+      getEditorLanguageDefinition: () => this.getSqlLanguageDefinition(this.db),
       fields: async (query: SQLQuery) => {
         if (!query?.dataset || !query?.table) {
           return [];
@@ -83,7 +88,7 @@ export class MssqlDatasource extends SqlDatasource {
       lookup: async (path?: string) => {
         if (!path) {
           const datasets = await this.fetchDatasets();
-          return datasets.map((d) => ({ name: d, completion: d }));
+          return datasets.map((d) => ({ name: d, completion: `${d}.` }));
         } else {
           const parts = path.split('.').filter((s: string) => s);
           if (parts.length > 2) {
@@ -97,7 +102,6 @@ export class MssqlDatasource extends SqlDatasource {
           }
         }
       },
-      functions: async () => AGGREGATE_FNS,
     };
   }
 }

public/app/plugins/datasource/mssql/sqlCompletionProvider.ts

@@ -1,18 +1,13 @@
-import { TableIdentifier } from '@grafana/experimental';
-import { AGGREGATE_FNS, OPERATORS } from 'app/features/plugins/sql/constants';
 import {
   ColumnDefinition,
-  CompletionItemKind,
-  CompletionItemPriority,
-  DB,
+  getStandardSQLCompletionProvider,
   LanguageCompletionProvider,
   LinkedToken,
-  SQLQuery,
-  StatementPlacementProvider,
-  SuggestionKindProvider,
   TableDefinition,
+  TableIdentifier,
   TokenType,
-} from 'app/features/plugins/sql/types';
+} from '@grafana/experimental';
+import { DB, SQLQuery } from 'app/features/plugins/sql/types';
 
 interface CompletionProviderGetterArgs {
   getColumns: React.MutableRefObject<(t: SQLQuery) => Promise<ColumnDefinition[]>>;
@@ -21,13 +16,17 @@ interface CompletionProviderGetterArgs {
 
 export const getSqlCompletionProvider: (args: CompletionProviderGetterArgs) => LanguageCompletionProvider =
   ({ getColumns, getTables }) =>
-  () => ({
-    triggerCharacters: ['.', ' ', '$', ',', '(', "'"],
+  (monaco, language) => ({
+    ...(language && getStandardSQLCompletionProvider(monaco, language)),
     tables: {
-      resolve: async () => {
-        return await getTables.current();
+      resolve: async (identifier) => {
+        return await getTables.current(identifier.table);
       },
       parseName: (token: LinkedToken) => {
+        if (!token) {
+          return { table: '' };
+        }
+
         let processedToken = token;
         let tablePath = processedToken.value;
 
@@ -36,6 +35,10 @@ export const getSqlCompletionProvider: (args: CompletionProviderGetterArgs) => L
           processedToken = processedToken.next;
         }
 
+        if (processedToken.value.endsWith('.')) {
+          tablePath = processedToken.value.slice(0, processedToken.value.length - 1);
+        }
+
         return { table: tablePath };
       },
     },
@@ -50,74 +53,8 @@ export const getSqlCompletionProvider: (args: CompletionProviderGetterArgs) => L
         return await getColumns.current({ table: `${schema}.${tableName}`, dataset: database, refId: 'A' });
       },
     },
-    supportedFunctions: () => AGGREGATE_FNS,
-    supportedOperators: () => OPERATORS,
-    customSuggestionKinds: customSuggestionKinds(getTables, getColumns),
-    customStatementPlacement,
   });
 
-export enum CustomStatementPlacement {
-  AfterDatabase = 'afterDatabase',
-}
-
-export enum CustomSuggestionKind {
-  TablesWithinDatabase = 'tablesWithinDatabase',
-}
-
-export const customStatementPlacement: StatementPlacementProvider = () => [
-  {
-    id: CustomStatementPlacement.AfterDatabase,
-    resolve: (currentToken, previousKeyword) => {
-      return Boolean(
-        currentToken?.is(TokenType.Delimiter, '.') ||
-          (currentToken?.is(TokenType.Whitespace) && currentToken?.previous?.is(TokenType.Delimiter, '.')) ||
-          (currentToken?.isNumber() && currentToken.value.endsWith('.'))
-      );
-    },
-  },
-];
-
-export const customSuggestionKinds: (
-  getTables: CompletionProviderGetterArgs['getTables'],
-  getFields: CompletionProviderGetterArgs['getColumns']
-) => SuggestionKindProvider = (getTables) => () =>
-  [
-    {
-      id: CustomSuggestionKind.TablesWithinDatabase,
-      applyTo: [CustomStatementPlacement.AfterDatabase],
-      suggestionsResolver: async (ctx) => {
-        const tablePath = ctx.currentToken ? getDatabaseName(ctx.currentToken) : '';
-        const t = await getTables.current(tablePath);
-
-        return t.map((table) => ({
-          label: table.name,
-          insertText: table.completion ?? table.name,
-          command: { id: 'editor.action.triggerSuggest', title: '' },
-          kind: CompletionItemKind.Field,
-          sortText: CompletionItemPriority.High,
-          range: {
-            ...ctx.range,
-            startColumn: ctx.range.endColumn,
-            endColumn: ctx.range.endColumn,
-          },
-        }));
-      },
-    },
-  ];
-
-export function getDatabaseName(token: LinkedToken) {
-  let processedToken = token;
-  let database = '';
-  while (processedToken?.previous && !processedToken.previous.isWhiteSpace()) {
-    processedToken = processedToken.previous;
-    database = processedToken.value + database;
-  }
-
-  database = database.trim();
-
-  return database;
-}
-
 export async function fetchColumns(db: DB, q: SQLQuery) {
   const cols = await db.fields(q);
   if (cols.length > 0) {
@@ -130,6 +67,6 @@ export async function fetchColumns(db: DB, q: SQLQuery) {
 }
 
 export async function fetchTables(db: DB, dataset?: string) {
-  const tables = await db.lookup(dataset);
-  return tables;
+  const tables = await db.lookup?.(dataset);
+  return tables || [];
 }

public/app/plugins/datasource/mysql/MySqlDatasource.ts

@@ -1,39 +1,41 @@
 import { DataSourceInstanceSettings, ScopedVars, TimeRange } from '@grafana/data';
+import { CompletionItemKind, LanguageDefinition, TableIdentifier } from '@grafana/experimental';
 import { TemplateSrv } from '@grafana/runtime';
 import { SqlDatasource } from 'app/features/plugins/sql/datasource/SqlDatasource';
-import { CompletionItemKind, DB, LanguageCompletionProvider, SQLQuery } from 'app/features/plugins/sql/types';
+import { DB, SQLQuery } from 'app/features/plugins/sql/types';
+import { formatSQL } from 'app/features/plugins/sql/utils/formatSQL';
 
 import MySQLQueryModel from './MySqlQueryModel';
 import { mapFieldsToTypes } from './fields';
 import { buildColumnQuery, buildTableQuery, showDatabases } from './mySqlMetaQuery';
-import { fetchColumns, fetchTables, getFunctions, getSqlCompletionProvider } from './sqlCompletionProvider';
+import { getSqlCompletionProvider } from './sqlCompletionProvider';
 import { MySQLOptions } from './types';
 
 export class MySqlDatasource extends SqlDatasource {
-  completionProvider: LanguageCompletionProvider | undefined;
+  sqlLanguageDefinition: LanguageDefinition | undefined;
 
   constructor(private instanceSettings: DataSourceInstanceSettings<MySQLOptions>) {
     super(instanceSettings);
-    this.completionProvider = undefined;
   }
 
   getQueryModel(target?: Partial<SQLQuery>, templateSrv?: TemplateSrv, scopedVars?: ScopedVars): MySQLQueryModel {
     return new MySQLQueryModel(target!, templateSrv, scopedVars);
   }
 
-  getSqlCompletionProvider(db: DB): LanguageCompletionProvider {
-    if (this.completionProvider !== undefined) {
-      return this.completionProvider;
+  getSqlLanguageDefinition(db: DB): LanguageDefinition {
+    if (this.sqlLanguageDefinition !== undefined) {
+      return this.sqlLanguageDefinition;
     }
 
     const args = {
-      getColumns: { current: (query: SQLQuery) => fetchColumns(db, query) },
-      getTables: { current: (dataset?: string) => fetchTables(db, { dataset }) },
-      fetchMeta: { current: (path?: string) => this.fetchMeta(path) },
-      getFunctions: { current: () => getFunctions() },
+      getMeta: { current: (identifier?: TableIdentifier) => this.fetchMeta(identifier) },
     };
-    this.completionProvider = getSqlCompletionProvider(args);
-    return this.completionProvider;
+    this.sqlLanguageDefinition = {
+      id: 'sql',
+      completionProvider: getSqlCompletionProvider(args),
+      formatter: formatSQL,
+    };
+    return this.sqlLanguageDefinition;
   }
 
   async fetchDatasets(): Promise<string[]> {
@@ -56,28 +58,20 @@ export class MySqlDatasource extends SqlDatasource {
     return mapFieldsToTypes(fields);
   }
 
-  async fetchMeta(path?: string) {
+  async fetchMeta(identifier?: TableIdentifier) {
     const defaultDB = this.instanceSettings.jsonData.database;
-    path = path?.trim();
-    if (!path && defaultDB) {
+    if (!identifier?.schema && defaultDB) {
       const tables = await this.fetchTables(defaultDB);
-      return tables.map((t) => ({ name: t, completion: t, kind: CompletionItemKind.Class }));
-    } else if (!path) {
+      return tables.map((t) => ({ name: t, completion: `${defaultDB}.${t}`, kind: CompletionItemKind.Class }));
+    } else if (!identifier?.schema && !defaultDB) {
       const datasets = await this.fetchDatasets();
       return datasets.map((d) => ({ name: d, completion: `${d}.`, kind: CompletionItemKind.Module }));
     } else {
-      const parts = path.split('.').filter((s: string) => s);
-      if (parts.length > 2) {
-        return [];
-      }
-      if (parts.length === 1 && !defaultDB) {
-        const tables = await this.fetchTables(parts[0]);
+      if (!identifier?.table && !defaultDB) {
+        const tables = await this.fetchTables(identifier?.schema);
         return tables.map((t) => ({ name: t, completion: t, kind: CompletionItemKind.Class }));
-      } else if (parts.length === 1 && defaultDB) {
-        const fields = await this.fetchFields({ dataset: defaultDB, table: parts[0] });
-        return fields.map((t) => ({ name: t.value, completion: t.value, kind: CompletionItemKind.Field }));
-      } else if (parts.length === 2 && !defaultDB) {
-        const fields = await this.fetchFields({ dataset: parts[0], table: parts[1] });
+      } else if (identifier?.table && identifier.schema) {
+        const fields = await this.fetchFields({ dataset: identifier.schema, table: identifier.table });
         return fields.map((t) => ({ name: t.value, completion: t.value, kind: CompletionItemKind.Field }));
       } else {
         return [];
@@ -96,9 +90,8 @@ export class MySqlDatasource extends SqlDatasource {
       validateQuery: (query: SQLQuery, range?: TimeRange) =>
         Promise.resolve({ query, error: '', isError: false, isValid: true }),
       dsID: () => this.id,
-      lookup: (path?: string) => this.fetchMeta(path),
-      getSqlCompletionProvider: () => this.getSqlCompletionProvider(this.db),
-      functions: async () => getFunctions(),
+      functions: () => ['VARIANCE', 'STDDEV'],
+      getEditorLanguageDefinition: () => this.getSqlLanguageDefinition(this.db),
     };
   }
 }

public/app/plugins/datasource/mysql/functions.ts

@@ -1,20 +0,0 @@
-export const FUNCTIONS = [
-  {
-    id: 'STDDEV',
-    name: 'STDDEV',
-    description: `STDDEV(
-    expression
-  )
-
-  Returns the standard deviation of non-NULL input values, or NaN if the input contains a NaN.`,
-  },
-  {
-    id: 'VARIANCE',
-    name: 'VARIANCE',
-    description: `VARIANCE(
-    expression
-  )
-
-  Returns the variance of non-NULL input values, or NaN if the input contains a NaN.`,
-  },
-];

public/app/plugins/datasource/mysql/sqlCompletionProvider.ts

@@ -1,284 +1,22 @@
-import { AGGREGATE_FNS, OPERATORS } from 'app/features/plugins/sql/constants';
 import {
-  Aggregate,
-  ColumnDefinition,
-  CompletionItemKind,
-  CompletionItemPriority,
-  DB,
+  getStandardSQLCompletionProvider,
   LanguageCompletionProvider,
-  LinkedToken,
-  MetaDefinition,
-  PositionContext,
-  SQLQuery,
-  StatementPlacementProvider,
-  StatementPosition,
-  SuggestionKindProvider,
   TableDefinition,
-  TokenType,
-} from 'app/features/plugins/sql/types';
-
-import { FUNCTIONS } from './functions';
+  TableIdentifier,
+} from '@grafana/experimental';
 
 interface CompletionProviderGetterArgs {
-  getColumns: React.MutableRefObject<(t: SQLQuery) => Promise<ColumnDefinition[]>>;
-  getTables: React.MutableRefObject<(d?: string) => Promise<TableDefinition[]>>;
-  fetchMeta: React.MutableRefObject<(d?: string) => Promise<MetaDefinition[]>>;
-  getFunctions: React.MutableRefObject<(d?: string) => Aggregate[]>;
+  getMeta: React.MutableRefObject<(t?: TableIdentifier) => Promise<TableDefinition[]>>;
 }
 
 export const getSqlCompletionProvider: (args: CompletionProviderGetterArgs) => LanguageCompletionProvider =
-  ({ getColumns, getTables, fetchMeta, getFunctions }) =>
-  () => ({
-    triggerCharacters: ['.', ' ', '$', ',', '(', "'"],
-    supportedFunctions: () => getFunctions.current(),
-    supportedOperators: () => OPERATORS,
-    customSuggestionKinds: customSuggestionKinds(getTables, getColumns, fetchMeta),
-    customStatementPlacement,
+  ({ getMeta }) =>
+  (monaco, language) => ({
+    ...(language && getStandardSQLCompletionProvider(monaco, language)),
+    tables: {
+      resolve: getMeta.current,
+    },
+    columns: {
+      resolve: getMeta.current,
+    },
   });
-
-export enum CustomStatementPlacement {
-  AfterDataset = 'afterDataset',
-  AfterFrom = 'afterFrom',
-  AfterSelect = 'afterSelect',
-}
-
-export enum CustomSuggestionKind {
-  TablesWithinDataset = 'tablesWithinDataset',
-}
-
-export enum Direction {
-  Next = 'next',
-  Previous = 'previous',
-}
-
-const TRIGGER_SUGGEST = 'editor.action.triggerSuggest';
-
-enum Keyword {
-  Select = 'SELECT',
-  Where = 'WHERE',
-  From = 'FROM',
-}
-
-export const customStatementPlacement: StatementPlacementProvider = () => [
-  {
-    id: CustomStatementPlacement.AfterDataset,
-    resolve: (currentToken, previousKeyword) => {
-      return Boolean(
-        currentToken?.is(TokenType.Delimiter, '.') ||
-          (currentToken?.is(TokenType.Whitespace) && currentToken?.previous?.is(TokenType.Delimiter, '.'))
-      );
-    },
-  },
-  {
-    id: CustomStatementPlacement.AfterFrom,
-    resolve: (currentToken, previousKeyword) => {
-      return Boolean(isAfterFrom(currentToken));
-    },
-  },
-  {
-    id: CustomStatementPlacement.AfterSelect,
-    resolve: (token, previousKeyword) => {
-      const is =
-        isDirectlyAfter(token, Keyword.Select) ||
-        (isAfterSelect(token) && token?.previous?.is(TokenType.Delimiter, ','));
-      return Boolean(is);
-    },
-  },
-];
-
-export const customSuggestionKinds: (
-  getTables: CompletionProviderGetterArgs['getTables'],
-  getFields: CompletionProviderGetterArgs['getColumns'],
-  fetchMeta: CompletionProviderGetterArgs['fetchMeta']
-) => SuggestionKindProvider = (getTables, _, fetchMeta) => () =>
-  [
-    {
-      id: CustomSuggestionKind.TablesWithinDataset,
-      applyTo: [CustomStatementPlacement.AfterDataset],
-      suggestionsResolver: async (ctx) => {
-        const tablePath = ctx.currentToken ? getTablePath(ctx.currentToken) : '';
-        const t = await getTables.current(tablePath);
-        return t.map((table) => suggestion(table.name, table.completion ?? table.name, CompletionItemKind.Field, ctx));
-      },
-    },
-    {
-      id: 'metaAfterSelect',
-      applyTo: [CustomStatementPlacement.AfterSelect],
-      suggestionsResolver: async (ctx) => {
-        const path = getPath(ctx.currentToken, Direction.Next);
-        const t = await fetchMeta.current(path);
-        return t.map((meta) => {
-          const completion = meta.kind === CompletionItemKind.Class ? `${meta.completion}.` : meta.completion;
-          return suggestion(meta.name, completion!, meta.kind, ctx);
-        });
-      },
-    },
-    {
-      id: 'metaAfterSelectFuncArg',
-      applyTo: [StatementPosition.AfterSelectFuncFirstArgument],
-      suggestionsResolver: async (ctx) => {
-        const path = getPath(ctx.currentToken, Direction.Next);
-        const t = await fetchMeta.current(path);
-        return t.map((meta) => {
-          const completion = meta.kind === CompletionItemKind.Class ? `${meta.completion}.` : meta.completion;
-          return suggestion(meta.name, completion!, meta.kind, ctx);
-        });
-      },
-    },
-    {
-      id: 'metaAfterFrom',
-      applyTo: [CustomStatementPlacement.AfterFrom],
-      suggestionsResolver: async (ctx) => {
-        // TODO: why is this triggering when isAfterFrom is false
-        if (!isAfterFrom(ctx.currentToken)) {
-          return [];
-        }
-        const path = ctx.currentToken?.value || '';
-        const t = await fetchMeta.current(path);
-        return t.map((meta) => suggestion(meta.name, meta.completion!, meta.kind, ctx));
-      },
-    },
-    {
-      id: `MYSQL${StatementPosition.WhereKeyword}`,
-      applyTo: [StatementPosition.WhereKeyword],
-      suggestionsResolver: async (ctx) => {
-        const path = getPath(ctx.currentToken, Direction.Previous);
-        const t = await fetchMeta.current(path);
-        return t.map((meta) => {
-          const completion = meta.kind === CompletionItemKind.Class ? `${meta.completion}.` : meta.completion;
-          return suggestion(meta.name, completion!, meta.kind, ctx);
-        });
-      },
-    },
-    {
-      id: StatementPosition.WhereComparisonOperator,
-      applyTo: [StatementPosition.WhereComparisonOperator],
-      suggestionsResolver: async (ctx) => {
-        if (!isAfterWhere(ctx.currentToken)) {
-          return [];
-        }
-        const path = getPath(ctx.currentToken, Direction.Previous);
-        const t = await fetchMeta.current(path);
-        const sugg = t.map((meta) => {
-          const completion = meta.kind === CompletionItemKind.Class ? `${meta.completion}.` : meta.completion;
-          return suggestion(meta.name, completion!, meta.kind, ctx);
-        });
-        return sugg;
-      },
-    },
-  ];
-
-function getPath(token: LinkedToken | null, direction: Direction) {
-  let path = token?.value || '';
-  const fromValue = keywordValue(token, Keyword.From, direction);
-  if (fromValue) {
-    path = fromValue;
-  }
-  return path;
-}
-
-export function getTablePath(token: LinkedToken) {
-  let processedToken = token;
-  let tablePath = '';
-  while (processedToken?.previous && !processedToken.previous.isWhiteSpace()) {
-    processedToken = processedToken.previous;
-    tablePath = processedToken.value + tablePath;
-  }
-
-  tablePath = tablePath.trim();
-  return tablePath;
-}
-
-function suggestion(label: string, completion: string, kind: CompletionItemKind, ctx: PositionContext) {
-  return {
-    label,
-    insertText: completion,
-    command: { id: TRIGGER_SUGGEST, title: '' },
-    kind,
-    sortText: CompletionItemPriority.High,
-    range: {
-      ...ctx.range,
-      startColumn: ctx.range.endColumn,
-      endColumn: ctx.range.endColumn,
-    },
-  };
-}
-
-function isAfterSelect(token: LinkedToken | null) {
-  return isAfterKeyword(token, Keyword.Select);
-}
-
-function isAfterFrom(token: LinkedToken | null) {
-  return isDirectlyAfter(token, Keyword.From);
-}
-
-function isAfterWhere(token: LinkedToken | null) {
-  return isAfterKeyword(token, Keyword.Where);
-}
-
-function isAfterKeyword(token: LinkedToken | null, keyword: string) {
-  if (!token?.is(TokenType.Keyword)) {
-    let curToken = token;
-    while (true) {
-      if (!curToken) {
-        return false;
-      }
-      if (curToken.is(TokenType.Keyword, keyword)) {
-        return true;
-      }
-      if (curToken.isKeyword()) {
-        return false;
-      }
-      curToken = curToken?.previous || null;
-    }
-  }
-  return false;
-}
-
-function isDirectlyAfter(token: LinkedToken | null, keyword: string) {
-  return token?.is(TokenType.Whitespace) && token?.previous?.is(TokenType.Keyword, keyword);
-}
-
-function keywordValue(token: LinkedToken | null, keyword: Keyword, direction: Direction) {
-  let next = token;
-  while (next) {
-    if (next.is(TokenType.Keyword, keyword)) {
-      return tokenValue(next);
-    }
-    next = next[direction];
-  }
-  return false;
-}
-
-function tokenValue(token: LinkedToken | null): string | undefined {
-  const ws = token?.next;
-  if (ws?.isWhiteSpace()) {
-    const v = ws.next;
-    const delim = v?.next;
-    if (!delim?.is(TokenType.Delimiter)) {
-      return v?.value;
-    }
-    return `${v?.value}${delim?.value}${delim.next?.value}`;
-  }
-  return undefined;
-}
-
-export async function fetchColumns(db: DB, q: SQLQuery) {
-  const cols = await db.fields(q);
-  if (cols.length > 0) {
-    return cols.map((c) => {
-      return { name: c.value, type: c.value, description: c.value };
-    });
-  } else {
-    return [];
-  }
-}
-
-export async function fetchTables(db: DB, q: Partial<SQLQuery>) {
-  const tables = await db.lookup(q.dataset);
-  return tables;
-}
-
-export function getFunctions(): Aggregate[] {
-  return [...AGGREGATE_FNS, ...FUNCTIONS];
-}

public/app/plugins/datasource/postgres/datasource.ts

@@ -1,11 +1,10 @@
 import { DataSourceInstanceSettings, ScopedVars } from '@grafana/data';
-import { AGGREGATE_FNS } from 'app/features/plugins/sql/constants';
+import { LanguageDefinition } from '@grafana/experimental';
 import { SqlDatasource } from 'app/features/plugins/sql/datasource/SqlDatasource';
-import { DB, LanguageCompletionProvider, SQLQuery, SQLSelectableValue } from 'app/features/plugins/sql/types';
+import { DB, SQLQuery, SQLSelectableValue } from 'app/features/plugins/sql/types';
+import { formatSQL } from 'app/features/plugins/sql/utils/formatSQL';
 import { TemplateSrv } from 'app/features/templating/template_srv';
 
-import { FUNCTIONS } from '../mysql/functions';
-
 import { PostgresQueryModel } from './PostgresQueryModel';
 import { getSchema, getTimescaleDBVersion, getVersion, showTables } from './postgresMetaQuery';
 import { fetchColumns, fetchTables, getSqlCompletionProvider } from './sqlCompletionProvider';
@@ -13,7 +12,7 @@ import { getFieldConfig, toRawSql } from './sqlUtil';
 import { PostgresOptions } from './types';
 
 export class PostgresDatasource extends SqlDatasource {
-  completionProvider: LanguageCompletionProvider | undefined = undefined;
+  sqlLanguageDefinition: LanguageDefinition | undefined = undefined;
 
   constructor(instanceSettings: DataSourceInstanceSettings<PostgresOptions>) {
     super(instanceSettings);
@@ -40,19 +39,21 @@ export class PostgresDatasource extends SqlDatasource {
     return tables.fields.table.values.toArray().flat();
   }
 
-  getSqlCompletionProvider(db: DB): LanguageCompletionProvider {
-    if (this.completionProvider !== undefined) {
-      return this.completionProvider;
+  getSqlLanguageDefinition(db: DB): LanguageDefinition {
+    if (this.sqlLanguageDefinition !== undefined) {
+      return this.sqlLanguageDefinition;
     }
 
     const args = {
       getColumns: { current: (query: SQLQuery) => fetchColumns(db, query) },
       getTables: { current: () => fetchTables(db) },
-      //TODO: Add aggregate functions
-      getFunctions: { current: () => [...AGGREGATE_FNS, ...FUNCTIONS] },
     };
-    this.completionProvider = getSqlCompletionProvider(args);
-    return this.completionProvider;
+    this.sqlLanguageDefinition = {
+      id: 'pgsql',
+      completionProvider: getSqlCompletionProvider(args),
+      formatter: formatSQL,
+    };
+    return this.sqlLanguageDefinition;
   }
 
   async fetchFields(query: SQLQuery): Promise<SQLSelectableValue[]> {
@@ -74,7 +75,7 @@ export class PostgresDatasource extends SqlDatasource {
       init: () => Promise.resolve(true),
       datasets: () => Promise.resolve([]),
       tables: () => this.fetchTables(),
-      getSqlCompletionProvider: () => this.getSqlCompletionProvider(this.db),
+      getEditorLanguageDefinition: () => this.getSqlLanguageDefinition(this.db),
       fields: async (query: SQLQuery) => {
         if (!query?.table) {
           return [];
@@ -89,7 +90,6 @@ export class PostgresDatasource extends SqlDatasource {
         const tables = await this.fetchTables();
         return tables.map((t) => ({ name: t, completion: t }));
       },
-      functions: async () => AGGREGATE_FNS,
     };
   }
 }

public/app/plugins/datasource/postgres/sqlCompletionProvider.ts

@@ -1,14 +1,11 @@
-import { TableIdentifier } from '@grafana/experimental';
-import { AGGREGATE_FNS, OPERATORS } from 'app/features/plugins/sql/constants';
 import {
   ColumnDefinition,
-  DB,
+  getStandardSQLCompletionProvider,
   LanguageCompletionProvider,
-  SQLQuery,
   TableDefinition,
-} from 'app/features/plugins/sql/types';
-
-import { FUNCTIONS } from '../mysql/functions';
+  TableIdentifier,
+} from '@grafana/experimental';
+import { DB, SQLQuery } from 'app/features/plugins/sql/types';
 
 interface CompletionProviderGetterArgs {
   getColumns: React.MutableRefObject<(t: SQLQuery) => Promise<ColumnDefinition[]>>;
@@ -17,8 +14,8 @@ interface CompletionProviderGetterArgs {
 
 export const getSqlCompletionProvider: (args: CompletionProviderGetterArgs) => LanguageCompletionProvider =
   ({ getColumns, getTables }) =>
-  () => ({
-    triggerCharacters: ['.', ' ', '$', ',', '(', "'"],
+  (monaco, language) => ({
+    ...(language && getStandardSQLCompletionProvider(monaco, language)),
     tables: {
       resolve: async () => {
         return await getTables.current();
@@ -29,8 +26,6 @@ export const getSqlCompletionProvider: (args: CompletionProviderGetterArgs) => L
         return await getColumns.current({ table: t?.table, refId: 'A' });
       },
     },
-    supportedFunctions: () => [...AGGREGATE_FNS, ...FUNCTIONS],
-    supportedOperators: () => OPERATORS,
   });
 
 export async function fetchColumns(db: DB, q: SQLQuery) {
@@ -45,6 +40,6 @@ export async function fetchColumns(db: DB, q: SQLQuery) {
 }
 
 export async function fetchTables(db: DB) {
-  const tables = await db.lookup();
-  return tables;
+  const tables = await db.lookup?.();
+  return tables || [];
 }