Python 封装pandas连接数据库方法

用户11081884

发布于 2026-07-20 18:53:14

530

在数据分析工作中，经常需要将数据库中的数据读取到pandas中进行处理，或者将pandas DataFrame写入数据库。通过封装这些操作可以提高代码复用性，简化开发流程。本文将介绍如何封装pandas连接数据库的常用方法。

封装实现

下面是一个封装了pandas连接MySQL数据库的类，提供了读取和写入数据的方法：

import pandas as pd
from sqlalchemy import create_engine
from sqlalchemy.exc import SQLAlchemyError

class PandasDBConnector:
    def __init__(self, db_type='mysql', host='localhost', port=3306, 
                 user='root', password='', database='test'):
        """
        初始化数据库连接参数
        
        参数:
            db_type: 数据库类型(mysql/postgresql/sqlite等)
            host: 数据库主机地址
            port: 数据库端口
            user: 用户名
            password: 密码
            database: 数据库名
        """
        self.db_type = db_type
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.engine = None
        
    def connect(self):
        """创建数据库连接引擎"""
        try:
            connection_str = f"{self.db_type}://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}"
            self.engine = create_engine(connection_str)
            print("数据库连接成功")
            return True
        except SQLAlchemyError as e:
            print(f"数据库连接失败: {e}")
            return False
            
    def read_table(self, table_name, query=None, chunksize=None):
        """
        从数据库读取表数据到DataFrame
        
        参数:
            table_name: 表名
            query: 自定义SQL查询语句
            chunksize: 分块读取大小
            
        返回:
            DataFrame或迭代器
        """
        if not self.engine:
            print("请先建立数据库连接")
            return None
            
        try:
            if query:
                return pd.read_sql(query, self.engine, chunksize=chunksize)
            else:
                return pd.read_sql_table(table_name, self.engine, chunksize=chunksize)
        except SQLAlchemyError as e:
            print(f"读取数据失败: {e}")
            return None
            
    def write_table(self, df, table_name, if_exists='fail', index=False, chunksize=None):
        """
        将DataFrame写入数据库表
        
        参数:
            df: 要写入的DataFrame
            table_name: 表名
            if_exists: 表存在时的处理方式(fail/replace/append)
            index: 是否写入索引
            chunksize: 分块写入大小
        """
        if not self.engine:
            print("请先建立数据库连接")
            return False
            
        try:
            df.to_sql(
                table_name, 
                self.engine, 
                if_exists=if_exists, 
                index=index, 
                chunksize=chunksize
            )
            print(f"数据成功写入表 {table_name}")
            return True
        except SQLAlchemyError as e:
            print(f"写入数据失败: {e}")
            return False
            
    def close(self):
        """关闭数据库连接"""
        if self.engine:
            self.engine.dispose()
            print("数据库连接已关闭")

使用示例

if __name__ == '__main__':
    # 创建连接器实例
    db_conn = PandasDBConnector(
        db_type='mysql',
        host='localhost',
        user='root',
        password='your_password',
        database='test_db'
    )
    
    # 连接数据库
    if db_conn.connect():
        # 读取整个表
        df = db_conn.read_table('users')
        print(df.head())
        
        # 使用自定义查询读取数据
        query = "SELECT * FROM orders WHERE amount > 1000"
        orders_df = db_conn.read_table('orders', query=query)
        print(orders_df.head())
        
        # 创建新DataFrame并写入数据库
        new_data = pd.DataFrame({
            'product_id': [101, 102, 103],
            'product_name': ['Laptop', 'Phone', 'Tablet'],
            'price': [999.99, 699.99, 299.99]
        })
        
        db_conn.write_table(new_data, 'products', if_exists='append')
        
        # 关闭连接
        db_conn.close()