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

Code before:

def save_indirect_host_entries_of_job(job: Job) -> None:
    "Once we have a job and we know that we want to do indirect host processing, this is called"
    job_event_queries = fetch_job_event_query(job)
    records = build_indirect_host_data(job, job_event_queries)
    IndirectManagedNodeAudit.objects.bulk_create(records)
    job.event_queries_processed = True


def cleanup_old_indirect_host_entries() -> None:
    """
    We assume that indirect host audit results older than one week have already been collected for analysis
    and can be cleaned up
    """
    limit = now() - timedelta(days=settings.INDIRECT_HOST_AUDIT_RECORD_MAX_AGE_DAYS)
    IndirectManagedNodeAudit.objects.filter(created__lt=limit).delete()


@task(queue=get_task_queuename)
def save_indirect_host_entries(job_id: int, wait_for_events: bool = True) -> None:
    with transaction.atomic():
        """
        Pre-emptively set the job marker to 'events processed'. This prevents other instances from running the
        same task.
        """
        job = Job.objects.select_for_update().get(id=job_id)

        if job.event_queries_processed is True:
            # this can mean one of two things:
            # 1. another instance has already processed the events of this job
            # 2. the artifacts_handler has not yet been called for this job
            return

        job.event_queries_processed = True
        job.save(update_fields=['event_queries_processed'])

    try:
        job = Job.objects.get(id=job_id)
    except Job.DoesNotExist:
        logger.debug(f'Job {job_id} seems to be deleted, bailing from save_indirect_host_entries')
        return

    if wait_for_events:
        # Gate running this task on the job having all events processed, not just EOF or playbook_on_stats
        current_events = 0
        for _ in range(10):
            current_events = job.job_events.count()
            if current_events >= job.emitted_events:
                break
            logger.debug(f'Waiting for job_id={job_id} to finish processing events, currently {current_events} < {job.emitted_events}')
            time.sleep(0.2)

Review comment (bug):
It looks like this needs to be combined with the lines right after it:

```python
    try:
        job = Job.objects.get(id=job_id)
    except Job.DoesNotExist:
        logger.debug(f'Job {job_id} seems to be deleted, bailing from save_indirect_host_entries')
        return
```

Otherwise it'll result in unhandled exceptions. But doing the exact same thing is a fine, a `return` will exit the context managers as well.

Code after:

def save_indirect_host_entries_of_job(job: Job) -> None:
    "Once we have a job and we know that we want to do indirect host processing, this is called"
    job_event_queries = fetch_job_event_query(job)
    records = build_indirect_host_data(job, job_event_queries)
    IndirectManagedNodeAudit.objects.bulk_create(records)
    job.event_queries_processed = True


def cleanup_old_indirect_host_entries() -> None:
    """
    We assume that indirect host audit results older than one week have already been collected for analysis
    and can be cleaned up
    """
    limit = now() - timedelta(days=settings.INDIRECT_HOST_AUDIT_RECORD_MAX_AGE_DAYS)
    IndirectManagedNodeAudit.objects.filter(created__lt=limit).delete()


@task(queue=get_task_queuename)
def save_indirect_host_entries(job_id: int, wait_for_events: bool = True) -> None:
    try:
        job = Job.objects.get(id=job_id)
    except Job.DoesNotExist:
        logger.debug(f'Job {job_id} seems to be deleted, bailing from save_indirect_host_entries')
        return

    if wait_for_events:
        # Gate running this task on the job having all events processed, not just EOF or playbook_on_stats
        current_events = 0
        for _ in range(10):
            current_events = job.job_events.count()
            if current_events >= job.emitted_events:
                break
            logger.debug(f'Waiting for job_id={job_id} to finish processing events, currently {current_events} < {job.emitted_events}')
            time.sleep(0.2)
        else:
            logger.warning(f'Event count {current_events} < {job.emitted_events} for job_id={job_id}, delaying processing of indirect host tracking')
            return

    with transaction.atomic():
        """
        Pre-emptively set the job marker to 'events processed'. This prevents other instances from running the
        same task.
        """
        try:
            job = Job.objects.select_for_update().get(id=job_id)
        except job.DoesNotExist:
            logger.debug(f'Job {job_id} seems to be deleted, bailing from save_indirect_host_entries')
            return

        if job.event_queries_processed is True: