import datetime import pathlib from typing import Dict from anaxi.constants import PATHS 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 ) -> "CosmosDBDatabaseConfig": 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 ) -> "PostgresDatabaseConfig": 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 class StreamConfig: config_root_key: str = "streams" def __init__( self, stream_name: str, cosmos_database_id: str, cosmos_container_name: str, cutoff_timestamp: datetime.datetime, postgres_database: str, postgres_schema_name: str, ) -> "StreamConfig": self.stream_name = stream_name self.cosmos_database_id = cosmos_database_id self.cosmos_container_name = cosmos_container_name self.cutoff_timestamp = cutoff_timestamp self.postgres_database = postgres_database self.postgres_schema_name = postgres_schema_name @classmethod def from_dict(cls, a_dict) -> "StreamConfig": return StreamConfig(**a_dict) @classmethod def from_yaml(cls, path: pathlib.Path) -> Dict[str, "StreamConfig"]: yaml_contents = read_yaml(path) streams_contents = yaml_contents[StreamConfig.config_root_key] try: streams_configs = { stream_id: StreamConfig.from_dict(stream_config) for stream_id, stream_config in streams_contents.items() } except KeyError as e: logger.error( "Error reading streams config yaml file. The file seems malformed. Review the example in the repository." ) raise e return streams_configs def get_postgres_database_config_from_file(postgres_database): logger.info("Reading Postgres config file...") postgres_configs = PostgresDatabaseConfig.from_yaml(PATHS.postgres_config_file_path) logger.info(f"Found file with {len(postgres_configs)} entries.") try: relevant_postgres_config = postgres_configs[postgres_database] except KeyError as e: logger.error( f"Couldn't find a config entry for database with id: {postgres_database}" ) raise e return relevant_postgres_config def get_cosmos_database_config_from_file(cosmos_db_id): logger.info("Reading Cosmos DB config file...") cosmos_db_configs = CosmosDBDatabaseConfig.from_yaml( PATHS.cosmos_db_config_file_path ) logger.info(f"Found file with {len(cosmos_db_configs)} entries.") try: relevant_cosmos_db_config = cosmos_db_configs[cosmos_db_id] except KeyError as e: logger.error( f"Couldn't find a config entry for database with id: {cosmos_db_id}" ) raise e return relevant_cosmos_db_config def get_stream_config_from_file(stream_id): logger.info("Reading streams config file...") streams_configs = StreamConfig.from_yaml(PATHS.streams_config_file_path) logger.info(f"Found file with {len(streams_configs)} entries.") try: stream_config = streams_configs[stream_id] except KeyError: logger.error( f"Could not find the stream {stream_id} in the stream configurations file. You might be missing an entry in {PATHS.streams_config_file_path} or have made a typo." ) return stream_config