diff --git a/superset/migrations/versions/96e99fb176a0_add_import_mixing_to_saved_query.py b/superset/migrations/versions/96e99fb176a0_add_import_mixing_to_saved_query.py new file mode 100644 index 0000000000..2dc38e214a --- /dev/null +++ b/superset/migrations/versions/96e99fb176a0_add_import_mixing_to_saved_query.py @@ -0,0 +1,93 @@ +# 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") 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 bffa983c2c..5600212c82 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 @@ -86,10 +86,10 @@ add_uuids_by_dialect = { } -def add_uuids(table_name, session, batch_size=default_batch_size): +def add_uuids(model, table_name, session, batch_size=default_batch_size): """Populate columns with pre-computed uuids""" bind = op.get_bind() - objects_query = session.query(models[table_name]) + objects_query = session.query(model) count = objects_query.count() # silently skip if the table is empty (suitable for db initialization) @@ -168,7 +168,7 @@ def upgrade(): bind = op.get_bind() session = db.Session(bind=bind) - for table_name in models.keys(): + for table_name, model in models.items(): try: with op.batch_alter_table(table_name) as batch_op: batch_op.add_column( @@ -180,7 +180,7 @@ def upgrade(): # ignore collumn update errors so that we can run upgrade multiple times pass - add_uuids(table_name, session) + add_uuids(model, table_name, session) try: # add uniqueness constraint diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py index f7f8f86dbe..6c805b2a79 100644 --- a/superset/models/sql_lab.py +++ b/superset/models/sql_lab.py @@ -39,7 +39,7 @@ from sqlalchemy.engine.url import URL from sqlalchemy.orm import backref, relationship from superset import security_manager -from superset.models.helpers import AuditMixinNullable, ExtraJSONMixin +from superset.models.helpers import AuditMixinNullable, ExtraJSONMixin, ImportMixin from superset.models.tags import QueryUpdater from superset.sql_parse import CtasMethod, ParsedQuery, Table from superset.utils.core import QueryStatus, user_label @@ -162,7 +162,7 @@ class Query(Model, ExtraJSONMixin): security_manager.raise_for_access(query=self) -class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin): +class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin, ImportMixin): """ORM model for SQL query""" __tablename__ = "saved_query" @@ -186,6 +186,15 @@ class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin): rows = Column(Integer, nullable=True) last_run = Column(DateTime, nullable=True) + export_parent = "database" + export_fields = [ + "db_id", + "schema", + "label", + "description", + "sql", + ] + def __repr__(self) -> str: return str(self.label)