openai-sql/nogui/database/postgresql.py
2023-05-03 16:28:37 +02:00

69 lines
1.9 KiB
Python

import psycopg2
import os
class Postgres:
def __init__(self, ip, port, username, password, database):
self.db_ip = ip
self.db_port = port
self.db_username = username
self.db_password = password
self.db_name = database
self.conn = psycopg2.connect(
host=self.db_ip,
port=self.db_port,
user=self.db_username,
password=self.db_password,
database=self.db_name
)
self.tables = []
fetch = self.fetchall("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';")
for item in fetch:
self.tables.append(item[0])
self.tableschema = {}
for table in self.tables:
fetch = self.fetchall(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table}';")
self.tableschema.update({table: fetch})
def __str__(self):
return f'PostgreSQL Server: {self.db_ip}:{self.db_port} as {self.db_username} on database {self.db_name}'
def fetchall(self, query):
cur = self.conn.cursor()
cur.execute(query)
result = cur.fetchall()
cur.close()
return result
def fetchone(self, query):
cur = self.conn.cursor()
cur.execute(query)
result = cur.fetchone()
cur.close()
return result
def fetchmany(self, query, size):
cur = self.conn.cursor()
cur.execute(query)
result = cur.fetchmany(size)
cur.close()
return result
def execute(self, query):
cur = self.conn.cursor()
cur.execute(query)
self.conn.commit()
cur.close()
def executemany(self, query, values):
cur = self.conn.cursor()
cur.executemany(query, values)
self.conn.commit()
cur.close()
def closeconnection(self):
self.conn.close()