ronantakizawa/codereview-bench huggingface dataset:
Repository: ansible/awx
File: awx/main/tests/live/tests/test_cleanup_task.py (Python)

Code before:
@pytest.mark.parametrize(
    'worktype',
    ('remote', 'local'),
)
def test_tagless_image(podman_image_generator, worktype):
    """
    Ensure podman images on Control and Hybrid nodes are deleted during cleanup.
    """
    podman_image_generator()

    dangling_image = next((image for image in get_podman_images() if image.get('Dangling', False)), None)
    assert dangling_image

    with (
        (
            mock.patch('awx.main.tasks.receptor.run_until_complete', lambda *args, **kwargs: run_until_complete(*args, worktype='local', ttl=None, **kwargs))
            if worktype == 'local'
            else nullcontext()
        ),
        (
            mock.patch('awx.main.tasks.system.CleanupImagesAndFilesHelper.get_execution_instances', lambda: [Instance.objects.me()])
            if worktype == 'local'
            else nullcontext()
        ),
    ):
        _cleanup_images_and_files(image_prune=True)

    for image in get_podman_images():
        assert image['Id'] != dangling_image['Id']


Review comment (suggestion):
My only nit pick about this test is now that this call with run _both_ the control & execution forms of the logic. Meaning it will run `ansible_runner.cleanup.run_cleanup(runner_cleanup_kwargs)` before the command over receptor. So if you're testing that the CLI args do not error, this test is sufficient. But it doesn't test that the command ran over receptor _actually cleans up the images_ because this would have already been done in the control form of the logic. To address this, I am pretty sure you'd have to break up `_cleanup_images_and_files` into 2 methods or something like that. I think that would be better for cleanness and also better testing. But as for the first goal of regression testing that we don't get errors from the execution node command, this is doing the job already.

Code after:
@pytest.mark.parametrize(
    'worktype',
    ('remote', 'local'),
)
def test_tagless_image(podman_image_generator, worktype: str):
    """
    Ensure podman images on Control and Hybrid nodes are deleted during cleanup.
    """
    podman_image_generator()

    dangling_image = next((image for image in get_podman_images() if image.get('Dangling', False)), None)
    assert dangling_image

    instance_me = Instance.objects.me()

    match worktype:
        case 'local':
            CleanupImagesAndFiles.run_local(instance_me, image_prune=True)
        case 'remote':
            with (
                mock.patch(
                    'awx.main.tasks.receptor.run_until_complete', lambda *args, **kwargs: run_until_complete(*args, worktype='local', ttl=None, **kwargs)
                ),
                mock.patch('awx.main.tasks.system.CleanupImagesAndFiles.get_execution_instances', lambda: [Instance.objects.me()]),
            ):
                CleanupImagesAndFiles.run_remote(instance_me, image_prune=True)
        case _:
            raise ValueError(f'worktype "{worktype}" not supported.')

    for image in get_podman_images():
        assert image['Id'] != dangling_image['Id']
Repository: ansible/awx
File: awx/main/utils/analytics_proxy.py (Python)

Code before:
'''
Proxy requests Analytics requests
'''

import time

from enum import Enum

from typing import Optional, Any

import requests

DEFAULT_OIDC_ENDPOINT = 'https://sso.redhat.com/auth/realms/redhat-external'


class TokenError(requests.RequestException):
    '''
    Raised when token generation request fails.

    Useful for differentiating request failure for make_request() vs.
    other requests issued to get a token i.e.:

    try:
      client = OIDCClient(...)
      client.make_request(...)
    except TokenGenerationError as e:
        print(f"Token generation failed due to {e.__cause__}")
    except requests.RequestException:
        print("API request failed)
    '''

    def __init__(self, message="Token generation request failed", response=None):
        super().__init__(message)
        self.response = response  # Store the response for debugging


