Skip to content

Commit

Permalink
Merge branch 'main' into cbeauchesne/crash-investigation
Browse files Browse the repository at this point in the history
  • Loading branch information
juanjux authored Jul 25, 2024
2 parents 42a0a3d + 584e00b commit 0fffe4f
Show file tree
Hide file tree
Showing 4 changed files with 106 additions and 16 deletions.
82 changes: 72 additions & 10 deletions ddtrace/contrib/botocore/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,11 @@

from ddtrace import config
from ddtrace.contrib.trace_utils import with_traced_module
from ddtrace.internal.utils.deprecations import DDTraceDeprecationWarning
from ddtrace.llmobs._integrations import BedrockIntegration
from ddtrace.settings.config import Config
from ddtrace.vendor import wrapt
from ddtrace.vendor.debtcollector import deprecate

from ...constants import SPAN_KIND
from ...ext import SpanKind
Expand Down Expand Up @@ -77,20 +79,30 @@
)


def get_version():
def _get_version():
# type: () -> str
return __version__


def get_version():
deprecate(
"get_version is deprecated",
message="get_version is deprecated",
removal_version="3.0.0",
category=DDTraceDeprecationWarning,
)
return _get_version()


def patch():
if getattr(botocore.client, "_datadog_patch", False):
return
botocore.client._datadog_patch = True

botocore._datadog_integration = BedrockIntegration(integration_config=config.botocore)
wrapt.wrap_function_wrapper("botocore.client", "BaseClient._make_api_call", patched_api_call(botocore))
wrapt.wrap_function_wrapper("botocore.client", "BaseClient._make_api_call", _patched_api_call(botocore))
Pin().onto(botocore.client.BaseClient)
wrapt.wrap_function_wrapper("botocore.parsers", "ResponseParser.parse", patched_lib_fn)
wrapt.wrap_function_wrapper("botocore.parsers", "ResponseParser.parse", _patched_lib_fn)
Pin().onto(botocore.parsers.ResponseParser)
_PATCHED_SUBMODULES.clear()

Expand All @@ -103,7 +115,7 @@ def unpatch():
unwrap(botocore.client.BaseClient, "_make_api_call")


def patch_submodules(submodules):
def _patch_submodules(submodules):
# type: (Union[List[str], bool]) -> None
if isinstance(submodules, bool) and submodules:
_PATCHED_SUBMODULES.clear()
Expand All @@ -112,7 +124,17 @@ def patch_submodules(submodules):
_PATCHED_SUBMODULES.update(submodules)


def patched_lib_fn(original_func, instance, args, kwargs):
def patch_submodules(submodules):
deprecate(
"patch_submodules is deprecated",
message="patch_submodules is deprecated",
removal_version="3.0.0",
category=DDTraceDeprecationWarning,
)
return _patch_submodules(submodules)


def _patched_lib_fn(original_func, instance, args, kwargs):
pin = Pin.get_from(instance)
if not pin or not pin.enabled() or not config.botocore["instrument_internals"]:
return original_func(*args, **kwargs)
Expand All @@ -124,8 +146,18 @@ def patched_lib_fn(original_func, instance, args, kwargs):
return original_func(*args, **kwargs)


def patched_lib_fn(original_func, instance, args, kwargs):
deprecate(
"patched_lib_fn is deprecated",
message="patched_lib_fn is deprecated",
removal_version="3.0.0",
category=DDTraceDeprecationWarning,
)
return _patched_lib_fn(original_func, instance, args, kwargs)


@with_traced_module
def patched_api_call(botocore, pin, original_func, instance, args, kwargs):
def _patched_api_call(botocore, pin, original_func, instance, args, kwargs):
if not pin or not pin.enabled():
return original_func(*args, **kwargs)

Expand Down Expand Up @@ -153,7 +185,7 @@ def patched_api_call(botocore, pin, original_func, instance, args, kwargs):
if endpoint_name == "bedrock-runtime" and operation.startswith("InvokeModel"):
patching_fn = patched_bedrock_api_call
else:
patching_fn = PATCHING_FUNCTIONS.get(endpoint_name, patched_api_call_fallback)
patching_fn = PATCHING_FUNCTIONS.get(endpoint_name, _patched_api_call_fallback)

return patching_fn(
original_func=original_func,
Expand All @@ -164,7 +196,17 @@ def patched_api_call(botocore, pin, original_func, instance, args, kwargs):
)


def prep_context_injection(ctx, endpoint_name, operation, trace_operation, params):
def patched_api_call(botocore, pin, original_func, instance, args, kwargs):
deprecate(
"patched_api_call is deprecated",
message="patched_api_call is deprecated",
removal_version="3.0.0",
category=DDTraceDeprecationWarning,
)
return _patched_api_call(botocore, pin, original_func, instance, args, kwargs)


def _prep_context_injection(ctx, endpoint_name, operation, trace_operation, params):
cloud_service = None
injection_function = None
schematization_function = schematize_cloud_messaging_operation
Expand All @@ -189,7 +231,17 @@ def prep_context_injection(ctx, endpoint_name, operation, trace_operation, param
)


