diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2021-02-08 10:31:05 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2021-02-08 10:31:05 +0000 |
commit | 6884720fae8a2622b14e93d9e35ca5fcc2283b40 (patch) | |
tree | df6f736bb623cdd7932bbe2256101a6ac4ef7f35 /tests/features/db_utils.py | |
parent | Initial commit. (diff) | |
download | pgcli-6884720fae8a2622b14e93d9e35ca5fcc2283b40.tar.xz pgcli-6884720fae8a2622b14e93d9e35ca5fcc2283b40.zip |
Adding upstream version 3.1.0.upstream/3.1.0
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'tests/features/db_utils.py')
-rw-r--r-- | tests/features/db_utils.py | 78 |
1 files changed, 78 insertions, 0 deletions
diff --git a/tests/features/db_utils.py b/tests/features/db_utils.py new file mode 100644 index 0000000..f57bc3b --- /dev/null +++ b/tests/features/db_utils.py @@ -0,0 +1,78 @@ +from psycopg2 import connect +from psycopg2.extensions import AsIs + + +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) + + # ISOLATION_LEVEL_AUTOCOMMIT = 0 + # Needed for DB creation. + cn.set_isolation_level(0) + + with cn.cursor() as cr: + cr.execute("drop database if exists %s", (AsIs(dbname),)) + cr.execute("create database %s", (AsIs(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, database=dbname, password=password, port=port + ) + + print("Created connection: {0}.".format(cn.dsn)) + return cn + + +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) + + # ISOLATION_LEVEL_AUTOCOMMIT = 0 + # Needed for DB drop. + cn.set_isolation_level(0) + + with cn.cursor() as cr: + cr.execute("drop database if exists %s", (AsIs(dbname),)) + + close_cn(cn) + + +def close_cn(cn=None): + """ + Close connection. + :param connection: psycopg2.connection + """ + if cn: + cn.close() + print("Closed connection: {0}.".format(cn.dsn)) |