def _now(reason: str):
    '''

Review comment (security):
```suggestion
DEFAULT_OIDC_ENDPOINT = 'https://sso.redhat.com/auth/realms/redhat-external/protocol/openid-connect/token'
```

Code after:
'''
Proxy requests Analytics requests
'''

import time

from enum import Enum

from typing import Optional, Any

import requests

DEFAULT_OIDC_ENDPOINT = 'https://sso.redhat.com/auth/realms/redhat-external/protocol/openid-connect/token'


class TokenError(requests.RequestException):
    '''
    Raised when token generation request fails.

    Useful for differentiating request failure for make_request() vs.
    other requests issued to get a token i.e.:

    try:
      client = OIDCClient(...)
      client.make_request(...)
    except TokenGenerationError as e:
        print(f"Token generation failed due to {e.__cause__}")
    except requests.RequestException:
        print("API request failed)
    '''

    def __init__(self, message="Token generation request failed", response=None):
        super().__init__(message)
        self.response = response  # Store the response for debugging


def _now(reason: str):
    '''
Repository: ansible/awx
File: awx/main/tests/unit/models/test_jobs.py (Python)

Code before:
        with open(filepath, 'w') as f:
            f.write('not valid json!')
            f.flush()
            new_modification_time = time.time() + 3600
            os.utime(filepath, (new_modification_time, new_modification_time))

    finish_fact_cache(hosts, fact_cache, last_modified)

    bulk_update.assert_not_called()


def test_finish_job_fact_cache_clear(hosts, mocker, ref_time, tmpdir):
    fact_cache = os.path.join(tmpdir, 'facts')
    last_modified, _ = start_fact_cache(hosts, fact_cache, timeout=0)

    bulk_update = mocker.patch('django.db.models.query.QuerySet.bulk_update')

    os.remove(os.path.join(fact_cache, hosts[1].name))
    finish_fact_cache(hosts, fact_cache, last_modified)

    for host in (hosts[0], hosts[2], hosts[3]):
        assert host.ansible_facts == {"a": 1, "b": 2}
        assert host.ansible_facts_modified == ref_time
    assert hosts[1].ansible_facts == {}
    assert hosts[1].ansible_facts_modified > ref_time
    bulk_update.assert_called_once_with([hosts[1]], ['ansible_facts', 'ansible_facts_modified'])


Review comment (suggestion):
Something seems strange with the git history here. Why does it show the entire file as changed? Manually look through the first lines, it's not. I'm worried that there is some git corruption in this branch.

Code after:
        with open(filepath, 'w') as f:
            f.write('not valid json!')
            f.flush()
            new_modification_time = time.time() + 3600
            os.utime(filepath, (new_modification_time, new_modification_time))

    finish_fact_cache(hosts, fact_cache, last_modified)

    bulk_update.assert_not_called()


def test_finish_job_fact_cache_clear(hosts, mocker, ref_time, tmpdir):
    fact_cache = os.path.join(tmpdir, 'facts')
    last_modified, _ = start_fact_cache(hosts, fact_cache, timeout=0)

    bulk_update = mocker.patch('django.db.models.query.QuerySet.bulk_update')

    os.remove(os.path.join(fact_cache, hosts[1].name))
    finish_fact_cache(hosts, fact_cache, last_modified)

    for host in (hosts[0], hosts[2], hosts[3]):
        assert host.ansible_facts == {"a": 1, "b": 2}
        assert host.ansible_facts_modified == ref_time
    assert hosts[1].ansible_facts == {}
    assert hosts[1].ansible_facts_modified > ref_time
    bulk_update.assert_called_once_with([hosts[1]], ['ansible_facts', 'ansible_facts_modified'])
Repository: ansible/awx
File: awx/main/utils/db.py (Python)

Code before:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.


from awx.settings.application_name import set_application_name
from django.conf import settings
from django.db import transaction


def set_connection_name(function):
    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)


def bulk_update_sorted_by_id(model, objects, fields, batch_size=1000):
    # Filter out objects with None ID
    objects = [obj for obj in objects if obj.id is not None]
    # If there are no valid objects, return early
    if not objects:
        return
    # Sort objects by their ID
    sorted_objects = sorted(objects, key=lambda obj: obj.id)

    # Perform the bulk update within an atomic transaction
    with transaction.atomic():
        model.objects.bulk_update(sorted_objects, fields, batch_size=batch_size)


