|
| 1 | +from contextlib import contextmanager |
1 | 2 | from typing import Dict |
2 | 3 | from typing import List |
3 | 4 | from typing import Optional |
4 | 5 | from typing import Union |
5 | 6 |
|
| 7 | +from ddtrace.constants import SPAN_KIND |
| 8 | +from ddtrace.contrib import trace_utils |
| 9 | +from ddtrace.ext import SpanKind |
| 10 | +from ddtrace.ext import SpanTypes |
| 11 | +from ddtrace.ext import db |
6 | 12 | from ddtrace.ext import net |
7 | 13 | from ddtrace.ext import redis as redisx |
8 | 14 | from ddtrace.internal import core |
| 15 | +from ddtrace.internal.constants import COMPONENT |
| 16 | +from ddtrace.internal.schema import schematize_cache_operation |
9 | 17 | from ddtrace.internal.utils.formats import stringify_cache_args |
10 | 18 |
|
11 | 19 |
|
|
30 | 38 | ROW_RETURNING_COMMANDS = SINGLE_KEY_COMMANDS + MULTI_KEY_COMMANDS |
31 | 39 |
|
32 | 40 |
|
33 | | -def _extract_conn_tags(conn_kwargs): |
34 | | - """Transform redis conn info into dogtrace metas""" |
35 | | - try: |
36 | | - conn_tags = { |
37 | | - net.TARGET_HOST: conn_kwargs["host"], |
38 | | - net.TARGET_PORT: conn_kwargs["port"], |
39 | | - net.SERVER_ADDRESS: conn_kwargs["host"], |
40 | | - redisx.DB: conn_kwargs.get("db") or 0, |
41 | | - } |
42 | | - client_name = conn_kwargs.get("client_name") |
43 | | - if client_name: |
44 | | - conn_tags[redisx.CLIENT_NAME] = client_name |
45 | | - return conn_tags |
46 | | - except Exception: |
47 | | - return {} |
48 | | - |
49 | | - |
50 | 41 | def determine_row_count(redis_command: str, result: Optional[Union[List, Dict, str]]) -> int: |
51 | 42 | empty_results = [b"", [], {}, None] |
52 | 43 | # result can be an empty list / dict / string |
@@ -82,3 +73,75 @@ async def _run_redis_command_async(ctx: core.ExecutionContext, func, args, kwarg |
82 | 73 | if redis_command not in ROW_RETURNING_COMMANDS: |
83 | 74 | rowcount = None |
84 | 75 | core.dispatch("redis.async_command.post", [ctx, rowcount]) |
| 76 | + |
| 77 | + |
| 78 | +def _extract_conn_tags(conn_kwargs) -> Dict[str, str]: |
| 79 | + try: |
| 80 | + conn_tags = { |
| 81 | + net.TARGET_HOST: conn_kwargs["host"], |
| 82 | + net.TARGET_PORT: conn_kwargs["port"], |
| 83 | + net.SERVER_ADDRESS: conn_kwargs["host"], |
| 84 | + redisx.DB: conn_kwargs.get("db") or 0, |
| 85 | + } |
| 86 | + client_name = conn_kwargs.get("client_name") |
| 87 | + if client_name: |
| 88 | + conn_tags[redisx.CLIENT_NAME] = client_name |
| 89 | + return conn_tags |
| 90 | + except Exception: |
| 91 | + return {} |
| 92 | + |
| 93 | + |
| 94 | +def _build_tags(query, pin, instance, integration_name): |
| 95 | + ret = dict() |
| 96 | + ret[SPAN_KIND] = SpanKind.CLIENT |
| 97 | + ret[COMPONENT] = integration_name |
| 98 | + ret[db.SYSTEM] = redisx.APP |
| 99 | + if query is not None: |
| 100 | + span_name = schematize_cache_operation(redisx.RAWCMD, cache_provider=redisx.APP) # type: ignore[operator] |
| 101 | + ret[span_name] = query |
| 102 | + if pin.tags: |
| 103 | + # PERF: avoid Span.set_tag to avoid unnecessary checks |
| 104 | + for key, value in pin.tags.items(): |
| 105 | + ret[key] = value |
| 106 | + # some redis clients do not have a connection_pool attribute (ex. aioredis v1.3) |
| 107 | + if hasattr(instance, "connection_pool"): |
| 108 | + for key, value in _extract_conn_tags(instance.connection_pool.connection_kwargs).items(): |
| 109 | + ret[key] = value |
| 110 | + return ret |
| 111 | + |
| 112 | + |
| 113 | +@contextmanager |
| 114 | +def _instrument_redis_execute_pipeline(pin, config_integration, cmds, instance): |
| 115 | + cmd_string = resource = "\n".join(cmds) |
| 116 | + if config_integration.resource_only_command: |
| 117 | + resource = "\n".join([cmd.split(" ")[0] for cmd in cmds]) |
| 118 | + |
| 119 | + with core.context_with_data( |
| 120 | + "redis.execute_pipeline", |
| 121 | + span_name=schematize_cache_operation(redisx.CMD, cache_provider=redisx.APP), |
| 122 | + resource=resource, |
| 123 | + service=trace_utils.ext_service(pin, config_integration), |
| 124 | + span_type=SpanTypes.REDIS, |
| 125 | + pin=pin, |
| 126 | + measured=True, |
| 127 | + tags=_build_tags(cmd_string, pin, instance, config_integration.integration_name), |
| 128 | + ) as ctx: |
| 129 | + core.dispatch("redis.execute_pipeline", [ctx, pin, config_integration, None, instance, cmd_string]) |
| 130 | + yield ctx.span |
| 131 | + |
| 132 | + |
| 133 | +@contextmanager |
| 134 | +def _instrument_redis_cmd(pin, config_integration, instance, args): |
| 135 | + query = stringify_cache_args(args, cmd_max_len=config_integration.cmd_max_length) |
| 136 | + with core.context_with_data( |
| 137 | + "redis.command", |
| 138 | + span_name=schematize_cache_operation(redisx.CMD, cache_provider=redisx.APP), |
| 139 | + pin=pin, |
| 140 | + service=trace_utils.ext_service(pin, config_integration), |
| 141 | + span_type=SpanTypes.REDIS, |
| 142 | + resource=query.split(" ")[0] if config_integration.resource_only_command else query, |
| 143 | + measured=True, |
| 144 | + tags=_build_tags(query, pin, instance, config_integration.integration_name), |
| 145 | + ) as ctx: |
| 146 | + core.dispatch("redis.execute_pipeline", [ctx, pin, config_integration, args, instance, query]) |
| 147 | + yield ctx |
0 commit comments