Skip to content

Postgresql python

python
import psycopg2
from datetime import datetime


class PostgreSQLConnector:

    def __init__(self, dbname, user, password, host, port):
        self.dbname = dbname
        self.user = user
        self.password = password
        self.host = host
        self.port = port
        self.conn = None
        self.cur = None

    def connect(self):
        self.conn = psycopg2.connect(dbname=self.dbname, user=self.user, password=self.password,
                                     host=self.host, port=self.port)
        self.cur = self.conn.cursor()

    def disconnect(self):
        if self.cur:
            self.cur.close()
        if self.conn:
            self.conn.close()

    def execute_query(self, sql, params=None):
        self.cur.execute(sql, params or ())
        self.conn.commit()

    def fetch_one(self, sql, params=None):
        self.cur.execute(sql, params or ())
        return self.cur.fetchone()

    def fetch_all(self, sql, params=None):
        self.cur.execute(sql, params or ())
        return self.cur.fetchall()

    def query_data(self, cols, table_name):
        sql = f"SELECT {cols} FROM {table_name}"
        for item in self.fetch_all(sql):
            item = list(map(lambda i: datetime.isoformat(i, sep=' ', timespec='seconds') if isinstance(i, datetime) else i, item))
            yield dict(zip(list(map(str.strip, cols.split(','))), item))

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()

Released under the MIT License.