feat: add endpoint to export saved queries using new format (#11447)

* Add UUID to saved_query

* Reuse function from previous migration

* Point to new head

* feat: add backend to export saved queries using new format

superset/charts/api.py

@@ -175,6 +175,7 @@ class ChartRestApi(BaseSupersetModelRestApi):
     apispec_parameter_schemas = {
         "screenshot_query_schema": screenshot_query_schema,
         "get_delete_ids_schema": get_delete_ids_schema,
+        "get_export_ids_schema": get_export_ids_schema,
     }
     """ Add extra schemas to the OpenAPI components schema section """
     openapi_spec_methods = openapi_spec_methods_override
@@ -733,9 +734,7 @@ class ChartRestApi(BaseSupersetModelRestApi):
             content:
               application/json:
                 schema:
-                  type: array
-                  items:
-                    type: integer
+                  $ref: '#/components/schemas/get_export_ids_schema'
           responses:
             200:
               description: A zip file with chart(s), dataset(s) and database(s) as YAML

superset/databases/api.py

@@ -167,6 +167,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
 
     apispec_parameter_schemas = {
         "database_schemas_query_schema": database_schemas_query_schema,
+        "get_export_ids_schema": get_export_ids_schema,
     }
     openapi_spec_tag = "Database"
     openapi_spec_component_schemas = (
@@ -682,9 +683,7 @@ class DatabaseRestApi(BaseSupersetModelRestApi):
             content:
               application/json:
                 schema:
-                  type: array
-                  items:
-                    type: integer
+                  $ref: '#/components/schemas/get_export_ids_schema'
           responses:
             200:
               description: A zip file with database(s) and dataset(s) as YAML

superset/datasets/api.py

@@ -166,6 +166,9 @@ class DatasetRestApi(BaseSupersetModelRestApi):
     allowed_rel_fields = {"database", "owners"}
     allowed_distinct_fields = {"schema"}
 
+    apispec_parameter_schemas = {
+        "get_export_ids_schema": get_export_ids_schema,
+    }
     openapi_spec_component_schemas = (DatasetRelatedObjectsResponse,)
 
     @expose("/", methods=["POST"])
@@ -360,9 +363,7 @@ class DatasetRestApi(BaseSupersetModelRestApi):
             content:
               application/json:
                 schema:
-                  type: array
-                  items:
-                    type: integer
+                  $ref: '#/components/schemas/get_export_ids_schema'
           responses:
             200:
               description: Dataset export

superset/models/sql_lab.py

@@ -188,7 +188,6 @@ class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin, ImportMixin):
 
     export_parent = "database"
     export_fields = [
-        "db_id",
         "schema",
         "label",
         "description",

superset/queries/saved_queries/api.py

@@ -15,9 +15,12 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
+from datetime import datetime
+from io import BytesIO
 from typing import Any
+from zipfile import ZipFile
 
-from flask import g, Response
+from flask import g, Response, send_file
 from flask_appbuilder.api import expose, protect, rison, safe
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from flask_babel import ngettext
@@ -32,6 +35,7 @@ from superset.queries.saved_queries.commands.exceptions import (
     SavedQueryBulkDeleteFailedError,
     SavedQueryNotFoundError,
 )
+from superset.queries.saved_queries.commands.export import ExportSavedQueriesCommand
 from superset.queries.saved_queries.filters import (
     SavedQueryAllTextFilter,
     SavedQueryFavoriteFilter,
@@ -39,6 +43,7 @@ from superset.queries.saved_queries.filters import (
 )
 from superset.queries.saved_queries.schemas import (
     get_delete_ids_schema,
+    get_export_ids_schema,
     openapi_spec_methods_override,
 )
 from superset.views.base_api import BaseSupersetModelRestApi, statsd_metrics
@@ -50,6 +55,7 @@ class SavedQueryRestApi(BaseSupersetModelRestApi):
     datamodel = SQLAInterface(SavedQuery)
 
     include_route_methods = RouteMethod.REST_MODEL_VIEW_CRUD_SET | {
+        RouteMethod.EXPORT,
         RouteMethod.RELATED,
         RouteMethod.DISTINCT,
         "bulk_delete",  # not using RouteMethod since locally defined
@@ -114,6 +120,7 @@ class SavedQueryRestApi(BaseSupersetModelRestApi):
 
     apispec_parameter_schemas = {
         "get_delete_ids_schema": get_delete_ids_schema,
+        "get_export_ids_schema": get_export_ids_schema,
     }
     openapi_spec_tag = "Queries"
     openapi_spec_methods = openapi_spec_methods_override
@@ -183,3 +190,62 @@ class SavedQueryRestApi(BaseSupersetModelRestApi):
             return self.response_404()
         except SavedQueryBulkDeleteFailedError as ex:
             return self.response_422(message=str(ex))
+
+    @expose("/export/", methods=["GET"])
+    @protect()
+    @safe
+    @statsd_metrics
+    @rison(get_export_ids_schema)
+    def export(self, **kwargs: Any) -> Response:
+        """Export saved queries
+        ---
+        get:
+          description: >-
+            Exports multiple saved queries and downloads them as YAML files
+          parameters:
+          - in: query
+            name: q
+            content:
+              application/json:
+                schema:
+                  $ref: '#/components/schemas/get_export_ids_schema'
+          responses:
+            200:
+              description: A zip file with saved query(ies) and database(s) as YAML
+              content:
+                application/zip:
+                  schema:
+                    type: string
+                    format: binary
+            400:
+              $ref: '#/components/responses/400'
+            401:
+              $ref: '#/components/responses/401'
+            404:
+              $ref: '#/components/responses/404'
+            500:
+              $ref: '#/components/responses/500'
+        """
+        requested_ids = kwargs["rison"]
+        timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
+        root = f"saved_query_export_{timestamp}"
+        filename = f"{root}.zip"
+
+        buf = BytesIO()
+        with ZipFile(buf, "w") as bundle:
+            try:
+                for file_name, file_content in ExportSavedQueriesCommand(
+                    requested_ids
+                ).run():
+                    with bundle.open(f"{root}/{file_name}", "w") as fp:
+                        fp.write(file_content.encode())
+            except SavedQueryNotFoundError:
+                return self.response_404()
+        buf.seek(0)
+
+        return send_file(
+            buf,
+            mimetype="application/zip",
+            as_attachment=True,
+            attachment_filename=filename,
+        )

superset/queries/saved_queries/commands/export.py

@@ -0,0 +1,92 @@
+# 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.
+# isort:skip_file
+
+import json
+import logging
+from typing import Iterator, List, Tuple
+
+import yaml
+
+from superset.commands.base import BaseCommand
+from superset.queries.saved_queries.commands.exceptions import SavedQueryNotFoundError
+from superset.queries.saved_queries.dao import SavedQueryDAO
+from superset.utils.dict_import_export import IMPORT_EXPORT_VERSION, sanitize
+from superset.models.sql_lab import SavedQuery
+
+logger = logging.getLogger(__name__)
+
+
+class ExportSavedQueriesCommand(BaseCommand):
+    def __init__(self, query_ids: List[int]):
+        self.query_ids = query_ids
+
+        # this will be set when calling validate()
+        self._models: List[SavedQuery] = []
+
+    @staticmethod
+    def export_saved_query(query: SavedQuery) -> Iterator[Tuple[str, str]]:
+        # build filename based on database, optional schema, and label
+        database_slug = sanitize(query.database.database_name)
+        schema_slug = sanitize(query.schema)
+        query_slug = sanitize(query.label) or str(query.uuid)
+        file_name = f"queries/{database_slug}/{schema_slug}/{query_slug}.yaml"
+
+        payload = query.export_to_dict(
+            recursive=False,
+            include_parent_ref=False,
+            include_defaults=True,
+            export_uuids=True,
+        )
+        payload["version"] = IMPORT_EXPORT_VERSION
+        payload["database_uuid"] = str(query.database.uuid)
+
+        file_content = yaml.safe_dump(payload, sort_keys=False)
+        yield file_name, file_content
+
+        # include database as well
+        file_name = f"databases/{database_slug}.yaml"
+
+        payload = query.database.export_to_dict(
+            recursive=False,
+            include_parent_ref=False,
+            include_defaults=True,
+            export_uuids=True,
+        )
+        # TODO (betodealmeida): move this logic to export_to_dict once this
+        # becomes the default export endpoint
+        if "extra" in payload:
+            try:
+                payload["extra"] = json.loads(payload["extra"])
+            except json.decoder.JSONDecodeError:
+                logger.info("Unable to decode `extra` field: %s", payload["extra"])
+
+        payload["version"] = IMPORT_EXPORT_VERSION
+
+        file_content = yaml.safe_dump(payload, sort_keys=False)
+        yield file_name, file_content
+
+    def run(self) -> Iterator[Tuple[str, str]]:
+        self.validate()
+
+        for query in self._models:
+            yield from self.export_saved_query(query)
+
+    def validate(self) -> None:
+        self._models = SavedQueryDAO.find_by_ids(self.query_ids)
+        if len(self._models) != len(self.query_ids):
+            raise SavedQueryNotFoundError()

superset/queries/saved_queries/schemas.py

@@ -31,3 +31,4 @@ openapi_spec_methods_override = {
 }
 
 get_delete_ids_schema = {"type": "array", "items": {"type": "integer"}}
+get_export_ids_schema = {"type": "array", "items": {"type": "integer"}}

tests/queries/saved_queries/api_tests.py

@@ -17,7 +17,9 @@
 # isort:skip_file
 """Unit tests for Superset"""
 import json
+from io import BytesIO
 from typing import Optional
+from zipfile import is_zipfile
 
 import pytest
 import prison
@@ -680,3 +682,50 @@ class TestSavedQueryApi(SupersetTestCase):
         uri = f"api/v1/saved_query/?q={prison.dumps(saved_query_ids)}"
         rv = self.delete_assert_metric(uri, "bulk_delete")
         assert rv.status_code == 404
+
+    @pytest.mark.usefixtures("create_saved_queries")
+    def test_export(self):
+        """
+        Saved Query API: Test export
+        """
+        admin = self.get_user("admin")
+        sample_query = (
+            db.session.query(SavedQuery).filter(SavedQuery.created_by == admin).first()
+        )
+
+        self.login(username="admin")
+        argument = [sample_query.id]
+        uri = f"api/v1/saved_query/export/?q={prison.dumps(argument)}"
+        rv = self.client.get(uri)
+        assert rv.status_code == 200
+        buf = BytesIO(rv.data)
+        assert is_zipfile(buf)
+
+    @pytest.mark.usefixtures("create_saved_queries")
+    def test_export_not_found(self):
+        """
+        Saved Query API: Test export
+        """
+        max_id = db.session.query(func.max(SavedQuery.id)).scalar()
+
+        self.login(username="admin")
+        argument = [max_id + 1, max_id + 2]
+        uri = f"api/v1/saved_query/export/?q={prison.dumps(argument)}"
+        rv = self.client.get(uri)
+        assert rv.status_code == 404
+
+    @pytest.mark.usefixtures("create_saved_queries")
+    def test_export_not_allowed(self):
+        """
+        Saved Query API: Test export
+        """
+        admin = self.get_user("admin")
+        sample_query = (
+            db.session.query(SavedQuery).filter(SavedQuery.created_by == admin).first()
+        )
+
+        self.login(username="gamma")
+        argument = [sample_query.id]
+        uri = f"api/v1/saved_query/export/?q={prison.dumps(argument)}"
+        rv = self.client.get(uri)
+        assert rv.status_code == 404

tests/queries/saved_queries/commands_tests.py

@@ -0,0 +1,109 @@
+# 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.
+
+from unittest.mock import patch
+
+import yaml
+
+from superset import db, security_manager
+from superset.models.sql_lab import SavedQuery
+from superset.queries.saved_queries.commands.exceptions import SavedQueryNotFoundError
+from superset.queries.saved_queries.commands.export import ExportSavedQueriesCommand
+from superset.utils.core import get_example_database
+from tests.base_tests import SupersetTestCase
+
+
+class TestExportSavedQueriesCommand(SupersetTestCase):
+    def setUp(self):
+        self.example_database = get_example_database()
+        self.example_query = SavedQuery(
+            database=self.example_database,
+            created_by=self.get_user("admin"),
+            sql="SELECT 42",
+            label="The answer",
+            schema="schema1",
+            description="Answer to the Ultimate Question of Life, the Universe, and Everything",
+        )
+        db.session.add(self.example_query)
+        db.session.commit()
+
+    def tearDown(self):
+        db.session.delete(self.example_query)
+        db.session.commit()
+
+    @patch("superset.queries.saved_queries.filters.g")
+    def test_export_query_command(self, mock_g):
+        mock_g.user = security_manager.find_user("admin")
+
+        command = ExportSavedQueriesCommand(query_ids=[self.example_query.id])
+        contents = dict(command.run())
+
+        expected = [
+            "queries/examples/schema1/the_answer.yaml",
+            "databases/examples.yaml",
+        ]
+        assert expected == list(contents.keys())
+
+        metadata = yaml.safe_load(contents["queries/examples/schema1/the_answer.yaml"])
+        assert metadata == {
+            "schema": "schema1",
+            "label": "The answer",
+            "description": "Answer to the Ultimate Question of Life, the Universe, and Everything",
+            "sql": "SELECT 42",
+            "uuid": str(self.example_query.uuid),
+            "version": "1.0.0",
+            "database_uuid": str(self.example_database.uuid),
+        }
+
+    @patch("superset.queries.saved_queries.filters.g")
+    def test_export_query_command_no_access(self, mock_g):
+        """Test that users can't export datasets they don't have access to"""
+        mock_g.user = security_manager.find_user("gamma")
+
+        command = ExportSavedQueriesCommand(query_ids=[self.example_query.id])
+        contents = command.run()
+        with self.assertRaises(SavedQueryNotFoundError):
+            next(contents)
+
+    @patch("superset.queries.saved_queries.filters.g")
+    def test_export_query_command_invalid_dataset(self, mock_g):
+        """Test that an error is raised when exporting an invalid dataset"""
+        mock_g.user = security_manager.find_user("admin")
+
+        command = ExportSavedQueriesCommand(query_ids=[-1])
+        contents = command.run()
+        with self.assertRaises(SavedQueryNotFoundError):
+            next(contents)
+
+    @patch("superset.queries.saved_queries.filters.g")
+    def test_export_query_command_key_order(self, mock_g):
+        """Test that they keys in the YAML have the same order as export_fields"""
+        mock_g.user = security_manager.find_user("admin")
+
+        command = ExportSavedQueriesCommand(query_ids=[self.example_query.id])
+        contents = dict(command.run())
+
+        metadata = yaml.safe_load(contents["queries/examples/schema1/the_answer.yaml"])
+        assert list(metadata.keys()) == [
+            "schema",
+            "label",
+            "description",
+            "sql",
+            "uuid",
+            "version",
+            "database_uuid",
+        ]