Simple BigQuery shell using ibis

bqshell.py

import pandas as pd
import ibis
import click

from prompt_toolkit import prompt
from prompt_toolkit.layout.lexers import PygmentsLexer
from pygments.lexers.sql import PostgresLexer


@click.command()
@click.argument('project')
@click.argument('dataset')
def main(project, dataset):
    lexer = PygmentsLexer(PostgresLexer)
    con = ibis.bigquery.connect(project, dataset)
    while True:
        query = prompt('bqsh> ', lexer=lexer)
        try:
            cursor = con.raw_sql(query)
            result = pd.DataFrame(cursor.fetchall(), columns=cursor.columns)
        except Exception as e:
            print(e)
        else:
            print(result)


if __name__ == '__main__':
    main()