Add ability to dispatch SPL2 via SPL2 Notebooks

out/extension.js

@@ -14,10 +14,11 @@ const splunkCustomRESTHandler = require('./customRESTHandler.js')
 const splunkSpec = require("./spec.js");
 const reload = require("./commands/reload.js");
 
-const notebookSerializers = require('./notebooks/serializers');
-const notebookControllers = require('./notebooks/controller');
+const { SplunkNotebookSerializer } = require('./notebooks/serializers');
+const { SplunkController } = require('./notebooks/controller');
+const { Spl2Controller } = require('./notebooks/spl2/controller');
 const notebookCommands = require('./notebooks/commands');
-const notebookProviders = require('./notebooks/provider');
+const { CellResultCountStatusBarProvider } = require('./notebooks/provider');
 
 //const { transpileModule } = require("typescript");
 //const { AsyncLocalStorage } = require("async_hooks");
@@ -226,11 +227,15 @@ function activate(context) {
     }));
 
     // Notebook
-    context.subscriptions.push(vscode.workspace.registerNotebookSerializer('splunk-notebook', new notebookSerializers.SplunkNotebookSerializer(), {transientCellMetadata: {inputCollapsed: true, outputCollapsed: true}, transientOutputs: false}));
-	let controller = new notebookControllers.SplunkController()
-    context.subscriptions.push(controller)
-    context.subscriptions.push(vscode.notebooks.registerNotebookCellStatusBarItemProvider('splunk-notebook', new notebookProviders.CellResultCountStatusBarProvider(splunkOutputChannel)));
-    notebookCommands.registerNotebookCommands(controller, splunkOutputChannel, context)
+    context.subscriptions.push(vscode.workspace.registerNotebookSerializer('splunk-notebook', new SplunkNotebookSerializer(), {transientCellMetadata: {inputCollapsed: true, outputCollapsed: true}, transientOutputs: false}));
+	context.subscriptions.push(vscode.workspace.registerNotebookSerializer('spl2-notebook', new SplunkNotebookSerializer(), {transientCellMetadata: {inputCollapsed: true, outputCollapsed: true}, transientOutputs: false}));
+    const controller = new SplunkController();
+    context.subscriptions.push(controller);
+    const spl2Controller = new Spl2Controller();
+    context.subscriptions.push(spl2Controller);
+    context.subscriptions.push(vscode.notebooks.registerNotebookCellStatusBarItemProvider('splunk-notebook', new CellResultCountStatusBarProvider(splunkOutputChannel)));
+    context.subscriptions.push(vscode.notebooks.registerNotebookCellStatusBarItemProvider('spl2-notebook', new CellResultCountStatusBarProvider(splunkOutputChannel)));
+    notebookCommands.registerNotebookCommands([controller, spl2Controller], splunkOutputChannel, context);
 }
 exports.activate = activate;
 

out/notebooks/commands.ts

@@ -3,9 +3,9 @@ import { SplunkController } from './controller';
 import { VIZ_TYPES } from './visualizations';
 import { getClient, getJobSearchLog, getSearchJobBySid } from './splunk';
 
