Merge branch 'master' into go-to-uv

Author: s3rius

taskiq/instrumentation.py

@@ -115,7 +115,7 @@ def uninstrument_broker(self, broker: AsyncBroker) -> None:
 
     def instrumentation_dependencies(self) -> Collection[str]:
         """This function tells which library this instrumentor instruments."""
-        return ("taskiq >= 0.0.1",)
+        return ("taskiq >= 0.0.0",)
 
     @classmethod
     def _start_listen_with_initialize(cls, args: WorkerArgs) -> None:

taskiq/middlewares/opentelemetry_middleware.py

@@ -1,6 +1,9 @@
 import logging
 from contextlib import AbstractContextManager
-from typing import Any, TypeVar
+from importlib.metadata import version
+from typing import Any, Dict, Optional, Tuple, TypeVar
+
+from packaging.version import Version, parse
 
 try:
     import opentelemetry  # noqa: F401
@@ -10,6 +13,7 @@
         "Please install 'taskiq[opentelemetry]'.",
     ) from exc
 
+
 from opentelemetry import context as context_api
 from opentelemetry import trace
 from opentelemetry.metrics import Meter, MeterProvider, get_meter
@@ -27,6 +31,16 @@
 # Taskiq Context key
 CTX_KEY = "__otel_task_span"
 
+# unlike pydantic v2, v1 includes CTX_KEY by default
+# excluding it here
+PYDANTIC_VER = parse(version("pydantic"))
+IS_PYDANTIC1 = Version("2.0") > PYDANTIC_VER
+if IS_PYDANTIC1:
+    if TaskiqMessage.__exclude_fields__:  # type: ignore[attr-defined]
+        TaskiqMessage.__exclude_fields__.update(CTX_KEY)  # type: ignore
+    else:
+        TaskiqMessage.__exclude_fields__ = {CTX_KEY}  # type: ignore
+
 # Taskiq Context attributes
 TASKIQ_CONTEXT_ATTRIBUTES = [
     "_retries",
@@ -95,7 +109,9 @@ def attach_context(
 
     if ctx_dict is None:
         ctx_dict = {}
-        setattr(message, CTX_KEY, ctx_dict)
+        # use object.__setattr__ directly
+        # to skip pydantic v1 setattr
+        object.__setattr__(message, CTX_KEY, ctx_dict)
 
     ctx_dict[(message.task_id, is_publish)] = (span, activation, token)
 

tests/api/test_scheduler.py

@@ -16,7 +16,7 @@ async def test_successful() -> None:
     @broker.task(schedule=[{"time": datetime.now(timezone.utc) - timedelta(seconds=1)}])
     def _() -> None: ...
 
-    msg = await asyncio.wait_for(broker.queue.get(), 1)
+    msg = await asyncio.wait_for(broker.queue.get(), 2)
     assert msg
 
     scheduler_task.cancel()

tests/opentelemetry/test_auto_instrumentation.py

@@ -1,3 +1,5 @@
+import asyncio
+
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.trace import SpanKind, StatusCode
 
@@ -6,17 +8,20 @@
 
 
 class TestTaskiqAutoInstrumentation(TestBase):
-    async def test_auto_instrument(self) -> None:
+    def test_auto_instrument(self) -> None:
         TaskiqInstrumentor().instrument()
 
-        broker = InMemoryBroker(await_inplace=True)
+        async def test() -> None:
+            broker = InMemoryBroker(await_inplace=True)
+
+            @broker.task
+            async def task_add(a: float, b: float) -> float:
+                return a + b
 
-        @broker.task
-        async def task_add(a: float, b: float) -> float:
-            return a + b
+            await task_add.kiq(1, 2)
+            await broker.wait_all()
 
-        await task_add.kiq(1, 2)
-        await broker.wait_all()
+        asyncio.run(test())
 
         spans = self.sorted_spans(self.memory_exporter.get_finished_spans())
         self.assertEqual(len(spans), 2)
@@ -25,15 +30,15 @@ async def task_add(a: float, b: float) -> float:
 
         self.assertEqual(
             consumer.name,
-            "execute/tests.test_auto_instrumentation:task_add",
+            "execute/tests.opentelemetry.test_auto_instrumentation:task_add",
             f"{consumer._end_time}:{producer._end_time}",
         )
         self.assertEqual(consumer.kind, SpanKind.CONSUMER)
         self.assertSpanHasAttributes(
             consumer,
             {
                 "taskiq.action": "execute",
-                "taskiq.task_name": "tests.test_auto_instrumentation:task_add",
+                "taskiq.task_name": "tests.opentelemetry.test_auto_instrumentation:task_add",  # noqa: E501
             },
         )
 
@@ -43,14 +48,14 @@ async def task_add(a: float, b: float) -> float:
 
         self.assertEqual(
             producer.name,
-            "send/tests.test_auto_instrumentation:task_add",
+            "send/tests.opentelemetry.test_auto_instrumentation:task_add",
         )
         self.assertEqual(producer.kind, SpanKind.PRODUCER)
         self.assertSpanHasAttributes(
             producer,
             {
                 "taskiq.action": "send",
-                "taskiq.task_name": "tests.test_auto_instrumentation:task_add",
+                "taskiq.task_name": "tests.opentelemetry.test_auto_instrumentation:task_add",  # noqa: E501
             },
         )
 

tests/opentelemetry/test_tasks.py

@@ -10,6 +10,7 @@
 from opentelemetry.trace import Span, SpanKind, StatusCode
 from wrapt import wrap_function_wrapper
 
+from taskiq import TaskiqResult
 from taskiq.instrumentation import TaskiqInstrumentor
 from taskiq.middlewares import opentelemetry_middleware
 
@@ -26,11 +27,14 @@ def tearDown(self) -> None:
         super().tearDown()
         TaskiqInstrumentor().uninstrument_broker(broker)
 
-    async def test_task(self) -> None:
+    def test_task(self) -> None:
         TaskiqInstrumentor().instrument_broker(broker)
 
-        await task_add.kiq(1, 2)
-        await broker.wait_all()
+        async def test() -> None:
+            await task_add.kiq(1, 2)
+            await broker.wait_all()
+
+        asyncio.run(test())
 
         spans = self.sorted_spans(self.memory_exporter.get_finished_spans())
         self.assertEqual(len(spans), 2)
@@ -72,11 +76,14 @@ async def test_task(self) -> None:
         self.assertEqual(consumer.parent.span_id, producer.context.span_id)
         self.assertEqual(consumer.context.trace_id, producer.context.trace_id)
 
-    async def test_task_raises(self) -> None:
+    def test_task_raises(self) -> None:
         TaskiqInstrumentor().instrument_broker(broker)
 
-        await task_raises.kiq()
-        await broker.wait_all()
+        async def test() -> None:
+            await task_raises.kiq()
+            await broker.wait_all()
+
+        asyncio.run(test())
 
         spans = self.sorted_spans(self.memory_exporter.get_finished_spans())
         self.assertEqual(len(spans), 2)
@@ -130,7 +137,7 @@ async def test_task_raises(self) -> None:
         self.assertEqual(consumer.parent.span_id, producer.context.span_id)
         self.assertEqual(consumer.context.trace_id, producer.context.trace_id)
 
-    async def test_uninstrument(self) -> None:
+    def test_uninstrument(self) -> None:
         TaskiqInstrumentor().instrument_broker(broker)
         TaskiqInstrumentor().uninstrument_broker(broker)
 
@@ -143,18 +150,21 @@ async def test() -> None:
         spans = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans), 0)
 
