From 9734ade7612ac1507ed2b0c169743b3d63ee3d9e Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Thu, 18 Feb 2021 09:41:07 -0800 Subject: [PATCH] fix: add missing columns --- ...0de1855_add_uuid_column_to_import_mixin.py | 28 ++---- .../c501b7c653a3_add_missing_uuid_column.py | 96 +++++++++++++++++++ 2 files changed, 106 insertions(+), 18 deletions(-) create mode 100644 superset/migrations/versions/c501b7c653a3_add_missing_uuid_column.py diff --git a/superset/migrations/versions/b56500de1855_add_uuid_column_to_import_mixin.py b/superset/migrations/versions/b56500de1855_add_uuid_column_to_import_mixin.py index 492930088813b..7bef33bc87260 100644 --- a/superset/migrations/versions/b56500de1855_add_uuid_column_to_import_mixin.py +++ b/superset/migrations/versions/b56500de1855_add_uuid_column_to_import_mixin.py @@ -31,7 +31,6 @@ from alembic import op from sqlalchemy.dialects.mysql.base import MySQLDialect from sqlalchemy.dialects.postgresql.base import PGDialect -from sqlalchemy.exc import OperationalError from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import load_only from sqlalchemy_utils import UUIDType @@ -169,26 +168,19 @@ def upgrade(): session = db.Session(bind=bind) for table_name, model in models.items(): - try: - with op.batch_alter_table(table_name) as batch_op: - batch_op.add_column( - sa.Column( - "uuid", UUIDType(binary=True), primary_key=False, default=uuid4, - ), - ) - except OperationalError: - # ignore collumn update errors so that we can run upgrade multiple times - pass + with op.batch_alter_table(table_name) as batch_op: + batch_op.add_column( + sa.Column( + "uuid", UUIDType(binary=True), primary_key=False, default=uuid4, + ), + ) add_uuids(model, table_name, session) - try: - # add uniqueness constraint - with op.batch_alter_table(table_name) as batch_op: - # batch mode is required for sqllite - batch_op.create_unique_constraint(f"uq_{table_name}_uuid", ["uuid"]) - except OperationalError: - pass + # add uniqueness constraint + with op.batch_alter_table(table_name) as batch_op: + # batch mode is required for sqllite + batch_op.create_unique_constraint(f"uq_{table_name}_uuid", ["uuid"]) # add UUID to Dashboard.position_json slice_uuid_map = { diff --git a/superset/migrations/versions/c501b7c653a3_add_missing_uuid_column.py b/superset/migrations/versions/c501b7c653a3_add_missing_uuid_column.py new file mode 100644 index 0000000000000..f8b252ed0cba3 --- /dev/null +++ b/superset/migrations/versions/c501b7c653a3_add_missing_uuid_column.py @@ -0,0 +1,96 @@ +# 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. +"""add missing uuid column + +Revision ID: c501b7c653a3 +Revises: 070c043f2fdb +Create Date: 2021-02-18 09:13:00.028317 + +""" + +# revision identifiers, used by Alembic. +revision = "c501b7c653a3" +down_revision = "070c043f2fdb" + +import logging +from uuid import uuid4 + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import mysql +from sqlalchemy.engine.reflection import Inspector +from sqlalchemy.orm import load_only +from sqlalchemy_utils import UUIDType + +from superset import db +from superset.migrations.versions.b56500de1855_add_uuid_column_to_import_mixin import ( + add_uuids, + models, + update_dashboards, +) + + +def has_uuid_column(table_name, bind): + inspector = Inspector.from_engine(bind) + columns = {column["name"] for column in inspector.get_columns(table_name)} + has_uuid_column = "uuid" in columns + if has_uuid_column: + logging.info("Table %s already has uuid column, skipping...", table_name) + else: + logging.info("Table %s doesn't have uuid column, adding...", table_name) + return has_uuid_column + + +def upgrade(): + bind = op.get_bind() + session = db.Session(bind=bind) + + for table_name, model in models.items(): + # this script adds missing uuid columns + if has_uuid_column(table_name, bind): + continue + + with op.batch_alter_table(table_name) as batch_op: + batch_op.add_column( + sa.Column( + "uuid", UUIDType(binary=True), primary_key=False, default=uuid4, + ), + ) + add_uuids(model, table_name, session) + + # add uniqueness constraint + with op.batch_alter_table(table_name) as batch_op: + # batch mode is required for sqllite + batch_op.create_unique_constraint(f"uq_{table_name}_uuid", ["uuid"]) + + # add UUID to Dashboard.position_json; this function is idempotent + # so we can call it for all objects + slice_uuid_map = { + slc.id: slc.uuid + for slc in session.query(models["slices"]) + .options(load_only("id", "uuid")) + .all() + } + update_dashboards(session, slice_uuid_map) + + +def downgrade() -> None: + """ + This script fixes b56500de1855_add_uuid_column_to_import_mixin.py by adding any + uuid columns that might have been skipped. There's no downgrade. + """ + pass