Add instrument_celery method (#322)

Author: Kludex

.github/workflows/main.yml

@@ -63,6 +63,8 @@ jobs:
           echo "203.0.113.0 logfire.dev" | sudo tee -a /etc/hosts
           echo "203.0.113.0 logfire-api.pydantic.dev" | sudo tee -a /etc/hosts
           echo "203.0.113.0 logfire.pydantic.dev" | sudo tee -a /etc/hosts
+      - name: Setup services
+        run: docker-compose up -d
       - name: set up python
         uses: actions/setup-python@v5
         with:

Makefile

@@ -30,6 +30,10 @@ lint:
 test:
 	rye run coverage run -m pytest
 
+.PHONY: test-integration  # Run the integration tests
+test-integration:
+	rye run coverage run -m pytest -m 'integration'
+
 .PHONY: generate-stubs  # Generate stubs for logfire-api
 generate-stubs:
 	rye run generate-stubs

compose.yml

@@ -0,0 +1,6 @@
+services:
+  redis:
+    image: redis:latest
+    container_name: redis
+    ports:
+      - "6379:6379"

docs/integrations/celery.md

@@ -1,6 +1,7 @@
 # Celery
 
-The [OpenTelemetry Instrumentation Celery][opentelemetry-celery] package can be used to instrument [Celery][celery].
+The [`logfire.instrument_celery()`][logfire.Logfire.instrument_celery] method will create a span for every task
+executed by your Celery workers.
 
 ## Installation
 
@@ -27,14 +28,13 @@ Below we have a minimal example using Celery. You can run it with `celery -A tas
 import logfire
 from celery import Celery
 from celery.signals import worker_process_init
-from opentelemetry.instrumentation.celery import CeleryInstrumentor
 
 
 logfire.configure()
 
 @worker_process_init.connect(weak=False)
 def init_celery_tracing(*args, **kwargs):
-    CeleryInstrumentor().instrument()
+    logfire.instrument_celery()
 
 app = Celery("tasks", broker="pyamqp://localhost//")  # (1)!
 
@@ -47,7 +47,8 @@ add.delay(42, 50)
 
 1. Install `pyamqp` with `pip install pyamqp`.
 
-You can read more about the Celery OpenTelemetry package [here][opentelemetry-celery].
+The keyword arguments of [`logfire.instrument_celery()`][logfire.Logfire.instrument_celery] are passed to the
+[`CeleryInstrumentor().instrument()`][opentelemetry.instrumentation.celery.CeleryInstrumentor] method.
 
 [celery]: https://docs.celeryq.dev/en/stable/
 [opentelemetry-celery]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/celery/celery.html

logfire/__init__.py

@@ -28,6 +28,7 @@
 instrument_anthropic = DEFAULT_LOGFIRE_INSTANCE.instrument_anthropic
 instrument_asyncpg = DEFAULT_LOGFIRE_INSTANCE.instrument_asyncpg
 instrument_httpx = DEFAULT_LOGFIRE_INSTANCE.instrument_httpx
+instrument_celery = DEFAULT_LOGFIRE_INSTANCE.instrument_celery
 instrument_requests = DEFAULT_LOGFIRE_INSTANCE.instrument_requests
 instrument_psycopg = DEFAULT_LOGFIRE_INSTANCE.instrument_psycopg
 instrument_django = DEFAULT_LOGFIRE_INSTANCE.instrument_django

logfire/_internal/formatter.py

@@ -203,7 +203,7 @@ def _fstring_chunks(
             if isinstance(node_value, ast.Constant):
                 # These are the parts of the f-string not enclosed by `{}`, e.g. 'foo ' in f'foo {bar}'
                 value = node_value.value
-                assert type(value) is str  # noqa
+                assert type(value) is str
                 result.append({'v': value, 't': 'lit'})
                 new_template += value
             else:

logfire/_internal/integrations/celery.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from opentelemetry.instrumentation.celery import CeleryInstrumentor
+
+if TYPE_CHECKING:
+    from typing_extensions import TypedDict, Unpack
+
+    class CeleryInstrumentKwargs(TypedDict, total=False):
+        skip_dep_check: bool
+
+
+def instrument_celery(**kwargs: Unpack[CeleryInstrumentKwargs]) -> None:
+    """Instrument the `celery` module so that spans are automatically created for each task.
+
+    See the `Logfire.instrument_celery` method for details.
+    """
+    return CeleryInstrumentor().instrument(**kwargs)  # type: ignore[reportUnknownMemberType]

logfire/_internal/main.py

@@ -75,6 +75,7 @@
     from starlette.websockets import WebSocket
     from typing_extensions import Unpack
 
+    from .integrations.celery import CeleryInstrumentKwargs
     from .integrations.flask import FlaskInstrumentKwargs
     from .integrations.psycopg import PsycopgInstrumentKwargs
     from .integrations.pymongo import PymongoInstrumentKwargs
@@ -1049,6 +1050,18 @@ def instrument_httpx(self, **kwargs: Any):
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_httpx(**kwargs)
 
+    def instrument_celery(self, **kwargs: Unpack[CeleryInstrumentKwargs]) -> None:
+        """Instrument `celery` so that spans are automatically created for each task.
+
+        Uses the
+        [OpenTelemetry Celery Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/celery/celery.html)
+        library.
+        """
+        from .integrations.celery import instrument_celery
+
+        self._warn_if_not_initialized_for_instrumentation()
+        return instrument_celery(**kwargs)
+
     def instrument_django(
         self,
         is_sql_commentor_enabled: bool | None = None,

pyproject.toml

@@ -113,6 +113,7 @@ dev-dependencies = [
     "opentelemetry-instrumentation-psycopg2",
     "opentelemetry-instrumentation-redis",
     "opentelemetry-instrumentation-pymongo",
+    "opentelemetry-instrumentation-celery",
     "eval-type-backport",
     "requests-mock",
     "inline-snapshot",
@@ -133,6 +134,7 @@ dev-dependencies = [
     "anthropic>=0.27.0",
     "sqlmodel",
     "mypy>=1.10.0",
+    "celery>=5.4.0",
 ]
 
 [tool.rye.scripts]
@@ -198,11 +200,20 @@ quote-style = "single"
 typeCheckingMode = "strict"
 reportUnnecessaryTypeIgnoreComment = true
 reportMissingTypeStubs = false
-exclude = ["docs/**/*.py", "site/**/*.py", ".venv", "venv*", "ignoreme", "out", "logfire-api"]
+exclude = [
+    "docs/**/*.py",
+    "site/**/*.py",
+    ".venv",
+    "venv*",
+    "ignoreme",
+    "out",
+    "logfire-api",
+]
 venvPath = ".venv"
 
 [tool.pytest.ini_options]
 xfail_strict = true
+markers = ['integration: mark a test as an integration test']
 filterwarnings = [
     "error",
     # fastapi uses deprecated pydantic functions

requirements-dev.lock

@@ -6,12 +6,13 @@
 #   features: []
 #   all-features: false
 #   with-sources: false
-#   generate-hashes: false
 
 -e file:.
 aiohttp==3.9.5
 aiosignal==1.3.1
     # via aiohttp
+amqp==5.2.0
+    # via kombu
 annotated-types==0.7.0
     # via pydantic
 anthropic==0.31.2
@@ -31,10 +32,13 @@ attrs==23.2.0
     # via aiohttp
 babel==2.15.0
     # via mkdocs-material
+billiard==4.2.0
+    # via celery
 black==24.4.2
     # via inline-snapshot
 blinker==1.8.2
     # via flask
+celery==5.4.0
 certifi==2024.7.4
     # via httpcore
     # via httpx
@@ -45,12 +49,22 @@ charset-normalizer==3.3.2
     # via requests
 click==8.1.7
     # via black
+    # via celery
+    # via click-didyoumean
+    # via click-plugins
+    # via click-repl
     # via flask
     # via inline-snapshot
     # via mkdocs
     # via mkdocstrings
     # via typer
     # via uvicorn
+click-didyoumean==0.3.1
+    # via celery
+click-plugins==1.1.1
+    # via celery
+click-repl==0.3.0
+    # via celery
 cloudpickle==3.0.0
 colorama==0.4.6
     # via griffe
@@ -130,6 +144,8 @@ jinja2==3.1.4
     # via mkdocstrings
 jiter==0.5.0
     # via anthropic
+kombu==5.3.7
+    # via celery
 loguru==0.7.2
 markdown==3.6
     # via mkdocs
@@ -184,6 +200,7 @@ opentelemetry-api==1.25.0
     # via opentelemetry-instrumentation-aiohttp-client
     # via opentelemetry-instrumentation-asgi
     # via opentelemetry-instrumentation-asyncpg
+    # via opentelemetry-instrumentation-celery
     # via opentelemetry-instrumentation-dbapi
     # via opentelemetry-instrumentation-django
     # via opentelemetry-instrumentation-fastapi
@@ -209,6 +226,7 @@ opentelemetry-instrumentation==0.46b0
     # via opentelemetry-instrumentation-aiohttp-client
     # via opentelemetry-instrumentation-asgi
     # via opentelemetry-instrumentation-asyncpg
+    # via opentelemetry-instrumentation-celery
     # via opentelemetry-instrumentation-dbapi
     # via opentelemetry-instrumentation-django
     # via opentelemetry-instrumentation-fastapi
@@ -228,6 +246,7 @@ opentelemetry-instrumentation-asgi==0.46b0
     # via opentelemetry-instrumentation-fastapi
     # via opentelemetry-instrumentation-starlette
 opentelemetry-instrumentation-asyncpg==0.46b0
+opentelemetry-instrumentation-celery==0.46b0
 opentelemetry-instrumentation-dbapi==0.46b0
     # via opentelemetry-instrumentation-psycopg
     # via opentelemetry-instrumentation-psycopg2
@@ -256,6 +275,7 @@ opentelemetry-semantic-conventions==0.46b0
     # via opentelemetry-instrumentation-aiohttp-client
     # via opentelemetry-instrumentation-asgi
     # via opentelemetry-instrumentation-asyncpg
+    # via opentelemetry-instrumentation-celery
     # via opentelemetry-instrumentation-dbapi
     # via opentelemetry-instrumentation-django
     # via opentelemetry-instrumentation-fastapi
@@ -299,6 +319,8 @@ platformdirs==4.2.2
 pluggy==1.5.0
     # via pytest
 pre-commit==3.7.1
+prompt-toolkit==3.0.47
+    # via click-repl
 protobuf==4.25.3
     # via googleapis-common-protos
     # via logfire
@@ -330,6 +352,7 @@ pytest==8.3.1
 pytest-django==4.8.0
 pytest-pretty==1.2.0
 python-dateutil==2.9.0.post0
+    # via celery
     # via ghp-import
     # via pandas
 python-dotenv==1.0.1
@@ -409,19 +432,26 @@ typing-extensions==4.12.2
     # via sqlalchemy
     # via typer
 tzdata==2024.1
+    # via celery
     # via pandas
 urllib3==2.2.2
     # via requests
 uvicorn==0.30.3
     # via fastapi
 uvloop==0.19.0
     # via uvicorn
+vine==5.1.0
+    # via amqp
+    # via celery
+    # via kombu
 virtualenv==20.26.3
     # via pre-commit
 watchdog==4.0.1
     # via mkdocs
 watchfiles==0.22.0
     # via uvicorn
+wcwidth==0.2.13
+    # via prompt-toolkit
 websockets==12.0
     # via uvicorn
 werkzeug==3.0.3