|
22 | 22 | from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor |
23 | 23 | from opentelemetry.trace import Span, StatusCode |
24 | 24 | from redis import Redis |
| 25 | +from redis.cluster import RedisCluster |
25 | 26 | from requests import PreparedRequest, Response |
26 | 27 | from sqlalchemy import Engine |
27 | 28 | from fastapi import status |
@@ -59,16 +60,28 @@ def response_hook(span: Span, request: PreparedRequest, response: Response): |
59 | 60 | span.set_status(StatusCode.ERROR if response.status_code >= 400 else StatusCode.OK) |
60 | 61 |
|
61 | 62 |
|
62 | | -def redis_request_hook(span: Span, instance: Redis, args, kwargs): |
| 63 | +def redis_request_hook(span: Span, instance: Union[Redis|RedisCluster], args, kwargs): |
63 | 64 | """ |
64 | 65 | Redis Request Hook |
65 | 66 | """ |
66 | 67 |
|
| 68 | + # In cluster mode, the instance can be of two types: |
| 69 | + # - redis.asyncio.cluster.RedisCluster |
| 70 | + # - redis.cluster.RedisCluster |
| 71 | + # Instead of checking the type, we check if the instance has a nodes_manager attribute. |
67 | 72 | try: |
68 | | - connection_kwargs: dict = instance.connection_pool.connection_kwargs |
69 | | - host = connection_kwargs.get("host") |
70 | | - port = connection_kwargs.get("port") |
71 | | - db = connection_kwargs.get("db") |
| 73 | + db = "" |
| 74 | + if hasattr(instance, 'nodes_manager'): |
| 75 | + default_node = instance.nodes_manager.default_node |
| 76 | + if not default_node: |
| 77 | + return |
| 78 | + host = default_node.host |
| 79 | + port = default_node.port |
| 80 | + else: |
| 81 | + connection_kwargs: dict = instance.connection_pool.connection_kwargs |
| 82 | + host = connection_kwargs.get("host") |
| 83 | + port = connection_kwargs.get("port") |
| 84 | + db = connection_kwargs.get("db") |
72 | 85 | span.set_attributes( |
73 | 86 | { |
74 | 87 | SpanAttributes.DB_INSTANCE: f"{host}/{db}", |
|
0 commit comments