Transfer the data to the kernel

schema/metadata-form.json

@@ -19,6 +19,10 @@
                 "title": "-"
               }
             ]
+          },
+          "/sql-cell/variable": {
+            "title": "Variable",
+            "type": "string"
           }
         }
       },

src/index.ts

@@ -5,7 +5,11 @@ import {
   JupyterFrontEnd,
   JupyterFrontEndPlugin
 } from '@jupyterlab/application';
-import { ICommandPalette, IToolbarWidgetRegistry } from '@jupyterlab/apputils';
+import {
+  ICommandPalette,
+  ISessionContext,
+  IToolbarWidgetRegistry
+} from '@jupyterlab/apputils';
 import { IEditorServices } from '@jupyterlab/codeeditor';
 import { IDefaultFileBrowser } from '@jupyterlab/filebrowser';
 import { IMetadataFormProvider } from '@jupyterlab/metadataform';
@@ -14,7 +18,13 @@ import {
   NotebookActions,
   NotebookPanel
 } from '@jupyterlab/notebook';
-import { Contents, ContentsManager } from '@jupyterlab/services';
+import {
+  Contents,
+  ContentsManager,
+  Kernel,
+  KernelMessage,
+  Session
+} from '@jupyterlab/services';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 import { ITranslator, nullTranslator } from '@jupyterlab/translation';
 import {
@@ -28,6 +38,8 @@ import { FieldProps } from '@rjsf/utils';
 import { CustomContentFactory } from './cellfactory';
 import { requestAPI } from './handler';
 import { CommandIDs, SQL_MIMETYPE, SqlCell, objectEnum } from './common';
+import * as injectCode from './injectedCode';
+import { ISqlCellInjection, SqlCellInjection } from './kernelInjection';
 import { Databases } from './sidepanel';
 import { DatabaseSelect, SqlWidget, SqlSwitchWidget } from './widget';
 
@@ -43,11 +55,12 @@ const plugin: JupyterFrontEndPlugin<void> = {
   id: '@jupyter/sql-cell:plugin',
   description: 'Add the commands to the registry.',
   autoStart: true,
-  requires: [INotebookTracker],
+  requires: [INotebookTracker, ISqlCellInjection],
   optional: [ICommandPalette, IDefaultFileBrowser],
   activate: (
     app: JupyterFrontEnd,
     tracker: INotebookTracker,
+    injection: ISqlCellInjection,
     commandPalette: ICommandPalette,
     fileBrowser: IDefaultFileBrowser | null
   ) => {
@@ -73,15 +86,28 @@ const plugin: JupyterFrontEndPlugin<void> = {
           console.error('The database has not been set.');
         }
         const date = new Date();
+        const kernel = tracker.currentWidget?.sessionContext.session?.kernel;
         const source = activeCell?.model.sharedModel.getSource();
         requestAPI<any>('execute', {
           method: 'POST',
           body: JSON.stringify({ query: source, id: database_id })
         })
           .then(data => {
-            Private.saveData(path, data.data, date, fileBrowser)
-              .then(dataPath => console.log(`Data saved ${dataPath}`))
-              .catch(undefined);
+            const variable = SqlCell.getMetadata(activeCell.model, 'variable');
+            if (kernel && injection.status && variable) {
+              const future = Private.transferDataToKernel(
+                kernel,
+                data.data,
+                variable
+              );
+              future.done.then(reply => {
+                console.log('REPLY', reply);
+              });
+            } else {
+              Private.saveData(path, data.data, date, fileBrowser)
+                .then(dataPath => console.log(`Data saved ${dataPath}`))
+                .catch(undefined);
+            }
           })
           .catch(reason => {
             console.error(reason);
@@ -251,12 +277,72 @@ const metadataForm: JupyterFrontEndPlugin<void> = {
   }
 };
 
+/*
+ * A plugin to inject a function in the notebook kernel.
+ */
+const kernelFunctionInjector: JupyterFrontEndPlugin<ISqlCellInjection> = {
+  id: '@jupyter/sql-cell:kernel-injection',
+  description: 'A JupyterLab extension to inject a function in notebook kernel',
+  autoStart: true,
+  provides: ISqlCellInjection,
+  requires: [INotebookTracker],
+  activate: (app: JupyterFrontEnd, tracker: INotebookTracker) => {
+    let sessionContext: ISessionContext | undefined = undefined;
+    const injection = new SqlCellInjection();
+
+    /**
+     * Triggered when the current notebook or current kernel changes.
+     */
+    const onKernelChanged = async (
+      _sessionContext: ISessionContext,
+      kernelChange: Session.ISessionConnection.IKernelChangedArgs
+    ) => {
+      injection.status = false;
+      const kernel = kernelChange.newValue;
+      if (kernel) {
+        kernel.info.then(info => {
+          let code = '';
+          if (info.language_info.name === 'python') {
+            code = injectCode.PYTHON_CODE;
+          }
+          if (!code) {
+            return;
+          }
+          const content: KernelMessage.IExecuteRequestMsg['content'] = {
+            code: code
+          };
+          const future = kernel.requestExecute(content);
+          future.done.then(reply => {
+            injection.status = reply.content.status === 'ok';
+          });
+        });
+      }
+    };
+
+    tracker.currentChanged.connect((_, panel) => {
+      sessionContext?.kernelChanged.disconnect(onKernelChanged);
+      sessionContext = panel?.sessionContext;
+      const kernel = sessionContext?.session?.kernel;
+      if (sessionContext && kernel) {
+        onKernelChanged(sessionContext, {
+          name: 'kernel',
+          oldValue: null,
+          newValue: kernel
+        });
+      }
+      sessionContext?.kernelChanged.connect(onKernelChanged);
+    });
+
+    return injection;
+  }
+};
+
 /**
  * The notebook toolbar widget.
  */
 const notebookToolbarWidget: JupyterFrontEndPlugin<void> = {
   id: '@jupyter/sql-cell:notebook-toolbar',
-  description: 'A JupyterLab extension to add a widget in the notebook tools',
+  description: 'A JupyterLab extension to add a widget in the Notebook toolbar',
   autoStart: true,
   requires: [INotebookTracker, IToolbarWidgetRegistry],
   optional: [ISettingRegistry],
@@ -298,11 +384,37 @@ export default [
   cellFactory,
   databasesList,
   metadataForm,
+  kernelFunctionInjector,
   notebookToolbarWidget,
   plugin
 ];
 
 namespace Private {
+  /**
+   * Call the function to transfer the data on the kernel.
+   *
+   * @param kernel - kernel on which to transfer the data.
+   * @param data - data to transfer to the kernel (mst be serializable).
+   * @returns the code execution future.
+   */
+  export function transferDataToKernel(
+    kernel: Kernel.IKernelConnection,
+    data: any,
+    variable?: string
+  ): Kernel.IShellFuture<
+    KernelMessage.IExecuteRequestMsg,
+    KernelMessage.IExecuteReplyMsg
+  > {
+    data = JSON.stringify(data).replace(/"/gi, '\\"');
+    const variableStr = variable ? `, "${variable}"` : '';
+    const code = `_sql_transfer_data("${data}"${variableStr})`;
+    const content: KernelMessage.IExecuteRequestMsg['content'] = {
+      code: code,
+      stop_on_error: true
+    };
+    return kernel.requestExecute(content, false);
+  }
+
   /**
    * Save data in a CSV file.
    *

src/injectedCode.ts

@@ -0,0 +1,19 @@
+/**
+ * The default Python function injected in kernel.
+ */
+export const PYTHON_CODE = `
+def _sql_transfer_data(data, variable='sql_result'):
+    import json
+    try:
+        import pandas
+        PANDAS_AVAILABLE = True
+    except ModuleNotFoundError:
+        PANDAS_AVAILABLE = False
+
+    result = json.loads(data)
+    if PANDAS_AVAILABLE:
+      result = pandas.DataFrame.from_records(result)
+
+    globals()[variable] = result
+
+`;

src/kernelInjection.ts

@@ -0,0 +1,42 @@
+import { Token } from '@lumino/coreutils';
+import { ISignal, Signal } from '@lumino/signaling';
+
+/**
+ * The kernel injection token.
+ */
+export const ISqlCellInjection = new Token<ISqlCellInjection>(
+  '@jupyter/sql-cell:kernel-injection',
+  'A boolean, whether the function has been injected in the kernel or not'
+);
+
+/**
+ * The kernel injection status interface.
+ */
+export interface ISqlCellInjection {
+  /**
+   * Whether the current kernel has the function to populate data.
+   */
+  status: boolean;
+  /**
+   * A signal emitted when the status changes.
+   */
+  readonly statusChanged: ISignal<this, boolean>;
+}
+
+/**
+ * The kernel injection status class.
+ */
+export class SqlCellInjection implements ISqlCellInjection {
+  /**
+   * Getter and setter of the status.
+   */
+  get status(): boolean {
+    return this._status;
+  }
+  set status(value: boolean) {
+    this._status = value;
+    this.statusChanged.emit(value);
+  }
+  private _status = false;
+  readonly statusChanged = new Signal<this, boolean>(this);
+}