mirror of https://github.com/dbcli/pgcli
88 lines
2.0 KiB
Python
88 lines
2.0 KiB
Python
from psycopg import connect
|
|
|
|
|
|
def create_db(
|
|
hostname="localhost", username=None, password=None, dbname=None, port=None
|
|
):
|
|
"""Create test database.
|
|
|
|
:param hostname: string
|
|
:param username: string
|
|
:param password: string
|
|
:param dbname: string
|
|
:param port: int
|
|
:return:
|
|
|
|
"""
|
|
cn = create_cn(hostname, password, username, "postgres", port)
|
|
|
|
cn.autocommit = True
|
|
with cn.cursor() as cr:
|
|
cr.execute(f"drop database if exists {dbname}")
|
|
cr.execute(f"create database {dbname}")
|
|
|
|
cn.close()
|
|
|
|
cn = create_cn(hostname, password, username, dbname, port)
|
|
return cn
|
|
|
|
|
|
def create_cn(hostname, password, username, dbname, port):
|
|
"""
|
|
Open connection to database.
|
|
:param hostname:
|
|
:param password:
|
|
:param username:
|
|
:param dbname: string
|
|
:return: psycopg2.connection
|
|
"""
|
|
cn = connect(
|
|
host=hostname, user=username, dbname=dbname, password=password, port=port
|
|
)
|
|
|
|
print(f"Created connection: {cn.info.get_parameters()}.")
|
|
return cn
|
|
|
|
|
|
def pgbouncer_available(hostname="localhost", password=None, username="postgres"):
|
|
cn = None
|
|
try:
|
|
cn = create_cn(hostname, password, username, "pgbouncer", 6432)
|
|
return True
|
|
except:
|
|
print("Pgbouncer is not available.")
|
|
finally:
|
|
if cn:
|
|
cn.close()
|
|
return False
|
|
|
|
|
|
def drop_db(hostname="localhost", username=None, password=None, dbname=None, port=None):
|
|
"""
|
|
Drop database.
|
|
:param hostname: string
|
|
:param username: string
|
|
:param password: string
|
|
:param dbname: string
|
|
"""
|
|
cn = create_cn(hostname, password, username, "postgres", port)
|
|
|
|
# Needed for DB drop.
|
|
cn.autocommit = True
|
|
|
|
with cn.cursor() as cr:
|
|
cr.execute(f"drop database if exists {dbname}")
|
|
|
|
close_cn(cn)
|
|
|
|
|
|
def close_cn(cn=None):
|
|
"""
|
|
Close connection.
|
|
:param connection: psycopg2.connection
|
|
"""
|
|
if cn:
|
|
cn_params = cn.info.get_parameters()
|
|
cn.close()
|
|
print(f"Closed connection: {cn_params}.")
|