New Intearctive Experience

src/extension/content/export.ts

@@ -1,6 +1,5 @@
 import { commands, NotebookCell, NotebookCellKind, Uri, window, workspace } from 'vscode';
-import { isKustoNotebook } from '../kernel/provider';
-import { registerDisposable } from '../utils';
+import { isKustoNotebook, registerDisposable } from '../utils';
 import { getConnectionInfoFromDocumentMetadata } from '../kusto/connections/notebookConnection';
 import { updateCache } from '../cache';
 import { ICodeCell, IMarkdownCell, INotebookContent } from './jupyter';

src/extension/index.ts

@@ -18,6 +18,7 @@ import { AzureAuthenticatedConnection } from './kusto/connections/azAuth';
 import KustoClient from 'azure-kusto-data/source/client';
 import { registerConnection } from './kusto/connections/baseConnection';
 import { AppInsightsConnection } from './kusto/connections/appInsights';
+import { CellCodeLensProvider } from './interactive/cells';
 
 let client: LanguageClient;
 
@@ -44,6 +45,7 @@ export async function activate(context: ExtensionContext) {
     monitorJupyterCells();
     registerInteractiveExperience();
     registerExportCommand();
+    CellCodeLensProvider.regsiter();
 }
 
 export async function deactivate(): Promise<void> {

src/extension/interactive/cells.ts

@@ -0,0 +1,63 @@
+import {
+    CancellationToken,
+    CodeLens,
+    CodeLensProvider,
+    EventEmitter,
+    FoldingRange,
+    languages,
+    Range,
+    TextDocument
+} from 'vscode';
+import { FoldingRangesProvider } from '../languageServer';
+import { IDisposable } from '../types';
+import { disposeAllDisposables, isNotebookCell, registerDisposable } from '../utils';
+
+export class CellCodeLensProvider implements CodeLensProvider, IDisposable {
+    private readonly _onDidChangeCodeLenses = new EventEmitter<void>();
+    private readonly disposables: IDisposable[] = [];
+    constructor() {
+        FoldingRangesProvider.instance.onDidChange(() => this._onDidChangeCodeLenses.fire(), this, this.disposables);
+    }
+    public static regsiter() {
+        const provider = new CellCodeLensProvider();
+        registerDisposable(languages.registerCodeLensProvider({ language: 'kusto' }, provider));
+        registerDisposable(provider);
+    }
+    public dispose() {
+        disposeAllDisposables(this.disposables);
+    }
+    public async provideCodeLenses(document: TextDocument, _token: CancellationToken): Promise<CodeLens[]> {
+        if (isNotebookCell(document)) {
+            return [];
+        }
+        const ranges = await FoldingRangesProvider.instance.getRanges(document);
+        const codelenses: CodeLens[] = [];
+        ranges.forEach((item) => {
+            const index = getStartOfCode(document, item);
+            if (!index) {
+                return;
+            }
+            codelenses.push(
+                new CodeLens(new Range(index, 0, item.end, 0), {
+                    title: 'Run Query',
+                    command: 'kusto.executeSelectedQuery',
+                    arguments: [document, index, item.end]
+                })
+            );
+        });
+        return codelenses;
+    }
+}
+
+function isComment(document: TextDocument, index: number) {
+    const line = document.lineAt(index);
+    return line.isEmptyOrWhitespace || line.text.trim().startsWith('//');
+}
+
+function getStartOfCode(document: TextDocument, range: FoldingRange): number | undefined {
+    for (let index = range.start; index <= range.end; index++) {
+        if (!isComment(document, index)) {
+            return index;
+        }
+    }
+}

src/extension/kernel/interactive.ts

@@ -1,13 +1,89 @@
-import { TextEditor } from 'vscode';
+import {
+    NotebookCellData,
+    NotebookCellKind,
+    NotebookDocument,
+    NotebookRange,
+    Range,
+    TextDocument,
+    Uri,
+    ViewColumn,
+    workspace,
+    WorkspaceEdit
+} from 'vscode';
 import { commands } from 'vscode';
-import { isKustoFile, registerDisposable } from '../utils';
+import { registerDisposable } from '../utils';
+import { Kernel } from './provider';
 
 export function registerInteractiveExperience() {
-    registerDisposable(commands.registerTextEditorCommand('kusto.executeSelectedQuery', executeSelectedQuery));
+    registerDisposable(commands.registerCommand('kusto.executeSelectedQuery', executeSelectedQuery));
 }
 
-async function executeSelectedQuery(editor: TextEditor) {
-    if (!isKustoFile(editor.document)) {
+type INativeInteractiveWindow = { notebookUri: Uri; inputUri: Uri };
+const documentInteractiveDocuments = new WeakMap<TextDocument, Promise<NotebookDocument | undefined>>();
+async function executeSelectedQuery(document: TextDocument, start: number, end: number) {
+    if (!documentInteractiveDocuments.has(document)) {
+        documentInteractiveDocuments.set(document, getNotebookDocument());
+    }
+    const notebook = await documentInteractiveDocuments.get(document);
+    if (!notebook) {
         return;
     }
+    // Ensure its visible.
+    await commands.executeCommand('interactive.open', undefined, notebook.uri, undefined);
+    const cell = await createCell(notebook, document, start, end);
+    Kernel.instance.executeInteractive([cell], document, Kernel.instance.interactiveController);
+}
+
+async function getNotebookDocument() {
+    // eslint-disable-next-line prefer-const
+    let notebookUri: Uri | undefined;
+    let isSelected = false;
+    const selected = new Promise<void>((resolve) => {
+        const disposable = Kernel.instance.interactiveController.onDidChangeSelectedNotebooks(
+            ({ notebook, selected }) => {
+                if (!selected) {
+                    return;
+                }
+                if (!notebookUri) {
+                    notebookUri = notebook.uri;
+                    isSelected = true;
+                    return resolve();
+                }
+                if (notebook.uri.toString() !== notebookUri.toString()) {
+                    return;
+                }
+                isSelected = true;
+                resolve();
+            }
+        );
+        registerDisposable(disposable);
+    });
+    const info = (await commands.executeCommand(
+        'interactive.open',
+        ViewColumn.Beside,
+        undefined,
+        'kustoInteractive'
+    )) as INativeInteractiveWindow;
+    if (!isSelected) {
+        notebookUri = info.notebookUri;
+        await Promise.all([selected, commands.executeCommand('notebook.selectKernel')]);
+    }
+
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    return workspace.notebookDocuments.find((item) => item.uri.toString() === notebookUri!.toString());
+}
+async function createCell(notebook: NotebookDocument, document: TextDocument, start: number, end: number) {
+    const text = document.getText(new Range(document.lineAt(start).range.start, document.lineAt(end).range.end));
+    const edit = new WorkspaceEdit();
+    const cell = new NotebookCellData(NotebookCellKind.Code, text.trim(), 'kusto');
+    cell.metadata = {
+        interactiveWindowCellMarker: document.lineAt(start).text,
+        interactive: {
+            file: document.uri.fsPath,
+            line: start
+        }
+    };
+    edit.replaceNotebookCells(notebook.uri, new NotebookRange(notebook.cellCount, notebook.cellCount), [cell]);
+    await workspace.applyEdit(edit);
+    return notebook.cellAt(notebook.cellCount - 1);
 }