Review comment (suggestion):
This appears inconsistent with your use of the method `bulk_update_sorted_by_id(host_list, fields=['ansible_facts', 'ansible_facts_modified'])`.

So the next implied question would be "how should this get `Host`?"

You can pass in `Host`, the class, that would be fine if you do that (use the current method arg spec). Just update the method calls if you do.

The other alternative would be that after you define `sorted_objects`, you set `model = type(sorted_objects[0])`. Some might find this hacky, but I'm totally okay with it. I do worse import workarounds myself. All the time.

Code after:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.


from awx.settings.application_name import set_application_name
from django.conf import settings


def set_connection_name(function):
    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)


def bulk_update_sorted_by_id(model, objects, fields, batch_size=1000):
    """
    Perform a sorted bulk update on model instances to avoid database deadlocks.

    This function was introduced to prevent deadlocks observed in the AWX Controller
    when concurrent jobs attempt to update different fields on the same `main_hosts` table.
    Specifically, deadlocks occurred when one process updated `last_job_id` while another
    simultaneously updated `ansible_facts`.

    By sorting updates ID, we ensure a consistent update order,
    which helps avoid the row-level locking contention that can lead to deadlocks
    in PostgreSQL when multiple processes are involved.

    Returns:
        int: The number of rows affected by the update.
    """
    objects = [obj for obj in objects if obj.id is not None]
    if not objects:
        return 0  # Return 0 when nothing is updated

    sorted_objects = sorted(objects, key=lambda obj: obj.id)
    return model.objects.bulk_update(sorted_objects, fields, batch_size=batch_size)
Repository: ansible/awx
File: awx/main/utils/db.py (Python)

Code before:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.


from awx.settings.application_name import set_application_name
from django.conf import settings
from django.db import transaction


def set_connection_name(function):
    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)


def bulk_update_sorted_by_id(model, objects, fields, batch_size=1000):
    # Filter out objects with None ID
    objects = [obj for obj in objects if obj.id is not None]
    # If there are no valid objects, return early
    if not objects:
        return
    # Sort objects by their ID
    sorted_objects = sorted(objects, key=lambda obj: obj.id)

    # Perform the bulk update within an atomic transaction
    with transaction.atomic():
        model.objects.bulk_update(sorted_objects, fields, batch_size=batch_size)


Review comment (suggestion):
The bulk update is a single statement, a transaction should not be necessary at all. The task system should usually, as a best practice, operate in auto-commit mode. If you don't, the read-write flow behaves in a way that's not intuitive and makes cross-process errors more likely due to uncommited changes in another client's transaction.

So I think the best course of action is to just not do this in a transaction.

Code after:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.


from awx.settings.application_name import set_application_name
from django.conf import settings


def set_connection_name(function):
    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)


def bulk_update_sorted_by_id(model, objects, fields, batch_size=1000):
    """
    Perform a sorted bulk update on model instances to avoid database deadlocks.

    This function was introduced to prevent deadlocks observed in the AWX Controller
    when concurrent jobs attempt to update different fields on the same `main_hosts` table.
    Specifically, deadlocks occurred when one process updated `last_job_id` while another
    simultaneously updated `ansible_facts`.

    By sorting updates ID, we ensure a consistent update order,
    which helps avoid the row-level locking contention that can lead to deadlocks
    in PostgreSQL when multiple processes are involved.

    Returns:
        int: The number of rows affected by the update.
    """
    objects = [obj for obj in objects if obj.id is not None]
    if not objects:
        return 0  # Return 0 when nothing is updated

    sorted_objects = sorted(objects, key=lambda obj: obj.id)
    return model.objects.bulk_update(sorted_objects, fields, batch_size=batch_size)
Repository: ansible/awx
File: awx/main/utils/db.py (Python)

Code before:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.


from awx.settings.application_name import set_application_name
from django.conf import settings
from django.db import transaction


def set_connection_name(function):
    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)


