mirror of https://github.com/apache/superset.git
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
This commit is contained in:
parent
88db2cc0ab
commit
4675ca31c5
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
)
|
|
@ -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.
|
|
@ -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"
|
|
@ -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",
|
||||
]
|
|
@ -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.
|
|
@ -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
|
||||
|
|
@ -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.
|
|
@ -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.common.example_data_generator.consts import US_STATES
|
||||
from tests.common.example_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),
|
||||
]
|
||||
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):
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -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:
|
|
@ -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
|
|
@ -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.
|
|
@ -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:
|
||||
...
|
|
@ -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.
|
|
@ -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())
|
|
@ -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)
|
|
@ -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.
|
|
@ -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:
|
||||
...
|
|
@ -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({})
|
|
@ -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
|
|
@ -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"
|
|
@ -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
|
|
@ -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()
|
||||
|
|
|
@ -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 (
|
||||
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 @@ BIRTH_NAMES_TBL_NAME = "birth_names"
|
|||
|
||||
|
||||
@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())
|
||||
|
|
Loading…
Reference in New Issue