data-anaxi/anaxi/config.py

import logging
import pathlib
from typing import Dict

from anaxi.file_persistence import read_yaml
from anaxi.logging import get_anaxi_logger

logger = get_anaxi_logger(__name__)


class CosmosDBDatabaseConfig:

    config_root_key: str = "cosmos-databases"

    def __init__(self, host: str, database_id: str, master_key: str) -> None:
        self.host = host
        self.database_id = database_id
        self.master_key = master_key

    @classmethod
    def from_dict(cls, a_dict) -> "CosmosDBDatabaseConfig":
        return CosmosDBDatabaseConfig(**a_dict)

    @classmethod
    def from_yaml(cls, path: pathlib.Path) -> Dict[str, "CosmosDBDatabaseConfig"]:
        yaml_contents = read_yaml(path)

        cosmos_db_contents = yaml_contents[CosmosDBDatabaseConfig.config_root_key]

        try:
            database_configs = {
                db_id: CosmosDBDatabaseConfig.from_dict(db_config)
                for db_id, db_config in cosmos_db_contents.items()
            }
        except KeyError as e:
            logger.error(
                "Error reading Cosmos DB config yaml file. The file seems malformed. Review the example in the repository."
            )
            raise e

        return database_configs


class PostgresDatabaseConfig:

    config_root_key: str = "postgres-databases"

    def __init__(self, host: str, port: int, dbname: str, user: str, password) -> None:
        self.host = host
        self.port = port
        self.dbname = dbname
        self.user = user
        self.password = password

    @classmethod
    def from_dict(cls, a_dict) -> "PostgresDatabaseConfig":
        return PostgresDatabaseConfig(**a_dict)

    @classmethod
    def from_yaml(cls, path: pathlib.Path) -> Dict[str, "PostgresDatabaseConfig"]:
        yaml_contents = read_yaml(path)

        postgres_contents = yaml_contents[PostgresDatabaseConfig.config_root_key]

        try:
            database_configs = {
                db_name: PostgresDatabaseConfig.from_dict(db_config)
                for db_name, db_config in postgres_contents.items()
            }
        except KeyError as e:
            logger.error(
                "Error reading Postgres config yaml file. The file seems malformed. Review the example in the repository."
            )
            raise e

        return database_configs