Revert "refactor(example_data): replace the way the birth_names data …

…is loaded to DB (apache#18060)"

This reverts commit 4675ca3.

superset/utils/database.py

@@ -39,10 +39,7 @@ def get_or_create_db(
     from superset.models import core as models
 
     database = (
-        db.session.query(models.Database)
-        .filter_by(database_name=database_name)
-        .autoflush(False)
-        .first()
+        db.session.query(models.Database).filter_by(database_name=database_name).first()
     )
 
     # databases with a fixed UUID
@@ -75,14 +72,3 @@ 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/example_data/data_generator/base_generator.py → tests/common/example_data_generator/base_generator.py

@@ -14,6 +14,16 @@
 #  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/consts/us_states.py → tests/common/example_data_generator/birth_names/__init__.py

@@ -31,56 +31,3 @@
 #  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/data_generator/birth_names/birth_names_generator.py → tests/common/example_data_generator/birth_names/birth_names_generator.py

@@ -14,28 +14,66 @@
 #  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, TYPE_CHECKING
+from typing import Any, Dict, Iterable
 
-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
+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
 
-if TYPE_CHECKING:
-    from tests.example_data.data_generator.string_generator import StringGenerator
+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),
+    ]
+)
 
 
 class BirthNamesGenerator(ExampleDataGenerator):

tests/example_data/data_generator/birth_names/birth_names_generator_factory.py → tests/common/example_data_generator/birth_names/birth_names_generator_factory.py

@@ -14,14 +14,23 @@
 #  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.example_data.data_generator.birth_names.birth_names_generator import (
+from tests.common.example_data_generator.birth_names.birth_names_generator import (
     BirthNamesGenerator,
 )
-from tests.example_data.data_generator.string_generator_factory import (
+from tests.common.example_data_generator.string_generator_factory import (
     StringGeneratorFactory,
 )
 

tests/example_data/data_generator/string_generator.py → tests/common/example_data_generator/string_generator.py

@@ -14,6 +14,15 @@
 #  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 → tests/common/example_data_generator/string_generator_factory.py

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

tests/example_data/data_generator/tests/test_string_generator.py → tests/common/example_data_generator/tests/test_string_generator.py

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