diff --git a/backend/kale/common/jputils.py b/backend/kale/common/jputils.py index a7df6cc76..d77a4a364 100644 --- a/backend/kale/common/jputils.py +++ b/backend/kale/common/jputils.py @@ -13,14 +13,19 @@ # limitations under the License. import os +import re import sys import time +import json import signal import logging import nbformat +import requests import threading +import ipykernel from queue import Empty +from notebook import notebookapp from jupyter_client.kernelspec import get_kernel_spec from kale.common.utils import remove_ansi_color_sequences from nbconvert.preprocessors.execute import ExecutePreprocessor @@ -309,3 +314,40 @@ def signal_handler(_signal, _frame): log.newline(lines=3) log.info("%s Successfully ran user code %s", "-" * 10, "-" * 10) return result + + +def get_notebook_path(): + """Returns the asb path of the Notebook or None if it cannot be determined. + + NOTE: works only when the security is token-based or there is no password + """ + log.info("Retrieving absolute path of the active notebook") + connection_file = os.path.basename(ipykernel.get_connection_file()) + try: + kernel_id = re.search(r"kernel-(.+?)\.json", connection_file).group(1) + except AttributeError: + log.error("Could not load kernel id.") + return None + + for srv in notebookapp.list_running_servers(): + sess_url = "%sapi/sessions" % srv["url"] + # No token and no password + if srv["token"] == "" and not srv['password']: + req = requests.get(sess_url) + else: + token_query_param = "?token=%s" % srv["token"] + req = requests.get(sess_url + token_query_param) + if req.status_code != 200: + log.error("Session request failed with status code %s" + % req.status_code) + return None + try: + sessions = json.loads(req.text) + except json.JSONDecodeError as e: + log.error("Could not parse session response: %s" % e) + return None + for sess in sessions: + if sess["kernel"]["id"] == kernel_id: + return os.path.join(srv["notebook_dir"], + sess["notebook"]["path"]) + return None diff --git a/backend/setup.py b/backend/setup.py index a2be4467c..d4b04aeac 100644 --- a/backend/setup.py +++ b/backend/setup.py @@ -46,6 +46,7 @@ 'jupyter-core >= 4.6.0', 'nbconvert >= 5.6.1, < 6.0.0', 'ipykernel >= 5.1.4', + 'notebook >= 6.0.0', 'packaging > 20', 'ml_metadata == 0.24.0', 'progress >= 1.5',