added support for jupyter notebook in VSCode

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": [

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<string, string>();
+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<CompletionList>(
+            'vscode.executeCompletionItemProvider',
+            vdocUri,
+            position,
+            context.triggerCharacter
+          );
+        }
+        else {
+          // console.log("Sending .sql file contents to LSP server")
+          return await next(document, position, context, token);
+        }
+      }
     }
   }
 

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",