diff --git a/icon_service_base/database/create_config.py b/icon_service_base/database/create_config.py new file mode 100644 index 0000000..600aaee --- /dev/null +++ b/icon_service_base/database/create_config.py @@ -0,0 +1,33 @@ +from pathlib import Path +from typing import Optional, TypeVar + +from confz import BaseConfig, FileSource +from loguru import logger + +T = TypeVar("T", bound=BaseConfig) + + +class NoConfigSourceError(Exception): + pass + + +def create_config( + config_class: type[T], + config_folder: Optional[Path | str] = None, + config_file: str = "", +) -> T: + if config_class.CONFIG_SOURCES is not None or config_folder is not None: + config_sources = None + if config_folder is not None: + config_sources = FileSource(Path(config_folder) / config_file) + return config_class(config_sources=config_sources) + else: + error_msg = ( + "No 'database_config' folder found in the root directory. Please ensure " + "that a 'database_config' folder exists in your project's root directory. " + "Alternatively, you can provide a different config folder by passing its " + "path to the constructor." + ) + logger.error(error_msg) + + raise NoConfigSourceError(error_msg) diff --git a/icon_service_base/database/influxdb_session.py b/icon_service_base/database/influxdb_session.py index 3627761..3216592 100644 --- a/icon_service_base/database/influxdb_session.py +++ b/icon_service_base/database/influxdb_session.py @@ -4,19 +4,19 @@ from pathlib import Path from types import TracebackType from typing import Optional -from confz import FileSource -from influxdb_client import ( # type: ignore +from influxdb_client import ( Bucket, BucketRetentionRules, BucketsApi, InfluxDBClient, WriteApi, ) -from influxdb_client.client.write_api import SYNCHRONOUS # type: ignore -from influxdb_client.rest import ApiException # type: ignore +from influxdb_client.client.write_api import SYNCHRONOUS +from influxdb_client.rest import ApiException from loguru import logger from icon_service_base.database.config import InfluxDBConfig +from icon_service_base.database.create_config import create_config class InfluxDBSession: @@ -53,21 +53,11 @@ class InfluxDBSession: conf_folder: Path | str def __init__(self, config_folder: Optional[Path | str] = None) -> None: - config_folder = config_folder or getattr(self, "conf_folder", None) - if InfluxDBConfig.CONFIG_SOURCES is not None or config_folder is not None: - config_sources = None - if config_folder is not None: - config_sources = FileSource( - Path(config_folder) / "influxdb_config.yaml" - ) - self._config = InfluxDBConfig(config_sources=config_sources) - else: - logger.error( - "No config folder given. Please provide a config folder either by " - "passing it to the constructor or by setting the 'conf_folder' " - "attribute." - ) - return + self._config = create_config( + InfluxDBConfig, + config_folder=config_folder, + config_file="influxdb_config.yaml", + ) self.url = self._config.url self.token = str(self._config.token) diff --git a/icon_service_base/database/postgres_session.py b/icon_service_base/database/postgres_session.py index fca19bd..784b4a9 100644 --- a/icon_service_base/database/postgres_session.py +++ b/icon_service_base/database/postgres_session.py @@ -13,6 +13,7 @@ from loguru import logger from sqlmodel import Session, SQLModel, create_engine from icon_service_base.database.config import OperationMode, PostgreSQLConfig +from icon_service_base.database.create_config import create_config def json_loads_or_return_input(input_string: str) -> dict[str, Any] | Any: @@ -132,21 +133,11 @@ class PostgresDatabaseSession(Session): def __init__(self, config_folder: Optional[Path | str] = None) -> None: """Initializes a new session bound to the database engine.""" - config_folder = config_folder or getattr(self, "conf_folder", None) - if PostgreSQLConfig.CONFIG_SOURCES is not None or config_folder is not None: - config_sources = None - if config_folder is not None: - config_sources = FileSource( - Path(config_folder) / f"postgres_{OperationMode().environment}.yaml" - ) - self._config = PostgreSQLConfig(config_sources=config_sources) - else: - logger.error( - "No config folder given. Please provide a config folder either by " - "passing it to the constructor or by setting the 'conf_folder' " - "attribute." - ) - return + self._config = create_config( + PostgreSQLConfig, + config_folder=config_folder, + config_file=f"postgres_{OperationMode().environment}.yaml", + ) super().__init__( bind=create_engine(