python链接pgsql

分享两个python连接pgsql的方式

psycopg2连接

import csv
import json
from urllib import parse

from geopandas import GeoDataFrame
import psycopg2
from typing import Literal
from pandas import read_sql
from pathlib import Path


def translation(url):
    return parse.quote_plus(url)


class PostgresSQLConnection(object):
    def __init__(
            self,
            user: str,
            pwd: str,
            dbname: str,
            host: str,
            port: str = '5432'
    ):
        self.connect = psycopg2.connect(
            dbname=dbname,
            user=translation(user),
            password=translation(pwd),
            host=host,
            port=port
        )
        self.cursor = self.connect.cursor()

    def table_name(self):
        """
        获取数据库表格
        """
        sql = "SELECT table_name FROM INFORMATION_SCHEMA.TABLES where table_schema='public' and table_type='BASE TABLE' and table_name not like 'us%' and table_name not like 'spatial%'"
        table_list = self.fetchall(sql)
        table_name_list = [table[0] for table in table_list]
        return table_name_list

    def table_title(self, table_name):
        """
        获取表格标题
        """
        sql = 'SELECT column_name FROM information_schema.columns WHERE table_name=\'{}\';'.format(table_name)
        table_list = self.fetchall(sql)
        title_list = [table[0] for table in table_list]
        return title_list

    def get_table_data(self, table_name: str, *cols, is_pandas: bool = False):
        """
        获取表数据，原生sql数据
        :param table_name: 表名
        :param cols: 表标题
        :param is_pandas: 是否使用 DataFrame
        """
        if len(cols) == 0:
            col = '*'
        else:
            col = ", ".join(cols)
        sql = f'SELECT {col} FROM  "{table_name}"'

        if is_pandas:
            return read_sql(sql, self.connect)
        return self.fetchall(sql)

    def get_geojson(self, table_name: str, *cols, is_pandas: bool = False, geom='shape'):
        """
        获取 geojson
        """
        if len(cols) == 0:
            col = '*'
        else:
            col = ", ".join(cols)
        sql = f'select {col}, st_asgeojson({geom}) as geojson,st_srid({geom}) as srid from "{table_name}"'
        if is_pandas:
            return GeoDataFrame.from_postgis(sql, self.connect, geom_col=geom)
        return self.fetchall(sql)

    def get_srid(self, table_name, *, pattern='public', geom='shape'):
        """
        获取表格的srid，坐标系
        """
        # sql = f"select st_srid({geom}) from \"{table_name}\";"
        sql = f"SELECT Find_SRID('{pattern}', '{table_name}', '{geom}');"
        data = self.fetchall(sql)[0][0]
        return data

    def modify_space_field(self, _type, srid, table_name, *, geom='shape'):
        """
        修改 shape 字段类型
        """
        sql = f'ALTER TABLE {table_name} ALTER COLUMN {geom} TYPE geometry({_type},{srid});'
        self.set_sql(sql)

    def get_max_num(self, col, table_name, *, _t: Literal['max', 'min'] = 'max'):
        """
        获取最大值，最小值
        """
        sql = f'select {_t}(t.{col}) from {table_name} t;'
        data = self.fetchall(sql)
        return int(data[0][0])

    def get_where_line(self, table_name, *col, conditions: dict, is_pandas: bool = False):
        """
        获取指定的行数
        :param col: 表头
        :param table_name: 表名
        :param conditions: 搜索条件
        :param is_pandas: 是否是 pandas
        """
        con = [f"{i}='{v}'" for i, v in conditions.items()]
        sql = f'SELECT {",".join(col)} FROM "{table_name}" WHERE {" and ".join(con)};'
        if is_pandas:
            return read_sql(sql, self.connect)
        return self.fetchall(sql)

    def delete_table_data(self, table_name):
        """
        清空表格
        :param table_name:
        :return:
        """
        sql = f'truncate table {table_name}'
        self.set_sql(sql)

    def fetchall(self, sql):
        self.cursor.execute(sql)
        return self.cursor.fetchall()

    def copy_to(self, file: str, table_name: str, geom='shape', *, is_geom=False):
        """
        导出表数据
        :param file: csv文件保存路径
            ./njq/table_name
        :param table_name: 表名
            'table_name'
        :param no_field: 不包含某个字段
        """
        data = Path(file) / 'data.txt'
        setting = Path(file) / 'setting.json'

        data_fp = open(data, 'w', newline='', encoding='utf-8')
        setting_fp = open(setting, 'w', newline='', encoding='utf-8')

        columns = self.table_title(table_name)
        if is_geom is True:
            geom_type = self.get_geom_type(table_name, geom)
            srid = self.get_srid(table_name, geom=geom)
            data = {
                'geom_type': geom_type,
                'srid': srid,
                'columns': columns
            }
        else:
            if geom in columns:
                columns.remove(geom)
            data = {
                'geom_type': None,
                'srid': None,
                'columns': columns
            }

        self.cursor.copy_to(data_fp, table_name, sep=',', columns=columns)

        setting_fp.write(json.dumps(data, ensure_ascii=False, indent=4))
        data_fp.close()

    def copy_from(self, file, table_name, geom: str = 'shape', is_geom=False):
        """
        导入数据
        :param file: 本地路径
            './njq'
        :param table_name: 表名
            'table_name'
        """
        data = Path(file) / table_name / 'data.txt'
        setting = Path(file) / table_name / 'setting.json'

        data_fp = open(data, 'r', newline='', encoding='utf-8')
        setting_fp = open(setting, 'r', newline='', encoding='utf-8')

        json_data = json.loads(setting_fp.read())

        if is_geom:
            self.delete_table_data(table_name)
            self.modify_space_field(
                _type=json_data['geom_type'],
                srid=json_data['srid'],
                table_name=table_name,
                geom=geom
            )
        self.cursor.copy_from(data_fp, table_name, columns=json_data['columns'], sep=',')
        self.commit()

        setting_fp.close()
        data_fp.close()

    def get_geom_type(self, table_name, geom):
        sql = f'select st_asgeojson({geom}) as json from {table_name} t limit 1'
        return json.loads(self.fetchall(sql)[0][0])['type']

    def init_postgis(self):
        """
        初始化 postgis 插件
        """
        sql = """
        create extension postgis;
        create extension postgis_topology;
        create extension postgis_sfcgal;
        create extension fuzzystrmatch;
        create extension postgis_tiger_geocoder;
        create extension address_standardizer;
        create extension address_standardizer_data_us;
        """
        self.set_sql(sql)

    def set_sql(self, query: str, vars=None):
        """
        写入sql
        """
        self.cursor.execute(query, vars=vars)
        self.commit()

    def set_sqls(self, sql, vars_list=None):
        self.cursor.executemany(sql, vars_list)
        self.commit()

    def rollback(self):
        self.connect.rollback()

    def commit(self):
        self.connect.commit()

    def close(self):
        self.cursor.close()
        self.connect.close()

sqlalchemy连接

class PostgresSQL:
    def __init__(
            self,
            user,
            pwd,
            dbname,
            host,
            *,
            port: int = 5432,
            is_connect: bool = True,
            echo=False

    ):
        url = f"postgresql+psycopg2://{translation(user)}:{translation(pwd)}@{host}:{port}/{dbname}"
        self.engine = create_engine(url, echo=echo)

        self.conn = None
        if is_connect:
            self.conn = self.connect()
            
    def connect(self):
    	return self.engine.connect()

    def commit(self):
        self.conn.commit()

    def rollback(self):
        self.conn.rollback()

    def close(self):
        self.conn.close()

    def engine_clear(self):
        self.engine.clear_compiled_cache()