Support notebooks

.github/actions/common-build/action.yml

@@ -17,3 +17,6 @@ runs:
     - run: npm install
       shell: bash  # Explicitly define the shell
       working-directory: ./webview_panels
+    - run: npm run postinstall
+      shell: bash
+

.gitignore

@@ -6,4 +6,6 @@ dist
 .DS_Store
 .idea
 .history
-__pycache__/
+__pycache__/
+
+src/lib/notebooks/

.vscodeignore

@@ -3,9 +3,17 @@
 out/**
 src/**
 **/node_modules
+!dist/node_modules/**
 .gitignore
 vsc-extension-quickstart.md
 **/tsconfig.json
 **/.eslintrc.json
 **/*.map
 **/*.ts
+.github
+lineage_panel
+new_lineage_panel
+!new_lineage_panel/dist
+webview_panels
+!webview_panels/dist
+query_panel

altimate_notebook_kernel.py

@@ -0,0 +1,206 @@
+import json
+import re
+from datetime import  datetime
+import jupyter_client
+import queue
+
+class CustomDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(object_hook=self.object_hook, *args, **kwargs)
+    def object_hook(self, obj):
+        for key, value in obj.items():
+            if value == 'null':
+                obj[key] = None
+        return obj
+
+
+# Notebook kernel which will responsible for creating kernel executors for each notebook
+# should shut down kernel after notebook is closed
+# also store the cell outputs/data and use it for further executions
+# TODO: implement as required - extract to separate file
+
+
+# a method which will be called once a notebook is open, create an instance of this class
+# save doc uri as unique identifier for the notebook
+# initialize a jupyter kernel executor on notebook open and shutdown on close
+# store results of each cell execution and update after each execution
+# handle cell deletions
+# destroy this instance when notebook is closed
+
+class JupyterKernelExecutor:
+    def __init__(self):
+        self.kernel_manager = jupyter_client.KernelManager()
+        self.kernel_manager.start_kernel()
+        self.kernel_client = self.kernel_manager.client()
+        self.kernel_client.start_channels()
+        print('session pid', self.kernel_client.session.pid)
+        # print('ip', self.kernel_manager.connection_file.split('-')[1])
+        # Get the connection file
+        connection_file = self.kernel_manager.connection_file
+        # Load connection info
+        connection_info = jupyter_client.find_connection_file(connection_file)
+        with open(connection_info) as f:
+            connection_data = json.load(f)
+
+        # Extract WebSocket URL
+        websocket_url = f"ws://{connection_data['ip']}:{connection_data['shell_port']}/api/kernels/{self.kernel_manager.kernel_id}/channels"
+        print(connection_data, self.kernel_manager.kernel_name, self.kernel_manager.kernel_spec.name)
+
+    def execute(self, code, user_expressions=None):
+        self.kernel_client.wait_for_ready()
+
+        # print("Executing code:", code, user_expressions)
+        # Execute the code
+        self.kernel_client.execute(code, silent=False, store_history=True, user_expressions=user_expressions)
+
+        # Capture and return the output
+        output = []
+        start_time = datetime.now()
+        while True:
+            try:
+                msg = self.kernel_client.get_iopub_msg(timeout=1)
+                def datetime_converter(o):
+                    if isinstance(o, datetime):
+                        return o.__str__()
+
+                # Print formatted JSON
+                # print("msg", msg)
+
+                if msg['msg_type'] == 'stream':
+                    # for stdout
+                    output.append({'mime': 'text/plain', 'value': msg['content']['text']})
+                    break
+                elif msg['msg_type'] == 'comm_open':
+                    state = msg['content']['data']['state']
+                    state['model_id'] = msg['content']['comm_id']
+                    # Handle comm_open messages
+                    output.append({'mime': 'application/vnd.jupyter.widget-view+json', 'value': state})
+                elif msg['msg_type'] == 'execute_result' or msg['msg_type'] == 'display_data':
+                    # Flag to check if any key other than 'text/plain' exists
+                    other_keys_exist = False
+
+                    # Iterate over the keys in msg['content']['data']
+                    for key, value in msg['content']['data'].items():
+                        # Check if the key is not 'text/plain'
+                        if key != 'text/plain':
+                            # Append the dictionary to the output list
+                            output.append({'mime': key, 'value': value})
+                            other_keys_exist = True
+
+                    # If no other keys exist, add the 'text/plain' value
+                    if not other_keys_exist:
+                        output.append({'mime': 'text/plain', 'value': msg['content']['data']['text/plain']})
+                    break
+                elif msg['msg_type'] == 'error':
+                    output.append({'mime': 'text/plain', 'value': '\n'.join(msg['content']['traceback'])})
+                    break
+                elif msg['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
+                    break
+            except queue.Empty:
+                if (datetime.now() - start_time).total_seconds() > 30:  # Timeout after 30 seconds
+                    break
+        return output
+
+    def shutdown(self):
+        # Shutdown the kernel client and kernel manager
+        self.kernel_client.stop_channels()
+        self.kernel_manager.shutdown_kernel()
+        del self.kernel_client
+
+class AltimateNotebookKernel:
+    def __init__(self, doc_uri):
+        """
+        Initialize the AltimateNotebookKernel instance.
+
+        Parameters:
+        doc_uri (str): The unique identifier for the notebook.
+        """
+        self.doc_uri = doc_uri
+        self.kernel_executor = self.initialize_kernel_executor()
+        self.cell_results = {}
+
+    def get_session_id(self):
+        return self.kernel_executor.kernel_client.session.pid
+
+    def close_notebook(self):
+        """
+        Method to be called when the notebook is closed.
+        Shuts down the Jupyter kernel executor.
+        """
+        if self.kernel_executor:
+            self.shutdown_kernel_executor()
+            self.kernel_executor = None
+            print(f"Notebook {self.doc_uri} closed and kernel shut down.")
+
+    def get_connection_file(self):
+        return self.kernel_executor.kernel_manager.connection_file
+
+    def initialize_kernel_executor(self):
+        """
+        Initializes the Jupyter kernel executor.
+
+        Returns:
+        kernel_executor: The initialized kernel executor.
+        """
+        kernel_executor = JupyterKernelExecutor()
+        return kernel_executor
+
+    def shutdown_kernel_executor(self):
+        """
+        Shuts down the Jupyter kernel executor.
+        """
+        # Placeholder for actual kernel shutdown logic
+        print("Kernel executor shut down.")
+
+    def get_sql_result_by_cell(self, cell_id):
+        code = f"cell_{cell_id}"
+        return self.kernel_executor.execute(code)
+
+    def store_sql_result(self, cell_id, result):
+        """
+        Stores the result of a cell execution.
+
+        Parameters:
+        cell_id (str): The unique identifier for the cell.
+        result: The result of the cell execution.
+        """
+        resultJson =  json.loads(result)
+
+        # Construct the code to store the result in the Jupyter kernel
+        code = f"""
+        cell_{cell_id} = {resultJson}
+        """
+        self.execute_python(code)
+
+    def execute_python(self, code):
+        """
+        Executes a cell and stores the result.
+
+        Parameters:
+        cell_id (str): The unique identifier for the cell.
+        code (str): The code to be executed in the cell.
+
+        Returns:
+        result: The result of the cell execution.
+        """
+        response = self.kernel_executor.execute(code)
+        return response
+
+    def delete_cell(self, cell_id):
+        """
+        Handles cell deletion.
+
+        Parameters:
+        cell_id (str): The unique identifier for the cell to be deleted.
+        """
+        if cell_id in self.cell_results:
+            del self.cell_results[cell_id]
+            print(f"Cell {cell_id} deleted.")
+
+    def destroy_instance(self):
+        """
+        Destroys the instance when the notebook is closed.
+        """
+        self.close_notebook()
+        self.cell_results.clear()
+        print(f"Instance for notebook {self.doc_uri} destroyed.")