Add celery_app parameter to job_queue_size for priority queue support

Fixes RabbitMQ PRECONDITION_FAILED error when querying queues configured
with custom arguments like x-max-priority. The celery_app parameter allows
extracting queue arguments from the app's task_queues configuration and
passing them to queue_declare().

The celery_app and broker_url parameters are mutually exclusive - a
ValueError is raised if both are provided.

Author: mrrooijen

hirefire_resource/macro/celery.py

@@ -25,6 +25,26 @@ class ChannelError(Exception):
 from hirefire_resource.errors import MissingQueueError
 
 
+def _get_queue_arguments_from_app(app, queues):
+    """
+    Extract queue arguments from Celery app configuration for specified queues.
+
+    Args:
+        app: Celery app instance with task_queues configuration
+        queues: List of queue names to extract arguments for
+
+    Returns:
+        dict: Mapping of queue name to queue_arguments dict
+    """
+    queue_args = {}
+    task_queues = getattr(app.conf, "task_queues", None) or []
+    for q in task_queues:
+        queue_name = getattr(q, "name", None)
+        if queue_name and queue_name in queues:
+            queue_args[queue_name] = getattr(q, "queue_arguments", None)
+    return queue_args
+
+
 def mitigate_connection_reset_error(retries=10, delay=1):
     """
     Decorator to retry a function when ConnectionResetError occurs.
@@ -210,7 +230,7 @@ async def async_job_queue_latency(*queues, broker_url=None):
 
 
 @mitigate_connection_reset_error()
-def job_queue_size(*queues, broker_url=None):
+def job_queue_size(*queues, broker_url=None, celery_app=None):
     """
     Calculates the total job queue size across the specified queues using Celery with either Redis
     or RabbitMQ (AMQP) as the broker.
@@ -229,22 +249,31 @@ def job_queue_size(*queues, broker_url=None):
           using a workaround, such as a separate queue for scheduled tasks that forwards tasks ready
           to run to the relevant regular queues. When using RabbitMQ (AMQP), consider using the
           Delayed Message Plugin.
+        - For RabbitMQ queues with custom arguments (e.g., x-max-priority for priority queues),
+          pass your configured Celery app via the `celery_app` parameter. This allows the function
+          to extract and use the correct queue arguments when querying RabbitMQ.
 
     Args:
         *queues (str): Names of the queues for size measurement.
-        broker_url (str, optional): The broker URL. Defaults in the following order:
+        broker_url (str, optional): The broker URL. Cannot be used together with `celery_app`.
+            Defaults in the following order:
             - Passed argument `broker_url`.
             - Environment variables `AMQP_URL`, `RABBITMQ_URL`, `RABBITMQ_BIGWIG_URL`,
               `CLOUDAMQP_URL`, `REDIS_TLS_URL`, `REDIS_URL`, `REDISTOGO_URL`, `REDISCLOUD_URL`,
               `OPENREDIS_URL`.
             - "amqp://guest:guest@localhost:5672" if AMQP is available, otherwise
               "redis://localhost:6379/0".
+        celery_app (Celery, optional): A configured Celery app instance. Cannot be used together
+            with `broker_url`. When provided, the function uses this app's connection and extracts
+            queue arguments from celery_app.conf.task_queues. This is required for RabbitMQ queues
+            with custom arguments like x-max-priority.
 
     Returns:
         int: The cumulative job queue size across the specified queues.
 
     Raises:
         MissingQueueError: If no queue names are provided.
+        ValueError: If both `broker_url` and `celery_app` are provided.
 
     Examples:
         >>> job_queue_size("celery")
@@ -255,43 +284,57 @@ def job_queue_size(*queues, broker_url=None):
         42
         >>> job_queue_size("celery", broker_url="redis://localhost:6379/0")
         42
