IBMi.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IBMi Jupyter Notebook Extension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import pyodbc, getpass, time, matplotlib\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plot\n",
    "from IPython.display import display, HTML, Image, Javascript\n",
    "from IPython.core.magic import register_cell_magic, needs_local_scope\n",
    "\n",
    "global _config\n",
    "_config = {'host': '', 'user': '', 'pwd': ''}\n",
    "\n",
    "# disable pandas df display limit, trust user to limit result set in query\n",
    "pd.set_option('display.max_rows', None)\n",
    "pd.set_option('display.max_columns', None)\n",
    "pd.set_option('display.width', None)\n",
    "pd.set_option('display.max_colwidth', None)\n",
    "\n",
    "# TODO: port? protocol? driver?\n",
    "_config['host'] = input('Enter host: ').strip()\n",
    "_config['user'] = input('Enter user: ').strip()\n",
    "_config['pwd'] = getpass.getpass('Enter password: ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Establish Connection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Connecting to {} as {} ...\".format(_config['host'], _config['user']))\n",
    "try:\n",
    "    conn = pyodbc.connect(driver='{IBM i Access ODBC Driver}', \n",
    "      system=_config['host'], uid=_config['user'], pwd=_config['pwd'])\n",
    "    print('Successfully connected!')\n",
    "except pyodbc.InterfaceError as e:\n",
    "    print('Could not connect :(\\n{}'.format(e))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Splitting cell content into SQL statements delimited by ';' (when not in string)\n",
    "def get_statements(sql):\n",
    "    stmts = []\n",
    "    i = 0\n",
    "    stmt = ''\n",
    "    q = ''\n",
    "    for c in sql:\n",
    "        if c in ('\\'', '\"'):\n",
    "            stmt += c\n",
    "            q = '' if c == q else (c if q == '' else q)\n",
    "        elif c == ';':\n",
    "            stmts.append(stmt)\n",
    "            stmt = ''\n",
    "        else:\n",
    "            stmt += c\n",
    "    if stmt != '':\n",
    "        stmts.append(stmt)\n",
    "    return stmts\n",
    "\n",
    "\n",
    "# flag validation and sanitizing\n",
    "def get_flag(line):\n",
    "    flag = line.lower().strip()\n",
    "    if   len(flag) == 0:                return ''\n",
    "    elif flag[0] != '-':                raise Exception('Invalid flag passed {}'.format(flag))\n",
    "    elif flag in ['-c', '-csv']:        return '-c'\n",
    "    elif flag in ['-h', '-html']:       return '-h'\n",
    "    elif flag in ['-j', '-json']:       return '-j'\n",
    "    elif flag in ['-p', '-prep']:       return '-p'\n",
    "    elif flag in ['-l', '-tex']:        return '-l'\n",
    "    elif flag in ['-s', '-str']:        return '-s'\n",
    "    elif flag in ['-x','-xml']:         return '-x'\n",
    "    elif flag in ['-plb', '-plotbar']:  return '-plb'\n",
    "    elif flag in ['-plp', '-plotpie']:  return '-plp'\n",
    "    elif flag in ['-pll', '-plotline']: return '-pll'\n",
    "    raise Exception('Invalid flag passed {}'.format(flag))\n",
    "\n",
    "\n",
    "# convert dataframe to XML ... modified from\n",
    "# https://stackoverflow.com/questions/18574108/how-do-convert-a-pandas-dataframe-to-xml\n",
    "def df_to_xml(df):\n",
    "    def row_to_xml(row):\n",
    "        xml = ['<row>']\n",
    "        for i, col_name in enumerate(row.index):\n",
    "            xml.append('<col name=\"{0}\">{1}</col>'.format(col_name, row.iloc[i]))\n",
    "        xml.append('</row>')\n",
    "        return ''.join(xml)\n",
    "    return '<rows>'+(''.join(df.apply(row_to_xml, axis=1)))+'</rows>'\n",
    "\n",
    "\n",
    "# Handle outputting result set dependent on flag passed\n",
    "def output_rs(rows, cursor, flag):\n",
    "    cols = [col[0] for col in cursor.description]\n",
    "    df = pd.DataFrame((tuple(row) for row in rows), columns=cols)\n",
    "    if   flag == '':     return df\n",
    "    elif flag == '-c':   return df.to_csv(index=False)\n",
    "    elif flag == '-h':   return df.to_html()\n",
    "    elif flag == '-j':   return df.to_json(orient='records')\n",
    "    elif flag == '-l':   return df.to_latex(index=False)\n",
    "    elif flag == '-s':   return df.to_string(index=False)\n",
    "    elif flag == '-x':   return df_to_xml(df)\n",
    "    elif flag == '-plb': return plot_bar(df, len(df.columns))\n",
    "    elif flag == '-plp': return plot_pie(df, len(df.columns))\n",
    "    elif flag == '-pll': return plot_line(df, len(df.columns))\n",
    "    raise Exception('Invalid flag passed {}'.format(flag))\n",
    "\n",
    "\n",
    "# plot bar graph from dataframe\n",
    "def plot_bar(df, col_count):\n",
    "    if col_count == 1:\n",
    "        df.index = df.index + 1\n",
    "        df.plot(kind='bar')\n",
    "        plot.plot()\n",
    "    elif col_count == 2:\n",
    "        df.plot(kind='bar', x=df.columns.values[0], y=df.columns.values[1])\n",
    "        plot.plot()\n",
    "    elif col_count == 3:\n",
    "        pivoted = pandas.pivot_table(df, \n",
    "                            values=df.columns.values[2], \n",
    "                            columns=df.columns.values[0], \n",
    "                            index=df.columns.values[1])\n",
    "        pivoted.plot.bar()\n",
    "    else:\n",
    "        raise Exception('cannot generate bar plot with > 3 column(s)')\n",
    "\n",
    "\n",
    "# plot pie graph from dataframe        \n",
    "def plot_pie(df, col_count):\n",
    "    if col_count == 1:\n",
    "        df.index = df.index + 1\n",
    "        pie = df.plot(kind='pie', y=df.columns.values[0], autopct='%.2f')\n",
    "    elif col_count == 2:\n",
    "        pie = df.plot(kind='pie', y=df.columns.values[1], labels=df[df.columns.values[0]].tolist(), \n",
    "            autopct='%.2f')\n",
    "    else:\n",
    "        raise Exception('cannot generate pie plot with > 2 column(s)')\n",
    "    plot.legend(loc='right', bbox_to_anchor=(1.5, 0.50))\n",
    "    plot.show()\n",
    "\n",
    "\n",
    "# plot line graph from dataframe\n",
    "def plot_line(df, col_count):\n",
    "    if col_count == 1:\n",
    "        df.index = df.index + 1\n",
    "        df.plot(kind='line')\n",
    "    elif col_count == 2:\n",
    "        df.plot(kind='line', x=df.columns.values[0], y=df.columns.values[1])\n",
    "    elif col_count == 3:\n",
    "        pivoted = pandas.pivot_table(df, \n",
    "                            values=df.columns.values[2], \n",
    "                            columns=df.columns.values[0], \n",
    "                            index=df.columns.values[1])\n",
    "        pivoted.plot()\n",
    "    else:\n",
    "        raise Exception('cannot generate line plot with > 3 column(s)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Cell Magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# invoke with %%ibmi\n",
    "@needs_local_scope \n",
    "@register_cell_magic\n",
    "def ibmi(line, cell):\n",
    "    if conn is None:\n",
    "        print('Not connected.')\n",
    "        return None\n",
    "    i = 0\n",
    "    flag = get_flag(line)\n",
    "    stmts = get_statements(cell.strip())\n",
    "    for stmt in stmts:\n",
    "        sql = ' '.join([x.replace('{','{{').replace('}','}}').replace('\\n','').strip() for x in stmt.split('\\n')])\n",
    "        i += 1\n",
    "        print('Executing statement {} of {} ...'.format(i, len(stmts)))\n",
    "        sql = sql.format(**globals())\n",
    "        cursor = conn.cursor()\n",
    "        try:\n",
    "            start_time = time.time()\n",
    "            affected = cursor.execute(sql)\n",
    "            while 1:\n",
    "                try:\n",
    "                    rows = cursor.fetchall()\n",
    "                    print('    Fetched {} row(s) in {} second(s)'.format(len(rows), round(time.time()-start_time,3)))\n",
    "                    if len(rows) > 0:\n",
    "                        result = output_rs(rows, cursor, flag)\n",
    "                        if not result is None: \n",
    "                            display(result)\n",
    "                except pyodbc.ProgrammingError as e:  # no result set, just an update,insert,create,etc statement\n",
    "                    print('    Statement executed successfully.\\n    {} row(s) affected by statement.'.format(affected.rowcount))\n",
    "                if not cursor.nextset(): break\n",
    "        except Exception as e:\n",
    "            print('Unexpected error occurred\\n' + str(e))\n",
    "        finally:\n",
    "            cursor.close()\n",
    "del ibmi"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}