def bulk_update_sorted_by_id(model, objects, fields, batch_size=1000):
    # Filter out objects with None ID
    objects = [obj for obj in objects if obj.id is not None]
    # If there are no valid objects, return early
    if not objects:
        return
    # Sort objects by their ID
    sorted_objects = sorted(objects, key=lambda obj: obj.id)

    # Perform the bulk update within an atomic transaction
    with transaction.atomic():
        model.objects.bulk_update(sorted_objects, fields, batch_size=batch_size)


Review comment (bug):
It would be really good to add a docstring basically explaining the bug that you're fixing and how this addresses it. That establishes the reasons for the method's existence.

Code after:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.


from awx.settings.application_name import set_application_name
from django.conf import settings


def set_connection_name(function):
    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)


def bulk_update_sorted_by_id(model, objects, fields, batch_size=1000):
    """
    Perform a sorted bulk update on model instances to avoid database deadlocks.

    This function was introduced to prevent deadlocks observed in the AWX Controller
    when concurrent jobs attempt to update different fields on the same `main_hosts` table.
    Specifically, deadlocks occurred when one process updated `last_job_id` while another
    simultaneously updated `ansible_facts`.

    By sorting updates ID, we ensure a consistent update order,
    which helps avoid the row-level locking contention that can lead to deadlocks
    in PostgreSQL when multiple processes are involved.

    Returns:
        int: The number of rows affected by the update.
    """
    objects = [obj for obj in objects if obj.id is not None]
    if not objects:
        return 0  # Return 0 when nothing is updated

    sorted_objects = sorted(objects, key=lambda obj: obj.id)
    return model.objects.bulk_update(sorted_objects, fields, batch_size=batch_size)
Repository: ansible/awx
File: awx/main/utils/db.py (Python)

Code before:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.


from awx.settings.application_name import set_application_name
from django.conf import settings
from django.db import transaction


def set_connection_name(function):
    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)


def bulk_update_sorted_by_id(model, objects, fields, batch_size=1000):
    # Filter out objects with None ID
    objects = [obj for obj in objects if obj.id is not None]
    # If there are no valid objects, return early
    if not objects:
        return
    # Sort objects by their ID
    sorted_objects = sorted(objects, key=lambda obj: obj.id)

    # Perform the bulk update within an atomic transaction
    with transaction.atomic():
        model.objects.bulk_update(sorted_objects, fields, batch_size=batch_size)


Review comment (suggestion):
https://docs.djangoproject.com/en/5.2/ref/models/querysets/#bulk-update

bulk_update has a return value. I don't care what the value is, but you should `return` it.

Code after:
# Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.


from awx.settings.application_name import set_application_name
from django.conf import settings


def set_connection_name(function):
    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)


def bulk_update_sorted_by_id(model, objects, fields, batch_size=1000):
    """
    Perform a sorted bulk update on model instances to avoid database deadlocks.

    This function was introduced to prevent deadlocks observed in the AWX Controller
    when concurrent jobs attempt to update different fields on the same `main_hosts` table.
    Specifically, deadlocks occurred when one process updated `last_job_id` while another
    simultaneously updated `ansible_facts`.

    By sorting updates ID, we ensure a consistent update order,
    which helps avoid the row-level locking contention that can lead to deadlocks
    in PostgreSQL when multiple processes are involved.

    Returns:
        int: The number of rows affected by the update.
    """
    objects = [obj for obj in objects if obj.id is not None]
    if not objects:
        return 0  # Return 0 when nothing is updated

    sorted_objects = sorted(objects, key=lambda obj: obj.id)
    return model.objects.bulk_update(sorted_objects, fields, batch_size=batch_size)
Repository: ansible/awx
File: awx/main/migrations/0156_capture_mesh_topology.py (Python)

