diff --git a/migrations/versions/14c78107b748_new_bank_migration.py b/migrations/versions/14c78107b748_new_bank_migration.py index a544c8ed..e2d1449f 100644 --- a/migrations/versions/14c78107b748_new_bank_migration.py +++ b/migrations/versions/14c78107b748_new_bank_migration.py @@ -9,6 +9,9 @@ import sqlalchemy as sa from datetime import datetime +from bank.system import SlurmAccount +from bank.exceptions import AccountNotFoundError + DATE_FORMAT = "%Y-%m-%d" # revision identifiers, used by Alembic. @@ -23,11 +26,6 @@ def upgrade(): conn = op.get_bind() - # Determine total number of proposals and investments in old db schema - num_proposals = conn.execute("SELECT count(*) FROM proposal").fetchall() - num_prop_archive = conn.execute("SELECT count(*) FROM proposal_archive").fetchall() - num_proposals_old = num_proposals[0][0] + num_prop_archive[0][0] - num_investments = conn.execute("SELECT count(*) FROM investor").fetchall() num_inv_archive = conn.execute("SELECT count(*) FROM investor_archive").fetchall() num_investments_old = num_investments[0][0] + num_inv_archive[0][0] @@ -41,8 +39,8 @@ def upgrade(): "SET percent_notified=100 " "WHERE percent_notified is NULL") - # Drop unused column and rename tabled - op.drop_column('proposal', 'proposal_type') + # Drop the proposal archive and rename the combined table + op.drop_table('proposal_archive') op.rename_table('proposal', '_proposal_old') # Concatenate investor tabled with archive @@ -59,14 +57,10 @@ def upgrade(): "SET rollover_sus=-9 " "WHERE rollover_sus is NULL") - # Drop unused column and rename table - op.drop_column('investor', 'proposal_type') + # Drop the investor archive and renamed the combined table + op.drop_table('investor_archive') op.rename_table('investor', '_investment_old') - # Drop old tables after concatenation - for table in ('investor_archive', 'proposal_archive'): - op.drop_table(table) - # Create Account Table account_table = op.create_table( 'account', @@ -119,18 +113,37 @@ def upgrade(): # For each account, create new schema proposals and investments from entries in old schema accounts = conn.execute("SELECT * from account").fetchall() for account in accounts: + + # Attempt to set up the SLURM account, skipping if it does not exist + try: + slurm_acct = SlurmAccount(account.name) + except AccountNotFoundError: + print(f"SLURM Account for {account.name} does not exist, skipping...\n") + continue + old_investments = conn.execute(f"SELECT * from _investment_old WHERE _investment_old.account='{account.name}'").fetchall() old_proposals = conn.execute(f"SELECT * from _proposal_old WHERE _proposal_old.account='{account.name}'").fetchall() new_proposals = [] for prop in old_proposals: + # Determine whether the proposal has more than a standard allocation, creating a floating alloc if not allocs = [] + total_sus = 0 + for cluster_name in ['smp','mpi','gpu','htc']: + total_sus += getattr(prop, cluster_name) allocs.append({'proposal_id': prop.id, 'cluster_name': cluster_name, - 'service_units_total': getattr(prop, cluster_name)} + 'service_units_total': getattr(prop, cluster_name), + 'service_units_used': slurm_acct.get_cluster_usage_total(cluster_name, in_hours=True)} ) + # Empty the allocations and establish a standard allocation + if total_sus <= 25000: + allocs = [{'proposal_id': prop.id, + 'cluster_name': 'all_clusters', + 'service_units_total': getattr(prop, cluster_name)}] + new_proposal = {'id': prop.id, 'account_id': account.id, 'start_date': datetime.strptime(prop.start_date, DATE_FORMAT), @@ -156,10 +169,6 @@ def upgrade(): op.bulk_insert(investment_table, new_investments) - # Make sure there is the same number of proposals/investments as the old db schema - num_new_proposals = conn.execute("SELECT count(*) FROM proposal").fetchall() - assert num_new_proposals[0][0] == num_proposals_old, "The number of new proposals must equal the number of old proposals" - num_new_investments = conn.execute("SELECT count(*) FROM investment").fetchall() assert num_new_investments[0][0] == num_investments_old, "The number of new investments must equal the number of old investments" diff --git a/migrations/versions/a6c26a203bd1_system_refactor_migration.py b/migrations/versions/a6c26a203bd1_system_refactor_migration.py deleted file mode 100644 index 02bb9633..00000000 --- a/migrations/versions/a6c26a203bd1_system_refactor_migration.py +++ /dev/null @@ -1,109 +0,0 @@ -"""System refactor migration - -Revision ID: a6c26a203bd1 -Revises: -Create Date: 2022-02-21 15:01:03.939907 -""" - -import sqlalchemy as sa -from alembic import op - -# revision identifiers, used by Alembic. -revision = 'a6c26a203bd1' -down_revision = None -branch_labels = None -depends_on = None - -# New enum value introduced for proposal types -PROPOSAL_TYPE_ENUM_NAME = 'proposalenum' -PROPOSAL_TYPE = 'Proposal' -CLASS_TYPE = 'Class' -UNKNOWN_TYPE = 'Unknown' - -# Map old percent notified enum values to their new int representations -PERCENT_NOTIFIED_MAPPER = { - 0: 0, - 1: 25, - 2: 50, - 3: 75, - 4: 90, - 5: 100} - - -def upgrade(): - """Upgrade the database schema to the next version""" - - proposal_type_col_type = sa.Enum(UNKNOWN_TYPE, PROPOSAL_TYPE, CLASS_TYPE, name=PROPOSAL_TYPE_ENUM_NAME) - with op.batch_alter_table("proposal", recreate='always') as proposal_table: - proposal_table.alter_column('account', new_column_name='account_name', existing_type=sa.TEXT(), type_=sa.String(), nullable=False) - proposal_table.alter_column('start', existing_type=sa.DATE(), nullable=False) - proposal_table.alter_column('end', existing_type=sa.DATE(), nullable=False) - proposal_table.alter_column('percent_notified', existing_type=sa.INTEGER(), nullable=False) - proposal_table.alter_column('proposal_type', existing_type=sa.INTEGER(), type_=proposal_type_col_type, nullable=True) - - # Change percent notified values from enum type values to actual notification percentage - op.execute( - 'UPDATE proposal SET proposal_type = (CASE ' - f'WHEN proposal_type = 0 THEN "{PROPOSAL_TYPE}" ' - f'WHEN proposal_type = 1 THEN "{CLASS_TYPE}" ' - 'END)') - - # Change percent notified values from enum type values to actual notification percentage - op.execute( - 'UPDATE proposal SET percent_notified = (CASE ' - f'WHEN percent_notified = 0 THEN {PERCENT_NOTIFIED_MAPPER[0]} ' - f'WHEN percent_notified = 1 THEN {PERCENT_NOTIFIED_MAPPER[1]} ' - f'WHEN percent_notified = 2 THEN {PERCENT_NOTIFIED_MAPPER[2]} ' - f'WHEN percent_notified = 3 THEN {PERCENT_NOTIFIED_MAPPER[3]} ' - f'WHEN percent_notified = 4 THEN {PERCENT_NOTIFIED_MAPPER[4]} ' - f'WHEN percent_notified = 5 THEN {PERCENT_NOTIFIED_MAPPER[5]} ' - 'END)') - - with op.batch_alter_table("proposal_archive", recreate='always') as p_archive_table: - p_archive_table.alter_column('account', new_column_name='account_name', existing_type=sa.TEXT(), type_=sa.String(), nullable=False) - p_archive_table.alter_column('end', existing_type=sa.DATE(), nullable=False) - p_archive_table.alter_column('gpu', existing_type=sa.INTEGER(), nullable=False) - p_archive_table.alter_column('gpu_usage', existing_type=sa.INTEGER(), nullable=False) - p_archive_table.alter_column('htc', existing_type=sa.INTEGER(), nullable=False) - p_archive_table.alter_column('htc_usage', existing_type=sa.INTEGER(), nullable=False) - p_archive_table.alter_column('mpi', existing_type=sa.INTEGER(), nullable=False) - p_archive_table.alter_column('mpi_usage', existing_type=sa.INTEGER(), nullable=False) - p_archive_table.alter_column('smp', existing_type=sa.INTEGER(), nullable=False) - p_archive_table.alter_column('smp_usage', existing_type=sa.INTEGER(), nullable=False) - p_archive_table.alter_column('start', existing_type=sa.DATE(), nullable=False) - p_archive_table.add_column(sa.Column('proposal_type', proposal_type_col_type, nullable=True)) - - # The original table does not track the proposal type, so we fill in missing values before settings nullable=False - op.execute(f'UPDATE proposal_archive SET proposal_type = "{UNKNOWN_TYPE}"') - with op.batch_alter_table("proposal_archive", recreate='always') as p_archive_table: - p_archive_table.alter_column('proposal_type', nullable=False) - - with op.batch_alter_table("investor", recreate='always') as investor_table: - investor_table.alter_column('account', new_column_name='account_name', existing_type=sa.TEXT(), type_=sa.String(), nullable=False) - investor_table.alter_column('current_sus', existing_type=sa.INTEGER(), nullable=False) - investor_table.alter_column('end', existing_type=sa.DATE(), nullable=False) - investor_table.alter_column('rollover_sus', type=sa.INTEGER(), nullable=False) - investor_table.alter_column('service_units', existing_type=sa.INTEGER(), nullable=False) - investor_table.alter_column('start', existing_type=sa.DATE(), nullable=False) - investor_table.alter_column('withdrawn_sus', existing_type=sa.INTEGER(), nullable=False) - investor_table.drop_column('proposal_type') - - with op.batch_alter_table("investor_archive", recreate='always') as inv_archive_table: - inv_archive_table.alter_column('account', new_column_name='account_name', existing_type=sa.TEXT(), type_=sa.String(), nullable=False) - inv_archive_table.alter_column('current_sus', existing_type=sa.INTEGER(), nullable=False) - inv_archive_table.alter_column('end', existing_type=sa.DATE(), nullable=False) - inv_archive_table.alter_column('exhaustion_date', existing_type=sa.DATE(), nullable=False) - inv_archive_table.alter_column('service_units', existing_type=sa.INTEGER(), nullable=False) - inv_archive_table.alter_column('start', existing_type=sa.DATE(), nullable=False) - inv_archive_table.drop_column('proposal_id') - inv_archive_table.drop_column('investor_id') - - -def downgrade(): - """Downgrade the database schema to the previous version""" - - raise NotImplementedError( - 'Rollbacks for this version are not supported. ' - 'Data has been dropped in the upgrade that cannot be recovered. ' - 'Restore a file system snapshot instead.' - )