feat: new import commands for dataset and databases (#11670)

* feat: commands for importing databases and datasets

* Refactor code

superset/databases/commands/importers/__init__.py

@@ -0,0 +1,16 @@
+# 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.

superset/databases/commands/importers/v1/__init__.py

@@ -0,0 +1,116 @@
+# 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 typing import Any, Dict, List
+
+from marshmallow import Schema, validate
+from marshmallow.exceptions import ValidationError
+from sqlalchemy.orm import Session
+
+from superset import db
+from superset.commands.base import BaseCommand
+from superset.commands.exceptions import CommandInvalidError
+from superset.commands.importers.v1.utils import (
+    load_metadata,
+    load_yaml,
+    METADATA_FILE_NAME,
+)
+from superset.databases.commands.importers.v1.utils import import_database
+from superset.databases.schemas import ImportV1DatabaseSchema
+from superset.datasets.commands.importers.v1.utils import import_dataset
+from superset.datasets.schemas import ImportV1DatasetSchema
+from superset.models.core import Database
+
+schemas: Dict[str, Schema] = {
+    "databases/": ImportV1DatabaseSchema(),
+    "datasets/": ImportV1DatasetSchema(),
+}
+
+
+class ImportDatabasesCommand(BaseCommand):
+
+    """Import databases"""
+
+    # pylint: disable=unused-argument
+    def __init__(self, contents: Dict[str, str], *args: Any, **kwargs: Any):
+        self.contents = contents
+        self._configs: Dict[str, Any] = {}
+
+    def _import_bundle(self, session: Session) -> None:
+        # first import databases
+        database_ids: Dict[str, int] = {}
+        for file_name, config in self._configs.items():
+            if file_name.startswith("databases/"):
+                database = import_database(session, config, overwrite=True)
+                database_ids[str(database.uuid)] = database.id
+
+        # import related datasets
+        for file_name, config in self._configs.items():
+            if (
+                file_name.startswith("datasets/")
+                and config["database_uuid"] in database_ids
+            ):
+                config["database_id"] = database_ids[config["database_uuid"]]
+                # overwrite=False prevents deleting any non-imported columns/metrics
+                import_dataset(session, config, overwrite=False)
+
+    def run(self) -> None:
+        self.validate()
+
+        # rollback to prevent partial imports
+        try:
+            self._import_bundle(db.session)
+            db.session.commit()
+        except Exception as exc:
+            db.session.rollback()
+            raise exc
+
+    def validate(self) -> None:
+        exceptions: List[ValidationError] = []
+
+        # verify that the metadata file is present and valid
+        try:
+            metadata = load_metadata(self.contents)
+        except ValidationError as exc:
+            exceptions.append(exc)
+            metadata = None
+
+        for file_name, content in self.contents.items():
+            prefix = file_name.split("/")[0]
+            schema = schemas.get(f"{prefix}/")
+            if schema:
+                try:
+                    config = load_yaml(file_name, content)
+                    schema.load(config)
+                    self._configs[file_name] = config
+                except ValidationError as exc:
+                    exc.messages = {file_name: exc.messages}
+                    exceptions.append(exc)
+
+        # validate that the type declared in METADATA_FILE_NAME is correct
+        if metadata:
+            type_validator = validate.Equal(Database.__name__)
+            try:
+                type_validator(metadata["type"])
+            except ValidationError as exc:
+                exc.messages = {METADATA_FILE_NAME: {"type": exc.messages}}
+                exceptions.append(exc)
+
+        if exceptions:
+            exception = CommandInvalidError("Error importing database")
+            exception.add_list(exceptions)
+            raise exception

superset/databases/commands/importers/v1/utils.py

@@ -0,0 +1,42 @@
+# 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.
+
+import json
+from typing import Any, Dict
+
+from sqlalchemy.orm import Session
+
+from superset.models.core import Database
+
+
+def import_database(
+    session: Session, config: Dict[str, Any], overwrite: bool = False
+) -> Database:
+    existing = session.query(Database).filter_by(uuid=config["uuid"]).first()
+    if existing:
+        if not overwrite:
+            return existing
+        config["id"] = existing.id
+
+    # TODO (betodealmeida): move this logic to import_from_dict
+    config["extra"] = json.dumps(config["extra"])
+
+    database = Database.import_from_dict(session, config, recursive=False)
+    if database.id is None:
+        session.flush()
+
+    return database

superset/databases/schemas.py

@@ -408,3 +408,24 @@ class DatabaseRelatedDashboards(Schema):
 class DatabaseRelatedObjectsResponse(Schema):
     charts = fields.Nested(DatabaseRelatedCharts)
     dashboards = fields.Nested(DatabaseRelatedDashboards)
+
+
+class ImportV1DatabaseExtraSchema(Schema):
+    metadata_params = fields.Dict(keys=fields.Str(), values=fields.Raw())
+    engine_params = fields.Dict(keys=fields.Str(), values=fields.Raw())
+    metadata_cache_timeout = fields.Dict(keys=fields.Str(), values=fields.Integer())
+    schemas_allowed_for_csv_upload = fields.List(fields.String)
+
+
+class ImportV1DatabaseSchema(Schema):
+    database_name = fields.String(required=True)
+    sqlalchemy_uri = fields.String(required=True)
+    cache_timeout = fields.Integer(allow_none=True)
+    expose_in_sqllab = fields.Boolean()
+    allow_run_async = fields.Boolean()
+    allow_ctas = fields.Boolean()
+    allow_cvas = fields.Boolean()
+    allow_csv_upload = fields.Boolean()
+    extra = fields.Nested(ImportV1DatabaseExtraSchema)
+    uuid = fields.UUID(required=True)
+    version = fields.String(required=True)