Code before:
            field=models.CharField(
                choices=[
                    ('control', 'Control plane node'),
                    ('execution', 'Execution plane node'),
                    ('hybrid', 'Controller and execution'),
                    ('hop', 'Message-passing node, no execution capability'),
                ],
                default='hybrid',
                max_length=16,
            ),
        ),
        migrations.CreateModel(
            name='InstanceLink',
            fields=[
                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
                ('source', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='+', to='main.Instance')),
                ('target', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='reverse_peers', to='main.Instance')),
            ],
            options={
                'unique_together': {('source', 'target')},
            },
        ),
        migrations.AddField(
            model_name='instance',
            name='peers',
            field=models.ManyToManyField(through='main.InstanceLink', through_fields=('source', 'target'), to='main.Instance'),
        ),
    ]


Review comment (suggestion):
Both 0156 and 0187 account for the accounting different - 3 fields changed previously versus 5 fields changed here. This is because both of these are _prior_ versions of the `Instance.peers` relationship. The changes to these could be reverted. But I don't think it would matter either way. I might revert because that would be more clear to me.

Code after:
            field=models.CharField(
                choices=[
                    ('control', 'Control plane node'),
                    ('execution', 'Execution plane node'),
                    ('hybrid', 'Controller and execution'),
                    ('hop', 'Message-passing node, no execution capability'),
                ],
                default='hybrid',
                max_length=16,
            ),
        ),
        migrations.CreateModel(
            name='InstanceLink',
            fields=[
                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
                ('source', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='+', to='main.Instance')),
                ('target', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='reverse_peers', to='main.Instance')),
            ],
            options={
                'unique_together': {('source', 'target')},
            },
        ),
        migrations.AddField(
            model_name='instance',
            name='peers',
            field=models.ManyToManyField(through='main.InstanceLink', to='main.Instance'),
        ),
    ]
Repository: apache/airflow
File: providers/fab/src/airflow/providers/fab/www/extensions/init_wsgi_middlewares.py (Python)

Code before:
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from __future__ import annotations

from typing import TYPE_CHECKING

from airflow.configuration import conf
from airflow.exceptions import AirflowConfigException
from werkzeug.middleware.proxy_fix import ProxyFix

if TYPE_CHECKING:
    from flask import Flask


def init_wsgi_middleware(flask_app: Flask) -> None:
    """Handle X-Forwarded-* headers and base_url support."""

    # Apply ProxyFix middleware
    if conf.getboolean("webserver", "ENABLE_PROXY_FIX"):
        flask_app.wsgi_app = ProxyFix(  # type: ignore
            flask_app.wsgi_app,
            x_for=conf.getint("webserver", "PROXY_FIX_X_FOR", fallback=1),
            x_proto=conf.getint("webserver", "PROXY_FIX_X_PROTO", fallback=1),
            x_host=conf.getint("webserver", "PROXY_FIX_X_HOST", fallback=1),
            x_port=conf.getint("webserver", "PROXY_FIX_X_PORT", fallback=1),
            x_prefix=conf.getint("webserver", "PROXY_FIX_X_PREFIX", fallback=1),
        )


Review comment (suggestion):
We removed this config because it was no longer used. Please update `airflow-core/src/airflow/cli/commands/config_command.py` as well to no longer mark this config as removed.

Code after:
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from __future__ import annotations

from typing import TYPE_CHECKING

from werkzeug.middleware.proxy_fix import ProxyFix

from airflow.configuration import conf

if TYPE_CHECKING:
    from flask import Flask


def init_wsgi_middleware(flask_app: Flask) -> None:
    """Handle X-Forwarded-* headers and base_url support."""
    # Apply ProxyFix middleware
    if conf.getboolean("fab", "ENABLE_PROXY_FIX"):
        flask_app.wsgi_app = ProxyFix(  # type: ignore
            flask_app.wsgi_app,
            x_for=conf.getint("fab", "PROXY_FIX_X_FOR", fallback=1),
            x_proto=conf.getint("fab", "PROXY_FIX_X_PROTO", fallback=1),
            x_host=conf.getint("fab", "PROXY_FIX_X_HOST", fallback=1),
            x_port=conf.getint("fab", "PROXY_FIX_X_PORT", fallback=1),
            x_prefix=conf.getint("fab", "PROXY_FIX_X_PREFIX", fallback=1),
        )