fix: Support datetime_format during import (#36679)

superset/datasets/schemas.py

@@ -257,6 +257,7 @@ class ImportV1ColumnSchema(Schema):
     expression = fields.String(allow_none=True)
     description = fields.String(allow_none=True)
     python_date_format = fields.String(allow_none=True)
+    datetime_format = fields.String(allow_none=True)
 
 
 class ImportMetricCurrencySchema(Schema):

tests/unit_tests/datasets/commands/importers/v1/import_test.py

@@ -17,6 +17,7 @@
 # pylint: disable=import-outside-toplevel, unused-argument, unused-import, invalid-name
 
 import copy
+import io
 import re
 import uuid
 from typing import Any
@@ -28,25 +29,29 @@ from flask_appbuilder.security.sqla.models import Role, User
 from pytest_mock import MockerFixture
 from sqlalchemy.orm.session import Session
 
-from superset import db
+from superset import db, security_manager
 from superset.commands.dataset.exceptions import (
     DatasetForbiddenDataURI,
 )
-from superset.commands.dataset.importers.v1.utils import validate_data_uri
+from superset.commands.dataset.importers.v1.utils import (
+    import_dataset,
+    validate_data_uri,
+)
 from superset.commands.exceptions import ImportFailedError
+from superset.connectors.sqla.models import SqlaTable, TableColumn
+from superset.datasets.schemas import ImportV1DatasetSchema
+from superset.models.core import Database
 from superset.utils import json
 from superset.utils.core import override_user
+from tests.integration_tests.fixtures.importexport import (
+    dataset_config as dataset_fixture,
+)
 
 
 def test_import_dataset(mocker: MockerFixture, session: Session) -> None:
     """
     Test importing a dataset.
     """
-    from superset import security_manager
-    from superset.commands.dataset.importers.v1.utils import import_dataset
-    from superset.connectors.sqla.models import SqlaTable
-    from superset.models.core import Database
-
     mocker.patch.object(security_manager, "can_access", return_value=True)
 
     engine = db.session.get_bind()
@@ -247,11 +252,6 @@ def test_import_dataset_no_folder(mocker: MockerFixture, session: Session) -> No
     """
     Test importing a dataset that was exported without folders.
     """
-    from superset import security_manager
-    from superset.commands.dataset.importers.v1.utils import import_dataset
-    from superset.connectors.sqla.models import SqlaTable
-    from superset.models.core import Database
-
     mocker.patch.object(security_manager, "can_access", return_value=True)
 
     engine = db.session.get_bind()
@@ -327,11 +327,6 @@ def test_import_dataset_duplicate_column(
     """
     Test importing a dataset with a column that already exists.
     """
-    from superset import security_manager
-    from superset.commands.dataset.importers.v1.utils import import_dataset
-    from superset.connectors.sqla.models import SqlaTable, TableColumn
-    from superset.models.core import Database
-
     mocker.patch.object(security_manager, "can_access", return_value=True)
 
     engine = db.session.get_bind()
@@ -452,12 +447,6 @@ def test_import_column_extra_is_string(mocker: MockerFixture, session: Session)
     """
     Test importing a dataset when the column extra is a string.
     """
-    from superset import security_manager
-    from superset.commands.dataset.importers.v1.utils import import_dataset
-    from superset.connectors.sqla.models import SqlaTable
-    from superset.datasets.schemas import ImportV1DatasetSchema
-    from superset.models.core import Database
-
     mocker.patch.object(security_manager, "can_access", return_value=True)
 
     engine = db.session.get_bind()
@@ -537,12 +526,6 @@ def test_import_dataset_extra_empty_string(
     """
     Test importing a dataset when the extra field is an empty string.
     """
-    from superset import security_manager
-    from superset.commands.dataset.importers.v1.utils import import_dataset
-    from superset.connectors.sqla.models import SqlaTable
-    from superset.datasets.schemas import ImportV1DatasetSchema
-    from superset.models.core import Database
-
     mocker.patch.object(security_manager, "can_access", return_value=True)
 
     engine = db.session.get_bind()
@@ -603,14 +586,6 @@ def test_import_column_allowed_data_url(
     """
     Test importing a dataset when using data key to fetch data from a URL.
     """
-    import io
-
-    from superset import security_manager
-    from superset.commands.dataset.importers.v1.utils import import_dataset
-    from superset.connectors.sqla.models import SqlaTable
-    from superset.datasets.schemas import ImportV1DatasetSchema
-    from superset.models.core import Database
-
     mock_urlopen.return_value = io.StringIO("col1\nvalue1\nvalue2\n")
 
     mocker.patch.object(security_manager, "can_access", return_value=True)
@@ -677,12 +652,6 @@ def test_import_dataset_managed_externally(
     """
     Test importing a dataset that is managed externally.
     """
-    from superset import security_manager
-    from superset.commands.dataset.importers.v1.utils import import_dataset
-    from superset.connectors.sqla.models import SqlaTable
-    from superset.models.core import Database
-    from tests.integration_tests.fixtures.importexport import dataset_config
-
     mocker.patch.object(security_manager, "can_access", return_value=True)
 
     engine = db.session.get_bind()
@@ -692,7 +661,7 @@ def test_import_dataset_managed_externally(
     db.session.add(database)
     db.session.flush()
 
-    config = copy.deepcopy(dataset_config)
+    config = copy.deepcopy(dataset_fixture)
     config["is_managed_externally"] = True
     config["external_url"] = "https://example.org/my_table"
     config["database_id"] = database.id
@@ -702,6 +671,36 @@ def test_import_dataset_managed_externally(
     assert sqla_table.external_url == "https://example.org/my_table"
 
 
+def test_import_dataset_column_datetime_format(
+    mocker: MockerFixture,
+    session: Session,
+) -> None:
+    """
+    Test importing a dataset with a column including a datetime format.
+    """
+    mocker.patch.object(security_manager, "can_access", return_value=True)
+
+    engine = db.session.get_bind()
+    SqlaTable.metadata.create_all(engine)  # pylint: disable=no-member
+
+    database = Database(database_name="my_database", sqlalchemy_uri="sqlite://")
+    db.session.add(database)
+    db.session.flush()
+
+    config = copy.deepcopy(dataset_fixture)
+    for column in config["columns"]:
+        column["datetime_format"] = "%Y-%m-%d"
+
+    schema = ImportV1DatasetSchema()
+    dataset_config = schema.load(config)
+
+    dataset_config["database_id"] = database.id
+
+    sqla_table = import_dataset(dataset_config)
+    for column in sqla_table.columns:
+        assert column.datetime_format == "%Y-%m-%d"
+
+
 def test_import_dataset_without_owner_permission(
     mocker: MockerFixture,
     session: Session,
@@ -709,12 +708,6 @@ def test_import_dataset_without_owner_permission(
     """
     Test importing a dataset that is managed externally.
     """
-    from superset import security_manager
-    from superset.commands.dataset.importers.v1.utils import import_dataset
-    from superset.connectors.sqla.models import SqlaTable
-    from superset.models.core import Database
-    from tests.integration_tests.fixtures.importexport import dataset_config
-
     mock_can_access = mocker.patch.object(
         security_manager, "can_access", return_value=True
     )
@@ -726,7 +719,7 @@ def test_import_dataset_without_owner_permission(
     db.session.add(database)
     db.session.flush()
 
-    config = copy.deepcopy(dataset_config)
+    config = copy.deepcopy(dataset_fixture)
     config["database_id"] = database.id
 
     import_dataset(config)