03-mysql.py

# -*- coding: utf-8 -*-
#
#  03-mysql.py
#

"""
Add MySQL magics. Needs MySQL credentials to be set in the environment
for this to work, and the "mysqlclient" package to be installed.

Add this file to the folder:

~/.ipython/profile_default/startup/

Then you can use SQL in a cell by using %%mysql in the first line. Some
options:

-d         Display the resulting data frame (default: do not display)
-o name    Save the frame as "<name>" (default: "df")
-i name    Set the column <name> as the index for the data frame
-t         Transpose the data frame
"""

from os import environ

JUPYTER_MYSQL_USER = environ.get("JUPYTER_MYSQL_USER", "root")
JUPYTER_MYSQL_DB = environ.get("JUPYTER_MYSQL_DB")
JUPYTER_MYSQL_HOST = environ.get("JUPYTER_MYSQL_HOST", "localhost")
JUPYTER_MYSQL_PASSWORD = environ.get("JUPYTER_MYSQL_PASSWORD", "")

try:
    import MySQLdb

    from IPython.core.magic import Magics, magics_class, cell_magic
    from IPython.core.magic_arguments import magic_arguments, argument, parse_argstring
    from IPython.core.interactiveshell import InteractiveShell

    import pandas as pd

    @magics_class
    class MySQLMagics(Magics):
        def __init__(self, shell: InteractiveShell):
            super().__init__(shell)
            self._conn = None

        @property
        def client(self):
            "Lazy load a MySQL client."
            if not self._conn:
                self._conn = MySQLdb.connect(
                    host=JUPYTER_MYSQL_HOST,
                    user=JUPYTER_MYSQL_USER,
                    passwd=JUPYTER_MYSQL_PASSWORD,
                    db=JUPYTER_MYSQL_DB,
                )

            return self._conn

        @magic_arguments()
        @argument("--output", "-o", help="Store result in a variable, defaults to df")
        @argument("--index", "-i", help="Set the given field as an index")
        @argument("--display", "-d", help="Display the output", action="store_true")
        @argument("--transpose", "-t", help="Transpose the output", action="store_true")
        @cell_magic
        def mysql(self, line, cell):
            args = parse_argstring(self.mysql, line)

            lower = cell.strip().lower()
            if any(lower.startswith(s) for s in ["select", "desc"]):
                # it's a query! return a data frame
                return self.execute_pandas(cell, args)

            else:
                # it's something else (insert, update, create, alter, etc.)
                return self.execute_raw(cell)

        def execute_pandas(self, cell, args):
            df = pd.read_sql(cell, self.client)

            if args.index:
                df.set_index(args.index, inplace=True)

            if args.transpose:
                df = df.T

            if args.output is not None:
                self.shell.push({args.output: df})
            else:
                self.shell.push({"df": df})

            if args.display:
                display(df)

        def execute_raw(self, cell):
            cursor = self.client.cursor()
            cursor.execute(cell)

    ip = get_ipython()
    ip.register_magics(MySQLMagics)

except ImportError:
    raise