"""db.py. Database level utility class for testing against the artist_info schema. """ from functools import wraps import sys from artist import config from artist.connectors.mysql import db_session from artist.models import artist as artist_model CREATE_ARTIST_INFO = """ CREATE TABLE `artist_info` ( artist_id INTEGER NOT NULL PRIMARY KEY, artist_type TEXT NOT NULL DEFAULT 'artist', isni_id TEXT DEFAULT NULL, name TEXT DEFAULT NULL, url TEXT DEFAULT NULL, orchard_country INTEGER DEFAULT NULL, unique_artist_id INTEGER DEFAULT NULL, vendor_id INTEGER DEFAULT NULL, last_updated DATETIME ); """ CREATE_RELEASE_ARTIST = """ CREATE TABLE release_artist ( release_artist_id INTEGER NOT NULL PRIMARY KEY, release_id INTEGER NOT NULL, artist_info_id INTEGER, artist_name TEXT NOT NULL, role TEXT ); """ CREATE_ARTIST_IDENTIFIER = """ CREATE TABLE artist_identifier ( id INTEGER NOT NULL PRIMARY KEY, artist_info_id INTEGER NOT NULL, store_id INTEGER NOT NULL, identifier TEXT NOT NULL, FOREIGN KEY(artist_info_id) REFERENCES artist_info(artist_id) ); """ CREATE_TRACK = """ CREATE TABLE track ( id INTEGER NOT NULL PRIMARY KEY, release_id INTEGER NOT NULL ); """ CREATE_TRACK_ARTIST = """ CREATE TABLE track_artist ( id INTEGER NOT NULL PRIMARY KEY, artist_info_id INTEGER, name TEXT NOT NULL, type TEXT NOT NULL, track_id INTEGER NOT NULL ); """ CREATE_TRACK_WRITER = """ CREATE TABLE track_writer ( track_writer_id INTEGER NOT NULL PRIMARY KEY, writer_name TEXT NOT NULL, artist_info_id INTEGER, unique_track_id INTEGER NOT NULL ); """ CREATE_RELEASES = """ CREATE TABLE `releases` ( release_id INTEGER NOT NULL PRIMARY KEY, artist_id INTEGER, subaccount_id INTEGER DEFAULT NULL, genre_id INTEGER, last_updated DATETIME ); """ CREATE_GENRE = """ CREATE TABLE `genre` ( genre_id INTEGER NOT NULL PRIMARY KEY, genre TEXT ); """ CREATE_VENDOR = """ CREATE TABLE `vendor` ( vendor_id INTEGER NOT NULL PRIMARY KEY, country INTEGER ); """ CREATE_COUNTRY = """ CREATE TABLE `country` ( id INTEGER NOT NULL PRIMARY KEY, country_code STRING ); """ DROP_ARTIST_INFO = """ DROP TABLE IF EXISTS `artist_info`; """ DROP_RELEASE_ARTIST = """ DROP TABLE IF EXISTS release_artist; """ DROP_ARTIST_IDENTIFIER = """ DROP TABLE IF EXISTS artist_identifier; """ DROP_TRACK = """ DROP TABLE IF EXISTS track; """ DROP_TRACK_ARTIST = """ DROP TABLE IF EXISTS track_artist; """ DROP_TRACK_WRITER = """ DROP TABLE IF EXISTS track_writer; """ DROP_RELEASES = """ DROP TABLE IF EXISTS `releases`; """ DROP_GENRE = """ DROP TABLE IF EXISTS `genre`; """ DROP_VENDOR = """ DROP TABLE IF EXISTS `vendor`; """ DROP_COUNTRY = """ DROP TABLE IF EXISTS `country`; """ def test_schema(function): """Test schema. Decorator that creates the test DB schema. before a function call and tears the schema down after the function call has finished. This just creates the schema and does not seed data. Indvidual test cases can use factories to seed data as needed. Args: Function (func): the function to be called after creating the test schema. Returns: Function: The decorated function. """ @wraps(function) def call_function_within_db_context(*args, **kwargs): create_artist_info_table() create_release_artist_table() create_track_artist_table() create_track_writer_table() create_artist_identifier_table() create_track_table() create_releases_table() create_genre_table() create_vendor_table() create_country_table() try: function_return = function(*args, **kwargs) finally: drop_artist_info_table() drop_release_artist_table() drop_track_table() drop_track_artist_table() drop_track_writer_table() drop_artist_identifier_table() drop_releases_table() drop_genre_table() drop_vendor_table() drop_country_table() return function_return return call_function_within_db_context def create_artist_info_table(): """Create the artist_info table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_ARTIST_INFO) def create_release_artist_table(): """Create the release_artist table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_RELEASE_ARTIST) def create_track_table(): """Create the track table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_TRACK) def create_track_artist_table(): """Create the track_artist table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_TRACK_ARTIST) def create_track_writer_table(): """Create the track_writer table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_TRACK_WRITER) def create_artist_identifier_table(): """Create the artist_identifier table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_ARTIST_IDENTIFIER) def create_releases_table(): """Create the releases table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_RELEASES) def create_genre_table(): """Create the genre table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_GENRE) def create_vendor_table(): """Create the vendor table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_VENDOR) def create_country_table(): """Create the country table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(CREATE_COUNTRY) def drop_artist_info_table(): """Drop the artist_info table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_ARTIST_INFO) def drop_release_artist_table(): """Drop the release_artist table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_RELEASE_ARTIST) def drop_artist_identifier_table(): """Drop the artist_identifier table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_ARTIST_IDENTIFIER) def drop_track_table(): """Drop the track table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_TRACK) def drop_track_artist_table(): """Drop the track_artist table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_TRACK_ARTIST) def drop_track_writer_table(): """Drop the track_writer table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_TRACK_WRITER) def drop_releases_table(): """Drop the releases table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_RELEASES) def drop_genre_table(): """Drop the genre table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_GENRE) def drop_vendor_table(): """Drop the vendor table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_VENDOR) def drop_country_table(): """Drop the country table.""" with db_session() as session: _exit_if_not_test_environment(session) session.execute(DROP_COUNTRY) def seed_models(models): """Save the given model(s) to the DB.""" if not hasattr(models, '__iter__'): models = [models] with db_session() as session: _exit_if_not_test_environment(session) for model in models: session.merge(model) session.commit() def seed_artist(item): """Save the item to the DB.""" with db_session() as session: _exit_if_not_test_environment(session) sql = ('INSERT INTO artist_info ' '(artist_id, name, vendor_id, last_updated) ' 'VALUES ({0}, \'{1}\', {2}, \'{3}\')'.format( item['artist_id'], item['name'], item['vendor_id'], item['last_updated'])) session.execute(sql) session.commit() def seed_release(item): """Save the item to the DB.""" with db_session() as session: _exit_if_not_test_environment(session) sql = ('INSERT INTO releases ' '(release_id, artist_id, subaccount_id, genre_id, last_updated) ' 'VALUES ({0}, {1}, {2}, {3}, \'{4}\')'.format( item.get('release_id'), item.get('artist_id'), item.get('subaccount_id', 'NULL'), item.get('genre_id', 'NULL'), item.get('last_updated'))) session.execute(sql) session.commit() def seed_genre(item): """Save the item to the DB.""" with db_session() as session: _exit_if_not_test_environment(session) sql = ('INSERT INTO genre (genre_id, genre) ' 'VALUES ({0}, \'{1}\')'.format( item['genre_id'], item['genre'])) session.execute(sql) session.commit() def get_artist_by_id(artist_id): """Get an artist from the db based on id. Args: id (int): artist id Returns: Artist: artist model instance """ with db_session() as session: _exit_if_not_test_environment(session) query_object = session.query(artist_model.Artist).filter_by( artist_id=artist_id) return query_object.first() def _exit_if_not_test_environment(session): """For safety, only run tests in test environment pointed to sqlite. Exit immediately if not in test environment or not pointed to sqlite. """ if config.ENVIRONMENT != config.TEST_ENVIRONMENT: sys.exit('Environment must be set to {}.'.format( config.TEST_ENVIRONMENT)) if 'sqlite' not in session.bind.url.drivername: sys.exit('Tests must point to sqlite database.')