"""Utility functions to support database interactions in tests.""" import sys from functools import wraps from sqlalchemy.sql import text from product_staging import config from product_staging.connectors import db CREATE_BULK_SESSIONS_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_sessions ( bulk_session_id TEXT PRIMARY KEY, vendor_uuid TEXT NOT NULL, subaccount_id INT, metadata_status TEXT, asset_status TEXT, metadata_error_report_json TEXT, created_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, created_by TEXT NOT NULL, updated_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_by TEXT NOT NULL, asset_error_report TEXT, asset_report_json TEXT, is_cancelled BOOLEAN NOT NULL DEFAULT FALSE ); """) DROP_BULK_SESSIONS_TABLE = text(""" DROP TABLE IF EXISTS bulk_sessions; """) CREATE_BULK_SESSION_METADATA_FILES_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_metadata_files ( bulk_session_metadata_file_id CHAR(36) NOT NULL PRIMARY KEY, bulk_session_id CHAR(36) NOT NULL, file_status TEXT, upload_token TEXT, s3_filename VARCHAR(45), original_filename VARCHAR(127), created_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, created_by CHAR(36), json_file_name CHAR(55), total_products INT UNSIGNED NULL DEFAULT NULL ) """) DROP_BULK_SESSION_METADATA_FILES_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_metadata_files; """) CREATE_BULK_SESSION_ASSET_FILES_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_asset_files ( bulk_session_asset_file_id CHAR(36) NOT NULL PRIMARY KEY, bulk_session_id CHAR(36) NOT NULL, file_status TEXT, ingestion_status TEXT, latest_ingestion_execution_arn VARCHAR(255), file_size INT UNSIGNED, upload_token TEXT, s3_filename VARCHAR(90), original_filename VARCHAR(127), error_message TEXT, file_type TEXT, created_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, created_by CHAR(36), bulk_session_asset_cloud_transfer_job_id CHAR(36) ) """) DROP_BULK_SESSION_ASSET_FILES_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_asset_files; """) CREATE_BULK_SESSION_ASSETS_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_assets ( bulk_session_asset_id CHAR(36) NOT NULL PRIMARY KEY, bulk_session_id CHAR(36) NOT NULL, original_filename VARCHAR(127) NOT NULL, required BOOLEAN NOT NULL DEFAULT TRUE, bulk_session_asset_file_id CHAR(36) DEFAULT NULL, product_code VARCHAR(127) NOT NULL, volume INT UNSIGNED DEFAULT NULL, sequence_number INT UNSIGNED DEFAULT NULL, product_name VARCHAR(127) DEFAULT NULL, track_name VARCHAR(127) DEFAULT NULL )""") DROP_BULK_SESSION_ASSETS_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_assets; """) CREATE_BULK_SESSION_INGESTIONS_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_ingestions ( bulk_session_ingestion_id CHAR(36) NOT NULL PRIMARY KEY, bulk_session_id CHAR(36) NOT NULL, ingestion_status TEXT, assets_required BOOLEAN NOT NULL DEFAULT 0, submit_products BOOLEAN NOT NULL DEFAULT 0, send_notifications BOOLEAN NOT NULL DEFAULT 1, created_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, created_by TEXT NOT NULL, updated_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_by TEXT NOT NULL, completed_on TIMESTAMP NULL DEFAULT NULL, json_file_name CHAR(61) NULL, total_products INT UNSIGNED NULL DEFAULT NULL, ingestion_failure_file CHAR(55) NULL DEFAULT NULL, FOREIGN KEY (bulk_session_id) REFERENCES bulk_sessions (bulk_session_id) ) """) DROP_BULK_SESSION_INGESTIONS_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_ingestions; """) CREATE_BULK_SESSION_INGESTION_PRODUCTS_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_ingestion_products ( bulk_session_ingestion_product_id CHAR(36) NOT NULL PRIMARY KEY, bulk_session_ingestion_project_id CHAR(36) NULL, bulk_session_ingestion_id CHAR(36) NOT NULL, product_id INT UNSIGNED NOT NULL, upc CHAR(20) NOT NULL, product_code CHAR(128) NOT NULL, submit_status CHAR(20) NOT NULL DEFAULT 'default', ingestion_status CHAR(20) NOT NULL DEFAULT 'in_progress', metadata_ingestion_status CHAR(20) DEFAULT NULL, created_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (bulk_session_ingestion_id) REFERENCES bulk_session_ingestions (bulk_session_ingestion_id), FOREIGN KEY (bulk_session_ingestion_project_id) REFERENCES bulk_session_ingestion_projects (bulk_session_ingestion_project_id) ) """) DROP_BULK_SESSION_INGESTION_PRODUCTS_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_ingestion_products; """) CREATE_BULK_SESSION_INGESTION_TRACKS_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_ingestion_tracks ( bulk_session_ingestion_track_id CHAR(36) NOT NULL PRIMARY KEY, bulk_session_ingestion_product_id CHAR(36) NOT NULL, volume INT UNSIGNED NOT NULL, sequence_number INT UNSIGNED NOT NULL, track_id INT UNSIGNED, isrc CHAR(16), ingestion_status CHAR(20) NOT NULL DEFAULT 'in_progress', created_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (bulk_session_ingestion_product_id) REFERENCES bulk_session_ingestion_products (bulk_session_ingestion_product_id) ) """) DROP_BULK_SESSION_INGESTION_TRACKS_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_ingestion_tracks; """) CREATE_BULK_SESSION_INGESTION_PROJECTS_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_ingestion_projects ( bulk_session_ingestion_project_id CHAR(36) NOT NULL PRIMARY KEY, bulk_session_ingestion_id CHAR(36) NOT NULL, vendor_uuid CHAR(36) NOT NULL, project_code VARCHAR(20) NOT NULL, project_id INT NULL, ingestion_status CHAR(20) NOT NULL DEFAULT 'in_progress', created_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (bulk_session_ingestion_id) REFERENCES bulk_session_ingestions (bulk_session_ingestion_id) ) """) DROP_BULK_SESSION_INGESTION_PROJECTS_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_ingestion_projects; """) CREATE_BULK_SESSION_INGESTION_EXECUTIONS_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_ingestion_executions ( bulk_session_ingestion_execution_id CHAR(36) NOT NULL PRIMARY KEY, bulk_session_ingestion_id CHAR(36) NOT NULL, execution_arn VARCHAR(255) NOT NULL, created_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (bulk_session_ingestion_id) REFERENCES bulk_session_ingestions (bulk_session_ingestion_id) ) """) DROP_BULK_SESSION_INGESTION_EXECUTIONS_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_ingestion_executions; """) CREATE_BULK_SESSION_INGESTION_EXECUTION_LOG_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_ingestion_execution_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, bulk_session_ingestion_execution_id CHAR(36) NOT NULL, identity_id CHAR(36) NOT NULL, log_message VARCHAR(1024) NOT NULL, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) DROP_BULK_SESSION_INGESTION_EXECUTION_LOG_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_ingestion_execution_log; """) CREATE_BULK_SESSION_ASSET_CLOUD_TRANSFER_JOBS_TABLE = text(""" CREATE TABLE IF NOT EXISTS bulk_session_asset_cloud_transfer_jobs ( bulk_session_asset_cloud_transfer_job_id CHAR(36) PRIMARY KEY NOT NULL, bulk_session_id CHAR(36) NOT NULL, execution_arn VARCHAR(255) DEFAULT NULL, assets_info_json CHAR(57), total_assets INT DEFAULT 0, job_status TEXT DEFAULT NULL, job_type TEXT NOT NULL, created_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP, created_by CHAR(36), FOREIGN KEY (bulk_session_id) REFERENCES bulk_sessions (bulk_session_id) ); """) DROP_BULK_SESSION_ASSET_CLOUD_TRANSFER_JOBS_TABLE = text(""" DROP TABLE IF EXISTS bulk_session_asset_cloud_transfer_jobs; """) async def execute_sql(*args): """Run a series of queries safely within the test environment.""" async with db.db_session() as session: _exit_if_not_test_environment(session) for query in args: await session.execute(query) 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. Individual test cases can use factories to seed data as needed. Args: function (func): function to be called after creating the test schema. Returns: Function: The decorated function. """ @wraps(function) async def call_function_within_db_context(*args, **kwargs): await execute_sql( CREATE_BULK_SESSIONS_TABLE, CREATE_BULK_SESSION_METADATA_FILES_TABLE, CREATE_BULK_SESSION_ASSET_FILES_TABLE, CREATE_BULK_SESSION_ASSETS_TABLE, CREATE_BULK_SESSION_INGESTIONS_TABLE, CREATE_BULK_SESSION_INGESTION_PRODUCTS_TABLE, CREATE_BULK_SESSION_INGESTION_TRACKS_TABLE, CREATE_BULK_SESSION_INGESTION_PROJECTS_TABLE, CREATE_BULK_SESSION_INGESTION_EXECUTIONS_TABLE, CREATE_BULK_SESSION_INGESTION_EXECUTION_LOG_TABLE, CREATE_BULK_SESSION_ASSET_CLOUD_TRANSFER_JOBS_TABLE, ) try: function_return = await function(*args, **kwargs) finally: await execute_sql( DROP_BULK_SESSION_ASSET_CLOUD_TRANSFER_JOBS_TABLE, DROP_BULK_SESSION_INGESTION_EXECUTION_LOG_TABLE, DROP_BULK_SESSION_INGESTION_EXECUTIONS_TABLE, DROP_BULK_SESSION_INGESTION_PROJECTS_TABLE, DROP_BULK_SESSION_INGESTION_TRACKS_TABLE, DROP_BULK_SESSION_INGESTION_PRODUCTS_TABLE, DROP_BULK_SESSION_INGESTIONS_TABLE, DROP_BULK_SESSION_ASSETS_TABLE, DROP_BULK_SESSION_ASSET_FILES_TABLE, DROP_BULK_SESSION_METADATA_FILES_TABLE, DROP_BULK_SESSIONS_TABLE, ) return function_return return call_function_within_db_context 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.") @db.db_session_wrap async def create_multiple(instances: list[db.base_model], session=None): assert session session.add_all(instances) await session.flush()