-export async function registerNotebookCommands(controller: SplunkController, outputChannel: vscode.OutputChannel, context: vscode.ExtensionContext) {
+export async function registerNotebookCommands(controllers: SplunkController[], outputChannel: vscode.OutputChannel, context: vscode.ExtensionContext) {
     context.subscriptions.push(vscode.commands.registerCommand('splunk.notebooks.addVisualizationPreference', (cell) => { 
-		addVisualizationPreference(controller, cell)
+		controllers.forEach((controller) => addVisualizationPreference(controller, cell));
 	}))
 
 	context.subscriptions.push(vscode.commands.registerCommand('splunk.notebooks.openJobInspector', (sid) => { 

out/notebooks/controller.ts

@@ -11,19 +11,17 @@ import {
 import { splunkMessagesToOutputItems } from './utils';
 
 export class SplunkController {
-    readonly controllerId = 'splunk-notebook-controller';
-    readonly notebookType = 'splunk-notebook';
-    readonly label = 'SPL Note';
-    readonly supportedLanguages = ['markdown', 'splunk_search', 'splunk-spl-meta'];
+    protected controllerId: string;
+    protected notebookType: string;
+    protected label: string;
+    protected supportedLanguages: string[];
 
-    private readonly _controller: vscode.NotebookController;
+    protected _controller: vscode.NotebookController;
     private _executionOrder = 0;
     private _interrupted = false;
     private _tokens = {};
     private _lastjob = undefined;
 
-    readonly rendererScriptId = 'splunk-visualization-script';
-
     readonly _spl_meta_help = `
     SPL-META Manual
 
@@ -40,7 +38,17 @@ export class SplunkController {
     _lastjob: Contains the search id (sid) of the last query
     `;
 
-    constructor() {
+    constructor(
+            controllerId = 'splunk-notebook-controller',
+            notebookType = 'splunk-notebook',
+            label = 'SPL Note',
+            supportedLanguages = ['markdown', 'splunk_search', 'splunk-spl-meta']
+        ) {
+        this.controllerId = controllerId;
+        this.notebookType = notebookType;
+        this.label = label;
+        this.supportedLanguages = supportedLanguages;
+
         this._controller = vscode.notebooks.createNotebookController(
             this.controllerId,
             this.notebookType,
@@ -61,7 +69,7 @@ export class SplunkController {
         this._execute([cell], notebookDocument, this._controller);
     }
 
-    private _execute(
+    protected _execute(
         cells: vscode.NotebookCell[],
         _notebook: vscode.NotebookDocument,
         _controller: vscode.NotebookController
@@ -122,22 +130,24 @@ export class SplunkController {
         execution.end(true, Date.now());
     }
 
-    private async _doExecution(cell: vscode.NotebookCell): Promise<void> {
+    protected _startExecution(cell: vscode.NotebookCell): vscode.NotebookCellExecution {
         this._interrupted = false;
         console.log(cell);
         const execution = this._controller.createNotebookCellExecution(cell);
         execution.executionOrder = ++this._executionOrder;
         execution.start(Date.now());
+        return execution;
+    }
+
+    private async _doExecution(cell: vscode.NotebookCell): Promise<void> {
+        const execution = this._startExecution(cell);
 
         let query = cell.document.getText().trim().replace(/^\s+|\s+$/g, '');
 
         const service = getClient()
 
         let jobs = service.jobs();
 
-        let activeThemeKind = vscode.window.activeColorTheme.kind;
-        let backgroundColor = new vscode.ThemeColor('notebook.editorBackground');
-
         const tokenRegex = /\$([a-zA-Z0-9_.|]*?)\$/g;
 
         var controller = this;
@@ -173,7 +183,10 @@ export class SplunkController {
             execution.end(false, Date.now());
             return;
         }
-
+        await this._finishExecution(job, cell, execution);
+    }
+
+    protected async _finishExecution(job: any, cell: vscode.NotebookCell, execution: vscode.NotebookCellExecution) {
         let sid = job['sid'];
         this._lastjob = sid;
         this._tokens['_lastjob'] = this._lastjob;
@@ -199,7 +212,8 @@ export class SplunkController {
             if (job.properties().isDone == true) {
                 jobComplete = true;
                 continue;
-            } 
+            }
+            execution.replaceOutput([new vscode.NotebookCellOutput([], { job: job.properties() })]);
             wait(1000);
         }        
 
@@ -214,6 +228,8 @@ export class SplunkController {
 
         if (!this._interrupted) {
             let results: any = await getSearchJobResults(job);
+            let activeThemeKind = vscode.window.activeColorTheme.kind;
+            let backgroundColor = new vscode.ThemeColor('notebook.editorBackground');    
 
             execution.replaceOutput([
                 new vscode.NotebookCellOutput(

out/notebooks/provider.ts

@@ -15,7 +15,7 @@ export class CellResultCountStatusBarProvider implements vscode.NotebookCellStat
 
         const items: vscode.NotebookCellStatusBarItem[] = [];
 
-        if (cell.document.languageId !== "splunk_search") {
+        if (cell.document.languageId !== "splunk_search" && cell.document.languageId !== "splunk_spl2") {
             return items
         }
 

out/notebooks/spl2/controller.ts

@@ -0,0 +1,51 @@
+import * as vscode from 'vscode';
+
+import {
+    dispatchSpl2Module,
+    getClient,
+} from '../splunk';
+import { SplunkController } from '../controller';
+import { splunkMessagesToOutputItems } from '../utils';
+
+export class Spl2Controller extends SplunkController {
+    constructor() {
+        super('spl2-notebook-controller', 'spl2-notebook', 'SPL2 Note', ['splunk_spl2']);
+        this._controller.executeHandler = this._execute.bind(this);
+    }
+
+    protected _execute(
+        cells: vscode.NotebookCell[],
+        _notebook: vscode.NotebookDocument,
+        _controller: vscode.NotebookController
+    ): void {
+        for (let cell of cells) {
+            if (cell.document.languageId == 'splunk_spl2') {
+                this._doSpl2Execution(cell);
+            }
+        }
+    }
+
+    private async _doSpl2Execution(cell: vscode.NotebookCell): Promise<void> {
+        const execution = super._startExecution(cell);
+
+        const spl2Module = cell.document.getText().trim();
+        const service = getClient();
+
+        let job;
+        try {
+            job = await dispatchSpl2Module(service, spl2Module);
+            await super._finishExecution(job, cell, execution);
+        } catch (failedResponse) {
+            let outputItems: vscode.NotebookCellOutputItem[] = [];
+            if (!failedResponse.data || !failedResponse.data.messages) {
+                outputItems = [vscode.NotebookCellOutputItem.error(failedResponse)];
+            } else {
+                const messages = failedResponse.data.messages;
+                outputItems = splunkMessagesToOutputItems(messages);
+            }
+
+            execution.replaceOutput([new vscode.NotebookCellOutput(outputItems)]);
+            execution.end(false, Date.now());
+        }
+    }
+}