Skip to content

Commit

Permalink
chore(trace_utils): move implementation details to internal
Browse files Browse the repository at this point in the history
  • Loading branch information
mabdinur committed Jan 17, 2025
1 parent 3aa6921 commit 9029ce2
Show file tree
Hide file tree
Showing 22 changed files with 922 additions and 842 deletions.
2 changes: 1 addition & 1 deletion ddtrace/_trace/trace_handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
from ddtrace.constants import SPAN_MEASURED_KEY
from ddtrace.contrib import trace_utils
from ddtrace.contrib.internal.botocore.constants import BOTOCORE_STEPFUNCTIONS_INPUT_KEY
from ddtrace.contrib.trace_utils import _set_url_tag
from ddtrace.contrib.internal.trace_utils import _set_url_tag
from ddtrace.ext import SpanKind
from ddtrace.ext import db
from ddtrace.ext import http
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/_trace/utils_redis.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from ddtrace.constants import SPAN_KIND
from ddtrace.constants import SPAN_MEASURED_KEY
from ddtrace.contrib import trace_utils
from ddtrace.contrib.redis_utils import _extract_conn_tags
from ddtrace.contrib.internal.redis_utils import _extract_conn_tags
from ddtrace.ext import SpanKind
from ddtrace.ext import SpanTypes
from ddtrace.ext import db
Expand Down
4 changes: 2 additions & 2 deletions ddtrace/appsec/_handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
from ddtrace.appsec._asm_request_context import get_blocked
from ddtrace.appsec._constants import SPAN_DATA_NAMES
from ddtrace.contrib import trace_utils
from ddtrace.contrib.trace_utils import _get_request_header_user_agent
from ddtrace.contrib.trace_utils import _set_url_tag
from ddtrace.contrib.internal.trace_utils import _get_request_header_user_agent
from ddtrace.contrib.internal.trace_utils import _set_url_tag
from ddtrace.ext import SpanTypes
from ddtrace.ext import http
from ddtrace.internal import core
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/appsec/_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ def get_rules() -> str:


def _set_headers(span: Span, headers: Any, kind: str, only_asm_enabled: bool = False) -> None:
from ddtrace.contrib.trace_utils import _normalize_tag_name
from ddtrace.contrib.internal.trace_utils import _normalize_tag_name

for k in headers:
if isinstance(k, tuple):
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/appsec/_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ def parse_response_body(raw_body):

from ddtrace.appsec import _asm_request_context
from ddtrace.appsec._constants import SPAN_DATA_NAMES
from ddtrace.contrib.trace_utils import _get_header_value_case_insensitive
from ddtrace.contrib.internal.trace_utils import _get_header_value_case_insensitive

if not raw_body:
return
Expand Down
15 changes: 15 additions & 0 deletions ddtrace/contrib/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,19 @@
from ddtrace._trace import trace_handlers # noqa:F401
from ddtrace.internal.utils.deprecations import DDTraceDeprecationWarning
from ddtrace.internal.utils.importlib import func_name # noqa:F401
from ddtrace.internal.utils.importlib import module_name # noqa:F401
from ddtrace.internal.utils.importlib import require_modules # noqa:F401
from ddtrace.vendor.debtcollector import deprecate


def __getattr__(name):
if name in ("trace_handlers", "func_name", "module_name", "require_modules"):
deprecate(
("%s.%s is deprecated" % (__name__, name)),
category=DDTraceDeprecationWarning,
removal_version="3.0.0",
)

if name in globals():
return globals()[name]
raise AttributeError("%s has no attribute %s", __name__, name)
2 changes: 1 addition & 1 deletion ddtrace/contrib/internal/aioredis/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from ddtrace.constants import SPAN_MEASURED_KEY
from ddtrace.contrib import trace_utils
from ddtrace.contrib.redis_utils import ROW_RETURNING_COMMANDS
from ddtrace.contrib.redis_utils import _run_redis_command_async
from ddtrace.contrib.internal.redis_utils import _run_redis_command_async
from ddtrace.contrib.redis_utils import determine_row_count
from ddtrace.ext import SpanKind
from ddtrace.ext import SpanTypes
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/contrib/internal/aredis/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from ddtrace import config
from ddtrace._trace.utils_redis import _instrument_redis_cmd
from ddtrace._trace.utils_redis import _instrument_redis_execute_pipeline
from ddtrace.contrib.redis_utils import _run_redis_command_async
from ddtrace.contrib.internal.redis_utils import _run_redis_command_async
from ddtrace.internal.schema import schematize_service_name
from ddtrace.internal.utils.formats import CMD_MAX_LEN
from ddtrace.internal.utils.formats import asbool
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/contrib/internal/django/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
from ddtrace.constants import SPAN_KIND
from ddtrace.contrib import dbapi
from ddtrace.contrib import trace_utils
from ddtrace.contrib.trace_utils import _get_request_header_user_agent
from ddtrace.contrib.internal.trace_utils import _get_request_header_user_agent
from ddtrace.ext import SpanKind
from ddtrace.ext import SpanTypes
from ddtrace.ext import db
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/contrib/internal/mysql/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
from ddtrace.appsec._iast.constants import VULN_SQL_INJECTION
from ddtrace.contrib.dbapi import TracedConnection
from ddtrace.contrib.trace_utils import _convert_to_string
from ddtrace.contrib.internal.trace_utils import _convert_to_string
from ddtrace.ext import db
from ddtrace.ext import net
from ddtrace.internal.schema import schematize_database_operation
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/contrib/internal/mysqldb/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from ddtrace.constants import SPAN_KIND
from ddtrace.constants import SPAN_MEASURED_KEY
from ddtrace.contrib.dbapi import TracedConnection
from ddtrace.contrib.trace_utils import _convert_to_string
from ddtrace.contrib.internal.trace_utils import _convert_to_string
from ddtrace.contrib.trace_utils import ext_service
from ddtrace.ext import SpanKind
from ddtrace.ext import SpanTypes
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/contrib/internal/pymysql/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

