1"""empty message 2 3Revision ID: v1.3.0.a 4Revises: v1.2.0.a 5Create Date: 2020-02-14 16:23:04.800808 6 7""" 8import json 9 10from alembic import op 11import sqlalchemy as sa 12 13from sqlalchemy.exc import SQLAlchemyError 14from sqlalchemy.ext.declarative import declarative_base 15from sqlalchemy import orm 16 17# revision identifiers, used by Alembic. 18revision = "v1.3.0.a" 19down_revision = "v1.2.0.a" 20branch_labels = None 21depends_on = None 22 23# Model definition 24MAX_INDEXED_STRING_LENGTH = 512 25MAX_STRING_LENGTH = 2048 26BaseModel = declarative_base() 27 28 29class TrialModel(BaseModel): 30 __tablename__ = "trials" 31 trial_id = sa.Column(sa.Integer, primary_key=True) 32 number = sa.Column(sa.Integer) 33 34 35class TrialSystemAttributeModel(BaseModel): 36 __tablename__ = "trial_system_attributes" 37 trial_system_attribute_id = sa.Column(sa.Integer, primary_key=True) 38 trial_id = sa.Column(sa.Integer, sa.ForeignKey("trials.trial_id")) 39 key = sa.Column(sa.String(MAX_INDEXED_STRING_LENGTH)) 40 value_json = sa.Column(sa.String(MAX_STRING_LENGTH)) 41 42 43def upgrade(): 44 bind = op.get_bind() 45 session = orm.Session(bind=bind) 46 47 with op.batch_alter_table("trials") as batch_op: 48 batch_op.add_column(sa.Column("number", sa.Integer(), nullable=True, default=None)) 49 50 try: 51 number_records = ( 52 session.query(TrialSystemAttributeModel) 53 .filter(TrialSystemAttributeModel.key == "_number") 54 .all() 55 ) 56 mapping = [ 57 {"trial_id": r.trial_id, "number": json.loads(r.value_json)} for r in number_records 58 ] 59 session.bulk_update_mappings(TrialModel, mapping) 60 61 stmt = ( 62 sa.delete(TrialSystemAttributeModel) 63 .where(TrialSystemAttributeModel.key == "_number") 64 .execution_options(synchronize_session=False) 65 ) 66 session.execute(stmt) 67 session.commit() 68 except SQLAlchemyError as e: 69 session.rollback() 70 raise e 71 finally: 72 session.close() 73 74 75def downgrade(): 76 bind = op.get_bind() 77 session = orm.Session(bind=bind) 78 79 try: 80 number_attrs = [] 81 trials = session.query(TrialModel).all() 82 for trial in trials: 83 number_attrs.append( 84 TrialSystemAttributeModel( 85 trial_id=trial.trial_id, key="_number", value_json=json.dumps(trial.number) 86 ) 87 ) 88 session.bulk_save_objects(number_attrs) 89 session.commit() 90 except SQLAlchemyError as e: 91 session.rollback() 92 raise e 93 finally: 94 session.close() 95 96 with op.batch_alter_table("trials") as batch_op: 97 batch_op.drop_column("number") 98