Spaces:
Paused
Paused
| # 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 | |
| from pandas import DataFrame | |
| from sqlalchemy import text | |
| from sqlalchemy.inspection import inspect | |
| from tests.common.logger_utils import log | |
| 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) -> str | None: | |
| return inspect(self._db_engine).default_schema_name | |
| def _take_data_types(self, table: Table) -> dict[str, str] | None: | |
| if metadata_table := table.table_metadata: | |
| types = metadata_table.types | |
| if types: | |
| return types | |
| return None | |
| def remove_table(self, table_name: str) -> None: | |
| with self._db_engine.begin() as conn: | |
| conn.execute(text(f"DROP TABLE IF EXISTS {table_name}")) | |
| class TableToDfConvertor(ABC): | |
| def convert(self, table: Table) -> DataFrame: ... | |