SQLAlchemy REPL Shell

sqlacmd.py

import cmd
import argparse
from sqlalchemy import create_engine, text
from sqlalchemy.exc import ProgrammingError
from tabulate import tabulate

class SQLShell(cmd.Cmd):
    intro = ("Welcome to SQLAlchemy REPL Shell. Type 'exit' to quit. \n"
             "This shell does not execute queries immediately. Press CTRL+D or type 'EOF' to execute. \n")
    _default_prompt = "SQL> "

    def __init__(self, connection_string):
        super().__init__()
        self.engine = create_engine(connection_string)
        self.connection = self.engine.connect()
        self.buffer = []
        self.enter_count = 0
        self.prompt = self._default_prompt


    def display_query_result(self, result):
        rows = [row for row in result]
        if rows:
            headers = rows[0].keys()
            rows_data = []
            for r in rows:
                rows_data.append([r[k] for k in headers])
            table = tabulate(rows_data, headers, tablefmt="fancy_grid")
            print(table)
        else:
            print("No rows returned.")

    def default(self, line):
        if line.strip().lower() == "exit":
            self.do_exit("")
        elif line == 'EOF':
            print('')
            self.execute_query()
        else:
            self.buffer.append(line.strip())
            self.prompt = ".... "

    def emptyline(self):
        pass

    def execute_query(self):
        if len(self.buffer) == 0:
            return
        sql_query = ' '.join(self.buffer)
        print(f'Executing {sql_query}')
        self.buffer.clear()
        self.enter_count = 0
        result = None
        try:
            result = self.connection.execute(text(sql_query))
        except ProgrammingError as e:
            print(f"Error: {e}")
        else:
            print("Query executed successfully.")
        if result is not None and result.returns_rows:
            self.display_query_result(result.mappings())
        self.prompt = self._default_prompt

    def do_exit(self, _):
        self.connection.close()
        self.engine.dispose()
        print("Goodbye!")
        return True

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="REPL shell to any SQLAlchemy supported databases")
    parser.add_argument('url', help='sqlalchemy connection url')

    args = parser.parse_args()
    shell = SQLShell(args.url)
    shell.cmdloop()