From 5e95d4638cf7e9f271e1e6f02b1bcc2959ed6d25 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 18 May 2021 20:51:33 -0700
Subject: [PATCH] fix: import dataset with extra; Vertica URI (#14698)

* fix: import dataset with extra; Vertica URI

* Fix lint
---
 superset/databases/schemas.py                    | 3 +--
 superset/datasets/commands/export.py             | 2 +-
 superset/datasets/commands/importers/v1/utils.py | 2 +-
 superset/datasets/schemas.py                     | 2 +-
 tests/databases/api_tests.py                     | 5 +++--
 tests/datasets/commands_tests.py                 | 2 +-
 tests/fixtures/importexport.py                   | 2 +-
 7 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/superset/databases/schemas.py b/superset/databases/schemas.py
index fd42c6343..a23659fc0 100644
--- a/superset/databases/schemas.py
+++ b/superset/databases/schemas.py
@@ -16,7 +16,6 @@
 # under the License.
 import inspect
 import json
-import urllib.parse
 from typing import Any, Dict
 
 from flask import current_app
@@ -550,6 +549,6 @@ class ImportV1DatabaseSchema(Schema):
     def validate_password(self, data: Dict[str, Any], **kwargs: Any) -> None:
         """If sqlalchemy_uri has a masked password, password is required"""
         uri = data["sqlalchemy_uri"]
-        password = urllib.parse.urlparse(uri).password
+        password = make_url(uri).password
         if password == PASSWORD_MASK and data.get("password") is None:
             raise ValidationError("Must provide a password for the database")
diff --git a/superset/datasets/commands/export.py b/superset/datasets/commands/export.py
index 64946f48f..84ae6c486 100644
--- a/superset/datasets/commands/export.py
+++ b/superset/datasets/commands/export.py
@@ -31,7 +31,7 @@ from superset.utils.dict_import_export import EXPORT_VERSION
 
 logger = logging.getLogger(__name__)
 
-JSON_KEYS = {"params", "template_params", "extra"}
+JSON_KEYS = {"params", "template_params"}
 
 
 class ExportDatasetsCommand(ExportModelsCommand):
diff --git a/superset/datasets/commands/importers/v1/utils.py b/superset/datasets/commands/importers/v1/utils.py
index 48e1f797f..c21c66ff1 100644
--- a/superset/datasets/commands/importers/v1/utils.py
+++ b/superset/datasets/commands/importers/v1/utils.py
@@ -38,7 +38,7 @@ logger = logging.getLogger(__name__)
 CHUNKSIZE = 512
 VARCHAR = re.compile(r"VARCHAR\((\d+)\)", re.IGNORECASE)
 
-JSON_KEYS = {"params", "template_params", "extra"}
+JSON_KEYS = {"params", "template_params"}
 
 
 type_map = {
diff --git a/superset/datasets/schemas.py b/superset/datasets/schemas.py
index 6f6cc8ada..a857efec8 100644
--- a/superset/datasets/schemas.py
+++ b/superset/datasets/schemas.py
@@ -162,7 +162,7 @@ class ImportV1DatasetSchema(Schema):
     template_params = fields.Dict(allow_none=True)
     filter_select_enabled = fields.Boolean()
     fetch_values_predicate = fields.String(allow_none=True)
-    extra = fields.Dict(allow_none=True)
+    extra = fields.String(allow_none=True)
     uuid = fields.UUID(required=True)
     columns = fields.List(fields.Nested(ImportV1ColumnSchema))
     metrics = fields.List(fields.Nested(ImportV1MetricSchema))
diff --git a/tests/databases/api_tests.py b/tests/databases/api_tests.py
index d7404f5f4..4d550095f 100644
--- a/tests/databases/api_tests.py
+++ b/tests/databases/api_tests.py
@@ -1314,7 +1314,7 @@ class TestDatabaseApi(SupersetTestCase):
         masked_database_config = database_config.copy()
         masked_database_config[
             "sqlalchemy_uri"
-        ] = "postgresql://username:XXXXXXXXXX@host:12345/db"
+        ] = "vertica+vertica_python://hackathon:XXXXXXXXXX@host:5433/dbname?ssl=1"
 
         buf = BytesIO()
         with ZipFile(buf, "w") as bundle:
@@ -1341,7 +1341,8 @@ class TestDatabaseApi(SupersetTestCase):
         )
         assert database.database_name == "imported_database"
         assert (
-            database.sqlalchemy_uri == "postgresql://username:XXXXXXXXXX@host:12345/db"
+            database.sqlalchemy_uri
+            == "vertica+vertica_python://hackathon:XXXXXXXXXX@host:5433/dbname?ssl=1"
         )
         assert database.password == "SECRET"
 
diff --git a/tests/datasets/commands_tests.py b/tests/datasets/commands_tests.py
index 84c108909..4df7365a5 100644
--- a/tests/datasets/commands_tests.py
+++ b/tests/datasets/commands_tests.py
@@ -315,7 +315,7 @@ class TestImportDatasetsCommand(SupersetTestCase):
         assert dataset.template_params == "{}"
         assert dataset.filter_select_enabled
         assert dataset.fetch_values_predicate is None
-        assert dataset.extra is None
+        assert dataset.extra == "dttm > sysdate() -10 "
 
         # database is also imported
         assert str(dataset.database.uuid) == "b8a1ccd3-779d-4ab7-8ad8-9ab119d7fe89"
diff --git a/tests/fixtures/importexport.py b/tests/fixtures/importexport.py
index 66eb3f603..04aaaa394 100644
--- a/tests/fixtures/importexport.py
+++ b/tests/fixtures/importexport.py
@@ -375,7 +375,7 @@ dataset_config: Dict[str, Any] = {
     "template_params": {},
     "filter_select_enabled": True,
     "fetch_values_predicate": None,
-    "extra": None,
+    "extra": "dttm > sysdate() -10 ",
     "metrics": [
         {
             "metric_name": "count",