"""Add contributed values table Revision ID: 159ba85908fd Revises: d5d88ac1d291 Create Date: 2019-11-01 15:39:50.970246 """ from alembic import op import sqlalchemy as sa from sqlalchemy import orm from sqlalchemy.dialects import postgresql from qcfractal.storage_sockets.models.sql_base import MsgpackExt from qcfractal.storage_sockets.models.collections_models import ContributedValuesORM import numpy as np # revision identifiers, used by Alembic. revision = "159ba85908fd" down_revision = "d5d88ac1d291" branch_labels = None depends_on = None def migrate_contributed_values_data(): bind = op.get_bind() session = orm.Session(bind=bind) # Dataset and reaction datasets tables ds_ids_data = session.execute("select id, contributed_values_data from dataset;").fetchall() print(f"Migrating datasets with ids: {[ds[0] for ds in ds_ids_data]}") rds_ids_data = session.execute("select id, contributed_values_data from reaction_dataset;").fetchall() print(f"Migrating reaction datasets with ids: {[ds[0] for ds in rds_ids_data]}") ds_ids_data.extend(rds_ids_data) for ds in ds_ids_data: (ds_id, ds_contrib) = ds if ds_contrib is None: continue for key, dict_values in ds_contrib.items(): idx, vals = [], [] for key, value in dict_values["values"].items(): idx.append(key) vals.append(value) dict_values["values"] = np.array(vals) dict_values["index"] = np.array(idx) cv = ContributedValuesORM(**dict_values) cv.collection_id = ds_id session.add(cv) session.commit() def upgrade(): # rename old column with data op.alter_column("dataset", "contributed_values", new_column_name="contributed_values_data") op.alter_column("reaction_dataset", "contributed_values", new_column_name="contributed_values_data") # ### commands auto generated by Alembic - please adjust! ### op.create_table( "contributed_values", sa.Column("name", sa.String(), nullable=False), sa.Column("collection_id", sa.Integer(), nullable=False), sa.Column("citations", sa.JSON(), nullable=True), sa.Column("theory_level", sa.JSON(), nullable=False), sa.Column("theory_level_details", sa.JSON(), nullable=True), sa.Column("comments", sa.String(), nullable=True), sa.Column("values", MsgpackExt(), nullable=False), sa.Column("index", MsgpackExt(), nullable=False), sa.Column("external_url", sa.String(), nullable=True), sa.Column("doi", sa.String(), nullable=True), sa.Column("units", sa.String(), nullable=False), sa.Column("values_structure", sa.JSON(), nullable=True, default=lambda: {}), sa.ForeignKeyConstraint(["collection_id"], ["collection.id"], ondelete="cascade"), sa.PrimaryKeyConstraint("name", "collection_id"), ) op.alter_column("contributed_values", "values_structure", server_default=None, nullable=False) migrate_contributed_values_data() op.drop_column("dataset", "contributed_values_data") op.drop_column("reaction_dataset", "contributed_values_data") def downgrade(): # ### Won't work on production data because data will be lost ### op.add_column( "reaction_dataset", sa.Column("contributed_values", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True), ) op.add_column( "dataset", sa.Column("contributed_values", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True), ) op.drop_table("contributed_values")