| # 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_import_mixing_to_saved_query |
| |
| Revision ID: 96e99fb176a0 |
| Revises: 585b0b1a7b18 |
| Create Date: 2020-10-21 21:09:55.945956 |
| |
| """ |
| |
| import os |
| from uuid import uuid4 |
| |
| import sqlalchemy as sa |
| from alembic import op |
| from sqlalchemy.exc import OperationalError |
| from sqlalchemy.ext.declarative import declarative_base |
| from sqlalchemy_utils import UUIDType |
| |
| from superset import db |
| from superset.migrations.versions.b56500de1855_add_uuid_column_to_import_mixin import ( |
| add_uuids, |
| ) |
| |
| # revision identifiers, used by Alembic. |
| revision = "96e99fb176a0" |
| down_revision = "585b0b1a7b18" |
| |
| |
| Base = declarative_base() |
| |
| |
| class ImportMixin: |
| id = sa.Column(sa.Integer, primary_key=True) |
| uuid = sa.Column(UUIDType(binary=True), primary_key=False, default=uuid4) |
| |
| |
| class SavedQuery(Base, ImportMixin): |
| __tablename__ = "saved_query" |
| |
| |
| default_batch_size = int(os.environ.get("BATCH_SIZE", 200)) |
| |
| |
| def upgrade(): |
| bind = op.get_bind() |
| session = db.Session(bind=bind) |
| |
| # Add uuid column |
| try: |
| with op.batch_alter_table("saved_query") as batch_op: |
| batch_op.add_column( |
| sa.Column( |
| "uuid", UUIDType(binary=True), primary_key=False, default=uuid4, |
| ), |
| ) |
| except OperationalError: |
| # Ignore column update errors so that we can run upgrade multiple times |
| pass |
| |
| add_uuids(SavedQuery, "saved_query", session) |
| |
| try: |
| # Add uniqueness constraint |
| with op.batch_alter_table("saved_query") as batch_op: |
| # Batch mode is required for sqllite |
| batch_op.create_unique_constraint("uq_saved_query_uuid", ["uuid"]) |
| except OperationalError: |
| pass |
| |
| |
| def downgrade(): |
| bind = op.get_bind() |
| session = db.Session(bind=bind) |
| |
| # Remove uuid column |
| with op.batch_alter_table("saved_query") as batch_op: |
| batch_op.drop_constraint("uq_saved_query_uuid", type_="unique") |
| batch_op.drop_column("uuid") |