rookeries/api/tasks/db.py

100 lines
3.3 KiB
Python

"""
Tasks for working with a Docker Compose setup for testing and development.
:copyright: Copyright 2013-2016, Dorian Pula <dorian.pula@amber-penguin-software.ca>
:license: AGPL v3+
"""
import contextlib
import os
import sys
import traceback
import arrow
import invoke as inv
import sqlalchemy
from sqlalchemy import exc as sql_error
import sqlalchemy_utils
from rookeries import database as rookeries_db
from tasks.utils import environment
from tests.fixtures import sample_site
DEFAULT_DB_CONNECTION = 'postgresql+psycopg2://admin:password@localhost:5432/rookeries'
def get_db_connection_uri(db_connection, db_name):
db_connection = os.environ.get(rookeries_db.EXTERNAL_DB_URI_PARAM, db_connection)
return '/'.join((db_connection.rsplit('/', maxsplit=1)[0], db_name))
@inv.task
def bootstrap(ctx, db_connection=DEFAULT_DB_CONNECTION, db_name='rookeries', reset_db=False):
"""
Bootstrap the database with sample code.
:param ctx: Context of the invoke task.
:param db_connection: The database connection string. Gets overridden with the environment variable
ROOKERIES_SQLALCHEMY_DATABASE_URI.
:param db_name: The name of the database to create.
:param reset_db: Flag to reset a database and tables if they exist.
"""
db_connection = get_db_connection_uri(db_connection, db_name)
if not sqlalchemy_utils.database_exists(db_connection):
sqlalchemy_utils.create_database(db_connection)
if sqlalchemy_utils.database_exists(db_connection):
if not reset_db:
return
sqlalchemy_utils.drop_database(db_connection)
with test_database(db_connection=db_connection, keep_test_db=True, test_db_name=db_name):
pass
@contextlib.contextmanager
def test_database(db_connection=DEFAULT_DB_CONNECTION, keep_test_db=False,
test_db_name='rookeries_{}'.format(arrow.now().format('YYYY_MM_DD_HHmmss'))):
"""
Context for creating and managing a test database.
:param db_connection: The database connection string. Gets overridden with the environment variable
ROOKERIES_SQLALCHEMY_DATABASE_URI.
:param keep_test_db: Keep the test data and database around.
:param test_db_name: The name of the test database.
"""
try:
db_connection = os.environ.get(rookeries_db.EXTERNAL_DB_URI_PARAM, db_connection)
db_connection = '/'.join((db_connection.rsplit('/', maxsplit=1)[0], test_db_name))
if not sqlalchemy_utils.database_exists(db_connection):
sqlalchemy_utils.create_database(db_connection)
db_engine = sqlalchemy.create_engine(db_connection, echo=True)
rookeries_db.ModelBase.metadata.create_all(db_engine)
sample_site.populate_database(db_engine)
modified_environ = {
'ROOKERIES_SQLALCHEMY_DATABASE_URI': db_connection,
}
with environment(modified_environ):
pass
yield
except sql_error.OperationalError as err:
print(f'Unable to connect to the database because of "{err}" \n Stack: \n')
traceback.print_tb(sys.exc_info()[2])
raise err
except Exception as err:
print(f'Unable to run server test because of "{err}" \n Stack: \n')
traceback.print_tb(sys.exc_info()[2])
raise err
finally:
if not keep_test_db:
sqlalchemy_utils.drop_database(db_connection)