import { v4 as uuidv4 } from 'uuid'; import { DataSourceInstanceSettings, ScopedVars } from '@grafana/data'; import { LanguageDefinition } from '@grafana/plugin-ui'; import { TemplateSrv } from '@grafana/runtime'; import { COMMON_FNS, DB, FuncParameter, MACRO_FUNCTIONS, SQLQuery, SQLSelectableValue, SqlDatasource, formatSQL, } from '@grafana/sql'; import { getSchema, getSchemaAndName, showDatabases } from './MSSqlMetaQuery'; import { MSSqlQueryModel } from './MSSqlQueryModel'; import { fetchColumns, fetchTables, getSqlCompletionProvider } from './sqlCompletionProvider'; import { getIcon, getRAQBType, toRawSql } from './sqlUtil'; import { MssqlOptions } from './types'; export class MssqlDatasource extends SqlDatasource { sqlLanguageDefinition: LanguageDefinition | undefined = undefined; constructor(instanceSettings: DataSourceInstanceSettings<MssqlOptions>) { super(instanceSettings); } getQueryModel(target?: SQLQuery, templateSrv?: TemplateSrv, scopedVars?: ScopedVars): MSSqlQueryModel { return new MSSqlQueryModel(target, templateSrv, scopedVars); } async fetchDatasets(): Promise<string[]> { const datasets = await this.runSql<{ name: string[] }>(showDatabases(), { refId: 'datasets' }); return datasets.fields.name?.values.flat() ?? []; } async fetchTables(dataset?: string): Promise<string[]> { // We get back the table name with the schema as well. like dbo.table const tables = await this.runSql<{ schemaAndName: string[] }>(getSchemaAndName(dataset), { refId: 'tables' }); return tables.fields.schemaAndName?.values.flat() ?? []; } async fetchFields(query: SQLQuery): Promise<SQLSelectableValue[]> { if (!query.table) { return []; } const [_, table] = query.table.split('.'); const schema = await this.runSql<{ column: string; type: string }>(getSchema(query.dataset, table), { refId: `columns-${uuidv4()}`, }); const result: SQLSelectableValue[] = []; for (let i = 0; i < schema.length; i++) { const column = schema.fields.column.values[i]; const type = schema.fields.type.values[i]; result.push({ label: column, value: column, type, icon: getIcon(type), raqbFieldType: getRAQBType(type) }); } return result; } 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.sqlLanguageDefinition = { id: 'sql', completionProvider: getSqlCompletionProvider(args), formatter: formatSQL, }; return this.sqlLanguageDefinition; } getFunctions = (): ReturnType<DB['functions']> => { const columnParam: FuncParameter = { name: 'Column', required: true, options: (query) => this.fetchFields(query), }; return [...MACRO_FUNCTIONS(columnParam), ...COMMON_FNS.map((fn) => ({ ...fn, parameters: [columnParam] }))]; }; getDB(): DB { if (this.db !== undefined) { return this.db; } return { init: () => Promise.resolve(true), datasets: () => this.fetchDatasets(), tables: (dataset?: string) => this.fetchTables(dataset), getEditorLanguageDefinition: () => this.getSqlLanguageDefinition(this.db), fields: async (query: SQLQuery) => { if (!query?.dataset || !query?.table) { return []; } return this.fetchFields(query); }, validateQuery: (query) => Promise.resolve({ isError: false, isValid: true, query, error: '', rawSql: query.rawSql }), dsID: () => this.id, dispose: (_dsID?: string) => {}, toRawSql, functions: () => this.getFunctions(), lookup: async (path?: string) => { if (!path) { const datasets = await this.fetchDatasets(); return datasets.map((d) => ({ name: d, completion: `${d}.` })); } else { const parts = path.split('.').filter((s: string) => s); if (parts.length > 2) { return []; } if (parts.length === 1) { const tables = await this.fetchTables(parts[0]); return tables.map((t) => ({ name: t, completion: t })); } else { return []; } } }, }; } }