很多时候，我们用Python处理数据，需要连接到Mysql、Postgresql等数据库，获取表数据，再构建pandas的DataFrame进行进一步处理。但是查询数据库结果集是没有表字段名称的，我们希望构建的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):
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'))
