diff --git a/package.json b/package.json index e6b1d12a..98e0f2d7 100644 --- a/package.json +++ b/package.json @@ -43,7 +43,10 @@ "vscode": "^1.45.1" }, "activationEvents": [ - "onLanguage:sql" + "onLanguage:sql", + "onLanguage:python", + "onNotebook:kusto-notebook", + "onNotebook:jupyter-notebook" ], "contributes": { "commands": [ diff --git a/packages/client/extension.ts b/packages/client/extension.ts index 892ccbd7..3e3c27c2 100644 --- a/packages/client/extension.ts +++ b/packages/client/extension.ts @@ -1,5 +1,16 @@ import * as path from 'path' -import { ExtensionContext, commands, window as Window } from 'vscode' +import { + languages, + workspace, + ExtensionContext, + commands, + window as Window, + TextDocument, + NotebookDocument, + NotebookCellKind, + Uri, + CompletionList +} from 'vscode' import { LanguageClient, LanguageClientOptions, @@ -9,7 +20,100 @@ import { import { ExecuteCommandParams } from 'vscode-languageserver-protocol' import { rebuild } from './rebuild' +const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; +const SQL = 'sql' +const FILE_EXTENSION = '.sql' +const EMBED_SCHEME = 'sql-language-server' +const MAGICS_DETECTED = [ + '%%sql', + '%sql', + '%%sparksql', + '%sparksql', + '%%trino', + '%trino' +] + +const SELECTORS = [ + { language: SQL, scheme: EMBED_SCHEME }, + { language: SQL, scheme: NOTEBOOK_CELL_SCHEME }, + { language: SQL, scheme: 'file', pattern: `**/*${FILE_EXTENSION}`} +] + +function isSqlMagic(text: string): boolean { + return MAGICS_DETECTED.some(magic => text.startsWith(magic)); +} + +function commentSqlCellMagic(text: string): string { + return text.replace('%%', '--'); +} + +export type IDisposable = { + dispose: () => void; +}; + +const disposables: IDisposable[] = []; + +export function registerDisposable(disposable: IDisposable) { + disposables.push(disposable); +} + +export function monitorJupyterCells() { + registerDisposable(workspace.onDidOpenNotebookDocument(updateSqlCellsOfDocument)); + registerDisposable(workspace.onDidChangeTextDocument((e) => updateSqlCells(e.document))); + workspace.notebookDocuments.forEach(updateSqlCellsOfDocument); +} + +function isJupyterNotebook(document?: NotebookDocument) { + return document?.notebookType === 'jupyter-notebook'; +} + +async function updateSqlCellsOfDocument(document?: NotebookDocument) { + if (!document || !isJupyterNotebook(document)) { + return; + } + await Promise.all( + document + .getCells() + .filter((item) => item.kind === NotebookCellKind.Code) + .map((item) => updateSqlCells(item.document)) + ); +} + +async function updateSqlCells(textDocument: TextDocument) { + const notebookDocument = getNotebookDocument(textDocument); + if (!notebookDocument || !isJupyterNotebook(notebookDocument)) { + return; + } + if (textDocument.languageId !== 'python') { + return; + } + if (!isSqlMagic(textDocument.lineAt(0).text)) { + return; + } + await languages.setTextDocumentLanguage(textDocument, SQL); +} + + +const virtualDocumentContents = new Map(); +function provideTextDocumentContent(uri: Uri): string { + let lookupUri = uri.path.slice(1) // remove front slash prefix + lookupUri = lookupUri.slice(0, -FILE_EXTENSION.length); // remove file extension + const text = virtualDocumentContents.get(lookupUri); + // console.log(`lookup vcd ${lookupUri} : ${text}`); + return text; +} + +export function getNotebookDocument(document: TextDocument | NotebookDocument): NotebookDocument | undefined { + return workspace.notebookDocuments.find((item) => item.uri.path === document.uri.path); +} + export function activate(context: ExtensionContext) { + // console.log("sql-language-server extension activated") + monitorJupyterCells() + workspace.registerTextDocumentContentProvider(EMBED_SCHEME, { + provideTextDocumentContent: uri => provideTextDocumentContent(uri) + }); + // Using the location of the javacript file built by `npm run prepublish` let serverModule = context.asAbsolutePath(path.join('packages', 'server', 'dist', 'cli.js')) let execArgs = ['up', '--method', 'node-ipc'] @@ -23,11 +127,42 @@ export function activate(context: ExtensionContext) { } let clientOptions: LanguageClientOptions = { - documentSelector: [{scheme: 'file', language: 'sql', pattern: '**/*.sql'}], + documentSelector: SELECTORS, diagnosticCollectionName: 'sqlLanguageServer', synchronize: { configurationSection: 'sqlLanguageServer', // fileEvents: workspace.createFileSystemWatcher('**/.sqllsrc.json') + }, + middleware: { + provideCompletionItem: async (document, position, context, token, next) => { + let originalUri = document.uri.toString(); + if ( originalUri.startsWith(EMBED_SCHEME)) { + // console.log("Sending modified cell magic text to LSP server") + return await next(document, position, context, token); + } + else if (isSqlMagic(document.getText())) { + // console.log("Handling a cell containing sql magic") + const text = commentSqlCellMagic(document.getText()); + // console.log(`set vdc content ${originalUri} : ${text}`) + virtualDocumentContents.set(originalUri, text); + const encodedUri = encodeURIComponent(originalUri); + const vdocUriString = `${EMBED_SCHEME}://sql/${encodedUri}${FILE_EXTENSION}`; + const vdocUri = Uri.parse(vdocUriString); + // Invoke completion, this will call us back again + // but with a virutal document + // with a properly commented out magic + return await commands.executeCommand( + 'vscode.executeCompletionItemProvider', + vdocUri, + position, + context.triggerCharacter + ); + } + else { + // console.log("Sending .sql file contents to LSP server") + return await next(document, position, context, token); + } + } } } diff --git a/packages/client/package.json b/packages/client/package.json index ce76aa5e..54fb7dd6 100644 --- a/packages/client/package.json +++ b/packages/client/package.json @@ -14,7 +14,9 @@ "vscode": "^1.52.0" }, "activationEvents": [ - "onLanguage:sql" + "onLanguage:sql", + "onLanguage:python", + "onNotebook:jupyter-notebook" ], "dependencies": { "electron-rebuild": "^1.11.0",