deepnote-toolkit/deepnote_toolkit/runtime_initialization.py at 2571c85e8b779d62fcc5de746556807fcddb4f36 · deepnote/deepnote-toolkit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
""" This module contains functions to set up the Deepnote runtime environment. """

import builtins

import IPython.core.page
import psycopg2.extensions
import psycopg2.extras

from deepnote_toolkit.runtime_patches import apply_runtime_patches

from .dataframe_utils import add_formatters
from .execute_post_start_hooks import execute_post_start_hooks
from .logging import LoggerManager
from .output_middleware import add_output_middleware
from .set_integrations_env import set_integration_env
from .set_notebook_path import set_notebook_path
from .sql.spark_sql_magic import SparkSql
from .sql.sql_utils import configure_sqlparse_limits


def init_deepnote_runtime():
    """
    This function initializes the runtime environment for Deepnote.
    """

    logger = LoggerManager().get_logger()

    logger.debug("Initializing Deepnote runtime environment started.")

    try:
        apply_runtime_patches()
    except Exception as e:
        logger.error("Failed to apply runtime patches with a error: %s", e)

    # Register sparksql magic
    try:
        IPython.get_ipython().register_magics(SparkSql)
    except Exception as e:
        logger.error("Failed to register sparksql magic with a error: %s", e)

    # Apply custom formatters for DataFrames
    try:
        logger.debug("Adding custom formatters for DataFrames.")
        add_formatters()
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to load custom dataFrame formatter with a error: %s", e)

    # Add output middleware
    try:
        logger.debug("Adding output middleware.")
        add_output_middleware()
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to add output middleware with a error: %s", e)

    # Disable sqlparse grouping limits for large analytical queries
    try:
        logger.debug("Configuring sqlparse limits.")
        configure_sqlparse_limits()
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to configure sqlparse limits with error: %s", e)

    # Set up psycopg2 to make long-running queries interruptible by SIGINT (interrupt kernel)
    try:
        logger.debug("Setting psycopg2.")
        psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to set psycopg2 with error: %s", e)

    # Set up IPython page printer
    try:
        logger.debug("Setting IPython page printer.")
        set_ipython_page_printer()
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to set IPython page printer with error: %s", e)

    # Set the notebook kernel's working directory
    try:
        logger.debug("Setting kernel working directory.")
        set_notebook_path()
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to set kernel working directory with error: %s", e)

    try:
        logger.debug("Setting integration environment.")
        set_integration_env()
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to set integration environment with error: %s", e)

    try:
        logger.debug("Executing post start hooks.")
        execute_post_start_hooks()
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to execute post start hooks with error: %s", e)

    logger.debug("Initializing Deepnote runtime environment finished.")


def set_ipython_page_printer():
    """
    This code changes how IPython handles the display of specific output, e.g. when running some ipython magic command such as `print?`.
    More info here: https://stackoverflow.com/questions/53498226/what-is-the-meaning-of-exclamation-and-question-marks-in-jupyter-notebook

    Instead of the default behavior (which in Jupyter opens a sidebar UI), it just prints the entire content directly to the console.
    """

    def page_printer(data, start=0, screen_lines=0, pager_cmd=None):
        if builtins.isinstance(data, builtins.dict):
            data = data["text/plain"]
        builtins.print(data)

    IPython.core.page.page = page_printer