Merge branch 'avara1986/fix_circular_imports' into avara1986/APPSEC-56946_cmdi_secure_mark_check

Author: avara1986

ddtrace/appsec/__init__.py

@@ -4,7 +4,6 @@
 
 
 _APPSEC_TO_BE_LOADED = True
-_IAST_TO_BE_LOADED = True
 
 
 def load_appsec():
@@ -17,16 +16,6 @@ def load_appsec():
         _APPSEC_TO_BE_LOADED = False
 
 
-def load_iast():
-    """Lazily load the iast module listeners."""
-    from ddtrace.appsec._iast._listener import iast_listen
-
-    global _IAST_TO_BE_LOADED
-    if _IAST_TO_BE_LOADED:
-        iast_listen()
-        _IAST_TO_BE_LOADED = False
-
-
 def load_common_appsec_modules():
     """Lazily load the common module patches."""
     from ddtrace.settings.asm import config as asm_config

ddtrace/appsec/_iast/__init__.py

@@ -37,11 +37,15 @@ def wrapped_function(wrapped, instance, args, kwargs):
 from ddtrace.internal.module import ModuleWatchdog
 from ddtrace.settings.asm import config as asm_config
 
+from ._listener import iast_listen
 from ._overhead_control_engine import OverheadControl
 
 
 log = get_logger(__name__)
 
+_IAST_TO_BE_LOADED = True
+_iast_propagation_enabled = False
+
 oce = OverheadControl()
 
 
@@ -79,9 +83,6 @@ def ddtrace_iast_flask_patch():
     exec(compiled_code, module.__dict__)  # nosec B102
 
 
-_iast_propagation_enabled = False
-
-
 def enable_iast_propagation():
     """Add IAST AST patching in the ModuleWatchdog"""
     # DEV: These imports are here to avoid _ast.ast_patching import in the top level
@@ -141,3 +142,11 @@ def disable_iast_propagation():
     "enable_iast_propagation",
     "disable_iast_propagation",
 ]
+
+
+def load_iast():
+    """Lazily load the iast module listeners."""
+    global _IAST_TO_BE_LOADED
+    if _IAST_TO_BE_LOADED:
+        iast_listen()
+        _IAST_TO_BE_LOADED = False

ddtrace/appsec/_iast/processor.py

@@ -15,7 +15,7 @@
 @dataclass(eq=False)
 class AppSecIastSpanProcessor(SpanProcessor):
     def __post_init__(self) -> None:
-        from ddtrace.appsec import load_iast
+        from . import load_iast
 
         load_iast()
 

docs/configuration.rst

@@ -269,7 +269,7 @@ Traces
 
    DD_TRACE_PROPAGATION_STYLE:
      default: |
-         ``datadog,tracecontext``
+         ``datadog,tracecontext,baggage``
 
      description: |
          Comma separated list of propagation styles used for extracting trace context from inbound request headers and injecting trace context into outbound request headers.
@@ -278,7 +278,7 @@ Traces
 
          Overridden by ``DD_TRACE_PROPAGATION_STYLE_INJECT`` for injection.
 
-         The supported values are ``datadog``, ``b3multi``, ``tracecontext``, and ``none``.
+         The supported values are ``datadog``, ``b3multi``, ``tracecontext``, ``baggage``, and ``none``.
 
          When checking inbound request headers we will take the first valid trace context in the order provided.
          When ``none`` is the only propagator listed, propagation is disabled.
@@ -293,6 +293,7 @@ Traces
        v1.7.0: The ``b3multi`` propagation style was added and ``b3`` was deprecated in favor it.
        v1.7.0: Added support for ``tracecontext`` W3C headers. Changed the default value to ``DD_TRACE_PROPAGATION_STYLE="tracecontext,datadog"``.
        v2.6.0: Updated default value to ``datadog,tracecontext``.
+       v2.16.0: Updated default value to ``datadog,tracecontex,baggage``.
 
    DD_TRACE_SPAN_TRACEBACK_MAX_SIZE:
       type: Integer

tests/contrib/aws_lambda/test_aws_lambda.py

@@ -1,3 +1,5 @@
+import contextlib
+
 import pytest
 
 from ddtrace.contrib.internal.aws_lambda.patch import patch
@@ -11,7 +13,6 @@
 from tests.contrib.aws_lambda.handlers import manually_wrapped_handler
 from tests.contrib.aws_lambda.handlers import static_handler
 from tests.contrib.aws_lambda.handlers import timeout_handler
-from tests.utils import flaky
 from tests.utils import override_env
 
 
@@ -46,13 +47,15 @@ def get_env(env=None):
     return {**common_env, **env}
 
 
-@pytest.fixture(autouse=True)
-def setup():
-    yield
-    unpatch()
+@contextlib.contextmanager
+def override_env_and_patch(env):
+    # patching and unpatching must be done while the environment is set
+    with override_env(env):
+        patch()
+        yield
+        unpatch()
 
 
-@flaky(1744053478)
 @pytest.mark.parametrize("customApmFlushDeadline", [("-100"), ("10"), ("100"), ("200")])
 @pytest.mark.snapshot
 def test_timeout_traces(context, customApmFlushDeadline):
@@ -64,9 +67,7 @@ def test_timeout_traces(context, customApmFlushDeadline):
         }
     )
 
-    with override_env(env):
-        patch()
-
+    with override_env_and_patch(env):
         datadog(timeout_handler)({}, context())
 
 
@@ -83,9 +84,7 @@ def test_continue_on_early_trace_ending(context):
         }
     )
 
-    with override_env(env):
-        patch()
-
+    with override_env_and_patch(env):
         datadog(finishing_spans_early_handler)({}, context())
 
 
@@ -98,11 +97,8 @@ async def test_file_patching(context):
         }
     )
 
-    with override_env(env):
-        patch()
-
+    with override_env_and_patch(env):
         result = datadog(handler)({}, context())
-
         assert result == {"success": True}
 
 
@@ -117,11 +113,8 @@ async def test_module_patching(mocker, context):
         }
     )
 
-    with override_env(env):
-        patch()
-
+    with override_env_and_patch(env):
         result = manually_wrapped_handler({}, context())
-
         assert result == {"success": True}
 
 
@@ -135,7 +128,6 @@ async def test_module_patching(mocker, context):
     ],
 )
 @pytest.mark.snapshot
-@flaky(1744053478, reason="Did not receive expected traces: 'aws.lambda' for [handler3-instance_handler_with_code]")
 def test_class_based_handlers(context, handler, function_name):
     env = get_env(
         {
@@ -144,8 +136,6 @@ def test_class_based_handlers(context, handler, function_name):
         }
     )
 
-    with override_env(env):
-        patch()
-
+    with override_env_and_patch(env):
         result = datadog(handler)({}, context())
         assert result == {"success": True}