Python连接数据库将结果转换为DataFrame(列名和表字段一致)

很多时候，我们用Python处理数据，需要连接到Mysql、Postgresql等数据库，获取表数据，再构建pandas的DataFrame进行进一步处理。但是查询数据库结果集是没有表字段名称的，我们希望构建的DataFrame的列名和表字段一样。

直接上代码

这里以Postgresql数据库为例，Mysql数据库差不多，其他的自行改造。

先封装一个查询类，查询返回的结果是一个字典，head是表列名,data是表数据，再用DataFrame构造数据结构。

import psycopg2
import pandas as pd

class db_pg:
    def __init__(self, host, db, user, pwd, port):
        self.host = host
        self.db = db
        self.user = user
        self.pwd = pwd
        self.port = port
        self._conn = self._connect()
        self._cursor = self._conn.cursor()

    def _connect(self):
        return psycopg2.connect(
            database=self.db,
            user=self.user,
            password=self.pwd,
            host=self.host,
            port=self.port)

    def select(self, sqlCode):
        self.common(sqlCode)
        col_names = []
        result = {}
        column_count = len(self._cursor.description)
        for i in range(column_count):
            desc = self._cursor.description[i]
            col_names.append(desc[0])
        data = self._cursor.fetchall()
        result['head'] = col_names
        result['data'] = data
        return result

    def close(self):
        self._cursor.close()
        self._conn.close()

    def common(self, sqlCode):
        try:
            self._cursor.execute(sqlCode)
        except Exception as e:
            print(e)
            self._conn.rollback()
            self._cursor.execute(sqlCode)
        self._conn.commit()

    def __del__(self):
        self.close()

db_conn = {
        'host': "******",
        'db'  : "******",
        'user': "******",
        'pwd' :"******",
        'port': "******"
    }
pg_conn = db_pg(host=db_conn['host'],db=db_conn['db'],user=db_conn['user'],pwd=db_conn['pwd'],port=db_conn['port'])
rs            = pg_conn.select("select * from test")
rs_df       = pd.DataFrame(list(rs.get('data')),columns=rs.get('head'))

运行示例

更多