"""Add table for namespaces and replace namespace field in module provider with foreign key

Revision ID: d72c5339c21b
Revises: acd5e83c690f
Create Date: 2022-09-28 06:14:12.046484

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = 'd72c5339c21b'
down_revision = 'acd5e83c690f'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('namespace',
        sa.Column('id', sa.Integer(), nullable=False, autoincrement=True),
        sa.Column('namespace', sa.String(length=128), nullable=True),
        sa.PrimaryKeyConstraint('id')
    )

    with op.batch_alter_table('module_provider') as module_provider_batch:
        module_provider_batch.add_column(sa.Column('namespace_id', sa.Integer(), nullable=True))
        module_provider_batch.create_foreign_key('fk_module_provider_namespace_id_namespace_id', 'namespace', ['namespace_id'], ['id'], onupdate='CASCADE', ondelete='CASCADE')

    # Migrate namespace column of module provider to new namespace table
    c = op.get_bind()
    module_providers = c.execute(f"""SELECT id, namespace FROM module_provider""")
    created_namespaces = {}
    for row in module_providers:

        module_provider_id, namespace_name = row
        # If namespace has not yet been created, create it and store
        # ID in created namespaces
        if namespace_name not in created_namespaces:
            namespace_insert_res = c.execute(
                sa.sql.text("""
                    INSERT INTO namespace(namespace) VALUES(:namespace_name)"""
                ),
                namespace_name=namespace_name,
            )
            created_namespaces[namespace_name] = namespace_insert_res.inserted_primary_key[0]

        # Update module provider namespace ID
        c.execute(
            sa.sql.text("""UPDATE module_provider SET namespace_id=:namespace_id WHERE id=:module_provider_id"""),
            namespace_id=created_namespaces[namespace_name],
            module_provider_id=module_provider_id
        )

    with op.batch_alter_table('module_provider') as module_provider_batch:
        module_provider_batch.drop_column('namespace')
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('module_provider') as module_provider_batch:
        module_provider_batch.add_column(sa.Column('namespace', sa.VARCHAR(length=1024), nullable=True))

    c = op.get_bind()
    c.execute(
        sa.sql.text("""
            UPDATE module_provider SET namespace=(
                SELECT namespace FROM namespace WHERE id=module_provider.namespace_id
            )
        """)
    )

    with op.batch_alter_table('module_provider') as module_provider_batch:
        module_provider_batch.drop_constraint('fk_module_provider_namespace_id_namespace_id', type_='foreignkey')
        module_provider_batch.drop_column('namespace_id')

    op.drop_table('namespace')
    # ### end Alembic commands ###