from ddtrace import config
from ddtrace.contrib.dbapi import TracedConnection
from ddtrace.contrib.trace_utils import _convert_to_string
from ddtrace.contrib.internal.trace_utils import _convert_to_string
from ddtrace.ext import db
from ddtrace.ext import net
from ddtrace.internal.schema import schematize_database_operation
Expand Down
2 changes: 1 addition & 1 deletion ddtrace/contrib/internal/redis/asyncio_patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from ddtrace._trace.utils_redis import _instrument_redis_cmd
from ddtrace._trace.utils_redis import _instrument_redis_execute_async_cluster_pipeline
from ddtrace._trace.utils_redis import _instrument_redis_execute_pipeline
from ddtrace.contrib.redis_utils import _run_redis_command_async
from ddtrace.contrib.internal.redis_utils import _run_redis_command_async
from ddtrace.internal.utils.formats import stringify_cache_args
from ddtrace.trace import Pin

Expand Down
84 changes: 84 additions & 0 deletions ddtrace/contrib/internal/redis_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
from typing import Dict
from typing import List
from typing import Optional
from typing import Union

from ddtrace.ext import net
from ddtrace.ext import redis as redisx
from ddtrace.internal import core
from ddtrace.internal.utils.formats import stringify_cache_args


SINGLE_KEY_COMMANDS = [
"GET",
"GETDEL",
"GETEX",
"GETRANGE",
"GETSET",
"LINDEX",
"LRANGE",
"RPOP",
"LPOP",
"HGET",
"HGETALL",
"HKEYS",
"HMGET",
"HRANDFIELD",
"HVALS",
]
MULTI_KEY_COMMANDS = ["MGET"]
ROW_RETURNING_COMMANDS = SINGLE_KEY_COMMANDS + MULTI_KEY_COMMANDS


def _extract_conn_tags(conn_kwargs):
"""Transform redis conn info into dogtrace metas"""
try:
conn_tags = {
net.TARGET_HOST: conn_kwargs["host"],
net.TARGET_PORT: conn_kwargs["port"],
net.SERVER_ADDRESS: conn_kwargs["host"],
redisx.DB: conn_kwargs.get("db") or 0,
}
client_name = conn_kwargs.get("client_name")
if client_name:
conn_tags[redisx.CLIENT_NAME] = client_name
return conn_tags
except Exception:
return {}


def determine_row_count(redis_command: str, result: Optional[Union[List, Dict, str]]) -> int:
empty_results = [b"", [], {}, None]
# result can be an empty list / dict / string
if result not in empty_results:
if redis_command == "MGET":
# only include valid key results within count
result = [x for x in result if x not in empty_results]
return len(result)
elif redis_command == "HMGET":
# only include valid key results within count
result = [x for x in result if x not in empty_results]
return 1 if len(result) > 0 else 0
else:
return 1
else:
return 0


async def _run_redis_command_async(ctx: core.ExecutionContext, func, args, kwargs):
parsed_command = stringify_cache_args(args)
redis_command = parsed_command.split(" ")[0]
rowcount = None
result = None
try:
result = await func(*args, **kwargs)
return result
except BaseException:
rowcount = 0
raise
finally:
if rowcount is None:
rowcount = determine_row_count(redis_command=redis_command, result=result)
if redis_command not in ROW_RETURNING_COMMANDS:
rowcount = None
core.dispatch("redis.async_command.post", [ctx, rowcount])
Loading

0 comments on commit 9029ce2

Please sign in to comment.