Add ability to enable/disable the SDK (#26)

Revision 3

Changes:
	- Removed enable flag from recorder. Now it checks the
	global SDK configuration to determine if the SDK is disabled.
        Global SDK configuration no longer modifies the recorder.
	- Tests modified to reflect change.

Author: chanchiem

aws_xray_sdk/core/recorder.py

@@ -5,6 +5,7 @@
 import platform
 import time
 
+import aws_xray_sdk
 from aws_xray_sdk.version import VERSION
 from .models.segment import Segment, SegmentContextManager
 from .models.subsegment import Subsegment, SubsegmentContextManager
@@ -65,7 +66,6 @@ def __init__(self):
             self._context = Context()
             self._sampler = DefaultSampler()
 
-        self._enabled = True
         self._emitter = UDPEmitter()
         self._sampling = True
         self._max_trace_back = 10
@@ -79,7 +79,7 @@ def __init__(self):
         if type(self.sampler).__name__ == 'DefaultSampler':
             self.sampler.load_settings(DaemonConfig(), self.context)
 
-    def configure(self, enabled=None, sampling=None, plugins=None,
+    def configure(self, sampling=None, plugins=None,
                   context_missing=None, sampling_rules=None,
                   daemon_address=None, service=None,
                   context=None, emitter=None, streaming=None,
@@ -90,16 +90,6 @@ def configure(self, enabled=None, sampling=None, plugins=None,
 
         Configure needs to run before patching thrid party libraries
         to avoid creating dangling subsegment.
-        :param bool enabled: If not enabled, the recorder automatically
-            generates DummySegments for every segment creation, whether
-            through patched extensions nor middlewares, and thus
-            not send any Segments out to the Daemon. May be set through an
-            environmental variable, where the environmental variable will
-            always take precedence over hardcoded configurations. The environment
-            variable is set as a case-insensitive string boolean. If the environment
-            variable exists but is an invalid string boolean, this enabled flag
-            will automatically be set to true. If no enabled flag is given and the
-            env variable is not set, then it will also default to being enabled.
         :param bool sampling: If sampling is enabled, every time the recorder
             creates a segment it decides whether to send this segment to
             the X-Ray daemon. This setting is not used if the recorder
@@ -145,9 +135,9 @@ class to have your own implementation of the streaming process.
             by auto-capture. Lower this if a single document becomes too large.
         :param bool stream_sql: Whether SQL query texts should be streamed.
 
-        Environment variables AWS_XRAY_ENABLED, AWS_XRAY_DAEMON_ADDRESS, AWS_XRAY_CONTEXT_MISSING
+        Environment variables AWS_XRAY_DAEMON_ADDRESS, AWS_XRAY_CONTEXT_MISSING
         and AWS_XRAY_TRACING_NAME respectively overrides arguments
-        enabled, daemon_address, context_missing and service.
+        daemon_address, context_missing and service.
         """
 
         if sampling is not None:
@@ -176,8 +166,6 @@ class to have your own implementation of the streaming process.
             self.max_trace_back = max_trace_back
         if stream_sql is not None:
             self.stream_sql = stream_sql
-        if enabled is not None:
-            self.enabled = enabled
 
         if plugins:
             plugin_modules = get_plugin_modules(plugins)
@@ -236,7 +224,7 @@ def begin_segment(self, name=None, traceid=None,
         # To disable the recorder, we set the sampling decision to always be false.
         # This way, when segments are generated, they become dummy segments and are ultimately never sent.
         # The call to self._sampler.should_trace() is never called either so the poller threads are never started.
-        if not self.enabled:
+        if not aws_xray_sdk.global_sdk_config.sdk_enabled():
             sampling = 0
 
         # we respect the input sampling decision
@@ -420,7 +408,7 @@ def record_subsegment(self, wrapped, instance, args, kwargs, name,
         # In the case when the SDK is disabled, we ensure that a parent segment exists, because this is usually
         # handled by the middleware. We generate a dummy segment as the parent segment if one doesn't exist.
         # This is to allow potential segment method calls to not throw exceptions in the captured method.
-        if not self.enabled:
+        if not aws_xray_sdk.global_sdk_config.sdk_enabled():
             try:
                 self.current_segment()
             except SegmentNotFoundException:

aws_xray_sdk/sdk_config.py

@@ -1,5 +1,4 @@
 import os
-import aws_xray_sdk.core
 
 
 class InvalidParameterTypeException(Exception):
@@ -12,8 +11,24 @@ class InvalidParameterTypeException(Exception):
 class SDKConfig(object):
     """
     Global Configuration Class that defines SDK-level configuration properties.
-    It is recommended to only use the recorder to set this configuration's enabled
-    flag to maintain thread safety.
+
+    Enabling/Disabling the SDK:
+        By default, the SDK is enabled unless if an environment variable AWS_XRAY_SDK_ENABLED
+            is set. If it is set, it needs to be a valid string boolean, otherwise, it will default
+            to true. If the environment variable is set, all calls to set_sdk_enabled() will
+            prioritize the value of the environment variable.
+        Disabling the SDK affects the recorder, patcher, and middlewares in the following ways:
+        For the recorder, disabling automatically generates DummySegments for subsequent segments
+            and DummySubsegments for subsegments created and thus not send any traces to the daemon.
+        For the patcher, module patching will automatically be disabled. The SDK must be disabled
+            before calling patcher.patch() method in order for this to function properly.
+        For the middleware, no modification is made on them, but since the recorder automatically
+            generates DummySegments for all subsequent calls, they will not generate segments/subsegments
+            to be sent.
+
+    Environment variables:
+        "AWS_XRAY_SDK_ENABLED" - If set to 'false' disables the SDK and causes the explained above
+            to occur.
     """
     XRAY_ENABLED_KEY = 'AWS_XRAY_SDK_ENABLED'
     __SDK_ENABLED = str(os.getenv(XRAY_ENABLED_KEY, 'true')).lower() != 'false'
@@ -28,13 +43,13 @@ def sdk_enabled(cls):
     @classmethod
     def set_sdk_enabled(cls, value):
         """
-        Modifies the enabled flag if the "AWS_XRAY_ENABLED" environment variable is not set,
+        Modifies the enabled flag if the "AWS_XRAY_SDK_ENABLED" environment variable is not set,
         otherwise, set the enabled flag to be equal to the environment variable. If the
         env variable is an invalid string boolean, it will default to true.
 
         :param bool value: Flag to set whether the SDK is enabled or disabled.
 
-        Environment variables AWS_XRAY_ENABLED overrides argument value.
+        Environment variables AWS_XRAY_SDK_ENABLED overrides argument value.
         """
         # Environment Variables take precedence over hardcoded configurations.
         if cls.XRAY_ENABLED_KEY in os.environ:
@@ -47,6 +62,3 @@ def set_sdk_enabled(cls, value):
                 raise InvalidParameterTypeException(
                     "Invalid parameter type passed into set_sdk_enabled(). Defaulting to True..."
                 )
-
-        # Modify all key paths.
-        aws_xray_sdk.core.xray_recorder.configure(enabled=cls.__SDK_ENABLED)

tests/ext/aiohttp/test_middleware.py

@@ -109,9 +109,8 @@ def recorder(loop):
     patcher.start()
 
     xray_recorder.clear_trace_entities()
-    sdk_enabled_state = aws_xray_sdk.global_sdk_config.sdk_enabled()
     yield xray_recorder
-    aws_xray_sdk.global_sdk_config.set_sdk_enabled(sdk_enabled_state)
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
     xray_recorder.clear_trace_entities()
     patcher.stop()
 

tests/ext/flask/test_flask.py

@@ -1,6 +1,7 @@
 import pytest
 from flask import Flask, render_template_string
 
+import aws_xray_sdk
 from aws_xray_sdk.ext.flask.middleware import XRayMiddleware
 from aws_xray_sdk.core.context import Context
 from aws_xray_sdk.core.models import http
@@ -51,6 +52,7 @@ def cleanup():
     recorder.clear_trace_entities()
     yield
     recorder.clear_trace_entities()
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
 
 
 def test_ok():
@@ -146,7 +148,7 @@ def test_sampled_response_header():
 
 
 def test_disabled_sdk():
-    recorder.configure(enabled=False)
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)
     path = '/ok'
     app.get(path)
     segment = recorder.emitter.pop()

tests/test_lambda_context.py

@@ -1,9 +1,9 @@
 import os
 
 import aws_xray_sdk
+import pytest
 from aws_xray_sdk.core import lambda_launcher
 from aws_xray_sdk.core.models.subsegment import Subsegment
-from aws_xray_sdk.core.models.dummy_entities import DummySubsegment
 
 
 TRACE_ID = '1-5759e988-bd862e3fe1be46a994272793'
@@ -12,7 +12,12 @@
 
 os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] = HEADER_VAR
 context = lambda_launcher.LambdaContext()
-aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
+
+
+@pytest.fixture(autouse=True)
+def setup():
+    yield
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
 
 
 def test_facade_segment_generation():

tests/test_patcher.py

@@ -35,7 +35,6 @@ def construct_ctx():
     """
     pre_run_modules = set(module for module in sys.modules.keys())
 
-    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
     xray_recorder.configure(service='test', sampling=False, context=Context())
     xray_recorder.clear_trace_entities()
     xray_recorder.begin_segment('name')

tests/test_recorder.py

@@ -23,6 +23,7 @@ def construct_ctx():
     xray_recorder.clear_trace_entities()
     yield
     xray_recorder.clear_trace_entities()
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
 
 
 def test_default_runtime_context():
@@ -175,14 +176,15 @@ def test_in_segment_exception():
 
 
 def test_default_enabled():
+    assert aws_xray_sdk.global_sdk_config.sdk_enabled()
     segment = xray_recorder.begin_segment('name')
     subsegment = xray_recorder.begin_subsegment('name')
     assert type(xray_recorder.current_segment()) is Segment
     assert type(xray_recorder.current_subsegment()) is Subsegment
 
 
 def test_disable_is_dummy():
-    xray_recorder.configure(enabled=False)
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)
     segment = xray_recorder.begin_segment('name')
     subsegment = xray_recorder.begin_subsegment('name')
     assert type(xray_recorder.current_segment()) is DummySegment

tests/test_sdk_config.py

@@ -16,6 +16,7 @@ def cleanup():
     yield
     if XRAY_ENABLED_KEY in os.environ:
         del os.environ[XRAY_ENABLED_KEY]
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
 
 
 def test_enable_key():
@@ -28,6 +29,7 @@ def test_default_enabled():
 
 def test_env_var_precedence():
     os.environ[XRAY_ENABLED_KEY] = "true"
+    # Env Variable takes precedence. This is called to activate the internal check
     aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)
     assert aws_xray_sdk.global_sdk_config.sdk_enabled() is True
     os.environ[XRAY_ENABLED_KEY] = "false"
@@ -46,7 +48,8 @@ def test_env_var_precedence():
 
 def test_env_enable_case():
     os.environ[XRAY_ENABLED_KEY] = "TrUE"
-    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)  # Env Variable takes precedence. This is called to activate the internal check
+    # Env Variable takes precedence. This is called to activate the internal check
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
     assert aws_xray_sdk.global_sdk_config.sdk_enabled() is True
 
     os.environ[XRAY_ENABLED_KEY] = "true"
@@ -64,7 +67,8 @@ def test_env_enable_case():
 
 def test_invalid_env_string():
     os.environ[XRAY_ENABLED_KEY] = "INVALID"
-    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)  # Env Variable takes precedence. This is called to activate the internal check
+    # Env Variable takes precedence. This is called to activate the internal check
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
     assert aws_xray_sdk.global_sdk_config.sdk_enabled() is True
 
     os.environ[XRAY_ENABLED_KEY] = "1.0"