dask-sql_jupyter.py

from IPython.core.magic import register_cell_magic, needs_local_scope
from sql_formatter.core import format_sql
import time, warnings

warnings.filterwarnings("ignore")

@register_cell_magic
@needs_local_scope
def sql(line, cell, local_ns):
    sql_statement = cell.format(**local_ns)
    t0 = time.time()
    res = c.sql(sql_statement)
    if "CREATE OR REPLACE" in sql_statement:
        table = sql_statement.split("CREATE OR REPLACE")[1]
        table = table.replace("TABLE", "").replace("VIEW", "").split()[0].strip()
        res = c.sql(f"SELECT * FROM {table} LIMIT 5").compute()
        print(f"Execution time: {time.time() - t0:.2f}s")
        return res
    else:
        res = res.compute()
        print(f"Execution time: {time.time() - t0:.2f}s")
        return res


# avoid name conflicts for automagic to work on line magics.
del sql


# formatting SQL into the next cell
example_sql = """
create or replace table mytable as -- mytable example
seLecT a.asdf, b.qwer, -- some comment here
c.asdf, -- some comment there
b.asdf2 frOm table1 as a leFt join 
table2 as b -- and here a comment
    on a.asdf = b.asdf  -- join this way
    inner join table3 as c
on a.asdf=c.asdf
whEre a.asdf= 1 -- comment this
anD b.qwer =2 and a.asdf<=1 --comment that
or b.qwer>=5
groUp by a.asdf
"""
get_ipython().set_next_input(format_sql(example_sql))