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

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/common/example_data_generator/birth_names/__init__.py → 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/common/example_data_generator/base_generator.py → 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/common/example_data_generator/birth_names/birth_names_generator.py → 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 tests.common.example_data_generator.base_generator import ExampleDataGenerator
-from tests.common.example_data_generator.consts import US_STATES
-from tests.common.example_data_generator.string_generator import StringGenerator
+from typing import Any, Dict, Iterable, TYPE_CHECKING
 
-NUM_GIRLS = "num_girls"
-NUM_BOYS = "num_boys"
-STATE = "state"
-NUM = "num"
-NAME = "name"
-GENDER = "gender"
-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.birth_names import (
+    BOY,
+    DS,
+    GENDER,
+    GIRL,
+    NAME,
+    NUM,
+    NUM_BOYS,
+    NUM_GIRLS,
+    STATE,
 )
+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/common/example_data_generator/birth_names/birth_names_generator_factory.py → 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,
 )