def test_get_instance_exist_extra_args(self, pooling_mock):
PostgreSQLPool.get_instance(host="test_host",
user="******",
password="******",
database="test_db",
database_args={"ssl_ca": "file"})
assert pooling_mock.mock_calls == []
def test_get_instance_different(self, pooling_mock):
inst_1 = PostgreSQLPool.get_instance(host="test_host",
user="******",
password="******",
database="test_db")
inst_2 = PostgreSQLPool.get_instance(host="test_host",
user="******",
password="******",
database="test_db2")
assert inst_1 is not None and inst_2 is not None
assert inst_1 != inst_2
def test_get_instance_wrong_chars(self, pooling_mock):
PostgreSQLPool.get_instance(host="test_host",
user="******",
password="******",
database="test_db")
assert pooling_mock.mock_calls == [
call.SimpleConnectionPool(maxconn=5,
minconn=5,
host='test_host',
database='test_db',
user='******',
password='******')
]
def close(self) -> None:
super(PostgreSQLHelper, self).close()
self.logger.info("Closing connection to database!")
self.cursor.close()
if self.pooled:
PostgreSQLPool.get_instance(
host=self.host,
user=self.user,
database=self.database,
database_args=self.database_args).putconn(self.db_conn,
key=id(self))
else:
self.db_conn.close()
def test_get_instance_not_exist(self, pooling_mock):
mock = pooling_mock.SimpleConnectionPool
PostgreSQLPool.get_instance(host="test_host",
user="******",
password="******",
database="test_db")
assert mock.mock_calls == [
call(maxconn=5,
minconn=5,
host='test_host',
database='test_db',
user='******',
password='******')
]
def test_get_instance_not_exist_extra_args(self, pooling_mock):
PostgreSQLPool.get_instance(host="test_host2",
user="******",
password="******",
database="test_db",
database_args={"ssl_ca": "file"})
assert pooling_mock.mock_calls == [
call.SimpleConnectionPool(maxconn=5,
minconn=5,
host='test_host2',
database='test_db',
user='******',
password='******',
ssl_ca='file')
]
def __init__(self,
host: str = os.environ.get('DB_URL'),
user: str = os.environ.get('DB_USER'),
password: str = os.environ.get('DB_PASSWORD'),
database: str = os.environ.get('DB_NAME'),
pooled: bool = True,
database_args: dict = None):
self.logger = logging.getLogger(__name__)
self.host = str(host) if host is not None else 'localhost'
self.user = str(user) if user is not None else 'root'
self.database = str(database) if database is not None else 'default'
if password is None:
self.logger.warning("No password provided to database. Using "
"default insecure password 'root'")
password = '******'
self.database_args = database_args \
if database_args is not None else dict()
self.pooled = pooled
try:
self.logger.info(
"Connecting to database %s at %s "
"with username %s. Pooled: %s. Extra args: %s", database, host,
user, pooled, database_args)
if self.pooled:
self.db_conn = PostgreSQLPool.get_instance(
host=self.host,
user=self.user,
password=str(password),
database=self.database,
database_args=self.database_args).getconn(key=id(self))
else:
self.db_conn = psycopg2.connect(host=self.host,
user=self.user,
password=str(password),
database=self.database,
**self.database_args)
self.cursor = self.db_conn.cursor()
except (InterfaceError, ValueError, DatabaseError, PoolError,
ProgrammingError) as err:
self.logger.critical("Unable to connect to database!",
exc_info=True)
raise ConnectionError("\nSomething went wrong when connecting "
"to postgresql: {err}\n\n".format(err=err)) \
from err