-    async def test_baggage(self) -> None:
+    def test_baggage(self) -> None:
         TaskiqInstrumentor().instrument_broker(broker)
 
+        async def test() -> TaskiqResult[Any]:
+            task = await task_returns_baggage.kiq()
+            return await task.wait_result(timeout=2)
+
         ctx = baggage.set_baggage("key", "value")
         context.attach(ctx)
 
-        task = await task_returns_baggage.kiq()
-        result = await task.wait_result(timeout=2)
+        result = asyncio.run(test())
 
         self.assertEqual(result.return_value, {"key": "value"})
 
-    async def test_task_not_instrumented_does_not_raise(self) -> None:
+    def test_task_not_instrumented_does_not_raise(self) -> None:
         def _retrieve_context_wrapper_none_token(
             wrapped: Callable[
                 [Any],
@@ -178,9 +188,11 @@ def _retrieve_context_wrapper_none_token(
 
         TaskiqInstrumentor().instrument_broker(broker)
 
-        task = await task_add.kiq(1, 2)
-        result = await task.wait_result(timeout=2)
+        async def test() -> TaskiqResult[float]:
+            task = await task_add.kiq(1, 2)
+            return await task.wait_result(timeout=2)
 
+        result = asyncio.run(test())
         spans = self.sorted_spans(self.memory_exporter.get_finished_spans())
         self.assertEqual(len(spans), 2)