feat: add UUID column to saved_query for export/import (#11397)

* Add UUID to saved_query

* Reuse function from previous migration

* Point to new head

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")

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

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)