Merge branch 'master' into feat/add-otel-instrumentation

Author: soapun

poetry.lock

@@ -34,7 +34,7 @@ pydantic = ">=1.0,<=3.0"
 importlib-metadata = "*"
 pycron = "^3.0.0"
 taskiq_dependencies = ">=1.3.1,<2"
-anyio = ">=3"
+anyio = ">=4"
 packaging = ">=19"
 # For opentelemetry instrumentation
 opentelemetry-api = { version = "^1.38.0", optional = true }
@@ -113,6 +113,7 @@ multi_line_output = 3
 
 [tool.pytest.ini_options]
 log_level = 'INFO'
+anyio_mode = "auto"
 
 [tool.coverage.run]
 omit = [
@@ -131,6 +132,7 @@ requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.ruff]
+target-version="py39"
 # List of enabled rulsets.
 # See https://docs.astral.sh/ruff/rules/ for more information.
 lint.select = [
@@ -159,6 +161,7 @@ lint.select = [
     "ERA", # Checks for commented out code
     "PL",  # PyLint checks
     "RUF", # Specific to Ruff checks
+    "FA102",  # Future annotations
 ]
 lint.ignore = [
     "D105",    # Missing docstring in magic method
@@ -183,7 +186,7 @@ line-length = 88
     "SLF001", # Private member accessed
     "S311",   # Standard pseudo-random generators are not suitable for security/cryptographic purposes
     "D101",   # Missing docstring in public class
-    "D102",   # Missing docstring in public method
+    "D102",  # Missing docstring in public method
 ]
 
 [tool.ruff.lint.pydocstyle]

pyproject.toml

@@ -3,9 +3,11 @@
 from .prometheus_middleware import PrometheusMiddleware
 from .simple_retry_middleware import SimpleRetryMiddleware
 from .smart_retry_middleware import SmartRetryMiddleware
+from .taskiq_admin_middleware import TaskiqAdminMiddleware
 
 __all__ = (
     "PrometheusMiddleware",
     "SimpleRetryMiddleware",
     "SmartRetryMiddleware",
+    "TaskiqAdminMiddleware",
 )

taskiq/middlewares/__init__.py

@@ -44,7 +44,7 @@ def __init__(
         logger.debug("Initializing metrics")
 
         try:
-            from prometheus_client import Counter, Histogram
+            from prometheus_client import Counter, Histogram  # noqa: PLC0415
         except ImportError as exc:
             raise ImportError(
                 "Cannot initialize metrics. Please install 'taskiq[metrics]'.",
@@ -85,7 +85,7 @@ def startup(self) -> None:
         This function starts prometheus server.
         It starts it only in case if it's a worker process.
         """
-        from prometheus_client import start_http_server
+        from prometheus_client import start_http_server  # noqa: PLC0415
 
         if self.broker.is_worker_process:
             try:

taskiq/middlewares/prometheus_middleware.py

@@ -1,12 +1,14 @@
 import asyncio
-from datetime import UTC, datetime
+from datetime import datetime, timezone
 from logging import getLogger
-from typing import Any
+from typing import Any, Optional
 from urllib.parse import urljoin
 
 import aiohttp
 
-from taskiq import TaskiqMessage, TaskiqMiddleware, TaskiqResult
+from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.message import TaskiqMessage
+from taskiq.result import TaskiqResult
 
 __all__ = ("TaskiqAdminMiddleware",)
 
@@ -36,7 +38,7 @@ def __init__(
         url: str,
         api_token: str,
         timeout: int = 5,
-        taskiq_broker_name: str | None = None,
+        taskiq_broker_name: Optional[str] = None,
     ) -> None:
         super().__init__()
         self.url = url
@@ -48,7 +50,7 @@ def __init__(
 
     @staticmethod
     def _now_iso() -> str:
-        return datetime.now(UTC).replace(tzinfo=None).isoformat()
+        return datetime.now(timezone.utc).replace(tzinfo=None).isoformat()
 
     def _get_client(self) -> aiohttp.ClientSession:
         """Create and cache session."""
@@ -118,6 +120,7 @@ async def post_send(self, message: TaskiqMessage) -> None:
             {
                 "args": message.args,
                 "kwargs": message.kwargs,
+                "labels": message.labels,
                 "queuedAt": self._now_iso(),
                 "taskName": message.task_name,
                 "worker": self.__ta_broker_name,
@@ -139,6 +142,7 @@ async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
             {
                 "args": message.args,
                 "kwargs": message.kwargs,
+                "labels": message.labels,
                 "startedAt": self._now_iso(),
                 "taskName": message.task_name,
                 "worker": self.__ta_broker_name,

taskiq/middlewares/taskiq_admin_middleware.py

@@ -7,7 +7,6 @@
 from tests.utils import AsyncQueueBroker
 
 
-@pytest.mark.anyio
 async def test_successful() -> None:
     broker = AsyncQueueBroker()
     kicked = 0
@@ -28,7 +27,6 @@ def test_func() -> None:
     assert kicked == desired_kicked
 
 
-@pytest.mark.anyio
 async def test_cancelation() -> None:
     broker = AsyncQueueBroker()
     kicked = 0

tests/api/test_receiver_task.py

@@ -2,15 +2,12 @@
 import contextlib
 from datetime import datetime, timedelta, timezone
 
-import pytest
-
 from taskiq import TaskiqScheduler
 from taskiq.api import run_scheduler_task
 from taskiq.schedule_sources import LabelScheduleSource
 from tests.utils import AsyncQueueBroker
 
 
-@pytest.mark.anyio
 async def test_successful() -> None:
     broker = AsyncQueueBroker()
     scheduler = TaskiqScheduler(broker, sources=[LabelScheduleSource(broker)])
@@ -26,7 +23,6 @@ def _() -> None:
     scheduler_task.cancel()
 
 
-@pytest.mark.anyio
 async def test_cancelation() -> None:
     broker = AsyncQueueBroker()
     scheduler = TaskiqScheduler(broker, sources=[LabelScheduleSource(broker)])

tests/api/test_scheduler.py

@@ -8,7 +8,6 @@
 from taskiq.state import TaskiqState
 
 
-@pytest.mark.anyio
 async def test_inmemory_success() -> None:
     broker = InMemoryBroker()
     test_val = uuid.uuid4().hex
@@ -23,7 +22,6 @@ async def task() -> str:
     assert not broker._running_tasks
 
 
-@pytest.mark.anyio
 async def test_cannot_listen() -> None:
     broker = InMemoryBroker()
 
@@ -32,7 +30,6 @@ async def test_cannot_listen() -> None:
             pass
 
 
-@pytest.mark.anyio
 async def test_startup() -> None:
     broker = InMemoryBroker()
     test_value = uuid.uuid4().hex
@@ -51,7 +48,6 @@ async def _c_startup(state: TaskiqState) -> None:
     assert broker.state.from_client == test_value
 
 
-@pytest.mark.anyio
 async def test_shutdown() -> None:
     broker = InMemoryBroker()
     test_value = uuid.uuid4().hex
@@ -70,7 +66,6 @@ async def _c_startup(state: TaskiqState) -> None:
     assert broker.state.from_client == test_value
 
 
-@pytest.mark.anyio
 async def test_execution() -> None:
     broker = InMemoryBroker()
     test_value = uuid.uuid4().hex
@@ -87,7 +82,6 @@ async def test_task() -> str:
     assert result.return_value == test_value
 
 
-@pytest.mark.anyio
 async def test_inline_awaits() -> None:
     broker = InMemoryBroker(await_inplace=True)
     slept = False
@@ -104,7 +98,6 @@ async def test_task() -> None:
     assert not broker._running_tasks
 
 
-@pytest.mark.anyio
 async def test_wait_all() -> None:
     broker = InMemoryBroker()
     slept = False

tests/brokers/test_inmemory.py

@@ -1,8 +1,6 @@
 from datetime import datetime
 from typing import List, Union
 
-import pytest
-
 from taskiq import InMemoryBroker, ScheduleSource
 from taskiq.cli.scheduler.run import get_all_schedules
 from taskiq.scheduler.scheduled_task import ScheduledTask
@@ -20,7 +18,6 @@ async def get_schedules(self) -> List[ScheduledTask]:
         return self.schedules
 
 
-@pytest.mark.anyio
 async def test_get_schedules_success() -> None:
     """Tests that schedules are returned correctly."""
     schedules1 = [
@@ -62,7 +59,6 @@ async def test_get_schedules_success() -> None:
     ]
 
 
-@pytest.mark.anyio
 async def test_get_schedules_error() -> None:
     """Tests that if source returned an error, empty list will be returned."""
     source1 = DummySource(

tests/cli/scheduler/test_updater.py

@@ -56,7 +56,6 @@ def get_message(
     )
 
 
-@pytest.mark.anyio
 @pytest.mark.parametrize(
     "state,meta",
     [
@@ -82,7 +81,6 @@ async def test_func(tes_val: ProgressTracker[Any] = TaskiqDepends()) -> None:
     assert progress.state == state
 
 
-@pytest.mark.anyio
 async def test_progress_tracker_ctx_none() -> None:
     broker = InMemoryBroker()
 
@@ -98,7 +96,6 @@ async def test_func() -> None:
     assert progress is None
 
 
-@pytest.mark.anyio
 @pytest.mark.parametrize(
     "state,meta",
     [