Skip to content
7 changes: 6 additions & 1 deletion logfire/_internal/cli/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,12 @@ def instrument_packages(installed_otel_packages: set[str], instrument_pkg_map: d

def instrument_package(import_name: str):
instrument_attr = f'instrument_{import_name}'
getattr(logfire, instrument_attr)()

if import_name in ('starlette', 'fastapi', 'flask'):
module = importlib.import_module(f'logfire._internal.integrations.{import_name}')
getattr(module, instrument_attr)(logfire.DEFAULT_LOGFIRE_INSTANCE)
else:
getattr(logfire, instrument_attr)()


def find_recommended_instrumentations_to_install(
Expand Down
80 changes: 44 additions & 36 deletions logfire/_internal/integrations/fastapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ def find_mounted_apps(app: FastAPI) -> list[FastAPI]:

def instrument_fastapi(
logfire_instance: Logfire,
app: FastAPI,
app: FastAPI | None = None,
*,
capture_headers: bool = False,
request_attributes_mapper: Callable[
Expand Down Expand Up @@ -78,39 +78,53 @@ def instrument_fastapi(
'meter_provider': logfire_instance.config.get_meter_provider(),
**opentelemetry_kwargs,
}
FastAPIInstrumentor.instrument_app(
app,
excluded_urls=excluded_urls,
server_request_hook=_server_request_hook(opentelemetry_kwargs.pop('server_request_hook', None)),
**opentelemetry_kwargs,
)
if app is None:
FastAPIInstrumentor().instrument(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If it seems too hard to also apply the other fastapi instrumentation, at least comment about it

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't have the app instance, how would I go about it to make the custom logic work?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if request.app not in registry then check for a 'global' FastAPIInstrumentation next.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You mean modifying the patch function below?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes

excluded_urls=excluded_urls,
server_request_hook=_server_request_hook(opentelemetry_kwargs.pop('server_request_hook', None)),
**opentelemetry_kwargs,
)

registry = patch_fastapi()
if app in registry: # pragma: no cover
raise ValueError('This app has already been instrumented.')
@contextmanager
def uninstrument_context():
yield
FastAPIInstrumentor().uninstrument()

return uninstrument_context()
else:
FastAPIInstrumentor.instrument_app(
app,
excluded_urls=excluded_urls,
server_request_hook=_server_request_hook(opentelemetry_kwargs.pop('server_request_hook', None)),
**opentelemetry_kwargs,
)

mounted_apps = find_mounted_apps(app)
mounted_apps.append(app)
registry = patch_fastapi()
if app in registry: # pragma: no cover
raise ValueError('This app has already been instrumented.')

for _app in mounted_apps:
registry[_app] = FastAPIInstrumentation(
logfire_instance,
request_attributes_mapper or _default_request_attributes_mapper,
)
mounted_apps = find_mounted_apps(app)
mounted_apps.append(app)

@contextmanager
def uninstrument_context():
# The user isn't required (or even expected) to use this context manager,
# which is why the instrumenting and patching has already happened before this point.
# It exists mostly for tests, and just in case users want it.
try:
yield
finally:
for _app in mounted_apps:
del registry[_app]
FastAPIInstrumentor.uninstrument_app(_app)
for _app in mounted_apps:
registry[_app] = FastAPIInstrumentation(
logfire_instance,
request_attributes_mapper or _default_request_attributes_mapper,
)

return uninstrument_context()
@contextmanager
def uninstrument_context():
# The user isn't required (or even expected) to use this context manager,
# which is why the instrumenting and patching has already happened before this point.
# It exists mostly for tests, and just in case users want it.
try:
yield
finally:
for _app in mounted_apps:
del registry[_app]
FastAPIInstrumentor.uninstrument_app(_app)

return uninstrument_context()


@lru_cache # only patch once
Expand Down Expand Up @@ -151,13 +165,7 @@ class FastAPIInstrumentation:
def __init__(
self,
logfire_instance: Logfire,
request_attributes_mapper: Callable[
[
Request | WebSocket,
dict[str, Any],
],
dict[str, Any] | None,
],
request_attributes_mapper: Callable[[Request | WebSocket, dict[str, Any]], dict[str, Any] | None],
):
self.logfire_instance = logfire_instance.with_settings(custom_scope_suffix='fastapi')
self.request_attributes_mapper = request_attributes_mapper
Expand Down
46 changes: 33 additions & 13 deletions logfire/_internal/integrations/flask.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,18 @@
" pip install 'logfire[flask]'"
)

from logfire import Logfire
from logfire._internal.utils import maybe_capture_server_headers
from logfire.integrations.flask import CommenterOptions, RequestHook, ResponseHook


def instrument_flask(
app: Flask,
logfire_instance: Logfire,
app: Flask | None = None,
*,
capture_headers: bool,
enable_commenter: bool,
commenter_options: CommenterOptions | None,
capture_headers: bool = False,
enable_commenter: bool = False,
commenter_options: CommenterOptions | None = None,
excluded_urls: str | None = None,
request_hook: RequestHook | None = None,
response_hook: ResponseHook | None = None,
Expand All @@ -41,12 +43,30 @@ def instrument_flask(
warn_at_user_stacklevel('exclude_urls is deprecated; use excluded_urls instead', DeprecationWarning)
excluded_urls = excluded_urls or kwargs.pop('exclude_urls', None)

FlaskInstrumentor().instrument_app( # type: ignore[reportUnknownMemberType]
app,
enable_commenter=enable_commenter,
commenter_options=commenter_options,
excluded_urls=excluded_urls,
request_hook=request_hook,
response_hook=response_hook,
**kwargs,
)
if app is None:
FlaskInstrumentor().instrument(
enable_commenter=enable_commenter,
commenter_options=commenter_options,
excluded_urls=excluded_urls,
request_hook=request_hook,
response_hook=response_hook,
**{
'tracer_provider': logfire_instance.config.get_tracer_provider(),
'meter_provider': logfire_instance.config.get_meter_provider(),
**kwargs,
},
)
else:
FlaskInstrumentor().instrument_app( # type: ignore[reportUnknownMemberType]
app,
enable_commenter=enable_commenter,
commenter_options=commenter_options,
excluded_urls=excluded_urls,
request_hook=request_hook,
response_hook=response_hook,
**{
'tracer_provider': logfire_instance.config.get_tracer_provider(),
'meter_provider': logfire_instance.config.get_meter_provider(),
**kwargs,
},
)
36 changes: 24 additions & 12 deletions logfire/_internal/integrations/starlette.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

def instrument_starlette(
logfire_instance: Logfire,
app: Starlette,
app: Starlette | None = None,
*,
record_send_receive: bool = False,
capture_headers: bool = False,
Expand All @@ -35,14 +35,26 @@ def instrument_starlette(
See the `Logfire.instrument_starlette` method for details.
"""
maybe_capture_server_headers(capture_headers)
StarletteInstrumentor().instrument_app(
app,
server_request_hook=server_request_hook,
client_request_hook=client_request_hook,
client_response_hook=client_response_hook,
**{ # type: ignore
'tracer_provider': tweak_asgi_spans_tracer_provider(logfire_instance, record_send_receive),
'meter_provider': logfire_instance.config.get_meter_provider(),
**kwargs,
},
)
if app is None:
StarletteInstrumentor().instrument(
server_request_hook=server_request_hook,
client_request_hook=client_request_hook,
client_response_hook=client_response_hook,
**{
'tracer_provider': tweak_asgi_spans_tracer_provider(logfire_instance, record_send_receive),
'meter_provider': logfire_instance.config.get_meter_provider(),
**kwargs,
},
)
else:
StarletteInstrumentor().instrument_app(
app,
server_request_hook=server_request_hook,
client_request_hook=client_request_hook,
client_response_hook=client_response_hook,
**{ # type: ignore
'tracer_provider': tweak_asgi_spans_tracer_provider(logfire_instance, record_send_receive),
'meter_provider': logfire_instance.config.get_meter_provider(),
**kwargs,
},
)
17 changes: 3 additions & 14 deletions logfire/_internal/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,7 @@
from enum import Enum
from functools import cached_property
from time import time
from typing import (
TYPE_CHECKING,
Any,
Callable,
Literal,
TypeVar,
Union,
overload,
)
from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar, Union, overload

import opentelemetry.context as context_api
import opentelemetry.trace as trace_api
Expand Down Expand Up @@ -1572,18 +1564,15 @@ def instrument_flask(

self._warn_if_not_initialized_for_instrumentation()
return instrument_flask(
self,
app,
capture_headers=capture_headers,
enable_commenter=enable_commenter,
commenter_options=commenter_options,
excluded_urls=excluded_urls,
request_hook=request_hook,
response_hook=response_hook,
**{
'tracer_provider': self._config.get_tracer_provider(),
'meter_provider': self._config.get_meter_provider(),
**kwargs,
},
**kwargs,
)

def instrument_starlette(
Expand Down
33 changes: 33 additions & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
)
from logfire._internal.config import LogfireCredentials, sanitize_project_name
from logfire.exceptions import LogfireConfigError
from logfire.testing import TestExporter
from tests.import_used_for_tests import run_script_test


Expand Down Expand Up @@ -1534,6 +1535,38 @@ async def test_instrument_packages_aiohttp_client() -> None:
AioHttpClientInstrumentor().uninstrument()


def test_instrument_web_frameworks(exporter: TestExporter) -> None:
try:
instrument_packages(
{
'opentelemetry-instrumentation-starlette',
'opentelemetry-instrumentation-fastapi',
'opentelemetry-instrumentation-flask',
},
{
'opentelemetry-instrumentation-starlette': 'starlette',
'opentelemetry-instrumentation-fastapi': 'fastapi',
'opentelemetry-instrumentation-flask': 'flask',
},
)

from fastapi import FastAPI
from flask import Flask
from starlette.applications import Starlette

assert getattr(Starlette(), '_is_instrumented_by_opentelemetry', False) is True
assert getattr(FastAPI(), '_is_instrumented_by_opentelemetry', False) is True
assert getattr(Flask(__name__), '_is_instrumented_by_opentelemetry', False) is True
finally:
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.flask import FlaskInstrumentor
from opentelemetry.instrumentation.starlette import StarletteInstrumentor

StarletteInstrumentor().uninstrument()
FastAPIInstrumentor().uninstrument()
FlaskInstrumentor().uninstrument()


def test_split_args_action() -> None:
parser = argparse.ArgumentParser()
parser.add_argument('--foo', action=SplitArgs)
Expand Down
Loading