superset2/superset/migrations/versions/4736ec66ce19_.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""empty message

Revision ID: 4736ec66ce19
Revises: f959a6652acd
Create Date: 2017-10-03 14:37:01.376578

"""

import logging

import sqlalchemy as sa
from alembic import op

from superset.utils.core import (
    generic_find_fk_constraint_name,
    generic_find_fk_constraint_names,
    generic_find_uq_constraint_name,
)

# revision identifiers, used by Alembic.
revision = "4736ec66ce19"
down_revision = "f959a6652acd"

conv = {
    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    "uq": "uq_%(table_name)s_%(column_0_name)s",
}

# Helper table for database migrations using minimal schema.
datasources = sa.Table(
    "datasources",
    sa.MetaData(),
    sa.Column("id", sa.Integer, primary_key=True),
    sa.Column("datasource_name", sa.String(255)),
)


def upgrade():
    bind = op.get_bind()
    insp = sa.engine.reflection.Inspector.from_engine(bind)

    # Add the new less restrictive uniqueness constraint.
    with op.batch_alter_table("datasources", naming_convention=conv) as batch_op:
        batch_op.create_unique_constraint(
            "uq_datasources_cluster_name", ["cluster_name", "datasource_name"]
        )

    # Augment the tables which have a foreign key constraint related to the
    # datasources.datasource_name column.
    for foreign in ["columns", "metrics"]:
        with op.batch_alter_table(foreign, naming_convention=conv) as batch_op:

            # Add the datasource_id column with the relevant constraints.
            batch_op.add_column(sa.Column("datasource_id", sa.Integer))

            batch_op.create_foreign_key(
                "fk_{}_datasource_id_datasources".format(foreign),
                "datasources",
                ["datasource_id"],
                ["id"],
            )

        # Helper table for database migration using minimal schema.
        table = sa.Table(
            foreign,
            sa.MetaData(),
            sa.Column("id", sa.Integer, primary_key=True),
            sa.Column("datasource_name", sa.String(255)),
            sa.Column("datasource_id", sa.Integer),
        )

        # Migrate the existing data.
        for datasource in bind.execute(datasources.select()):
            bind.execute(
                table.update()
                .where(table.c.datasource_name == datasource.datasource_name)
                .values(datasource_id=datasource.id)
            )

        with op.batch_alter_table(foreign, naming_convention=conv) as batch_op:

            # Drop the datasource_name column and associated constraints. Note
            # due to prior revisions (1226819ee0e3, 3b626e2a6783) there may
            # incorectly be multiple duplicate constraints.
            names = generic_find_fk_constraint_names(
                foreign, {"datasource_name"}, "datasources", insp
            )

            for name in names:
                batch_op.drop_constraint(
                    name or "fk_{}_datasource_name_datasources".format(foreign),
                    type_="foreignkey",
                )

            batch_op.drop_column("datasource_name")

    try:
        # Drop the old more restrictive uniqueness constraint.
        with op.batch_alter_table("datasources", naming_convention=conv) as batch_op:
            batch_op.drop_constraint(
                generic_find_uq_constraint_name(
                    "datasources", {"datasource_name"}, insp
                )
                or "uq_datasources_datasource_name",
                type_="unique",
            )
    except Exception as ex:
        logging.warning(
            "Constraint drop failed, you may want to do this "
            "manually on your database. For context, this is a known "
            "issue around undeterministic contraint names on Postgres "
            "and perhaps more databases through SQLAlchemy."
        )
        logging.exception(ex)


def downgrade():
    bind = op.get_bind()
    insp = sa.engine.reflection.Inspector.from_engine(bind)

    # Add the new more restrictive uniqueness constraint which is required by
    # the foreign key constraints. Note this operation will fail if the
    # datasources.datasource_name column is no longer unique.
    with op.batch_alter_table("datasources", naming_convention=conv) as batch_op:
        batch_op.create_unique_constraint(
            "uq_datasources_datasource_name", ["datasource_name"]
        )

    # Augment the tables which have a foreign key constraint related to the
    # datasources.datasource_id column.
    for foreign in ["columns", "metrics"]:
        with op.batch_alter_table(foreign, naming_convention=conv) as batch_op:

            # Add the datasource_name column with the relevant constraints.
            batch_op.add_column(sa.Column("datasource_name", sa.String(255)))

            batch_op.create_foreign_key(
                "fk_{}_datasource_name_datasources".format(foreign),
                "datasources",
                ["datasource_name"],
                ["datasource_name"],
            )

        # Helper table for database migration using minimal schema.
        table = sa.Table(
            foreign,
            sa.MetaData(),
            sa.Column("id", sa.Integer, primary_key=True),
            sa.Column("datasource_name", sa.String(255)),
            sa.Column("datasource_id", sa.Integer),
        )

        # Migrate the existing data.
        for datasource in bind.execute(datasources.select()):
            bind.execute(
                table.update()
                .where(table.c.datasource_id == datasource.id)
                .values(datasource_name=datasource.datasource_name)
            )

        with op.batch_alter_table(foreign, naming_convention=conv) as batch_op:

            # Drop the datasource_id column and associated constraint.
            batch_op.drop_constraint(
                "fk_{}_datasource_id_datasources".format(foreign), type_="foreignkey"
            )

            batch_op.drop_column("datasource_id")

    with op.batch_alter_table("datasources", naming_convention=conv) as batch_op:

        # Prior to dropping the uniqueness constraint, the foreign key
        # associated with the cluster_name column needs to be dropped.
        batch_op.drop_constraint(
            generic_find_fk_constraint_name(
                "datasources", {"cluster_name"}, "clusters", insp
            )
            or "fk_datasources_cluster_name_clusters",
            type_="foreignkey",
        )

        # Drop the old less restrictive uniqueness constraint.
        batch_op.drop_constraint(
            generic_find_uq_constraint_name(
                "datasources", {"cluster_name", "datasource_name"}, insp
            )
            or "uq_datasources_cluster_name",
            type_="unique",
        )

        # Re-create the foreign key associated with the cluster_name column.
        batch_op.create_foreign_key(
            "fk_{}_datasource_id_datasources".format(foreign),
            "clusters",
            ["cluster_name"],
            ["cluster_name"],
        )