"""Make Job.execution_id non-nullable

Revision ID: cd2d53191748
Revises: d4af139b6f53
Create Date: 2021-08-01 07:18:29.549984

"""

import uuid
from typing import Any

import sqlalchemy as sa
from alembic import op
from sqlalchemy import Column, ForeignKey, String, text
from sqlalchemy.orm import Session, declarative_base

# revision identifiers, used by Alembic.
revision = "cd2d53191748"
down_revision = "d4af139b6f53"
branch_labels = None
depends_on = None


Base: Any = declarative_base()


class Job(Base):
    __tablename__ = "job"

    id = Column(String, nullable=False, primary_key=True)


class Execution(Base):
    __tablename__ = "execution"

    id = Column(String, nullable=False, primary_key=True)
    args = Column(String, nullable=False)
    job_id = Column(String, ForeignKey("job.id", deferrable=True), nullable=False, index=True)


def upgrade():
    # Create stub executions for root jobs (parent_id is null) that have no execution.
    session = Session(bind=op.get_bind())
    try:
        job_ids = session.execute(
            text(
                """
            select job.id
            from job
            left join execution on execution.job_id = job.id
            where
              job.parent_id is null and
              execution.id is null
        """
            )
        ).fetchall()
        for (job_id,) in job_ids:
            session.add(Execution(id=str(uuid.uuid4()), args='"Stub Execution"', job_id=job_id))
        session.commit()
    finally:
        session.close()

    # Perform data migration to backfill Job.execution_id
    # using a recursive query through all parent_job pointers.
    op.execute(
        """
        -- Use a temp table to speed up the update.
        drop table if exists tmp_ancestors;
        """
    )

    op.execute(
        """
        create table tmp_ancestors as
        with recursive ancestors as (
            -- Get root jobs and their execution.
            select j.id as job_id, e.id as exec_id
            from job j
            join execution e on e.job_id = j.id
            union
            -- recurse.
            select j2.id, a.exec_id
            from job j2
            join ancestors a on a.job_id = j2.parent_id
        )
        select job_id, exec_id
        from ancestors;
        """
    )

    op.execute(
        """
        create index ix_tmp_ancestors_job_id on tmp_ancestors(job_id);
        """
    )

    op.execute(
        """
        update job
        set execution_id = (
            select a.exec_id
            from tmp_ancestors a
            where a.job_id = job.id
        );
        """
    )

    op.execute(
        """
        drop table tmp_ancestors;
        """
    )

    op.execute("commit")

    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("job", schema=None) as batch_op:
        batch_op.alter_column("execution_id", existing_type=sa.VARCHAR(), nullable=False)

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("job", schema=None) as batch_op:
        batch_op.alter_column("execution_id", existing_type=sa.VARCHAR(), nullable=True)

    # ### end Alembic commands ###