+        >>> # For priority queues, pass your configured Celery app:
+        >>> job_queue_size("celery", celery_app=celery_app)
+        42
     """
     if not queues:
         raise MissingQueueError()
 
-    broker_url = (
-        broker_url
-        or os.environ.get("AMQP_URL")
-        or os.environ.get("RABBITMQ_URL")
-        or os.environ.get("RABBITMQ_BIGWIG_URL")
-        or os.environ.get("CLOUDAMQP_URL")
-        or os.environ.get("REDIS_TLS_URL")
-        or os.environ.get("REDIS_URL")
-        or os.environ.get("REDISTOGO_URL")
-        or os.environ.get("REDISCLOUD_URL")
-        or os.environ.get("OPENREDIS_URL")
-    )
-
-    if not broker_url:
-        if AMQP_AVAILABLE:
-            broker_url = "amqp://guest:guest@localhost:5672"
-        else:
-            broker_url = "redis://localhost:6379/0"
-
-    app = Celery(broker=broker_url)
+    if celery_app is not None and broker_url is not None:
+        raise ValueError(
+            "Cannot specify both 'celery_app' and 'broker_url'. "
+            "Use 'celery_app' to pass your configured Celery app (recommended for priority queues), "
+            "or 'broker_url' for simple setups."
+        )
+
+    if celery_app is None:
+        broker_url = (
+            broker_url
+            or os.environ.get("AMQP_URL")
+            or os.environ.get("RABBITMQ_URL")
+            or os.environ.get("RABBITMQ_BIGWIG_URL")
+            or os.environ.get("CLOUDAMQP_URL")
+            or os.environ.get("REDIS_TLS_URL")
+            or os.environ.get("REDIS_URL")
+            or os.environ.get("REDISTOGO_URL")
+            or os.environ.get("REDISCLOUD_URL")
+            or os.environ.get("OPENREDIS_URL")
+        )
+
+        if not broker_url:
+            if AMQP_AVAILABLE:
+                broker_url = "amqp://guest:guest@localhost:5672"
+            else:
+                broker_url = "redis://localhost:6379/0"
+
+        celery_app = Celery(broker=broker_url)
+
+    # Extract queue arguments from app configuration (for priority queues, etc.)
+    queue_args = _get_queue_arguments_from_app(celery_app, queues)
 
     try:
-        with app.connection_or_acquire() as connection:
+        with celery_app.connection_or_acquire() as connection:
             with connection.channel() as channel:
-                worker_task_count = _job_queue_size_worker(app, queues)
-                broker_task_count = _job_queue_size_broker(channel, queues)
+                worker_task_count = _job_queue_size_worker(celery_app, queues)
+                broker_task_count = _job_queue_size_broker(channel, queues, queue_args)
                 return worker_task_count + broker_task_count
 
     except OperationalError:
         return 0
 
 
-async def async_job_queue_size(*queues, broker_url=None):
+async def async_job_queue_size(*queues, broker_url=None, celery_app=None):
     """
     Asynchronously calculates the total job queue size across the specified queues using Celery with
     either Redis or RabbitMQ (AMQP) as the broker.
@@ -311,22 +354,29 @@ async def async_job_queue_size(*queues, broker_url=None):
           using a workaround, such as a separate queue for scheduled tasks that forwards tasks ready
           to run to the relevant regular queues. When using RabbitMQ (AMQP), consider using the
           Delayed Message Plugin.
+        - For RabbitMQ queues with custom arguments (e.g., x-max-priority for priority queues),
+          pass your configured Celery app via the `celery_app` parameter.
 
     Args:
         *queues (str): Names of the queues for size measurement.
-        broker_url (str, optional): The broker URL. Defaults in the following order:
+        broker_url (str, optional): The broker URL. Cannot be used together with `celery_app`.
+            Defaults in the following order:
             - Passed argument `broker_url`.
             - Environment variables `AMQP_URL`, `RABBITMQ_URL`, `RABBITMQ_BIGWIG_URL`,
               `CLOUDAMQP_URL`, `REDIS_TLS_URL`, `REDIS_URL`, `REDISTOGO_URL`, `REDISCLOUD_URL`,
               `OPENREDIS_URL`.
             - "amqp://guest:guest@localhost:5672" if AMQP is available, otherwise
               "redis://localhost:6379/0".
+        celery_app (Celery, optional): A configured Celery app instance. Cannot be used together
+            with `broker_url`. When provided, the function uses this app's connection and extracts
+            queue arguments from celery_app.conf.task_queues.
 
     Returns:
         int: The cumulative job queue size across the specified queues.
 
     Raises:
         MissingQueueError: If no queue names are provided.
+        ValueError: If both `broker_url` and `celery_app` are provided.
 
     Examples:
         >>> await async_job_queue_size("celery")
@@ -335,11 +385,13 @@ async def async_job_queue_size(*queues, broker_url=None):
         85
         >>> await async_job_queue_size("celery", broker_url="amqp://user:password@host:5672")
         42
