Skip to content

Commit 24b95f1

Browse files
Kludexalexmojaki
Kludex
and
alexmojaki
authored
Add capture_statement to Redis instrumentation (#355)
Co-authored-by: Alex Hall <[email protected]>
1 parent 4a4dea0 commit 24b95f1

File tree

5 files changed

+206
-21
lines changed

5 files changed

+206
-21
lines changed

‎docs/integrations/pymongo.md‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ import logfire
3030
from pymongo import MongoClient
3131

3232
logfire.configure()
33-
logfire.instrument_pymongo(capture_statement=True) # (1)!
33+
logfire.instrument_pymongo()
3434

3535
client = MongoClient()
3636
db = client["database"]
@@ -39,12 +39,10 @@ collection.insert_one({"name": "MongoDB"})
3939
collection.find_one()
4040
```
4141

42-
1. The `capture_statement` parameter is set to `True` to capture the executed statements.
42+
!!! info
43+
You can pass `capture_statement=True` to `logfire.instrument_pymongo()` to capture the queries.
4344

44-
This is the default behavior on other OpenTelemetry instrumentation packages, but it's
45-
disabled by default in PyMongo.
46-
47-
---
45+
By default, it is set to `False` to avoid capturing sensitive information.
4846

4947
The keyword arguments of `logfire.instrument_pymongo()` are passed to the `PymongoInstrumentor().instrument()` method of the OpenTelemetry pymongo Instrumentation package, read more about it [here][opentelemetry-pymongo].
5048

‎docs/integrations/redis.md‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,18 @@ Install `logfire` with the `redis` extra:
1010

1111
## Usage
1212

13-
Let's see a minimal example below:
13+
Let's setup a container with Redis and run a Python script that connects to the Redis server to
14+
demonstrate how to use **Logfire** with Redis.
1415

15-
<!-- TODO(Marcelo): Create a secret gist with a docker-compose. -->
16+
### Setup a Redis Server Using Docker
17+
18+
First, we need to initialize a Redis server. This can be easily done using Docker with the following command:
19+
20+
```bash
21+
docker run --name redis -p 6379:6379 -d redis:latest
22+
```
23+
24+
### Run the Python script
1625

1726
```py title="main.py"
1827
import logfire
@@ -22,11 +31,9 @@ import redis
2231
logfire.configure()
2332
logfire.instrument_redis()
2433

25-
# This will report a span with the default settings
2634
client = redis.StrictRedis(host="localhost", port=6379)
27-
client.get("my-key")
35+
client.set("my-key", "my-value")
2836

29-
# This will report a span with the default settings
3037
async def main():
3138
client = redis.asyncio.Redis(host="localhost", port=6379)
3239
await client.get("my-key")
@@ -37,6 +44,11 @@ if __name__ == "__main__":
3744
asyncio.run(main())
3845
```
3946

47+
!!! info
48+
You can pass `capture_statement=True` to `logfire.instrument_redis()` to capture the Redis command.
49+
50+
By default, it is set to `False` given that Redis commands can contain sensitive information.
51+
4052
The keyword arguments of `logfire.instrument_redis()` are passed to the `RedisInstrumentor().instrument()` method of the OpenTelemetry Redis Instrumentation package, read more about it [here][opentelemetry-redis].
4153

4254
[redis]: https://redis.readthedocs.io/en/stable/

‎logfire/_internal/integrations/redis.py‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
from __future__ import annotations
22

3+
import functools
34
from typing import TYPE_CHECKING, Any
45

56
from opentelemetry.instrumentation.redis import RedisInstrumentor
67

8+
from logfire._internal.constants import ATTRIBUTES_MESSAGE_KEY
9+
from logfire._internal.utils import truncate_string
10+
711
if TYPE_CHECKING:
812
from opentelemetry.trace import Span
913
from redis import Connection
@@ -21,9 +25,32 @@ class RedisInstrumentKwargs(TypedDict, total=False):
2125
skip_dep_check: bool
2226

2327

24-
def instrument_redis(**kwargs: Unpack[RedisInstrumentKwargs]) -> None:
28+
def instrument_redis(capture_statement: bool = False, **kwargs: Unpack[RedisInstrumentKwargs]) -> None:
2529
"""Instrument the `redis` module so that spans are automatically created for each operation.
2630
2731
See the `Logfire.instrument_redis` method for details.
32+
33+
Args:
34+
capture_statement: Whether to capture the statement being executed. Defaults to False.
35+
**kwargs: Additional keyword arguments to pass to the `RedisInstrumentor.instrument` method.
2836
"""
29-
RedisInstrumentor().instrument(**kwargs) # type: ignore[reportUnknownMemberType]
37+
request_hook = kwargs.pop('request_hook', None)
38+
if capture_statement:
39+
request_hook = _capture_statement_hook(request_hook)
40+
41+
RedisInstrumentor().instrument(request_hook=request_hook, **kwargs) # type: ignore[reportUnknownMemberType]
42+
43+
44+
def _capture_statement_hook(request_hook: RequestHook | None = None) -> RequestHook:
45+
truncate_value = functools.partial(truncate_string, max_length=20, middle='...')
46+
47+
def _capture_statement(
48+
span: Span, instance: Connection, command: tuple[object, ...], *args: Any, **kwargs: Any
49+
) -> None:
50+
str_command = list(map(str, command))
51+
span.set_attribute('db.statement', ' '.join(str_command))
52+
span.set_attribute(ATTRIBUTES_MESSAGE_KEY, ' '.join(map(truncate_value, str_command)))
53+
if request_hook is not None:
54+
request_hook(span, instance, command, *args, **kwargs)
55+
56+
return _capture_statement

‎logfire/_internal/main.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1212,17 +1212,21 @@ def instrument_pymongo(self, **kwargs: Unpack[PymongoInstrumentKwargs]) -> None:
12121212
self._warn_if_not_initialized_for_instrumentation()
12131213
return instrument_pymongo(**kwargs)
12141214

1215-
def instrument_redis(self, **kwargs: Unpack[RedisInstrumentKwargs]) -> None:
1215+
def instrument_redis(self, capture_statement: bool = False, **kwargs: Unpack[RedisInstrumentKwargs]) -> None:
12161216
"""Instrument the `redis` module so that spans are automatically created for each operation.
12171217
12181218
Uses the
12191219
[OpenTelemetry Redis Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/redis/redis.html)
12201220
library, specifically `RedisInstrumentor().instrument()`, to which it passes `**kwargs`.
1221+
1222+
Args:
1223+
capture_statement: Set to `True` to capture the statement in the span attributes.
1224+
kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods.
12211225
"""
12221226
from .integrations.redis import instrument_redis
12231227

12241228
self._warn_if_not_initialized_for_instrumentation()
1225-
return instrument_redis(**kwargs)
1229+
return instrument_redis(capture_statement=capture_statement, **kwargs)
12261230

12271231
def instrument_mysql(
12281232
self,

‎tests/otel_integrations/test_redis.py‎

Lines changed: 150 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,155 @@
1-
from redis import Redis
1+
from typing import Any, Iterator
2+
3+
import pytest
4+
from inline_snapshot import snapshot
5+
from opentelemetry.instrumentation.redis import RedisInstrumentor
6+
from opentelemetry.trace import Span
7+
from redis import Connection, Redis
8+
from testcontainers.redis import RedisContainer
29

310
import logfire
11+
from logfire.testing import TestExporter
12+
13+
14+
@pytest.fixture(scope='module', autouse=True)
15+
def redis_container() -> Iterator[RedisContainer]:
16+
with RedisContainer('redis:latest') as redis:
17+
yield redis
18+
19+
20+
@pytest.fixture
21+
def redis(redis_container: RedisContainer) -> Redis:
22+
return redis_container.get_client() # type: ignore
23+
424

25+
@pytest.fixture
26+
def redis_port(redis_container: RedisContainer) -> str:
27+
return redis_container.get_exposed_port(6379)
528

6-
# TODO real test
7-
def test_instrument_redis():
8-
original = Redis.execute_command # type: ignore
9-
assert Redis.execute_command is original # type: ignore
29+
30+
@pytest.fixture(autouse=True)
31+
def uninstrument_redis():
32+
try:
33+
yield
34+
finally:
35+
RedisInstrumentor().uninstrument() # type: ignore
36+
37+
38+
def test_instrument_redis(redis: Redis, redis_port: str, exporter: TestExporter):
1039
logfire.instrument_redis()
11-
assert Redis.execute_command is not original # type: ignore
40+
41+
redis.set('my-key', 123)
42+
43+
assert exporter.exported_spans_as_dict() == snapshot(
44+
[
45+
{
46+
'name': 'SET',
47+
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
48+
'parent': None,
49+
'start_time': 1000000000,
50+
'end_time': 2000000000,
51+
'attributes': {
52+
'logfire.span_type': 'span',
53+
'logfire.msg': 'SET ? ?',
54+
'db.statement': 'SET ? ?',
55+
'db.system': 'redis',
56+
'db.redis.database_index': 0,
57+
'net.peer.name': 'localhost',
58+
'net.peer.port': redis_port,
59+
'net.transport': 'ip_tcp',
60+
'db.redis.args_length': 3,
61+
},
62+
}
63+
]
64+
)
65+
66+
67+
def test_instrument_redis_with_capture_statement(redis: Redis, redis_port: str, exporter: TestExporter):
68+
logfire.instrument_redis(capture_statement=True)
69+
70+
redis.set('my-key', 123)
71+
72+
assert exporter.exported_spans_as_dict() == snapshot(
73+
[
74+
{
75+
'name': 'SET',
76+
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
77+
'parent': None,
78+
'start_time': 1000000000,
79+
'end_time': 2000000000,
80+
'attributes': {
81+
'logfire.span_type': 'span',
82+
'logfire.msg': 'SET my-key 123',
83+
'db.statement': 'SET my-key 123',
84+
'db.system': 'redis',
85+
'db.redis.database_index': 0,
86+
'net.peer.name': 'localhost',
87+
'net.peer.port': redis_port,
88+
'net.transport': 'ip_tcp',
89+
'db.redis.args_length': 3,
90+
},
91+
}
92+
]
93+
)
94+
95+
96+
def test_instrument_redis_with_big_capture_statement(redis: Redis, redis_port: str, exporter: TestExporter):
97+
logfire.instrument_redis(capture_statement=True)
98+
99+
redis.set('k' * 100, 'x' * 1000)
100+
101+
assert exporter.exported_spans_as_dict() == snapshot(
102+
[
103+
{
104+
'name': 'SET',
105+
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
106+
'parent': None,
107+
'start_time': 1000000000,
108+
'end_time': 2000000000,
109+
'attributes': {
110+
'logfire.span_type': 'span',
111+
'db.system': 'redis',
112+
'db.redis.database_index': 0,
113+
'net.peer.name': 'localhost',
114+
'net.peer.port': redis_port,
115+
'net.transport': 'ip_tcp',
116+
'db.redis.args_length': 3,
117+
'db.statement': 'SET kkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkk xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
118+
'logfire.msg': 'SET kkkkkkkk...kkkkkkkk xxxxxxxx...xxxxxxxx',
119+
},
120+
}
121+
]
122+
)
123+
124+
125+
def test_instrument_redis_with_request_hook(redis: Redis, redis_port: str, exporter: TestExporter):
126+
def request_hook(span: Span, instance: Connection, *args: Any, **kwargs: Any) -> None:
127+
span.set_attribute('potato', 'tomato')
128+
129+
logfire.instrument_redis(request_hook=request_hook, capture_statement=True)
130+
131+
redis.set('my-key', 123)
132+
133+
assert exporter.exported_spans_as_dict() == snapshot(
134+
[
135+
{
136+
'name': 'SET',
137+
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
138+
'parent': None,
139+
'start_time': 1000000000,
140+
'end_time': 2000000000,
141+
'attributes': {
142+
'logfire.span_type': 'span',
143+
'logfire.msg': 'SET my-key 123',
144+
'db.statement': 'SET my-key 123',
145+
'db.system': 'redis',
146+
'db.redis.database_index': 0,
147+
'net.peer.name': 'localhost',
148+
'net.peer.port': redis_port,
149+
'net.transport': 'ip_tcp',
150+
'db.redis.args_length': 3,
151+
'potato': 'tomato',
152+
},
153+
}
154+
]
155+
)

0 commit comments

Comments
 (0)