data-anaxi/anaxi/cosmos_tools.py

import azure.cosmos.cosmos_client as cosmos_client
from azure.cosmos.container import ContainerProxy

from anaxi.config import CosmosDBDatabaseConfig


def create_cosmos_client_from_config(
    config: CosmosDBDatabaseConfig,
) -> cosmos_client.CosmosClient:
    return cosmos_client.CosmosClient(
        url=config.host,
        credential={"masterKey": config.master_key},
    )


def get_container_client(
    config: CosmosDBDatabaseConfig, container_name: str
) -> ContainerProxy:
    cosmos_client = create_cosmos_client_from_config(config)
    database_client = cosmos_client.get_database_client(database=config.database_id)
    container_client = database_client.get_container_client(container=container_name)

    return container_client