Using cloud workload identity with Prefect Managed work pools

import asyncio
import sys

from prefect.client.cloud import get_cloud_client
from prefect.client.orchestration import get_client
from prefect.client.schemas.actions import WorkPoolUpdate
from prefect.utilities.collections import deep_merge


async def update_work_pool_template(work_pool_name: str):
    async with get_cloud_client() as cloud_client, get_client() as client:
        metadata = await cloud_client.read_worker_metadata()
        work_pool = await client.read_work_pool(work_pool_name)
        work_pool_type = work_pool.type

        # Dict has top-level keys containing worker package names
        # for example: { "prefect": { "prefect-agent": ...}, "prefect-push-work-pools": {"prefect-managed": ...}}
        for key, value in metadata.items():
            if work_pool_type in value.keys():
                new_base_template = value[work_pool_type]["default_base_job_configuration"]
                break
        else:
            raise ValueError(
                f"Worker metadata does not contain configuration for work pool type '{work_pool_type}'"
            )

        existing_template = work_pool.base_job_template
        merged_template = deep_merge(
            new_base_template,
            existing_template,
        )
        await client.update_work_pool(
            work_pool_name,
            work_pool=WorkPoolUpdate(base_job_template=merged_template),
        )


if __name__ == "__main__":
    if len(sys.argv) != 2:
        raise SystemExit(f"Usage: {sys.argv[0]} <work-pool-name>")

    # Set your profile using `prefect profile use <PROFILE_NAME>` to target the correct account/workspace
    asyncio.run(update_work_pool_template(sys.argv[1]))