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'))

執行示例

更多