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