Sets the default agent pool for the entirety of build definitions within an AZDO org

set_default_pool_for_org.py

#!/usr/bin/env python3

"""
Description:

Sets the default agent pool for the entirety of build definitions within an org

Usage instructions:

python set_default_pool_for_org.py

Requires the following runtime env vars:

* POOL_TO_USE (This should be the pool name)
* AZDO_PERSONAL_ACCESS_TOKEN
* AZDO_ORG_SERVICE_URL
* DRY_RUN (Defaults to true)

PyPI dependencies:

azure-devops==6.0.0b4
"""

import os
from typing import Sequence
from azure.devops.connection import Connection
from azure.devops.v6_0.build import BuildClient
from azure.devops.v6_0.build import BuildDefinitionReference
from azure.devops.v6_0.build.models import BuildDefinition
from azure.devops.v6_0.task_agent import TaskAgentPool
from azure.devops.v6_0.task_agent import TaskAgentClient
from azure.devops.v6_0.core import CoreClient
from azure.devops.v6_0.core import TeamProjectReference
from msrest.authentication import BasicAuthentication
from threading import Thread


def update_porject_default_agent_pool(
    project_name: str,
    pool_to_use: TaskAgentPool,
    azdo_build_client: BuildClient,
    azdo_task_agent_client: TaskAgentClient,
    dry_run: bool,
) -> None:
    queue_to_use, *_ = azdo_task_agent_client.get_agent_queues_for_pools(
        pool_ids=[pool_to_use.id], project=project_name
    )

    build_def_ref: BuildDefinitionReference
    for build_def_ref in azdo_build_client.get_definitions(
        project_name,
    ):
        build_def: BuildDefinition = azdo_build_client.get_definition(
            project_name, definition_id=build_def_ref.id
        )
        if build_def.queue.pool.id == pool_to_use.id:
            continue

        build_def.queue = queue_to_use

        print(f"🔨 Updating build definition {build_def.id} of {project_name}...")

        if not dry_run:
            updated_def: BuildDefinition = azdo_build_client.update_definition(
                build_def, project_name, build_def.id
            )
            assert updated_def.queue.pool.id == pool_to_use.id


def main() -> None:
    credentials = BasicAuthentication("", os.environ["AZDO_PERSONAL_ACCESS_TOKEN"])
    connection = Connection(
        base_url=os.environ["AZDO_ORG_SERVICE_URL"],
        creds=credentials,
    )

    azdo_core_client: CoreClient = connection.clients_v6_0.get_core_client()
    azdo_build_client: BuildClient = connection.clients_v6_0.get_build_client()
    azdo_task_agent_client: TaskAgentClient = (
        connection.clients_v6_0.get_task_agent_client()
    )

    pools: Sequence[TaskAgentPool] = azdo_task_agent_client.get_agent_pools()

    pool_to_use: TaskAgentPool = next(
        p for p in pools if p.name == os.environ["POOL_TO_USE"]
    )
    print(f"ℹ️ Pool to use: ID {pool_to_use.id} - NAME {pool_to_use.name}")

    projects: Sequence[TeamProjectReference] = azdo_core_client.get_projects()

    for project in projects:
        print(f"ℹ️ Updating definitions of {project.name} project...")
        t = Thread(
            target=update_porject_default_agent_pool,
            kwargs={
                "project_name": project.name,
                "pool_to_use": pool_to_use,
                "azdo_build_client": azdo_build_client,
                "azdo_task_agent_client": azdo_task_agent_client,
                "dry_run": os.environ.get("DRY_RUN", "true").lower() == "true"
            },
        )
        t.start()


if __name__ == "__main__":
    main()