Python操作PostGre的简单封装

---

一、安装依赖

pip install psycopg2 numpy 

---

二、配置文件

utils.config.py

import os


############### 233 PostGre Configuration ###############
POSTGRE_HOST = "192.168.0.233"
POSTGRE_PORT = 5432
POSTGRE_USER = "usss"
POSTGRE_PWD = "pwww"
POSTGRE_DB = "dbbb"

---

三、实现类

utils.PostGreOp.py

# encoding: utf-8

# sql导出
# mysqldump -u root -p db_name > test_db.sql

import psycopg2
import config.config as cf
import json
import numpy as np



class PostGreOp(object):
    def __init__(self, host=cf.POSTGRE_HOST, port=cf.POSTGRE_PORT, username=cf.POSTGRE_USER, password=cf.POSTGRE_PWD, db=cf.POSTGRE_DB):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.db = db

    # 增删改
    def operate(self, sql):
        db = psycopg2.connect(host=self.host, port=self.port, user=self.username, password=self.password, database=self.db)
        cur = db.cursor()
        try:
            # 执行sql语句
            cur.execute(sql)
            op_id = cur.lastrowid
            cur.close()
            # 提交到数据库执行
            db.commit()
        except Exception as e:
            print(e)

            op_id = None

            cur.close()
            # Rollback in case there is any error
            db.rollback()
        # 关闭数据库连接
        db.close()

        return op_id


    # 查
    def select(self, sql):
        db = psycopg2.connect(host=self.host, port=self.port, user=self.username, password=self.password, database=self.db)
        cur = db.cursor()
        results = None
        try:
            # 执行sql语句
            cur.execute(sql)
            # 获取所有记录列表
            results = cur.fetchall()
            # print(results)
        except Exception as e:
            print(e)
            # 关闭数据库连接
        db.close()
        return results

    @classmethod
    def escape_str(cls, text):
        '''
        string类型数据导入时有可能出现单双引号等需要转义的字段，也可能出现nan这样的字段，需要先处理一下
        :return:
        '''
        if cls.isNaNo(text):
            return 'null'
        else:
        	return "'" + str(text).replace("'", "''") + "'"
            # return "'" + json.dumps(str(text), ensure_ascii=False)[1:-1] + "'"

    @classmethod
    def escape_num(cls, text):
        '''
        转一下整数，报错的话说明传入的不是数字，有sql注入风险
        :return:
        '''
        if str(text) == '0':
            return '0'
        elif cls.isNaNo(text):
            return 'null'
        else:
            # return json.dumps(str(text), ensure_ascii=False)
            try:
                int(text)
                return str(text)
            except Exception as e:
                raise Exception('传入不是数字，有sql注入风险')


        # if cls.isNaNo(text):
        #     return 'null'
        # else:
        #     return json.dumps(str(text), ensure_ascii=False)

    @classmethod
    def isNaNo(cls, sth):
        '''
        NaN、None或者空字符串返回True，其他情况返回False
        '''
        if not sth:
            return True
        if isinstance(sth, float):
            if np.isnan(sth):
                return True
        return False


if __name__ == '__main__':
    # 创建数据库连接
    psg = PostGreOp()

    # 执行sql【查询db下所有表】
    table_name_list = psg.select('''select * from pg_tables''')
    print(table_name_list)
    # [
    #   ('openalex', 'authors_counts_by_year', 'psss', None, False, False, False, False),
    #   ('openalex', 'authors_ids', 'psss', None, False, False, False, False),
    #   ('openalex', 'authors_x_concepts', 'psss', None, False, False, False, False),
    #   ('openalex', 'concepts_ancestors', 'psss', None, False, False, False, False),
    #   ('pg_catalog', 'pg_statistic', 'postgres', None, True, False, False, False),
    #   ('pg_catalog', 'pg_type', 'postgres', None, True, False, False, False),
    #   ('openalex', 'concepts_counts_by_year', 'psss', None, False, False, False, False),
    #   ('openalex', 'authors', 'psss', None, True, False, False, False),
    #   ('openalex', 'concepts_ids', 'psss', None, False, False, False, False),
    #   ('openalex', 'concepts_related_concepts', 'psss', None, False, False, False, False)
    # ]