refactor(example_data): replace the way the birth_names data is loaded to DB (#18060)

* refactor: replace the way the birth_names data is loaded to DB

* fix failed unit test

* fix failed unit test

* fix failed tests

* fix pass wrong flag of support datetime type

* remove unused fixture

superset/utils/database.py

@@ -39,7 +39,10 @@ def get_or_create_db(
     from superset.models import core as models
 
     database = (
-        db.session.query(models.Database).filter_by(database_name=database_name).first()
+        db.session.query(models.Database)
+        .filter_by(database_name=database_name)
+        .autoflush(False)
+        .first()
     )
 
     # databases with a fixed UUID
@@ -72,3 +75,14 @@ def get_example_database() -> Database:
 def get_main_database() -> Database:
     db_uri = current_app.config["SQLALCHEMY_DATABASE_URI"]
     return get_or_create_db("main", db_uri)
+
+
+# TODO - the below method used by tests so should move there but should move together
+# with above function... think of how to refactor it
+def remove_database(database: Database) -> None:
+    # pylint: disable=import-outside-toplevel
+    from superset import db
+
+    session = db.session
+    session.delete(database)
+    session.commit()

tests/conftest.py

@@ -0,0 +1,103 @@
+#  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.
+#
+#  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 __future__ import annotations
+
+from typing import Callable, TYPE_CHECKING
+from unittest.mock import MagicMock, Mock, PropertyMock
+
+from pytest import fixture
+
+from tests.example_data.data_loading.pandas.pandas_data_loader import PandasDataLoader
+from tests.example_data.data_loading.pandas.pands_data_loading_conf import (
+    PandasLoaderConfigurations,
+)
+from tests.example_data.data_loading.pandas.table_df_convertor import (
+    TableToDfConvertorImpl,
+)
+
+SUPPORT_DATETIME_TYPE = "support_datetime_type"
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Engine
+
+    from superset.connectors.sqla.models import Database
+    from tests.example_data.data_loading.base_data_loader import DataLoader
+    from tests.example_data.data_loading.pandas.pandas_data_loader import (
+        TableToDfConvertor,
+    )
+
+pytest_plugins = "tests.fixtures"
+
+PRESTO = "presto"
+BACKEND_PROPERTY_VALUE = "sqlite"
+
+
+@fixture(scope="session")
+def example_db_provider() -> Callable[[], Database]:
+    def mock_provider() -> Mock:
+        mock = MagicMock()
+        type(mock).backend = PropertyMock(return_value=BACKEND_PROPERTY_VALUE)
+        return mock
+
+    return mock_provider
+
+
+@fixture(scope="session")
+def example_db_engine(example_db_provider: Callable[[], Database]) -> Engine:
+    return example_db_provider().get_sqla_engine()
+
+
+@fixture(scope="session")
+def pandas_loader_configuration(support_datetime_type,) -> PandasLoaderConfigurations:
+    return PandasLoaderConfigurations.make_from_dict(
+        {SUPPORT_DATETIME_TYPE: support_datetime_type}
+    )
+
+
+@fixture(scope="session")
+def support_datetime_type(example_db_provider: Callable[[], Database]) -> bool:
+    return example_db_provider().backend != PRESTO
+
+
+@fixture(scope="session")
+def table_to_df_convertor(
+    pandas_loader_configuration: PandasLoaderConfigurations,
+) -> TableToDfConvertor:
+    return TableToDfConvertorImpl(
+        not pandas_loader_configuration.support_datetime_type,
+        pandas_loader_configuration.strftime,
+    )
+
+
+@fixture(scope="session")
+def data_loader(
+    example_db_engine: Engine,
+    pandas_loader_configuration: PandasLoaderConfigurations,
+    table_to_df_convertor: TableToDfConvertor,
+) -> DataLoader:
+    return PandasDataLoader(
+        example_db_engine, pandas_loader_configuration, table_to_df_convertor
+    )

tests/consts/__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.

tests/consts/birth_names.py

@@ -0,0 +1,26 @@
+#  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.
+TABLE_NAME = "birth_names"
+NUM_GIRLS = "num_girls"
+NUM_BOYS = "num_boys"
+STATE = "state"
+NUM = "num"
+NAME = "name"
+GENDER = "gender"
+DS = "ds"
+GIRL = "girl"
+BOY = "boy"

tests/consts/us_states.py

@@ -31,3 +31,56 @@
 #  KIND, either express or implied.  See the License for the
 #  specific language governing permissions and limitations
 #  under the License.
+US_STATES = [
+    "AL",
+    "AK",
+    "AZ",
+    "AR",
+    "CA",
+    "CO",
+    "CT",
+    "DE",
+    "FL",
+    "GA",
+    "HI",
+    "ID",
+    "IL",
+    "IN",
+    "IA",
+    "KS",
+    "KY",
+    "LA",
+    "ME",
+    "MD",
+    "MA",
+    "MI",
+    "MN",
+    "MS",
+    "MO",
+    "MT",
+    "NE",
+    "NV",
+    "NH",
+    "NJ",
+    "NM",
+    "NY",
+    "NC",
+    "ND",
+    "OH",
+    "OK",
+    "OR",
+    "PA",
+    "RI",
+    "SC",
+    "SD",
+    "TN",
+    "TX",
+    "UT",
+    "VT",
+    "VA",
+    "WA",
+    "WV",
+    "WI",
+    "WY",
+    "other",
+]

tests/example_data/__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.

tests/example_data/data_generator/base_generator.py

@@ -14,16 +14,6 @@
 #  KIND, either express or implied.  See the License for the
 #  specific language governing permissions and limitations
 #  under the License.
-#
-#  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 abc import ABC, abstractmethod
 from typing import Any, Dict, Iterable
 

tests/example_data/data_generator/birth_names/__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.

tests/example_data/data_generator/birth_names/birth_names_generator.py

@@ -14,66 +14,28 @@
 #  KIND, either express or implied.  See the License for the
 #  specific language governing permissions and limitations
 #  under the License.
-#
-#  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.
-#
-#  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.
-#
-#  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 __future__ import annotations
 
 from datetime import datetime
 from random import choice, randint
-from typing import Any, Dict, Iterable
+from typing import Any, Dict, Iterable, TYPE_CHECKING
 
-from tests.common.example_data_generator.base_generator import ExampleDataGenerator
+from tests.consts.birth_names import (
-from tests.common.example_data_generator.consts import US_STATES
+    BOY,
-from tests.common.example_data_generator.string_generator import StringGenerator
+    DS,
-
+    GENDER,
-NUM_GIRLS = "num_girls"
+    GIRL,
-NUM_BOYS = "num_boys"
+    NAME,
-STATE = "state"
+    NUM,
-NUM = "num"
+    NUM_BOYS,
-NAME = "name"
+    NUM_GIRLS,
-GENDER = "gender"
+    STATE,
-DS = "ds"
-GIRL = "girl"
-BOY = "boy"
-
-from collections import OrderedDict
-
-BIRTH_NAMES_COLUMNS = OrderedDict(
-    [
-        (DS, datetime),
-        (GENDER, str),
-        (NAME, str),
-        (NUM, int),
-        (STATE, str),
-        (NUM_BOYS, int),
-        (NUM_GIRLS, int),
-    ]
 )
+from tests.consts.us_states import US_STATES
+from tests.example_data.data_generator.base_generator import ExampleDataGenerator
+
+if TYPE_CHECKING:
+    from tests.example_data.data_generator.string_generator import StringGenerator
 
 
 class BirthNamesGenerator(ExampleDataGenerator):

tests/example_data/data_generator/birth_names/birth_names_generator_factory.py

@@ -14,23 +14,14 @@
 #  KIND, either express or implied.  See the License for the
 #  specific language governing permissions and limitations
 #  under the License.
-#
-#  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 __future__ import annotations
 
 from abc import ABC, abstractmethod
 
-from tests.common.example_data_generator.birth_names.birth_names_generator import (
+from tests.example_data.data_generator.birth_names.birth_names_generator import (
     BirthNamesGenerator,
 )
-from tests.common.example_data_generator.string_generator_factory import (
+from tests.example_data.data_generator.string_generator_factory import (
     StringGeneratorFactory,
 )
 

tests/example_data/data_generator/string_generator.py

@@ -14,15 +14,6 @@
 #  KIND, either express or implied.  See the License for the
 #  specific language governing permissions and limitations
 #  under the License.
-#
-#  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 random import choices, randint
 
 

tests/example_data/data_generator/string_generator_factory.py

@@ -16,7 +16,7 @@
 #  under the License.
 import string
 
-from tests.common.example_data_generator.string_generator import StringGenerator
+from tests.example_data.data_generator.string_generator import StringGenerator
 
 
 class StringGeneratorFactory:

tests/example_data/data_generator/tests/test_string_generator.py

@@ -16,11 +16,11 @@
 #  under the License.
 from unittest.mock import Mock, patch
 
-from tests.common.example_data_generator.string_generator import StringGenerator
+from tests.example_data.data_generator.string_generator import StringGenerator
 
 
-@patch("tests.common.example_data_generator.string_generator.choices")
+@patch("tests.example_data.data_generator.string_generator.choices")
-@patch("tests.common.example_data_generator.string_generator.randint")
+@patch("tests.example_data.data_generator.string_generator.randint")
 def test_string_generator(randint_mock: Mock, choices_mock: Mock):
     letters = "abcdets"
     min_len = 3

tests/example_data/data_loading/__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.

tests/example_data/data_loading/base_data_loader.py

@@ -0,0 +1,33 @@
+#  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 __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from tests.common.example_data.data_loading.data_definitions.types import Table
+
+
+class DataLoader(ABC):
+    @abstractmethod
+    def load_table(self, table: Table) -> None:
+        ...
+
+    @abstractmethod
+    def remove_table(self, table_name: str) -> None:
+        ...

tests/example_data/data_loading/data_definitions/__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.

tests/example_data/data_loading/data_definitions/birth_names.py

@@ -0,0 +1,64 @@
+#  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 sqlalchemy import DateTime, Integer, String
+
+from tests.consts.birth_names import (
+    DS,
+    GENDER,
+    NAME,
+    NUM,
+    NUM_BOYS,
+    NUM_GIRLS,
+    STATE,
+    TABLE_NAME,
+)
+from tests.example_data.data_loading.data_definitions.types import (
+    TableMetaData,
+    TableMetaDataFactory,
+)
+
+BIRTH_NAMES_COLUMNS = {
+    DS: DateTime,
+    GENDER: String(16),
+    NAME: String(255),
+    NUM: Integer,
+    STATE: String(10),
+    NUM_BOYS: Integer,
+    NUM_GIRLS: Integer,
+}
+
+BIRTH_NAMES_COLUMNS_WITHOUT_DATETIME = {
+    DS: String(255),
+    GENDER: String(16),
+    NAME: String(255),
+    NUM: Integer,
+    STATE: String(10),
+    NUM_BOYS: Integer,
+    NUM_GIRLS: Integer,
+}
+
+
+class BirthNamesMetaDataFactory(TableMetaDataFactory):
+    _datetime_type_support: bool
+
+    def __init__(self, datetime_type_support: bool = True):
+        self._datetime_type_support = datetime_type_support
+
+    def make(self) -> TableMetaData:
+        if self._datetime_type_support:
+            return TableMetaData(TABLE_NAME, BIRTH_NAMES_COLUMNS.copy())
+        return TableMetaData(TABLE_NAME, BIRTH_NAMES_COLUMNS_WITHOUT_DATETIME.copy())

tests/example_data/data_loading/data_definitions/types.py

@@ -0,0 +1,53 @@
+#  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.
+#
+#  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 abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, Optional
+
+from sqlalchemy.types import TypeEngine
+
+
+@dataclass
+class TableMetaData:
+    table_name: str
+    types: Optional[Dict[str, TypeEngine]]
+
+
+@dataclass
+class Table:
+    table_name: str
+    table_metadata: TableMetaData
+    data: Iterable[Dict[Any, Any]]
+
+
+class TableMetaDataFactory(ABC):
+    @abstractmethod
+    def make(self) -> TableMetaData:
+        ...
+
+    def make_table(self, data: Iterable[Dict[Any, Any]]) -> Table:
+        metadata = self.make()
+        return Table(metadata.table_name, metadata, data)

tests/example_data/data_loading/pandas/__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.

tests/example_data/data_loading/pandas/pandas_data_loader.py

@@ -0,0 +1,89 @@
+#  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.
+#
+#  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 __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, TYPE_CHECKING
+
+from pandas import DataFrame
+from sqlalchemy.inspection import inspect
+
+from tests.example_data.data_loading.base_data_loader import DataLoader
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Engine
+
+    from tests.example_data.data_loading.data_definitions.types import Table
+    from tests.example_data.data_loading.pandas.pands_data_loading_conf import (
+        PandasLoaderConfigurations,
+    )
+
+
+class PandasDataLoader(DataLoader):
+    _db_engine: Engine
+    _configurations: PandasLoaderConfigurations
+    _table_to_df_convertor: TableToDfConvertor
+
+    def __init__(
+        self,
+        db_engine: Engine,
+        config: PandasLoaderConfigurations,
+        table_to_df_convertor: TableToDfConvertor,
+    ) -> None:
+        self._db_engine = db_engine
+        self._configurations = config
+        self._table_to_df_convertor = table_to_df_convertor
+
+    def load_table(self, table: Table) -> None:
+        df = self._table_to_df_convertor.convert(table)
+        df.to_sql(
+            table.table_name,
+            self._db_engine,
+            if_exists=self._configurations.if_exists,
+            chunksize=self._configurations.chunksize,
+            index=self._configurations.index,
+            dtype=self._take_data_types(table),
+            method=self._configurations.method,
+            schema=self._detect_schema_name(),
+        )
+
+    def _detect_schema_name(self) -> Optional[str]:
+        return inspect(self._db_engine).default_schema_name
+
+    def _take_data_types(self, table: Table) -> Optional[Dict[str, str]]:
+        if metadata_table := table.table_metadata:
+            if types := metadata_table.types:
+                return types
+        return None
+
+    def remove_table(self, table_name: str) -> None:
+        self._db_engine.execute(f"DROP TABLE IF EXISTS {table_name}")
+
+
+class TableToDfConvertor(ABC):
+    @abstractmethod
+    def convert(self, table: Table) -> DataFrame:
+        ...

tests/example_data/data_loading/pandas/pands_data_loading_conf.py

@@ -0,0 +1,64 @@
+#  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 __future__ import annotations
+
+from typing import Any, Dict
+
+default_pandas_data_loader_config = {
+    "if_exists": "replace",
+    "chunksize": 500,
+    "index": False,
+    "method": "multi",
+    "strftime": "%Y-%m-%d %H:%M:%S",
+    "support_datetime_type": False,
+}
+
+
+class PandasLoaderConfigurations:
+    if_exists: str
+    chunksize: int
+    index: bool
+    method: str
+    strftime: str
+    support_datetime_type: bool
+
+    def __init__(
+        self,
+        *,
+        if_exists: str,
+        chunksize: int,
+        index: bool,
+        method: str,
+        strftime: str,
+        support_datetime_type: bool,
+    ):
+        self.if_exists = if_exists
+        self.chunksize = chunksize
+        self.index = index
+        self.method = method
+        self.strftime = strftime
+        self.support_datetime_type = support_datetime_type
+
+    @classmethod
+    def make_from_dict(cls, _dict: Dict[str, Any]) -> PandasLoaderConfigurations:
+        copy_dict = default_pandas_data_loader_config.copy()
+        copy_dict.update(_dict)
+        return PandasLoaderConfigurations(**copy_dict)  # type: ignore
+
+    @classmethod
+    def make_default(cls) -> PandasLoaderConfigurations:
+        return cls.make_from_dict({})

tests/example_data/data_loading/pandas/table_df_convertor.py

@@ -0,0 +1,46 @@
+#  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 __future__ import annotations
+
+from typing import Optional, TYPE_CHECKING
+
+from pandas import DataFrame
+
+from tests.example_data.data_loading.pandas.pandas_data_loader import TableToDfConvertor
+
+if TYPE_CHECKING:
+    from tests.example_data.data_loading.data_definitions.types import Table
+
+
+class TableToDfConvertorImpl(TableToDfConvertor):
+    convert_datetime_to_str: bool
+    _time_format: Optional[str]
+
+    def __init__(
+        self, convert_ds_to_datetime: bool, time_format: Optional[str] = None
+    ) -> None:
+        self.convert_datetime_to_str = convert_ds_to_datetime
+        self._time_format = time_format
+
+    def convert(self, table: Table) -> DataFrame:
+        df_rv = DataFrame(table.data)
+        if self._should_convert_datetime_to_str():
+            df_rv.ds = df_rv.ds.dt.strftime(self._time_format)
+        return df_rv
+
+    def _should_convert_datetime_to_str(self) -> bool:
+        return self.convert_datetime_to_str and self._time_format is not None

tests/fixtures/__init__.py

@@ -0,0 +1,18 @@
+#  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.
+
+pytest_plugins = "tests.fixtures.birth_names"

tests/fixtures/birth_names.py

@@ -0,0 +1,51 @@
+#  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 __future__ import annotations
+
+from typing import Callable, TYPE_CHECKING
+
+from pytest import fixture
+
+from tests.example_data.data_generator.birth_names.birth_names_generator_factory import (
+    BirthNamesGeneratorFactory,
+)
+from tests.example_data.data_loading.data_definitions.birth_names import (
+    BirthNamesMetaDataFactory,
+)
+
+if TYPE_CHECKING:
+    from tests.example_data.data_generator.birth_names.birth_names_generator import (
+        BirthNamesGenerator,
+    )
+    from tests.example_data.data_loading.data_definitions.types import Table
+
+
+@fixture(scope="session")
+def birth_names_data_generator() -> BirthNamesGenerator:
+    return BirthNamesGeneratorFactory.make()
+
+
+@fixture(scope="session")
+def birth_names_table_factory(
+    birth_names_data_generator: BirthNamesGenerator, support_datetime_type: bool,
+) -> Callable[[], Table]:
+    def _birth_names_table_factory() -> Table:
+        return BirthNamesMetaDataFactory(support_datetime_type).make_table(
+            data=birth_names_data_generator.generate()
+        )
+
+    return _birth_names_table_factory

tests/integration_tests/conftest.py

@@ -14,19 +14,23 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# isort:skip_file
+from __future__ import annotations
+
 import functools
-from typing import Any
+from typing import Any, Callable, Generator, Optional, TYPE_CHECKING
+from unittest.mock import patch
 
 import pytest
 from sqlalchemy.engine import Engine
-from unittest.mock import patch
 
-from tests.integration_tests.test_app import app
 from superset import db
 from superset.extensions import feature_flag_manager
 from superset.utils.core import json_dumps_w_dates
-from superset.utils.database import get_example_database
+from superset.utils.database import get_example_database, remove_database
+from tests.integration_tests.test_app import app
+
+if TYPE_CHECKING:
+    from superset.connectors.sqla.models import Database
 
 CTAS_SCHEMA_NAME = "sqllab_test_db"
 ADMIN_SCHEMA_NAME = "admin_database"
@@ -82,6 +86,36 @@ def drop_from_schema(engine: Engine, schema_name: str):
         engine.execute(f"DROP VIEW IF EXISTS {schema_name}.{tv[0]}")
 
 
+@pytest.fixture(scope="session")
+def example_db_provider() -> Callable[[], Database]:  # type: ignore
+    class _example_db_provider:
+        _db: Optional[Database] = None
+
+        def __call__(self) -> Database:
+            with app.app_context():
+                if self._db is None:
+                    self._db = get_example_database()
+                    self._load_lazy_data_to_decouple_from_session()
+
+                return self._db
+
+        def _load_lazy_data_to_decouple_from_session(self) -> None:
+            self._db.get_sqla_engine()  # type: ignore
+            self._db.backend  # type: ignore
+
+        def remove(self) -> None:
+            if self._db:
+                with app.app_context():
+                    remove_database(self._db)
+
+    _instance = _example_db_provider()
+
+    yield _instance
+
+    # TODO - can not use it until referenced objects will be deleted.
+    # _instance.remove()
+
+
 def setup_presto_if_needed():
     backend = app.config["SQLALCHEMY_EXAMPLES_URI"].split("://")[0]
     database = get_example_database()

tests/integration_tests/fixtures/birth_names_dashboard.py

@@ -14,12 +14,9 @@
 # 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, Optional
+from typing import Callable, List, Optional
 
-import pandas as pd
 import pytest
-from pandas import DataFrame
-from sqlalchemy import DateTime, String
 
 from superset import ConnectorRegistry, db
 from superset.connectors.sqla.models import SqlaTable
@@ -28,9 +25,8 @@ from superset.models.dashboard import Dashboard
 from superset.models.slice import Slice
 from superset.utils.core import get_example_default_schema
 from superset.utils.database import get_example_database
-from tests.common.example_data_generator.birth_names.birth_names_generator_factory import (
+from tests.example_data.data_loading.base_data_loader import DataLoader
-    BirthNamesGeneratorFactory,
+from tests.example_data.data_loading.data_definitions.types import Table
-)
 from tests.integration_tests.dashboard_utils import create_table_metadata
 from tests.integration_tests.test_app import app
 
@@ -38,31 +34,13 @@ BIRTH_NAMES_TBL_NAME = "birth_names"
 
 
 @pytest.fixture(scope="session")
-def load_birth_names_data():
+def load_birth_names_data(
-    with app.app_context():
+    birth_names_table_factory: Callable[[], Table], data_loader: DataLoader
-        database = get_example_database()
+):
-        df = _get_dataframe(database)
+    birth_names_table: Table = birth_names_table_factory()
-        dtype = {
+    data_loader.load_table(birth_names_table)
-            "ds": DateTime if database.backend != "presto" else String(255),
-            "gender": String(16),
-            "state": String(10),
-            "name": String(255),
-        }
-
-        df.to_sql(
-            BIRTH_NAMES_TBL_NAME,
-            database.get_sqla_engine(),
-            if_exists="replace",
-            chunksize=500,
-            dtype=dtype,
-            index=False,
-            method="multi",
-            schema=get_example_default_schema(),
-        )
     yield
-    with app.app_context():
+    data_loader.remove_table(birth_names_table.table_name)
-        engine = get_example_database().get_sqla_engine()
-        engine.execute("DROP TABLE IF EXISTS birth_names")
 
 
 @pytest.fixture()
@@ -137,15 +115,3 @@ def _cleanup(dash_id: int, slices_ids: List[int]) -> None:
     for slice_id in slices_ids:
         db.session.query(Slice).filter_by(id=slice_id).delete()
     db.session.commit()
-
-
-def _get_dataframe(database: Database) -> DataFrame:
-    data = _get_birth_names_data()
-    df = pd.DataFrame.from_dict(data)
-    if database.backend == "presto":
-        df.ds = df.ds.dt.strftime("%Y-%m-%d %H:%M:%S")
-    return df
-
-
-def _get_birth_names_data() -> List[Dict[Any, Any]]:
-    return list(BirthNamesGeneratorFactory.make().generate())