def patched_api_call_fallback(original_func, instance, args, kwargs, function_vars):
def prep_context_injection(ctx, endpoint_name, operation, trace_operation, params):
deprecate(
"prep_context_injection is deprecated",
message="prep_context_injection is deprecated",
removal_version="3.0.0",
category=DDTraceDeprecationWarning,
)
return _prep_context_injection(ctx, endpoint_name, operation, trace_operation, params)


def _patched_api_call_fallback(original_func, instance, args, kwargs, function_vars):
# default patched api call that is used generally for several services / operations
params = function_vars.get("params")
trace_operation = function_vars.get("trace_operation")
Expand All @@ -212,7 +264,7 @@ def patched_api_call_fallback(original_func, instance, args, kwargs, function_va
) as ctx, ctx.get_item("instrumented_api_call"):
core.dispatch("botocore.patched_api_call.started", [ctx])
if args and config.botocore["distributed_tracing"]:
prep_context_injection(ctx, endpoint_name, operation, trace_operation, params)
_prep_context_injection(ctx, endpoint_name, operation, trace_operation, params)

try:
result = original_func(*args, **kwargs)
Expand All @@ -230,3 +282,13 @@ def patched_api_call_fallback(original_func, instance, args, kwargs, function_va
else:
core.dispatch("botocore.patched_api_call.success", [ctx, result])
return result


def patched_api_call_fallback(original_func, instance, args, kwargs, function_vars):
deprecate(
"patched_api_call_fallback is deprecated",
message="patched_api_call_fallback is deprecated",
removal_version="3.0.0",
category=DDTraceDeprecationWarning,
)
return _patched_api_call_fallback(original_func, instance, args, kwargs, function_vars)
28 changes: 25 additions & 3 deletions ddtrace/contrib/consul/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
from ddtrace import config
from ddtrace.internal.constants import COMPONENT
from ddtrace.internal.schema.span_attribute_schema import SpanDirection
from ddtrace.internal.utils.deprecations import DDTraceDeprecationWarning
from ddtrace.vendor.debtcollector import deprecate
from ddtrace.vendor.wrapt import wrap_function_wrapper as _w

from ...constants import ANALYTICS_SAMPLE_RATE_KEY
Expand All @@ -22,11 +24,21 @@
_KV_FUNCS = ["put", "get", "delete"]


def get_version():
def _get_version():
# type: () -> str
return getattr(consul, "__version__", "")


def get_version():
deprecate(
"get_version is deprecated",
message="get_version is deprecated",
removal_version="3.0.0",
category=DDTraceDeprecationWarning,
)
return _get_version()


def patch():
if getattr(consul, "__datadog_patch", False):
return
Expand All @@ -36,7 +48,7 @@ def patch():
pin.onto(consul.Consul.KV)

for f_name in _KV_FUNCS:
_w("consul", "Consul.KV.%s" % f_name, wrap_function(f_name))
_w("consul", "Consul.KV.%s" % f_name, _wrap_function(f_name))


def unpatch():
Expand All @@ -48,7 +60,7 @@ def unpatch():
_u(consul.Consul.KV, f_name)


def wrap_function(name):
def _wrap_function(name):
def trace_func(wrapped, instance, args, kwargs):
pin = Pin.get_from(instance)
if not pin or not pin.enabled():
Expand Down Expand Up @@ -83,3 +95,13 @@ def trace_func(wrapped, instance, args, kwargs):
return wrapped(*args, **kwargs)

return trace_func


def wrap_function(name):
deprecate(
"wrap_function is deprecated",
message="get_version is deprecated",
removal_version="3.0.0",
category=DDTraceDeprecationWarning,
)
return _wrap_function(name)
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
---
deprecations:
- |
botocore: All methods in botocore/patch.py except ``patch()`` and ``unpatch()`` are deprecated and will be removed in version 3.0.0.
- |
consul: All methods in consul/patch.py except ``patch()`` and ``unpatch()`` are deprecated and will be removed in version 3.0.0.
6 changes: 3 additions & 3 deletions tests/contrib/botocore/test.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,8 @@
from ddtrace.constants import ERROR_MSG
from ddtrace.constants import ERROR_STACK
from ddtrace.constants import ERROR_TYPE
from ddtrace.contrib.botocore.patch import _patch_submodules
from ddtrace.contrib.botocore.patch import patch
from ddtrace.contrib.botocore.patch import patch_submodules
from ddtrace.contrib.botocore.patch import unpatch
from ddtrace.internal.compat import PYTHON_VERSION_INFO
from ddtrace.internal.datastreams.processor import PROPAGATION_KEY_BASE_64
Expand Down Expand Up @@ -76,7 +76,7 @@ class BotocoreTest(TracerTestCase):
@mock_sqs
def setUp(self):
patch()
patch_submodules(True)
_patch_submodules(True)

self.session = botocore.session.get_session()
self.session.set_credentials(access_key="access-key", secret_key="secret-key")
Expand All @@ -103,7 +103,7 @@ def tearDown(self):
@mock_ec2
@mock_s3
def test_patch_submodules(self):
patch_submodules(["s3"])
_patch_submodules(["s3"])
ec2 = self.session.create_client("ec2", region_name="us-west-2")
Pin(service=self.TEST_SERVICE, tracer=self.tracer).onto(ec2)

Expand Down

0 comments on commit 0fffe4f

Please sign in to comment.