diff --git a/superset/utils/database.py b/superset/utils/database.py index 146412d57ceb0..1375611ab877e 100644 --- a/superset/utils/database.py +++ b/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() diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000000000..6350d3235c5b6 --- /dev/null +++ b/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 + ) diff --git a/tests/consts/__init__.py b/tests/consts/__init__.py new file mode 100644 index 0000000000000..87f6e4d72b350 --- /dev/null +++ b/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. diff --git a/tests/consts/birth_names.py b/tests/consts/birth_names.py new file mode 100644 index 0000000000000..540ce6f7de852 --- /dev/null +++ b/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" diff --git a/tests/common/example_data_generator/birth_names/__init__.py b/tests/consts/us_states.py similarity index 75% rename from tests/common/example_data_generator/birth_names/__init__.py rename to tests/consts/us_states.py index e5379060d1f17..741ef816cac39 100644 --- a/tests/common/example_data_generator/birth_names/__init__.py +++ b/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", +] diff --git a/tests/example_data/__init__.py b/tests/example_data/__init__.py new file mode 100644 index 0000000000000..87f6e4d72b350 --- /dev/null +++ b/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. diff --git a/tests/common/example_data_generator/__init__.py b/tests/example_data/data_generator/__init__.py similarity index 100% rename from tests/common/example_data_generator/__init__.py rename to tests/example_data/data_generator/__init__.py diff --git a/tests/common/example_data_generator/base_generator.py b/tests/example_data/data_generator/base_generator.py similarity index 72% rename from tests/common/example_data_generator/base_generator.py rename to tests/example_data/data_generator/base_generator.py index ffd30e99c92d9..023b929091439 100644 --- a/tests/common/example_data_generator/base_generator.py +++ b/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 diff --git a/tests/example_data/data_generator/birth_names/__init__.py b/tests/example_data/data_generator/birth_names/__init__.py new file mode 100644 index 0000000000000..87f6e4d72b350 --- /dev/null +++ b/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. diff --git a/tests/common/example_data_generator/birth_names/birth_names_generator.py b/tests/example_data/data_generator/birth_names/birth_names_generator.py similarity index 55% rename from tests/common/example_data_generator/birth_names/birth_names_generator.py rename to tests/example_data/data_generator/birth_names/birth_names_generator.py index 92deef3e2cde1..2b68abbd4f12b 100644 --- a/tests/common/example_data_generator/birth_names/birth_names_generator.py +++ b/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): diff --git a/tests/common/example_data_generator/birth_names/birth_names_generator_factory.py b/tests/example_data/data_generator/birth_names/birth_names_generator_factory.py similarity index 78% rename from tests/common/example_data_generator/birth_names/birth_names_generator_factory.py rename to tests/example_data/data_generator/birth_names/birth_names_generator_factory.py index 20abebdc56c22..0eba3f1c322aa 100644 --- a/tests/common/example_data_generator/birth_names/birth_names_generator_factory.py +++ b/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, ) diff --git a/tests/common/example_data_generator/consts.py b/tests/example_data/data_generator/consts.py similarity index 100% rename from tests/common/example_data_generator/consts.py rename to tests/example_data/data_generator/consts.py diff --git a/tests/common/example_data_generator/string_generator.py b/tests/example_data/data_generator/string_generator.py similarity index 78% rename from tests/common/example_data_generator/string_generator.py rename to tests/example_data/data_generator/string_generator.py index 101f49c9c47b6..103ceefcb52e1 100644 --- a/tests/common/example_data_generator/string_generator.py +++ b/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 diff --git a/tests/common/example_data_generator/string_generator_factory.py b/tests/example_data/data_generator/string_generator_factory.py similarity index 95% rename from tests/common/example_data_generator/string_generator_factory.py rename to tests/example_data/data_generator/string_generator_factory.py index caca7e2a20e42..2cad2f785a226 100644 --- a/tests/common/example_data_generator/string_generator_factory.py +++ b/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: diff --git a/tests/common/example_data_generator/tests/__init__.py b/tests/example_data/data_generator/tests/__init__.py similarity index 100% rename from tests/common/example_data_generator/tests/__init__.py rename to tests/example_data/data_generator/tests/__init__.py diff --git a/tests/common/example_data_generator/tests/test_string_generator.py b/tests/example_data/data_generator/tests/test_string_generator.py similarity index 85% rename from tests/common/example_data_generator/tests/test_string_generator.py rename to tests/example_data/data_generator/tests/test_string_generator.py index 7088b752d5e29..65fdd18ce6b3f 100644 --- a/tests/common/example_data_generator/tests/test_string_generator.py +++ b/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.common.example_data_generator.string_generator.randint") +@patch("tests.example_data.data_generator.string_generator.choices") +@patch("tests.example_data.data_generator.string_generator.randint") def test_string_generator(randint_mock: Mock, choices_mock: Mock): letters = "abcdets" min_len = 3 diff --git a/tests/example_data/data_loading/__init__.py b/tests/example_data/data_loading/__init__.py new file mode 100644 index 0000000000000..87f6e4d72b350 --- /dev/null +++ b/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. diff --git a/tests/example_data/data_loading/base_data_loader.py b/tests/example_data/data_loading/base_data_loader.py new file mode 100644 index 0000000000000..770150c14fabd --- /dev/null +++ b/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: + ... diff --git a/tests/example_data/data_loading/data_definitions/__init__.py b/tests/example_data/data_loading/data_definitions/__init__.py new file mode 100644 index 0000000000000..87f6e4d72b350 --- /dev/null +++ b/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. diff --git a/tests/example_data/data_loading/data_definitions/birth_names.py b/tests/example_data/data_loading/data_definitions/birth_names.py new file mode 100644 index 0000000000000..64c1a50d49bd4 --- /dev/null +++ b/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()) diff --git a/tests/example_data/data_loading/data_definitions/types.py b/tests/example_data/data_loading/data_definitions/types.py new file mode 100644 index 0000000000000..e393019e0192f --- /dev/null +++ b/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) diff --git a/tests/example_data/data_loading/pandas/__init__.py b/tests/example_data/data_loading/pandas/__init__.py new file mode 100644 index 0000000000000..87f6e4d72b350 --- /dev/null +++ b/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. diff --git a/tests/example_data/data_loading/pandas/pandas_data_loader.py b/tests/example_data/data_loading/pandas/pandas_data_loader.py new file mode 100644 index 0000000000000..8f50bdbb6ce91 --- /dev/null +++ b/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: + ... diff --git a/tests/example_data/data_loading/pandas/pands_data_loading_conf.py b/tests/example_data/data_loading/pandas/pands_data_loading_conf.py new file mode 100644 index 0000000000000..1c43adc9316e9 --- /dev/null +++ b/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({}) diff --git a/tests/example_data/data_loading/pandas/table_df_convertor.py b/tests/example_data/data_loading/pandas/table_df_convertor.py new file mode 100644 index 0000000000000..3b3c6ed57787a --- /dev/null +++ b/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 diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py new file mode 100644 index 0000000000000..18d5e2f245a23 --- /dev/null +++ b/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" diff --git a/tests/fixtures/birth_names.py b/tests/fixtures/birth_names.py new file mode 100644 index 0000000000000..bcf2a5aa91b8a --- /dev/null +++ b/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 diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index ec22df22d4fb5..2765a1d909375 100644 --- a/tests/integration_tests/conftest.py +++ b/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() diff --git a/tests/integration_tests/fixtures/birth_names_dashboard.py b/tests/integration_tests/fixtures/birth_names_dashboard.py index 7d9616f57f895..685cf43f581fa 100644 --- a/tests/integration_tests/fixtures/birth_names_dashboard.py +++ b/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.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 ( - BirthNamesGeneratorFactory, -) +from tests.example_data.data_loading.base_data_loader import DataLoader +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 @@ @pytest.fixture(scope="session") -def load_birth_names_data(): - with app.app_context(): - database = get_example_database() - df = _get_dataframe(database) - dtype = { - "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(), - ) +def load_birth_names_data( + birth_names_table_factory: Callable[[], Table], data_loader: DataLoader +): + birth_names_table: Table = birth_names_table_factory() + data_loader.load_table(birth_names_table) yield - with app.app_context(): - engine = get_example_database().get_sqla_engine() - engine.execute("DROP TABLE IF EXISTS birth_names") + data_loader.remove_table(birth_names_table.table_name) @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())