Repository: ansible/awx
File: awx/main/tasks/jobs.py (Python)

Code before:
# Django flags
from flags.state import flag_enabled

# Workload Identity
from ansible_base.lib.workload_identity.controller import AutomationControllerJobScope

logger = logging.getLogger('awx.main.tasks.jobs')


def populate_claims_for_workload(workload):
    """
    Extract JWT claims from a Job workload for the aap_controller_automation_job scope.
    Missing or null attributes are omitted from the returned claims dictionary.
    """

    def safe_get(*attrs):
        """Safely traverse object attributes and dictionary keys."""
        current = workload
        for attr in attrs:
            if current is None:
                return None
            if isinstance(current, dict):
                current = current.get(attr)
            else:
                current = getattr(current, attr, None)
        return current

    claim_mappings = [
        (AutomationControllerJobScope.CLAIM_JOB_ID, ('id',)),
        (AutomationControllerJobScope.CLAIM_JOB_NAME, ('name',)),
        (AutomationControllerJobScope.CLAIM_JOB_TYPE, ('job_type',)),
        (AutomationControllerJobScope.CLAIM_LAUNCH_TYPE, ('launch_type',)),
        (AutomationControllerJobScope.CLAIM_PLAYBOOK_NAME, ('playbook',)),
        (AutomationControllerJobScope.CLAIM_LAUNCHED_BY_USER_NAME, ('launched_by', 'name')),
        (AutomationControllerJobScope.CLAIM_LAUNCHED_BY_USER_ID, ('launched_by', 'id')),
        (AutomationControllerJobScope.CLAIM_ORGANIZATION_NAME, ('organization', 'name')),
        (AutomationControllerJobScope.CLAIM_ORGANIZATION_ID, ('organization', 'id')),
        (AutomationControllerJobScope.CLAIM_INVENTORY_NAME, ('inventory', 'name')),
        (AutomationControllerJobScope.CLAIM_INVENTORY_ID, ('inventory', 'id')),
        (AutomationControllerJobScope.CLAIM_EXECUTION_ENVIRONMENT_NAME, ('execution_environment', 'name')),
        (AutomationControllerJobScope.CLAIM_EXECUTION_ENVIRONMENT_ID, ('execution_environment', 'id')),
        (AutomationControllerJobScope.CLAIM_PROJECT_NAME, ('project', 'name')),
        (AutomationControllerJobScope.CLAIM_PROJECT_ID, ('project', 'id')),
        (AutomationControllerJobScope.CLAIM_JOB_TEMPLATE_NAME, ('job_template', 'name')),
        (AutomationControllerJobScope.CLAIM_JOB_TEMPLATE_ID, ('job_template', 'id')),
        (AutomationControllerJobScope.CLAIM_UNIFIED_JOB_TEMPLATE_NAME, ('unified_job_template', 'name')),
        (AutomationControllerJobScope.CLAIM_UNIFIED_JOB_TEMPLATE_ID, ('unified_job_template', 'id')),
        (AutomationControllerJobScope.CLAIM_INSTANCE_GROUP_NAME, ('instance_group', 'name')),
        (AutomationControllerJobScope.CLAIM_INSTANCE_GROUP_ID, ('instance_group', 'id')),
    ]


Review comment (refactor):
I feel like we can trim this a little bit. Python has a builtin `.get` for dictionaries. Think we can refactor this so we check if `current` is a dict once and then use the builtin method?

Code after:

# Django flags
from flags.state import flag_enabled

# Workload Identity
from ansible_base.lib.workload_identity.controller import AutomationControllerJobScope

logger = logging.getLogger('awx.main.tasks.jobs')


def populate_claims_for_workload(unified_job) -> dict:
    """
    Extract JWT claims from a Controller workload for the aap_controller_automation_job scope.
    """

    # Related objects in the UnifiedJob model, applies to all job types
    organization = getattr_dne(unified_job, 'organization')
    ujt = getattr_dne(unified_job, 'unified_job_template')
    instance_group = getattr_dne(unified_job, 'instance_group')

    claims = {
        AutomationControllerJobScope.CLAIM_JOB_ID: unified_job.id,
        AutomationControllerJobScope.CLAIM_JOB_NAME: unified_job.name,
        AutomationControllerJobScope.CLAIM_LAUNCH_TYPE: unified_job.launch_type,
    }

    # Related objects in the UnifiedJob model, applies to all job types
    # null cases are omitted because of OIDC
    if organization := getattr_dne(unified_job, 'organization'):
        claims[AutomationControllerJobScope.CLAIM_ORGANIZATION_NAME] = organization.name
        claims[AutomationControllerJobScope.CLAIM_ORGANIZATION_ID] = organization.id

    if ujt := getattr_dne(unified_job, 'unified_job_template'):
        claims[AutomationControllerJobScope.CLAIM_UNIFIED_JOB_TEMPLATE_NAME] = ujt.name
        claims[AutomationControllerJobScope.CLAIM_UNIFIED_JOB_TEMPLATE_ID] = ujt.id

    if instance_group := getattr_dne(unified_job, 'instance_group'):
        claims[AutomationControllerJobScope.CLAIM_INSTANCE_GROUP_NAME] = instance_group.name
        claims[AutomationControllerJobScope.CLAIM_INSTANCE_GROUP_ID] = instance_group.id

    # Related objects on concrete models, may not be valid for type of unified_job
    if inventory := getattr_dne(unified_job, 'inventory', None):
        claims[AutomationControllerJobScope.CLAIM_INVENTORY_NAME] = inventory.name
        claims[AutomationControllerJobScope.CLAIM_INVENTORY_ID] = inventory.id

    if execution_environment := getattr_dne(unified_job, 'execution_environment', None):
        claims[AutomationControllerJobScope.CLAIM_EXECUTION_ENVIRONMENT_NAME] = execution_environment.name
        claims[AutomationControllerJobScope.CLAIM_EXECUTION_ENVIRONMENT_ID] = execution_environment.id

    if project := getattr_dne(unified_job, 'project', None):
        claims[AutomationControllerJobScope.CLAIM_PROJECT_NAME] = project.name