9027a1881082_unique_sirets.py

"""unique-sirets

Revision ID: 9027a1881082
Revises: aba5d96052ab
Create Date: 2021-09-27 15:41:27.345669

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = "9027a1881082"
down_revision = "aba5d96052ab"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index("company_siren_key", "company")
    op.add_column("company", sa.Column("short_sirets", sa.ARRAY(sa.Integer())))
    op.execute("CREATE EXTENSION IF NOT EXISTS intarray")
    op.execute(
        """
        UPDATE company SET short_sirets = a.short_sirets FROM (
            SELECT id,
                array_agg(substring(siret, 10)::int) AS short_sirets
            FROM (SELECT id, unnest(sirets) as siret FROM company) b
            GROUP BY id
        ) a WHERE a.id = company.id;
        """
    )
    op.execute("CREATE EXTENSION IF NOT EXISTS btree_gist")
    op.execute(
        """
        ALTER TABLE company ADD CONSTRAINT only_one_company_per_siret
        EXCLUDE USING GIST (
            siren WITH =,
            short_sirets WITH &&
        )
        """
    )
    op.drop_column("company", "sirets")
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_unique_constraint("company_siren_key", "company", ["siren"])
    op.add_column(
        "company", sa.Column("sirets", sa.ARRAY(sa.VARCHAR(length=14)))
    )
    op.execute(
        """
        UPDATE company SET sirets = a.sirets FROM (
            SELECT id,
                array_agg(siren::text || short_siret::text) AS sirets
            FROM (SELECT id, siren, unnest(short_sirets) as short_siret FROM company) b
            GROUP BY id, siren
        ) a WHERE a.id = company.id;
        """
    )
    op.drop_constraint("only_one_company_per_siret", "company")
    op.drop_column("company", "short_sirets")
    # ### end Alembic commands ###