-        >>> await async_job_queue_size("celery", broker_url="redis://localhost:6379/0")
+        >>> await async_job_queue_size("celery", celery_app=celery_app)
         42
     """
     loop = asyncio.get_event_loop()
-    func = functools.partial(job_queue_size, *queues, broker_url=broker_url)
+    func = functools.partial(
+        job_queue_size, *queues, broker_url=broker_url, celery_app=celery_app
+    )
     return await loop.run_in_executor(None, func)
 
 
@@ -399,22 +451,28 @@ def _job_queue_size_worker(app, queues):
     return sum(worker_data.get(queue, 0) for queue in queues)
 
 
-def _job_queue_size_broker(channel, queues):
+def _job_queue_size_broker(channel, queues, queue_args=None):
+    if queue_args is None:
+        queue_args = {}
+
     if hasattr(channel, "_size"):
-        fn = _job_queue_size_redis
+        return sum(_job_queue_size_redis(channel, queue) for queue in queues)
     else:
-        fn = _job_queue_size_rabbitmq
-
-    return sum(fn(channel, queue) for queue in queues)
+        return sum(
+            _job_queue_size_rabbitmq(channel, queue, queue_args.get(queue))
+            for queue in queues
+        )
 
 
 def _job_queue_size_redis(channel, queue):
     return channel.client.llen(queue)
 
 
-def _job_queue_size_rabbitmq(channel, queue):
+def _job_queue_size_rabbitmq(channel, queue, arguments=None):
     try:
-        return channel.queue_declare(queue=queue, passive=True).message_count
+        return channel.queue_declare(
+            queue=queue, passive=True, arguments=arguments
+        ).message_count
     except ChannelError:
         return 0
 

tests/hirefire_resource/macro/test_celery.py

@@ -1,9 +1,9 @@
 import math
 from datetime import datetime, timedelta, timezone
-from unittest.mock import patch
 
 import pytest
 from celery import Celery
+from kombu import Queue
 
 from hirefire_resource.errors import MissingQueueError
 from hirefire_resource.macro.celery import (
@@ -215,3 +215,133 @@ async def test_job_queue_size_with_jobs_async(celery_app):
         )
         == 10
     )
+
+
+# Tests for priority queues (RabbitMQ only)
+
+
+@pytest.fixture
+def priority_celery_app():
+    """Create a Celery app with priority queue configuration."""
+    broker_url = "amqp://guest:guest@localhost:5672"
+    app = Celery(broker=broker_url)
+
+    # Configure queues with x-max-priority
+    queue_arguments = {"x-max-priority": 10}
+    app.conf.task_queues = [
+        Queue("priority_queue", queue_arguments=queue_arguments),
+    ]
+    app.conf.task_queue_max_priority = 10
+    app.conf.task_default_priority = 5
+
+    return app
+
+
+@pytest.fixture
+def setup_priority_queue(priority_celery_app):
+    """Create the priority queue in RabbitMQ with x-max-priority argument."""
+    with priority_celery_app.connection_or_acquire() as connection:
+        channel = connection.default_channel
+        # Delete queue if it exists (to start fresh)
+        channel.queue_delete(queue="priority_queue")
+        # Create queue WITH x-max-priority argument
+        channel.queue_declare(
+            queue="priority_queue",
+            durable=True,
+            auto_delete=False,
+            arguments={"x-max-priority": 10},
+        )
+
+    yield priority_celery_app
+
+    # Cleanup: delete the queue after test
+    with priority_celery_app.connection_or_acquire() as connection:
+        channel = connection.default_channel
+        channel.queue_delete(queue="priority_queue")
+
+
+def test_job_queue_size_priority_queue_with_broker_url(setup_priority_queue):
+    """
+    Test job_queue_size with broker_url on a priority queue.
+
+    Note: In our test environment (RabbitMQ 4.2.2, py-amqp 5.3.1), this works fine
+    because passive=True declarations don't validate arguments. However, some
+    RabbitMQ versions or configurations may return PRECONDITION_FAILED.
+
+    For guaranteed compatibility with priority queues, use celery_app parameter instead.
+    """
+    priority_celery_app = setup_priority_queue
+    broker_url = priority_celery_app.conf.broker_url
+
+    # Add tasks to the queue
+    priority_celery_app.send_task("test_task", queue="priority_queue")
+    priority_celery_app.send_task("test_task", queue="priority_queue")
+
+    # Using broker_url (no queue arguments passed)
+    result = job_queue_size("priority_queue", broker_url=broker_url)
+
+    # In our test environment this works, but may fail in other environments
+    assert result == 2
+
+
+def test_job_queue_size_priority_queue_with_celery_app_returns_correct_count(
+    setup_priority_queue,
+):
+    """
+    Test that job_queue_size returns the correct count when passed the Celery app
+    that has the queue configuration with x-max-priority.
+
+    This is the recommended approach for priority queues - by passing the celery_app,
+    we extract the queue arguments and pass them to queue_declare.
+    """
+    priority_celery_app = setup_priority_queue
+
+    # Add tasks to the queue
+    priority_celery_app.send_task("test_task", queue="priority_queue")
+    priority_celery_app.send_task("test_task", queue="priority_queue")
+    priority_celery_app.send_task("test_task", queue="priority_queue")
+
+    # Recommended: Pass the celery_app parameter so queue arguments are extracted
+    result = job_queue_size("priority_queue", celery_app=priority_celery_app)
+
+    # This should return the correct count
+    assert result == 3
+
+
+def test_job_queue_size_raises_error_when_both_broker_url_and_celery_app_provided(
+    setup_priority_queue,
+):
+    """
+    Test that job_queue_size raises ValueError when both broker_url and celery_app
+    are provided, since they are mutually exclusive.
+    """
+    priority_celery_app = setup_priority_queue
+
+    with pytest.raises(ValueError) as exc_info:
+        job_queue_size(
+            "priority_queue",
+            broker_url="amqp://guest:guest@localhost:5672",
+            celery_app=priority_celery_app,
+        )
+
+    assert "Cannot specify both" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_async_job_queue_size_raises_error_when_both_broker_url_and_celery_app_provided(
+    setup_priority_queue,
+):
+    """
+    Test that async_job_queue_size raises ValueError when both broker_url and celery_app
+    are provided, since they are mutually exclusive.
+    """
+    priority_celery_app = setup_priority_queue
+
+    with pytest.raises(ValueError) as exc_info:
+        await async_job_queue_size(
+            "priority_queue",
+            broker_url="amqp://guest:guest@localhost:5672",
+            celery_app=priority_celery_app,
+        )
+
+    assert "Cannot specify both" in str(exc_info.value)