pyqt SQL Server 数据库查询-优化2

1、增加导出数据功能
2、增加删除表里数据功能

import sys
import pyodbc
from PyQt6.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QListWidget, QLineEdit, QPushButton, \
    QTableWidget, QTableWidgetItem, QLabel, QMessageBox
from PyQt6.QtGui import QFont
from PyQt6.QtCore import Qt
import pandas as pd
from datetime import datetime


class DatabaseQueryApp(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.connect_to_database()

    def initUI(self):
        # 设置窗口字体
        font = QFont()
        font.setPointSize(18)
        self.setFont(font)

        # 创建布局
        main_layout = QVBoxLayout()

        # 表列表、字段列表和输入框布局
        list_input_layout = QHBoxLayout()

        # 表列表部分
        table_label = QLabel("表:")
        list_input_layout.addWidget(table_label)
        self.table_list = QListWidget()
        # 调整显示表的控件大小
        self.table_list.setFixedHeight(100)
        # 设置 QListWidget 选中项整行变蓝色
        self.table_list.setStyleSheet("QListWidget::item:selected { background-color: blue; color: white; }")
        self.table_list.itemClicked.connect(self.show_table_columns)
        list_input_layout.addWidget(self.table_list)

        # 字段列表部分
        column_label = QLabel("字段:")
        list_input_layout.addWidget(column_label)
        self.column_list = QListWidget()
        self.column_list.setFixedHeight(100)
        self.column_list.setStyleSheet("QListWidget::item:selected { background-color: blue; color: white; }")
        list_input_layout.addWidget(self.column_list)

        # 查询输入部分
        value_label = QLabel("查询值:")
        list_input_layout.addWidget(value_label)
        self.value_input = QLineEdit()
        self.query_button = QPushButton('查询')

        # 获取当前按钮的大小
        current_width = self.query_button.sizeHint().width()
        current_height = self.query_button.sizeHint().height()
        # 设置按钮大小为原来的两倍
        button_size = (current_width * 2, current_height * 2)
        self.query_button.setFixedSize(*button_size)
        self.query_button.clicked.connect(self.execute_query)
        list_input_layout.addWidget(self.value_input)
        list_input_layout.addWidget(self.query_button)

        # 导出按钮
        self.export_button = QPushButton('导出')
        self.export_button.setFixedSize(*button_size)
        self.export_button.clicked.connect(self.export_table)
        list_input_layout.addWidget(self.export_button)

        # 清空按钮
        self.clear_button = QPushButton('清空')
        self.clear_button.setFixedSize(*button_size)
        self.clear_button.clicked.connect(self.clear_table)
        list_input_layout.addWidget(self.clear_button)

        main_layout.addLayout(list_input_layout)

        # 自定义查询输入部分
        custom_query_layout = QHBoxLayout()
        custom_query_label = QLabel("自定义查询:")
        custom_query_layout.addWidget(custom_query_label)
        self.custom_query_input = QLineEdit()
        self.custom_execute_button = QPushButton('执行')
        self.custom_execute_button.setFixedSize(*button_size)
        self.custom_execute_button.clicked.connect(self.execute_custom_query)
        custom_query_layout.addWidget(self.custom_query_input)
        custom_query_layout.addWidget(self.custom_execute_button)

        main_layout.addLayout(custom_query_layout)

        # 查询结果表格部分
        result_label = QLabel("查询结果:")
        main_layout.addWidget(result_label)
        self.result_table = QTableWidget()
        self.result_table.setColumnCount(0)
        self.result_table.setRowCount(0)
        # 设置 QTableWidget 选择行为为整行选择
        self.result_table.setSelectionBehavior(QTableWidget.SelectionBehavior.SelectRows)
        # 设置 QTableWidget 选中行整行变蓝色
        self.result_table.setStyleSheet("QTableWidget::item:selected { background-color: blue; color: white; }")

        main_layout.addWidget(self.result_table)

        self.setLayout(main_layout)
        self.setWindowTitle('数据库查询工具')
        self.setGeometry(300, 300, 1200, 900)
        self.setStyleSheet("""
            QWidget {
                background-color: #f0f0f0;
            }
            QLabel {
                font-weight: bold;
            }
            QPushButton {
                background-color: #4CAF50;
                color: white;
                padding: 10px 20px;
                border: none;
                border-radius: 5px;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
            QLineEdit {
                padding: 8px;
                border: 1px solid #ccc;
                border-radius: 5px;
            }
        """)
        self.show()

    def connect_to_database(self):
        try:
            # 使用指定的连接字符串
            connection_string = 'DRIVER={SQL Server};SERVER=LEGENDLI;DATABASE=testbase;UID=sa;PWD=1'
            self.conn = pyodbc.connect(connection_string)
            cursor = self.conn.cursor()
            cursor.execute("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE = 'BASE TABLE'")
            tables = cursor.fetchall()
            for table in tables:
                self.table_list.addItem(table[0])
        except Exception as e:
            self.show_error_message(f"数据库连接错误: {e}")

    def show_table_columns(self, item):
        table_name = item.text()
        try:
            cursor = self.conn.cursor()
            cursor.execute(f"SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = '{table_name}'")
            columns = cursor.fetchall()
            self.column_list.clear()
            for column in columns:
                self.column_list.addItem(column[0])
        except Exception as e:
            self.show_error_message(f"获取字段信息错误: {e}")

    def execute_query(self):
        selected_table_items = self.table_list.selectedItems()
        selected_column_items = self.column_list.selectedItems()
        if not selected_table_items:
            self.show_error_message("请选择一个表。")
            return

        table_name = selected_table_items[0].text()
        value = self.value_input.text()

        if not selected_column_items or not value:
            query = f"SELECT * FROM {table_name}"
        else:
            column_name = selected_column_items[0].text()
            query = f"SELECT * FROM {table_name} WHERE {column_name} = '{value}'"

        try:
            cursor = self.conn.cursor()
            cursor.execute(query)
            results = cursor.fetchall()
            headers = [description[0] for description in cursor.description]

            self.result_table.setColumnCount(len(headers))
            self.result_table.setRowCount(len(results))
            self.result_table.setHorizontalHeaderLabels(headers)

            for row_index, row_data in enumerate(results):
                for col_index, col_data in enumerate(row_data):
                    item = QTableWidgetItem(str(col_data))
                    self.result_table.setItem(row_index, col_index, item)

        except Exception as e:
            self.show_error_message(f"查询错误: {e}")

    def execute_custom_query(self):
        query = self.custom_query_input.text().strip().lower()
        # 检查是否包含危险指令
        if "drop" in query or "DROP" in query \
                or 'update' in query or 'UPDATE' in query \
                or 'delete' in query or 'DELETE' in query \
                or 'truncate' in query or 'TRUNCATE' in query:
            self.show_error_message("不允许执行删除表或删除数据库的指令。")
            return

        try:
            cursor = self.conn.cursor()
            cursor.execute(query)
            results = cursor.fetchall()
            headers = [description[0] for description in cursor.description]

            self.result_table.setColumnCount(len(headers))
            self.result_table.setRowCount(len(results))
            self.result_table.setHorizontalHeaderLabels(headers)

            for row_index, row_data in enumerate(results):
                for col_index, col_data in enumerate(row_data):
                    item = QTableWidgetItem(str(col_data))
                    self.result_table.setItem(row_index, col_index, item)

        except Exception as e:
            self.show_error_message(f"查询错误: {e}")

    def export_table(self):
        selected_table_items = self.table_list.selectedItems()
        if not selected_table_items:
            self.show_error_message("请选择一个表。")
            return

        table_name = selected_table_items[0].text()
        try:
            # 导出数据到 Excel
            query = f"SELECT * FROM {table_name}"
            df = pd.read_sql(query, self.conn)
            now = datetime.now().strftime("%Y%m%d%H%M%S")
            file_name = f"{table_name}_{now}.xlsx"
            df.to_excel(file_name, index=False)
            self.show_success_message(f"数据已导出到 {file_name}。")
        except Exception as e:
            self.show_error_message(f"导出数据时出错: {e}")

    def clear_table(self):
        selected_table_items = self.table_list.selectedItems()
        if not selected_table_items:
            self.show_error_message("请选择一个表。")
            return

        table_name = selected_table_items[0].text()
        reply = QMessageBox.question(self, '确认', f'是否真的要清空表 {table_name} 的数据？',
                                     QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
                                     QMessageBox.StandardButton.No)
        if reply == QMessageBox.StandardButton.Yes:
            try:
                cursor = self.conn.cursor()
                cursor.execute(f"DELETE FROM {table_name}")
                self.conn.commit()
                self.show_success_message(f"表 {table_name} 的数据已清空。")
            except Exception as e:
                self.show_error_message(f"清空表数据时出错: {e}")

    def show_error_message(self, message):
        msg_box = QMessageBox()
        msg_box.setIcon(QMessageBox.Icon.Critical)
        msg_box.setText(message)
        msg_box.setWindowTitle("错误提示")
        msg_box.exec()

    def show_success_message(self, message):
        msg_box = QMessageBox()
        msg_box.setIcon(QMessageBox.Icon.Information)
        msg_box.setText(message)
        msg_box.setWindowTitle("成功提示")
        msg_box.exec()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = DatabaseQueryApp()
    sys.exit(app.exec())