data-anaxi/anaxi/processes.py

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