af8b2e3c91b2_biosample_omics_association.py

"""Add many-to-many relationship between biosample and omics_processing

Revision ID: af8b2e3c91b2
Revises: dcc0a41b60af
Create Date: 2023-07-10 18:18:46.785564

"""

from typing import Optional

import sqlalchemy as sa
from alembic import op
from sqlalchemy import Column, ForeignKey, String, Table
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, relationship

Base = declarative_base()

# revision identifiers, used by Alembic.
revision: str = "af8b2e3c91b2"
down_revision: Optional[str] = "dcc0a41b60af"
branch_labels: Optional[str] = None
depends_on: Optional[str] = None


class Biosample(Base):
    __tablename__ = "biosample"
    id = Column(String, primary_key=True)


class OmicsProcessing(Base):
    __tablename__ = "omics_processing"
    id = Column(String, primary_key=True)
    biosample_id = Column(String, ForeignKey("biosample.id"), nullable=True)
    biosample = relationship("Biosample", backref="omics_processing")


biosample_input_association = Table(
    "biosample_input_association",
    Base.metadata,
    Column("biosample_id", String),
    Column("omics_processing_id", String),
)


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "biosample_input_association",
        sa.Column("omics_processing_id", sa.String(), nullable=False),
        sa.Column("biosample_id", sa.String(), nullable=False),
        sa.ForeignKeyConstraint(
            ["biosample_id"],
            ["biosample.id"],
            name=op.f("fk_biosample_input_association_biosample_id_biosample"),
        ),
        sa.ForeignKeyConstraint(
            ["omics_processing_id"],
            ["omics_processing.id"],
            name=op.f("fk_biosample_input_association_omics_processing_id_omics_processing"),
        ),
        sa.PrimaryKeyConstraint(
            "omics_processing_id", "biosample_id", name=op.f("pk_biosample_input_association")
        ),
    )

    session = Session(bind=op.get_bind())
    mappings = []
    for omics_processing in session.query(OmicsProcessing):
        if omics_processing.biosample_id:
            mappings.append(
                {
                    "biosample_id": omics_processing.biosample_id,
                    "omics_processing_id": omics_processing.id,
                }
            )
    op.bulk_insert(biosample_input_association, mappings)

    op.drop_constraint(
        "fk_omics_processing_biosample_id_biosample", "omics_processing", type_="foreignkey"
    )
    op.drop_column("omics_processing", "biosample_id")
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "omics_processing",
        sa.Column("biosample_id", sa.VARCHAR(), autoincrement=False, nullable=True),
    )
    op.create_foreign_key(
        "fk_omics_processing_biosample_id_biosample",
        "omics_processing",
        "biosample",
        ["biosample_id"],
        ["id"],
    )
    op.drop_table("biosample_input_association")
    # ### end Alembic commands ###