使用prompt_toolkit实现mysql交互式终端工具

1. prompt_toolkit

prompt_toolkit是一个专门用于实现交互式终端工具的python库，它支持代码补全，自动提示，查询历史，语法高亮，可以让你快速的实现一个可交互的终端工具，这在需要交互的场景下非常受欢迎，本文将带你实现一个mysql客户端交互工具，实现与mysql服务器的交互。

首先，进行安装

pip install prompt_toolkit

2. prompt 接收用户输入

既然是交互式工具，就必须能够做到与用户交互，prompt函数可以接收用户在终端的输入

from prompt_toolkit import prompt

def main():
    text = prompt("> ")
    print('You entered:', text)

if __name__ == '__main__':
    main()

启动程序后，在终端输入hello world， 得到"You entered:hello world" 的输出同时程序运行结束。实践中，与用户的交互是持续的，直到用户输入退出的命令，因此，需要在程序中使用while循环不停的接收用户的输入

from prompt_toolkit import prompt

def main():
    while True:
        text = prompt("> ")
        if text == 'quit':
            print("退出程序")
            break
        print('You entered:', text)

if __name__ == '__main__':
    main()

2. 创建数据库连接

为了能够与数据库进行交互，需要在程序里创建一个数据库的连接，为了降低难度，创建数据库连接所需要的参数直接硬编码在脚本里， 我使用mysql-connector库与mysql进行交互，安装命令

pip install mysql-connector

创建mysql的连接

import mysql.connector

mydb = mysql.connector.connect(
    host="101.201.37.248",  # 数据库主机地址
    user="manager",  # 数据库用户名
    passwd="3243Wdj!$dkwf44",  # 数据库密码
    port=1348,
)

mycursor = mydb.cursor(dictionary=True)

3. 实现对数据库的交互

from prompt_toolkit import prompt
import mysql.connector

mydb = mysql.connector.connect(
    host="101.201.37.248",  # 数据库主机地址
    user="manager",  # 数据库用户名
    passwd="3243Wdj!$dkwf44",  # 数据库密码
    port=1348,
)

mycursor = mydb.cursor(dictionary=True)

def main():
    while True:
        text = prompt("> ")
        if text == 'quit':
            print("退出程序")
            break

        try:
            mycursor.execute(text)
            datas = mycursor.fetchall()
        except Exception as e:
            print(e)
        else:
            print(datas)

if __name__ == '__main__':
    main()

启动程序，在终端里输入命令并回车

show databases

程序将执行show databases命令并返回结果

[{'Database': 'information_schema'}, {'Database': 'automated_testing'}, {'Database': 'stumanager'}]

上面的逻辑是不完善的，我们所执行的sql语句大致可分为两类，一类是有返回值的比如select语句，一类是没有返回值的比如insert语句，我使用了一个简单的方法来区分它们

def main():
    while True:
        text = prompt("> ")
        if text == 'quit':
            print("退出程序")
            break

        try:
            mycursor.execute(text)
            rowcount = mycursor.rowcount
            if rowcount >= 0:
                mydb.commit()
                print(f"{rowcount}行数据受影响")
            else:
                data = mycursor.fetchall()
                print(data)
        except Exception as e:
            print(e)

mycursor.rowcount是受到sql语句影响的条数，如果这个值大于等于0说明是在执行修改删除动作，如果是-1，则表示是查询类操作。

4. 加入语法高亮

使用语法高亮能让我们的终端工具更加漂亮，编写的sql语句中的关键字更加突出

from pygments.lexers import MySqlLexer
from prompt_toolkit.lexers import PygmentsLexer

def main():
    while True:
        text = prompt("> ", lexer=PygmentsLexer(MySqlLexer))

5. 加入自动补全功能

下面是代码改动的部分

from prompt_toolkit.completion import Completer, Completion


class MySimpleCompleter(Completer):
    def get_completions(self, document, complete_event):
        completions = ['create', 'select', 'insert', 'drop',
                               'delete', 'from', 'where', 'table']

        word_before_cursor = document.get_word_before_cursor()
        completions = [c for c in completions if c.startswith(word_before_cursor)]

        for completion in completions:
            yield Completion(completion, start_position=-len(word_before_cursor))
            
def main():
    while True:
        text = prompt("> ", lexer=PygmentsLexer(MySqlLexer), completer=MySimpleCompleter())

6. 保存历史命令

下面是代码改动部分

from prompt_toolkit.history import FileHistory

def main():
    history = FileHistory('my_sql.txt')
    while True:
        text = prompt("> ", lexer=PygmentsLexer(MySqlLexer), completer=MySimpleCompleter(), history=history)
        

在终端输入的命令都会记录在my_sql.txt里，下一次启动程序时，通过键盘的上上下键就可以选择之前输入过的命令，实现快速输入

7. 使用表格展示查询数据

虽然前面实现了数据的查询，但是展示的数据很难阅读，接下来的目标是使用表格展示他们，实现

from prettytable import PrettyTable

def show_data(lst):
    if not lst:
        print("未查到数据")
        return

    data = lst[0]
    coumns = list(data.keys())
    table = PrettyTable(coumns)
    for row in lst:
        row = list(row.values())
        row = [str(item) for item in row]
        table.add_row(row)

    print(table)
    

8. 完整代码与效果

import mysql.connector
from prompt_toolkit import prompt
from pygments.lexers import MySqlLexer
from prompt_toolkit.lexers import PygmentsLexer
from prompt_toolkit.completion import Completer, Completion
from prompt_toolkit.history import FileHistory
from prettytable import PrettyTable


mydb = mysql.connector.connect(
    host="",  # 数据库主机地址
    user="",  # 数据库用户名
    passwd="",  # 数据库密码
    port=3306,
)

mycursor = mydb.cursor(dictionary=True)

class MySimpleCompleter(Completer):
    def get_completions(self, document, complete_event):
        completions = ['create', 'select', 'insert', 'drop',
                               'delete', 'from', 'where', 'table']

        word_before_cursor = document.get_word_before_cursor()
        completions = [c for c in completions if c.startswith(word_before_cursor)]

        for completion in completions:
            yield Completion(completion, start_position=-len(word_before_cursor))

def show_data(lst):
    if not lst:
        print("未查到数据")
        return

    data = lst[0]
    coumns = list(data.keys())
    table = PrettyTable(coumns)
    for row in lst:
        row = list(row.values())
        row = [str(item) for item in row]
        table.add_row(row)

    print(table)

def main():
    history = FileHistory('my_sql.txt')
    while True:
        text = prompt("> ", lexer=PygmentsLexer(MySqlLexer), completer=MySimpleCompleter(), history=history)
        if text == 'quit':
            print("退出程序")
            break

        try:
            mycursor.execute(text)
            rowcount = mycursor.rowcount
            if rowcount >= 0:
                mydb.commit()
                print(f"{rowcount}行数据受影响")
            else:
                data = mycursor.fetchall()
                show_data(data)
        except Exception as e:
            print(e)

if __name__ == '__main__':
    main()

效果图