from psycopg2.sql import SQL from anaxi.config import CosmosDBDatabaseConfig, PostgresDatabaseConfig from anaxi.constants import PATHS from anaxi.cosmos_tools import create_cosmos_client_from_config from anaxi.logging import get_anaxi_logger from anaxi.postgres_tools import simply_query logger = get_anaxi_logger(__name__) def run_cosmos_db_healthcheck_process(cosmos_db_id: str) -> None: 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 logger.info("Creating client...") cosmos_client = create_cosmos_client_from_config(relevant_cosmos_db_config) logger.info("Client created.") logger.info("Sending a SELECT 1;") response = cosmos_client.query_databases(query="SELECT 1") logger.info(f"Response: {response.next()}") return def run_postgres_healthcheck_process(postgres_database: str) -> None: 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 logger.info("Connecting and sending a SELECT 1...") query_result = simply_query(config=relevant_postgres_config, query=SQL("SELECT 1;")) logger.info(f"Response: {query_result}") return