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

* feat: commands for importing databases and datasets

* Refactor code

superset/cli.py

@@ -304,6 +304,9 @@ def import_datasources(path: str, sync: str, recursive: bool) -> None:
     from superset.datasets.commands.importers.v0 import ImportDatasetsCommand
 
     sync_array = sync.split(",")
+    sync_columns = "columns" in sync_array
+    sync_metrics = "metrics" in sync_array
+
     path_object = Path(path)
     files: List[Path] = []
     if path_object.is_file():
@@ -316,7 +319,7 @@ def import_datasources(path: str, sync: str, recursive: bool) -> None:
         files.extend(path_object.rglob("*.yml"))
     contents = {path.name: open(path).read() for path in files}
     try:
-        ImportDatasetsCommand(contents, sync_array).run()
+        ImportDatasetsCommand(contents, sync_columns, sync_metrics).run()
     except Exception:  # pylint: disable=broad-except
         logger.exception("Error when importing dataset")
 

superset/commands/importers/exceptions.py

@@ -0,0 +1,23 @@
+# 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 superset.commands.exceptions import CommandException
+
+
+class IncorrectVersionError(CommandException):
+    status = 422
+    message = "Import has incorrect version"

superset/commands/importers/v1/__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/commands/importers/v1/utils.py

@@ -0,0 +1,67 @@
+# 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 logging
+from typing import Any, Dict
+
+import yaml
+from marshmallow import fields, Schema, validate
+from marshmallow.exceptions import ValidationError
+
+from superset.commands.importers.exceptions import IncorrectVersionError
+
+METADATA_FILE_NAME = "metadata.yaml"
+IMPORT_VERSION = "1.0.0"
+
+logger = logging.getLogger(__name__)
+
+
+class MetadataSchema(Schema):
+    version = fields.String(required=True, validate=validate.Equal(IMPORT_VERSION))
+    type = fields.String(required=True)
+    timestamp = fields.DateTime()
+
+
+def load_yaml(file_name: str, content: str) -> Dict[str, Any]:
+    """Try to load a YAML file"""
+    try:
+        return yaml.safe_load(content)
+    except yaml.parser.ParserError:
+        logger.exception("Invalid YAML in %s", METADATA_FILE_NAME)
+        raise ValidationError({file_name: "Not a valid YAML file"})
+
+
+def load_metadata(contents: Dict[str, str]) -> Dict[str, str]:
+    """Apply validation and load a metadata file"""
+    if METADATA_FILE_NAME not in contents:
+        # if the contents ahve no METADATA_FILE_NAME this is probably
+        # a original export without versioning that should not be
+        # handled by this command
+        raise IncorrectVersionError(f"Missing {METADATA_FILE_NAME}")
+
+    metadata = load_yaml(METADATA_FILE_NAME, contents[METADATA_FILE_NAME])
+    try:
+        MetadataSchema().load(metadata)
+    except ValidationError as exc:
+        # if the version doesn't match raise an exception so that the
+        # dispatcher can try a different command version
+        if "version" in exc.messages:
+            raise IncorrectVersionError(exc.messages["version"][0])
+
+        # otherwise we raise the validation error
+        exc.messages = {METADATA_FILE_NAME: exc.messages}
+        raise exc
+
+    return metadata

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)

superset/datasets/commands/importers/v0.py

@@ -282,9 +282,19 @@ class ImportDatasetsCommand(BaseCommand):
     in Superset.
     """
 
-    def __init__(self, contents: Dict[str, str], sync: Optional[List[str]] = None):
+    def __init__(
+        self,
+        contents: Dict[str, str],
+        sync_columns: bool = False,
+        sync_metrics: bool = False,
+    ):
         self.contents = contents
-        self.sync = sync
+
+        self.sync = []
+        if sync_columns:
+            self.sync.append("columns")
+        if sync_metrics:
+            self.sync.append("metrics")
 
     def run(self) -> None:
         self.validate()